arXiv:2606.01372cs.LGcs.AI2026-06

让神经网络学会模运算规则,实现零样本泛化。

BRo-JEPA: Learning Modular Transformations in Latent Space

论文配图:BRo-JEPA: Learning Modular Transformations in Latent Space
图 1 · 摘自论文原文
  • 用旋转操作建模模运算,让隐空间具循环结构。
  • 在未见运算上零样本准确率达99.44%(MNIST)和94.35%(EMNIST)。
  • 适合研究世界模型、可解释表示与数学推理的读者。

我们以MNIST(或EMNIST字母)为状态,模算术运算为动作,构建类JEPA的世界模型,探究神经网络能否从视觉输入中学习代数规则。标准监督与JEPA基线在已见运算上表现良好,但无法可靠外推至未见运算。本文提出BRo-JEPA,通过块旋转预测器将算术运算表示为旋转,使隐空间呈现模运算的循环结构。通过旋转对齐隐表示,模型能严格实现零样本泛化。最佳基线仅达54.54%(MNIST)与25.13%(EMNIST)零样本准确率,而使用ResNet-18编码器的BRo-JEPA分别达到99.44%与94.35%,且仅在±1原始操作上训练。结果表明,当隐变换编码底层模结构时,世界模型可学习代数规则。代码已开源。

原文摘要 · Abstract (English)

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

世界模型模运算零样本隐空间

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。