Technical Summary: Turn-Averaged SAEs for Feature Discovery and Long-Context Attribution

1. Problem Statement

Sparse Autoencoders (SAEs) have emerged as a primary tool for extracting interpretable features from language models, facilitating circuit analysis, behavioral steering, and safety monitoring. However, standard SAE architectures operate on individual token activations. This per-token granularity creates two significant bottlenecks for interpretability at scale:

1. Scaling Limitations: The number of active features scales linearly with context length. In long-context scenarios or multi-turn conversations, the resulting data volume makes downstream analysis (such as attribution graphs) cumbersome and difficult to inspect.
2. Loss of High-Level Context: Many interpretability goals involve understanding why a model produced a specific response, requiring the identification of high-level properties (e.g., topic, persona, safety traits) that are distributed across many tokens. Standard per-token SAEs often learn low-level detectors for specific words or syntactic patterns. Because the training objective reconstructs each position independently, there is no pressure for the model to separate these high-level, span-wide characteristics from token-level noise.

2. Methodology

2.1 Turn-Averaged SAEs

The core proposal is to train SAEs on the mean hidden state across a span of tokens (specifically, a single Human or Assistant turn) rather than on individual token activations.

• Input Transformation: Given a turn $T$ with tokens $x_t$, the input is the average hidden state $\bar{x} = \frac{1}{|T|} \sum_{t \in T} x_t$.
• Architecture: The SAE architecture remains identical to standard implementations (using BatchTopK for sparsity). Since the mean of vectors in $\mathbb{R}^{d_{model}}$ retains the same dimensionality as individual vectors, no structural changes are required.
• Intuition: Averaging acts as a low-pass filter. It smooths out high-frequency, token-specific details (like specific word choices or syntactic markers) while preserving low-frequency components that correspond to broad concepts like topic, style, and behavioral traits.

2.2 Hybrid Architectures

To address the limitation that averaging discards per-token detail, the authors propose two hybrid approaches to capture both high-level structure and fine-grained residuals:

1. Decoupled SAE: Trains two independent SAEs. One reconstructs the turn-averaged mean ($\bar{x}$), and the other reconstructs the residual ($x_t - \bar{x}$). While this separates the signals, the models are causally disconnected, complicating attribution analysis.
2. Nested SAE (Matryoshka-style): Trains a single SAE with a shared encoder and decoder but a specialized loss function. The feature space is divided into an inner partition $[0:h]$ and an outer partition $[h:d_{sae}]$.
   • Loss Function: $L = \sum_t \| \text{dec}(\text{enc}(x_t)) - x_t \|^2 + \alpha \cdot \| \text{dec}_h(\text{enc}(\bar{x})[0:h]) - \bar{x} \|^2$.
   • The inner partition is forced to reconstruct the turn-averaged mean, while the full set of features reconstructs the per-token activations. This allows a single model to learn both high-level turn characteristics and token-level details simultaneously.

2.3 Attribution Graph Adaptation

The authors adapt attribution graphs to the turn level. Instead of nodes representing features at specific token positions, nodes represent features active on a specific turn and layer.

• Edge Weights: Calculated by averaging gradient attribution across the source turn and normalizing by turn length ($1/|T_A|$ and $1/|T_B|$) to ensure weights are agnostic to turn length.
• Intervention Validation: Edge weights are validated by perturbing source features along their decoder directions across all tokens in a turn and measuring the change in target feature pre-activations.

3. Key Contributions

1. Turn-Averaged SAEs: Demonstrated that training on mean hidden states produces features that capture high-level turn characteristics (topic, style, safety traits) more completely than per-token features.
2. Nested SAE Architecture: Introduced a single-model architecture that jointly learns turn-averaged and per-token features using a Matryoshka-like loss, achieving high reconstruction quality for both signal components.
3. Scalable Attribution: Adapted attribution graphs to the turn level, enabling compact analysis of multi-turn conversations. The authors validated these graphs through intervention experiments and sufficiency metrics (completeness and replacement).

4. Results

4.1 Feature Quality and Reconstruction

• Reconstruction: The Nested SAE achieves a per-token Normalized Mean Squared Error (nMSE) of 0.167, closely matching the pure per-token SAE (0.162), while also achieving a turn-averaged nMSE of 0.116. In contrast, a standard per-token SAE performs poorly on turn-averaged inputs (nMSE 0.355), indicating the mean is out-of-distribution for token-trained models.
• Qualitative Discovery: On examples involving mathematical errors, enthusiastic technical styles, and regulatory compliance, turn-averaged features successfully identified high-level behavioral properties (e.g., "incorrect answers," "enthusiastic technical style") that per-token features missed or only partially captured.
• Long-Context Generalization: Turn-averaged SAEs trained on chat turns (average length ~150 tokens) successfully generalized to documents 150x longer (24k–32k tokens) from LongBench v2, identifying high-level concepts like "macroeconomics" and "computer graphics" that per-token features failed to detect.

4.2 Quantitative Evaluation

Using an LLM-as-a-judge framework to compare feature descriptions against structured summaries of turns:

• Discrimination (10-way matching): Per-token features (max-pooled) outperformed turn-averaged features (95.0% vs. 73.9%), as they rely on specific token-level details to distinguish texts.
• Coverage (Pairwise and 5-way ranking): Turn-averaged features significantly outperformed per-token features (winning 87.9% of pairwise comparisons). They provided more complete descriptions of a turn's high-level characteristics.
• Embedding Metric: Turn-averaged features achieved the highest cosine similarity (0.663) to structured summary fields, confirming their superior coverage of semantic content.

4.3 Attribution Analysis

• Graph Compactness: Turn-averaged graphs reduced node counts from ~128K (per-token) to ~5K for a 10-turn conversation, making inspection feasible.
• Case Studies:
  • System Prompt Attribution: Turn-averaged graphs correctly identified which system prompt paragraphs influenced specific user queries, showing attribution shifts comparable to per-token graphs.
  • Degeneration Tracing: In a 14-turn conversation where the model degraded into repetitive text, the turn-averaged graph traced a coherent causal chain: "AI safety disclaimers" $\to$ "Nuclear weapons" $\to$ "Adversarial prompt detection" $\to$ "Degenerate output." Per-token graphs failed to find the adversarial detection feature, tracing instead to unrelated concepts like "chemical compounds."
• Intervention & Sufficiency: Edge weights showed significant correlation with intervention effects. The Nested SAE achieved the highest "completeness" and "replacement" sufficiency scores, outperforming both pure per-token and post-hoc aggregated per-token features.

5. Significance and Claims

The paper claims that turn-averaged SAEs make interpretability techniques practical at long context lengths. By shifting the unit of analysis from tokens to spans, the method:

• Produces features that describe high-level characteristics (topic, persona, safety) more completely than token-level features.
• Enables the construction of compact, interpretable attribution graphs for multi-turn conversations that would otherwise be too large to analyze.
• Demonstrates that modifying the training signal (via averaging) can compel SAEs to learn features at a desired granularity (high-level concepts) independent of token count.
• Shows that these features generalize to spans far longer than those seen during training, suggesting robustness for frontier model analysis.

The authors note limitations, including the current scale of experiments (7B model, LMSYS-Chat-1M corpus) and the reliance on LLMs for evaluation metrics, identifying the application to frontier models and the creation of ground-truth benchmarks as necessary future work.