vjepa2.1-vit-base-384

V-JEPA 2.1 ViT-B/16 384 — HuggingFace port

A HuggingFace-format conversion of Meta AI's V-JEPA 2.1 ViT-B/16 video encoder and predictor, operating at 384x384 resolution.

The weights are Meta's, copied without modification. This repository provides the transformers-compatible packaging plus a documented numerical validation against the original implementation.

An equivalent community port already exists (Dev-Jahn/vjepa2.1-vitb-fpc64-384). This repository adds an independently reproduced conversion together with the validation results below.

Provenance

The only structural change is that the fused QKV projection of each attention block is split into separate query / key / value matrices, following the convention used by transformers. This is a re-parameterization, not a change of weights. It is also convenient downstream: PEFT adapters apply to independent q, k, v projections as in the original LoRA formulation, rather than to a fused projection.

If you want to reproduce the comparisons below, note that torch.hub.load against the current main of facebookresearch/vjepa2 does not work: src/hub/backbones.py ships with VJEPA_BASE_URL = "http://localhost:8300" and the real URL commented out above it. Download the checkpoint directly from https://dl.fbaipublicfiles.com/vjepa2/vjepa2_1_vitb_dist_vitG_384.pt and build the reference modules from app/vjepa_2_1/models/, which is what the shipped scripts do.

Architecture

Total: 109.7 M.

Distilled from ViT-G, with a single distillation output (n_output_distillation = 1). The predictor takes the last encoder representation and projects to the 1664-dim teacher space.

The encoder carries a separate patch embedding for single images (img_temporal_dim_size = 1), learnable image/video modality embeddings, and four per-layer normalizations used for hierarchical feature extraction.

The four per-layer normalizations sit at encoder layers [2, 5, 8, 11]. Because this checkpoint was trained with a single distillation output, only the last of them (norms_block.3) carries a trained affine transform; the other three are at PyTorch's default LayerNorm initialization (weight = 1, bias = 0) in Meta's checkpoint as well. They are not dead — a LayerNorm with identity affine still normalizes — and the out_layers path uses them exactly as the reference implementation does. If you need intermediate features with trained affine parameters, use the ViT-g or ViT-G checkpoints.

The attentive pooler used by VJEPA21ForVideoClassification defaults to num_pooler_heads = 16 and num_pooler_layers = 3 — three self-attention blocks plus one cross-attention block, which is num_probe_blocks: 4 and classifier.num_heads: 16 in every frozen-probe config under configs/eval_2_1/. The pooler is always trained from scratch, so this only sets the architecture of the probe you train.

RoPE is interpolatable: positions are rescaled by (pretrained_grid_size - 1) / (n_patches - 1) with pretrained_grid_size = 256 / patch_size, following app/vjepa_2_1/models/utils/modules.py in the reference repository. This is what allows the model to accept resolutions other than 384 and non-square inputs.

Usage

Feature extraction

import torch
from transformers import AutoModel, AutoVideoProcessor

CKPT = "apiantonio/vjepa2.1-vit-base-384"

processor = AutoVideoProcessor.from_pretrained(CKPT, trust_remote_code=True)
model = AutoModel.from_pretrained(CKPT, trust_remote_code=True, dtype=torch.float32).eval().cuda()

# list_of_clips: list of clips, each a list of HxWx3 uint8 frames
inputs = processor(list_of_clips, return_tensors="pt").to("cuda")

with torch.no_grad():
    out = model(**inputs, skip_predictor=True)

feats = out.last_hidden_state       # (B, T/2 * 24 * 24, 768)

pixel_values_videos accepts (B, C, T, H, W), (B, T, C, H, W) — the layout produced by the video processor and by the transformers V-JEPA 2 classes — and (B, T, H, W, C). The channel axis is identified by matching config.in_chans; when more than one axis could match, the channels-first reading wins and a warning is emitted. A clip of exactly 3 frames is the realistic ambiguous case.

Token grid: (frames / 2) x 24 x 24. A 64-frame clip yields 32 x 24 x 24 = 18,432 tokens, so attention cost grows quadratically with clip length — use model.gradient_checkpointing_enable() and small batches for long clips.

Intermediate features

out = model(
    **inputs,
    skip_predictor=True,
    out_layers=model.config.encoder_hierarchical_layers,   # [2, 5, 8, 11]
)
levels = out.multilevel_hidden_states   # 4 tensors, each normalized with its own LayerNorm

This reproduces the out_layers recipe of the reference frozen-evaluation probes. Levels come back in network order regardless of the order they were requested in, so zipping them with encoder_hierarchical_layers is correct. output_hidden_states=True is also available and returns the raw, un-normalized output of the embedding layer plus every transformer layer.

