BRo-JEPA: Learning Modular Arithmetic in Latent Space
For researchers in world models and symbolic reasoning, this work shows that latent world models can learn symbolic transformation rules when the architecture aligns with the problem's structure, though the task is narrow (modular arithmetic on MNIST).
The paper investigates whether neural networks can learn abstract algebraic rules using a JEPA-style latent world model on MNIST digits with modular arithmetic operations. The proposed block-rotation predictor achieves 99.46% zero-shot and rollout accuracy, demonstrating strong generalization to unseen operations.
Can neural networks learn abstract algebraic rules, or do they merely memorize training patterns? We investigate this using MNIST digits as states and modular arithmetic operations as actions in a JEPA-style latent world model. Standard supervised baselines and JEPA models with additive operation embeddings fit seen operations but fail to extrapolate reliably to unseen ones. To bridge this gap, we introduce a block-rotation predictor that imposes the circular structure of modulo-10 arithmetic in latent space. This enables strong zero-shot generalization, with the best ResNet-based JEPA block-rotation model achieving 99.46\% zero-shot and 99.46\% rollout accuracy. Our results suggest that latent world models can learn symbolic transformation rules when architecture matches the structure of the problem. Our code can be \href{https://github.com/DL-World-Models/mnist-math}{accessed here}.