GoKawiilThe Annotated JEPA Elon Litman
This post is a step-by-step, annotated, from-scratch walkthrough of Joint Embedding Predictive Architectures, or JEPAs. The goal is to do for JEPA what The Annotated Transformer did for the Transformer: build the full object, explain every moving part, and end with a working training loop. JEPA is Yann LeCun's proposed answer to a fundamental question in self-supervised learning: how do you train a model to understand the world without labels, without collapsing to trivial solutions, and without wasting capacity on irrelevant details?
The answer, elegant in principle and subtle in practice, is prediction in representation or latent space. To keep the discussion concrete, the main running example is I‑JEPA, the image instantiation Why images and video rather than text? LeCun argues that language is already a highly compressed, discrete representation of knowledge; predicting the next token requires modeling human communication patterns, not physical reality. Visual prediction, by contrast, demands understanding of persistence, occlusion, and dynamics. JEPA is designed for domains where pixel-level reconstruction wastes capacity on irrelevant details, a problem that does not arise in the same way for discrete tokens. We return to this near the end., introduced as a self-supervised method that learns semantic image representations by predicting representations of masked regions from visible context See Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture (2023). I-JEPA is a non-generative approach that avoids hand-crafted data augmentations entirely.. We will build I‑JEPA from scratch, then discuss its extension to video with V‑JEPA and V‑JEPA 2 See V-JEPA: Latent Video Prediction for Visual Representation Learning (2024), and V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning (2025)., and then examine LeJEPA, the latest attempt to replace engineering heuristics with a distributional regularizer See LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics (2025).. What follows is meant to be pedagogical. The implementation omits FlashAttention, gradient checkpointing, mixed-precision, and the batching strategies that make large-scale training feasible. These are engineering choices that would dominate a production codebase but are trivially separable from the mathematics.
The Problem Self-supervised representation learning asks: how do you learn useful features without labels? You need an objective that captures meaningful structure and, without labels, finding one that actually works is the central difficulty of the field. JEPA's answer: train by prediction, but predict in representation space. Why should this work at all, though? Suppose you see part of an image, the context \(x\), and want to learn representations. Somewhere else in the image is a target region \(y\) that you cannot see. An encoder maps \(y\) to a representation \(s_y\). A predictor takes your encoding of the context and outputs \(\hat{s}_y\), its guess at what \(s_y\) should be. Training minimizes the distance \(D(\hat{s}_y, s_y)\). Now ask: when can the predictor succeed? Only when the context encoding \(s_x\) contains enough information to determine what \(s_y\) must be. If you saw the hood of a car, predicting the representation of the wheels requires that your encoding of the hood captures this is a car. If you saw a face, predicting the representation of the hair requires that your encoding captures identity, pose, and lighting. The predictor cannot hallucinate structure that the context encoding lacks. This is the forcing function. The context encoder must learn to extract features from \(x\) that are predictive of \(y\)'s representation. These are exactly the semantic, structural features: object identity, spatial relationships, physical constraints. Pixel-level noise in \(x\) does not help predict \(s_y\), so the encoder learns to ignore it. What remains is what generalizes. The target encoder has a complementary pressure. Its output \(s_y\) must be predictable from context. If \(s_y\) encoded random high-frequency texture, no amount of context would help predict it. So the target encoder learns to output representations that capture the shared structure between \(x\) and \(y\), the structure that makes prediction possible, rather than idiosyncratic details of \(y\) alone. LeCun's position paper frames this as an energy-based formulation See A Path Towards Autonomous Machine Intelligence (2022), OpenReview. LeCun frames JEPA as the foundation for world models that plan in latent space.: encode \(x\) and \(y\) into representations, predict one from the other, define energy as prediction error in that abstract space. The architecture factors into components we can implement and analyze: two encoders, one predictor, one distance function.
The JEPA Template A JEPA starts from paired, semantically related views of the world. In the most general form, imagine triples \((x, y, z)\): \(x\) is what you observe, \(y\) is what you want to predict, and \(z\) is an optional latent capturing unknown factors that make the prediction multimodal. The pairs \((x, y)\) are drawn from some joint distribution over observations. In image pretraining, \(x\) might be visible patches and \(y\) the masked patches of the same image; in video, \(x\) is a clip prefix and \(y\) is the continuation; in cross-modal settings, \(x\) could be audio and \(y\) the corresponding video. The only structural requirement is that knowing \(x\) should constrain what \(y\) can be. LeCun's formulation adds an optional latent \(z\) to handle multimodality in the predictive relationship. When multiple values of \(y\) are consistent with the same \(x\), the predictor conditions on \(z\) to select among them. This matters for temporal prediction, where the future is genuinely uncertain, less so for masked image modeling, where the context typically determines the target up to noise. The template has three parts. Encode both observations into a shared representation space: $$ s_x = f_\theta(x), \quad s_y = f_{\bar\theta}(y) $$ The encoders \(f_\theta\) and \(f_{\bar\theta}\) can differ in architecture or parameter sharing. When \(x\) and \(y\) live in different modalities, they must differ; when they are the same modality, weight sharing is a design choice that trades inductive bias against flexibility. Predict the target representation from the context representation: $$\widehat{s}_y = g_\phi(s_x, z)$$ Minimize prediction error in representation space: $$ \mathcal{L}(\theta, \bar\theta, \phi) = D(\widehat{s}_y, s_y) $$ The rest of this post instantiates these abstractions: defining \(x\) and \(y\) concretely for images, choosing architectures for the encoders and predictor, specifying the distance function \(D\), and implementing the mechanism that prevents the trivial solution where all representations collapse to a constant.
Prediction vs. Reconstruction Why predict in representation space rather than pixel space? Consider two image patches that are semantically identical, same object, same meaning, but differ in hundreds of pixels due to lighting, texture, or JPEG artifacts. A reconstruction objective must explain all those differences. The model wastes capacity modeling high-entropy noise. With JEPA, the target encoder can learn to output representations that discard nuisance details. When you regress to that representation, you push the model to capture what the encoder preserves, not what pixels happen to contain. I‑JEPA makes this explicit: the architecture resembles generative models, but the loss lives in embedding space, not input space Analogous to perceptual losses (LPIPS) vs. pixel losses (MSE). Representation-space prediction lets the model focus on semantically meaningful features..
Collapse The easiest way to minimize a matching loss? Output the same vector for everything. Then, loss goes to zero. But the representations become useless. This is the central failure mode of all joint embedding methods, and I‑JEPA addresses it directly: they use an asymmetric design between encoders. The literature offers many anti-collapse strategies SimCLR uses contrastive negatives. VICReg uses variance-invariance-covariance regularization. BYOL and SimSiam use stop-gradient. DINO/MoCo use EMA teachers. Each trades off differently.: contrastive negatives (expensive), explicit covariance constraints (tricky to tune), stop-gradient (surprisingly effective), or teacher-student encoders with exponential moving average updates (the I‑JEPA choice). I‑JEPA trains the context encoder and predictor by gradient descent, but updates the target encoder by exponential moving average (EMA) of the context encoder weights. Here is why this prevents collapse. Consider what happens if both encoders are trained by gradient descent. The loss is \(D(\hat{s}_y, s_y)\). If the target encoder can change freely, it will learn to output representations that are easy to predict, which means constant. The context encoder follows. Both converge to outputting the same vector for all inputs: loss zero, representations useless. EMA breaks this co-adaptation. The target encoder updates as $$\theta_{\text{target}} \leftarrow m \theta_{\text{target}} + (1-m) \theta_{\text{context}}$$ with \(m\) close to 1. This means the target encoder lags behind the context encoder. The target representations are stable on the timescale of a gradient step, so the context encoder cannot exploit fast-moving targets to find degenerate solutions. The context encoder must actually learn to predict what the (slowly evolving) target encoder outputs, which forces it to extract meaningful features. Later we discuss LeJEPA, which tries to make collapse prevention principled rather than heuristic. For now, we focus on the classic I‑JEPA design.
Specializing JEPA to images A JEPA needs paired views \(x\) and \(y\) that share semantics. For I‑JEPA, both come from the same image: \(x\) is a large context block of patches, and \(y\) consists of several smaller target blocks elsewhere in the image. The goal is to predict the representations of the target blocks from the context block. There are three key design choices in I‑JEPA that you should understand before writing any code. The first choice is to represent the image as a sequence of patch tokens, following Vision Transformers See Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (2020). The ViT architecture splits images into fixed-size patches and processes them as a sequence.. A 224×224 image with 16×16 patches yields 196 tokens arranged in a 14×14 grid. The second choice is the masking strategy. The paper states that to guide I‑JEPA toward semantic representations, it is crucial to predict target blocks that are sufficiently large in scale, and to use a context block that is sufficiently informative and spatially distributed. In the method section, the paper specifies typical mask sampling ranges: it commonly uses \(M = 4\) target blocks, with an aspect ratio range \((0.75, 1.5)\) and a scale range \((0.15, 0.2)\). The context block is sampled with a scale range \((0.85, 1.0)\) and unit aspect ratio, and any overlap between targets and context is removed to keep the prediction task non-trivial. The third choice is where masking happens for targets. Target blocks are obtained by masking the output of the target encoder, not the input. The target encoder sees the full image and produces high-level representations; we then select which representations to predict. The paper says this distinction is crucial to ensure the target representations are of a high semantic level This is computationally heavier than masking inputs for both views, but necessary for correctness. Since ViT uses global attention, a patch's representation depends on the full image context. Running the target encoder only on target patches would change what those representations mean.. These are not arbitrary implementation details. They are the entire point of I‑JEPA's inductive bias, replacing hand-designed multi-view augmentations with a structured prediction task.
The full Image‑JEPA architecture We will implement three networks. The context encoder (also called the student) is a Vision Transformer that processes only the visible context patches and produces a set of context embeddings. The target encoder (also called the teacher) is another Vision Transformer that processes the full image and produces embeddings for all patches. Its weights are not updated by gradients; they are updated by exponential moving average of the context encoder weights. The predictor is a smaller Transformer that takes the context embeddings and a set of mask tokens corresponding to the target patch positions, and produces predicted embeddings for those target positions. The method description says the predictor takes as input the context encoder output plus a mask token for each patch to predict, with the mask token parameterized by a shared learnable vector and an added positional embedding. The loss is mean squared error between predicted patch representations and the target encoder's patch representations for the target blocks. This is the distance function \(D\) from the JEPA template: $$D(\hat{s}_y, s_y) = \|\hat{s}_y - s_y\|^2_2$$
A PyTorch implementation The implementation is broken into pieces; you should be able to read it from top to bottom. Throughout, we will use these conventions: images have shape \((B, C, H, W)\) and patch tokens have shape \((B, N, D)\), where \(N\) is the number of patches and \(D\) is the embedding dimension. Whenever we select a subset of patches, we will represent that subset as a vector of indices into the \(N\) patches and we will implement absolute positional embeddings for simplicity. Imports and small helpers import math import random from dataclasses import dataclass from typing import List , Optional , Tuple import torch import torch.nn as nn import torch.nn.functional as F Patch embedding A standard Vision Transformer creates patch tokens with a strided convolution. Each patch becomes a vector in \(\mathbb{R}^D\). class PatchEmbed ( nn . Module ): """ Converts an image to a sequence of patch embeddings. Input: (B, C, H, W) Output: (B, N, D) where N = (H / P) * (W / P) """ def __init__ ( self , img_size : int = 224 , patch_size : int = 16 , in_chans : int = 3 , embed_dim : int = 768 , ): super () . __init__ () assert img_size % patch_size == 0 , "img_size must be divisible by patch_size" self . img_size = img_size self . patch_size = patch_size self . grid_size = img_size // patch_size self . num_patches = self . grid_size * self . grid_size self . proj = nn . Conv2d ( in_channels = in_chans , out_channels = embed_dim , kernel_size = patch_size , stride = patch_size , bias = True , ) def forward ( self , x : torch . Tensor ) -> torch . Tensor : x = self . proj ( x ) # (B, D, G, G) x = x . flatten ( 2 ) . transpose ( 1 , 2 ) # (B, N, D) return x The only non-obvious line is the reshape. After the convolution you have \((B, D, G, G)\). Flattening dimensions 2 and 3 yields \((B, D, N)\). Transposing gives the usual token format \((B, N, D)\). The Transformer block I‑JEPA uses Vision Transformers for the encoders and for the predictor. We will implement a standard pre-norm Transformer block, which is close to what most practical ViTs use Pre-norm (LayerNorm before attention/MLP) is more stable for training deep transformers than post-norm (LayerNorm after). See Xiong et al., On Layer Normalization in the Transformer Architecture (2020).. First, multi-head self-attention: class MultiHeadSelfAttention ( nn . Module ): def __init__ ( self , dim : int , num_heads : int , attn_drop : float = 0.0 , proj_drop : float = 0.0 , ): super () . __init__ () assert dim % num_heads == 0 , "dim must be divisible by num_heads" self . dim = dim self . num_heads = num_heads self . head_dim = dim // num_heads self . scale = self . head_dim ** - 0.5 self . qkv = nn . Linear ( dim , 3 * dim , bias = True ) self . attn_drop = nn . Dropout ( attn_drop ) self . proj = nn . Linear ( dim , dim , bias = True ) self . proj_drop = nn . Dropout ( proj_drop ) def forward ( self , x : torch . Tensor ) -> torch . Tensor : """ x: (B, T, D) """ B , T , D = x . shape qkv = self . qkv ( x ) # (B, T, 3D) qkv = qkv . reshape ( B , T , 3 , self . num_heads , self . head_dim ) qkv = qkv . permute ( 2 , 0 , 3 , 1 , 4 ) # (3, B, H, T, Hd) q , k , v = qkv [ 0 ], qkv [ 1 ], qkv [ 2 ] # (B, H, T, Hd) each attn = ( q @ k . transpose ( - 2 , - 1 )) * self . scale # (B, H, T, T) attn = attn . softmax ( dim =- 1 ) attn = self . attn_drop ( attn ) out = attn @ v # (B, H, T, Hd) out = out . transpose ( 1 , 2 ) . reshape ( B , T , D ) # (B, T, D) out = self . proj ( out ) out = self . proj_drop ( out ) return out Then the MLP: class MLP ( nn . Module ): def __init__ ( self , dim : int , hidden_dim : int , drop : float = 0.0 ): super () . __init__ () self . fc1 = nn . Linear ( dim , hidden_dim ) self . act = nn . GELU () self . fc2 = nn . Linear ( hidden_dim , dim ) self . drop = nn . Dropout ( drop ) def forward ( self , x : torch . Tensor ) -> torch . Tensor : x = self . fc1 ( x ) x = self . act ( x ) x = self . drop ( x ) x = self . fc2 ( x ) x = self . drop ( x ) return x Then the Transformer block: class TransformerBlock ( nn . Module ): def __init__ ( self , dim : int , num_heads : int , mlp_ratio : float = 4.0 , attn_drop : float = 0.0 , drop : float = 0.0 , ): super () . __init__ () self . norm1 = nn . LayerNorm ( dim ) self . attn = MultiHeadSelfAttention ( dim = dim , num_heads = num_heads , attn_drop = attn_drop , proj_drop = drop , ) self . norm2 = nn . LayerNorm ( dim ) self . mlp = MLP ( dim = dim , hidden_dim = int ( dim * mlp_ratio ), drop = drop ) def forward ( self , x : torch . Tensor ) -> torch . Tensor : x = x + self . attn ( self . norm1 ( x )) x = x + self . mlp ( self . norm2 ( x )) return x You can now build a ViT encoder by stacking these blocks. A ViT encoder that can process a subset of patches I‑JEPA's efficiency comes from letting the context encoder process only visible context patches. To model that, we write an encoder that can optionally select patches by index before running the Transformer: class ViTEncoder ( nn . Module ): """ Minimal ViT-style patch encoder. This version does not use a CLS token. The I-JEPA method section describes patch-level representations and does not require CLS during pretraining. """ def __init__ ( self , img_size : int = 224 , patch_size : int = 16 , in_chans : int = 3 , embed_dim : int = 768 , depth : int = 12 , num_heads : int = 12 , mlp_ratio : float = 4.0 , drop : float = 0.0 , attn_drop : float = 0.0 , ): super () . __init__ () self . patch_embed = PatchEmbed ( img_size , patch_size , in_chans , embed_dim ) self . num_patches = self . patch_embed . num_patches self . pos_embed = nn . Parameter ( torch . zeros ( 1 , self . num_patches , embed_dim )) nn . init . trunc_normal_ ( self . pos_embed , std = 0.02 ) self . blocks = nn . ModuleList ([ TransformerBlock ( dim = embed_dim , num_heads = num_heads , mlp_ratio = mlp_ratio , attn_drop = attn_drop , drop = drop , ) for _ in range ( depth ) ]) self . norm = nn . LayerNorm ( embed_dim ) def forward ( self , images : torch . Tensor , patch_indices : Optional [ torch . Tensor ] = None ) -> torch . Tensor : """ images: (B, C, H, W) patch_indices: None or a 1D tensor of indices into the patch sequence Returns: (B, T, D) """ x = self . patch_embed ( images ) # (B, N, D) if patch_indices is None : x = x + self . pos_embed else : if patch_indices . dim () != 1 : raise ValueError ( "patch_indices must be a 1D tensor in this minimal implementation" ) x = x [:, patch_indices , :] + self . pos_embed [:, patch_indices , :] for blk in self . blocks : x = blk ( x ) x = self . norm ( x ) return x There is a deliberate limitation: patch_indices is shared across the batch. In real training, each image has different masks, so you either vectorize with padding and attention masks or you loop over samples. For pedagogy, we will loop. Block mask sampling Now we implement the core I‑JEPA view construction: sample multiple target blocks and one context block, then remove overlap. The paper describes this procedure and provides typical ranges: four target blocks with scale in the range \((0.15, 0.2)\) and aspect ratio \((0.75, 1.5)\), one context block with scale \((0.85, 1.0)\) and unit aspect ratio, and then remove overlap between target and context. We will implement a mask generator that returns patch indices: def _sample_block_on_grid ( grid_size : int , scale_range : Tuple [ float , float ], aspect_ratio_range : Tuple [ float , float ], ) -> Tuple [ int , int , int , int ]: """ Samples a rectangle on a grid_size x grid_size patch grid. Returns (top, left, height, width) in patch coordinates. """ G = grid_size N = G * G scale = random . uniform ( scale_range [ 0 ], scale_range [ 1 ]) aspect = random . uniform ( aspect_ratio_range [ 0 ], aspect_ratio_range [ 1 ]) area = scale * N w = int ( round ( math . sqrt ( area * aspect ))) h = int ( round ( math . sqrt ( area / aspect ))) w = max ( 1 , min ( G , w )) h = max ( 1 , min ( G , h )) top = random . randint ( 0 , G - h ) left = random . randint ( 0 , G - w ) return top , left , h , w def _rect_to_indices ( grid_size : int , top : int , left : int , h : int , w : int ) -> torch . Tensor : """ Converts a rectangle to flat patch indices in row major order. """ rows = torch . arange ( top , top + h ) cols = torch . arange ( left , left + w ) rr , cc = torch . meshgrid ( rows , cols , indexing = "ij" ) idx = rr * grid_size + cc return idx . flatten () Now the generator: @dataclass class IJEPAMasks : context_idx : torch . Tensor target_blocks : List [ torch . Tensor ] class IJEPAMaskGenerator : def __init__ ( self , grid_size : int , num_targets : int = 4 , target_scale : Tuple [ float , float ] = ( 0.15 , 0.20 ), target_aspect : Tuple [ float , float ] = ( 0.75 , 1.50 ), context_scale : Tuple [ float , float ] = ( 0.85 , 1.00 ), context_aspect : Tuple [ float , float ] = ( 1.00 , 1.00 ), ): self . grid_size = grid_size self . num_targets = num_targets self . target_scale = target_scale self . target_aspect = target_aspect self . context_scale = context_scale self . context_aspect = context_aspect def __call__ ( self ) -> IJEPAMasks : G = self . grid_size target_blocks : List [ torch . Tensor ] = [] target_union = torch . zeros ( G * G , dtype = torch . bool ) for _ in range ( self . num_targets ): top , left , h , w = _sample_block_on_grid ( G , self . target_scale , self . target_aspect ) idx = _rect_to_indices ( G , top , left , h , w ) target_blocks . append ( idx ) target_union [ idx ] = True c_top , c_left , c_h , c_w = _sample_block_on_grid ( G , self . context_scale , self . context_aspect ) context_idx = _rect_to_indices ( G , c_top , c_left , c_h , c_w ) keep = ~ target_union [ context_idx ] context_idx = context_idx [ keep ] return IJEPAMasks ( context_idx = context_idx , target_blocks = target_blocks ) This generator embodies the paper's method section description of how the targets and context are sampled and how overlap is removed. From a JEPA viewpoint, this is the definition of \(x\) and the definition of which parts of \(y\) we are trying to predict. The predictor and mask tokens In I‑JEPA, the predictor receives context embeddings plus mask tokens that mark the target positions we want it to predict. The paper describes the mask tokens as a shared learnable vector with positional embedding added. A minimal design is to implement the predictor as a small Transformer that processes a concatenated sequence: the input sequence is the context tokens followed by mask tokens, one per target patch position, each augmented with the positional embedding for that patch position. The output sequence is the same length, and we read out the final tokens corresponding to the target positions. class IJEPA_Predictor ( nn . Module ): def __init__ ( self , embed_dim : int , depth : int = 6 , num_heads : int = 12 , mlp_ratio : float = 4.0 , drop : float = 0.0 , attn_drop : float = 0.0 , ): super () . __init__ () self . mask_token = nn . Parameter ( torch . zeros ( 1 , 1 , embed_dim )) nn . init . trunc_normal_ ( self . mask_token , std = 0.02 ) self . blocks = nn . ModuleList ([ TransformerBlock ( dim = embed_dim , num_heads = num_heads , mlp_ratio = mlp_ratio , attn_drop = attn_drop , drop = drop , ) for _ in range ( depth ) ]) self . norm = nn . LayerNorm ( embed_dim ) def forward ( self , context_tokens : torch . Tensor , target_pos_embed : torch . Tensor ) -> torch . Tensor : """ context_tokens: (B, Nc, D) target_pos_embed: (B, Nt, D) Returns: (B, Nt, D) """ B , Nc , D = context_tokens . shape _ , Nt , _ = target_pos_embed . shape mask_tokens = self . mask_token . expand ( B , Nt , D ) + target_pos_embed x = torch . cat ([ context_tokens , mask_tokens ], dim = 1 ) for blk in self . blocks : x = blk ( x ) x = self . norm ( x ) pred_target = x [:, Nc :, :] return pred_target There are two conceptual reasons this works. First, self-attention lets each mask token attend to all context tokens, so information flows from what the model saw to what it must predict. Second, adding positional embeddings to the mask tokens tells the predictor where in the image the unknown region lives; without this, all target positions would be indistinguishable. Wiring the full model Now we implement the JEPA template exactly. The target encoder produces \(s_y\) for all patches. The mask generator chooses a context subset and several target subsets. The context encoder encodes only the context subset into \(s_x\). For each target block, the predictor outputs \(\widehat{s}_y\) at those positions. The loss is mean squared error between predicted embeddings and target embeddings on the target positions. The target encoder is updated by exponential moving average, rather than gradients: $$\theta_{\text{target}} \leftarrow m \theta_{\text{target}} + (1-m) \theta_{\text{context}}$$ I‑JEPA states this explicitly and notes that they find an EMA target encoder essential in their setting Here \(m\) is typically \(0.996\) to \(1.0\). This creates a slowly-evolving teacher that provides stable targets.. class IJEPAModel ( nn . Module ): def __init__ ( self , img_size : int = 224 , patch_size : int = 16 , in_chans : int = 3 , embed_dim : int = 768 , enc_depth : int = 12 , enc_heads : int = 12 , pred_depth : int = 6 , pred_heads : int = 12 , mlp_ratio : float = 4.0 , ): super () . __init__ () self . context_encoder = ViTEncoder ( img_size = img_size , patch_size = patch_size , in_chans = in_chans , embed_dim = embed_dim , depth = enc_depth , num_heads = enc_heads , mlp_ratio = mlp_ratio , ) self . target_encoder = ViTEncoder ( img_size = img_size , patch_size = patch_size , in_chans = in_chans , embed_dim = embed_dim , depth = enc_depth , num_heads = enc_heads , mlp_ratio = mlp_ratio , ) self . _init_teacher_from_student () self . predictor = IJEPA_Predictor ( embed_dim = embed_dim , depth = pred_depth , num_heads = pred_heads , mlp_ratio = mlp_ratio , ) G = self . context_encoder . patch_embed . grid_size self . mask_generator = IJEPAMaskGenerator ( grid_size = G ) @torch . no_grad () def _init_teacher_from_student ( self ) -> None : for p_t , p_s in zip ( self . target_encoder . parameters (), self . context_encoder . parameters ()): p_t . data . copy_ ( p_s . data ) for p in self . target_encoder . parameters (): p . requires_grad_ ( False ) @torch . no_grad () def ema_update_teacher ( self , momentum : float ) -> None : """ Teacher update: theta_bar <- m * theta_bar + (1 - m) * theta """ for p_t , p_s in zip ( self . target_encoder . parameters (), self . context_encoder . parameters ()): p_t . data . mul_ ( momentum ) . add_ ( p_s . data , alpha = ( 1.0 - momentum )) def forward ( self , images : torch . Tensor ) -> torch . Tensor : """ Computes the I-JEPA loss on a batch. This version loops over samples for clarity, since masks differ per image. """ B = images . shape [ 0 ] device = images . device with torch . no_grad (): target_tokens_all = self . target_encoder ( images , patch_indices = None ) # (B, N, D) pos_embed = self . context_encoder . pos_embed # (1, N, D) total_loss = 0.0 total_blocks = 0 for b in range ( B ): masks = self . mask_generator () context_idx = masks . context_idx . to ( device ) context_tokens = self . context_encoder ( images [ b : b + 1 ], patch_indices = context_idx ) # (1, Nc, D) for block_idx in masks . target_blocks : block_idx = block_idx . to ( device ) target_block = target_tokens_all [ b : b + 1 , block_idx , :] # (1, Nt, D) target_pos = pos_embed [:, block_idx , :] # (1, Nt, D) pred_block = self . predictor ( context_tokens , target_pos ) # (1, Nt, D) loss = F . mse_loss ( pred_block , target_block , reduction = "mean" ) total_loss += loss total_blocks += 1 return total_loss / max ( 1 , total_blocks ) This is the clearest embodiment of I‑JEPA's method description: the predictor is applied once per target block, conditioned on the target locations, and the loss averages over those predicted patch representations A subtlety: since target blocks have variable sizes, averaging MSE per block weights patches in smaller blocks more heavily than patches in larger blocks. To define energy strictly as average error per patch, you would sum squared errors and divide by total patch count. In practice, stochastic sampling averages this out.. The architecture is complete. What remains is training.
... continue reading