Predictor

out = model(**inputs)                       # runs the predictor
pred = out.predictor_output.last_hidden_state      # (B, N_target, 1664)
ctx = out.predictor_output.context_hidden_state    # (B, N_context, 1664)

Without explicit masks, context_mask and target_mask both default to all tokens. Pass lists of (B, K) index tensors for a real masked-prediction setup, and mask_index to select which of the eight learnable mask tokens is injected — the reference passes the index of the sequence-length group and defaults to 1. A single (context_mask, target_mask) pair is supported; more raises NotImplementedError rather than silently producing wrong shapes.

For the pre-training forward, pass masks to the encoder as well. Tokens are then dropped before the transformer layers, so attention only ever sees the context and RoPE receives the true token positions — this is z = encoder(clips, masks_enc) in the reference, and it is not the same thing as masking the encoder output:

out = model(
    **inputs,
    masks=[context_idx],          # (B, K_ctx)
    context_mask=[context_idx],
    target_mask=[target_idx],     # (B, K_tgt)
)

Classification head

from transformers import AutoModelForVideoClassification

model = AutoModelForVideoClassification.from_pretrained(
    CKPT, trust_remote_code=True, num_labels=101, label2id=label2id, id2label=id2label
)

The attentive pooler and the classifier are initialized from scratch, so the warning about newly initialized weights is expected and the video-classification pipeline tag is for discoverability only — the head must be trained before the model predicts anything meaningful. Multi-label targets work through config.problem_type = "multi_label_classification".

Parameter-efficient fine-tuning

LoRA and DoRA work out of the box with peft. Scope target_modules explicitly, because matching on bare names such as ["query", "key", "value", "proj"] also adapts the predictor, which VJEPA21ForVideoClassification never executes:

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=r".*vjepa21\.encoder\.layer\.\d+\.attention\.(query|key|value|proj)$",
    modules_to_save=["classifier", "pooler"],
)
model = get_peft_model(model, config)

Preprocessing

VJEPA21VideoProcessor reproduces EvalVideoTransform from evals/video_classification_frozen/utils.py: the short side is resized to 384, a 384x384 crop is taken, pixels are scaled to [0, 1] and normalized with ImageNet statistics — mean (0.485, 0.456, 0.406), std (0.229, 0.224, 0.225), confirmed at evals/video_classification_frozen/eval.py:411.

This differs from the stock VJEPA2VideoProcessor in two ways, both of which require a dedicated class. It resizes the short side to exactly the crop size rather than to crop_size * 256 / 224, and it disables torchvision's antialiasing so the resize matches the reference cv2.resize(..., cv2.INTER_LINEAR). The first point is worth being precise about: make_transforms only selects EvalVideoTransform when num_views_per_clip > 1, and falls back to the crop_size * 256 / 224 transform otherwise — but every V-JEPA 2.1 frozen-probe config under configs/eval_2_1/ sets num_views_per_segment: 3, so EvalVideoTransform is the transform the 2.1 evaluations actually use. The second point only shows up when downscaling: on a 1080p frame the antialiased and non-antialiased results differ by up to 144/255 per pixel, which is a visibly different image fed to a model that never saw antialiased inputs. With antialiasing disabled the two pipelines agree to 1/255, the 8-bit quantisation floor, at every resolution tested.

The reference transform takes several sliding crops along the long axis and averages the predictions; this processor takes a single centre crop. Multi-view aggregation, if you want it, is the caller's job.

Frame sampling is left to the caller. The reference configs for V-JEPA 2.1 — under configs/eval_2_1/, not configs/eval/, which holds the V-JEPA 2 settings — use 16 frames with frame_step = 4 and 8 segments for Kinetics-400 and COIN; 32 frames with frame_step = 2 and 4 segments for Jester and Diving48; and for Something-Something v2 either 16x4 on ViT-B or 64x2 with 2 segments on ViT-g and ViT-G. All of them use 3 spatial views per segment. config.frames_per_clip = 64 records the pre-training clip length and is not an inference constraint.

Data pipeline

The processor covers the spatial transform. Which frames a clip is built from, and what a decoder returns for them, are the caller's responsibility and both fail silently, so video_io.py ships a reference implementation of each and tests/test_frame_sampling.py checks any implementation against it.

from video_io import clip_indices, decode_frames

indices = clip_indices(video_len=len(video), frames_per_clip=16, frame_step=4, num_clips=1)
frames = decode_frames(path, indices[0])      # (16, H, W, 3) uint8, RGB
inputs = processor([list(frames)], return_tensors="pt")

