Mechanistic evaluation of Transformers and state space models

TL;DR

This work investigates how Transformers and state space models (SSMs) implement in-context retrieval on synthetic tasks, arguing that behavioural success alone obscures underlying mechanisms. By applying interchange interventions, it reveals that Transformers and the Based model solve associative recall (AR) via induction, while many SSMs rely on direct retrieval; Mamba uses short convolutions to enable induction-like associations. The authors introduce Associative Treecall (ATR) to test generality and show similar mechanistic patterns across architectures, with Mamba able to learn induction even without short convolutions on ATR. The study demonstrates the value of mechanistic evaluations for understanding architectural differences and guiding future model design with interpretability in mind.

Abstract

State space models (SSMs) for language modelling promise an efficient and performant alternative to quadratic-attention Transformers, yet show variable performance on recalling basic information from the context. While performance on synthetic tasks like Associative Recall (AR) can point to this deficiency, behavioural metrics provide little information as to \textit{why} -- on a mechanistic level -- certain architectures fail and others succeed. To address this, we conduct experiments on AR, and find that only Transformers and Based SSM models fully succeed at AR, with Mamba and DeltaNet close behind, while the other SSMs (H3, Hyena) fail. We then use causal interventions to explain why. We find that Transformers and Based learn to store key-value associations in-context using induction. By contrast, the SSMs seem to compute these associations only at the last state using a single layer. We further investigate the mechanism underlying the success of Mamba, and find novel evidence that Mamba \textit{does} implement induction: not via the SSM, but instead via short convolutions. Further experiments on a new hierarchical retrieval task, Associative Treecall (ATR), show that all architectures learn the same mechanism as they did for AR. Furthermore, we show that Mamba can learn Attention-like induction on ATR when short convolutions are removed. These results reveal that architectures with similar accuracy may still have substantive differences, motivating the adoption of mechanistic evaluations.

Figures (13)

• Figure 1: Associative mechanisms in Attention, Based, and Mamba: Key results from applying mechanistic metrics to understand how Attention and SSMs solve AR and our new retrieval task, Associative Treecall (ATR). Attention and Based both implement induction, but with different architectural components (\ref{['sec:ar-exp', 'sec:atr-exp']}). Mamba instead uses a single layer for both association and retrieval, but uses short convolution for association just like Based (\ref{['sec:ar-short-conv', 'sec:ar-mamba-induction']}). When convolutions are removed, Mamba implements induction, but only on ATR (\ref{['sec:atr-mamba']}).
• Figure 2: An example interhcange intervention where we corrupt the key (A) and attempt to restore it by intervening at the value token in an intermediate representation. We evaluate the downstream effect on the next-token prediction at the query (which should ideally predict the answer).
• Figure 3: Associative recall: Likelihood of correct answer without any interventions, after corrupting the key, and after restoring representations at the layer 1 block input with interchange intervention, on AR with vocabulary size $8192$ and key--value count of $32$. SSMs (except for Based) and Transformers learn different mechanisms.
• Figure 4: Ablating short convolution in Mamba and Based: Accuracy and interchange intervention results when ablating parameters of the short convolution component in Mamba and Based. Conv. size less than $2$, using long conv., or no conv., all lead to near-zero performance in both models.
• Figure 5: Interventions on short convolutions in Mamba: Likelihood of the correct answer for all Mamba checkpoints (varying model dimension and LR) on AR; (from left to right) original performance, after corrupting the key, restoring at the layer 0 short conv output at the value, same but at the next key, and same but at the layer 1 value token. (Some runs failed due to high LR.)

Theorems & Definitions (4)

• Definition B.1
• Definition B.2
• Definition B.3
• Definition B.4