"""K-treatment × 3-band probability head for the Collections surface. Where the dispute surface has one 3-class output (unlikely/ambiguous/ likely), Collections has K independent 3-class outputs, one per treatment option (settlement / payment_plan / soft_touch / no_offer). Each treatment is its own categorical distribution over the same three bands. Why this design over alternatives: - K separate `ProbabilityHead` instances would multiply parameters and complicate the model's __init__ wiring. The K outputs share most representation; one head emitting (K * 3) logits is cleaner. - K independent sigmoids ("probability customer responds") would collapse the calibration story to a continuous-attribute prediction on 350M — exactly the attractor-state pathology the doctrine warns about (liquid-finetuning-playbook §10). Categorical over 3 bands per treatment preserves softmax calibration. - A flat K*3-way softmax mixes treatments, breaking the per-treatment independence we want. We softmax along the band axis only. Shape contract: Input hidden_states: (B, T_total, D) Output logits: (B, K, num_bands) — K=4 treatments, 3 bands Loss target: (B, K) int64 — per-treatment band index """ from __future__ import annotations from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F @dataclass class MultiTreatmentProbabilityHeadConfig: """Multi-treatment probability head configuration. Attributes: name: Identifier for logging. num_treatments: K — number of treatment options. Collections ships K=4 (settlement / payment_plan / soft_touch / no_offer). num_bands: Bands per treatment. 3 for the {unlikely_respond, ambiguous, likely_respond} taxonomy. hidden_dim: Backbone hidden size (1024 for LFM2.5-350M). mlp_hidden: Intermediate hidden size of the 2-layer head MLP. 512 is a touch bigger than the dispute head's 256 because the output dim is 12 (K=4 × bands=3) instead of 3. dropout: Dropout between MLP layers. num_tx_positions: Number of leading transaction pseudo-tokens to mean-pool over (64 in Phase 1). band_class_weights: Optional per-band CE weights of length `num_bands`. Useful when the corpus has rare bands (collections no_offer-LIKELY is rare per Day 2 audit). If None, all bands are weighted equally. """ name: str num_treatments: int = 4 num_bands: int = 3 hidden_dim: int = 1024 mlp_hidden: int = 512 dropout: float = 0.1 num_tx_positions: int = 64 band_class_weights: list[float] | None = None class MultiTreatmentProbabilityHead(nn.Module): """K independent 3-band probability outputs over a shared pooled rep. Parameter count at defaults (D=1024, mlp_hidden=512, K=4, bands=3): ~530K. About 2x the dispute probability head; still immaterial against the 354M backbone. The head mean-pools over the 64 transaction positions (NOT the text tokens), MLPs to (K * num_bands), reshapes to (B, K, bands). Softmax is along the band axis only — each treatment's distribution is independent of the others. """ def __init__(self, config: MultiTreatmentProbabilityHeadConfig) -> None: super().__init__() self.config = config self.output_dim = config.num_treatments * config.num_bands self.mlp = nn.Sequential( nn.Linear(config.hidden_dim, config.mlp_hidden), nn.ReLU(), nn.Dropout(config.dropout), nn.Linear(config.mlp_hidden, self.output_dim), ) if config.band_class_weights is not None: if len(config.band_class_weights) != config.num_bands: raise ValueError( f"band_class_weights must have length {config.num_bands}, " f"got {len(config.band_class_weights)}", ) self.register_buffer( "band_weights", torch.tensor(config.band_class_weights, dtype=torch.float32), ) else: self.register_buffer("band_weights", None, persistent=False) 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). Returns: (B, D) pooled 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 head. Args: hidden_states: (B, T_total, D). Returns: (B, num_treatments, num_bands) raw logits. Caller applies softmax along the last axis for per-treatment probabilities or per-treatment cross-entropy for loss. """ pooled = self.pool(hidden_states) # (B, D) flat_logits = self.mlp(pooled) # (B, K * bands) return flat_logits.view( -1, self.config.num_treatments, self.config.num_bands, ) # (B, K, bands) def compute_loss( self, logits: torch.Tensor, targets: torch.Tensor, ) -> torch.Tensor: """Mean CE across treatments and batch. Args: logits: (B, K, num_bands) raw logits. targets: (B, K) int64 band indices in [0, num_bands). Returns: Scalar CE loss averaged over (batch × treatments). If `band_class_weights` was set, the loss is weighted per band. """ flat_logits = logits.reshape(-1, self.config.num_bands) flat_targets = targets.reshape(-1).long() weight = self.band_weights if self.band_weights is not None else None return F.cross_entropy( flat_logits, flat_targets, weight=weight, reduction="mean", ) @torch.no_grad() def score(self, logits: torch.Tensor) -> torch.Tensor: """Per-treatment likely-respond probability for the UI. Returns softmax(logits)[..., -1] — the probability mass on the last band (LIKELY_RESPOND) for each treatment. Args: logits: (B, K, num_bands). Returns: (B, K) in [0, 1] — P(likely_respond) per treatment. """ return F.softmax(logits, dim=-1)[..., -1] @torch.no_grad() def predict_band(self, logits: torch.Tensor) -> torch.Tensor: """Per-treatment argmax band. Args: logits: (B, K, num_bands). Returns: (B, K) int64 band indices in [0, num_bands). """ return logits.argmax(dim=-1) @torch.no_grad() def dominant_treatment(self, logits: torch.Tensor) -> torch.Tensor: """Index of the treatment with the highest LIKELY-band probability. Ties broken in argmax's deterministic order (lower index wins). Args: logits: (B, K, num_bands). Returns: (B,) int64 treatment indices in [0, K). """ return self.score(logits).argmax(dim=-1) def num_parameters(self) -> int: return sum(p.numel() for p in self.parameters())