BRo-JEPA: Learning Modular Transformations in Latent Space
Abstract
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 $\pm$1. Our results suggest that world models can learn algebraic rules when the latent transformations encode the underlying modular structure. Code is available \href{https://github.com/DL-World-Models/brojepa}{here}.
