Searcharxiv⌕ Search
Home/Archive/arXiv

arXiv · 2609.22674

Mask-Aware Execution for Efficient JEPA Training

Md Musfiqur Rahman Sanim·Zhihao Shu·Bahram Afsharmanesh·Amirali Mirian·Wei Niu·Gagan Agrawal
View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

cs.DCcs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Md Musfiqur Rahman Sanim, Zhihao Shu, Bahram Afsharmanesh, Amirali Mirian, Wei Niu, Gagan Agrawal. 2026-09-19. Mask-Aware Execution for Efficient JEPA Training. https://doi.org/10.1145/3838684.3846877

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Fast and Robust Information Spreading in the Noisy PULL Model

Efficient information spreading in stochastic multi agent systems is a core challenge when communication is noisy, bandwidth limited, and agents lack global coordination. Yet biological systems, including ant colonies and fish schools, routinely overcome these constraints. A small number of informed individuals can reliably guide large, uncoordinated populations using minimal, noisy signals. Motivated by these observations, we study how reliable information dissemination can be achieved in such bio inspired settings. A population of $n$ agents, each with a binary preference, includes a designated subset of source agents, and the goal is to converge to the majority preference among the sources. In the noisy $\mathcal{PULL}(h)$ model, each agent observes noisy messages from $h$ randomly sampled peers in every round. Prior work shows that convergence requires $Ω(n/h)$ rounds even under favorable conditions. We ask how far we can push simplicity, with no synchronization at the start time and minimal message size, without compromising convergence speed. We present a quasi self-stabilizing protocol using only 2-bit messages that converges from arbitrary initial states despite severe noise and initial asynchrony. It achieves optimal convergence time $O((n/h)\log n)$ with high probability, and in particular $O(\log n)$ time in the snapshot regime $h=Θ(n)$. A key subroutine is an even simpler 1-bit protocol assuming simultaneous start, based on a natural two phase listen then amplify mechanism. Together, our results show that simple, biologically inspired protocols can achieve optimal and robust information dissemination even in highly unreliable and uncoordinated systems.

cs.DC↗
arXiv2026

LOIP:Collaborative Lossless LLM Inference Serving with Offloading-based Pipeline Parallelism on Edge Devices

Providing lossless inference services of LLMs on edge devices remains challenging, especially given the extremely tight memory budgets. The existing offloading techniques inevitably introduce numerous loading bubbles, which further inflate the end-to-end latency of the entire inference pipeline. Meanwhile, dynamically fluctuating network bandwidth and diverse user request patterns pose additional obstacles to efficient lossless inference on edge devices. To address this, we propose LOIP, a collaborative lossless LLM inference system that employs an offloading-based interleaved pipeline parallelism to better overlap model offloading with computing and communicating. Specifically, LOIP first constructs an offloading-aware cost model to characterize inference latency and memory overhead under heterogeneous device capabilities and limited bandwidth. Based on this cost model, LOIP develops a fine-grained allocation scheduler that determines latency-efficient layer partitions across devices while explicitly accounting for offloading overhead, along with a unified memory architecture (UMA)-aware loading optimization using customized CUDA operators to reduce runtime loading overhead. LOIP further designs an online memory adaptation strategy to handle the increasing KV cache pressure and dynamic bandwidth fluctuations during inference. We implement LOIP with 2500+ lines of Python and 500+ lines of C++/CUDA code, and deploy it on five heterogeneous NVIDIA Jetson edge devices for lossless collaborative inference of LLaMA3.3-70B-Instruct. Extensive experiments demonstrate that LOIP achieves 8.8$\times$$\sim$20.3$\times$ speedups over the SOTA baselines under different bandwidth conditions and request patterns without compromising model accuracy.

cs.DC↗
arXiv2026

Fluid Notarization: Verifiable Evolution of Concurrently Edited Structured Documents

Traditional blockchain-based document notarization follows a snapshot-oriented model in which each document revision is represented as an independent state anchored on-chain through a cryptographic reference. While effective for immutable artifacts, this approach becomes inadequate when documents evolve through collaborative editing. Concurrent modifications create divergent document versions that must be reconciled outside the notarization layer, while even minor changes require generating and distributing new document snapshots. Conversely, collaborative replication frameworks such as CRDTs provide deterministic reconciliation of concurrent updates, but do not inherently provide independently verifiable evidence of when contributions were published. This paper introduces Fluid Notarization, a notarization paradigm in which document evolution itself becomes the object of notarization. Rather than certifying isolated states, Fluid Notarization certifies a graph of causally related evolution artifacts generated by a JSON-native delta-CRDT. The proposed model builds upon Melda, which represents document changes as compact, content-addressed deltas linked through causal dependencies. Blockchain notarization is reduced to recording identifiers of these evolution artifacts, while synchronization, reconstruction, and conflict resolution remain entirely off-chain. The resulting architecture combines two complementary guarantees: deterministic convergence provided by the CRDT and independently auditable proof-of-existence, provenance, and publication evidence provided by the blockchain. A prototype implementation and validation scenario based on collaboratively edited electronic health records demonstrate the feasibility of the approach and highlight the advantages of notarizing document evolution rather than successive document snapshots.

cs.DC↗