SAEVerbalizer: Generating Explanations for Sparse Autoencoder Features via Representation Verbalization

Overview§

SAEVerbalizer addresses the bottleneck of explaining sparse autoencoder (SAE) features in large language models (LLMs). Traditional explainability methods rely on external behavioral probes (e.g., dataset activation maximization) which are slow and yield shallow descriptions. SAEVerbalizer instead learns to verbalize the SAE decoder direction directly: it injects a feature direction into the LLM's residual stream and fine-tunes the subsequent layers to generate a natural-language explanation.

Method§

The core idea is representation verbalization. Given a pretrained LLM with layers $L_1, \dots, L_N$ and an SAE with decoder directions $\{\mathbf{d}_i\}$ (unit-normalized), the residual stream at layer $k$ is perturbed by adding a scaled feature direction:

$$ \mathbf{x}_k \leftarrow \mathbf{x}_k + \alpha \hat{\mathbf{d}}_i $$

where $\alpha$ is a scalar controlling the intervention strength (typically chosen to match the feature activation scale). Only layers $k+1, \dots, N$ are fine-tuned (optionally with LoRA adapters), while lower layers remain frozen. The training objective is the standard next-token cross-entropy between the generated explanation and a target explanation (e.g., from an existing auto-labeler like TCAV or GPT-4).

The training dataset consists of SAE features with known explanations. Each feature's decoder vector is injected into the middle of the LLM, and the fine-tuned upper layers are trained to produce the target text. This forces the model to 'read' the geometric direction and map it to linguistic meaning.

Training Details§

• Feature injection: At layer $k$, the residual stream is modified as above. The injection point is chosen to balance information retention (lower layers encode general language) and semantics (higher layers ready for generation).
• Parameter efficiency: Only the layers above $k$ are updated, and in practice a low-rank adapter (LoRA) is applied, enabling cheap training.
• Loss: Cross-entropy on the explanation tokens. No additional contrastive or auxiliary losses are required.

Transfer and Generalization§

Once trained, the verbalizer generalizes to unseen SAE features (not in training) from the same dictionary. It also transfers to features from a different SAE trained on the same LLM, and—with a lightweight adapter (a linear projection on the decoder direction)—to SAEs from other LLMs. Intervention experiments show that injecting multiple directions yields a blended explanation (e.g., 'cat' + 'dog' → 'pet'), while reversing a direction flips the semantic polarity.

Code Sketch§

The following pseudocode illustrates the training loop in PyTorch-style pseudo-code:

import torch
import torch.nn as nn

class SAEVerbalizer(nn.Module):
    def __init__(self, base_model, injection_layer_idx, adapter=None):
        super().__init__()
        self.base_model = base_model
        self.injection_layer_idx = injection_layer_idx
        self.adapter = adapter  # optional linear projection for cross-model transfer
        # Freeze all layers except those above injection_layer_idx
        for name, param in self.base_model.named_parameters():
            if int(name.split('.')[1]) <= injection_layer_idx:
                param.requires_grad = False

    def inject_feature(self, hidden_states, d_i, alpha):
        # d_i: decoder direction for feature i (normalized)
        if self.adapter is not None:
            d_i = self.adapter(d_i)
        return hidden_states + alpha * d_i.unsqueeze(0).unsqueeze(0)

    def forward(self, input_ids, d_i, alpha):
        outputs = self.base_model(
            input_ids,
            output_hidden_states=True,
            return_dict=True
        )
        hidden = outputs.hidden_states[self.injection_layer_idx]
        hidden = self.inject_feature(hidden, d_i, alpha)
        # Replace the original hidden state and continue forward through upper layers
        # (simplified; in practice you would use hooks or a custom forward)
        return self.base_model.forward_with_hidden_state(hidden, attention_mask=...)

# Training loop
model = SAEVerbalizer(load_llm(), injection_layer_idx=12)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
for feature_direction, target_explanation in dataloader:
    input_ids = tokenizer(target_explanation, return_tensors='pt').input_ids
    logits = model(input_ids[:, :-1], feature_direction, alpha)
    loss = nn.CrossEntropyLoss()(logits, input_ids[:, 1:])
    loss.backward()
    optimizer.step()

Key Takeaways§

• Eliminates the need for external behavioral probing by directly decoding the geometric meaning encoded in SAE directions.
• Achieves cross-dictionary and cross-model transfer, demonstrating that SAE directions share a common semantic geometry.
• The injection-based approach enables compositional and polarity-based interventions, providing a richer interpretability tool than static feature-labeling methods.