You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet

Q1. More speed results

A1. Thank you for your suggestion. We have added a comparison of the speed between Mamba2 and Flash Attention, as shown in the table below.

Q2. Error bars on experiments

A2. Thank you for your suggestion. Noting that some experiments (e.g., LLMs and Diffusion models) are computationally expensive, it is challenging to run multiple repetitions. To verify the robustness of our method, we conducted five experiments on the ImageNet-1k dataset. The results are shown below:

The results demonstrate that LightNet achieves stable performance on ImageNet-1k.

Q3. Results table not easy to parse

A3. Thank you for your suggestion. In the next version, we will bold the best results for better clarity.

Q4. Missing reference

A4. Thank you for your suggestion. We will address these issues in the revised version.

Q5. "Directionality" in additive decay

A5. Thank you for your suggestion. This describes a causal setting, such as in language modeling, where each time step depends only on previous ones. In this scenario, only the LHS exists without the RHS, which is reasonable.

Q6. Setup in Figure 2

A6. In Figure 2, we compare LightNet Attention with Linear Attention (2-scan) at the operator level. The input is randomly generated, and we vary the sequence length to measure the time consumed by each operator.

Q1. What does "global information" mean?

A1. Thank you for your question. Let me clarify the meaning of "global information" as described in lines 249–250. Based on Eq. (3) and the assumption $0\le a_t \le 1$, $a_t$ can take two forms:

• $a_t=\rho(t), 0\le \rho(t)\le 1;$ where at depends only on the current time step t.
• $a_t=\frac{\sum_{s=1}^{t-1} \delta(s)}{\sum_{s=1}^t \delta(s)}, \delta(s) \geq 0$, where at depends on all time steps from 1 to t, i.e., "global information."

We will clarify this distinction further in the updated manuscript.

---

Q2. Definition of the decay function

A2. Thank you for your suggestion. In Sec. 3, we use a simple linear recursion as a toy example to analyze the decay function. The implementation in real models differs. In Sec. 4, the decay function depends on $k_t$, indicating it is data-dependent—a standard approach in the Linear Attention community, as seen in [1], [2], and [3].

Q3. Figure 3 not clear

A3. Thank you for your suggestion. In Figure 3, we applied $ψ$ to $U$. We will update the figure for clarity in the revised version.

Q4. Typo in Eq. 10 & 11

A4. This was indeed a typo. The RHS should be $t$. These equations describe relative position encoding [4]:

$\mathbf{y}\_{n\_1, \ldots, n\_k}=\sum\_{m\_k \leq n\_k} \ldots \sum\_{m\_1 \leq n\_1} \mathbf{t}\_{n\_1-m\_1, \ldots, n\_k-m\_k} \mathbf{x}\_{m\_1, \ldots, m\_k}.$

---

Q5. Significance of Theorem 3.1

A5. Thank you for your suggestion. While the result of Theorem 3.1 is textbook material, its application to analyzing decay in Linear Attention is, to our best knowledge, novel. Current approaches to decay in Linear Attention predominantly use multiplicative decay. To highlight our contribution, we will emphasize Proposition 3.2, which represents the core of this work, in future versions.

Q6. Differences between multiplicative and additive decay

A6. Thank you for your suggestion. We believe our naming conventions are appropriate, as decay can be unified through the following framework:

• Decay: $\lambda_t=s_{t-1}/s_t$, where:
  • Multiplicative decay: $\log s_t=\log s_{t-1} + \delta(s)$.
  • Additive decay: $s_t=s_{t-1}+\delta(s)$.

Since addition in logarithmic space is equivalent to multiplication, this justifies the term "multiplicative decay." We will add this explanation in the next version of the paper.

---

Q7. Interpretation of experimental results

A7. Our main contributions are twofold:

1. Achieving a global receptive field in a single scan, leading to faster processing.
2. Matching the performance of multi-scan and Transformer-based models, as demonstrated in Tables 1, 2, and 3.

A single-scan baseline has limited value in tasks requiring a global receptive field, as it inherently performs worse. For example, GPT underperforms BERT in text classification [5], and Table 4 in [6] shows significant advantages of multi-scan models in image classification.

Additionally, we demonstrate that our model is competitive in language modeling—a domain where most Linear Attention models are applied. This is the significance of Tables 4 and 5.

---

Q8. More details about Table 1

A8. In Table 1, we replaced the token mixer with LightNet Attention. This controlled comparison, a standard approach in the field [3], [7], evaluates model performance end-to-end. The results validate the effectiveness of LightNet.

---

Q9. Updating model sizes in tables

