"""Per-position behavioral attribution head. For each transaction position in the customer's 64-tx history, predict a binary signal: did this transaction contribute to the surface's score? This is the "which transactions drove the verdict" output that makes scores inspectable in the UI. Why per-position binary (BIO-free, single-label): The doctrine (`liquid-models-architecture` Section 10, PII whitepaper pattern) uses BIO/BIOES tagging for entity recognition. Our task is simpler — there are no spans, no entity classes. Each transaction either contributed (1) or didn't (0). A single sigmoid per position is the minimal head shape that does this. The output is interpretable: position-level probabilities can be thresholded or top-k'd for the UI ("transactions 3, 14, 22, 31 ... contributed"). Co-training discipline (MANDATORY): The multi-head decision layer ADR-014 finding: co-training a token-level head (this attribution head) with sequence-level heads (the probability head) naively collapsed PII to 9.9% F1. The fix is per-batch homogeneous-head sampling — each minibatch contains examples training EITHER the probability head OR this attribution head, never both. The surface_trainer enforces this. Shape contract: Input hidden_states: (B, T_total, D) T_total = 64 + 1 + T_txt Output position_logits: (B, num_tx_positions) single logit per tx """ from __future__ import annotations from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F @dataclass class AttributionHeadConfig: """Attribution head configuration. Attributes: name: Identifier for logging and config dispatch. hidden_dim: Backbone hidden size (1024 for LFM2.5-350M). mlp_hidden: Width of the per-position MLP. Kept small (64) because the head runs at every of the 64 transaction positions; the parameter count is num_tx_positions * weight. dropout: Dropout between the two MLP layers. num_tx_positions: Number of transaction pseudo-token positions at the start of the sequence. 64 in Phase 1. pos_weight: BCE positive-class weight. Used when attribution labels are skewed (most positions don't contribute; only a few do). Doctrine for recall-critical detector-style heads (`data-distribution-doctrine` Section 9): class-weighted loss on rare labels. Default 5.0 reflects ~10-20% positive position rate in synthesis; tune per surface. """ name: str hidden_dim: int = 1024 mlp_hidden: int = 64 dropout: float = 0.1 num_tx_positions: int = 64 pos_weight: float = 5.0 class AttributionHead(nn.Module): """Per-position binary classifier over the transaction pseudo-tokens. Applied independently at each transaction position. The same MLP weights run at every position. Parameter count at defaults (hidden_dim=1024, mlp_hidden=64): ~66K. Position-independent weights mean the head generalizes across positions and across different sequence lengths (relevant only if we later extend beyond 64). """ def __init__(self, config: AttributionHeadConfig) -> None: super().__init__() self.config = config # The MLP is applied position-wise (Linear → ReLU → Dropout → # Linear). PyTorch's Linear broadcasts over leading dims, so # mlp(tx_hidden) where tx_hidden is (B, T, D) returns (B, T, 1). self.mlp = nn.Sequential( nn.Linear(config.hidden_dim, config.mlp_hidden), nn.ReLU(), nn.Dropout(config.dropout), nn.Linear(config.mlp_hidden, 1), ) # Register pos_weight as a buffer so it follows the module to # the correct device but isn't a trained parameter. self.register_buffer( "pos_weight", torch.tensor(config.pos_weight, dtype=torch.float32), ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Run the head at every transaction position. Args: hidden_states: (B, T_total, D) hidden states from the backbone. We slice the leading `num_tx_positions` positions and run the position-wise MLP. Returns: (B, num_tx_positions) raw logits, one per transaction position. Caller applies sigmoid for probabilities or BCE for loss. """ # (B, num_tx_positions, D); cast to head dtype since the # backbone runs in bf16 but the small MLP head stays in fp32 # for numerical stability. head_dtype = next(self.mlp.parameters()).dtype tx_slice = hidden_states[:, : self.config.num_tx_positions, :].to(head_dtype) # (B, num_tx_positions, 1) -> (B, num_tx_positions) return self.mlp(tx_slice).squeeze(-1) def compute_loss( self, logits: torch.Tensor, targets: torch.Tensor, ) -> torch.Tensor: """Weighted BCE-with-logits, summed over positions then meaned. Args: logits: (B, num_tx_positions) raw logits. targets: (B, num_tx_positions) float in {0.0, 1.0} per position. 1.0 = this transaction contributed to the surface's score. Returns: Scalar loss. Mean over (batch × positions). The pos_weight up-weights the positive class because most positions are negatives in attribution. """ return F.binary_cross_entropy_with_logits( logits, targets.float(), pos_weight=self.pos_weight, reduction="mean", ) @torch.no_grad() def probabilities(self, logits: torch.Tensor) -> torch.Tensor: """Per-position sigmoid probabilities. Args: logits: (B, num_tx_positions). Returns: (B, num_tx_positions) in [0, 1]. """ return torch.sigmoid(logits) @torch.no_grad() def top_k_positions( self, logits: torch.Tensor, k: int = 8, ) -> torch.Tensor: """Top-k contributing transaction indices for the UI. Args: logits: (B, num_tx_positions). k: how many positions to return per batch element. Returns: (B, k) int64 transaction indices, sorted by descending contribution probability. """ # (B, k) indices into [0, num_tx_positions) return logits.topk(k=k, dim=-1).indices def num_parameters(self) -> int: return sum(p.numel() for p in self.parameters())