Mechanistic evaluation of Transformers and state space models

TL;DR

This paper compares Transformers and state space models in language tasks, revealing that Transformers and Based SSMs succeed in associative recall by using induction, while others fail, highlighting the importance of mechanistic understanding.

Contribution

It provides a mechanistic analysis of different architectures' ability to perform associative recall, identifying how induction mechanisms contribute to success.

Findings

Transformers and Based SSMs fully succeed at associative recall.

Mamba implements induction via short convolutions, not SSMs.

Architectures with similar accuracy can have different underlying mechanisms.

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.