EquiCaps: Predictor-Free Pose-Aware Pre-Trained Capsule Networks

TL;DR

EquiCaps introduces a capsule-based, predictor-free approach for learning pose-aware representations that inherently handle transformations, outperforming prior methods in rotation prediction and demonstrating strong generalisation under complex geometric transformations.

Contribution

The paper presents EquiCaps, a novel capsule network architecture that eliminates the need for predictor modules to enforce equivariance, leveraging capsules' intrinsic pose-awareness for improved self-supervised learning.

Findings

Outperforms state-of-the-art equivariant methods on rotation prediction with R^2 of 0.78.

Maintains robust equivariance under combined geometric transformations.

Demonstrates superior generalisation capabilities in complex transformation scenarios.

Abstract

Learning self-supervised representations that are invariant and equivariant to transformations is crucial for advancing beyond traditional visual classification tasks. However, many methods rely on predictor architectures to encode equivariance, despite evidence that architectural choices, such as capsule networks, inherently excel at learning interpretable pose-aware representations. To explore this, we introduce EquiCaps (Equivariant Capsule Network), a capsule-based approach to pose-aware self-supervision that eliminates the need for a specialised predictor for enforcing equivariance. Instead, we leverage the intrinsic pose-awareness capabilities of capsules to improve performance in pose estimation tasks. To further challenge our assumptions, we increase task complexity via multi-geometric transformations to enable a more thorough evaluation of invariance and equivariance by introducing 3DIEBench-T, an extension of a 3D object-rendering benchmark dataset. Empirical results demonstrate that EquiCaps outperforms prior state-of-the-art equivariant methods on rotation prediction, achieving a supervised-level $R^2$ of 0.78 on the 3DIEBench rotation prediction benchmark and improving upon SIE and CapsIE by 0.05 and 0.04 $R^2$, respectively. Moreover, in contrast to non-capsule-based equivariant approaches, EquiCaps maintains robust equivariant performance under combined geometric transformations, underscoring its generalisation capabilities and the promise of predictor-free capsule architectures.