Publish PRISM weights and modeling code

README.md

@@ -0,0 +1,130 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: feature-extraction
+base_model:
+  - google/siglip2-so400m-patch14-384
+  - Qwen/Qwen3-Embedding-0.6B
+tags:
+  - video
+  - representation-learning
+  - view-invariant
+  - cross-view
+  - egocentric
+  - egoexo4d
+  - emnlp2026
+---
+
+# PRISM — EgoExo4D
+
+**P**redictive **R**ecomposition v**I**a **S**emantic Latent Deco**M**position for
+View-invariant Video Representation Learning — **EMNLP 2026 (Main Conference)**.
+
+[Paper][paper] · [Project page][project] · [Code][code]
+
+PRISM encodes a video clip into a **view-invariant** embedding: cosine similarity between
+two PRISM embeddings measures *what is happening*, not *where the camera is*. It gets there
+by refusing to produce a single embedding — each clip is decomposed into a view-invariant
+latent `z_vi` and a view-variant latent `z_vv`, and the split is enforced by recomposing
+across clips (`z_vi` of clip A with `z_vv` of clip B) and aligning the result with the
+matching recomposed language description. Semantics that leaked between the two streams
+cannot survive that swap.
+
+This checkpoint is trained on **EgoExo4D**.
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModel, AutoImageProcessor
+
+model = AutoModel.from_pretrained(
+    "litcoderr/prism", trust_remote_code=True, dtype=torch.bfloat16
+).eval().cuda()
+proc = AutoImageProcessor.from_pretrained(model.config.vision_backbone_name)
+
+frames = [...]                                              # list[PIL.Image], sampled at 4 fps
+pixel_values = proc(images=frames, return_tensors="pt").pixel_values[None]
+pixel_values = pixel_values.to("cuda", torch.bfloat16)      # (1, T, 3, 384, 384)
+valid_mask = torch.ones(pixel_values.shape[:2], dtype=torch.bool, device="cuda")
+
+emb = model.encode(pixel_values, valid_mask)                # (1, 512) — L2-normalized
+```
+
+| | |
+|:---|:---|
+| `model.encode(pixel_values, valid_mask)` | `(B, 512)` L2-normalized clip embedding — mean-pooled `z_vi` over valid frames |
+| `model.encode_streams(pixel_values, valid_mask)` | `{"z_vi_seq", "z_vv_seq"}`, each `(B, T, 512)` per-frame |
+
+**Inputs.** `pixel_values` is `(B, T, 3, 384, 384)`, frames sampled at 4 fps, up to
+`T = 128` (32 s), preprocessed by the SigLIP2 image processor. `valid_mask` is `(B, T)`
+bool marking real frames in a padded batch. Both encode paths run under `torch.no_grad()`
+and use the EMA target encoder θ̄.
+
+For batched encoding of a clip manifest, see `scripts/encode.sh` in the [code repo][code].
+
+## Results
+
+Cross-view semantic alignment (R@5 / top-1 / avg):
+
+| Method | <sub>EgoExo4D</sub><br>ego→exo | <br>exo→ego | <br>avg | <br>Recog. top-1 | <sub>EgoExoLearn</sub><br>Assoc. avg |
+|:---|:--:|:--:|:--:|:--:|:--:|
+| SigLIP2 | 35.1 | 19.7 | 27.4 | 13.9 | 26.6 |
+| ViewpointRosetta <sub>*prior SoTA*</sub> | 58.1 | 47.2 | 52.7 | 34.5 | 32.3 |
+| **PRISM** | **75.9** | **50.3** | **63.1** | **41.9** | **43.9** |
+
+Zero-shot transfer to AE2 (never trained on AE2 videos) beats AE2-trained models on
+temporal ordering — Kendall's τ **0.601** vs 0.562, phase progression **0.647** vs 0.480.
+Full tables and ablations are in the [paper][paper].
+
+## Architecture
+
+| Component | |
+|:---|:---|
+| Vision backbone | `google/siglip2-so400m-patch14-384` — **frozen**, not stored here |
+| Text backbone | `Qwen/Qwen3-Embedding-0.6B` — **frozen**, not stored here |
+| Decompositional Encoder θ | 4-layer Q-Former (2 queries → `z_vi`, `z_vv`) + 12-layer causal temporal stack per stream |
+| Compositional Latent Predictor φ | 4-layer causal transformer over `concat(z_vv, z_vi)`, with `cls_head` / `vi_head` / `vv_head` |
+| Target encoder θ̄ | EMA of θ, decay 0.998 |
+| Embedding dim | 512 |
+
+This repo holds **trained weights only** (0.97 GB, fp32): θ, φ, θ̄, and the logit scale. The
+two backbones are re-downloaded from the Hub when the model is constructed, so nothing about
+them is duplicated here.
+
+## Training
+
+EgoExo4D, 6 epochs, batch 4 × 7 GPUs (DDP), constant-with-warmup lr 7e-5 (10% warmup),
+weight decay 0.01, grad-clip 1.0, bf16, seed 42. Video sampled at 4 fps, clips capped at
+32 s / 128 frames, 384×384. Objective `L = 1.0 · L_decomp + 0.5 · L_temp`, InfoNCE
+all-gathered across ranks, sliding-shift augmentation on.
+
+Per-clip view-invariant / view-variant captions are a *provided input*; recomposed captions
+are generated online by a local vLLM server running `Qwen/Qwen3-1.7B`. Full recipe in the
+[code repo][code].
+
+## Limitations
+
+Trained on EgoExo4D — skill-centric, mostly indoor, ego/exo camera pairs. Domains far from
+that distribution are untested. Clips longer than 32 s are truncated to the first 128
+sampled frames. The training pipeline depends on decoupled view-invariant / view-variant
+captions, which most datasets do not ship.
+
+## Citation
+
+```bibtex
+@inproceedings{chee2026prism,
+  title     = {{PRISM}: Predictive Recomposition via Semantic Latent Decomposition
+               for View-invariant Video Representation Learning},
+  author    = {Chee, Youngchae and Lee, Hosu and Park, Sungjune and Kim, Junho and Ro, Yong Man},
+  booktitle = {Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing (EMNLP)},
+  year      = {2026}
+}
+```
+
+MIT licensed. The frozen backbones keep their own licenses.
+
+<!-- swap these when the links go live -->
+[paper]:   https://github.com/litcoderr/prism
+[project]: https://github.com/litcoderr/prism
+[code]:    https://github.com/litcoderr/prism

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "PRISMModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_prism.PRISMConfig",
+    "AutoModel": "modeling_prism.PRISMModel"
+  },
+  "d_z": 512,
+  "ema_decay": 0.998,
+  "infonce_all_gather": true,
+  "lambda_decomp": 1.0,
+  "lambda_temp": 0.5,
+  "logit_scale_init": 2.6592,
+  "max_frames": 128,
+  "mlp_ratio": 4.0,
+  "model_type": "prism",
+  "num_heads": 8,
+  "predictor_depth": 4,
+  "qformer_depth": 4,
+  "sliding_shift_aug": true,
+  "temporal_depth": 12,
+  "text_backbone_name": "Qwen/Qwen3-Embedding-0.6B",
+  "transformers_version": "4.57.6",
+  "use_ema": true,
+  "vision_backbone_name": "google/siglip2-so400m-patch14-384"
+}

configuration_prism.py

@@ -0,0 +1,54 @@
+"""PRISM configuration.
+
+Defaults reproduce the released checkpoint (SigLIP2-so400m / Qwen3-Embedding-0.6B).
+"""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class PRISMConfig(PretrainedConfig):
+    model_type = "prism"
+
+    def __init__(
+        self,
+        # ---- frozen backbones (loaded from the Hub by name) ----
+        vision_backbone_name: str = "google/siglip2-so400m-patch14-384",
+        text_backbone_name: str = "Qwen/Qwen3-Embedding-0.6B",
+        # ---- architecture ----
+        d_z: int = 512,
+        qformer_depth: int = 4,
+        temporal_depth: int = 12,
+        predictor_depth: int = 4,
+        num_heads: int = 8,
+        mlp_ratio: float = 4.0,
+        max_frames: int = 128,
+        logit_scale_init: float = 2.6592,
+        # ---- training objective ----
+        lambda_decomp: float = 1.0,
+        lambda_temp: float = 0.5,
+        infonce_all_gather: bool = True,
+        sliding_shift_aug: bool = True,
+        # ---- EMA target encoder ----
+        use_ema: bool = True,
+        ema_decay: float = 0.998,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vision_backbone_name = vision_backbone_name
+        self.text_backbone_name = text_backbone_name
+        self.d_z = d_z
+        self.qformer_depth = qformer_depth
+        self.temporal_depth = temporal_depth
+        self.predictor_depth = predictor_depth
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.max_frames = max_frames
+        self.logit_scale_init = logit_scale_init
+        self.lambda_decomp = lambda_decomp
+        self.lambda_temp = lambda_temp
+        self.infonce_all_gather = infonce_all_gather
+        self.sliding_shift_aug = sliding_shift_aug
+        self.use_ema = use_ema
+        self.ema_decay = ema_decay

ema.py

@@ -0,0 +1,48 @@
+"""EMA target encoder helpers.
+
+The target encoder ``θ̄`` is a deep copy of the (trainable) Decompositional
+Encoder ``θ``. It receives no gradient and is updated in place after every
+optimizer step:
+
+    θ̄ ← α · θ̄ + (1 - α) · θ
+
+It supplies the (stable) prediction targets for the temporal objective
+``L_temp`` and is the encoder used at inference (see the paper, §3.3).
+"""
+
+from __future__ import annotations
+
+import copy
+
+import torch
+from torch import nn
+
+
+def make_ema_copy(module: nn.Module) -> nn.Module:
+    """Deep-copy ``module`` for EMA use: ``requires_grad=False``, eval, fp32."""
+    ema = copy.deepcopy(module)
+    for p in ema.parameters():
+        p.requires_grad = False
+        p.data = p.data.float()
+    ema.eval()
+    return ema
+
+
+@torch.no_grad()
+def update_ema(ema_module: nn.Module, online_module: nn.Module, decay: float) -> None:
+    """In place ``θ̄ ← decay·θ̄ + (1-decay)·θ`` over matching parameters.
+
+    Online params may be bf16/fp16/fp32 (mixed-precision keeps fp32 masters);
+    EMA params stay fp32 for numerical stability across many steps. Buffers
+    (e.g. sinusoidal positional embeddings) are not trained and not updated.
+    """
+    for p_ema, p in zip(ema_module.parameters(), online_module.parameters()):
+        p_ema.data.mul_(decay).add_(p.data.float(), alpha=1.0 - decay)
+
+
+@torch.no_grad()
+def sync_ema_from_online(ema_module: nn.Module, online_module: nn.Module) -> None:
+    """Copy online → EMA (fp32). Used to warm-start the target encoder when a
+    checkpoint lacks EMA weights."""
+    for p_ema, p in zip(ema_module.parameters(), online_module.parameters()):
+        p_ema.data.copy_(p.data.float())

encoder.py

@@ -0,0 +1,150 @@
+"""Decompositional Encoder θ.
+
+Decomposes a video into a view-invariant stream ``z_vi`` and a view-variant
+stream ``z_vv`` (paper §3.1). Two sub-modules:
+
+  - ``QFormer``               : per-frame BLIP-2-style Q-Former with two learned
+                                queries (Q_vi, Q_vv) that attend to a single
+                                frame's frozen patch tokens.
+  - ``CausalTemporalEncoder`` : two parallel causal temporal streams (one per
+                                factor). Within a stream, frame ``t`` attends to
+                                frames ``≤ t``; the streams never attend to each
+                                other (cross-factor mixing is deferred to φ).
+
+``forward(patches) -> (z_vi, z_vv)`` with each ``(B, T, d_z)``.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .layers import QFormerBlock, TemporalBlock, build_sin_pos_embed, causal_mask
+
+
+class QFormer(nn.Module):
+    """Per-frame Q-Former with N=2 queries: query 0 → z_vi, query 1 → z_vv."""
+
+    def __init__(
+        self,
+        d_z: int = 512,
+        d_kv: int = 1024,
+        depth: int = 4,
+        num_heads: int = 8,
+        mlp_ratio: float = 4.0,
+    ):
+        super().__init__()
+        self.d_z = d_z
+        self.d_kv = d_kv
+        self.queries = nn.Parameter(torch.zeros(1, 2, d_z))
+        nn.init.normal_(self.queries, std=0.02)
+        self.blocks = nn.ModuleList(
+            [QFormerBlock(d_z, d_kv, num_heads, mlp_ratio) for _ in range(depth)]
+        )
+        self.final_norm = nn.LayerNorm(d_z)
+
+    def forward(self, patches: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """patches: ``(B, T, P, d_kv)`` → ``(z_vi, z_vv)`` each ``(B, T, d_z)``."""
+        B, T, P, d_kv = patches.shape
+        assert d_kv == self.d_kv, f"expected d_kv={self.d_kv}, got {d_kv}"
+        kv = patches.reshape(B * T, P, d_kv)
+        q = self.queries.expand(B * T, -1, -1).contiguous()
+        for block in self.blocks:
+            q = block(q, kv)
+        q = self.final_norm(q)
+        z_vi = q[:, 0, :].reshape(B, T, self.d_z)   # query 0 → view-invariant
+        z_vv = q[:, 1, :].reshape(B, T, self.d_z)   # query 1 → view-variant
+        return z_vi, z_vv
+
+
+class CausalTemporalEncoder(nn.Module):
+    """Two parallel causal temporal streams over z_vi and z_vv.
+
+    Each stream has its own (non-shared) stack of ``TemporalBlock``s. A shared
+    sinusoidal positional embedding is added before the blocks; a boolean
+    ``key_padding_mask`` ``(B, T)`` blocks padded frames in both streams.
+    """
+
+    def __init__(
+        self,
+        d_z: int = 512,
+        max_frames: int = 128,
+        depth: int = 12,
+        num_heads: int = 8,
+        mlp_ratio: float = 4.0,
+    ):
+        super().__init__()
+        self.d_z = d_z
+        self.max_frames = max_frames
+        self.register_buffer(
+            "pos_embed", build_sin_pos_embed(max_frames, d_z), persistent=False
+        )
+        self.blocks_vi = nn.ModuleList(
+            [TemporalBlock(d_z, num_heads, mlp_ratio) for _ in range(depth)]
+        )
+        self.blocks_vv = nn.ModuleList(
+            [TemporalBlock(d_z, num_heads, mlp_ratio) for _ in range(depth)]
+        )
+        self.norm_vi = nn.LayerNorm(d_z)
+        self.norm_vv = nn.LayerNorm(d_z)
+
+    def _pos_embed(self, T: int, device, dtype) -> torch.Tensor:
+        # Training never exceeds max_frames; T > max_frames only happens when
+        # encoding a full native-fps video longer than the cap, where we extend
+        # the (deterministic) sinusoidal PE on the fly.
+        if T <= self.max_frames:
+            return self.pos_embed[:, :T, :]
+        return build_sin_pos_embed(T, self.d_z).to(device=device, dtype=dtype)
+
+    def forward(
+        self,
+        z_vi_seq: torch.Tensor,
+        z_vv_seq: torch.Tensor,
+        key_padding_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """z_vi_seq, z_vv_seq: ``(B, T, d_z)``; ``key_padding_mask`` True = padded."""
+        B, T, _ = z_vi_seq.shape
+        pos = self._pos_embed(T, z_vi_seq.device, z_vi_seq.dtype)
+        zi = z_vi_seq + pos
+        zv = z_vv_seq + pos
+        attn_mask = causal_mask(T, device=z_vi_seq.device)
+        for blk in self.blocks_vi:
+            zi = blk(zi, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
+        for blk in self.blocks_vv:
+            zv = blk(zv, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
+        return self.norm_vi(zi), self.norm_vv(zv)
+
+
+class DecompositionalEncoder(nn.Module):
+    """θ: video patches → (z_vi, z_vv).
+
+    Composes the per-frame ``QFormer`` with the ``CausalTemporalEncoder``.
+    """
+
+    def __init__(
+        self,
+        d_z: int = 512,
+        d_kv: int = 1024,
+        qformer_depth: int = 4,
+        temporal_depth: int = 12,
+        num_heads: int = 8,
+        mlp_ratio: float = 4.0,
+        max_frames: int = 128,
+    ):
+        super().__init__()
+        self.qformer = QFormer(
+            d_z=d_z, d_kv=d_kv, depth=qformer_depth,
+            num_heads=num_heads, mlp_ratio=mlp_ratio,
+        )
+        self.temporal = CausalTemporalEncoder(
+            d_z=d_z, max_frames=max_frames, depth=temporal_depth,
+            num_heads=num_heads, mlp_ratio=mlp_ratio,
+        )
+
+    def forward(
+        self,
+        patches: torch.Tensor,
+        key_padding_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        z_vi_pre, z_vv_pre = self.qformer(patches)
+        return self.temporal(z_vi_pre, z_vv_pre, key_padding_mask=key_padding_mask)

layers.py

@@ -0,0 +1,128 @@
+"""Shared building blocks for PRISM.
+
+- ``QFormerBlock``  : one block of the per-frame Q-Former (self-attn over the
+  queries → cross-attn into the frozen patch tokens → FFN). Used by the
+  Decompositional Encoder's frame-level stage.
+- ``TemporalBlock`` : a pre-LN transformer block used by both the temporal
+  stage of the Decompositional Encoder (causal) and the Compositional
+  Latent Predictor.
+- ``build_sin_pos_embed`` / ``causal_mask`` : positional-embedding and masking
+  helpers.
+
+PyTorch's ``MultiheadAttention`` boolean masks follow the convention
+``True == blocked`` for both ``attn_mask`` and ``key_padding_mask``.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+
+def build_sin_pos_embed(num_positions: int, dim: int) -> torch.Tensor:
+    """1-D sinusoidal positional embedding, shape ``(1, num_positions, dim)``."""
+    pe = torch.zeros(num_positions, dim)
+    pos = torch.arange(0, num_positions, dtype=torch.float32).unsqueeze(1)
+    div_term = torch.exp(
+        torch.arange(0, dim, 2, dtype=torch.float32) * (-math.log(10000.0) / dim)
+    )
+    pe[:, 0::2] = torch.sin(pos * div_term)
+    pe[:, 1::2] = torch.cos(pos * div_term)
+    return pe.unsqueeze(0)
+
+
+def causal_mask(length: int, device: torch.device | None = None) -> torch.Tensor:
+    """Boolean causal mask ``(L, L)``; ``True`` blocks future positions."""
+    return torch.triu(
+        torch.ones(length, length, dtype=torch.bool, device=device), diagonal=1
+    )
+
+
+class QFormerBlock(nn.Module):
+    """BLIP-2-style block: query self-attn → query→patch cross-attn → FFN."""
+
+    def __init__(
+        self,
+        d_z: int,
+        d_kv: int,
+        num_heads: int = 8,
+        mlp_ratio: float = 4.0,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_z)
+        self.self_attn = nn.MultiheadAttention(
+            embed_dim=d_z, num_heads=num_heads, dropout=attn_drop, batch_first=True
+        )
+        self.norm2_q = nn.LayerNorm(d_z)
+        self.norm2_kv = nn.LayerNorm(d_kv)
+        self.cross_attn = nn.MultiheadAttention(
+            embed_dim=d_z,
+            num_heads=num_heads,
+            dropout=attn_drop,
+            kdim=d_kv,
+            vdim=d_kv,
+            batch_first=True,
+        )
+        self.norm3 = nn.LayerNorm(d_z)
+        hidden = int(d_z * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(d_z, hidden),
+            nn.GELU(),
+            nn.Dropout(proj_drop),
+            nn.Linear(hidden, d_z),
+            nn.Dropout(proj_drop),
+        )
+
+    def forward(self, q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(q)
+        sa, _ = self.self_attn(h, h, h, need_weights=False)
+        q = q + sa
+        hq = self.norm2_q(q)
+        hkv = self.norm2_kv(kv)
+        ca, _ = self.cross_attn(hq, hkv, hkv, need_weights=False)
+        q = q + ca
+        q = q + self.mlp(self.norm3(q))
+        return q
+
+
+class TemporalBlock(nn.Module):
+    """Pre-LN transformer block (self-attn + FFN).
+
+    Accepts a boolean ``attn_mask`` ``(L, L)`` (True = blocked) and a
+    ``key_padding_mask`` ``(B, L)`` (True = padded) on every forward.
+    """
+
+    def __init__(self, d_in: int, num_heads: int = 8, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_in)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=d_in, num_heads=num_heads, batch_first=True
+        )
+        self.norm2 = nn.LayerNorm(d_in)
+        hidden = int(d_in * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(d_in, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, d_in),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+        key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        h = self.norm1(x)
+        a, _ = self.attn(
+            h, h, h,
+            need_weights=False,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+        )
+        x = x + a
+        x = x + self.mlp(self.norm2(x))
+        return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d542adf5895ba690a825a920408faa97b7c4fba9876868db8ffc8bd38d6465bd
+size 970913516

modeling_prism.py

@@ -0,0 +1,454 @@
+"""PRISM model.
+
+Frozen SigLIP2 vision + frozen Qwen3-Embedding text + trainable Decompositional
+Encoder θ + Compositional Latent Predictor φ + EMA target encoder θ̄.
+
+One training ``forward`` returns:
+
+  - ``loss_decomp`` : symmetric InfoNCE between the compositional latent
+    ``s = φ(z_vv^B, z_vi^A)`` and the recomposed text embedding
+    ``e = Qwen3Embedding(compose(T_vi^A, T_vv^B))``, over the batch
+    (with DDP all-gather of ``s`` / ``e`` / the valid-pair mask). [paper §3.2]
+  - ``loss_temp_vi`` / ``loss_temp_vv`` : ``1 - cos(ẑ_t, z̄_{t+1})`` for each
+    stream, where the target ``z̄`` comes from the EMA encoder θ̄. [paper §3.3]
+  - ``loss = λ_decomp · loss_decomp + λ_temp · ½(loss_temp_vi + loss_temp_vv)``.
+
+Cross-pairing (one clip's view-variant stream with another clip's view-invariant
+stream) is done inside ``forward`` via a cyclic shift of the batch: the
+view-variant stream comes from clip ``i``, the view-invariant stream from clip
+``(i+1) mod B``. The trainer builds the matching recomposed caption with the
+same convention.
+
+At inference, ``encode`` returns an L2-normalized clip embedding (mean-pooled
+``z_vi`` over valid frames) from the EMA encoder.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+from transformers import AutoModel, PreTrainedModel
+
+logger = logging.getLogger(__name__)
+
+from .configuration_prism import PRISMConfig
+from .ema import make_ema_copy, sync_ema_from_online, update_ema
+from .encoder import DecompositionalEncoder
+# Unused here, but kept as a direct import: when this file is served as Hub remote
+# code, transformers ships only the relative imports named in *this* module, so
+# ``layers`` (used by encoder/predictor) has to be visible from here.
+from .layers import QFormerBlock, TemporalBlock  # noqa: F401
+from .predictor import CompositionalPredictor
+
+
+@dataclass
+class PRISMOutput:
+    loss: torch.Tensor
+    loss_decomp: torch.Tensor
+    loss_temp_vi: torch.Tensor
+    loss_temp_vv: torch.Tensor
+    n_valid_pairs: int
+
+
+# ---------------------------------------------------------------------------
+# Loss / distributed helpers
+# ---------------------------------------------------------------------------
+def _symmetric_infonce(
+    a: torch.Tensor, b: torch.Tensor, logit_scale: torch.Tensor
+) -> torch.Tensor:
+    """CLIP-style symmetric InfoNCE on L2-normalized features."""
+    a = F.normalize(a, dim=-1)
+    b = F.normalize(b, dim=-1)
+    scale = logit_scale.exp().clamp(max=100.0)
+    logits = scale * a @ b.t()
+    labels = torch.arange(a.shape[0], device=a.device)
+    return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels))
+
+
+def _all_gather_with_grad(x: torch.Tensor) -> torch.Tensor:
+    """CLIP-style all-gather: concat across ranks; own-rank slot keeps gradient.
+
+    Other ranks' tensors are detached for this rank's backward; DDP's gradient
+    all-reduce then distributes the gradient across ranks, making it equivalent
+    to a single forward over the full ``B * world_size`` batch. Single-process
+    → returns ``x`` unchanged.
+    """
+    if not dist.is_available() or not dist.is_initialized():
+        return x
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return x
+    rank = dist.get_rank()
+    gathered = [torch.empty_like(x) for _ in range(world_size)]
+    dist.all_gather(gathered, x.contiguous())
+    gathered[rank] = x  # own-rank slot keeps grad
+    return torch.cat(gathered, dim=0)
+
+
+def _all_gather_bool(x: torch.Tensor) -> torch.Tensor:
+    """Plain all-gather for boolean masks (no gradient)."""
+    if not dist.is_available() or not dist.is_initialized():
+        return x
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return x
+    gathered = [torch.empty_like(x) for _ in range(world_size)]
+    dist.all_gather(gathered, x.contiguous())
+    return torch.cat(gathered, dim=0)
+
+
+def _sample_shift_plan(
+    valid_a: torch.Tensor, valid_b: torch.Tensor, T: int
+) -> dict:
+    """Sliding-shift augmentation plan.
+
+    For each sample, place the shorter clip's valid frames at a random offset
+    within the longer clip's valid range. Returns gather indices and post-shift
+    valid masks for both sides; apply identically to online and EMA tensors so
+    predictor input and temporal target stay time-aligned.
+    """
+    B = valid_a.shape[0]
+    device = valid_a.device
+    t_a = valid_a.int().sum(dim=1)
+    t_b = valid_b.int().sum(dim=1)
+    max_off_a = torch.clamp(t_b - t_a, min=0)
+    max_off_b = torch.clamp(t_a - t_b, min=0)
+
+    rand = torch.rand(B, 2, device=device)
+    offset_a = (rand[:, 0] * (max_off_a.float() + 1.0)).long().clamp(max=max_off_a)
+    offset_b = (rand[:, 1] * (max_off_b.float() + 1.0)).long().clamp(max=max_off_b)
+
+    arange_T = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+    src_idx_a = arange_T - offset_a.unsqueeze(1)
+    src_idx_b = arange_T - offset_b.unsqueeze(1)
+    new_valid_a = (src_idx_a >= 0) & (src_idx_a < t_a.unsqueeze(1))
+    new_valid_b = (src_idx_b >= 0) & (src_idx_b < t_b.unsqueeze(1))
+    return {
+        "src_idx_a": src_idx_a.clamp(0, T - 1),
+        "src_idx_b": src_idx_b.clamp(0, T - 1),
+        "new_valid_a": new_valid_a,
+        "new_valid_b": new_valid_b,
+    }
+
+
+def _apply_shift(
+    z: torch.Tensor, src_idx: torch.Tensor, valid: torch.Tensor
+) -> torch.Tensor:
+    """Gather ``z[B, T, D]`` along T via ``src_idx[B, T]``; zero invalid positions."""
+    gather_idx = src_idx.unsqueeze(-1).expand(-1, -1, z.shape[-1])
+    return torch.gather(z, dim=1, index=gather_idx) * valid.unsqueeze(-1).to(z.dtype)
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint resolution (local directory or Hugging Face Hub repo)
+# ---------------------------------------------------------------------------
+_WEIGHTS_NAME = "model.safetensors"
+_HUB_KWARGS = ("revision", "cache_dir", "token", "force_download", "local_files_only", "proxies")
+
+
+def _resolve_weights(path_or_repo: str, **hub_kwargs) -> str:
+    """Path to the checkpoint's weights: a local directory, else a Hub repo id."""
+    local = Path(path_or_repo) / _WEIGHTS_NAME
+    if local.is_file():
+        return str(local)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=str(path_or_repo), filename=_WEIGHTS_NAME, **hub_kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+class PRISMModel(PreTrainedModel):
+    config_class = PRISMConfig
+    base_model_prefix = "prism"
+    # Frozen backbones are reloaded from the Hub in __init__ and excluded from
+    # the saved checkpoint (see ``state_dict``); silence the load-time warning.
+    _keys_to_ignore_on_load_missing = [r"^vision_model\.", r"^text_model\."]
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: PRISMConfig):
+        super().__init__(config)
+
+        # ---- Frozen vision tower ----
+        # CLIP keeps a CLS token in last_hidden_state; SigLIP / SigLIP2 do not.
+        # SigLIP2 weights use the SigLIP v1 architecture, so SiglipVisionModel
+        # handles both checkpoint families.
+        if "siglip" in config.vision_backbone_name.lower():
+            from transformers import SiglipVisionModel
+            self.vision_model = SiglipVisionModel.from_pretrained(config.vision_backbone_name)
+            self._vision_has_cls = False
+        else:
+            from transformers import CLIPVisionModel
+            self.vision_model = CLIPVisionModel.from_pretrained(config.vision_backbone_name)
+            self._vision_has_cls = True
+        d_v = int(self.vision_model.config.hidden_size)
+
+        # ---- Frozen Qwen3-Embedding text tower ----
+        self.text_model = AutoModel.from_pretrained(config.text_backbone_name)
+        d_t = int(self.text_model.config.hidden_size)
+
+        for p in self.vision_model.parameters():
+            p.requires_grad = False
+        for p in self.text_model.parameters():
+            p.requires_grad = False
+        self.vision_model.eval()
+        self.text_model.eval()
+
+        self.d_v = d_v
+        self.d_t = d_t
+
+        # ---- Trainable: Decompositional Encoder θ + Compositional Predictor φ ----
+        self.encoder = DecompositionalEncoder(
+            d_z=config.d_z, d_kv=d_v,
+            qformer_depth=config.qformer_depth, temporal_depth=config.temporal_depth,
+            num_heads=config.num_heads, mlp_ratio=config.mlp_ratio,
+            max_frames=config.max_frames,
+        )
+        self.predictor = CompositionalPredictor(
+            d_z=config.d_z, d_t=d_t, max_frames=config.max_frames,
+            depth=config.predictor_depth, num_heads=config.num_heads,
+            mlp_ratio=config.mlp_ratio,
+        )
+
+        # ---- EMA target encoder θ̄ ----
+        if config.use_ema:
+            self.target_encoder = make_ema_copy(self.encoder)
+
+        self.logit_scale = nn.Parameter(torch.tensor(config.logit_scale_init))
+
+    # -- keep trainable defaults; do not re-init the from-Hub backbones --
+    def _init_weights(self, module):  # noqa: D401
+        pass
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Load a weights-only PRISM checkpoint from a local directory or the Hub.
+
+        The frozen vision / text backbones are not stored in the checkpoint;
+        they are rebuilt from the Hub in ``__init__``. Only the trained weights
+        (encoder θ, predictor φ, target encoder θ̄, logit scale) are loaded. The
+        ``dtype`` / ``torch_dtype`` kwarg is honored; other HF loading kwargs
+        (device_map, sharding, ...) are not needed for this single-file ckpt.
+        """
+        from safetensors.torch import load_file
+
+        hub_kwargs = {k: kwargs.pop(k) for k in _HUB_KWARGS if kwargs.get(k) is not None}
+        config = kwargs.pop("config", None)
+        if not isinstance(config, PRISMConfig):
+            config = PRISMConfig.from_pretrained(pretrained_model_name_or_path, **hub_kwargs)
+        dtype = kwargs.pop("torch_dtype", None) or kwargs.pop("dtype", None)
+
+        model = cls(config)  # backbones materialized from the Hub
+        state = load_file(_resolve_weights(pretrained_model_name_or_path, **hub_kwargs))
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        bad_missing = [m for m in missing if not m.startswith(("vision_model.", "text_model."))]
+        if bad_missing:
+            logger.warning(f"missing non-backbone keys: {bad_missing[:8]}")
+        if unexpected:
+            logger.warning(f"unexpected keys: {unexpected[:8]}")
+        # Warm-start the target encoder if a checkpoint predates EMA weights.
+        if config.use_ema and any(k.startswith("target_encoder.") for k in bad_missing):
+            model.sync_ema_from_online()
+        if dtype is not None:
+            model = model.to(dtype)
+        return model
+
+    def train(self, mode: bool = True):
+        """Keep frozen backbones (and the EMA target encoder) in eval mode."""
+        super().train(mode)
+        self.vision_model.eval()
+        self.text_model.eval()
+        if getattr(self.config, "use_ema", False):
+            self.target_encoder.eval()
+        return self
+
+    def state_dict(self, *args, **kwargs):
+        """Exclude the frozen, from-Hub backbones from saved checkpoints."""
+        sd = super().state_dict(*args, **kwargs)
+        return type(sd)(
+            (k, v) for k, v in sd.items()
+            if not k.startswith(("vision_model.", "text_model."))
+        )
+
+    # ------------------------------------------------------------------
+    # Frozen backbone helpers
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def _encode_video(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """``(B, T, 3, H, W)`` → patch tokens ``(B, T, P, d_v)`` (CLS dropped for CLIP)."""
+        B, T = pixel_values.shape[:2]
+        x = pixel_values.reshape(B * T, *pixel_values.shape[2:])
+        seq = self.vision_model(pixel_values=x).last_hidden_state
+        if self._vision_has_cls:
+            seq = seq[:, 1:, :]
+        return seq.reshape(B, T, seq.shape[1], self.d_v)
+
+    @torch.no_grad()
+    def _encode_text(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """Qwen3-Embedding: last-token pool over a right-padded batch → L2-normed ``(N, d_t)``."""
+        last_hidden = self.text_model(
+            input_ids=input_ids, attention_mask=attention_mask
+        ).last_hidden_state
+        last_idx = (attention_mask.sum(dim=1) - 1).clamp(min=0)
+        rows = torch.arange(last_hidden.shape[0], device=last_hidden.device)
+        return F.normalize(last_hidden[rows, last_idx], dim=-1)
+
+    # ------------------------------------------------------------------
+    # EMA hooks (called by the trainer after each optimizer step)
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def update_ema(self) -> None:
+        if getattr(self.config, "use_ema", False):
+            update_ema(self.target_encoder, self.encoder, self.config.ema_decay)
+
+    @torch.no_grad()
+    def sync_ema_from_online(self) -> None:
+        if getattr(self.config, "use_ema", False):
+            sync_ema_from_online(self.target_encoder, self.encoder)
+
+    # ------------------------------------------------------------------
+    # Inference
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def encode_streams(
+        self, pixel_values: torch.Tensor, valid_mask: torch.Tensor | None = None
+    ) -> dict:
+        """Run θ̄ (or θ if ``use_ema=False``) → ``{z_vi_seq, z_vv_seq}`` each ``(B, T, d_z)``."""
+        patches = self._encode_video(pixel_values)
+        kpm = None if valid_mask is None else (~valid_mask)
+        enc = self.target_encoder if getattr(self.config, "use_ema", False) else self.encoder
+        z_vi, z_vv = enc(patches, key_padding_mask=kpm)
+        return {"z_vi_seq": z_vi, "z_vv_seq": z_vv}
+
+    @torch.no_grad()
+    def encode(
+        self, pixel_values: torch.Tensor, valid_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Clip embedding: L2-normalized mean-pool of ``z_vi`` over valid frames → ``(B, d_z)``."""
+        z_vi = self.encode_streams(pixel_values, valid_mask)["z_vi_seq"].float()
+        if valid_mask is None:
+            valid = torch.ones(z_vi.shape[:2], device=z_vi.device, dtype=z_vi.dtype)
+        else:
+            valid = valid_mask.to(z_vi.dtype)
+        denom = valid.sum(dim=1, keepdim=True).clamp(min=1.0)
+        emb = (z_vi * valid.unsqueeze(-1)).sum(dim=1) / denom
+        return F.normalize(emb, dim=-1)
+
+    # ------------------------------------------------------------------
+    # Training forward
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        valid_mask: torch.Tensor,
+        composed_input_ids: torch.Tensor,
+        composed_attention_mask: torch.Tensor,
+        valid_pair_mask: torch.Tensor | None = None,
+    ) -> PRISMOutput:
+        """
+        pixel_values:            ``(B, T, 3, H, W)`` — padded to T in the collator.
+        valid_mask:              ``(B, T)`` bool, True = real frame.
+        composed_input_ids:      ``(B, L)`` — tokenized recomposed caption per pair.
+        composed_attention_mask: ``(B, L)``.
+        valid_pair_mask:         ``(B,)`` bool — True if the composer succeeded.
+        """
+        B = pixel_values.shape[0]
+        device = pixel_values.device
+        use_ema = getattr(self.config, "use_ema", False)
+
+        # ---- Frozen encoders ----
+        patches = self._encode_video(pixel_values)                       # (B, T, P, d_v)
+        T = patches.shape[1]
+        e_text = self._encode_text(composed_input_ids, composed_attention_mask)  # (B, d_t)
+
+        # ---- Decompositional encoder θ (and EMA θ̄ for temporal targets) ----
+        kpm = ~valid_mask
+        z_vi_seq, z_vv_seq = self.encoder(patches, key_padding_mask=kpm)
+        if use_ema:
+            with torch.no_grad():
+                z_vi_seq_ema, z_vv_seq_ema = self.target_encoder(patches, key_padding_mask=kpm)
+
+        # ---- Cross-pairing: view-variant from clip i, view-invariant from (i+1) ----
+        shift = torch.roll(torch.arange(B, device=device), shifts=-1, dims=0)
+        z_vv = z_vv_seq                       # view-variant  (clip i)
+        z_vi = z_vi_seq[shift]                # view-invariant (clip i+1)
+        valid_vv = valid_mask
+        valid_vi = valid_mask[shift]
+        if use_ema:
+            z_vv_ema = z_vv_seq_ema
+            z_vi_ema = z_vi_seq_ema[shift]
+
+        # ---- Sliding-shift augmentation (training only) ----
+        if self.training and getattr(self.config, "sliding_shift_aug", True):
+            plan = _sample_shift_plan(valid_vv, valid_vi, T)
+            z_vv = _apply_shift(z_vv, plan["src_idx_a"], plan["new_valid_a"])
+            z_vi = _apply_shift(z_vi, plan["src_idx_b"], plan["new_valid_b"])
+            if use_ema:
+                z_vv_ema = _apply_shift(z_vv_ema, plan["src_idx_a"], plan["new_valid_a"])
+                z_vi_ema = _apply_shift(z_vi_ema, plan["src_idx_b"], plan["new_valid_b"])
+            valid_vv = plan["new_valid_a"]
+            valid_vi = plan["new_valid_b"]
+
+        # ---- Compositional predictor φ ----
+        out = self.predictor(z_vv, z_vi, valid_vv=valid_vv, valid_vi=valid_vi)
+        s = out["s"]                          # (B, d_t)   compositional latent
+        z_vi_pred = out["z_vi_pred"]          # (B, T, d_z) vi next-frame head
+        z_vv_pred = out["z_vv_pred"]          # (B, T, d_z) vv next-frame head
+        pair_valid = out["pair_valid"]        # (B, T)
+
+        # ---- L_decomp: InfoNCE(s, e_text) over valid pairs ----
+        if valid_pair_mask is None:
+            valid_pair_mask = torch.ones(B, dtype=torch.bool, device=device)
+        valid_pair_mask = valid_pair_mask & pair_valid.any(dim=1)
+
+        if self.config.infonce_all_gather:
+            s_g = _all_gather_with_grad(s)
+            e_g = _all_gather_with_grad(e_text)
+            valid_g = _all_gather_bool(valid_pair_mask)
+        else:
+            s_g, e_g, valid_g = s, e_text, valid_pair_mask
+
+        valid_idx = valid_g.nonzero(as_tuple=True)[0]
+        n_valid = int(valid_idx.numel())
+        if n_valid >= 2:
+            loss_decomp = _symmetric_infonce(s_g[valid_idx], e_g[valid_idx], self.logit_scale)
+        else:
+            loss_decomp = torch.zeros((), device=device)
+
+        # ---- L_temp: 1 - cos(prediction at t, EMA target at t+1), per stream ----
+        if T >= 2:
+            if use_ema:
+                tgt_vi = z_vi_ema[:, 1:T, :]
+                tgt_vv = z_vv_ema[:, 1:T, :]
+            else:
+                tgt_vi = z_vi.detach()[:, 1:T, :]
+                tgt_vv = z_vv.detach()[:, 1:T, :]
+            valid_next_vi = valid_vi[:, :T - 1] & valid_vi[:, 1:T]
+            valid_next_vv = valid_vv[:, :T - 1] & valid_vv[:, 1:T]
+            err_vi = 1.0 - F.cosine_similarity(z_vi_pred[:, :T - 1, :], tgt_vi, dim=-1)
+            err_vv = 1.0 - F.cosine_similarity(z_vv_pred[:, :T - 1, :], tgt_vv, dim=-1)
+            loss_temp_vi = err_vi[valid_next_vi].mean() if valid_next_vi.any() else torch.zeros((), device=device)
+            loss_temp_vv = err_vv[valid_next_vv].mean() if valid_next_vv.any() else torch.zeros((), device=device)
+        else:
+            loss_temp_vi = torch.zeros((), device=device)
+            loss_temp_vv = torch.zeros((), device=device)
+
+        loss_temp = 0.5 * (loss_temp_vi + loss_temp_vv)
+        loss = self.config.lambda_decomp * loss_decomp + self.config.lambda_temp * loss_temp
+
+        return PRISMOutput(
+            loss=loss,
+            loss_decomp=loss_decomp.detach(),
+            loss_temp_vi=loss_temp_vi.detach(),
+            loss_temp_vv=loss_temp_vv.detach(),
+            n_valid_pairs=n_valid,
+        )

predictor.py

@@ -0,0 +1,106 @@
+"""Compositional Latent Predictor φ.
+
+Takes the view-variant stream of one clip and the view-invariant stream of
+another, fuses them per time step, and runs a small causal transformer
+(paper §3.1–3.3). One token per frame::
+
+    token_t = concat(z_vv_t, z_vi_t) ∈ R^{2·d_z}
+
+with a standard causal mask (frame ``t`` attends to ``≤ t``) and a key-padding
+mask blocking frames where either stream is padded.
+
+Three linear heads on the final hidden state ``x ∈ R^{B×T×2·d_z}``:
+
+  - ``cls_head`` (2·d_z → d_t): at the last valid pair-frame → the compositional
+    semantic latent ``s`` aligned to language (``L_decomp``).
+  - ``vi_head`` (2·d_z → d_z): per t → ``ẑ_vi`` predicting the view-invariant
+    input stream's next frame (``L_temp``).
+  - ``vv_head`` (2·d_z → d_z): per t → ``ẑ_vv`` predicting the view-variant
+    input stream's next frame (``L_temp``).
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .layers import TemporalBlock, build_sin_pos_embed, causal_mask
+
+
+class CompositionalPredictor(nn.Module):
+    def __init__(
+        self,
+        d_z: int = 512,
+        d_t: int = 1024,
+        max_frames: int = 128,
+        depth: int = 4,
+        num_heads: int = 8,
+        mlp_ratio: float = 4.0,
+    ):
+        super().__init__()
+        self.d_z = d_z
+        self.d_t = d_t
+        self.max_frames = max_frames
+        self.d_in = 2 * d_z  # channel-concat [z_vv ‖ z_vi]
+
+        self.register_buffer(
+            "pos_embed", build_sin_pos_embed(max_frames, self.d_in), persistent=False
+        )
+        self.input_norm = nn.LayerNorm(self.d_in)
+        self.blocks = nn.ModuleList(
+            [TemporalBlock(self.d_in, num_heads, mlp_ratio) for _ in range(depth)]
+        )
+        self.final_norm = nn.LayerNorm(self.d_in)
+
+        self.cls_head = nn.Linear(self.d_in, d_t)   # → s   (compositional latent)
+        self.vi_head = nn.Linear(self.d_in, d_z)    # → ẑ_vi (next-frame, V-I)
+        self.vv_head = nn.Linear(self.d_in, d_z)    # → ẑ_vv (next-frame, V-V)
+
+    def forward(
+        self,
+        z_vv: torch.Tensor,
+        z_vi: torch.Tensor,
+        valid_vv: torch.Tensor,
+        valid_vi: torch.Tensor,
+    ) -> dict:
+        """
+        z_vv, z_vi: ``(B, T, d_z)`` — view-variant / view-invariant streams,
+            generally sourced from two different clips.
+        valid_vv, valid_vi: ``(B, T)`` bool — True where the frame is real.
+
+        Returns dict:
+            s:          ``(B, d_t)``    cls head at the last valid pair-frame.
+            z_vi_pred:  ``(B, T, d_z)`` vi head per t (predicts z_vi_{t+1}).
+            z_vv_pred:  ``(B, T, d_z)`` vv head per t (predicts z_vv_{t+1}).
+            pair_valid: ``(B, T)``      valid_vv & valid_vi.
+        """
+        B, T, _ = z_vv.shape
+        device = z_vv.device
+        assert T <= self.max_frames, f"T={T} > max_frames={self.max_frames}"
+        assert z_vi.shape == z_vv.shape, "z_vv and z_vi must match shape"
+
+        x = torch.cat([z_vv, z_vi], dim=-1)                  # (B, T, 2·d_z)
+        x = x + self.pos_embed[:, :T, :].to(dtype=x.dtype)
+        x = self.input_norm(x)
+
+        attn_mask = causal_mask(T, device=device)            # (T, T) True = blocked
+        pair_valid = valid_vv & valid_vi                     # (B, T)
+        key_padding_mask = ~pair_valid                       # True = blocked
+        for blk in self.blocks:
+            x = blk(x, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
+        x = self.final_norm(x)                               # (B, T, 2·d_z)
+
+        # cls head at the last valid pair-frame per sample. (Samples with no
+        # valid pair-frame are filtered out of the loss downstream.)
+        arange_T = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+        scored = arange_T.where(pair_valid, torch.full_like(arange_T, -1))
+        last_valid = scored.max(dim=1).values.clamp(min=0)   # (B,)
+        last_tokens = x[torch.arange(B, device=device), last_valid, :]
+        s = self.cls_head(last_tokens)                       # (B, d_t)
+
+        return {
+            "s": s,
+            "z_vi_pred": self.vi_head(x),                    # (B, T, d_z)
+            "z_vv_pred": self.vv_head(x),                    # (B, T, d_z)
+            "pair_valid": pair_valid,
+        }