mlsanityData to 14 Oct 2026

BRo-JEPA: Learning Modular Transformations in Latent Space

arXiv 2606.01372 ↗First posted 31 May 2026cs.LG cs.AI cs.CVcited by 0

On the map: World models → Latent predictive world models

Abstract (arXiv)

Can neural networks learn algebraic rules from visual inputs, or do they merely fit observed patterns? We study this question using MNIST (or EMNIST letters) as states and modular arithmetic operations as actions in a JEPA-style world model. Standard supervised and JEPA baselines with operation embeddings achieve high accuracy on seen operations but fail to extrapolate reliably to unseen operations. We propose BRo-JEPA, a world model with a block-rotation predictor that represents arithmetic operations as rotations, resulting in the cyclic structure of modular arithmetic in latent space. By applying actions as rotations, the BRo-JEPA predictor learns the rotation angles to align the latent representations with the underlying modular structure which enables strict zero-shot operation generalization. While our best block-rotation supervised baseline reaches only 54.54% zero-shot accuracy on MNIST and 25.13% on EMNIST, BRo-JEPA with a ResNet-18 encoder achieves 99.44% and 94.35% respectively, despite being trained only on the primitive operations ±1. Our results suggest that world models can learn algebraic rules when the latent transformations encode the underlying modular structure. Code is available hrefhttps://github.com/DL-World-Models/brojepahere.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Latent predictive world models, one of 11 directions in World models. Builds on 14 of its 18 references that are in the map. First posted on arXiv 31 May 2026.

Builds on

14 of 18 references.

Closest ancestors · 1

Same direction, newest first

  1. 2023Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture19 Jan

Other references · 13

  1. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  2. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  3. 2020Bootstrap your own latent: A new approach to self-supervised Learningcited by 4,750
  4. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  5. 2021VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learningcited by 1,057
  6. 2025V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planningcited by 809
  7. 2018Neural Arithmetic Logic Unitscited by 70
  8. 2021Improving Compositional Generalization with Latent Structure and Data Augmentationcited by 46
  9. 2020Neural Arithmetic Unitscited by 28
  10. 2022Compositional generalization through abstract representations in human and artificial neural networkscited by 22
  11. 2023Can neural networks do arithmetic? A survey on the elementary numerical skills of state-of-the-art deep learning modelscited by 18
  12. 2021A Primer for Neural Arithmetic Logic Modulescited by 7
1 more
  1. 2025Scaling can lead to compositional generalizationcited by 7

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025Learning Modular Exponentiation with Transformers
  2. 2024Unraveling Arithmetic in Large Language Models: The Role of Algebraic Structures
  3. 2024Fourier Circuits in Neural Networks and Transformers: A Case Study of Modular Arithmetic with Multiple Inputs
  4. 2026Low-Dimensional Execution Manifolds in Transformer Learning Dynamics: Evidence from Modular Arithmetic Tasks
  5. 2024Principled Understanding of Generalization for Generative Transformer Models in Arithmetic Reasoning Tasks
  6. 2022Transformers discover an elementary calculation system exploiting local attention and grid-like problem representation
  7. 2025LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics
  8. 2026Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency