Linear-Time Sequence Modeling with MLPs

We thank the reviewer for their thoughtful feedback and careful reading of our paper. We particularly appreciate the identification of the typo in Eq. 8, where indeed $q_i$ should be $q_j$.

 

Weakness #1

Thank you for pointing out the potential mismatch between our positioning and our experiments. We are drafting a version that makes our motivation and takeaways clearer.

 

Weakness #1a

We note that the copying task experiments already uses the single-head transformer.

 

Weakness #1b

In Figures 5 and 6, while the Quadratic CausalRN surpasses Transformer on the copying task, the observed differences are not significant. We think Transformer and the Quadratic CausalRN are generally on par. We don't think there is a fundamental mechanism that would make one significantly faster than the other. To be clear, we do think there is a fundamental reason for them to be on par, which is elaborated in the paper (infinitely-growing matrix-valued memory).

 

Weakness #2

In our current experiments, we stacked 12 CausalRN blocks as this is a standard configuration. We believe the quantitative results on language modeling and image classification already show the benefit of stacking blocks and residual connections.

 

Proposed Changes

We are working on a refined version of the paper incorporating the reviewer's feedback, especially better aligning our positioning and experiments.

Thank you for your insightful comments about our paper. We are glad you agree that our approach is logically presented and that our ablation study is thorough. Below, we address each of your concerns and questions.

 

Regarding Weakness #1 (Similarity to Linear Attention)

CausalRN is not a special case of linear attention [1] or its variants [2]. We will explain from a higher level and then present a mathematical proof.

High-level explanation:

Linear attention [1] maintains a matrix-valued memory state of shape (d, d), where d is the embedding size. CausalRN, however, maintains a vector-valued memory state of an arbitrary length. A typical choice is (4d,), but it can be larger or smaller (see Fig. 4(d)).

Proof that CausalRN cannot be a special case of linear attention [1]:

Let us start by expressing linear attention and CausalRN in their recurrent forms.

---

Linear Attention (see Algorithm 1 (forward pass) in [1]):

$ S_i &= S_{i-1} + \phi(k_i)^\top v_i, \newline S_i &= \sum_{j=1}^i \left( \phi(k_j) \otimes v_j \right), \newline o_i &= \phi(q_i) S_i, \newline o_i &= \phi(q_i) \sum_{j=1}^i \left( \phi(k_j) \otimes v_j \right). $

---

CausalRN (see Eq. 9 in our paper):

For the purpose of this discussion, we use the notation in [1], replace $\exp$ with $\phi$, and remove the bias terms $b_{\text{in}}$ in CausalRN.

$ o_i &= \phi(q_i) \odot \sum_{j=1}^i \phi(k_j), \newline &= \phi(q_i) \text{diag}\left( \sum_{j=1}^i \phi(k_j) \right), \newline &= \phi(q_i) \sum_{j=1}^i \text{diag}\left( \phi(k_j) \right). $

---

Notice that when $\phi$ is $\exp$:

• $\text{rank}\left( \text{diag}\left( \phi(k_j) \right) \right) \equiv d$,
• $\text{rank}\left( \phi(k_j) \otimes v_j \right) \equiv 1$.

Under the mild condition that the embedding size is $d > 1$, $\text{diag}\left( \phi(k_j) \right)$ cannot be rewritten into $\phi(k_j) \otimes v_j$. Therefore, CausalRN cannot be a special case of linear attention.

QED

Other Considerations:

• For a similar reason, CausalRN is not a special case of gated linear attention [2], even though many linear attention variants are.
• Quadratic CausalRN cannot be represented as either linear attention or gated linear attention, as the memory state is concatenative and grows with the number of tokens like self-attention.
• CausalRN uses the exponential function and a bias term, both are not found in linear attention.

Conclusion:

Indeed, the only similarity between CausalRN and linear attention is the use of partial sum to enforce causality. We believe what really sets architectures apart is their motivation and properties, and as we will discuss later, CausalRN is novel and interesting on both fronts. We hope this explanation resolves the reviewer's concern.

 

References:

[1] Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proceedings of ICML 2020

[2] Yang, S., Wang, B., Shen, Y., Panda, R., & Kim, Y. (2024). Gated Linear Attention Transformers with Hardware-Efficient Training. In Proceedings of ICML 2024

Regarding Weakness #2 and #3 (Low Performance & No Clear Contribution)

Our work is a scientific investigation in ML that aims to reveal surprising phenomena and explore new research directions. This type of paper, as elaborated by Nakkiran & Belkin in [3], sits between forming a complete theoretical framework and achieving state-of-the-art empirical results. We hope that the reviewer also recognizes the value of such contributions.

From a theoretical perspective, we contributed several surprising discoveries:

• The ability to train a quadratic number of shared-weight MLPs in linear time.
• The second ever architecture to match Transformers on the copying task [4].
• A naturally emerging gating mechanism. For contemporary architectures, the gating mechanism is usually a deliberate design choice.

We presented a breadth of evaluations covering language modeling, image classification, the copying task, interpretability, and necessary ablation studies. This allowed us to report several interesting phenomena that have not been observed by prior work:

