8. I-JEPA — Joint Embedding Predictive Architecture¶

NNx ships I-JEPA (Assran et al., CVPR 2023) as a TrainStepFn factory alongside the other self-supervised paradigms. JEPA predicts in latent space: no decoder, no pixel-reconstruction loss. It avoids representation collapse via an EMA target encoder whose weights are never touched by the optimizer.

The shipped path verifies the plumbing on a laptop: a small ViT for a few epochs on 32x32 images. It is not a SOTA reproduction.

1. Public surface¶

Symbol	Notes
`nnx.ViTNN`	Small Vision Transformer encoder. Patch-embed conv + learned pos embeds + CLS + N pre-norm blocks (RMSNorm + bidirectional MHA + SwiGLU). `forward(x, mask=None)` accepts an optional `BoolTensor[B, n_patches]` mask (True = keep), so masked patches never enter attention. Its RMSNorms reduce in at least FP32 and keep FP64 as FP64 (see the precision note below the table).
`nnx.ViTBlock`	Single pre-norm ViT block. Reused by `JEPAPredictor`.
`nnx.JEPAPredictor`	Small predictor module mapping `(context_embeds, context_positions, target_positions) -> predicted_target_embeds`. Uses its own (smaller) hidden width + position embeddings + a learned mask token. Keyword-only constructor: `JEPAPredictor(*, embed_dim, n_patches, predictor_dim=None, n_layers=2, n_heads=2, ffn_mult=4)`. `predictor_dim=None` resolves to `max(8, embed_dim // 2)`.
`nnx.build_target_encoder(source)`	Deep-copy `source`, freeze every param (`requires_grad=False`), pin to `eval()`. The factory function freezes again defensively.
`nnx.update_ema(source, target, momentum)`	In-place EMA update: `target ← momentum · target + (1 - momentum) · source`. Name-keyed against the target's params so a source with extra submodules (the typical "predictor under model.net" idiom) is fine.
`nnx.random_block_mask(*, n_patches, grid_size, …)`	Sample one rectangular block as the prediction target. All arguments are keyword-only. Returns `(context_mask, target_mask)` 1-D BoolTensors.
`nnx.jepa_train_step_factory(target_encoder, predictor, mask_fn, *, ema_momentum=0.996)`	Returns a `TrainStepFn` for `NNModel.train(..., train_step_fn=...)`. The predictor must be a submodule of `model.net`; the first step refuses it otherwise.
`nnx.jepa_objective(target_encoder, predictor, mask_fn, *, ema_momentum=0.996, nonfinite="fail")`	The same loss as an objective (`JEPAObjective`, §5) for `NNModel.train(..., objective=...)` / `Trainer.train(..., objective=...)`. The shared update engine owns the update, and the EMA target advances once per committed update.

Precision. The ViTNN RMSNorms, like those in the ViTBlocks that JEPAPredictor reuses, reduce in at least FP32 and keep FP64 as FP64. A .double() encoder, its EMA target and the predictor therefore run forward, backward and update_ema in FP64. The tests and float64_vit_predictor_step in example 16 check this.

2. How a step runs¶

1. mask_fn(n_patches, device) -> (ctx_mask_1d, tgt_mask_1d)
2. context_embeds = model.net(x, mask=ctx_mask)
3. with torch.no_grad():
       target_embeds = target_encoder(x)[:, target_positions, :]
4. predicted    = predictor(context_embeds, ctx_positions, tgt_positions)
5. loss         = MSE(predicted, target_embeds)
6. finalize_step(loss)              # NaN guard + grad-clip + optimizer.step()
7. update_ema(model.net, target_encoder, ema_momentum)

3. Quickstart¶

import torch
from torch.utils.data import DataLoader, TensorDataset

from nnx import (
    Devices, Losses, Nets,
    NNModel, NNModelParams, NNOptimParams, NNSchedulerParams, NNTrainParams, NNViTParams, Optims,
    JEPAPredictor,
    build_target_encoder, jepa_train_step_factory, random_block_mask, set_seed,
)

set_seed(0)

# Synthetic 32x32 batch — JEPA is self-supervised, labels are ignored.
loader = DataLoader(
    TensorDataset(torch.randn(64, 3, 32, 32), torch.zeros(64, dtype=torch.long)),
    batch_size=8,
)

