Sparse Feature Coactivation Reveals Causal Semantic Modules in Large Language Models

1. Problem Statement

Large Language Models (LLMs) demonstrate remarkable proficiency in relational reasoning (e.g., knowing the capital of a country or translating a word), yet the internal mechanisms enabling this remain opaque. Previous mechanistic interpretability efforts have faced two main limitations:

• Layer-wise Abstraction: Prior work often models capabilities as simple linear transformations or vector arithmetic at the layer level, lacking fine-grained insight into specific features and their interactions.
• Complexity of Circuit Tracing: Methods like cross-layer transcoders can trace information flow but often produce dense, complex computation graphs with hundreds of nodes that are difficult to interpret without manual grouping.

The authors aim to identify semantically coherent, context-consistent network components (modules) within LLMs that causally drive specific concepts (e.g., "China") and relations (e.g., "capital city") and demonstrate that manipulating these modules allows for precise, targeted control over model outputs.

2. Methodology

The authors propose a pipeline to extract and validate causal semantic modules using Sparse Autoencoders (SAEs) and feature coactivation. The process involves the following steps:

A. Data Collection & Feature Selection

• Model: Experiments primarily use Gemma 2 2B (with some validation on Gemma 2 9B).
• SAE Integration: Pre-trained SAEs (width 16k) are integrated into the model. For each prompt, the residual stream activations at every layer $\ell$ are mapped to sparse codes $\phi_\ell$.
• Top-k Filtering: To ensure tractability, only the top $k=5$ activated features per token are selected for each layer.

B. Graph Construction

• Nodes: Each node represents a specific feature at a specific layer $(\ell, i)$.
• Edges: Directed edges are drawn between features in adjacent layers if their activation patterns across the prompt tokens show a high temporal correlation (Pearson correlation $\rho > 0.9$). This creates an inter-layer feature network.

C. Pruning & Component Identification

• Density Pruning: To remove generic, polysemantic features, the graph is pruned using activation density scores from Neuronpedia. Features with activation density $> 0.01$ (i.e., too frequent) are removed.
• Connected Components: The remaining sparse graph is analyzed to find weakly connected components. These components represent groups of features that co-activate consistently for a specific task.
• Semantic Validation: Components are validated by checking if they yield predictable counterfactual outputs when manipulated. The authors manually select or combine components that maximize compositional changes in model outputs (e.g., ablating a "China" component shifts the answer to another country's capital).

D. Causal Intervention (Steering)

• Ablation: Activations of features within a identified component are set to zero during the forward pass.
• Amplification: Activations of features in a target component are boosted by a steering strength $\alpha$.
• Evaluation: The impact is measured via KL divergence between the original and perturbed token distributions and by checking if the model generates counterfactual answers (e.g., answering "Abuja" when prompted about China's capital).

3. Key Contributions

1. Discovery of Modular Semantic Units: The paper identifies that LLMs organize knowledge into distinct, sparse components representing concepts (entities/words) and relations (properties/transformations). These components are consistent across different contexts.
2. Causal Control via Composition: The authors demonstrate that these modules are composable. By ablating an in-prompt concept/relation and amplifying a target concept/relation, they can steer the model to generate compound counterfactual outputs (e.g., changing "China's capital" to "Nigeria's currency").
3. Hierarchical Organization: Analysis reveals a distinct layer-wise distribution:
   • Concepts: Concrete entities (e.g., "China," "love") emerge primarily in the first few layers.
   • Relations: Abstract relations (e.g., "capital city," "translation language") are concentrated in deeper layers.
4. Superiority over Single Features: The study shows that steering based on a group of co-activating features (a component) is significantly more effective and robust than steering based on a single "most causal" feature.

4. Experimental Results

The authors evaluated their method on three multi-relational prediction tasks:

• Country Facts: (Capital, Currency, Language) for 10 countries.
• Word Translation: Translating 11 English words into Spanish, French, and German.
• Verb Transformation: Synonym, antonym, past tense, capitalization, and first letter for 8 verbs.

Key Findings:

• High Steering Success:
  • Concept Steering: 48–96% success rate (e.g., changing the subject country while keeping the relation).
  • Relation Steering: 23–98% success rate (e.g., changing the relation while keeping the subject).
  • Composite Steering: 19–90% success rate (changing both subject and relation simultaneously).
• Specificity: Ablating components for one task (e.g., Country Facts) generally preserved performance on other tasks (e.g., Word Translation), indicating high specificity of the extracted modules.
• Failure Modes:
  • Verb Transformations: This task showed lower success rates (especially for synonyms/antonyms), likely because these relations are less objective and more context-dependent than facts like capitals or currencies.
  • High-Density Features: Preliminary experiments suggested that ablating high-density features leads to incoherent outputs, implying they serve critical, unexplored computational functions.

5. Significance and Implications

• Mechanistic Interpretability: This work advances the field by moving beyond "black box" layer analysis to identifying specific, interpretable, and manipulable functional modules. It provides a lightweight alternative to exhaustive circuit tracing.
• Model Control: The ability to reliably steer LLMs toward counterfactual responses by manipulating specific internal components offers a new paradigm for targeted model editing and safety alignment without retraining.
• Knowledge Representation: The findings support the hypothesis that LLMs store knowledge in a hierarchical, modular fashion, where concrete facts are established early in the network and abstract relational logic is applied in later stages.
• Scalability: The method relies on pre-trained SAEs and automated graph pruning, making it a scalable approach for analyzing and controlling future, larger models.

In conclusion, the paper establishes sparse feature coactivation as a powerful, efficient tool for uncovering the causal architecture of LLM reasoning, enabling precise manipulation of model behavior through the identification of semantic modules.