Mask-Aware Execution for Efficient JEPA Training

2026-09-18

AI Summary

本文针对JEPA训练效率低的问题，提出了一种名为M-JEPA的执行架构，通过分离与掩码相关的计算和路由，减少了计算、内存流量和同步开销，提高了训练速度。

Abstract

Joint Embedding Predictive Architectures (JEPAs) are becoming a core representation-learning primitive and a building block for latent world models across vision, video, audio, brain dynamics, and time series. Despite (potential of) wide deployment, current JEPA training pipelines are inefficient: each input is executed through multiple mask-specific branches, with redundant target-side work, and memory-bound token routing. These costs grow with the number of masks and limit GPU efficiency. We present M-JEPA, a mask-aware execution architecture that restructures JEPA training without changing the learning objective. M-JEPA separates mask-independent computation from mask-dependent routing, enabling shared context encoder execution, fused token routing and slicing with backward support, sparse target encoder execution over the union of target tokens, and masked patch embedding for sparse inputs. The resulting pipeline preserves training semantics while reducing computation, memory traffic, and synchronization overhead. We implement M-JEPA for five JEPA variants and evaluate it on NVIDIA A100 GPUs. Compared against the state-of-the-art baselines, M-JEPA achieves up to 1.7x end-to-end training speedup for 2-10 masks. Separately, with masked patch embedding, 4.75x patch-embedding speedup at high sparsity. These results show that execution restructuring, rather than changes to the JEPA objective, is a key lever for efficient JEPA training.

Problem

• JEPA
• mask-specific branches
• redundant work
• memory-bound token routing
• GPU efficiency

Innovation

• Mask-Aware Execution
• Shared Context Encoder
• Sparse Target Encoding
• Fused Token Routing