PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Or Shafran, Shaked Ronen, Omri Fahn, Shauli Ravfogel, Atticus Geiger, Mor Geva

arXiv:2602.02464v1cs.CL

TL;DR

Existing activation-decomposition methods often rely on isolated global directions, which are poorly suited to nonlinear or multi-dimensional concepts. The paper proposes scalable unsupervised Mixtures of Factor Analyzers that represent activations through local Gaussian regions and within-region variation, and reports stronger or competitive benchmark performance than several baselines. It concludes that local subspace structure is a promising unit for concept discovery and model control.

  • Problem

    Global-direction methods overlook nonlinear or multi-dimensional concepts that are dispersed across multiple directions.

  • Method

    MFA partitions activation space into Gaussian regions and models low-rank within-region variation, decomposing each activation into a region and local offset.

  • Results

    MFA outperforms unsupervised baselines, is competitive with supervised localization methods, and often steers better than SAE features.

  • Takeaways & Limitations

    Local subspace structure is a promising unit of analysis for scalable concept discovery and model control.

  • Takeaways & Limitations

    The method is potentially dual-use because isolating and manipulating internal mechanisms could evade safety measures or amplify undesirable behaviors.

Abstract

from arXiv · show

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.

1. Introduction

Global-direction methods assume concepts are linearly organized, but nonlinear and multi-dimensional concepts may span related regions. This work uses MFA to model local regions and within-region variation, yielding interpretable decompositions and strong localization and steering results.

  • Motivation: Global-direction decompositions impose geometric assumptions that overlook nonlinear or multi-dimensional concepts dispersed across multiple directions.Recovering such concepts requires post hoc assumptions about which directions form one representation.
  • Method: MFA partitions activation space into regions and learns a low-rank subspace capturing dominant within-region variation.Each activation is decomposed into a region and a within-region offset.
  • Findings: Concepts may be realized by constellations of neighboring Gaussians that jointly cover complex semantic regions.Broad Gaussians can span thematic topics, while narrow Gaussians can concentrate on constrained lexical patterns; larger component counts make Gaussians narrower.
  • Comparison: 75% of active SAE features were not directly interpretable from context on average.The comparison contrasts MFA’s region-and-local-variation decomposition with SAEs’ single dictionary of global directions.
  • Evaluation: MFA outperformed Desiderata-Based Masking on 5 out of 8 localization tasks and often achieved about a twofold steering gain in coherence and conceptual alignment.It was also often competitive with DAS, while MFA centroids steered better than SAE features in most settings.
  • Conclusion: The paper positions local subspace structure as a scalable unit for understanding how language models organize information.The approach models low-dimensional regions and scales to thousands of subspaces.

2. Preliminaries and Notation

Factor Analysis models observed correlations through a small number of latent factors and explicitly represents independent noise. MFA extends this model by giving different regions their own means and local low-dimensional geometries.

  • Factor Analysis: Factor Analysis explains correlations among observed variables with a small number of latent factors while modeling remaining variation as independent noise.It is a generative probabilistic model with a low-rank approximation.
  • Factor Analysis: The FA covariance combines shared variation from W with independent noise from Ψ as C = WW^⊤ + Ψ.The noise covariance Ψ is diagonal.
  • Factor Analysis: FA identifies the subspace span(W), whereas individual factor axes depend on the chosen rotation convention.Orthogonal rotations preserve the covariance.
  • Mixtures of Factor Analyzers: MFA extends FA by modeling representation space as local low-dimensional Factor Analyzers rather than one model of global variation.This supports factors of variation that differ across regions.
  • Mixtures of Factor Analyzers: Each MFA component has a mean µk defining its region and a component-specific matrix Wk defining its local subspace.A discrete latent variable ω identifies which component generated a sample.
  • Mixtures of Factor Analyzers: The overall MFA density weights Gaussian components according to πk and how well each component’s mean and subspace explain the sample.The mixture combines the component-specific distributions.

3. Mapping the Activation Space with Mixtures of Factor Analyzers

