"""Per-surface probability head: 3-class categorical with mean-pool. This is the headline head for every multi-surface in ADR 0003. One calibrated score per surface, expressed as a 3-class categorical distribution {unlikely, ambiguous, likely}. Why categorical instead of a single sigmoid scalar (the obvious choice): The `liquid-finetuning-playbook` Section 10 calls out that LFM2.5- 350M numeric outputs cluster at attractor states (0.05, 0.45, 0.85, 0.95). A continuous "friendly-fraud probability" trained on 350M will fight those attractors. The fix doctrine: use categorical labels (low / medium / high / critical) and let softmax produce a probability distribution. For dispute legitimacy: class 0 = unlikely (low friendly-fraud probability) class 1 = ambiguous (mid; analyst should review) class 2 = likely (high friendly-fraud probability) The "score" displayed in the UI is `softmax(logits)[2]` — the probability assigned to the `likely` class. Softmax is naturally calibrated within {0, 1} without fighting attractor states because the model picks between three peaks, not one continuous range. Why mean-pool over the tx positions only: LFM2.5's final layer is conv (kernel 3). Last-token pooling reads only a 3-token receptive field (`liquid-models-architecture` Section 3, 16). For sequence classification we mean-pool over all 64 transaction pseudo-token positions, which inherits global context from the 6 attention layers at L2/L5/L8/L10/L12/L14. Importantly: we mean-pool over the 64 *transaction* positions, NOT the entire (tx + sep + text) combined sequence. The text tokens are conditioning information; the score belongs to the customer's behavioral signature, which lives in the transaction positions. Shape contract: Input hidden_states: (B, T_total, D) where T_total = 64 + 1 + T_txt Output logits: (B, num_classes) default num_classes = 3 """ from __future__ import annotations from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F @dataclass class ProbabilityHeadConfig: """Probability head configuration. Attributes: name: Identifier for logging and config dispatch. num_classes: Categorical output dimension. Default 3 for the unlikely/ambiguous/likely scheme. Surfaces with a richer taxonomy (collections treatment, fraud pattern) override. hidden_dim: Backbone hidden size. 1024 for LFM2.5-350M, 2048 for LFM2.5-1.2B. mlp_hidden: Width of the 2-layer MLP head. 256 is comfortable for 350M; 30 KB at fp32, ~7 KB at int8. dropout: Dropout between the two MLP layers. 0.1 default. num_tx_positions: Number of transaction pseudo-token positions at the start of the sequence. 64 in Phase 1. The head mean-pools over `hidden_states[:, :num_tx_positions, :]`. """ name: str num_classes: int = 3 hidden_dim: int = 1024 mlp_hidden: int = 256 dropout: float = 0.1 num_tx_positions: int = 64 class ProbabilityHead(nn.Module): """Mean-pool over the transaction pseudo-tokens, then MLP, then softmax. Parameter count at default config (hidden_dim=1024, mlp_hidden=256, num_classes=3): ~263K params. Small relative to the LoRA adapter (~1M) and immaterial against the 350M backbone. The head is intentionally minimal. The work is done by the backbone + per-surface LoRA; the head is a thin calibration layer on top of the pooled representation. """ def __init__(self, config: ProbabilityHeadConfig) -> None: super().__init__() self.config = config self.mlp = nn.Sequential( nn.Linear(config.hidden_dim, config.mlp_hidden), nn.ReLU(), nn.Dropout(config.dropout), nn.Linear(config.mlp_hidden, config.num_classes), ) def pool(self, hidden_states: torch.Tensor) -> torch.Tensor: """Mean-pool over the leading `num_tx_positions` positions. Args: hidden_states: (B, T_total, D) hidden states from the backbone after attending over [tx pseudo-tokens, SEP, text tokens]. T_total = num_tx_positions + 1 + T_txt. Returns: (B, D) pooled vector over the transaction positions only. """ # (B, num_tx_positions, D) -> (B, 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) return tx_slice.mean(dim=1) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Run the full head. Args: hidden_states: (B, T_total, D). Returns: (B, num_classes) raw logits. Caller applies softmax for probabilities or cross-entropy for loss. """ # (B, D) pooled = self.pool(hidden_states) # (B, num_classes) return self.mlp(pooled) def compute_loss( self, logits: torch.Tensor, targets: torch.Tensor, ) -> torch.Tensor: """Cross-entropy over the categorical distribution. Args: logits: (B, num_classes) raw logits from `forward`. targets: (B,) int64 class indices in [0, num_classes). Returns: Scalar CE loss. """ return F.cross_entropy(logits, targets) @torch.no_grad() def score(self, logits: torch.Tensor) -> torch.Tensor: """Friendly-fraud probability for display in the UI. Returns softmax(logits)[..., -1] — the probability mass on the last class, which by convention is the "likely" / high-score class. For 3-class {unlikely, ambiguous, likely}, this is P(likely). Surfaces with different taxonomies can override which class index is the "headline" by passing it in. Args: logits: (B, num_classes). Returns: (B,) in [0, 1] — calibrated probability mass on the highest-severity class. """ # (B, num_classes) -> (B,) return F.softmax(logits, dim=-1)[..., -1] @torch.no_grad() def predict_band(self, logits: torch.Tensor) -> torch.Tensor: """Predicted class index. Argmax over softmax. Args: logits: (B, num_classes). Returns: (B,) int64 class indices. """ return logits.argmax(dim=-1) def num_parameters(self) -> int: return sum(p.numel() for p in self.parameters())