Two details are worth stating because they are easy to get wrong and produce no error.

The reference places the frames of a clip with np.linspace across a window of frames_per_clip * frame_step, so the effective stride is fpc*fstp/(fpc-1), not frame_step. Writing range(0, fpc*fstp, fstp) is the intuitive reading of "16 frames, step 4" and it selects a different frame 12 times out of 16. The resulting clip looks entirely plausible.

The decoder must return RGB. cv2.VideoCapture returns BGR, and decode_frames converts explicitly; it also decodes sequentially rather than seeking with CAP_PROP_POS_FRAMES, which snaps to the nearest keyframe on several codecs and quietly returns a neighbouring frame. decord, the backend the reference pipeline uses, is preferred when installed; the shipped decoder was checked against it and agrees exactly, max|Δ| = 0/255.

Long videos and multiple clips

The reference evaluation splits a video into num_segments partitions, takes one clip and num_views_per_segment spatial crops from each, and averages the resulting softmax distributions into a single prediction per video. That is variance reduction for classifying a short video, and it does not transfer to a long one: with the reference settings for this checkpoint, eight segments cover 98.7% of a ten-second video but only 14.2% of a two-minute one, leaving a blind gap of 386 frames — thirteen seconds during which nothing is observed.

clip_indices reproduces that partitioned sampling. When you need a score per position in time rather than one prediction per video, use dense_clip_indices instead, and map the results back onto frames:

from video_io import dense_clip_indices, temporal_coverage, clip_scores_to_frame_scores

clips = dense_clip_indices(video_len, frames_per_clip=16, frame_step=4)   # stride=window
print(temporal_coverage(clips, video_len))     # {'covered_fraction': 1.0, 'max_gap': 0, ...}

frame_scores = clip_scores_to_frame_scores(clips, scores, video_len, reduce="max")

stride controls overlap and therefore temporal resolution. temporal_coverage reports what a set of clips actually looks at, which is worth checking before trusting any per-frame metric.

For video-level predictions, aggregate_predictions reproduces the reference combination — the mean of the softmax outputs, not of the logits. The two are different estimators and can rank classes differently; on a two-view example they give [0.52, 0.19, 0.29] against [0.77, 0.10, 0.13].

Note that dense extraction has a storage cost. A two-minute video at 30 fps needs 57 non-overlapping clips, and a 16-frame clip at 384 is 4608 tokens, so caching token-level features costs about 0.8 GB at this hidden size. Pooling to one vector per clip is four orders of magnitude smaller, but forecloses any probe that consumes the token sequence.

Validation

All tests use float32 and torch.no_grad(). Where the port and the reference are compared, the attention kernel is matched on both sides: the reference encoder is built with the same use_sdpa setting as the port, and the reference predictor always runs SDPA because VisionTransformerPredictor.__init__ has no use_sdpa parameter and swallows it through **kwargs. Comparing across kernels measures the kernel, not the port.

1. Weight-level provenance

Every tensor in the published safetensors was traced back to a tensor in the official checkpoint and compared element-wise.

The last row matters: the conversion script calls load_state_dict(..., strict=False), which silently leaves unmatched parameters randomly initialized. Zero orphans means no such parameter survived — for the predictor as well as the encoder. Loading through AutoModel and through AutoModelForVideoClassification was also checked with output_loading_info=True: zero missing, zero unexpected, zero mismatched keys on the first, and only pooler.* and classifier.* newly initialized on the second.

2. Forward parity on the published weights, full resolution

The same input tensor was passed through Meta's official modules, loaded with the official checkpoint, and through this port, at 384x384.

* cosine_similarity computes a·b / (‖a‖·‖b‖); numerator and denominator follow different rounding paths, so the metric can print as 0.999999 even for bit-identical tensors. The absolute differences are exactly zero.

Summary. Bit-exact agreement with the reference implementation, encoder and predictor, on the published weights at full resolution. The predictor comparison is not vacuous here: the mask tokens of this checkpoint have a maximum magnitude of 5.738e-01, so the mask-token lookup is genuinely exercised. A comparison run with zero_init_mask_tokens=True on both sides would pass even with a broken index, because every mask token would be exactly zero.

Reproduce with verify_vjepa21_port.py --repo apiantonio/vjepa2.1-vit-base-384 --checks ABCDEFG, or with verify_big_models.py for the ViT-g and ViT-G variants, which never holds the reference and the port in memory at the same time. --frames controls the clip length; 32 and 64 reproduce the same zero.

3. Implementation parity, code paths not covered above

