V-JEPA 2: representaciones predictivas (Checkpoint C4·a)

Objetivo de maestría

usar un world model predictivo congelado como extractor de representaciones y demostrar, con un attentive probe, que esas representaciones baten a un baseline supervisado en una tarea downstream. Esto prueba que entiendes qué hace valioso a un world model: no genera bonito, representa bien.

2.1La idea operativa

V-JEPA 2 se entrenó prediciendo en latente sobre vídeo masivo, sin etiquetas. El resultado es un encoder cuyas features capturan estructura física y de movimiento. La forma de usarlo (y de evaluarlo) es el attentive probe: congelas el encoder y entrenas encima una cabeza ligera (en el paper, 4 bloques transformer, el último con cross-attention y un query token aprendido) para una tarea concreta. Si esa cabeza ligera sobre features congeladas bate a un modelo supervisado entrenado end-to-end con los mismos datos, has demostrado el valor de la representación.

2.2Laboratorio L2.1 — Inferencia y extracción de embeddings

bash

1# Setup del repo oficial
2git clone https://github.com/facebookresearch/vjepa2 && cd vjepa2
3uv venv && source .venv/bin/activate
4uv pip install -e .   # instala dependencias del repo
5
6# Descarga pesos (elige tamaño según VRAM; ViT-L entra holgado, ViT-g justo)
7wget https://dl.fbaipublicfiles.com/vjepa2/vitl-256.pt -P ./ckpts

python

1# lab_n4l2_embed.py — extrae embeddings de vídeo con V-JEPA 2 congelado
2# (también disponible vía HuggingFace transformers: AutoModel "facebook/vjepa2-vitl-fpc64-256")
3import torch
4from transformers import AutoVideoProcessor, AutoModel
5
6DEVICE = "cuda"
7repo = "facebook/vjepa2-vitl-fpc64-256"
8processor = AutoVideoProcessor.from_pretrained(repo)
9model = AutoModel.from_pretrained(repo, torch_dtype=torch.float16).to(DEVICE).eval()
10
11@torch.no_grad()
12def embed(video_frames):                      # video_frames: lista de PIL/np frames de un clip
13    inputs = processor(video_frames, return_tensors="pt").to(DEVICE)
14    out = model(**inputs)
15    # last_hidden_state: [B, num_patches_temporales*espaciales, hidden] -> features congeladas
16    return out.last_hidden_state              # esto es lo que alimenta al probe
17
18# Verifica: dos clips de la MISMA acción deberían tener embeddings más cercanos
19# (coseno) que dos clips de acciones distintas. Es tu sanity check de que las
20# representaciones capturan semántica de movimiento.

Líneas no triviales:

torch.no_grad() + .eval(): el encoder está congelado; no calculamos gradientes sobre él → ahorro de VRAM masivo (solo forward).
last_hidden_state: las features por patch espacio-temporal. El probe las agregará con cross-attention en un vector por clip.
float16: ViT-L en fp16 entra cómodo en 32 GB; ViT-g requiere más cuidado (offload o batch 1).

2.3Laboratorio L2.2 — Entrenar el attentive probe (el experimento de C4·a)

Replicamos el patrón del paper: backbone congelado + probe ligero entrenable. El reto genérico: una tarea de clasificación de acción en vídeo con un dataset público pequeño (p.ej. un subset de UCF101 o Jester), donde compararemos contra un baseline supervisado.

python

