Capsule Network Projectors are Equivariant and Invariant Learners

TL;DR

The paper addresses the challenge of learning representations that are simultaneously invariant and equivariant to viewpoint transformations in self-supervised learning. It introduces CapsIE, a Capsule Network–based projection that outputs invariant activations and equivariant poses, coupled with an entropy-based invariant objective and a learned predictor conditioned on known viewpoint transformations to enforce equivariance, formalized as for all $g\in G$, $f(\rho_X(g)\cdot x) = \rho_Y(g)\cdot f(x)$. CapsIE leverages a ResNet-18 backbone with a Capsule Network projector (using SRCaps) to produce compact, structured embeddings and reports state-of-the-art rotation-equivariance on 3DIEBench, while delivering competitive performance against supervised baselines and demonstrating robust generalization on Objaverse and MOVi-E. The work highlights the efficiency and generalization benefits of CapsNets for self-supervised, multi-task representations and suggests directions for scaling capsule capacity and exploring alternative routing schemes to further improve invariant performance.

Abstract

Learning invariant representations has been the long-standing approach to self-supervised learning. However, recently progress has been made in preserving equivariant properties in representations, yet do so with highly prescribed architectures. In this work, we propose an invariant-equivariant self-supervised architecture that employs Capsule Networks (CapsNets), which have been shown to capture equivariance with respect to novel viewpoints. We demonstrate that the use of CapsNets in equivariant self-supervised architectures achieves improved downstream performance on equivariant tasks with higher efficiency and fewer network parameters. To accommodate the architectural changes of CapsNets, we introduce a new objective function based on entropy minimisation. This approach, which we name CapsIE (Capsule Invariant Equivariant Network), achieves state-of-the-art performance on the equivariant rotation tasks on the 3DIEBench dataset compared to prior equivariant SSL methods, while performing competitively against supervised counterparts. Our results demonstrate the ability of CapsNets to learn complex and generalised representations for large-scale, multi-task datasets compared to previous CapsNet benchmarks. Code is available at https://github.com/AberdeenML/CapsIE.

Figures (5)

• Figure 1: Left: Schematic overview of the proposed CapsIE architecture. Representations are fed into a CapsNet projector, and the output embeddings $Z_{\text{act}}$ and $Z_{\text{pose}}$ correspond to invariant and equivariant embeddings, respectively. Right: Generalised view of a Capsule projection head. CNN feature maps are transformed via the primary capsules into poses $u_{i}$, represented by cylinders, and activations $a_{i}$, represented by circles. Poses are transformed to votes, which represent a lower-level capsule's prediction for each of the higher-level capsules. The routing process then determines how well these votes match the concept represented by the upper-level capsule, thereby creating the coupling coefficients. Coupling coefficients inform $u_{j}$ and $a_{j}$, the output of the capsule projector head.
• Figure 2: Visual depiction of the problem statement. Two images are represented by subscripts $0,1$ while view under transformation $g$ is given by $'$. Arrows represent the construction of embeddings from an encoder network. Top the invariance objective is to maximise the similarity between embeddings of views originating from the same image. Bottom the equivariance objective aim to learn the transformations $\rho_X(g)\cdot x$ applied to $x$.
• Figure 3: Simplified visual representation of CapsNet outputs. Vector of activations outputs the probability of each capsule being activated, whereas the pose matrix corresponds to the object pose in relation to the frame.
• Figure 4: Invariant and equivariant performance of encoder representations during training when varying the number of capsules.(left) Classification evaluation performance (top-1 %) and (right) Rotation prediction online evaluation performance ($R^2$) both learned encoder representations. This evaluation is co-optimised during training and hence differs from the downstream evaluation reported in the Table \ref{['tab:inv-benchmark']} and \ref{['tab:equi-benchmark']}. We observe that as capsule numbers increase so does performance; however, at low capsule numbers, later in training invariant objective takes precedence.
• Figure 5: The 3DIEBench dataset, one 3d model is used to create 50 different views in a synthetic environment, which are saved as images along with the latent values by which they are transformed.