V-JEPA 2 context encoder, EMA target encoder and predictor.
Clips use (B, T, C, H, W). A plain call returns the target encoder's
token map. Supplying both context_indices and target_indices runs
one masked representation-prediction step, which is useful for training or
checking the predictor in isolation. Indices are flattened tubelet
positions in row-major (time, height, width) order.
Used by 2
Constructors
1Properties
1Instance methods
7Return mean-pooled target-encoder representations (B, dim).
forward(x: Tensor, context_indices: Tensor | None = None, target_indices: Tensor | None = None, mask_index: int = 1)Run representation extraction or a masked prediction step.
Return the mean target representation for backbone consumers.
The target encoder's momentum at a training step.
Parameters
stepinttotal_stepsintReturns
floatLinear from config.ema[0] toward config.ema[1] over
total_steps * config.ema_schedule_scale steps.
Notes
The released pretraining configuration sets both ends of ema
to 0.99925, so this is a constant for the published runs —
unlike V-JEPA 1, which ramps. The schedule is still built from
the two endpoints, because a method that takes a step and a
horizon and reads neither cannot be given a different recipe.
Return the full (B, N, dim) token map.
trainable_parameters()Return context-encoder and predictor parameters only.
Update the non-gradient target encoder after an optimiser step.