File size: 5,573 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """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)
|