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

Authors: Zhen Qin, Yuxin Mao, Xuyang Shen, Dong Li, Jing Zhang, Yuchao Dai, Yiran Zhong
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: LightNet, multi-dimensional sequential modeling

Abstract:
Linear attention mechanisms have gained prominence in causal language models due to their linear computational complexity and enhanced speed. However, the inherent decay mechanism in linear attention presents challenges when applied to multi-dimensional sequence modeling tasks, such as image processing and multi-modal learning. In these scenarios, the utilization of sequential scanning to establish a global receptive field necessitates multiple scans for multi-dimensional data, thereby leading to inefficiencies. This paper identifies the inefficiency caused by a \enquote{multiplicative} linear recurrence and proposes an efficient alternative \enquote{additive} linear recurrence to avoid the issue, as it can handle multi-dimensional data within a single scan. We further develop an efficient multi-dimensional sequential modeling framework called LightNet based on the new recurrence. Moreover, we present two new multi-dimensional linear relative positional encoding methods, MD-TPE and MD-LRPE to enhance the model's ability to discern positional information in multi-dimensional scenarios. Our empirical evaluations across various tasks, including image classification, image generation, bidirectional language modeling, and autoregressive language modeling, demonstrate the efficacy of LightNet, showcasing its potential as a versatile and efficient solution for multi-dimensional sequential modeling.

Decision: reject
Reviewer scores: 5, 3, 3, 8

Reviews and discussion:
<thread>
creation date:
1730585530739

thread_type:
official_review

<thread_object>
creator: reviewer_4dh3

<content>
summary:
The authors start by exploring the space of linear recurrences, and when they can represent linear combinations.  They then define a constraint to keep this recurrence numerical stable, and two different ways to satisfy the constraint.  They call these “additive” and “multiplicative”.  

They show that the SSMs normally use a “multiplicative” linear recurrence, which stops information from “future” of the sequence affecting the current position, requiring two scans along a sequence, one forward and one backward.   Instead, they modify their “additive” linear recurrence to be a “global” summary of information present by modifying the normalization. They show this modified and novel “additive” linear recurrence is significantly computationally cheaper than a 2-scan approach (Figure). This is particularly useful when dealing with multi-dimensional arrays, where there is no natural scan order, and multiple scans in different orders might be required.    They also demonstrate that using one-scan is significantly faster in practice (Figure 2), but they do not compare speed for their other experiments.

However, their ‘additive’ linear recurrence loses the relative position of tokens, so they introduce two new forms of positional encoding they call Multi-Dimensional Toeplitz Position Encoding (MD-TPE), and Multi-Dimensional Linearized Relative Positional Encoding (MD-LRPE), where MD-TPE is applied after the input embedding and MD-LRPE is applied after each Q/K computation.   

Together, they use their one-scan attention and MD-LRPE blocks to construct a LightNet Attention (LNA) block, which alternates with a Gated Linear Unit (GLU) to construct what they call the LightNet.  

They empirically test the performance of LightNet in both multi-dimensional image classification and generation tasks, as well as at word modeling tasks.  

In image classification, they show that they perform similar to the system doing full softmax attention or multi-scan attention , while they only require their one scan system, which should be faster.  In Image generation, they demonstrate SOTA FID compared to the other methods.

In the bidirectional language modeling, they compare in 24-hour training regime, and LightNet significantly outperforms all other methods.  In autoregressive LM, they match MAMBA performance, which makes sense.

strengths:
Originality:
To my knowledge both the “additive” linear recurrence and corresponding improvements to the positional encodings system are novel, and important for non-linear regimes.  

Quality:
The technical claims are clear and appear to be solid. However, see note about it being hard to compare the computational efficiency of their methods vs. competitors due to not reporting timings in weaknesses section.  

Clarity:
The paper is clearly written and relatively easy to follow, with experiments described in sufficient detail that they could be reproduced, as well as network.  However, note that the appendices also appear to be missing, which includes one of the smaller proofs from the paper.

Significance:
They demonstrate SOTA on some of their tasks, so it is definitely the best method to use in at least certain situations.  This combined with their claims of gains in computational efficiency could make LightNet the network of choice in various situations.  It’s also surprising that their modification works as well as it does, so that itself may lead to more interesting results and modifications.

weaknesses:
It would be more informative if the timing information for the evaluation/training iteration per experiment were reported along with the results, as opposed to only being shown in Figure 2.  As it is, it is hard to judge what computational speed difference it makes in practice, which is one of the main attractions of their method.  Assuming the results in Figure 2 apply to the other methods, it would make it more impressive.  At least why this wasn't included should be addressed in the text.  

