Revisiting associative recall in modern recurrent models

TL;DR

The paper investigates why modern recurrent state-space models (SSMs) struggle or succeed relative to Transformers on associative recall benchmarks. It shows that optimization dynamics, particularly learning-rate sensitivity, largely explain performance gaps, challenging the view that expressivity alone governs recall capabilities. Width scaling benefits SSMs while depth scaling favors Transformers, and 1-layer Transformers exhibit induction-head–like dynamics without achieving MQAR success unless architectural tweaks are used. The work highlights optimization stability as a critical objective for future sequence-model design and provides practical guidance on when and how to scale SSMs versus Transformers, with DeltaNet and related approaches offering routes to improved stability.

Abstract

Despite the advantageous subquadratic complexity of modern recurrent deep learning models -- such as state-space models (SSMs) -- recent studies have highlighted their potential shortcomings compared to transformers on reasoning and memorization tasks. In this paper, we dive deeper into one of such benchmarks: associative recall (AR), which has been shown to correlate well with language modeling performance, and inspect in detail the effects of scaling and optimization issues in recently proposed token mixing strategies. We first demonstrate that, unlike standard transformers, the choice of learning rate plays a critical role in the performance of modern recurrent models: an issue that can severely affect reported performance in previous works and suggests further research is needed to stabilize training. Next, we show that recurrent and attention-based models exhibit contrasting benefits when scaling in width as opposed to depth, with attention being notably unable to solve AR when limited to a single layer. We then further inspect 1-layer transformers, revealing that despite their poor performance, their training dynamics surprisingly resemble the formation of induction heads, a phenomenon previously observed only in their 2-layer counterparts. Finally, through architectural ablations, we study how components affects Transformer and Mamba's performance and optimization stability.

Figures (12)

• Figure 1: Performance on MQAR (mean and relative max-min errors using 3 seeds) after an extensive learning rate grid search. Unlike attention, the window of suitable learning rates for Mamba and Hyena is relatively narrow. We compare our grid search with the one used by arora2023zoology (dashed vertical lines) to highlight how the suitable learning rate can be missed.
• Figure 2: Performance of 2-layers models in MQAR. We report the official results ( green stars) and the replication running the original code of arora2023zoology ( dotted blue line). While for replication, we used the learning rate grid by arora2023zoology, we note here that, due to high sensitivity to the learning rate (Fig, \ref{['fig:lr']}), tuning drastically affects performance. In solid orange, we provide results with a finer grid (cf. Fig.\ref{['fig:lr']}). Careful tuning of the learning rate gives a general improvement in the performance of recurrent models. This becomes especially crucial in Mamba, where the task becomes solvable at high sequence lengths $>>$ hidden size. The results show the mean and relative max-min errors for 3 seeds. Attention always solves the task (all curves overlap).
• Figure 3: Performance of 1-layer models on MQAR. We show how for recurrent models, scaling the width boosts performance. On the contrary, Attention models can no longer solve the task anymore as they do in the 2-layer setting, and performances are unaffected by the scaling in width. The results show the mean and relative max-min errors after 3 runs with different seeds.
• Figure 4: Scaling behavior (Seq len: 256, KV pairs: 64). Symbols with the same shape and color represent models of increasing dimension from 64 to 2048. We show that the scaling strategy, rather than the total number of parameters, is what primarily impacts performance. Specifically, recurrent models benefit from scaling in width, while attention-based models benefit from scaling in depth.
• Figure 5: Performance of a Transformer with RoPE and Mamba on the copy task following jelassi2024repeat implementation. This task also highlights the narrow window of suitable learning rates that allows Mamba to solve the task.