To cover the paths §2 does not reach, reference modules were built directly from app/vjepa_2_1/models/ and their weights copied into the port, then outputs compared. These runs use reduced hidden sizes so the suite stays fast: they establish that the code paths are equivalent, which is width-independent, and complement the full-resolution published-weight results in §2.

Both distillation regimes are covered: n_output_distillation = 1 is the ViT-B and ViT-L configuration, 4 the ViT-g and ViT-G one, where the encoder emits concatenated levels and the predictor fuses them through a 2-layer MLP.

The remaining sub-microvolt residuals on the predictor rows are a kernel difference, not an algorithmic one: this suite runs the port in eager mode while the reference predictor always runs SDPA. Matching the kernel removes them entirely, which is why §2 reports exactly zero. The port reorders tokens with gather where the reference uses stack; that is a pure permutation and contributes nothing, as the full-resolution result confirms.

The masked-encoder row is worth calling out. Passing masks drops tokens before the transformer layers, so attention sees only the context and RoPE receives the true positions. A control test checks that this is not equal to gathering the encoder output afterwards — otherwise the parity result would be meaningless.

tests/test_end_to_end_pipeline.py closes the loop with the preprocessing: one clip is pushed through the reference Resize -> crop -> ClipToTensor -> Normalize and through VJEPA21VideoProcessor, at 240p, 480p, 720p and 1080p. All four land on max|Δ| = 1.751e-02, which is one 8-bit grey level divided by the smallest normalisation std — the floor two different bilinear implementations can reach. A control test forces antialiasing back on and confirms the comparison then fails by 2.65, so the tolerance is tight enough to catch the regression it was written for.

Pushed through the model, that pixel-level residual costs:

This variant amplifies the preprocessing residual more than the larger ones: the same pixel-level difference of 1.751e-02 costs roughly twice the relative error here that it does on ViT-L, ViT-g and ViT-G, whose worst case is 7.4e-03 relative and 0.999148 cosine. That is expected — a narrower representation has less redundancy to absorb input noise — but it means the choice of resize backend matters most on the smallest checkpoint, which is the one most likely to be used for quick experiments.

4. Functional test suite

122 tests ship with the repository, green on both transformers==4.57.1 and transformers==5.14.1.

The first file covers configuration properties and validation, input layouts and the ambiguous-layout warning, hidden states and attentions, multi-level extraction and its ordering, the image branch, variable resolution, sdpa-vs-eager agreement, the masked encoder forward, the predictor with default, explicit and pre-masked contexts, mask_index selection, classification loss for both single- and multi-label targets, gradient checkpointing parity on the encoder and the predictor path, rejection of unknown keyword arguments, low-precision weight loading, the video processor's geometry and normalization constants, and a full save_pretrained / from_pretrained round-trip through the Auto classes. It needs neither weights nor the reference repository.

5. Properties relevant to downstream pipelines

tests/test_thesis_robustness.py checks assumptions an experimental pipeline tends to make implicitly. Measured on this checkpoint, float32 on an NVIDIA GPU unless stated otherwise.

Reduced precision, split by attention kernel, because the two do not behave the same way:

Repeated forwards of the same shape are bit-identical, so a cached feature is reproducible from its input. Changing the batch size is not guaranteed to be bit-exact on GPU: cuBLAS and the fused attention kernels select tiling and reduction order as a function of the input shape. Where the residual appears it is around 1e-05 relative and irrelevant to a probe, but it means a feature store should be built with a fixed batch size.

Read absolute differences with care: they track the magnitude of the activations, not the quality of the agreement, which is why the relative error is the meaningful number.

The last two rows of the first table are the reparameterisation property that makes LoRA and DoRA free at inference time. They are measured on a small model, since exact merging follows from the algebra rather than from the weights. DoRA is checked explicitly and not assumed to follow from LoRA: it decomposes the update into direction and magnitude, so merging has to fold the magnitude vector back in as well.

To refresh the first two rows on this checkpoint:

VJEPA21_CKPT=apiantonio/vjepa2.1-vit-base-384 python -m pytest tests/test_thesis_robustness.py -s -q

Precision guidance

float16 is several times more accurate than bfloat16 for inference with this model — ten mantissa bits against eight — and the gap is consistent across the whole family. But float16 is only safe with a fused attention kernel. The model was pre-trained in bfloat16, which shares float32's exponent range, so nothing ever constrained the attention logits to stay inside float16's maximum of 65504. SDPA and flash attention accumulate the softmax in float32 and are unaffected; the eager kernel computes q·kᵀ in float16 natively and can overflow to NaN. The residual stream itself stays well inside range — it peaks around 430 — so the failure is specific to the attention logits.

