mechanistic_interp/sequence_probe.py

"""Attention-pooling sequence probes over raw residual-stream hidden states.

Self-contained port of multilayer-sae's ``experiment/training/sequence_probe.py``
(no ``experiment.*`` dependencies) for use by ``train_probe_latent.py``.

Replaces the linear + max-pool probe (``probe/probing.py:LinearProbe`` trained by
``train_binary_probe``) with a per-layer probe that reads the *sequence* of
per-token hidden states directly, with no fixed pooling:

  * Operating on raw ``h_l`` (d_model) keeps the probe independent of any SAE
    dictionary.
  * A learned-query multi-head attention pool replaces max/mean pooling so no
    information is discarded by a hard reduction, and the probe co-adapts to
    *where* the concept lives in the token sequence (follows HyperSteer,
    arXiv:2506.03292).

``forward_logits`` gives one per-sequence logit per layer (for training on a
per-caption label); ``forward_token_logits`` gives a per-token logit per layer
(for detecting the concept live during generation) — both reuse the same weights.
"""
from __future__ import annotations

import torch
import torch.nn as nn


class _AttnPoolProbe(nn.Module):
    """Single-layer attention-pooling probe: (B, T, d_model) -> (B,) logit."""

    def __init__(
        self,
        d_model: int,
        d_probe: int = 256,
        n_heads: int = 4,
        n_ctx_blocks: int = 1,
        dropout: float = 0.0,
        spectral_norm: bool = False,
    ):
        super().__init__()
        self.in_proj = nn.Linear(d_model, d_probe)
        self.ctx_blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_probe,
                nhead=n_heads,
                dim_feedforward=4 * d_probe,
                dropout=dropout,
                batch_first=True,
                norm_first=True,
            )
            for _ in range(n_ctx_blocks)
        ])
        # Learned query that attends over the token sequence (attention pooling).
        self.query = nn.Parameter(torch.zeros(1, 1, d_probe))
        nn.init.normal_(self.query, std=0.02)
        self.pool_attn = nn.MultiheadAttention(
            d_probe, n_heads, dropout=dropout, batch_first=True
        )
        self.norm = nn.LayerNorm(d_probe)
        head = nn.Linear(d_probe, 1)
        self.head = nn.utils.spectral_norm(head) if spectral_norm else head

    def forward(
        self, x: torch.Tensor, key_padding_mask: torch.Tensor | None
    ) -> torch.Tensor:
        # x: (B, T, d_model); key_padding_mask: (B, T) bool, True = ignore (PAD).
        x = self.in_proj(x)                                    # (B, T, d_probe)
        for blk in self.ctx_blocks:
            x = blk(x, src_key_padding_mask=key_padding_mask)
        B = x.shape[0]
        q = self.query.expand(B, -1, -1)                       # (B, 1, d_probe)
        pooled, _ = self.pool_attn(
            q, x, x, key_padding_mask=key_padding_mask, need_weights=False
        )                                                      # (B, 1, d_probe)
        pooled = self.norm(pooled.squeeze(1))                  # (B, d_probe)
        return self.head(pooled).squeeze(-1)                   # (B,)

    def token_logits(
        self, x: torch.Tensor, key_padding_mask: torch.Tensor | None
    ) -> torch.Tensor:
        """Per-token logits (no attention pooling): (B, T, d_model) -> (B, T).

        Reuses in_proj + context blocks + norm + head so a single probe module
        supports both the per-sequence (attention-pooled) and per-token readouts.
        Used to detect the concept live during generation.
        """
        x = self.in_proj(x)                                    # (B, T, d_probe)
        for blk in self.ctx_blocks:
            x = blk(x, src_key_padding_mask=key_padding_mask)
        return self.head(self.norm(x)).squeeze(-1)             # (B, T)


class SequenceLayerProbes(nn.Module):
    """One attention-pooling probe per monitored layer, run on raw hidden states."""

    def __init__(
        self,
        layer_indices: list[int],
        d_model: int,
        d_probe: int = 256,
        n_heads: int = 4,
        n_ctx_blocks: int = 1,
        dropout: float = 0.0,
        spectral_norm: bool = False,
    ):
        super().__init__()
        self.layer_indices = list(layer_indices)
        self.d_model = d_model
        self.probes = nn.ModuleList([
            _AttnPoolProbe(
                d_model, d_probe=d_probe, n_heads=n_heads,
                n_ctx_blocks=n_ctx_blocks, dropout=dropout,
                spectral_norm=spectral_norm,
            )
            for _ in layer_indices
        ])
        self._idx = {l: i for i, l in enumerate(self.layer_indices)}

    def forward_logits(
        self,
        feats_seq: dict[int, torch.Tensor],     # l -> (B, T, d_model)
        key_padding_mask: torch.Tensor,         # (B, T) bool, True = PAD/ignore
    ) -> list[torch.Tensor]:
        # Probe runs in fp32 for numerical stability regardless of model dtype.
        w_dtype = self.probes[0].in_proj.weight.dtype
        return [
            self.probes[self._idx[l]](
                feats_seq[l].to(w_dtype), key_padding_mask
            )
            for l in self.layer_indices
        ]

    def forward_token_logits(
        self,
        feats_seq: dict[int, torch.Tensor],     # l -> (B, T, d_model)
        key_padding_mask: torch.Tensor,         # (B, T) bool, True = PAD/ignore
    ) -> list[torch.Tensor]:
        """Per-token logits per layer: list of (B, T). For live-generation detection."""
        w_dtype = self.probes[0].in_proj.weight.dtype
        return [
            self.probes[self._idx[l]].token_logits(feats_seq[l].to(w_dtype), key_padding_mask)
            for l in self.layer_indices
        ]

    def forward(
        self, feats_seq: dict[int, torch.Tensor], key_padding_mask: torch.Tensor
    ) -> list[torch.Tensor]:
        return [torch.sigmoid(z) for z in self.forward_logits(feats_seq, key_padding_mask)]


def sequence_layer_probes_from_checkpoint(
    path: str,
    device: torch.device | str | None = None,
) -> SequenceLayerProbes:
    """Load a SequenceLayerProbes checkpoint saved by train_probe_latent.py.

    The checkpoint stores both the state_dict and the construction meta
    (layer_indices, d_model, probe_dim, ...) so the architecture is rebuilt
    exactly. Handles a ``module.`` DDP prefix.
    """
    # Checkpoint holds only tensors + a plain dict of ints/lists, so the safe
    # weights_only loader suffices (no arbitrary-object unpickling).
    ckpt = torch.load(path, map_location="cpu", weights_only=True)
    meta = ckpt["meta"]
    sd = ckpt["state_dict"]
    if any(k.startswith("module.") for k in sd):
        sd = {k.replace("module.", "", 1): v for k, v in sd.items()}
    probes = SequenceLayerProbes(
        meta["layer_indices"], meta["d_model"],
        d_probe=meta["d_probe"], n_heads=meta["n_heads"],
        n_ctx_blocks=meta["n_ctx_blocks"], spectral_norm=meta["spectral_norm"],
    )
    probes.load_state_dict(sd, strict=True)
    if device is not None:
        dev = device if isinstance(device, torch.device) else torch.device(device)
        probes = probes.to(dev)
    return probes