# The trainable ViT context encoder: Nets.VIT builds ViTNN from NNViTParams.
model = NNModel(
    net_params=NNViTParams(
        input_dim=3 * 32 * 32, output_dim=64,   # in_channels * image_size**2; d_model
        dropout_prob=0.0, image_size=32, patch_size=4,
        d_model=64, n_layers=4, n_heads=4,
    ),
    params=NNModelParams(net=Nets.VIT, device=Devices.CPU, loss=Losses.CROSS_ENTROPY),
)

# EMA target + predictor (registered under model.net so the optimizer
# trains both encoder and predictor jointly).
target_encoder = build_target_encoder(model.net)
predictor = JEPAPredictor(embed_dim=model.net.d_model,
                          n_patches=model.net.n_patches,
                          predictor_dim=32, n_layers=2, n_heads=2)
model.net.add_module("_jepa_predictor", predictor)

# One random block per step, shared across the batch.
GRID = 32 // 4
def mask_fn(n_p, device):
    return random_block_mask(n_patches=n_p, grid_size=GRID, device=device)

run = model.train(
    params=NNTrainParams(
        n_epochs=3, train_loader=loader,
        optim=NNOptimParams(name=Optims.ADAM, max_lr=5e-4,
                            momentum=(0.9, 0.999), weight_decay=1e-4),
        scheduler=NNSchedulerParams(min_lr=1e-7, factor=0.5,
                                    patience=2, cooldown=1, threshold=1e-3),
    ),
    train_step_fn=jepa_train_step_factory(
        target_encoder=target_encoder,
        predictor=predictor,
        mask_fn=mask_fn,
        ema_momentum=0.996,
    ),
)

Storing and rebuilding the encoder. NNViTParams carries every ViTNN setting into run.yaml and the checkpoints; NNRun.load returns them as the run's net. NNModel.from_checkpoint rebuilds the encoder through Nets.VIT with its trained weights, with no caller-side swap. The recipe above registers the predictor under model.net so the optimizer trains it. This puts the predictor's weights in the same checkpoint. Rebuild the encoder without them:

from nnx import Checkpoints, NNCheckpoint

checkpoint = NNCheckpoint.load(run=run.id, type=Checkpoints.LAST)
encoder = NNModel.from_checkpoint(checkpoint, exclude_submodules=("_jepa_predictor",))

Every other weight still loads strictly, and a name the checkpoint does not hold is refused. Without exclude_submodules, the error names the extra submodules and the option of the entry point you called. The other two rebuild paths take the same option:

from nnx.bundles import export_bundle, reconstruct_bundle

export_bundle(run.id, "jepa-bundle")  # a run bundle of the checkpoint
encoder = reconstruct_bundle("jepa-bundle", exclude_submodules=("_jepa_predictor",)).model

model.save_pretrained("jepa-hub")  # a Hub save of the trained model, predictor attached
encoder = NNModel.from_pretrained("jepa-hub", exclude_submodules=("_jepa_predictor",))

Unlike from_pretrained(..., strict=False), which also accepts missing weights, exclude_submodules drops only the named submodules. An encoder rebuilt from a bundle this way serves inference; it cannot ReconstructedBundle.resume() the run, whose training state includes the predictor.

The full example with an optional CIFAR-10 download lives in examples/16_ijepa_image_plumbing.py.

4. Design notes¶

4.1. Why a separate `ViTNN` instead of reusing `TransformerNN`¶

TransformerNN (the LM path; see docs/lm.md) is a decoder-only LM stack. Its MultiHeadCausalAttention hard-codes a causal mask + RoPE on Q/K; both are wrong for vision tokens (bidirectional attention, learned absolute positions). A parameter in the causal path would be one that vision callers have to remember to flip. Instead, ViTNN ships a sibling _MultiHeadSelfAttention and reuses the parts that generalize (RMSNorm, SwiGLU).

4.2. Why a single block mask per step¶

The reference I-JEPA samples 4 target blocks per image. The shipped random_block_mask samples 1, which is enough for the demo. The mask is shared across the batch so ViTNN's forward can stack the kept patches into a rectangular tensor. For the 4-block recipe, compose four calls inside a custom mask_fn. The only hard constraint is that kept-counts per batch row must agree; otherwise the encoder will raise.

4.3. EMA update is name-keyed, not positional¶

update_ema walks target.named_parameters() and looks each name up on the source. The predictor must be a submodule of model.net, so a single optimizer picks up both encoder and predictor params. The predictor's params then appear on source.named_parameters() but not on the target, and the EMA correctly leaves them alone. Mismatches the other direction (target param missing on source) raise — that's a real bug.

