BRo-JEPA: 潜在空間におけるモジュラー変換の学習
BRo-JEPA: Learning Modular Transformations in Latent Space
視覚入力から代数規則を学習できるかという問いに対し、算術演算を潜在空間での回転として表現する世界モデルBRo-JEPAを提案し、未学習の演算へのゼロショット汎化を達成した。
著者: Divyansh Jha, Yuanfang Xie, Brennen Yu, Varan Mehra
分類: cs.LG, cs.AI, cs.CV
原文アブストラクト
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}.