| """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. |
| """ |
| |
| |
| |
| 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. |
| """ |
| |
| pooled = self.pool(hidden_states) |
| |
| 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. |
| """ |
| |
| 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()) |
|
|
