| """Small heads on ARROW's encoder tokens (DA3-Giant: per frame, a 3072-d token per 14x14 patch + an id token). |
| |
| They decide *what to query* at decode time, so they must be far cheaper than the decoder: |
| motion per-pixel logit "this pixel moves around this frame's time" (two linears per patch, pixel-shuffled |
| to 14x14 pixels) |
| sky per-pixel logit, same shape (sky pixels are never queried for geometry) |
| covisibility per-pixel logit "this pixel of frame i is visible in frame j", conditioned on frame j's id token |
| (FiLM on the patch feature); decides which pixels of a new frame reused geometry can cover |
| descriptor per-patch L2-normalised local descriptors + a GeM-pooled global descriptor (matching / retrieval) |
| All heads read the same tokens (B, F, N, C) with the patch grid (h, w); none touch the decoder. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| TOKEN_DIM = 3072 |
| PATCH = 14 |
|
|
|
|
| def pixel_shuffle(logits, h, w, patch=PATCH): |
| """(B, F, h*w, patch*patch) per-patch pixel logits -> (B, F, h*patch, w*patch).""" |
| B, Fr = logits.shape[:2] |
| x = logits.view(B, Fr, h, w, patch, patch).permute(0, 1, 2, 4, 3, 5) |
| return x.reshape(B, Fr, h * patch, w * patch) |
|
|
|
|
| class PatchMLP(nn.Module): |
| def __init__(self, out_dim, hidden=256, in_dim=TOKEN_DIM): |
| super().__init__() |
| self.net = nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, hidden), nn.GELU(), nn.Linear(hidden, out_dim)) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class DenseLogitHead(nn.Module): |
| """Per-pixel binary logits from patch tokens (motion, sky).""" |
|
|
| def __init__(self, hidden=256, patch=PATCH): |
| super().__init__() |
| self.patch = patch |
| self.mlp = PatchMLP(patch * patch, hidden) |
|
|
| def forward(self, tokens, h, w): |
| return pixel_shuffle(self.mlp(tokens), h, w, self.patch) |
|
|
|
|
| class CovisibilityHead(nn.Module): |
| """Per-pixel logit that a pixel of frame i is seen in frame j: FiLM(patch_i | id_j) -> 14x14 logits.""" |
|
|
| def __init__(self, hidden=256, patch=PATCH): |
| super().__init__() |
| self.patch = patch |
| self.source = PatchMLP(hidden, hidden) |
| self.film = nn.Sequential(nn.LayerNorm(TOKEN_DIM), nn.Linear(TOKEN_DIM, 2 * hidden)) |
| self.out = nn.Sequential(nn.GELU(), nn.Linear(hidden, patch * patch)) |
|
|
| def forward(self, tokens, id_tokens, pairs, h, w): |
| """tokens (B, F, N, C), id_tokens (B, F, C), pairs (P, 2) frame indices (i, j) -> (B, P, H, W).""" |
| f = self.source(tokens[:, pairs[:, 0]]) |
| scale, shift = self.film(id_tokens[:, pairs[:, 1]]).chunk(2, dim=-1) |
| f = f * (1 + scale[:, :, None]) + shift[:, :, None] |
| return pixel_shuffle(self.out(f), h, w, self.patch) |
|
|
|
|
| class DescriptorHead(nn.Module): |
| def __init__(self, dim=128, global_dim=256, hidden=512, gem_p=3.0): |
| super().__init__() |
| self.local = PatchMLP(dim, hidden) |
| self.global_proj = nn.Linear(dim, global_dim) |
| self.p = nn.Parameter(torch.tensor(float(gem_p))) |
|
|
| def forward(self, tokens): |
| local = F.normalize(self.local(tokens), dim=-1) |
| pooled = local.clamp_min(1e-6).pow(self.p).mean(2).pow(1 / self.p) |
| return local, F.normalize(self.global_proj(pooled), dim=-1) |
|
|
|
|
| class TokenHeads(nn.Module): |
| """The configured subset of heads; forward returns a dict keyed by head name.""" |
|
|
| def __init__(self, heads=("motion", "covisibility", "descriptor", "sky"), hidden=256, descriptor_dim=128): |
| super().__init__() |
| self.names = tuple(heads) |
| unknown = set(self.names) - {"motion", "covisibility", "descriptor", "sky"} |
| if unknown: |
| raise ValueError(f"Unknown heads: {sorted(unknown)}") |
| self.heads = nn.ModuleDict() |
| if "motion" in self.names: |
| self.heads["motion"] = DenseLogitHead(hidden) |
| if "sky" in self.names: |
| self.heads["sky"] = DenseLogitHead(hidden) |
| if "covisibility" in self.names: |
| self.heads["covisibility"] = CovisibilityHead(hidden) |
| if "descriptor" in self.names: |
| self.heads["descriptor"] = DescriptorHead(descriptor_dim) |
|
|
| def forward(self, tokens, id_tokens, h, w, pairs=None): |
| out = {} |
| for name in ("motion", "sky"): |
| if name in self.heads: |
| out[name] = self.heads[name](tokens, h, w) |
| if "covisibility" in self.heads and pairs is not None and len(pairs): |
| out["covisibility"] = self.heads["covisibility"](tokens, id_tokens, pairs, h, w) |
| if "descriptor" in self.heads: |
| out["descriptor_local"], out["descriptor_global"] = self.heads["descriptor"](tokens) |
| return out |
|
|