4.4. Loss reporting¶

JEPA has no classification metric. The returned NNEvaluationDataPoint reports the L2 loss in both .loss and .error so BEST checkpoint tracking and ReduceLROnPlateau have a signal to lock onto. The other classification fields stay zero.

4.5. Sharp edges¶

NNModelParams.mixed_precision=True is rejected by NNModel.train with PrecisionUnsupportedError when it resolves to FP16 on CUDA, before any step runs, because the JEPA step is full-precision only. Off CUDA the legacy flag falls back to full precision. AMP requires per-loss scaler bookkeeping that JEPA's custom step doesn't implement.
NNOptimParams.accumulate_grad_batches != 1 raises ValueError at the first step. The imperative JEPA step does not implement gradient accumulation. JEPA's reference recipe uses large batches instead. If you need accumulation (or mixed precision), train with jepa_objective (§5). Its EMA advances once per committed update.
Resume-from-checkpoint restores the EMA target encoder too. jepa_train_step_factory returns a JEPATrainStep, a checkpointable component named jepa.target_encoder (Concepts §14.1). Every checkpoint's training state carries the target encoder's weights. A stateful warm resume puts them back, frozen and in eval mode.

Four epochs and a 2 + 2 split with resume_from_run_id end with identical encoder, predictor, target encoder, optimizer and scheduler state. Build the resumed session's step with the factory as usual; the restore overwrites the freshly copied target.

The component is required. A current checkpoint written without it fails before anything is restored. A checkpoint written before component state existed resumes with a fresh target copy and a warning.

5. Objective mode¶

jepa_objective(target_encoder, predictor, mask_fn, ema_momentum=...) describes the same latent-prediction loss as an objective (Concepts §6.5):

from nnx import jepa_objective

objective = jepa_objective(target_encoder, predictor, mask_fn, ema_momentum=0.996)
model.train(params=NNTrainParams(..., optim=NNOptimParams(..., accumulate_grad_batches=4)), objective=objective)

Loss. One "latent_mse" term per microbatch. It is the sum of squared errors between the predicted and target embeddings at the masked positions. The sum is normalized by the number of those elements (rows × target patches × embedding width). Microbatches with different target masks therefore combine as (sum_1 + sum_2) / (count_1 + count_2). For the same masks (a fixed mask_fn), uneven microbatches give exactly the full-batch update; a random mask_fn draws new masks per microbatch.

The record carries the loss and metrics={"latent_mse": ...}, with no classification fields. Masks are checked before any forward pass. They must be complementary BoolTensor[n_patches] with at least one context and one target patch.
Update. The shared engine owns backward, accumulation, mixed precision, clipping and the optimizer step, with the objective's non-finite policy (nonfinite="fail" or "skip").
EMA, once per committed update. The target advances in the objective's after_update hook. This runs after the window's step, from the updated online weights, and before callbacks see the update event. It never runs per microbatch and never for a skipped window.

ema_momentum is fixed in [0, 1). ema_updates counts this run's EMA steps and continues across a stateful resume. A Trainer with separate encoder and predictor optimizers still advances it once per committed update.
Ownership, checked before any run is reserved. These rules apply:
- The predictor must be a submodule of model.net whose trainable parameters the run's optimizers own exactly once.
- The target encoder must stay outside model.net, share none of its parameters and be owned by no optimizer.
- Every target parameter needs a same-named, same-shaped online one.
- Callbacks that change the net's topology (checkpoint_transforms, e.g. QAT) are refused. The imperative step passed as objective= is refused too.
A callback that renames the net's parameters without declaring it is not detected. An example is torch.nn.utils.prune, which turns weight into weight_orig. update_ema then raises a KeyError naming the parameter at the next committed update.
Checkpoints and resume. JEPAObjective is a checkpointable component named jepa.objective. It holds the spec (loss, momentum), the predictor's reference, the EMA target's weights and the EMA counter. The reference is the predictor's name inside model.net; its weights are saved once, with the net, in LAST and in ModelCheckpoint snapshots.

Two epochs and one epoch plus a stateful resume end in the same encoder, predictor, target and counter. A saved state with a second predictor payload, another predictor reference or another spec is rejected before anything is restored.

examples/16_ijepa_image_plumbing.py --objective trains the demo this way, and its bounded objective_mode() helper checks one EMA step per committed update.