Explaining Modern Gated-Linear RNNs via a Unified Implicit Attention Formulation

Authors: Itamar Zimerman, Ameen Ali Ali, Lior Wolf
Primary area: interpretability and explainable AI
Keywords: Explainability, Interpretability, Gated-Linear RNNs, Attention-free, Mamba

Abstract:
Recent advances in efficient sequence modeling have led to attention-free layers, such as Mamba, RWKV, and various gated RNNs, all featuring sub-quadratic complexity in sequence length and excellent scaling properties, enabling the construction of a new type of foundation models. In this paper, we present a unified view of these models, formulating such layers as implicit causal self-attention layers. The formulation includes most of their sub-components and is not limited to a specific part of the architecture. The framework compares the underlying mechanisms on similar grounds for different layers and provides a direct means for applying explainability methods. Our experiments show that our attention matrices and attribution method outperform an alternative and a more limited formulation that was recently proposed for Mamba. For the other architectures for which our method is the first to provide such a view, our method is effective and competitive in the relevant metrics compared to the results obtained by state-of-the-art Transformer explainability methods. Our code is publicly available.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730073981624

thread_type:
official_review

<thread_object>
creator: reviewer_pbpq

<content>
summary:
This paper tries to provide an implicit self-attention formulation for most state-of-the art non-transformer sequence models (known as Gated Linear RNNs) such as Mamba, Griffin, RWKV, and ReNet. In this way, it can exploit techniques used in attention explainablity to explain these new models.

Compared to the closest work (Ali et al, 2024), which only formulates the S6 layer in Mamba, the main contribution of paper is:
- Formulating more layers with implicit self-attention and propose a unified and more accurate attention-based representation for all the SOTA gated linear RNNs.

Other contributions include:
- Introducing new explainablity technique for these models leveraging the self-attention formulation
- Showing performance of their explanations and attributions by purturbation testing and segmentation.
- Showing their proposed formulation can give attributions which can further used in some performance-enhancing techniques (based on in-context learning) for large language models.

strengths:
- Covering multiple models including most popular modern non-transformer sequence models.
- Evaluating the performance of the resulting attributions in multiple quantitate experiments and down-stream tasks (across both vision and NLP).
- Showing the impact of various sub-layers in the ablations study

weaknesses:
- In their ablation study, the authors could discuss the trade-off between the time explainability and more accurate formulation/explainability.

- The main baseline paper (Ali et al, 2024) has not been published yet. So, it is hard to evaluate this paper. Actually, the performance of the model in downstream tasks such as segmentation and attribution-based performance-enhancement helped me to have better evaluation of the proposed method.

questions:
mentioned above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730456548144

thread_type:
official_review

<thread_object>
creator: reviewer_rdjf

<content>
summary:
This paper writes out gated recurrent architectures such as Mamba, RKWV, and Griffin as causal self-attention layers. The main objective of this is to increase interpretability, which is tested through perturbation and segmentation tests.

strengths:
In my view, this is a pretty complete paper. From my understanding, the authors present an extension of the paper Ali et al. (2024) to include additional architectural components in the linearization and not just the S6 mechanism considered before. The authors show that the model improves interpretability through visualizations and a set of both perturbation and segmentation tests. The ablation gives quite a lot of strength to their arguments, but I am not so familiar with these types of explainability results so I am not able to comment on the details.

weaknesses:
While the authors have explored the interpretability side of things extensively, I was wondering if it would be worth comparing the performance of the linearized models compared to its recurrent counterparts when trained on some small datasets?

questions:
See above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691439628

thread_type:
official_review

<thread_object>
creator: reviewer_jknz

<content>
summary:
This paper introduces a unified framework that reformulates various gated recurrent neural network (RNN) architectures, such as Mamba, RWKV, and Griffin, into implicit causal self-attention layers. This reinterpretation aims to make these models more interpretable by constructing attention-like matrices for use in visual and NLP explainability tasks. Experimental evaluations demonstrate that this approach achieves competitive performance in robustness and attribution metrics, though prior work has already suggested that certain gated models, including Mamba, are attention-based.

strengths:
The paper offers a clear explanation of how implicit attention could be used to interpret gated RNNs, making it accessible to readers interested in explainability across model types.

By applying the framework to both NLP and vision tasks, the authors demonstrate its cross-domain relevance.

weaknesses:
Previous work has already conceptualized models like Mamba as attention-like, meaning that simply reinterpreting these gated RNNs under an implicit attention framework may not be largely novel.

