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

TL;DR

The paper introduces Mixtures of Factor Analyzers (MFA) to model language-model activation spaces as a collection of region-centered, low-dimensional subspaces, addressing limitations of global-direction approaches. By decomposing activations into region centroids and within-region variation, MFA captures nonlinear, multi-dimensional structures and enables both localization and causal steering. Large-scale experiments on Llama-3.1-8B and Gemma-2-2B show MFA outperforms unsupervised baselines in localization, competes with supervised methods, and provides stronger, more interpretable steering than sparse autoencoders. The authors advocate a local-geometry view of activation space as a scalable unit for concept discovery and control, and release code and trained MFAs to support further research.

Abstract

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.

Figures (11)

• Figure 1: MFA decomposes each activation into a region assignment and a within-region offset. Left: the region structure is modeled by Gaussian components (centroids $\boldsymbol\mu_k$), with complex concepts typically spanning multiple Gaussians -- here, the broader Emotions neighborhood is spanned by several interpretable Gaussians. Right: each component is equipped with a low-dimensional subspace that parameterizes structured within-region variation.
• Figure 2: Example MFA Gaussians in the activation space of Llama-3.1-8B, visualized in 3D using three loadings as axes. (Left) A broad region spanning multiple movie genres, where the loadings separate genre-related themes. (Right) A narrow region centered on the token National, where the loadings capture context-dependent usage.
• Figure 3: Characterizing MFA regions. (a) Broad vs. narrow regions differ across model families (Gemma skews narrow/token-driven; Llama stays mostly broad/semantic). (b) Semantic vs. syntactic loadings, split by broad/narrow components, become more semantic as $K$ increases, indicating more context-dependent within-region variation.
• Figure 4: MFA vs. SAE reconstructions. MFA reconstructs an activation by anchoring it to a region (centroid) and refining it with a region-specific direction, whereas SAEs reconstruct by accumulating many global dictionary features. Left: Llama-3.1-8B, layer 22; right: Gemma-2-2B, layer 18.
• Figure 5: Steering results across layers in Gemma-2-2B and Llama-3.1-8B of state-of-the-art SAEs, DiffMeans and 1K, 8K, 32K Gaussian MFAs. Across the majority of settings MFA significantly outperforms DiffMeans and SAEs.