On this checkpoint the eager kernel does overflow: float16 inference returns NaN. The same happens on ViT-B, ViT-L and ViT-g; ViT-G is the one variant tested where it does not. It is a property of the individual checkpoint's activation magnitudes, not of the architecture.

Since output_attentions=True forces the eager kernel — it is the only one that materializes attention probabilities — extracting attention maps in float16 is the concrete way to hit this.

For feature extraction, prefer float32, or float16 with SDPA. bfloat16 costs a worst-case per-token cosine similarity of 0.998602; if features feed a cosine-similarity or nearest-neighbour decision rule, that noise enters the decision directly. Storing float32 features as float16 is nearly free by comparison and halves the cache, which keeps the inference-precision and storage-precision decisions separate.

For training, the usual bfloat16 argument still applies: the wider exponent range is safer for gradients. These are two separate decisions and it is worth making them separately.

Compatibility

Requires transformers >= 4.50; tested on 4.57.1 and 5.14.1. sdpa and flash_attention_2 are supported; requesting output_attentions=True falls back to the eager kernel, the only one that materializes attention probabilities — see the precision guidance before combining that with float16. peft 0.19.1 was verified for LoRA and DoRA, including merge_and_unload and the resulting parameter count. Weights held in bfloat16 or float16 are supported.

Not verified

Stated explicitly so the scope of the validation above is not overread:

• Downstream benchmarks. No evaluation numbers from the V-JEPA 2.1 release were reproduced. This port is validated for numerical equivalence, not for task performance.
• Attentive pooler initialization. The pooler reproduces the reference AttentivePooler structurally, but the reference calls _rescale_blocks() after initialization, dividing attn.proj.weight and mlp.fc2.weight of layer i by sqrt(2*(i+1)). That is not reproduced here. It touches no published weight — the pooler is always trained from scratch — but a probe trained here starts from a different initialization scale than Meta's.
• Your own data pipeline. video_io.py reproduces the reference sampler and decoder and is tested against them, but if you write your own, only the contract tests in tests/test_frame_sampling.py stand between you and a silent mismatch.
• Aggregation policy. aggregate_predictions reproduces the reference combination and dense_clip_indices covers the whole timeline, but which policy suits a given task — mean over clips, max, top-k — is a modelling decision that is neither made nor evaluated here.
• Multiple mask pairs in the predictor. Only a single (context_mask, target_mask) pair is supported.
• Reduced precision downstream. bf16 and fp16 deviations are measured above, but their effect on any downstream metric was not assessed.

License and attribution

The weights are released by Meta AI under the license of the original V-JEPA 2 release (MIT). This repository redistributes them unmodified and adds no additional restrictions.

Credit where it is due:

• Meta AI — the V-JEPA 2 / V-JEPA 2.1 research and the pretrained weights.
• Dev-Jahn — the original transformers-compatible modeling code and the conversion script, which this repository builds on.

The contribution specific to this repository is the conversion of the B variant, the revised modeling code and video processor described above, and the numerical validation documented here.

Citation

@article{assran2025vjepa2,
  title={V-JEPA~2: Self-Supervised Video Models Enable Understanding, Prediction and Planning},
  author={Assran, Mahmoud and Bardes, Adrien and Fan, David and Garrido, Quentin and Howes, Russell and
Komeili, Mojtaba and Muckley, Matthew and Rizvi, Ammar and Roberts, Claire and Sinha, Koustuv and Zholus, Artem and
Arnaud, Sergio and Gejji, Abha and Martin, Ada and Robert Hogan, Francois and Dugas, Daniel and
Bojanowski, Piotr and Khalidov, Vasil and Labatut, Patrick and Massa, Francisco and Szafraniec, Marc and
Krishnakumar, Kapil and Li, Yong and Ma, Xiaodong and Chandar, Sarath and Meier, Franziska and LeCun, Yann and
Rabbat, Michael and Ballas, Nicolas},
  journal={arXiv preprint arXiv:2506.09985},
  year={2025}
}

@article{murlabadia2026vjepa2_1,
  title={V-JEPA 2.1: Unlocking Dense Features in Video Self-Supervised Learning},
  author={Mur-Labadia, Lorenzo and Muckley, Matthew and Bar, Amir and Assran, Mahmoud and
Sinha, Koustuv and Rabbat, Michael and LeCun, Yann and Ballas, Nicolas and Bardes, Adrien},
  journal={arXiv preprint arXiv:2603.14482},
  year={2026}
}

Please cite the original V-JEPA 2 work. If the format conversion itself was useful, a link back to Dev-Jahn/vjepa2-hf is appreciated.