| """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) |
| flat_logits = self.mlp(pooled) |
| return flat_logits.view( |
| -1, self.config.num_treatments, self.config.num_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()) |
|
|
