Towards Universality: Studying Mechanistic Similarity Across Language Model Architectures

Authors: Junxuan Wang, Xuyang Ge, Wentao Shu, Qiong Tang, Yunhua Zhou, Zhengfu He, Xipeng Qiu
Primary area: interpretability and explainable AI
Keywords: Mechanistic Interpretability, Sparse Autoencoders, Universality, State Space Models

Abstract:
The hypothesis of \textit{Universality} in interpretability suggests that different neural networks may converge to
implement similar algorithms on similar tasks. In this work, we investigate two mainstream architectures
for language modeling, namely Transformers and Mambas, to explore the extent of their mechanistic similarity.
We propose to use Sparse Autoencoders (SAEs) to isolate interpretable features from these models and show
that most features are similar in these two models. We also validate the correlation between feature similarity
and~\univ. We then delve into the circuit-level analysis of Mamba models
and find that the induction circuits in Mamba are structurally analogous to those in Transformers. We also identify a nuanced difference we call \emph{Off-by-One motif}: The information of one token is written into the 
SSM state in its next position. Whilst interaction between tokens in Transformers does not exhibit such trend.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730215507041

thread_type:
official_review

<thread_object>
creator: reviewer_jsst

<content>
summary:
This work investigates the similarity of sparse autoencoders (SAE) features and induction circuits between a Transformer-based and a comparable Mamba-based language model trained on the same dataset. Results show that many features from Transformers SAEs show a high max pairwise Pearson correlation with Mamba SAE features, with their depth along model layers roughly matching across the two models. The correlations found for cross-architecture comparison are compared to a neuron baseline and two skylines obtained from different models and SAE training seeds, showing that the cross-architecture comparison falls only slightly short of skylines. Authors further examine the correlation between cross-architectural matching features and their complexity, finding that features with the most overlap are generally simpler and more monosemantic. Finally, the authors briefly investigate the similarity between induction circuits in the same architectures using path patching, finding a similar mechanism mediated by the convolutional operation of the Mamba architecture. Notably, the information mixing necessary for the induction operation is performed earlier in Mamba ("off-by-one" mixing).

strengths:
The universality evaluation pursued in this paper is timely and relevant, given recent advances in non-Transformer architectures for language modeling. The baseline and skylines employed in this work are well-motivated and provide helpful reference points for the analysis. The analysis of feature correlation based on complexity is also interesting, showing convincing proof that most commonalities across architectures are found for simpler SAE features. Overall, the figures are designed clearly and compellingly to support the findings detailed in the main body of the paper.

