File size: 4,333 Bytes
a596b0a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """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,
}
|