DMT-JEPA: Discriminative Masked Targets for Joint-Embedding Predictive Architecture

TL;DR

DMT-JEPA addresses the limited local semantics in embedding-space masked modeling by introducing discriminative latent targets constructed from semantically similar neighboring patches. It employs Masked Semantic Neighboring to select neighbors and a Local Aggregation Target with cross-attention to produce dense latent targets, optimized via an L2 loss between predicted and target representations. Empirically, it outperforms I-JEPA and MAE across dense prediction tasks (ADE20K, COCO, DAVIS) and strengthens ImageNet classification performance, demonstrating improved local-semantic understanding without sacrificing efficiency. The method yields more interpretable attention maps and illustrates the value of latent-space discriminative targets for robust self-supervised learning and downstream transfer.

Abstract

The joint-embedding predictive architecture (JEPA) recently has shown impressive results in extracting visual representations from unlabeled imagery under a masking strategy. However, we reveal its disadvantages, notably its insufficient understanding of local semantics. This deficiency originates from masked modeling in the embedding space, resulting in a reduction of discriminative power and can even lead to the neglect of critical local semantics. To bridge this gap, we introduce DMT-JEPA, a novel masked modeling objective rooted in JEPA, specifically designed to generate discriminative latent targets from neighboring information. Our key idea is simple: we consider a set of semantically similar neighboring patches as a target of a masked patch. To be specific, the proposed DMT-JEPA (a) computes feature similarities between each masked patch and its corresponding neighboring patches to select patches having semantically meaningful relations, and (b) employs lightweight cross-attention heads to aggregate features of neighboring patches as the masked targets. Consequently, DMT-JEPA demonstrates strong discriminative power, offering benefits across a diverse spectrum of downstream tasks. Through extensive experiments, we demonstrate our effectiveness across various visual benchmarks, including ImageNet-1K image classification, ADE20K semantic segmentation, and COCO object detection tasks. Code is available at: \url{https://github.com/DMTJEPA/DMTJEPA}.

Figures (9)

• Figure 1: Illustration of the proposed novel masked modeling framework (DMT-JEPA), rooted in I-JEPA, designed for discriminative latent targets from neighboring information. The Masked Semantic Neighboring module computes the dense semantic similarity between the query patch and its neighboring patch based on representations from the target encoder $f_{\tilde{\theta}}(\cdot)$ to select semantically similar patches from the neighborhood. Then Local Aggregation Target module composed of a context patch aggregation head $h_\theta(\cdot)$ and a target patch aggregation head $h_{\tilde{\theta}}(\cdot)$, aggregates target features of selected patches using cross-attention to construct dense targets. Finally, the model is optimized by the average $L_2$ distance between the predicted dense representations and the target dense representation.
• Figure 2: Qualtitative visualization of learned attention maps using ViT-B/16 model. Columns for each sample denote the original image, attention maps from I-JEPA assran2023self, and attention maps from our DMT-JEPA. Our DMT-JEPA achieves much better attention maps.
• Figure 3.1: Qualtitative visualization of learned attention maps using ViT-B/16 model. Columns for each sample denote the original image, attention maps from target encoder in I-JEPA assran2023self, attention maps from target encoder in our DMT-JEPA, and attention maps from the local aggregation head in our DMT-JEPA. Our DMT-JEPA achieves much better attention maps.
• Figure 3.2: Qualtitative visualization of learned attention maps using ViT-B/16 model. Columns for each sample denote the original image, attention maps from target encoder in I-JEPA assran2023self, attention maps from target encoder in our DMT-JEPA, and attention maps from the local aggregation head in our DMT-JEPA. Our DMT-JEPA achieves much better attention maps.
• Figure 3.3: Qualtitative visualization of learned attention maps using ViT-B/16 model. Columns for each sample denote the original image, attention maps from target encoder in I-JEPA assran2023self, attention maps from target encoder in our DMT-JEPA, and attention maps from the local aggregation head in our DMT-JEPA. Our DMT-JEPA achieves much better attention maps.