encoder/src/model/heads/attribution_head.py

"""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())