Equivariant Representation Learning via Class-Pose Decomposition

Technical Summary: Equivariant Representation Learning via Class-Pose Decomposition

Problem Statement
Intelligent agents require rich representations that reflect the intrinsic structures of perceived data. While self-supervised approaches like auto-encoders and contrastive learning are general, they often fail to capture fundamental aspects hidden in data, such as the separation of intrinsic object classes from their poses. Recent work suggests that disentangling these factors requires specific biases or supervision. A fundamental geometric structure in many datasets is symmetry (e.g., rigid transformations acting on object poses or frame changes in mobile agent exploration). The challenge is to learn representations that are not only equivariant to these symmetries but also lossless, interpretable, and disentangled, preserving the complete geometry of the data.

Methodology
The authors propose a general framework that decomposes the latent space into two distinct components: an invariant factor representing intrinsic data classes and a symmetry component representing the pose.

1. Theoretical Foundation:
   The approach is grounded in group theory. The authors formalize symmetries as a group $G$ acting on the data space $X$. Under the assumption that the action is free (no non-identity element fixes a point), they rely on a theoretical result (Proposition 2.1) stating that there exists an equivariant isomorphism $X \simeq (X/G) \times G$. Here, $X/G$ represents the set of orbits (intrinsic classes), and $G$ represents the symmetry group (poses). This decomposition guarantees that the representation is lossless and that the geometry of the data is preserved.

2. Latent Space Decomposition:
   The latent space $Z$ is defined as the product $Z = E \times G$, where:

   • $E$ is an invariant set representing the classes (orbits).
   • $G$ is the symmetry group representing the poses.
     The group $G$ acts trivially on $E$ and via multiplication on itself.

3. Learning Objective:
   The model $\phi: X \to E \times G$ is trained using a dataset of triples $(x, g, y)$ where $y = g \cdot x$. The loss function encourages equivariance based on relative symmetry information:
   $$L(\theta; x, g, y) = d_E(\phi_E(y), \phi_E(x)) + d_G(\phi_G(y), g \cdot \phi_G(x))$$

   • The first term ($d_E$) enforces invariance, ensuring points in the same orbit map to the same class representation. To prevent "orbit collapse" and ensure injectivity (necessary for lossless representation), the authors employ an InfoNCE loss, encouraging latent points to spread apart.
   • The second term ($d_G$) enforces equivariance with respect to the group action on the pose component.

4. Implementation:
   To handle non-Euclidean groups (e.g., $SO(3)$), the authors parametrize the group $G$ using the exponential map from its Lie algebra $\mathfrak{g}$. The model outputs an element $v \in \mathfrak{g}$, which is mapped to $G$ via $e^v$. For $G=SO(3)$, this is computed efficiently using Rodrigues' formula.

Key Contributions

• Class-Pose Decomposition: A method for learning equivariant representations that explicitly separates intrinsic data classes from poses by decomposing the latent space.
• General Mathematical Formalism: A framework based on group theory that guarantees lossless, interpretable, and disentangled representations under mild assumptions (free group actions).
• Empirical Investigation: A comprehensive evaluation across datasets with various symmetries (translations, dilations, rotations) and applications to scene mapping.

Experimental Results
The authors evaluated their method on five datasets: Sprites, Shapes, Multi-Sprites, Chairs, and Apartments (simulating a mobile agent). They compared their approach against:

• MDP Homomorphisms (MDPH): Learns latent actions without prior knowledge of the group.
• Linear Models: Uses a linear latent space with matrix multiplication (non-free action).
• Equivariant Neural Rendering (ENR): Uses a discretized latent grid with approximate equivariance.

Findings:

• Geometry Preservation: Visualizations show that the proposed method correctly separates orbits in the invariant component $E$ and preserves the geometry of the group $G$ in the pose component. For the "Apartments" dataset, the model successfully extracted maps of multiple environments simultaneously.
• Performance: The method achieved near-perfect hit-rates on single-step predictions. Crucially, on longer trajectories (10 and 20 steps), the proposed method significantly outperformed baselines.
  • MDPH accumulated errors due to the lack of structural constraints in its learned latent action.
  • Linear models degraded performance because the action was not free (rotating a vector around itself has no effect), forcing information loss.
  • ENR suffered from interpolation errors due to latent discretization.
• Disentanglement: The framework naturally achieves disentanglement, allowing independent variation of symmetry factors by acting on the pose component.

Significance and Claims
The paper claims that by leveraging the abstract language of Lie group theory, their framework provides a general solution for equivariant representation learning applicable to arbitrary groups and data types. Unlike previous works that focus on specific scenarios (e.g., Euclidean representations or group convolutions) or learn unstructured latent dynamics, this approach guarantees isomorphic (lossless) representations. The authors emphasize that their method enables the recovery of intrinsic classes and poses, facilitating tasks like self-supervised localization and mapping for mobile agents.

Limitations and Future Work
The authors acknowledge that their formalism assumes the group of symmetries is known a priori and is not inferred from data. While viable for robotics where odometry provides symmetry information, this may be limiting in other domains. They suggest that learning the group structure from data is a potential, though less interpretable, direction for future research.