Technical Summary: Uncovering Language Model Processing Strategies with Non-Negative Per-Example Fisher Factorization

Problem Statement

Understanding the heuristics and algorithms that constitute a language model's behavior is critical for safe and reliable deployment. While previous work on Transformer circuits has successfully identified interpretable computational graphs for specific, pre-specified behaviors (e.g., indirect object identification), these approaches require researchers to define the behaviors of interest beforehand, introducing bias and risking the omission of unintuitive model strategies.

Recent unsupervised methods, such as gradient clustering, attempt to discover model behaviors by grouping examples based on the gradients of the loss with respect to model parameters. However, these methods suffer from two primary limitations:

1. Monogenic Assumption: They operate under the assumption that a single internal module drives the behavior for each example. In reality, model behaviors are often "polygenic," meaning they are influenced by multiple factors simultaneously (e.g., ambiguous contexts involving multiple possible continuations). Gradient clustering forces a single factor assignment, failing to capture this complexity.
2. Limited Scope of Gradients: Gradients of a single log-probability (or loss) capture only a narrow slice of the model's predictive distribution. They are biased toward the specific class being predicted, making them poorly suited for capturing factors that influence predictions across different classes, particularly in high-entropy or multi-class scenarios.

Methodology: NPEFF

The authors introduce NPEFF (Non-Negative Per-Example Fisher Factorization), an interpretability method designed to uncover polygenic behavioral factors without pre-specifying them.

1. Per-Example Fisher (PEF) Matrices

Instead of using loss gradients, NPEFF utilizes the Per-Example Fisher (PEF) matrix, $F(x)$, defined as:
$$F(x) = \mathbb{E}_{y \sim p_\theta(y|x)} [\nabla_\theta \log p_\theta(y|x) \nabla_\theta \log p_\theta(y|x)^T]$$
The PEF is a positive semi-definite (PSD) matrix that relates parameter perturbations to changes in the model's entire predictive distribution. Unlike gradients, which focus on a single outcome, the PEF captures sensitivity across all potential predictions, making it suitable for analyzing polygenic behaviors where multiple factors influence the distribution.

2. Efficient Computation via Random Projections

PEF matrices are computationally intractable for modern models due to their size (quintillions of entries). To address this, the authors employ random projections (Achlioptas, 2003) to approximate the PEF. They construct a low-rank matrix representation (LRM-PEF) and apply random projections to reduce dimensionality. Crucially, they demonstrate that these projections can be "reversed" using compressed sensing techniques (Donoho, 2006) to recover sparse parameter-space vectors, allowing the identification of specific model directions corresponding to behavioral factors.

3. Non-Negative Decomposition

The core of NPEFF is a decomposition algorithm that approximates a set of PEFs as a non-negative combination of rank-1 PSD matrices:
$$\min \sum_{i=1}^m \| G_i^T G_i - \sum_{j=1}^C W_{ij} h_j h_j^T \|_F^2 \quad \text{subject to} \quad W_{ij} \geq 0$$
Here, $h_j h_j^T$ represents the "pseudo-Fisher" of component $j$, and $W_{ij}$ are non-negative coefficients. This formulation allows each example to be explained by a weighted combination of multiple factors, directly modeling polygenicity. The algorithm alternates between updating coefficients (using multiplicative updates similar to NMF) and updating pseudo-Fisher vectors (using gradient descent).

4. G-NPEFF Variant

To reduce computational costs, the authors propose G-NPEFF, which applies the same decomposition algorithm to rank-1 matrices constructed from the gradient of the log-probability of the predicted class/token. While cheaper, G-NPEFF is shown to be less effective at capturing polygenic factors as the number of classes increases, as it remains biased toward the dominant predicted class.

Key Contributions and Results

1. Discovery of Interpretable Polygenic Factors

Through human and LLM evaluations on tasks like sentiment analysis (SST2), topic identification (Yahoo Answers Topics), and intent classification (CLINC150), the authors demonstrate that NPEFF components correspond to interpretable themes (e.g., "questions about rapidly losing weight" or "positive descriptions of actors").

• Polygenicity: Unlike gradient clustering, which recovers almost exclusively "monogenic" (single-factor) behaviors, NPEFF successfully recovers a significant fraction of polygenic factors. This is evidenced by a high percentage of components that are tuned to ground-truth labels but not to the model's specific predictions (LnP-tuned), indicating they capture factors influencing the distribution beyond the top prediction.
• Comparison: NPEFF outperforms baselines including gradient clustering, activation sparse autoencoders (SAEs), and SVD on gradients/activations in identifying these nuanced, multi-factor behaviors.

2. Causal Validation via Parameter Perturbations

The authors validate the causal importance of NPEFF components by reconstructing parameter-space vectors from projected pseudo-Fishers using compressed sensing. They apply perturbations along these directions and observe that:

• Perturbations significantly disrupt the model's predictions on examples where the component has a high coefficient.
• The disruption is selective, affecting the targeted examples far more than random examples.
• This confirms that the directions uncovered by NPEFF are causally linked to the model's processing of specific behavioral factors.

3. Analysis of In-Context Learning (ICL)

The authors apply NPEFF to analyze ICL. They construct a linear classifier based on zero-shot NPEFF component tunings. This classifier achieves accuracy comparable to 6-shot ICL and makes highly similar predictions.

• Finding: This suggests that much of the performance gain from ICL comes from the model re-weighting its existing zero-shot behavioral factors to adapt to the task formulation, rather than learning entirely new behaviors from the context.

4. Efficiency and Scalability

• G-NPEFF: The paper shows that for tasks with few classes, G-NPEFF performs similarly to full NPEFF, offering a cost-effective alternative.
• Ablations: Experiments on random projection sizes and component counts indicate that the method is robust. Even with aggressive dimensionality reduction, the decomposition retains significant information about model behavior.

Significance and Claims

The paper claims that NPEFF represents a novel and well-suited method for uncovering factors of polygenic model behaviors, addressing a critical gap left by gradient-based clustering methods.

• Methodological Novelty: By utilizing Per-Example Fisher matrices and a non-negative decomposition over rank-1 PSD matrices, NPEFF provides a framework that respects the multi-factor nature of language model processing.
• Causal Interpretability: The ability to map components back to parameter-space directions and selectively disrupt them provides strong causal evidence that these components represent genuine processing strategies used by the model.
• Practical Utility: The method enables the dissection of complex behaviors (like ICL) into their constituent factors, offering a path toward more targeted model modification and safety interventions.

The authors remain modest regarding the scope of their work, noting that while NPEFF identifies the directions in parameter space important for specific behaviors, constructing the explicit algorithms or computational circuits corresponding to these directions remains the purview of future transformer circuit research. They also caution that while their perturbation code allows for disruption, sophisticated modification of model behavior at scale carries risks of inadvertently removing rare but contextually appropriate behaviors.