[Paper Notes] JEPA: Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture


↓


⭕ EMA (Exponential moving avg.): Update target encoder parameters from the EMA of context encoders. This 'delays' the target encoder to prevent collapsing (a trick from the BYOL paper[2020], proven essential to training JEAs with ViT).

2. Model & High-Level Intuitions

2.1 Model Architecture

Input: randomly samples block masks from original image within certain aspect ratio changes, and apply mask for context image

2.1.2 Context

Context Encoder: ViT encodes context image to embedding


SxS_xSx​


Mask Token: an [1,D] randomly initialized shared learnable vector

$M$

, values are used where it is a masked pactc (colored pacthes in the figure).
Positional embedding: [1,D] sinusoidal embedding

$PE$


Predictor: standard ViT, inputs masked token, context embedding and positional embedding for

2.1.3 Target

Target Encoder: input original image and masked bbox to get corresponding embedding patches as

$y$

.

2.1.4 Loss & Training

The avg.

L2L_2L2​

distance between predicted

y^\hat{y}y^​

and target

$y$

.

Where:



fθf_{\theta}fθ​

— context encoder (trained), processes visible patches

xvx_vxv​

fξf_{\xi}fξ​

— target encoder (EMA), processes full image

$x$

gϕg_{\phi}gϕ​

— predictor (trained), takes context representations + positional embedding of masked position

$\mathcal{M}$

— set of masked patch positions

PEi,jPE_{i,j}PEi,j​

— positional embedding at masked position

$(i,j)$

Target encoder's EMA (parameters update not in loss by gradient descent, but direct update after each step):

The weight of a past context encoder snapshot

θt−k\theta_{t-k}θt−k​

decays as

αk\alpha^kαk

— exponentially in how many steps ago it was. That's where the name comes from.

So with

$\alpha =0.996$

:

1 step ago: weight

=0.9961=0.996= 0.996^1 = 0.996=0.9961=0.996


100 steps ago: weight

=0.996100≈0.67= 0.996^{100} \approx 0.67=0.996100≈0.67


1000 steps ago: weight

=0.9961000≈0.02= 0.996^{1000} \approx 0.02=0.9961000≈0.02

2.2 Intuitions

For anyone familiar with the encoder/decoder architecture pre-training, this paper's biggest innovation no doubt goes into it's self-supervision, the way of obtaining semantic information without labels. It jumped outside of 2 boxes:

• mask reconstruction is 1:1


• learning semantics needs labels.

2.2.1 Discover Image Semantics in the Context vs. Multi-Target Structure

It recognized the semantics exist in pixel structure and provide it to the model by multi context-to-target relationship (similar to data augmentation) combined with positional embedding. Essentially saying, given the context, here is what it misses in different locations, now the pixel & location variation of different mask patches in relation to the context image becomes the source of semantic information.

2.2.2 Target Encoder as Semantics Filter for Mask Reconstruction

While the context-to-target relation provides semantic, it learns low level semantic as it is equivalent to maximizing mutual information between the original and reconstructed image. But pixel data has high entropy from irrelevant details and reconstruction in pixel space never optimize towards compressing information.
Hence, I-JEPA's target encoder acts as a stochastic bottleneck, discarding unpredictable information from the target. In information theory,

Rather than maximizing

I(z^; x)I(\hat{z};\ x)I(z^; x)

directly. Since

fξ(x)f_{\xi}(x)fξ​(x)

already has low-level entropy compressed away, the predictor only needs to capture what's semantically predictable — the mutual information that survives the encoder bottleneck.