"""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)