MFA maps activation space into coherent regions, each with its own low-rank local geometry, and decomposes activations into region-level and within-region components. The approach supports scalable analysis and steering of local structure in language-model activations.

  • MFA decomposition: MFA partitions activation space into regions while modeling each region’s variation with a low-rank subspace.Each component is a Gaussian whose covariance is parameterized by learned low-rank directions.
  • Component assignment: Component responsibilities assign activations to the local subspaces that best explain them.Responsibilities are normalized component likelihoods, allowing activations to be expressed as mixtures of components.
  • MFA decomposition: Each activation combines responsibilities-weighted component centroids with responsibilities-weighted within-region latent coordinates.The resulting reconstruction is a single matrix multiplication using component means and loading axes.
  • Global versus local decomposition: The meaningful geometric object is each local subspace span(Wk), not an individual loading vector, because factorization within a component is rotationally invariant.This motivates analyzing local regions with their own low-rank geometry rather than isolated global directions.
  • Evaluation: The study trains MFAs on language-model residual-stream activations and evaluates their regions, local variation, reconstruction, localization, and steering.The experiments include Llama-3.1-8B and Gemma-2-2B, while comparing MFA with SAE decomposition and unsupervised and supervised benchmarks.
  • Steering with MFA: MFA steering separately moves activations toward a centroid and applies additive offsets within the centroid-centered local subspace.The interpolation strength α controls movement toward the centroid, while latent coordinates control the offset direction and magnitude.

4. Activation Structures Discovered by MFA

Large-scale MFAs reveal diverse local regions and low-dimensional variation in Gemma-2-2B and Llama-3.1-8B activations. Their components form semantic neighborhoods, with nonlinear concepts often distributed across multiple nearby Gaussians.

  • Experimental setup: 12 MFAs trained on 100M activations expose semantically coherent nonlinear manifolds as collections of diverse, locally linear regions.The models are Gemma-2-2B and Llama-3.1-8B, with components varying in scale and model layer.
  • Discovered regions: MFA components range from narrow regions focused on specific tokens or contexts to broad regions spanning comprehensive themes.Their local subspaces capture both semantic and syntactic differences.
  • Model differences: Larger K increases the share of narrow Gaussians, especially in Gemma, while Llama remains predominantly broad even at K=32,000.This model-dependent pattern suggests different activation-space notions of similarity.
  • Within-region variation: Across settings, narrow Gaussians skew more syntactic whereas broad Gaussians skew more semantic.Figure 3 separates broad and narrow components from semantic and syntactic loadings.
  • Multi-Gaussian concepts: Neighboring Gaussians tend to encode related meanings, and globally nonlinear concepts often appear as clusters of nearby components.These component constellations jointly cover broader semantic regions rather than isolating a concept in one component.

5. MFA vs. Dictionary Learning

MFA decomposes activations into a region-level centroid and local variation, whereas SAEs assemble reconstructions from many global dictionary features. This local decomposition yields more interpretable components and distinct reconstruction geometry.

  • Centroid versus loading: Table 1 contrasts centroid interventions, which promote broad themes, with loading interventions, which separate subthemes or refine genres into subgenres.The comparison covers both narrow and broad regions.
  • Reconstruction geometry: MFA reconstructions follow two segments—centroid first, then local variation—whereas SAE reconstructions often require many incremental additions.The comparison examines residual-stream activations from Gemma-2-2B and Llama-3.1-8B using 8K MFAs and corresponding SAEs.
  • Decomposition: MFA represents an activation through a region centroid and within-region variation, while SAEs use many active global dictionary features.For MFA, the two reconstruction terms are the centroid µk and local term Wkẑk; SAE terms are ajfj.
  • Interpretability: MFA achieves an average interpretability fraction of 0.96 ± 0.2, compared with 0.29 ± 0.2 for SAEs.The interpretability fraction weights labeled-interpretable features by their contribution magnitude; I denotes the interpretable-feature set.
  • Interpretability: SAEs rely on many features, with an average of 75% of active features not directly interpretable from context, while MFA uses a small set of interpretable features.MFA local features are judged by whether coordinates place activations in coherent local clusters, rather than by interpreting individual loadings in isolation.

6. Evaluations

