src/model/heads.py

"""Prediction heads for the transaction foundation model.

PerFeatureLMHeads: 15 linear projections for next-token prediction during
pretraining. Each head projects hidden_dim -> feature_vocab_size with weights
tied to the corresponding value embedding table (parameter identity, not copy).
No bias, matching LFM2's lm_head convention.

Weight tying semantics:
  - Forward through embedding: output = weight[token_id]  (index select)
  - Forward through LM head: logits = hidden @ weight.T   (matmul)
  - Backward: gradients flow through BOTH paths to the SAME Parameter.
    The optimizer updates it once, accumulating both contributions.
  - The feature-type embedding table is NOT tied to any head (it has no
    corresponding prediction target).

The backbone's final RMSNorm (embedding_norm) is applied to hidden states
BEFORE they reach these heads, matching LFM2's architecture. That
normalization lives in the backbone (F8), not here.
"""

import torch
import torch.nn as nn

from src.model.embedding import StructuredEmbedding


class PerFeatureLMHeads(nn.Module):
    """Per-feature prediction heads with tied embedding weights.

    Each of the 15 heads is a bias-free linear projection from hidden_dim to
    that feature's vocab_size. The weight matrix is the SAME nn.Parameter as
    the corresponding value embedding table.
    """

    def __init__(self, embedding: StructuredEmbedding) -> None:
        super().__init__()
        self.num_features = embedding.num_features
        self.hidden_dim = embedding.hidden_dim

        self.heads = nn.ModuleList()
        for f_idx in range(self.num_features):
            vocab_size = embedding.value_tables[f_idx].num_embeddings
            head = nn.Linear(embedding.hidden_dim, vocab_size, bias=False)
            # Tie weights: head.weight IS embedding.value_tables[f].weight.
            # Both the embedding lookup and the LM head projection operate on
            # the same underlying tensor. PyTorch deduplicates in parameters().
            head.weight = embedding.value_tables[f_idx].weight
            self.heads.append(head)

    def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:
        """Project hidden states to per-feature logits.

        Args:
            hidden_states: (B, S, D) from the backbone (after embedding_norm).

        Returns:
            List of 15 tensors. heads[f] has shape (B, S, vocab_size_f).
            During training, the loss function selects positions per feature:
            position p predicts feature (p+1) % num_features.
        """
        return [head(hidden_states) for head in self.heads]

    def get_vocab_sizes(self) -> list[int]:
        """Vocab size for each feature, useful for loss computation."""
        return [head.out_features for head in self.heads]