File size: 7,059 Bytes
1e05592 | 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 152 153 154 155 156 157 158 159 160 161 162 163 164 | """Dynamic-residual features for LKAlert-BD.
Day-1 diagnostic showed mean adjacent cosine distance ≈ 0.03 across all
caches — the per-frame Qwen3-VL belief is dynamically smooth. The whole
"belief is too invariant" observation is concrete evidence that the GRU
head can't recover motion residual on its own. This module builds explicit
hand-crafted features from the belief sequence so a small MLP can decide
how much motion residual is recoverable from the existing belief cache.
Features (all differentiable, all derivable from `beliefs_frame [B, T, D]`
and `valid_frames [B, T]`):
belief-pool channels (length D):
b_last : last valid belief
b_first : first valid belief
b_mean : valid-mean
b_max : valid-max-pool (per-dim)
delta_last : b_last - b_first (motion direction over the clip)
scalar dynamics (length 1 each):
mean_adj_cos_dist : mean cosine distance between adjacent valid frames
p95_adj_cos_dist : 95-percentile of same
max_norm_jump : max ||b_t - b_{t-1}|| / max_t ||b_t||
mean_norm_slope : (||b_last|| - ||b_first||) / max(1, T-1)
n_valid : number of valid frames
optional TTA channels (length 2 each, if `tta_means`/`tta_vars` provided):
tta_mean_last, tta_var_last
tta_mean_max, tta_var_max
tta_mean_first
tta_mean_slope (last - first) / valid steps
This module exposes only `build_features(...)`. It returns a single
dict and never makes architectural decisions for the caller.
"""
from __future__ import annotations
from typing import Dict, Optional
import torch
import torch.nn.functional as F
@torch.no_grad()
def _adjacent_cos_distance(b: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
"""[B,T,D] → [B,T-1] adjacent (1-cos), 0 where either side invalid."""
eps = 1e-6
bn = b / b.norm(dim=-1, keepdim=True).clamp(min=eps)
cos = (bn[:, 1:] * bn[:, :-1]).sum(dim=-1) # [B, T-1]
pair = (valid[:, 1:] & valid[:, :-1]).float() # [B, T-1]
return (1.0 - cos) * pair # 0 where invalid
def build_features(
beliefs: torch.Tensor, # [B, T, D]
valid: torch.Tensor, # [B, T] bool
tta_means: Optional[torch.Tensor] = None, # [B, T]
tta_vars: Optional[torch.Tensor] = None, # [B, T]
) -> Dict[str, torch.Tensor]:
"""Returns a dict of named feature tensors. All keys are length-axis B.
`pooled` is a single concatenated [B, F] tensor for downstream MLPs.
"""
B, T, D = beliefs.shape
valid_f = valid.float().unsqueeze(-1) # [B, T, 1]
n_valid = valid_f.sum(dim=1).squeeze(-1).clamp(min=1.0) # [B]
# last valid index per row (fallback to T-1 if all invalid)
pos = torch.arange(T, device=beliefs.device).unsqueeze(0).expand(B, T)
last_idx = (pos * valid.long()).max(dim=1).values # [B]
first_idx = (pos.masked_fill(~valid, T) ).min(dim=1).values # [B]
first_idx = first_idx.clamp(max=T - 1)
bidx = torch.arange(B, device=beliefs.device)
b_last = beliefs[bidx, last_idx] # [B, D]
b_first = beliefs[bidx, first_idx] # [B, D]
b_mean = (beliefs * valid_f).sum(dim=1) / n_valid.unsqueeze(-1) # [B, D]
# masked max
masked = beliefs.masked_fill(~valid.unsqueeze(-1), float("-inf"))
b_max = masked.max(dim=1).values # [B, D]
# if a row had no valid frames the max collapses to -inf — recover with mean
b_max = torch.where(b_max == float("-inf"), b_mean, b_max)
delta_last = b_last - b_first # [B, D]
# scalar dynamics
adj = _adjacent_cos_distance(beliefs, valid) # [B, T-1]
pair_count = (valid[:, 1:] & valid[:, :-1]).float().sum(dim=1).clamp(min=1.0)
mean_adj = adj.sum(dim=1) / pair_count # [B]
p95_adj = torch.quantile(adj, q=0.95, dim=1) # [B]
norm_t = beliefs.norm(dim=-1) # [B, T]
max_norm = norm_t.max(dim=1).values.clamp(min=1e-6) # [B]
diffs = (beliefs[:, 1:] - beliefs[:, :-1]).norm(dim=-1) # [B, T-1]
pair_mask = (valid[:, 1:] & valid[:, :-1]).float()
diffs = diffs * pair_mask
max_norm_jump = diffs.max(dim=1).values / max_norm # [B]
norm_last = beliefs[bidx, last_idx].norm(dim=-1)
norm_first = beliefs[bidx, first_idx].norm(dim=-1)
mean_norm_slope = (norm_last - norm_first) / n_valid.clamp(min=1.0) # [B]
out: Dict[str, torch.Tensor] = {
"b_last": b_last,
"b_first": b_first,
"b_mean": b_mean,
"b_max": b_max,
"delta_last": delta_last,
"mean_adj_cos_dist": mean_adj,
"p95_adj_cos_dist": p95_adj,
"max_norm_jump": max_norm_jump,
"mean_norm_slope": mean_norm_slope,
"n_valid": n_valid,
}
if tta_means is not None and tta_vars is not None:
# The qwen3vl4b cache stores tta as a clip-level scalar [B], not [B,T].
# Older caches store [B,T]. Handle both transparently.
if tta_means.dim() == 1:
out.update({
"tta_mean_last": tta_means,
"tta_var_last": tta_vars,
"tta_mean_first": tta_means,
"tta_mean_max": tta_means,
"tta_var_max": tta_vars,
"tta_mean_slope": torch.zeros_like(tta_means),
})
else:
tm = tta_means * valid.float()
tv = tta_vars * valid.float()
out.update({
"tta_mean_last": tta_means[bidx, last_idx],
"tta_var_last": tta_vars [bidx, last_idx],
"tta_mean_first": tta_means[bidx, first_idx],
"tta_mean_max": tm.max(dim=1).values,
"tta_var_max": tv.max(dim=1).values,
"tta_mean_slope": (tta_means[bidx, last_idx]
- tta_means[bidx, first_idx]) / n_valid,
})
# convenience: a single concatenated pooled tensor
pieces = [
out["b_last"], out["b_mean"], out["b_max"], out["delta_last"],
out["mean_adj_cos_dist"].unsqueeze(-1),
out["p95_adj_cos_dist"].unsqueeze(-1),
out["max_norm_jump"].unsqueeze(-1),
out["mean_norm_slope"].unsqueeze(-1),
]
if "tta_mean_last" in out:
pieces += [
out["tta_mean_last"].unsqueeze(-1),
out["tta_var_last"].unsqueeze(-1),
out["tta_mean_first"].unsqueeze(-1),
out["tta_mean_max"].unsqueeze(-1),
out["tta_var_max"].unsqueeze(-1),
out["tta_mean_slope"].unsqueeze(-1),
]
out["pooled"] = torch.cat(pieces, dim=-1) # [B, F]
return out
def feature_dim(belief_dim: int, with_tta: bool = True) -> int:
"""Returns F = D*4 + 4 (+6 if with_tta)."""
return belief_dim * 4 + 4 + (6 if with_tta else 0)
|