[Paper Notes] JEPA: Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
Input: randomly samples block masks from original image within certain aspect ratio changes, and apply mask for context image
2.1.2 Context
2.1.2 Context
Context Encoder: ViT encodes context image to embedding
SxS_xSx
Mask Token: an [1,D] randomly initialized shared learnable vector
MMM
, values are used where it is a masked pactc (colored pacthes in the figure).
Positional embedding: [1,D] sinusoidal embedding
PEPEPE
Predictor: standard ViT, inputs masked token, context embedding and positional embedding for
gϕ(M+PE+Sx)=y^g_{\phi}(M + PE + S_x) = \hat{y}gϕ(M+PE+Sx)=y^
2.1.3 Target
Target Encoder: input original image and masked bbox to get corresponding embedding patches as
yyy
.
2.1.4 Loss & Training
The avg.
L2L_2L2
distance between predicted
y^\hat{y}y^
and target
yyy
.
L=1∣M∣∑(i,j)∈M∣gϕ(fθ(xv), PEi,j)−fξ(x)i,j∣2\mathcal{L} = \frac{1}{|\mathcal{M}|} \sum_{(i,j) \in \mathcal{M}} \left| g_{\phi}(f_{\theta}(x_v),\ PE_{i,j}) - f_{\xi}(x)_{i,j} \right|^2L=∣M∣1(i,j)∈M∑∣gϕ(fθ(xv), PEi,j)−fξ(x)i,j∣2
Where:
fθf_{\theta}fθ
— context encoder (trained), processes visible patches
xvx_vxv
fξf_{\xi}fξ
— target encoder (EMA), processes full image
xxx
gϕg_{\phi}gϕ
— predictor (trained), takes context representations + positional embedding of masked position
M\mathcal{M}M
— set of masked patch positions
PEi,jPE_{i,j}PEi,j
— positional embedding at masked position
(i,j)(i,j)(i,j)
Target encoder's EMA (parameters update not in loss by gradient descent, but direct update after each step):
ξt=αξt−1+(1−α)θt−1\xi_t = \alpha \xi_{t-1} + (1-\alpha)\theta_{t-1}ξt=αξt−1+(1−α)θt−1
=(1−α)∑k=0tαkθt−k= (1-\alpha)\sum_{k=0}^{t} \alpha^k \theta_{t-k}=(1−α)k=0∑tαkθt−k
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
α=0.996\alpha=0.996α=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 informatio theory,
max I(z^; fξ(x))\max\ I(\hat{z};\ f_{\xi}(x))max I(z^; fξ(x))
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.
SOCIAL SHARE CARD GENERATOR