A9. Thank you for your suggestion. We will include model sizes in the revised manuscript. As a simple comparison, LightNet has 672 million parameters, while Dit has 675 million parameters.

---

Q10. Parameter sharing

A10. Thank you for your suggestion. This approach was first introduced in [4], so it is not our primary contribution. We will include proper citations in future versions.

---

Citations:

1. Albert Gu, & Tri Dao. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces.
2. Tri Dao, & Albert Gu. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.
3. Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, Yiran Zhong. (2024). HGRN2: Gated Linear RNNs with State Expansion.
4. Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, & Yiran Zhong. (2023). Toeplitz Neural Network for Sequence Modeling.
5. Jacob Devlin, Ming-Wei Chang, Kenton Lee, & Kristina Toutanova. (2018). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.
6. Benedikt Alkin, Maximilian Beck, Korbinian Pöppel, Sepp Hochreiter, & Johannes Brandstetter. (2024). Vision-LSTM: xLSTM as Generic Vision Backbone.
7. Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, & Xinggang Wang. (2024). Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model.

Q6. FLOPs and latency analysis against FlashAttention and Mamba

A6. We compared the results of LightNet with TNL's Lightning Attention [9] in Figure 2. As described in [9], Lightning Attention is faster than FlashAttention, and LightNet outperforms FlashAttention in speed. Additionally, Mamba2 employs data-dependent decay, making it slower than Lightning Attention. To verify this, we conducted further tests, and the results are as follows: The time unit is second.

Theoretical FLOPs Analysis

Let d denote the embedding dimension, n the sequence length, h the number of heads, and B the chunk size in Mamba2.

Softmax Attention FLOPs:

• qk multiplication: $2 n^2 d$
• Softmax: exp, sum, divide $3 n^2 h$.
• Sofmtax(qk)v multiplication: $2n^2d$.

Total: $4n^2d + 3hn^2$.

Mamba2:

• Intra block: $(n/B)\times (4 B^2 d+B^2 h)=4Bnd+nBh$.
• Inter block: $n/B\times \frac{2 B d^2}{h}=2nd^2/h.$
• State update: $(n/B) \times \left( \frac{2 B d^2}{h}+\frac{d^2}{h} \right)=2nd^2/h+nd^2/(Bh).$

Total (2 pass): $8nd^2/h + 8Bnd + 2Bnh + 2nd^2/(Bh).$

LightNet FLOPs:

• kv multiplication: $2nd^2/h$；
• q(kv) multiplication: $2nd^2/h$；

Total: $4nd^2/h.$

These results demonstrate that LightNet has theoretical and practical speed advantages.

Q7. Can models similar to Mamba utilize a 1-scan method?

A7. Thank you for your question. The answer is no. As discussed in our Proposition 3.2, Mamba-like models belong to the Multiplicative Decay category. Let me explain the specific reasons. Linear Attention has shown superior speed and competitive performance compared to Softmax Attention in one-dimensional language modeling. However, when dealing with multi-dimensional data, Softmax Attention can obtain results through a single parallel computation. In contrast, existing Linear Attention methods require at least a two-scan process to achieve the same receptive field as Softmax Attention due to the presence of multiplicative decay. LightNet, on the other hand, uses additive decay, which allows it to achieve results with just one scan.

Citations:

1. Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, Soham De. "Resurrecting Recurrent Neural Networks for Long Sequences." (2023).
2. Mega: Moving Average Equipped Gated Attention. Ma et al. ICLR 2023
3. Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length. Ma et al.
4. Rethinking and Improving Relative Position Encoding for Vision Transforme. Wu et al. ICCV 2021
5. Ma, Xin, Yaohui, Wang, Gengyun, Jia, Xinyuan, Chen, Ziwei, Liu, Yuan-Fang, Li, Cunjian, Chen, Yu, Qiao. "Latte: Latent diffusion transformer for video generation". arXiv preprint arXiv:2401.03048. (2024).
6. Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, Oskar van der Wal. "Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling." (2023).
7. Teven Le Scao, et al.. "BLOOM: A 176B-Parameter Open-Access Multilingual Language Model". CoRR abs/2211.05100. (2022).
8. William Peebles, Saining Xie. "Scalable Diffusion Models with Transformers." (2022).
9. Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, Yiran Zhong. "Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models." (2024).

Q1. The paper lacks clarity.

A1. Thank you for your question. We will address these issues in future versions. Here are the answers to your concerns:

