From Directions to Regions: Decomposing Activations in Language Models via Local Geometry

Activation decomposition methods in language models are tightly coupled to geometric assumptions on how concepts are realized in activation space, and this work introduces Mixture of Factor Analyzers (MFA) as a scalable, unsupervised alternative that models activation space as a collection of Gaussian regions with their local covariance structure. MFA decomposes activations into two compositional geometric objects: the region’s centroid in activation space, and the local variation from the centroid.

The gist

MFA maps the activation space into Gaussian regions, where each region is modeled with a low-dimensional subspace that captures structured within-region variation, thereby providing an interpretable decomposition of language model activations.

How it works

1. MFA models the space as a collection of local low-dimensional Factor Analyzers (FA), where each component models a different region of the representation space.

2. Each component k is defined by a mean centroid µk, which sets the center of its region, and a factor analysis model with a component-specific matrix Wk, which determines the orientation of its local low-dimensional subspace.

3. The generative model for an activation x conditioned on component k is given by x = µk + Wkzk + ϵ, where z is a latent factor vector and ϵ is noise with a diagonal covariance matrix Ψ.

4. The overall density of the activation space is represented as a mixture of these component densities: p(x) = Σ πk N (x µk, Ck), where πk are the mixture weights.

Decomposition and Interpretation

The decomposition yields two compositional geometric objects for any given activation x: a region in activation space defined by the assigned centroid (Rk(x) µk), and a within-region offset defined by the local variation captured by the component's subspace (Rk(x) zˆk). The reconstruction of an activation is then written as a linear product: x ≈ A b(x), where A is formed by concatenating the component means and loadings, and b(x) concatenates the responsibilities and latent coordinates. This structure moves the unit of analysis from isolated global directions to local regions with their own low-rank geometry.

Structural Analysis of Discovered Regions

Training large-scale MFAs on Llama-3.1-8B and Gemma-2-2B reveals rich structures in the activation space, characterized by two classes of Gaussians: narrow Gaussians that concentrate on a constrained lexical pattern (showing more syntactic variance) and broad Gaussians that encompass wide thematic topics (often exhibiting semantic local variation). Analyzing these regions shows that neighboring components tend to encode related semantics, suggesting concepts may be realized by constellations of nearby Gaussians. Furthermore, the analysis of loadings reveals that within-region variation often reflects both semantic and syntactic differences, with narrow Gaussians skewing more syntactic and broad Gaussians skewing more semantic.

Performance on Localization and Steering

MFA is evaluated on localization (MCQA and RAVEL) and steering benchmarks. On localization, MFA outperforms large-scale SAEs by large margins, beating supervised baselines like Desiderata-Based Masking (DBM) in several tasks. For causal steering, utilizing the MFA centroids steers better than SAE features in the majority of settings, typically exhibiting a twofold gain on coherence and conceptual alignment. The results indicate that absolute positions learned by the centroids are an effective unit for steering, while local offsets capture fine-grained shifts within the broader theme. MFA achieves an average interpretability fraction (IF) of 0.96 ± 0.2, significantly higher than SAEs' average IF of 0.29 ± 0.2, indicating that most of the high-contribution features in its decomposition are interpretable.

Reconstruction and Limitations

When comparing reconstruction error on a held-out validation set, MFA substantially outperforms the K-means baseline, which consistently has 1.3 − 1.5 times higher MSE, demonstrating that modeling within-region variation with a learned low-rank structure captures a considerable portion of the structured representation. However, a key limitation is that MFA explicitly models the activation distribution it is trained on. This means that when an activation lies in a region that is rare or out of distribution of the training set, MFA may assign it to the nearest available component even if none provides a good local fit, leading to high reconstruction error. Despite this, MFA isolates meaningful and useful features that generalize to tasks like steering and localization.

Conclusion

MFA proposes a local-geometry view of activation space, partitioning it into low-dimensional regions and modeling the intrinsic modes of variation within each region, offering a scalable approach to model control that generalizes across layers and models.

Here are the specific improvements that can be made to AI systems by implementing the concepts from this scientific paper, and what those improved systems will be able to do:

The core improvement lies in shifting from analyzing activations as isolated global directions (the limitation of Sparse Autoencoders or simple Global Direction methods) to analyzing them as a structured collection of local regions defined by low-dimensional Gaussian distributions.

Here are the specific improvements and resulting capabilities:

1. Implementation of Mixture of Factor Analyzers (MFA) for Activation Decomposition:

2. Improved Conceptual Disentanglement via Region Assignment and Local Subspace Analysis:

3. Enhanced Causal Localization Capabilities through Centroid Interventions:

4. Superior Controllability and Steering via Local Variation Manipulation:

Specific Improvements and Capabilities of the Improved AI System:

1. Implementation of MFA for Activation Decomposition (Regions + Local Subspaces):

2. The system will decompose any activation into two compositional geometric objects: a region (its centroid in activation space) and a within-region offset (the local variation parameterized by a low-dimensional subspace).

3. Improved Conceptual Disentanglement via Region Assignment and Local Subspace Analysis:

4. The AI will no longer rely on isolated global directions to represent concepts. Instead, it will discover clusters of nearby Gaussian regions, where neighboring Gaussians collectively form a unified semantic neighborhood (e.g., grouping diverse movie genres under one broad Gaussian).

5. Enhanced Causal Localization Capabilities through Centroid Interventions:

6. The system can perform precise causal localization by manipulating the region's centroid (absolute position in activation space). Intervening on the centroid reliably steers the model toward a broad, high-level semantic theme associated with that entire region (e.g., steering toward Sports by moving to the Sports Gaussian centroid).

7. Superior Controllability and Steering via Local Variation Manipulation:

8. The system gains fine-grained control over sub-concepts within a broad theme by intervening on the local subspace directions (the loadings, represented as within-region offsets). This allows for targeted shifts to specific subthemes or syntactic variations even when operating within a broad conceptual region (e.g., steering toward Basketball from the general Sports Gaussian centroid).

9. Improved Interpretability and Feature Relevance:

10. The system provides a quantifiable measure of feature interpretability (IF score, averaging 0.96 in the paper), ensuring that the features it isolates are coherent and meaningful to human annotators, rather than being artifacts of the training objective.

In summary, this research enables an AI system to move from simply identifying what a concept is represented by (a single direction) to understanding where a concept lives in the activation space (the region/centroid) and how it manifests internally (the local geometry/subspace), leading to systems that are more robust, controllable, and transparent for complex language understanding tasks.

Activation decomposition methods in language models are tightly coupled to geometric assumptions on how concepts are realized in activation space. Existing approaches search for individual global directions, implicitly assuming linear separability, which overlooks concepts with nonlinear or multi-dimensional structure. In this work, we leverage Mixture of Factor Analyzers (MFA) as a scalable, unsupervised alternative that models the activation space as a collection of Gaussian regions with their local covariance structure. MFA decomposes activations into two compositional geometric objects: the region's centroid in activation space, and the local variation from the centroid. We train large-scale MFAs for Llama-3.1-8B and Gemma-2-2B, and show they capture complex, nonlinear structures in activation space. Moreover, evaluations on localization and steering benchmarks show that MFA outperforms unsupervised baselines, is competitive with supervised localization methods, and often achieves stronger steering performance than sparse autoencoders. Together, our findings position local geometry, expressed through subspaces, as a promising unit of analysis for scalable concept discovery and model control, accounting for complex structures that isolated directions fail to capture.