There are no error bars on any results, which makes it harder to be sure that the differences are meaningful.  However, this could be understandable given the computational expense involved.  A comment on how much variability was seen if multiple runs were performed would be useful. 

Their results table a not that easy to parse - bolding winners on each table would make it quicker and easier to read.  

There are some places where an appendix is referenced, but there are no appendices. E.g., the proof of Equation 4 is missing.

questions:
There is still a “directionality” in the additive linear recurrence, as the terms before and after a specific position have different weightings (there is a sum to the LHS, but not the RHS of a term).  Could you provide more intuition about why this asymmetry doesn't matter in practice?

It’s not clear what the setup is in Figure 2 - you say you substitute the attention block, but in what setting/dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730697285334

thread_type:
official_review

<thread_object>
creator: reviewer_tjct

<content>
summary:
This paper proposes a formulation for the decay function of the linear attention models in order to enable single-pass processing of multi-dimensional signals (such as image) in transformers. The authors also propose two positional embedding functions for multi-dimensional data, extending the previous formulations in the literature to multi-dimensional settings.

strengths:
The paper tackles an important problem re: processing of multi-dimensional data using linear transformers and its performance implications. 

The paper also proposes two novel positional embedding formulations for multi-dimensional data which are interesting and to my limited knowledge novel.

The range of the empirical study is also quite wide and contains image classification, image generation and language modeling with lots of comparisons across different models.

weaknesses:
1- Overall, this paper is poorly written; the general argument is quite vague, and the derivation is unclear with vague notation. In particular:

A) What does it mean that $a_t$ "cannot provide global information"? What do authors mean by "global information" here? I can think of any monotonic function of $t$ to capture the global information re: the position, or any aggregation function (sum, product, etc) to capture the global information re: the value. 

B) In Sec 3. the decay function is presented as a function of $t$; however, in Sec 4. all of the sudden the decay is a function of the key variables themselves with no further explanation. Why should the decay be a function of keys and not simply their time indices? There's no explanation whatsoever. 

C) Figure 3 is not consistent with the formulation presented in Eqs. 7, 8 & 9. For instance, it's not clear if the $phi$ transformation is applied on $U$ or not and if so, why should that be the case.

D) What is "y" in Eq. 10 & 11? Also there's no intuitive explanation of where these formulations come from.

I'd highly recommend the authors to (I) rewrite Sec 4 and their preceding arguments in a clear manner, (II) provide a section or appendix re: any background relevant to their work because the current manuscript is not self-contained, and (III) make sure all the notations used in the paper are consistent and clearly defined.

2- The main contribution of the paper is not clear. In particular,

A) Theorem 3.1 is simply the textbook solution for the first-order non-homogeneous recurrence equations with variable coefficients (which is what Eq. 3 is, mathematically speaking), which is already a established result.

B) The distinction between the "multiplicative" and "additive" cases is not clear. The additive case is indeed a special case of the multiplicative case where one chooses $\rho$ to be the ratio proposed for the additive case. In other words, the proposed supposedly "additive" decay formulation is in fact multiplicative underneath.

Instead of using confusing terminology, I'd recommend authors to clearly state their contribution re: their choice of decay function and the theoretical intuitions behind it. 

3- The empirical results are not particularly conclusive. More precisely,

A) If the main contribution of the paper is achieving better results with a single pass, then it makes sense for the baselines to run on single pass setting as well and then measure the potential improvement the proposed method has achieved using the proposed decay formulation. For the current reported results, it's not exactly clear where the benefit comes from.

B) In table 1, it's really hard to pinpoint the main source of improvement as there are many other factors among these models that are not controlled.

C) In addition to controlling for the number of passes, the authors should also control for model size. For instance, in Table 3, the model sizes need to be reported.

D) Since language is not a multi-dimensional signal, it is not clear what hypothesis exactly the authors are trying to test by the language experiments. In particular, is the comparison against other linear attention models of the same size? Or is it against any model of the same size? Or is it against any model whatsoever (i.e. beating the leaderboard)? In other words, what exact part of the proposed framework are we testing here and how do we control for other factors that are NOT being tested? There are too many uncontrolled factors here: models' architectures, models' sizes, models' kernel embedding functions (for linear attention models), the decay formulation, etc. This means that the presented results are to be interpreted as leaderboard benchmark results, but is that your claim to beat the leaderboard for these tasks?

E) The "parameter sharing" contribution is first mentioned in the ablation study section and unlike what is claimed there, there is no single mention of it in the methodology sections before! Is it another contribution of the paper? If so, then it needs to clearly stated so and fully explained in the methodology section.