• The phase change phenomenon of Quadratic CausalRN (Figure 5).
• Convergence behavior in the State Expansion experiments (Figure 4).
• Attention-like interpretability results (Figure 7).

Our work opens several promising research directions for future investigation, such as:

• How does a matrix-valued memory state help with in-context retrieval?
• What are the roles of exponential functions in sequence modeling?
• How does the infinite width limit affect the behavior of (Causal) Relation Networks?

These additional new directions are out of scope for this paper, but we agree with the reviewer that these will be worthy research investigations for future work.

Given the extent of the contributions mentioned above, we respectfully ask the reviewer to reconsider the contribution rating.

 

Questions

1. What are the concrete benefits and practical use cases where CausalRNs would be preferred over existing architectures?

• CausalRN presents a new, alternative approach to sequence modeling.
• CausalRN provides a new perspective to advance our theoretical understanding of sequence modeling.
• CausalRN is not intended to replace Transformer or its variants in practical settings.

2. Beyond being an "interesting research object," can the authors suggest ways in which this could advance our theoretical understanding of sequence modeling?

By comparing and contrasting with existing architectures, we were able to filter out the mechanisms that don’t matter and highlight the mechanisms that do matter:

• Exponential functions are indispensable for the phase change phenomenon to occur on a non-Transformer architecture.
• We verified on a non-Transformer architecture that a matrix-valued memory state is needed for good in-context retrieval ability.
• We showed that by increasing the state size, the retrieval ability can be significantly increased. A linear model can approach Transformer-level retrieval with a large state size.

These findings can inform future research on efficient sequence modeling architecture design.

3. How does this architecture fundamentally differ from linear attention mechanisms, and is it a special case of linear attention?

Please see our response to Weakness #1. CausalRN is not a special case of linear attention.

 

References:

[1] Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proceedings of ICML 2020

[2] Yang, S., Wang, B., Shen, Y., Panda, R., & Kim, Y. (2024). Gated Linear Attention Transformers with Hardware-Efficient Training. In Proceedings of ICML 2024

[3] Nakkiran, P., & Belkin, M. (2022). Incentivizing Empirical Science in Machine Learning: Problems and Proposals. In ML Evaluation Standards Workshop at ICLR 2022.

[4] Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat After Me: Transformers Are Better Than State Space Models at Copying. In Proceedings of ICML 2024.

Thank you for your insightful comments and careful review of our paper. We are glad you agree that our approach is novel and our discoveries are interesting. We share the view that it is important and useful to explore the space of powerful, simplified models. Below, we address each of your questions.

 

Q: "There isn't quite enough detail of the entire architecture to be able to easily recreate it in their experiments. Was this a replacement for the attention layer, was no other layer used then 12 of the BiRN layers? They mention they can use residual connections - were they used? What tokenizer was used for the text experiment?"

A: We agree. In our newly submitted version, we have addressed these points in the implementation details and Section 4.3 Quantitative Results. To clarify: yes, BiRN/CausalRN is a direct replacement for the attention layer, and no other layers are used. We do apply residual connections. For text experiments, we used an ASCII-based character-level tokenizer with a vocabulary of 128.

 

Q: "Minor corrections."

A: Thank you for your careful reading! In our recently submitted version, we have corrected the typo in Equation 8, updated Figure 3 to be red-green colorblind friendly, and clarified which experiments have quadratic versus linear time complexity. If Figure 3 is still not accessible enough, please give us direct suggestions and we will implement them in our camera-ready version.

 

Q: "3c seems surprising - I got from the text that the advantage of the exp is computation, vs. training time or accuracy. Any comment on why the exp is more efficient for training? Would something simple like a bias term help"

A: To our knowledge, the theoretical explanation for why softmax or exponential functions appear indispensable in modern sequence modeling architectures remains an open question, and we welcome future research on this topic. One possible explanation is that exp(x) enables important tokens to dominate the feature vector after the reduction step. We have clarified this point in the new version.

 

Q: "Was positional encoding used in the text/image tasks?"

A: All experiments used a randomly initialized and trainable positional embedding scheme. While this was mentioned in the implementation details section, we have made it more explicit in our new version.

 

Q: "What was the relative computational cost in the experiments vs Transformer and Mamba?"

A: In our experiments, Linear CausalRNs are computationally lighter than Transformers and comparable to Mamba. The quadratic CausalRN, since it is not I/O-aware and requires O(n²) storage in VRAM, is significantly more resource-intensive than Transformers with Flash Attention. Nevertheless, the quadratic CausalRN remains computationally feasible. We were able to finish training a quadratic CausalRN with a context window of 512 within one and a half hour on the copying task.

 

Q: "I’m confused by the state expansion experiment - clarifying within the paper or an appendix would be useful."

A: We elaborated the state expansion experiment with more detail in the paper. In essence, we increased the number of hidden neurons in the MLPs and measured how it affects the memory capacity of CausalRN. We observed that, as the number of hidden neurons increases, the model's memorization capacity also grows. We observe an early convergence phenomenon, suggesting sequence modeling architectures may not require infinite memory states to achieve perfect retrieval.