1706_scope/models.py

"""
Neural network modules for SAE-based topic classification.
"""

import torch
import torch.nn as nn


class TopKSAE(nn.Module):
    """Sparse autoencoder that keeps only the top-k activations per token."""

    def __init__(self, d_model, d_sae, k):
        super().__init__()
        self.d_model = d_model
        self.d_sae = d_sae
        self.k = k
        # register_buffer so weights move with .to(device) and appear in state_dict
        self.register_buffer("W_enc", torch.zeros(d_sae, d_model))  # (d_sae, d_model)
        self.register_buffer("b_enc", torch.zeros(d_sae))

    def topk_indices(self, x):
        """Return indices of the top-k pre-activations for each token in x."""
        # W_enc is (d_sae, d_model), so x @ W_enc.T gives (n_tokens, d_sae)
        pre = x @ self.W_enc.T + self.b_enc
        return pre.topk(self.k, dim=-1).indices


class TopicClassifier(nn.Module):
    """Two-layer MLP that maps SAE feature vectors to topic logits."""

    def __init__(self, n_in, n_classes, hidden):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_in, hidden),
            nn.ReLU(),
            nn.Linear(hidden, n_classes),
        )

    def forward(self, x):
        return self.net(x)