Revisiting associative recall in modern recurrent models

TL;DR

This paper investigates associative recall in modern recurrent models, highlighting the importance of learning rate, contrasting scaling benefits of recurrent versus attention models, and analyzing training dynamics and architecture effects.

Contribution

It provides a detailed analysis of associative recall in recurrent models, emphasizing the role of learning rate, scaling differences, and architectural components, which were previously underexplored.

Findings

Learning rate critically affects recurrent model performance.

Attention models struggle with associative recall when limited to one layer.

Training dynamics of 1-layer transformers resemble induction head formation.

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.