The paper does not thoroughly compare its implicit attention framework with existing interpretability tools for gated RNNs.

questions:
A critical question is: given that previous work has noted the attention-like properties of models such as Mamba, what specific benefits does the implicit attention framework offer over these prior interpretations?

What are the current explainability methods or metrics for modern gated-linear RNNs and how's the comparison between them and attention matrices?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731179722644

thread_type:
official_review

<thread_object>
creator: reviewer_xswf

<content>
summary:
The paper studies the problem of sequence modelling. The authors aim to provide a unified framework of the recent attention-free methods such as Mamba and RWKV. The paper presents empirical results to validate the proposed unified framework.

strengths:
- The unified framework makes it easier to study and compare the different sequence modelling algorithms.
- It is important for the community to learn about such work.
- The experimental results are interesting.

weaknesses:
Although it is hard to evaluate such approaches empirically, interpretability-based metrics are not very conclusive in general.

questions:
- What is the best method to evaluate the unified framework beside the interpretability analysis?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731800259413

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Comment by Authors

comment:
We genuinely appreciate the reviewers' insightful and constructive feedback, which has significantly improved our manuscript and inspired further advancements in our research.

To begin, we deeply value the positive feedback provided by the reviewers. Notable examples include the following:
> This is a pretty complete paper (by Reviewer rDjF).

> It is important for the community to learn about such work (by Reviewer xsWF)

> The unified framework makes it easier to study and compare the different sequence modelling algorithms (by Reviewer xsWF).

> The paper offers a clear explanation of how implicit attention could be used to interpret gated RNNs, making it accessible to readers interested in explainability across model types (by Reviewer JKNZ).

In response to the concerns raised by reviewers JKNZ and PBPq, we conducted additional experiments and incorporated the results into the **revised manuscript** (see Table 3, with modifications highlighted in blue). **Specifically, we have included benchmarking against the recent MambaLRP** work by Jafari et al. (NeurIPS 2024) [1], which introduced an LRP-based XAI technique for Mamba models. As demonstrated in Table 3, **our method consistently and significantly outperforms this new baseline**, establishing a new SoTA for providing explanations for Mamba models. Notably, our approach achieves over a 10% improvement across all metrics, with an average gain exceeding 15%. We also note that our framework applies and was tested across multiple architectures, while MambaLRP was applied only to Mamba.

Moreover, in response to reviewer xsWF, we added a new section in Appendix E that discusses the insights arising from our formulation regarding the inner dynamics and parametrization of gated-linear RNNs. In short, our formulation reveals how these layers efficiently and implicitly parametrize thousands of data-dependent attention heads. These attention matrices exhibit properties such as locality and smoothness (enabled by Conv1D operations), alongside the capacity to capture sparse and complex features (via gating mechanisms and non-linear activations). In architectures like Mamba-2 and RetNet, these matrices are further refined through implicit, data-dependent normalization achieved via the LayerNorm applied at the end of the block.

Additional comments and responses are provided in the threads below. 

Please let us know if there are any remaining questions or concerns we can address.

___
[1] MambaLRP: Explaining Selective State Space Sequence Models. Jafari et al. NeurIPS 2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1732719377623

thread_type:
public_comment

<thread_object>
creator: ~da_xiao1

<content>
title:
Thanks for your work!

comment:
It is very insightful and helpful to me. A few questions on technical details:
1. I have some confusion on the shapes of tensors in Eq (9). E.g. why is W_x' of shape L * L? According to Eq (8), x is L * D. Shouldn't W_x' 
 = diag(SILU(Linear_2(x))) be D' * L * L? (I assume Linear_2.W has shape D * D', though I'm not sure what D' should be.) And what's the shape for the final resulting tensor H? It would be very helpful if you could explicitly give shapes for each tensor in Equations 9, 11, 17, 21, 24 and 34.

2. In Appendix E you mentioned: "This implies that these components collectively enable the model to efficiently and implicitly express thousands of data-dependent (where transformers of the same size typically have around a few dozen), complex attention heads." 
Why are there **thousands of** data-dependent heads in gated-linear RNN layers?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper provides a unified view of attention-free architectures as implicit causal self-attention layers, which they apply to explainability. They also show good results on semantic segmentation and in-context learning.

Overall, the reviewer's side on the acceptance. The main criticism is that the idea that these models exhibit attention-like behavior has been noted before, but the paper does a good job comparing. An interesting direction could have been exploring the tradeoffs between the complexity of the implicit attention formulation vs accuracy