# BRo-JEPA: Learning Modular Transformations in Latent Space

arXiv: https://arxiv.org/abs/2606.01372 · HTML: https://mlsanity.com/paper/2606.01372/

## 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.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 14 of 14)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [Neural Arithmetic Logic Units](https://mlsanity.com/paper/1808.00508/) (2018)
- [Improving Compositional Generalization with Latent Structure and Data Augmentation](https://mlsanity.com/paper/2112.07610/) (2021)
- [Neural Arithmetic Units](https://mlsanity.com/paper/2001.05016/) (2020)
- [Compositional generalization through abstract representations in human and artificial neural networks](https://mlsanity.com/paper/2209.07431/) (2022)
- [Can neural networks do arithmetic? A survey on the elementary numerical skills of state-of-the-art deep learning models](https://mlsanity.com/paper/2303.07735/) (2023)
- [A Primer for Neural Arithmetic Logic Modules](https://mlsanity.com/paper/2101.09530/) (2021)
- [Scaling can lead to compositional generalization](https://mlsanity.com/paper/2507.07207/) (2025)

## 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.