The evaluations test MFA for causal localization and steering against unsupervised, supervised, and SAE-based baselines. MFA generally improves localization over unsupervised methods, remains competitive with supervised methods, and outperforms steering baselines across most settings.

  • Evaluation design: MFA is evaluated on causal localization and steering benchmarks against PCA, SAEs, DBM, DAS, and DiffMeans.Localization uses RAVEL and MCQA settings, while steering compares MFA centroids with SAE features and supervised difference-in-means interventions.
  • 6.1. Localization: On RAVEL, MFA beats PCA and SAEs by 3–16 points, defeats DBM in 5 of 6 cases, and outperforms DAS on the Continent task.For Llama-3.1-8B, MFA comes within two points of DAS on the remaining RAVEL tasks.
  • 6.1. Localization: On MCQA, MFA outperforms SAEs by up to 15 points and, on Gemma, exceeds PCA while nearly matching DBM.The benchmark tests whether interventions change model answers for a positional pointer variable.
  • 6.1. Localization: Restricting MFA to centroids preserves RAVEL performance but reduces MCQA performance from 80% to 39%, indicating that some variables depend on within-region variation.Centroid-only RAVEL scores are Continent 86, Language 64, and Country 59.
  • 6.2. Causal Steering: MFA outperforms SAEs and DiffMeans across most steering settings, roughly doubling Gemma-2-2B’s median score and improving Llama-3.1-8B’s median by about one third.Layer 8 lowers the Llama average, and increasing MFA capacity produces no consistent gains.

7. Related Work

Related work characterizes language-model representations as structured, contextual, and locally low-dimensional rather than uniformly organized by a single global linear space. MFA builds on this geometry while addressing the practical challenge of using local subspaces for feature discovery and control.

  • Geometric Structures in LMs: Contextual language-model embeddings form usage-specific regions rather than a single global linear space.This extends earlier work on linear relations in language representations toward context-dependent geometry.
  • Geometric Structures in LMs: Modern language-model representations span many global directions but have low intrinsic dimension locally, motivating MFA’s region-and-subspace model.The paper presents MFA as a scalable feature-discovery method grounded in local geometry.
  • Feature Discovery in LMs: SAEs are the predominant unsupervised decomposition approach but assume activations are modeled by a global sparse dictionary of directions.This assumption conflicts with evidence that meaningful variation can be multidimensional and local.
  • Feature Discovery in LMs: Recent methods use concept-conditioned subspaces, partitioned representation subspaces, or task-relevant feature manifolds, but practical decomposition and steering remain open challenges.These approaches motivate subspaces as interpretation units without resolving the full practical problem addressed by MFA.

8. Conclusion

The paper presents MFA as a local-geometry approach to activation decomposition, modeling low-rank Gaussian regions and their local variation instead of isolated global directions. It reports stronger performance than unsupervised methods, competitiveness with supervised ones, and released models and code for further research.

  • MFA decomposes activation spaces into low-rank Gaussian regions and their local axes of variation.
  • The local-geometry view represents complex, nonlinear structures beyond what a single global set of directions can express.
  • MFA surpasses existing unsupervised methods on recent benchmarks and often exceeds supervised methods.
  • The authors release code and 12 trained MFAs for Gemma-2-2B and Llama-3.1-8B to support further community research.

Impact Statement

The work frames local geometry as a framework for decomposing and intervening on language-model activations. It also identifies dual-use risks because isolating and manipulating internal mechanisms could support safety evasion or undesirable behaviors.

  • The framework identifies regions and local low-rank structure rather than modeling activations as combinations of isolated global directions.
  • The decomposition supports interventions targeting either an activation-space region or specific within-region variations.
  • The same ability to isolate and manipulate internal mechanisms could potentially evade safety measures or amplify undesirable behaviors.

A. MFA Initialization and Training