• Softmax operation in line 280: This refers to the activation function used in the Key module, and it has no relation to Softmax Attention.
• W*{u1} and W{u2} in line 283*: These correspond to the weight matrices of the "low-rank output gate from TNL3" mentioned in lines 268 to 269.
• Line 247: This is not a typo. The function $\exp(-\exp(x))$ maps $\mathbb{R}$ to [0, 1], aligning with the definition of decay. Additionally, this is illustrated in Figure 1 of [1].
• Broken references and typos: We will fix these in the updated version.

---

Q2. Missing baselines

A2. Thank you for your suggestion. We will include discussions of [2] and [3] in subsequent revisions. Note that in [2] and [3], only the Mega-chunk variant exhibits linear complexity. Moreover, the Softmax function in equations 7–9 is merely an activation function derived from additive decay. To prevent confusion, we will clarify this point in future versions.

---

Q3. Missing baselines and related work on N-Dimensional position encoding

A3. Thank you for your suggestion. Most N-Dimensional position encoding methods involve operations on the attention matrix, which are incompatible with Linear Attention as we utilize the right-product trick. For instance, the official code for [4] includes:

attn = (q @ k.transpose(-2, -1))

# image relative position on keys
if self.rpe_k is not None:
  attn += self.rpe_k(q)

This approach clearly cannot be applied to Linear Attention, motivating the design of our own N-Dimensional position encoding.

---

Q4. Limited scope of experimental validation

A4. Thank you for your suggestion. To evaluate the proposed LightNet on higher-dimensional tasks, we replaced the Transformer in Latte [5] with LightNet. The results are shown below, demonstrating consistent performance improvements on both datasets. The metric used is FVD (lower is better).

---

Q5. Robustness of the results

A5. Thank you for your suggestion. Some experiments, such as those involving LLMs and Diffusion, are computationally expensive and difficult to repeat multiple times (notably, baselines [6], [7], [8] also ran only one trial). To assess robustness, we conducted five runs on the ImageNet-1k dataset. The results are shown below:

The results show that LightNet delivers stable performance on ImageNet-1k.

Q1. Saturation of Additive Decay for Long Sequences

A1. Thank you for your suggestion. From an intuitive perspective, the saturation of additive decay does not prevent distinguishing local information. This is comparable to Softmax Attention, which can still focus on local information in long texts. To validate this point, we compared the Needle-in-a-Haystack experiments of Mamba-2 3b [1] and LightNet 3b [2], with both models trained on the same 300b tokens. As shown in the figure below, LightNet significantly outperforms Mamba in this experiment, without any issues in distinguishing local information.

LightNet

Mamba2

Q2. Comparing LA and LightNet

A2. Thank you for your suggestion. We removed the additive decay from LightNet, replaced it with $\phi(x) = \mathrm{silu}(x)$, and trained LA models of size 1b and 3b. We then compared them to LightNet 1b (trained on the same 300b tokens). The table below presents the results on Commonsense Reasoning Tasks, showing that LightNet's average accuracy is nearly 2 points higher, demonstrating the effectiveness of additive decay.

Q3. The Role of Positional Encoding

A3. Thank you for your suggestion. Let me explain why RoPE is not suitable in this scenario. Assuming RoPE [3] is used, take a language model as an example:

\begin{aligned} \mathbf o_t^{\top} &= \sum_{s\le t}\mathbf q_t^{\top} \mathbf W_t^{\top} \left((\mathbf W_s {\exp(\mathbf k_s)})/\left( \sum_{j=1}^t \exp(\mathbf k_j) \right) \right)\mathbf v_s^\top \\\\ &= \sum_{s\le t}\mathbf q_t^{\top} \mathbf W_t^{\top} \mathrm{diag}\left\\{ \sum_{j=1}^t \exp(\mathbf k_j) \right\\}^{-1} \left(\mathbf W_s {\exp(\mathbf k_s)}\right)\mathbf v_s^\top \\\\ &\neq \sum_{s\le t}\mathbf q_t^{\top} \mathbf W_{t}^{\top}\mathbf W_s \mathrm{diag}\left\\{ \sum_{j=1}^t \exp(\mathbf k_j) \right\\}^{-1} {\exp(\mathbf k_s)}\mathbf v_s^\top \\\\ &= \sum_{s\le t}\mathbf q_t^{\top} \mathbf W_{t-s}^{\top} \mathrm{diag}\left\\{ \sum_{j=1}^t \exp(\mathbf k_j) \right\\}^{-1} {\exp(\mathbf k_s)}\mathbf v_s^\top. \end{aligned}

The inequality in the second-to-last step arises because block-diagonal matrices (RoPE matrices) and diagonal matrices are non-commutative. However, this issue does not occur in LRPE [4], which is implemented as follows:

$$f_{\mathrm{lrpe}}(\mathbf x,\Theta)=\mathrm{concat}([\mathbf x \odot\cos \Theta, \mathbf x \odot\sin \Theta], \mathrm{dim}=-1).$$

Thus, the computation becomes:

$$\begin{aligned} \mathbf o_t^{\top} &= \sum_{s\le t}[\mathbf q_t \odot\cos(t\Theta ), \mathbf q_t \odot\sin (t\Theta)]^{\top} \left[{\exp(\mathbf k_s)}/\left( \sum_{j=1}^t \exp(\mathbf k_j) \right) \odot\cos(s\Theta ), {\exp(\mathbf k_s)}/\left( \sum_{j=1}^t \exp(\mathbf k_j) \right) \odot\sin (s\Theta) \right]\mathbf v_s^\top \\\\ &= \sum_{s\le t}\mathbf q_t^\top \mathrm{diag}\\{\cos ((t-s)\Theta)\\} {\exp(\mathbf k_s)}/\left( \sum_{j=1}^t \exp(\mathbf k_j) \right) \mathbf v_s^\top . \end{aligned}$$

This shows that LRPE can capture relative positional information, which is why we chose LRPE.

Q4. 1b Result

A4. Refer to A2.

Q5. $O(\text{Sequence Length})$ Time Algorithm in Causal Setting

A5. Thank you for your suggestion. This is not particularly complicated. We discuss two cases separately:

Without LRPE:

\begin{aligned} &\mathbf {s}\_t=\mathbf {s}\_{t-1}+\exp \left(\mathbf {k}\_t\right), \overline{{k}}\_{{t}}=\exp \left({\mathbf k}\_t\right) / \mathbf {s}\_t, \\\\ &\lambda\_t = \mathbf s\_{t-1}/\mathbf s\_t , \\\\ &\mathbf {k v}\_t=\operatorname{diag}\left\\{\lambda\_t\right\\} \mathbf {k v}\_{t-1}+\overline{\mathbf {k}}\_{{t}} \mathbf {v}\_t^{\top}, \\\\ &\mathbf o\_t^\top = \mathbf {kv}\_t^\top \mathbf q\_t. \end{aligned}

With LRPE:

The recursive formulation becomes:

\begin{aligned} &\mathbf {s}\_t=\mathbf {s}\_{t-1}+\exp \left(\mathbf {k}\_t\right), \tilde{\mathbf {k}}\_{{t}}=\exp \left(\mathbf {k}\_t\right) / \mathbf {s}\_t, \\\\ &\lambda\_t = \mathbf s\_{t-1}/\mathbf s\_t ,\\\\ &\overline {\mathbf q}\_t = [\mathbf q\_t \odot\cos(t\Theta), \mathbf q\_t \odot \sin(t\Theta)], \overline {\mathbf k}\_t = [\tilde {\mathbf k}\_t \odot\cos(t\Theta), \tilde {\mathbf k}\_t \odot \sin(t\Theta)], \bar \lambda\_t = [\lambda\_t, \lambda\_t], \\\\ &\mathbf {k v}\_t=\operatorname{diag}\left\\{\lambda\_t\right\\} \mathbf {k v}\_{t-1}+\overline{\mathbf {k}}\_{{t}} {\mathbf v}\_t^{\top}, \\\\ &\mathbf o\_t^\top = {\mathbf kv}\_t^\top \bar {\mathbf q}\_t. \end{aligned}

In either case, after transformation, we convert it into a recursive formulation consistent with FLA [5], enabling $O(n)$complexity computation, where $n$ is the sequence length.

Q6. Explain "Parameter Tiling Strategy"

A6. What we actually meant was "parameter sharing strategy." Specifically, the nn.Linear parameters in the additive decay part are shared with the nn.Linear parameters in the key computation, a technique also used in HGRN2 [6].

Citations

1. Tri Dao, & Albert Gu. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.
2. gkamradt. GitHub - gkamradt/LLMTest_NeedleInAHaystack: Doing simple retrieval from LLM models at various context lengths to measure accuracy –- github.com.
3. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, & Yunfeng Liu. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding.
4. Zhen Qin, Weixuan Sun, Kaiyue Lu, Hui Deng, Dongxu Li, Xiaodong Han, Yuchao Dai, Lingpeng Kong, & Yiran Zhong. (2023). Linearized Relative Positional Encoding.
5. Yang, S., & Zhang, Y. (2024). FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism. https://github.com/sustcsonglin/flash-linear-attention
6. Zhen Qin, , Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, Yiran Zhong. "HGRN2: Gated Linear RNNs with State Expansion." (2024).