1# lab_n4l2_probe.py — attentive probe sobre features congeladas de V-JEPA 2
2import torch, torch.nn as nn
3
4class AttentiveProbe(nn.Module):
5    """Cabeza ligera: un query token aprendido agrega las features del clip por cross-attention."""
6    def __init__(self, dim, n_classes, n_heads=8, n_blocks=4):
7        super().__init__()
8        self.query = nn.Parameter(torch.randn(1, 1, dim))          # token de consulta aprendible
9        self.blocks = nn.ModuleList([
10            nn.TransformerEncoderLayer(dim, n_heads, batch_first=True) for _ in range(n_blocks-1)])
11        self.cross = nn.MultiheadAttention(dim, n_heads, batch_first=True)  # último bloque: cross-attn
12        self.head = nn.Linear(dim, n_classes)
13
14    def forward(self, feats):                  # feats: [B, N, dim] features congeladas del encoder
15        x = feats
16        for blk in self.blocks:
17            x = blk(x)                          # self-attention sobre los patches
18        q = self.query.expand(x.size(0), -1, -1)
19        pooled, _ = self.cross(q, x, x)         # el query token "lee" la info relevante del clip
20        return self.head(pooled.squeeze(1))     # logits de clase
21
22# Bucle de entrenamiento: SOLO el probe tiene gradientes; el encoder V-JEPA 2 está congelado.
23def train_probe(encoder, probe, loader, epochs=10, lr=1e-3):
24    opt = torch.optim.AdamW(probe.parameters(), lr=lr)   # <- solo params del probe
25    lossf = nn.CrossEntropyLoss()
26    for ep in range(epochs):
27        for frames, label in loader:
28            with torch.no_grad():
29                feats = encoder(frames)          # forward congelado (sin grad)
30            logits = probe(feats)                # forward del probe (con grad)
31            loss = lossf(logits, label.cuda())
32            opt.zero_grad(); loss.backward(); opt.step()
33    return probe

Líneas no triviales:

El query token + cross-attention es el corazón del attentive probe: en vez de hacer mean-pooling de las features (que pierde info), un token aprendible "pregunta" qué patches importan para la tarea. Es lo que usa el paper.
optim sobre probe.parameters() solo: el encoder no se entrena → estás midiendo la calidad de las representaciones, no re-entrenando el modelo. Esa es la esencia del probing.
with torch.no_grad() alrededor del encoder: confirma que el backbone está congelado y ahorra VRAM/cómputo.

2.4El baseline a batir y el criterio de C4(a)

Entrena dos modelos en la misma tarea/datos:

Probe sobre V-JEPA 2 congelado (lo de arriba).
Baseline supervisado: un modelo de vídeo entrenado end-to-end desde cero (o un ResNet/ViT de imagen aplicado por frames) con el mismo dataset y presupuesto de datos similar.

Compara accuracy. Si el probe (que solo entrena una cabeza ligera sobre features congeladas) iguala o supera al baseline supervisado entrenado entero, has demostrado el valor del world model: la representación pre-entrenada vale más que entrenar desde cero con tus datos. Ese es el checkpoint.

2.5CHECKPOINT C4(a) — criterio de aprobado

Extraes embeddings de V-JEPA 2 congelado y verificas que capturan semántica (clips de la misma acción más cercanos).
Entrenas un attentive probe sobre el encoder congelado para una tarea downstream.
El probe bate (o iguala) a un baseline supervisado entrenado end-to-end con datos comparables, con accuracy reportada y reproducible.
Sabes explicar por qué predecir en latente produce representaciones útiles y por qué el probing mide la representación, no el modelo.

Rúbrica: Nivel 3 si reproduces lo anterior; Nivel 4 si exploras V-JEPA 2-AC (action-conditioned) para una tarea de planificación, o usas las features para algo no trivial (RAG visual, detección).

2.6Ejercicios

E1. Varía el nº de bloques del probe (1, 2, 4). ¿Cuánta accuracy ganas? ¿A partir de cuántos bloques estás "entrenando un modelo" en vez de "probando representaciones"?

E2. Compara mean-pooling de features vs el query token + cross-attention. ¿Cuánto aporta el attentive probe sobre el pooling ingenuo?

E3. Toma features de V-JEPA 2 y de un encoder supervisado (p.ej. un ViT de ImageNet) y entrena el mismo probe sobre ambas. ¿Cuál representa mejor el movimiento? (V-JEPA debería ganar en tareas de movimiento como SSv2/Jester.)

2.7Trampas comunes

Descongelar el encoder sin querer → ya no mides la representación, y haces OOM.
Comparar contra un baseline con muchos más datos/cómputo → comparación injusta.
Usar ViT-g sin gestionar VRAM (offload/batch 1) → OOM.

2.8Referencias

facebookresearch/vjepa2 (repo, vjepa2_demo.py, scripts de probe). Paper V-JEPA 2 (Assran et al. 2025). Modelos HF facebook/vjepa2-*. Datasets: UCF101, Jester, SSv2.