BRo-JEPA: Learning Modular Arithmetic in Latent Space

Technical Summary: BRo-JEPA – Learning Modular Arithmetic in Latent Space

Problem Statement
The paper investigates a fundamental question in deep learning: do neural networks truly learn abstract algebraic rules, or do they merely memorize training patterns? While current models excel at mathematical tasks within seen ranges, they often fail to extrapolate to unseen numerical ranges or operations. The authors posit that standard supervised models and existing Joint Embedding Predictive Architecture (JEPA) models rely on statistical correlations rather than underlying arithmetic logic. Specifically, when tasked with modular arithmetic (e.g., $y = (x + k) \mod 10$) using MNIST digits as states, standard models fit seen operations but fail to generalize to unseen operations (strict zero-shot generalization).

Methodology
The authors utilize the JEPA framework, where a model predicts the latent representation of a target digit given a context digit image and an arithmetic operation (action), rather than predicting raw pixel outputs. The study compares several architectural variations:

1. Standard Baselines: Supervised models (MLP and ResNet-18) using cross-entropy loss and additive compositional action embeddings. These embeddings represent operations as scalar multiples of a base vector (e.g., $op(k) = k \cdot v$) to enforce arithmetic structure in the embedding space.
2. JEPA with Additive Embeddings: A JEPA model using the same additive embeddings, augmented with a "Composition Consistency Loss." This auxiliary loss enforces that applying a composed operation (e.g., $+2$) in one step should yield the same latent result as applying primitive operations sequentially (e.g., $+1$ then $+1$).
3. BRo-JEPA (Block-Rotation JEPA): The core contribution is a novel predictor that replaces the standard unconstrained Multi-Layer Perceptron (MLP) with an operation-conditioned block-rotation predictor.
   • Mechanism: Inspired by Rotary Position Embeddings (RoPE), this predictor treats the latent space dynamics as a series of 2D planar rotations. For a latent vector $z$ and operation $k$, the prediction is $z_p = R_{block}(k \cdot \theta) \cdot z$.
   • Inductive Bias: This architecture explicitly imposes the circular structure of modulo-10 arithmetic onto the latent space. The rotation angles $\theta$ can be fixed (set to $2\pi/10$) or learnable.
   • Variants: The authors test Single-Frequency Rotation (SFR), where all latent blocks share one angle, and Multi-Frequency Rotation (MFR), where each block has an independent angle.
   • Regularization: To prevent "latent collapse" (where different inputs map to indistinguishable representations), the model employs Variance-Invariance-Covariance Regularization (VICReg) and an Exponential Moving Average (EMA) target encoder.

Key Contributions

1. Testbed Definition: The introduction of MNIST modular arithmetic as a controlled testbed for studying rule learning in latent world models.
2. Failure Analysis: Demonstration that supervised baselines and JEPA models with additive embeddings, despite fitting seen operations, fail at strict zero-shot extrapolation. They capture operation signs but fail to encode magnitude or compositional geometry.
3. BRo-JEPA Proposal: The development of a JEPA-style latent dynamics model featuring an operation-conditioned block-rotation predictor that aligns with the cyclic structure of modulo arithmetic.
4. Geometric Analysis: Analysis showing that BRo-JEPA organizes latent representations into a structured geometric manifold (a circular manifold), enabling the model to internalize modular arithmetic topology.

Results
The experiments were conducted with training on operations $\{-1, +1\}$ and evaluation on unseen operations $\{\pm 2, \dots, \pm 9\}$.

• Baseline Performance: Supervised baselines achieved high accuracy on seen operations (~97-98%) but near-zero accuracy on unseen operations (2.42% for ResNet-18). PCA visualizations revealed that these models arbitrarily crammed unseen operations into the representation space of seen ones.
• JEPA with Additive Embeddings: While these models showed high rollout accuracy (applying $+1$ repeatedly to simulate $+k$), they failed at direct zero-shot prediction of unseen operations (accuracy ~10%). The composition consistency loss improved this to ~99% but is considered "weak zero-shot" because the model was indirectly exposed to composed operations during training.
• BRo-JEPA Performance:
  • Zero-Shot Accuracy: The ResNet-based BRo-JEPA with Multi-Frequency Rotation (MFR) achieved 99.46% zero-shot accuracy on unseen operations.
  • Rollout Accuracy: The model also achieved 99.46% accuracy in multi-step rollouts, confirming the stability of the learned latent dynamics.
  • Latent Structure: PCA visualizations confirmed that BRo-JEPA successfully coerced the latent space into a continuous elliptical/circular manifold where digit classes are arranged sequentially. The MFR variant produced better-separated class prototypes (lower off-diagonal cosine similarity) compared to SFR, leading to superior generalization.
  • Learnable Angles: When rotation angles were learnable, MFR converged rapidly to effective generalization, whereas SFR exhibited significant delays and instability, often converging to suboptimal harmonic solutions.

Significance and Claims
The paper claims that latent world models can learn symbolic transformation rules, but only when the architecture's inductive bias matches the structure of the problem. The results suggest that:

• Standard neural predictors lack the constraints necessary to preserve compositional geometry for unseen operations.
• Explicitly modeling the cyclic nature of modular arithmetic via block rotations in latent space enables strong zero-shot generalization without requiring exposure to the specific operations during training.
• This approach offers a promising direction for representation learning beyond arithmetic, suggesting that incorporating appropriate transformation structures into latent transition models is crucial for learning abstract rules.

The authors remain modest, noting that their results are specific to the highly controlled setting of MNIST modular arithmetic with a known cyclic group structure. They do not claim the architecture will automatically generalize to arbitrary symbolic reasoning tasks or unknown transformation groups without similar structural priors.