BRo-JEPA: Learning Modular Transformations in Latent Space

2026-05-31 · Divyansh Jha, Yuanfang Xie, Varan Mehra, Brennen Yu 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 $\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}.

📄 PDF Abstract BibTeX arXiv:2606.01372

Code (0)

등록된 구현이 없습니다.

Tasks

Zero-shot Generalization

SPLICE: Latent Diffusion over JEPA Embeddings for Conformal Time-Series Inpainting

2026-04-30 · Arnaud Zinflou arxiv

Generative models for time-series imputation achieve strong reconstruction accuracy, yet provide no finite-sample reliability guarantees, a critical limitation in power systems where imputed values inform dispatch and pl…

UWM-JEPA: Predictive World Models That Imagine in Belief Space

2026-05-25 · Santosh Kumar Radha, Oktay Goktas arxiv

World models for partially observed environments must imagine multiple compatible hidden futures and steer between them under counterfactual actions. Joint Embedding Predictive Architectures (JEPAs) do this in latent spa…

VLA-JEPA: Enhancing Vision-Language-Action Model with Latent World Model

2026-02-10 · Jingwen Sun, Wenyao Zhang, Zekun Qi, Shaojie Ren 외 arxiv

Pretraining Vision-Language-Action (VLA) policies on internet-scale video is appealing, yet current latent-action objectives often learn the wrong thing: they remain anchored to pixel variation rather than action-relevan…

VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

2026-01-20 · Yongchao Huang arxiv

Joint Embedding Predictive Architectures (JEPA) offer a scalable paradigm for self-supervised learning by predicting latent representations rather than reconstructing high-entropy observations. However, existing formulat…

Self-Supervised Learning Representation Learning

SC-JEPA: Stabilizing Latent Predictive Learning for Time-Series Anomaly Prediction

2026-02-04 · Yanan He, Yunshi Wen, Xin Wang, Tengfei Ma arxiv

Time-series anomaly prediction aims to forecast future system failures before they fully emerge, making latent predictive models such as JEPA a promising framework for capturing precursor dynamics. However, directly appl…