This section compares MFA initialization strategies and describes annotation procedures for evaluating whether LLMs can label loadings and Gaussian regions. K-Means generally provides a strong convergence and diversity trade-off, while random alternatives have distinct weaknesses.

  • Initialization Comparison: MFA training compares K-Means, fully random, and random-point initialization while holding hyperparameters fixed.Convergence is defined by a successive log-likelihood change below 10^-3.
  • Initialization Comparison: Fully random initialization often converges to poor solutions, with little pairwise centroid-distance variation and relatively higher NLL.
  • Initialization Comparison: Random-point initialization scales to very large datasets and converges fast but is more prone to local minima and produces less uniform centroid spread than K-Means.
  • Annotation Evaluation: The annotation evaluation uses Cohen’s κ and an Alternative Annotator Test comparing LLM–human agreement with leave-one-out human agreement.The test uses winning rate ω, replacement threshold ω ≥0.5, q = 0.05, and ϵ = 0.1.
  • Annotation Results: Across 58 sampled loadings, human agreement averaged κ = 0.29, while LLM agreement averaged κ = 0.44 and met the replacement criterion in all three comparisons.For Gaussian labeling, humans and the LLM achieved pairwise or mean κ = 1.0, with ω = 1.0 meeting the replacement criterion.
  • Latent-Space Evaluation: For a single activation, the evaluation tests whether its latent coordinates place it in the correct concept group using 10 nearest and 10 farthest within-component neighbors.

D. Reconstruction Analysis

The reconstruction analysis evaluates MFA and SAE reconstruction error on held-out activations, finding that MFA error improves with more components but remains higher than SAE error. MFA’s main scope boundary is its dependence on the activation distribution observed during training.

  • The experiment measures mean squared reconstruction error on a held-out validation set using 10 million Wikipedia activations.MFA reconstructs each activation using Eq. 11, while the SAE encodes and decodes it.
  • The comparison includes 12 MFAs for Gemma-2-2B and Llama-3.1-8B, matched-layer Gemmascope-65k and Llamascope SAEs, and a nearest-centroid K-Means baseline.
  • Increasing MFA components from 1K to 8K substantially reduces reconstruction error, while gains from 8K to 32K are comparatively smaller.
  • SAEs achieve lower reconstruction error across all settings because they permit more flexible, sample-specific reconstruction.
  • Limitation: MFA models the activation distribution of its training corpus, so rare or out-of-distribution activations may be assigned to a poorly fitting nearest component and yield high reconstruction error.Increasing K can expand training-distribution coverage but does not guarantee low error on unseen regions.

E. Benchmarking: Additional Details

The steering evaluation compares MFA with SAEs and DiffMeans across model layers, using method-specific intervention schemes and standardized generation settings. Additional examples examine centroid and loading interventions, while annotations distinguish semantic/syntactic loadings and broad/narrow Gaussians.

  • Intervention ablations: 32K-MFA additive centroid intervention scores 0.124 ± 0.14, compared with 0.24 ± 0.20 for interpolation, supporting interpolation for MFA.The ablation was conducted on Gemma-2-2B layer 18 using 100 randomly sampled features or Gaussians excluded from the final results.
  • Intervention ablations: SAE interpolation yields 0.177 ± 0.19 versus 0.195 ± .21 for additive intervention, while additive interventions remain the common SAE setting.The paper therefore uses the best-performing intervention method from the ablations for each decomposition.
  • Causal steering evaluation: MFA significantly outperforms SAEs and DiffMeans on Concept Score, while Fluency Scores remain broadly consistent with a slightly lower MFA upper bound.The result indicates stronger concept promotion without a corresponding fluency advantage.
  • Causal steering evaluation: MFA, SAEs, and DiffMeans are evaluated across layers in Gemma-2-2B and Llama-3.1-8B using Concept and Fluency Scores.Each condition generates 8 continuations from “I think that” with at most 50 new tokens, top-k=30, and top-p=0.3.
  • Qualitative steering examples: MFA loading evidence is qualitative because loadings define a local variation subspace rather than a unique semantic direction.A single loading may therefore fail to align with the region’s dominant variation direction.
  • Qualitative steering examples: Centroid interventions promote broad themes, whereas loading interventions isolate sharper subthemes and can reveal variation associated with local covariance structure.Examples are reported for Llama-3.1-8B layer 22 and Gemma-2-2B layer 18; the paper presents these observations as qualitative evidence.
Loading 2602.02464v1…