Resona: Improving Context Copying in Linear Recurrence Models with Retrieval

While Linear Recurrent Models (LRMs) demonstrate computational efficiency advantages, their fixed-size hidden states create an information bottleneck that limits effective utilization of long-context information. To address this, we propose Resona, a framework that enhances LRMs' in-context learning capabilities through a retrieval mechanism. This approach bypasses the hidden state constraints of traditional LRMs by enabling direct information flow from input contexts to model reasoning. Extensive experiments demonstrate that Resona significantly improves LRMs' performance without requiring major architectural modifications, offering plug-and-play usability. I recommend clear accept.

This paper proposes a new neural network module named Resona, that resembles an RAG procedure on a per-layer basis. The context is first chunked and embedded and relevant chunks are retrieved using queries corresponding to the hidden state of each token. The information retrieved chunks are then fused back into the model using cross-attention.

This paper is well-motivated as it is known that one of the critical bottlenecks of linear attention modules is that it could not retrieve perfectly from the context. While a hybrid model with softmax attention has been proposed to address this issue, this paper's more explicit way of fusing retrieval is a meaningful exploration.

While the paper is overall clear, there are two issues of the paper that need to be addressed (I will expand more in the reasons to reject):

1. While the 'formal' theory part has good intuition, the formulation of this theory is unfortunately incorrect.

2. There are some important experiment details that are not properly explained, which may hinder reproduction.

This article introduces a novel technique aimed at improving the contextual information handling capabilities of Linear Recurrent Models (LRMs), a recently proposed efficient alternative to transformer-based models based on the principles of recurrent networks. Despite their advantages in terms of computational and spatial efficiency, LRMs suffer from limitations when copying and integrating previous input fragments due to their fixed-size hidden states, thereby creating an information bottleneck. To address this, the authors propose extending LRMs with a method inspired by Retrieval-Augmented Generation (RAG). Unlike traditional RAG, which retrieves external documents, this approach encodes and retrieves embeddings of previous context passages within the input itself. These embeddings then influence the model's hidden state updates and subsequent layer outputs. The authors present rigorous experiments demonstrating significant performance improvements for LRMs such as Mamba, RWKV, and Hyena in tasks that heavily depend on context copying capabilities.

General Response

We first thank the reviewers for their comprehensive reviews as well as their mentions of strengths and areas of improvement for our work. We are happy to see that two reviewers express high enthusiasm for our work, while another has provided significant comments that can greatly improve it.

Given the shared nature of some questions/comments, we would like to provide these details in a general response to improve readability of our discussion and reduce repetition.

Computational overhead

We thank all the reviewers for mentioning this comment. We provide the following tables, which indicate pre-filling speed, generation speed as well as memory usage of a number of methods. In particular, we compare a Transformer using Flash-Attention, as well as various linear recurrent models both with and without Resona plugins.

Overall, we observe that the Resona plug-in is very lightweight and does not add significant overhead in computation, particularly for longer sequences.

Prefilling Time (ms)

Generation Time (ms)

All results are for generation of 128 tokens following the given pre-filling length

Memory Usage (GB)

Resona adaptations for specific architectures.

We thank the reviewers for pointing this out. It is true that different architectures utilize different hidden states, which can influence how Resona layers interact with them. We hope the following details provide more clarity on this front.

In the most simple setting, this hidden state is simply a vector (such as in Mamba), hence we can directly use it as queries for the Resona module.

In other settings, this can be a matrix (DeltaNet, GLA, etc.) and hence we simply take the hidden states and reshape it into a vector that is projected as queries.

Finally, as Reviewer UqFQ mentioned, we can also use the output $y$ of the previous layer as the queries. While we did not directly consider this setting in our initial manuscript, we have tried this and observed little difference compared to using $h$. Nevertheless, we do realize that for computational reasons, it can be desirable to use $y$ given that this is a vector and hence may lead to lower memory usage compared to using a matrix as the queries, which we believe can further improve our method. Our initial goal was to maintain a more general method that could be applied between architectures, but this is certainly a valid point to consider.

This paper presents Resona, a framework that augments Linear Recurrent Models with retrieval mechanisms to address their well-known issues in context copying and ICL. The main idea is to add a parallel cross-attention branch that retrieves relevant information from the input context, bypassing the bottleneck created by fixed-size hidden states.

Strengths:

• The motivation is clear and addresses a genuine limitation of LRMs
• The experimental eval is comprehensive across multiple archs(mamba, RWKV, Hyena) on both synthetic and real-world tasks
• The approach is practical and plug-and-play
• Authors provided detailed overhead analysis in their rebuttal

Weaknesses:

• The theory analysis (Lemma 3.1) seems flawed as one reviewer correctly pointed out
• Some implementation details were initially unclear around the "offline" chunk and search process and how hidden states are handled across different architectures
• perhaps lacks systematic ablation on some hyperparameters (though the authors partially addressed this during rebuttal with some results on eg chunk size)
• Missing important related work (fixable)

Overall it's a solid piece of work, that all Revs and I agree.