questions:
Please see my comments above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705396498

thread_type:
official_review

<thread_object>
creator: reviewer_ytrf

<content>
summary:
The paper addresses the question of how to efficiently model multi-dimensional data with modern architectures. The authors first categorize existing linear recurrence approaches into additive and multiplicative scans, demonstrating that the additive approach is significantly more efficient. By exploring the fundamental formulation of additive scans, they propose LightNet - a linear-attention-based model variant that can be implemented in a single scan, resulting in a highly efficient model. Empirical analysis shows that LightNet achieves competitive results in both language modeling and vision tasks.

strengths:
-	The paper addresses an important question: how to efficiently process multimodal data using modern DL models.

-	The proposed '1 Scan' strategy significantly reduces latency.

-	The results are promising, demonstrating improvements over the selected baselines.

weaknesses:
**W.1. The paper lacks clarity** and contains numerous typos in critical sections, making it difficult to understand key aspects of the methods and results. For example:

-	In the method section, line 280, a softmax operation is mentioned, yet the authors claim the transformer relies on linear attention. This inconsistency is especially problematic given the absence of a reference code.

-	In the method section, line 283, the variables W_{u1} and W_{u2} are referenced, but they do not appear to be introduced or explained.

-	In the method section, line 247, there seems to be a double exponent function. Is "exp(-exp(x))" correct, or is this a typo?
-	Broken references are present (lines 180, 401).

-	Unclear notation: In Section 4.1, equations (7-9), the distinction between bold and standard letters is not defined.

-	Additional typos: line 159 (“E.q 3” should be “Eq. 3”), line 413 (“modeling..” should be “modeling.”).

**W.2. Missing Baselines (and related work) in linear attention variants:** The paper does not discuss several existing transformer-based linear attention models that have demonstrated improved performance in certain contexts. Examples include Mega [2] and Megalodon [3], which are also based on gating and decay mechanisms. The authors should also include a discussion, possibly in a new subsection, on how LightNet (equations 7-9) compares to similar variants that omit the softmax.

**W.3. Missing Baselines and related work on N-Dimensional position encoding:** As far as I understand, the authors did not discuss or benchmark previously proposed N-D position encoding methods. Is this correct? Numerous approaches address this, see [1] as an example. 

**W.4. Limited scope of experimental validation**: While the paper claims to support n-dimensional data, experiments are conducted only on 1-D and 2-D cases.

**W.5. The statistical significance** of the results is not well-defined. Were the experiments run over multiple seeds? In some cases, such as the ablations without PE, performance gains are marginal. Could the authors clarify the robustness and significance of these results?

---

[1] Rethinking and Improving Relative Position Encoding for Vision Transforme. Wu et al. ICCV 2021