weaknesses:
**Novelty of the findings** While to my knowledge, this is the first work evaluating the similarity of SAE features between Transformers and Mamba models, other works such as [Paulo et al. 2024](https://arxiv.org/pdf/2404.05971) and [Sharma et al. 2024](https://arxiv.org/abs/2404.03646) already showed that many interpretability approaches such as logit lenses, steering vectors, and probes produce similar results across a variety of LLM architectures (Mamba, Transformer and RWKV). It is, hence, not particularly surprising that such findings extend to SAE decompositions of model activations.

**Generality of findings for larger models**  Authors experiment only with 2 tiny models with ~100M parameters each. This can be reasonable in light of the requirements (same training data, effort of training SAEs on each layer), but these systems are significantly smaller than those used by [Paulo et al. 2024](https://arxiv.org/pdf/2404.05971) for comparable cross-architecture interpretability experiments. Notably, larger checkpoints for both models used by the authors are publicly available, including the same training data for control, and could have been used to further prove the generality of the reported results. Importantly, without further experiments, it cannot be excluded that the limited capacity of tiny models might be the main motivation behind the high similarity features and circuits across the two architectures, and this could not be the case for more capable models with e.g. 1B or 7B parameters.

**Multiple Comparison and Correlation Analysis without Hypothesis Testing** The maximal correlation of feature activation patterns with other (24576 x # of layers) features is bound to be quite high due to the enormous amounts of comparisons. In Section 4.4, no hypothesis is formulated regarding the expected similarity of features found across the four tested variants, and consequently, no significance measure for the correlation coefficients is reported. As a result, conclusions regarding the similarity of Mamba and Pythia SAE features are ambiguous (e.g. the statement "[...] our neuron baseline almost exhibits zero sign of similarity between Mamba and Pythia" at line 268 does not agree with Figure 3a, where at least 15% of neurons exhibit a correlation > 0.4). To make the analysis more convincing, a clear hypothesis regarding the degree of similarity in resulting SAE features should have been formulated and tested for baseline, experiment, and skylines, each including a correction procedure such as the Bonferroni method to account for multiple comparisons.

**Minor formatting/clarification points:**

Line 135: The mention "$F_a$ and $F_b$ are some kinds of operation function." is too generic in this context. The purpose of these functions should be specified, and at least one example of functions used for this purpose should be provided.

Line 179: Broken reference.

Figure 1 is too tight with the inline text, making the distinction between caption and main body text unclear.

Section 4.1: title is ungrammatical. I imagine you meant something like "Searching for / In search of Interpretable Primitives".

Line 202: Clarify that you mean all features for all SAEs across all model layers (it becomes clear only from Figure 3c later in the paper)

Line 263: The acronym MPPC is never introduced alongside its meaning.

Figure 5: The mention "both Model Seed Variant and Cross-Arch SAE MPPC exhibit correlation, while one in SAE Seed Variant is weaker" in the caption is not very meaningful, since the trends for all three variants are pretty similar. For part (b), the mention "scores ranging from 1 (No) to 2 (Yes)" is confusing: it would be better to say "Distribution of MPCC for polysemantic (1) and monosemantic (2) auto-generated feature labels."

questions:
What justified the choice of plain SAEs over more performant variants such as [Gated](https://arxiv.org/abs/2404.16014) or [Top-K SAEs](https://openai.com/index/extracting-concepts-from-gpt-4/)? It is currently hard to gauge the impact of this choice on the obtained results, and whether findings could have been different if improved SAE variants were tested.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711714267

thread_type:
official_review

<thread_object>
creator: reviewer_eqto

<content>
summary:
This paper presents an exploration of the Transformer and Mamba models through mechanistic interpretability methodology. Despite the architectures being very different, the features and circuits of these models turn out to be very similar.

strengths:
- The paper is clearly written. 
- I appreciate the idea of using skylines, as it helps support the authors' claims. 
- The results are interesting and useful for further research.

weaknesses:
I couldn't identify any specific weaknesses. However, below are some suggestions that could enhance this work from my perspective:


- It would be interesting to explore more circuits through SAE, as suggested in [1] (see Section 4.3). However, it is unclear where SAE should be placed within the Mamba architecture to achieve similar features.
- While the Pearson correlation appears to be a natural choice for measuring feature similarity, it assumes that the feature space has linear properties. It might be worthwhile to explore other correlation measures, such as distance correlation, which could potentially yield better results.
- A clear statement clarifying that MPPC refers to the maximum Pearson correlation between models' features is needed to improve understanding.

[1] Interpreting Attention Layer Outputs with Sparse Autoencoders (Kissane et al.)

questions:
- While the heatmap matrix in Figure 3c is mainly diagonal, I can see that there is a cluster of features located in the last layer of the Pythia and distributed fairly uniformly in the middle layers of Mamba. Can the authors clarify the meanings of these features?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730730403144

thread_type:
official_review

<thread_object>
creator: reviewer_nup5

<content>
summary:
This paper studies the mechanistic similarity between language model structures (Mamba, RWKV, Transformer). The authors focus on their Universality, a hypothesized property that suggests different neural architectures implementing similar algorithms on similar tasks.

In the first part of the experiment section, they use Sparse Autoencoder (SAE) as the major tool for their analysis. The representations from two LM architectures are taken to train SAEs. The latent vectors in the SAEs, in which each dimension corresponds to various syntactic and semantic phenomena, exhibit mechanical similarities, and it is possible to find a matching between the vectors from different LM architectures.

In the second part, the authors study the induction behavior of two LM architectures. They found that the 17th layer of Mamba is the most important for the LM’s inductive ability.

I think this paper studies an important problem and is well executed. I found the experiments in this paper to be well implemented. My only concern is with the role of circuit analysis experiments. They are indeed very interesting but I’m not sure how they contribute to building a mechanistic analogy between SSMs and Transformers. Do Transformers have such layer-specific behavior when it comes to inductive ability? Is there a way to empirically verify the claims in Sec 6.2?

Minor:

Appendix D: How does the feature mapping between RWKV and Mamba look like?
Sec 6.1: Is the layer-17 phenomenon robust to random initializations? I.e., if one retrains the SSM with another seed, would layer 17 still be the key in induction?
Line 179: missing section reference.
Line 861: missing space between ‘Universality’ and ‘is’

strengths:
* This paper studies an important problem.
* It makes good use of sparse autoencoders for analysis.
* The experiments in this paper are well implemented.

weaknesses:
* The role of circuit analysis experiments is unclear.
* The claims made in Section 6.2 need to be empirically supported.
* (Minor) The size of LMs is limited. Only ~100m models are used in the experiments.

questions:
* Do Transformers have such layer-specific behavior when it comes to inductive ability?
* Is there a way to empirically verify the claims in Sec 6.2?
Appendix D: How does the feature mapping between RWKV and Mamba look like?
* Sec 6.1: Is the layer-17 phenomenon robust to random initializations? I.e., if one retrains the SSM with another seed, would layer 17 still be the key in induction?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730991374772

thread_type:
official_review

<thread_object>
creator: reviewer_4gec

<content>
summary:
The paper investigates the "universality hypothesis" in mechanistic interpretability, which suggests that different neural network architectures may converge to implement similar algorithms when tasked with analogous objectives. The authors focus on two mainstream architectures for language modeling: Transformers and Mambas. They propose using Sparse Autoencoders (SAEs) to extract interpretable features from these models and demonstrate that a significant portion of features are shared between the two architectures. The paper validates the correlation between feature similarity and universality and delves into the circuit-level analysis of Mamba models, finding structural analogies with Transformers, particularly in induction circuits. 

The paper's contributions include:
- Introduction of a novel metric to isolate and quantify feature universality in the context of architectural variations.
- Empirical evidence shows that Transformer and Mamba models learn similar features through the application of SAEs.
- Circuit analysis of Mamba models reveals structural analogies and nuanced differences compared to Transformer circuits.
- Support for the universality hypothesis by demonstrating cross-architecture feature similarity and identifying the "Off-by-One motif" in Mamba models.

strengths:
- The identification of the "Off-by-One motif" in Mamba models is a unique contribution that highlights nuanced differences between architectures.

- The introduction of a complexity-based interpretation for understanding feature similarity differences is innovative.

- The circuit-level analysis of Mamba models, revealing structural analogies with Transformers, is good and adds depth to the study. The validation of feature similarity and its correlation with universality further strengthens the study.


- The findings of this paper have implications for the field of neural network interpretability. By demonstrating that different architectures can converge to similar algorithms and features, the study provides valuable insights into the generalizability of mechanistic findings across models.

weaknesses:
- While the paper focuses on Transformers and Mambas, it would benefit from a broader examination of additional architectures. Including a more diverse set of models, such as recurrent neural networks (RNNs) or convolutional neural networks (CNNs), could strengthen the universality hypothesis by offering a more comprehensive understanding of feature similarity across a wider range of neural networks. This would enhance the generalizability of the findings.
	
- The paper utilizes OpenWebText for correlation analysis but does not discuss how the choice of dataset might affect the results. A more detailed examination of the potential biases and limitations introduced by the dataset choice would provide a clearer context for the findings and ensure that the results are not overly dependent on a specific dataset.

- The claims of feature similarity and universality would be more robust if supported by statistical significance tests. Including such tests would provide stronger evidence for the observed correlations and enhance the credibility of the conclusions.

- SAE-related technical gaps:
    - The paper does not include ablation studies on SAE hyperparameters (dictionary/code size, training duration, etc.). Conducting these studies would help to understand the sensitivity of the results to different hyperparameter settings and ensure the robustness of the findings.
    - There is no discussion of how SAE reconstruction quality relates to feature similarity. Addressing this relationship would provide insights into the effectiveness of SAEs in isolating interpretable features and validate the methodology used.

- The use of GPT-4 for complexity scoring lacks rigorous validation. The paper does not provide inter-rater reliability metrics or comparisons with human annotations, nor does it discuss potential biases in the automated scoring. 

- The paper provides a limited exploration of why the "Off-by-One" motif exists in Mamba models. A deeper investigation into the underlying reasons for this motif would enhance the understanding of the structural differences between the Mamba and Transformer models and provide more insights into the universality hypothesis.

questions:
1. Why did you choose OpenWebText as your primary dataset for analysis? How might the choice of OpenWebText as the dataset influence your results? Have you tested if the feature similarities hold across different domains (e.g., code, mathematics, or structured data)? Would analyzing domain-specific text reveal different patterns of architectural universality?

2. Have you performed any statistical significance tests to support your claims of feature similarity and universality?

3. How generalizable are your findings to other tasks beyond language modeling?

4. Can you provide more details on why the "Off-by-One" motif exists in Mamba models?

5. Is there a risk that the Sparse Autoencoder pre-processing itself may impose a degree of alignment between features in Transformers and Mambas? Could the sparsity constraint inadvertently enhance apparent similarity?

6. How do you expect your findings to scale to larger models? Did you observe whether model size impacts universality between architectures? Could smaller or larger versions of Transformers and Mambas exhibit different degrees of feature similarity?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733297676244

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all reviewers for their thoughtful feedback and constructive suggestions, which have greatly helped us refine this work. Below, we summarize and address the key points raised:

- We are grateful for the recognition from all four reviewers for the **significance** (reviewers 4GEc, nuP5, eqTo, and jsst), **soundness** (reviewers 4GEc, nuP5, eqTo, and jsst), and **novelty** (reviewers 4GEc, nuP5, and eqTo) of our work.
- Specifically, we appreciate the recognition of our **complexity-based perspective on feature similarity** as both innovative (reviewer 4GEc) and interesting (reviewer jsst).
- The **skylines introduced for comparison** were also highlighted as well-motivated and supportive (reviewers jsst and eqTo).
- Additionally, our **induction circuit universality analysis** was acknowledged as adding depth to the study (reviewer 4GEc).

We are particularly thankful for the reviewers’ insightful suggestions, which have significantly enhanced this work. Below, we address four common concerns:

1. **Generalization to larger models** (reviewers 4GEc, nuP5, and jsst):  
   Our initial experiments were conducted on ~100M parameter models, and we acknowledge concerns about scaling to modern, larger models. To address this, we conducted our main experiments along with corresponding baselines and skylines on 2.8B versions of both models. These results align with the trends reported in the original manuscript, demonstrating the robustness of our findings at larger scales.

2. **Statistical significance for robust comparison** (reviewers 4GEc and jsst):  
   To ensure the reliability of our results, we conducted statistical significance tests, confirming a probability of less than $5 \times 10^{-38}$ for random noise influence. Additionally, cross-architecture comparisons on complex features revealed greater similarities than cross-layer comparisons, highlighting the robustness and meaningfulness of our analysis.

3. **Ablation studies on SAE training settings** (reviewers 4GEc and jsst):  
   Reviewers expressed concerns about the robustness of our findings to variations in SAE training hyperparameters and architectures. We expanded our ablation studies to include sparsity coefficients, SAE size, and TopK SAEs. These additional experiments confirm the robustness of our conclusions across a broader range of training settings.

4. **Further discussions** (reviewers 4GEc, eqTo, and jsst):  
   During the discussion period, several inspiring ideas emerged that we plan to incorporate into future versions of this work. These include:  
   - **Implications of complexity-based MPPC analysis** for broader model comparison methods (from discussions with reviewer jsst).  
   - **Further details on the Mamba Off-by-One motif** (suggested by reviewer 4GEc).  
   - **Highlighting an interesting exception** in the final-layer representation observed in our depth-specialization analysis (suggested by reviewer eqTo).  

These insights add significant depth to our work, and we thank the reviewers for their valuable contributions. We also deeply value additional feedback not explicitly mentioned here, which we address in the individual discussion sections. We are committed to incorporating these suggestions to improve both the depth and clarity of our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper investigates the universality hypothesis, showing that Transformers and Mambas, despite architectural differences, learn similar features for language modeling tasks. Using Sparse Autoencoders (SAEs), the authors demonstrate cross-architecture feature similarity and identify a novel “Off-by-One motif” in Mamba models, providing new insights into their induction circuits. The study combines feature-level analysis and circuit-level comparisons to support claims of mechanistic universality.

The strengths of the paper lie in its novel application of SAEs to compare architectures, robust empirical results validated by statistical significance tests, and the depth of analysis that identifies meaningful similarities and differences. The primary weaknesses include limited experiments on larger models and questions about the robustness of the results. Additionally, reviewers raised concerns about potential biases introduced by SAE pre-processing and the generalization of findings beyond language modeling. Most of these issues were addressed during the rebuttal period: the authors provided results for larger models (2.8B parameters), performed statistical tests, and demonstrated that the SAE method did not artificially impose alignment. While broader generalization and deeper circuit validation remain future directions, the authors' responses sufficiently strengthen their claims.

Overall the paper makes a significant and timely contribution to neural network interpretability. Its findings are well-supported, and the concerns raised have been addressed to a reasonable extent. I recommend accepting this paper.