[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.

questions:
Q.1. Can the authors present FLOPs and latency analysis against highly optimized implementations, such as FlashAttention and Mamba?

Q.2. Selective mechanisms, such as Mamba and the selective transformer, emphasize empirical advantages. Can these (or similar variations) be incorporated with the 1-S approach?

See additional questions under Weaknesses (e.g., W.5).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720443773

thread_type:
official_review

<thread_object>
creator: reviewer_tk6m

<content>
summary:
1. The paper presents *LightNet*, a new variant of State Space Models (SSMs) incorporating an additive decay/selectivity parameter. Unlike Mamba or Linear Attention (LA) with decay, LightNet avoids recursion-based parallelization, achieving faster parallelization using only matrix multiplications, similar to LA without decay.

2. To address permutation invariance in non-causal settings (a characteristic it shares with LA), the authors introduce two multi-dimensional positional encodings MD-TPE and MD-LRPE.

3. The model is evaluated across diverse tasks, including autoregressive and bidirectional language modeling, image classification and generation.

strengths:
1. The authors identify that SSM-based methods like Mamba and LA with decay suffer from slow GPU performance due to recursion-based parallelization. In contrast vanilla LA is fully parallelizable via matrix multiplications but underperforms other SOTA SSMs.
 
2. To introduce decay to improve LA's performance while retaining full parallelizability, they derive the mathematical formulation of the linear operation induced by a decayed recurrence and using this they motivate an “Additive” decay mechanism, that allows for a decay with matrix multiplication-based parallelization.

3. They evaluate their model across a broad range of tasks and demonstrate near competitive performance against SOTA methods.

weaknesses:
**1. Saturation of Additive Decay for Long Sequences**

The proposed additive decay mechanism, defined as $g(t) = \sum_{s=1}^{t} \delta(s)$, tends to saturate for long sequences. As the decay accumulates $\delta$ values over the course of a long sequence, $g(t)/g(s)$ reaches a point where it can no longer distinguish between nearby tokens effectively. This may be problematic for language modeling as it has a strong local bias.

This saturation effect will likely be more pronounced with very long sequences or on large-scale models. **I am curious how LightNet might perform compared to Mamba in this setting**, which has the opposite issue: difficulty paying attention to very distant tokens despite selectivity. I suspect LightNet may struggle with paying attention locally at scale, potentially leading to underperformance on long sequences.

**2. Comparing LA and LightNet on Long Sequences**

In line with the reasoning above, I feel that LA and LightNet may have similar performance at scale. This is due to LightNet's decay factor saturating to 1. To better understand this, an ablation study comparing LightNet directly with Linear Attention, keeping all other factors constant, would be useful. Specifically, I suggest replacing the additive decay with vanilla LA’s $\phi(K)$ (substituting `softmax(K)` with $\phi(K)$). Additionally, scaling up model sizes could reveal whether the performance gap between LightNet and Linear Attention narrows, providing insight into the effectiveness of additive decay.


**3. It is unclear how much of the heavy lifting is being done by the position embeddings**. 

I believe there is a chance that the position embeddings may be doing a lot of the heavy lifting in getting good performance. For instance Toeplitz matrix in itself is a performant SSM [1] and may artificially boost the apparent efficacy of Additive Decay. Can authors maybe use 2D rope as the position embedding and see if they improve/compete with original LightNet, LightNet with no position embeddings and Linear Attention.

[1] Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers. Sukjun Hwang, Aakash Lahoti, Tri Dao, Albert Gu

questions:
1. In the paper you do claim that you have results on 1B language modeling as well. Could you please share them?
2. Could you also compare against vanilla LA especially on larger scales. This helps prove/disprove my hypothesis.
3. Could the authors please explain how is the method implemented through matrix multiplications in the causal setting? I am not able to figure out the vectorized formula for the same that computes in O(Sequence Length) time.
4. Could the authors please clarify what they mean by "parameter tiling strategy" in Line 266, and if and why this is not implemented for non-causal settings?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732479617408

thread_type:
official_comment

<thread_object>
creator: area_chair_esro

<content>
comment:
Dear Reviewers,

This is a gentle reminder that the authors have submitted their rebuttal, and the discussion period will conclude on November 26th AoE. To ensure a constructive and meaningful discussion, we kindly ask that you review the rebuttal as soon as possible and verify if your questions and comments have been adequately addressed.

We greatly appreciate your time, effort, and thoughtful contributions to this process.

Best regards,
AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1732785001634

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely appreciate all the reviewers for their thorough reviews and constructive feedback. Based on their comments, we have revised the manuscript, with modifications highlighted in **brown** font. Below, we summarize the main changes:

1. **Addressing minor errors and missing citations:** Following the reviewers' suggestions, we corrected typographical errors and added missing citations, making the paper more self-contained.
2. **Improved table visualization:** Following Reviewer 4dh3's suggestion, we enhanced the presentation of table results.
3. **Refined Section 3.1:** Following Reviewer TjCt's suggestion, we revised the recursive definitions of additive decay and multiplicative decay, and provided clearer explanations of global information.
4. **Revised Figure 3:** Following Reviewer TjCt's suggestion, we have corrected Figure 3.
5. **Expanded motivation for LightNet positional encoding:** Following Reviewers tk6m and YtrF's suggestions, we elaborated on the design motivation for LightNet’s positional encoding, reflected in Section 4.2 and Appendix A2.
6. **Added ablation studies:** Following Reviewers tk6m and YtrF's suggestions, we included additional ablation studies on removing additive decay/softmax activation function, detailed in Section 5.3, *Effectiveness of Additive Decay*.
7. **Extended discussion of related work:** Following Reviewer YtrF's suggestion, we added discussions on related works [1], [2].
8. **Comparison with Mamba retrieval capabilities:** Following Reviewer tk6m's suggestion, we included a comparison with Mamba’s retrieval performance, as shown in Figure 5.

Citations:

[1] Mega: Moving Average Equipped Gated Attention. Ma et al. ICLR 2023

[2] Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length. Ma et al.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a new form of additive linear attention to efficiently process multi-dimensional data. Further, the paper also proposes  new forms of positional embeddings for the same. Empirical evaluations show performance improvements over linear attention on many multi-dimensional settings. While the reviewers appreciated the theoretical motivations of the work, there were concerns regarding the clarity of exposition, lack of sufficient ablations and stress testing (eg, in long sequence tasks specifically).