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

## Abstract

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 “multiplicative” linear recurrence and proposes an efficient alternative “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. † † ✉ ✉ {}^{\textrm{{\char 0\relax}}} start_FLOATSUPERSCRIPT ✉ end_FLOATSUPERSCRIPT Indicates corresponding author (Email address: zhongyiran@gmail.com ).

## 1 Introduction

Linear attention has emerged as an effective alternative to softmax attention due to its linear computational complexity and enhanced processing speed, especially in causal language models (; ). The benefits of linear attention largely depend on its decay mechanism (; ; ), which prevents attention dilution () and facilitates global interaction among tokens. However, the decay mechanism presents two primary issues:
First, the decay mechanism is not easily applicable to high-dimensional inputs due to the need for multiple sequential scans to establish a global multi-dimensional receptive field, which reduces computational efficiency (; ).
Additionally, without the decay mechanism, linear attention lacks positional awareness during computations, leading to decreased performance (). In light of these challenges, we are investigating the feasibility of reducing sequential scans for multi-dimensional scenarios while preserving performance.

We first analyze the types of linear recurrence and divide them into two categories: *multiplicative* and *additive*. In multiplicative recurrence, the decay rate is dependent only on the current moment, making it impossible to obtain information about subsequent moments with a single scan. By taking image processing as an example, using multiplicative recurrence will require at least two scans to retrieve the global information (; ). Conversely, in additive recurrence, the decay rate depends on all moments through the summation of the importance score of each moment, enabling it to gather global information in a single scan.

It is important to note that in non-causal situations, additive recurrence is permutation-invariant, which means it lacks local precedence and therefore diminishes the capture of positional information. To overcome this limitation, we put forth a new approach to positional encoding called Multi-Dimensional Toeplitz Positional Encoding (MD-TPE). This method utilizes the mathematical properties of the Toeplitz matrix to embed relative positional information with linear time complexity, thus ensuring efficiency in multi-dimensional scenarios. Additionally, we expand the Linearized Relative Positional Encoding (LRPE) () to high-dimensional scenarios, resulting in the creation of Multi-Dimensional Linearized Relative Positional Encoding (MD-LRPE).

We then present LightNet, a new multi-dimensional linear attention model built on additive recurrence. LightNet features a pioneering decay mechanism, allowing for efficient single-scan processing of high-dimensional sequential data. Furthermore, it integrates highly effective multi-dimensional position encoding such as MD-TPE and MD-LRPE to precisely capture positional information.

We conduct several evaluations of the performance of our proposed LightNet on a range of tasks, including image generation, image classification, image generation, bidirectional language modeling, and autoregressive language modeling. LightNet performs comparably or better than its competitors across all tasks.

We summarize our main contributions as follows:

- •
We analyze the types of linear recurrence, dividing them into two types: *multiplicative* and *additive*, where the additive type can obtain global information in a single scan.
- •
We propose two multi-dimensional position encoding strategies, MD-TPE and MD-LRPE, to effectively capture positional information in multi-dimensional scenarios.
- •
We propose LightNet, a new multi-dimensional linear attention model that can process high-dimensional sequences in a single scan.
- •
We conduct thorough evaluations to assess the efficiency and efficacy of LightNet for multi-dimensional sequential modeling tasks. The LightNet demonstrates competitive performance in all scenarios.

## 2 Related Work

#### Linear Attention.

The linear attention mechanism has greatly advanced deep learning, particularly in natural language processing, by providing a scalable solution for long input sequences and reducing the computational demands of traditional attention models ; ; ; . However, despite its faster training speeds, linear attention’s performance still falls short of softmax attention due to the attention dilution issue (). The TNL/RetNet ; introduces a decay mechanism to address this problem. Additionally, GLA  incorporating gating mechanisms show the potential to enhance linear attention models.

#### State Space Model.

State Space Models (SSMs) are increasingly crucial in sequence modeling due to their structured approach to capturing temporal dynamics through latent variables.
The S4 model  enhances state space modeling for long sequences by leveraging structured spaces to improve computational efficiency and tackle complex dynamics.
With additional parameterizing and initializing diagonal state space strategy , the SSMs can achieve comparable performance to naive transformers.
Furthermore, the Gated State Space (GSS) model  introduces a gating mechanism to SSMs, which is particularly effective for long-range language modeling by allowing nuanced control over information flow.
The S5 model  reduces complexity using “scan” while maintaining the capability to handle intricate sequences.
However, directly extending the SSM to multi-dimensional input usually requires multiple sequential scans, which will reduce the computational efficiency ().

#### Linear RNN.

Linear RNNs employ element-wise recursion for sequence modeling, and due to their linear recursive form, they can be accelerated using parallel scans ().
At their core is the decay mechanism, where RWKV-4/LRU ; utilizes data-independent decay.
HGRN ; leverage data-dependent decay to enhance performance.
Linear RNNs have shown considerable potential in language modeling and long-sequence modeling tasks.

#### Multi-dimensional Tasks with Linear Complexity Model.

The development of linear attention in language models has led to its extension into multi-dimensional tasks. Building upon the cosFormer framework , VVT  explores a local prior of 2D linear attention and applies it to image classification tasks.
Vim  and Vision-RWKV  utilize a sequential scan mechanism to expand Mamba  and RWKV  for image classification.
Additionally, leveraging the benefits structure of the diffusion transformer  in image generation, several works have extended linear complexity models into 2D space ; ; ; to replace the traditional transformer architecture, achieving efficient image generation.
However, some of these tasks encounter issues with inadequate performance.
Moreover, frequent sequential scans can compromise the efficiency of the model.

## 3 Linear Recurrence in Multi-dimensional Space

In this section, we discuss the theoretical and practical computational complexity of linear recurrence (with decay) when dealing with high-dimensional data, and then analyze the types of linear recurrence. In subsequent discussions, we assume $n$ is the sequence length, $d$ is the embedding dimension, and $\mathbf{x}_{t}\in\mathbb{R}^{d}$ is the transpose of the $t$-th row of matrix $\mathbf{X}\in\mathbb{R}^{n\times d}$.

Figure: Figure 1: Processing time of 1 Scan and 2 Scan in relation to sequence length. 1 Scan is significantly faster than 2 Scan in both forward and backward passes. As the sequence length increases, the advantage of 1 Scan becomes more substantial.
Refer to caption: x1.png

### 3.1 Computational Complexity of Linear Recurrence

It has been proved that all linear attention, state space model, and linear RNN can be expressed using a linear recurrence formula (). We use linear attention with decay (; ) as an example. Below is the recursive form (*i.e.* a scan):

$$ $\mathbf{k}\mathbf{v}_{0}=\mathbf{0},\mathbf{k}\mathbf{v}_{t}=\lambda_{t} \mathbf{k}\mathbf{v}_{t-1}+\mathbf{k}_{t}\mathbf{v}_{t}^{\top},\mathbf{o}_{t}^ {\top}=\mathbf{q}_{t}^{\top}\mathbf{k}\mathbf{v}_{\mathbf{t}},t=1,\ldots,n.$ $$

Here, $0<\lambda_{t}\leq 1$ is the decay rate. Note that the above formula is for the causal scenario. When dealing with non-causal scenarios, a common practice in the literature is to perform causal computation twice (; ). We call this method “2 scan”:

$$ $\begin{gathered}\overset{\rightarrow}{\mathbf{kv}}_{0}=\mathbf{0},\overset{ \rightarrow}{\mathbf{kv}}_{t}=\lambda_{t}\overset{\rightarrow}{\mathbf{kv}}_{t -1}+\mathbf{k}_{t}\mathbf{v}_{t}^{\top},\overset{\rightarrow}{\mathbf{o}}_{t}^ {\top}=\mathbf{q}_{t}^{\top}\overset{\rightarrow}{\mathbf{kv}}_{t},\\ \overset{\leftarrow}{\mathbf{kv}}_{n+1}=\mathbf{0},\overset{\leftarrow}{ \mathbf{kv}}_{t}=\lambda_{t}\overset{\leftarrow}{\mathbf{kv}}_{t+1}+\mathbf{k} _{t}\mathbf{v}_{t}^{\top},\overset{\leftarrow}{\mathbf{o}}_{t}^{\top}=\mathbf{ q}_{t}^{\top}\overset{\leftarrow}{\mathbf{kv}}_{t},\\ \mathbf{o}_{t}=\overset{\rightarrow}{\mathbf{o}}_{t}+\overset{\leftarrow}{ \mathbf{o}}_{t}.\end{gathered}$ $$

When $\lambda_{t}=1$, *i.e.* there is no decay, the right product trick () can be used in this case. We call this method “1 scan”.

$$ $[\mathbf{KV}]=\mathbf{K}^{\top}\mathbf{V},\mathbf{O}=\mathbf{Q}[\mathbf{KV}].$ $$

Although both of the above formulas have a time complexity of $O(nd^{2})$, the “2 scan” version is significantly slower than the “1 scan” version. This is because causal computation requires block-level recursion (; ), whereas the second formula can be fully parallelized due to matrix multiplication (). We provide a speed comparison in Fig. 1, where the “2 scan” is implemented with Lightning Attention , the fastest linear attention implementation so far. It can be seen that the “2 scan” is several times slower than the “1 scan” in both forward and backward passes.

It is apparent that the need for multiple scans is mainly due to the presence of decay $\lambda_{t}$.
However, directly removing $\lambda_{t}$ would lead to degraded performance .
A natural question arises: can we retain $\lambda_{t}$ while only performing a single scan?
In the next section, we will discuss the types of linear recurrence and answer the question.

### 3.2 Types of Linear Recurrence

We first explore the representation range of linear recurrences by 1D linear recurrence:

$$ $y_{t}=a_{t}y_{t-1}+x_{t},y_{0}=0.$ (1) $$

Unroll the recursion equation of E.q 1, we obtain:

$$ $y_{t}=\sum_{s=1}^{t}\frac{A_{s}}{A_{t}}x_{s}\triangleq\sum_{s=1}^{t}c_{ts}x_{s },A_{t}=\left(\prod_{s=1}^{t}a_{s}\right)^{-1}.$ (2) $$

The detailed proof of the unrolling process can be found in Appendix A.1.
Note that $y_{t}$ is a linear combination of $x_{1},\ldots,x_{t}$.
A natural question arises: Can every linear combination $\sum_{s=1}^{t}c_{ts}x_{s}$ be represented as a linear recursion?
We now prove that a linear recursion representation is possible only when the coefficients $c_{ts}$ satisfy certain conditions.

###### Theorem 3.1 .

A linear recurrence $y_{t}=a_{t}y_{t-1}+x_{t},y_{0}=0$ is equivalent to a linear combination $y_{t}=\sum_{s=1}^{t}c_{ts}x_{s}$, iff $c_{ts}=\frac{g(s)}{g(t)}$, where $g(\cdot)$ is a function.

###### Proof of Theorem 3.1 .

$\Rightarrow$

Given a linear recurrence, we multiply it by $A_{t}=\left(\prod_{s=1}^{t}a_{s}\right)^{-1}$ and the following recurrence equation:

$$ $\displaystyle A_{t}{y_{t}}=A_{t}a_{t}{y_{t-1}}+A_{t}{x_{t}}$ $\displaystyle=A_{t-1}{y_{t-1}}+A_{t}x_{t}.$ $$

Unroll it, we get:

$$ ${A_{t}}{y_{t}}-{A_{t-1}}{y_{t-1}}={A_{t}}{x_{t}},\ldots,{A_{2}}{y_{2}}-{A_{1}} {y_{1}}={A_{2}}{x_{2}}.$ (3) $$

To derive an expression for $y_{t}$, we sum the recursive equations and obtain:

$$ ${A_{t}}{y_{t}}-{A_{1}}{y_{1}}=\sum_{s=2}^{t}{A_{s}}c{x_{s}},\\ {y_{t}}{A_{t}}=\sum_{s=1}^{t}{A_{s}}{x_{s}},y_{t}=\sum_{s=1}^{t}\frac{A_{s}}{A _{t}}x_{s}.$ (4) $$

By comparing the coefficients, we can obtain $c_{ts}={A_{s}}/{A_{t}}$.

$\Leftarrow$:

Given the linear combination $y_{t}=\sum_{s=1}^{t}c_{ts}x_{s}$ and $c_{ts}=\frac{g(s)}{g(t)}$, we define $a_{t}\triangleq\frac{g(t-1)}{g(t)}$. Then $y_{t}$ can be expressed as:

$$ $\displaystyle y_{t}$ $\displaystyle=\sum_{s=1}^{t}c_{ts}x_{s}=\sum_{s=1}^{t-1}c_{ts}x_{s}+c_{tt}x_{t }=\sum_{s=1}^{t-1}\frac{g(s)}{g(t)}x_{s}+\frac{g(t)}{g(t)}x_{t}$ $\displaystyle=\frac{g(t-1)}{g(t)}\sum_{s=1}^{t-1}\frac{g(s)}{g(t-1)}x_{s}+x_{t }=a_{t}\sum_{s=1}^{t-1}c_{t-1,s}x_{s}+x_{t}=a_{t}y_{t-1}+x_{t}.\qed$ $$

Based on the Theorem 3.1, for linear recurrence, we can directly discuss $g(t)$, as $a_{t}$ can be obtained through $\frac{g(t-1)}{g(t)}$. Intuitively, $g(t)$ can be interpreted as an importance score up to moment $t$, $c_{ts}=\frac{g(s)}{g(t)}$ can be interpreted as the ratio of the score at moment $s$ relative to moment $t$, and $a_{t}$ can be interpreted as the ratio of the previous moment’s score to moment $t$’s score.

Typically, to prevent numerical overflow, we assume $0\leq a_{t}=\frac{g(t-1)}{g(t)}\leq 1$. To meet this condition, we present the following two forms:

###### Proposition 3.2 .

For Linear Recurrence with $0\leq a_{t}\leq 1$, there are two forms:

1. Multiplicative: $g(t)=\frac{1}{\prod_{s=1}^{t}\rho(s)},a_{t}=\rho(t),0\leq\rho(t)\leq 1$;

2. Additive: $g(t)=\sum_{s=1}^{t}\delta(s),a_{t}=\frac{\sum_{s=1}^{t-1}\delta(s)}{\sum_{s=1}
^{t}\delta(s)},\delta(s)\geq 0$.

###### Proof of Proposition 3.1 .

The condition $0\leq\frac{g(t-1)}{g(t)}\leq 1$, is equivalent to $\rho(t)=\frac{g(t-1)}{g(t)}\leq 1$ or $\delta(t)=g(t)-g(t-1)\geq 0$.
By expanding the first formula, we obtain the multiplicative type. By expanding the second formula, we get the additive type.
∎

It can be observed that the typical linear attention with decay corresponds to the Multiplicative form, where $\rho(t)$ is utilized as $\text{Sigmoid}(\cdot)$ (), $\text{exp}(-\text{exp}(\cdot))$ (), or a fixed value (; ). Multiplicative requires multiple scans when dealing with high dimensions because $a_{t}$ itself cannot provide global information (as $a_{t}$ is only related to the moment $t$). However, for the Additive method, since the computation form is $a_{t}=\frac{\sum_{s=1}^{t-1}\delta(s)}{\sum_{s=1}^{t}\delta(s)}$, by modifying the denominator to $\Delta=\sum_{s=1}^{n}\delta(s)$ ($n$ is the sequence length), global information can be obtained through $a_{t}=\frac{\sum_{s=1}^{t-1}\delta(s)}{\Delta}$.

###### Theorem 3.1 .

###### Proof of Theorem 3.1 .

###### Proposition 3.2 .

###### Proof of Proposition 3.1 .

## 4 LightNet

Building upon the preceding analysis, we introduce a novel Linear Transformer architecture termed LightNet, designed to handle multi-dimensional data efficiently in a single scan.
An overview of its structure is depicted in Fig. 2.
LightNet comprises an Input Embedding, MD-TPE module, and several stacked LightNet Layers.

Figure: Figure 2: The network structure of LightNet: each LightNet model is comprised of an Input Embedding, MD-TPE, and a stack of multiple LightNet Layers. Each LightNet Layer consists of an LNA and a GLU, with the computation of LNA illustrated in the figure on the right.
Refer to caption: x3.png

### 4.1 LightNet Layer

The LightNet Layer is composed of a LightNet Attention (LNA) and a Gated Linear Unit (GLU) (). Within the LNA, an additive decay is employed, with the parameter $\delta$ implemented through the exponential function. Additionally, a parameter tiling strategy is utilized for both the key and decay, which has been empirically observed to enhance performance. This empirical evidence is detailed in Table 3. Furthermore, the integration of a low-rank output gate from TNL3 () and a normalization after linear attention () has been incorporated.

In causal settings, the LightNet Layer can be represented as follows:

$$ $\displaystyle\mathbf{s}_{t}=\mathbf{s}_{t-1}+\exp(\mathbf{k}_{t}),\mathbf{\bar {k}_{t}}$ $\displaystyle=\exp(\mathbf{k}_{t})/\mathbf{s}_{t},\mathbf{kv}_{t}=\mathrm{diag }\{1-\mathbf{\bar{k}_{t}}\}\mathbf{kv}_{t-1}+\mathbf{\bar{k}_{t}}\mathbf{v}_{t }^{\top},$ (5) $\displaystyle\mathbf{o}_{t}^{\top}$ $\displaystyle=\mathrm{Norm}[\mathbf{kv}_{t}^{\top}\phi(\mathbf{q}_{t})]\odot \psi(u_{t}).$ $$

In non-causal settings, the expression becomes:

$$ $\begin{gathered}\mathbf{s}=\sum_{t}\exp(\mathbf{k}_{t}),\mathbf{o}_{t}=\mathrm {Norm}\left[\phi(q_{t})\sum_{t}(\exp(k_{t})/s)^{\top}\mathbf{v}_{t}\right] \odot\psi(u_{t}),\\ \mathbf{O}=\mathrm{Norm}\left[\phi(\mathbf{Q})(\mathrm{Softmax}(\mathbf{K})^{ \top}\mathbf{V})\right]\odot\psi(\mathbf{U}).\end{gathered}$ (6) $$

where $\mathbf{X}$ is the input of LNA:

$$ $\mathbf{Q}=\mathbf{X}\mathbf{W}_{q},\mathbf{K}=\mathbf{X}\mathbf{W}_{k}, \mathbf{V}=\mathbf{X}\mathbf{W}_{v},\mathbf{U}=\mathbf{X}\mathbf{W}_{u1} \mathbf{W}_{u2},\phi=\mathrm{Swish},\psi=\mathrm{Sigmoid}.$ (7) $$

### 4.2 Multi Dimension Position Encoding

It is noted that additive recurrence does not have a locality prior like multiplicative recurrence and is permutation invariant in non-causal scenarios, as shown in E.q 6. Therefore, it is necessary to introduce new positional encoding.
To tackle this challenge, we introduce two novel relative positional encoding methods, MD-TPE (Multi-Dimensional Toeplitz Positional Encoding) and expand the existing LRPE () to the high-dimensional context as MD-LRPE (Multi-Dimensional Linearized Relative Positional Encoding). This enhancement enables the management of relative positional relationships in any dimension.

#### MD-TPE.

Given multi-dimension input $\mathbf{x}_{n_{1},\ldots,n_{k}},1\leq n_{s}\leq N_{s},s=1,\ldots k$, we use the following equation to capture positional information:

$$ $\displaystyle\mathbf{y}_{n_{1},\ldots,n_{k}}=\sum_{m_{k}\leq n_{k}}\ldots\sum_ {m_{1}\leq n_{1}}\mathbf{y}_{n_{1}-m_{1},\ldots,n_{k}-m_{k}}\mathbf{x}_{m_{1}, \ldots,m_{k}}.$ (8) $$

However, the time complexity of implementing the aforementioned method is $O(N\log N)$,where $N=\prod_{s=1}^{k}n_{s}$, making it inefficient.
To address this, we simplified the above formula by performing toeplitz matrix production for each dimension separately and using SSM for parameterization (), we denotes $e$ as the hidden dimension of SSM below:

$$ $\displaystyle\mathbf{y}_{n_{1},\ldots,n_{k}}=\sum_{s=1}^{k}\sum_{m_{s}=1}^{n_{ s}}\mathbf{y}_{n_{s}-m_{s}}\mathbf{x}_{n_{1},\ldots,m_{s},\ldots,n_{k}}=\sum_{ s=1}^{k}\sum_{m_{s}=1}^{n_{s}}\sum_{t=1}^{e}\lambda_{t}^{n_{s}-m_{s}}\mathbf{x }_{n_{1},\ldots,m_{s},\ldots,n_{k}}.$ (9) $$

By using a scan approach, the above calculation becomes linear in complexity, $O(Ne)$.

#### MD-LRPE.

Given $\mathbf{x}_{t}\in\mathbb{R}^{d},\mathbf{x}\in\{\mathbf{q},\mathbf{k}\}$, LRPE transforms it through the matrix $\mathbf{W}_{t}$ to $\mathbf{W}_{t}\mathbf{x}_{t},\mathbf{x}\in\{\mathbf{q},\mathbf{k}\}$, and it holds that:

$$ $\displaystyle(\mathbf{W}_{s}\mathbf{q}_{s})^{\mathrm{H}}(\mathbf{W}_{t}\mathbf {k}_{t})=\mathbf{q}_{s}^{\mathrm{H}}\mathbf{W}_{s}^{\mathrm{H}}\mathbf{W}_{t} \mathbf{k}_{t}=\mathbf{q}_{s}^{\mathrm{H}}\mathbf{W}_{t-s}\mathbf{k}_{t}.$ (10) $$

We choose the complex version of LRPE, where:

$$ $\displaystyle\mathbf{W}_{t}=\mathrm{diag}\{\exp(it\theta_{1}),\ldots,\exp(it \theta_{d})\}.$ (11) $$

To generalize to higher dimensions, *i.e.*, given $\mathbf{x}_{n_{1},\ldots,n_{k}}\in\mathbb{R}^{d},\mathbf{x}\in\{\mathbf{q},
\mathbf{k}\}$, we divide the $d$ features into $k$ groups, each group has $d/k$ features, with the $s$-th group’s features corresponding to dimension $n_{s}$, $s\in[1,k]$. Specifically, we define:

$$ $\displaystyle\mathbf{W}_{n_{1},\ldots,n_{k}}=\mathrm{diag}\{[\Theta_{1},\ldots ,\Theta_{k}]\},\Theta_{s}=\exp(in_{k}\theta_{j}),sd/k<j\leq(s+1)d/k,\theta_{j} =10000^{-2j/d}.$ (12) $$

Thus:

$$ $\displaystyle\mathbf{W}_{n_{1},\ldots,n_{k}}^{\mathbf{H}}\mathbf{W}_{m_{1}, \ldots,m_{k}}=\mathbf{W}_{m_{1}-n_{1},\ldots,m_{k}-n_{k}}$ (13) $$

Then:

$$ $\begin{gathered}(\mathbf{W}_{n_{1},\ldots,n_{k}}\mathbf{q}_{n_{1},\ldots,n_{k} })^{\mathrm{H}}(\mathbf{W}_{m_{1},\ldots,m_{k}}\mathbf{k}_{m_{1},\ldots,m_{k}} )=\mathbf{q}_{n_{1},\ldots,n_{k}}^{\mathrm{H}}\mathbf{W}_{s}^{\mathrm{H}} \mathbf{W}_{t}\mathbf{k}_{m_{1},\ldots,m_{k}}\\ =\mathbf{q}_{n_{1},\ldots,n_{k}}^{\mathrm{H}}\mathbf{W}_{m_{1}-n_{1},\ldots,m_ {k}-n_{k}}\mathbf{k}_{m_{1},\ldots,m_{k}}.\end{gathered}$ (14) $$

## 5 Experiments

We comprehensively evaluate the substitutability of our LightNet in performance, scalability, flexibility, and efficiency.
We validate the effectiveness of our model on various multi-dimensional sequential modeling tasks.
We also test the proposed ability of LightNet to serve as a language model.

**Table 1: Performance comparison for image classification task on ImageNet1k. “S.A.” represents Softmax Attention, “M.S.” denotes multiple scans, and “O.S.” signifies one scan.**
| Model | Category | Tiny | Small | Base |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Acc (%) $\uparrow$ | Params (M) | Acc (%) $\uparrow$ | Params (M) | Acc (%) $\uparrow$ | Params (M) |  |  |
| DeiT | S.A. | 72.20 | 5.7 | 79.90 | 22.00 | 81.80 | 86.00 |
| HGRN | M.S. | 74.40 | 6.1 | 80.09 | 23.70 | - | - |
| Vim | M.S. | 76.10 | 7.0 | 80.50 | 26.00 | - | - |
| V-RWKV | M.S. | 75.10 | 6.2 | 80.10 | 23.80 | 82.00 | 93.70 |
| LightNet | O.S. | 74.46 | 6.0 | 80.12 | 22.64 | 81.90 | 87.74 |
| LightNet w/o TPE | O.S. | 73.97 | 6.0 | 79.65 | 22.54 | 81.45 | 87.54 |
| LightNet w/o LRPE | O.S. | 74.02 | 6.0 | 79.54 | 22.63 | 81.72 | 87.69 |

### 5.1 Setting

**Table 2: Performance comparison for image generation task on ImageNet-1k. LightNet-XL/2 achieves state-of-the-art FID with or without classifier-free guidance (-G) .**
| Model | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Precision$\uparrow$ | Recall$\uparrow$ |
| --- | --- | --- | --- | --- | --- |
| CDM | 4.88 | - | 158.71 | - | - |
| LDM-8 | 15.51 | - | 79.03 | 0.65 | 0.63 |
| LDM-8-G | 7.76 | - | 209.52 | 0.84 | 0.35 |
| LDM-4 | 10.56 | - | 103.49 | 0.71 | 0.62 |
| LDM-4-G | 3.60 | - | 247.67 | 0.87 | 0.48 |
| DiT-XL/2 | 9.62 | 6.85 | 121.50 | 0.67 | 0.67 |
| DiT-XL/2-G | 2.27 | 4.60 | 278.24 | 0.83 | 0.57 |
| LightNet-XL/2 | 5.35 | 5.93 | 171.18 | 0.73 | 0.65 |
| LightNet-XL/2-G | 2.18 | 4.58 | 281.85 | 0.83 | 0.58 |

#### Image Classification.

We trained our LightNet model for image classification on the ImageNet-1K dataset . Our approach modifies the network architecture and training protocols of DeiT , substituting its Transformer Layers with our proprietary LightNet Layers.

#### Image Generation.

We build our model upon the latent diffusion model ; and use our proposed LightNet as the denoising network.
We adjust the model size across various configurations (S, B, L, XL) and patch sizes (8, 4, 2), consistent with DiT .
Experiments are conducted on the ImageNet dataset  at a resolution of $256\times 256$.
Each model is trained over 0.4M steps with a batch size of 256 to assess scaling capabilities.
For the largest model variant, training is extended to 0.8M steps with a batch size of 1024, as opposed to the 7M steps in DiT, to enhance generative performance.

**Table 3: Performance comparison on Wikitext-103. $\downarrow$ means lower is better. We adopted the configuration of HGRN for Wikitext-103, and we can observe that LightNet significantly outperforms all other methods.**
| Model | PPL<br>(val) $\downarrow$ | PPL<br>(test) $\downarrow$ | Params<br>(M) |
| --- | --- | --- | --- |
| Attn-based |  |  |  |
| Transformer | 24.40 | 24.78 | 44.65 |
| FLASH | 25.92 | 26.70 | 42.17 |
| 1+elu | 27.44 | 28.05 | 44.65 |
| Performer | 62.50 | 63.16 | 44.65 |
| cosFormer | 26.53 | 27.06 | 44.65 |
| RNN-based |  |  |  |
| S4 | 38.34 | 39.66 | 45.69 |
| DSS | 39.39 | 41.07 | 45.73 |
| GSS | 29.61 | 30.74 | 43.84 |
| RWKV-4 | 24.31 | 25.07 | 46.23 |
| LRU | 29.86 | 31.12 | 46.24 |
| HGRN | 24.14 | 24.82 | 46.25 |
| FFT-based |  |  |  |
| TNN | 23.98 | 24.67 | 48.68 |
| LightNet | 23.09 | 23.75 | 45.07 |

#### Bidirectional Language Modeling.

We utilize Cramming-BERT () as our pipeline, employing a 24-hour training regime to pre-train on the Pile dataset, subsequently finetuning on the GLUE benchmark ().
During pre-training, we follow established guidelines by setting a learning rate of 1e-3, a sequence length of 128, and a batch size of 8192.
In the finetuning phase, we experiment with learning rates from the set {5e-5, 4e-5, 3e-5, 2e-5} and determine the optimal outcome by finetuning over 5 epochs.

#### Autoregressive Language Modeling.

We evaluate two capabilities: perplexity (PPL) and zero-shot reasoning ability.
The perplexity of the 44M model is assessed on the Wikitext-103 dataset , and the 380M model’s perplexity is tested on the Pile dataset, consuming 10 billion tokens .
For large language model experiments, we train LightNet models at scales of 150M, 350M, 1B, and 3B using 100 billion tokens sampled from subsets of the Pile ().
These models are then evaluated on commonsense reasoning tasks using the lm-eval-harness (). Detailed training hyperparameters are listed in Table 9.

### 5.2 Results

#### Image Classification.

As shown in Table 1, the proposed LightNet shows competitive performance on the ImageNet-1k dataset.
It can be observed that using only a single sequential scan, LightNet can achieve comparable performance to models with naive attention and multiple sequential scans.
At the same time, the speed advantage of using a single sequential scan is shown in Fig. 1.

#### Image Generation.

The image generation results are presented in Table 2.
Our proposed LightNet demonstrates superior performance, achieving a lower Fréchet Inception Distance (FID) and a higher Inception Score (IS) than DiT  with fewer training steps (0.8M steps vs 7M steps).
Additionally, LightNet exhibits commendable scaling capabilities, as illustrated in Fig. 3.

**Table 4: Performance comparison Pile for large-scale language modeling.. We trained under the 10 billion token subset of Pile, and it can be seen that LightNet’s PPL is better than LLaMA’s.**
| Model | PPL $\downarrow$ | Params |
| --- | --- | --- |
| LLaMA | 4.62 | 385M |
| TNL | 4.62 | 379M |
| Mamba | 4.59 | 385M |
| LightNet | 4.59 | 379M |
| LightNet w/o TPE | 4.69 | 379M |
| LightNet w/o LRPE | 4.69 | 379M |
| LightNet no share | 4.76 | 385M |

#### Bidirectional Language Modeling.

As shown in Table 5, LightNet outperforms Crammed Bert () on the GLUE dataset, demonstrating its superior capability in handling natural language understanding tasks. Despite BERT-Base () achieving comparable performance, it is noteworthy that LightNet does so with a significantly lower computational cost, having been trained on a single A100 for 24 hours.

#### Autoregressive Language Modeling.

In the Wikitext-103 dataset, as depicted in Table 3, LightNet surpasses all competitors on both the validation and test datasets. Regarding large-scale datasets, as illustrated in Table 4, LightNet exhibits superior perplexity (PPL) compared to LLaMA () and TNL (), and matches the performance of Mamba (). The ability of LightNet to achieve high performance with reduced parameter complexity underscores its potential for scalability and broader application across various large-scale data scenarios.

Figure: Figure 3: Scaling up the LightNet enhances the FID during every stages of training. We present the FID-50K across training iterations for twelve LightNet models. Enhancing the LightNet backbone results in improved generative models for all sizes of models and patches.
Refer to caption: x4.png

**Table 5: Performance Scores on GLUE Benchmark. We utilize the Cramming-BERT 24-hour training configuration and observe that LightNet outperforms Crammed BERT and achieves comparable results to BERT-Base, which is trained with more GPU hours.**
| Model | MNLI | SST-2 | STSB | RTE | QNLI | QQP | MRPC | CoLA | GLUE |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| BERT-Base (Fully trained) | 83.2 / 83.4 | 91.9 | 86.7 | 59.2 | 90.6 | 87.7 | 89.3 | 56.5 | 80.9 |
| BERT-Base (No Pretrain) | 34.1 / 34.1 | 79.9 | 17.8 | 47.3 | 50.0 | 68.6 | 77.9 | - | 45.5 |
| Crammed BERT | 83.9 / 84.1 | 92.2 | 84.6 | 53.8 | 89.5 | 87.3 | 87.5 | 44.5 | 78.6 |
| LightNet | 83.3 / 83.5 | 92.9 | 86.3 | 55.6 | 89.1 | 87.7 | 88.5 | 52.6 | 79.9 |
| LightNet w/o TPE | 82.1 / 82.9 | 92.4 | 79.4 | 57.8 | 89.2 | 87.7 | 83.8 | 44.1 | 77.7 |

### 5.3 Ablation Studies

**Table 6: Ablation studies on image generation for LightNet-B/2 Configurations. We compare the performance of FID under different training steps.**
| Model | 50K | 100K | 150K | 200K | 250K | 300K | 350K | 400K |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LightNet-B/2 | 104.19 | 74.27 | 59.60 | 51.22 | 45.70 | 41.65 | 38.60 | 36.45 |
| LightNet-B/2 w/o TPE | 105.86 | 77.64 | 64.81 | 57.24 | 51.98 | 48.12 | 44.90 | 42.74 |
| LightNet-B/2 w/o LRPE | 132.17 | 82.99 | 67.79 | 59.02 | 52.88 | 48.41 | 44.94 | 42.37 |

Effectiveness of Parameters Sharing.
As discussed in Sec. 4.2, we employ a parameter sharing strategy between decay and key, and the performance comparison is presented in Table 4.
The results demonstrate that employing independent parameters for decay and key leads to performance deterioration, highlighting the significance of parameter sharing.

Effectiveness of MD-TPE.
The proposed MD-TPE provides relative positional information under linear complexity.
We thus explore the effectiveness of the MD-TPE across all tasks, shown in Table 1,4,5,6.
We can observe that removing MD-TPE results in significant performance degradation, particularly for image generation, which highly depends on the relative position of the image content.
Similarly, performance comparison in language modeling tasks also confirms the effectiveness of MD-TPE when reduced to a single dimension.

Effectiveness of MD-LRPE.
LRPE has already proven its effectiveness in the field of language modeling. Therefore, when faced with higher-dimensional inputs, the contributions of its extension, MD-LRPE should be systematically validated. To this end, we conduct numerous ablation experiments, and the results, as shown in Table 1,4,5,6, demonstrate the effectiveness of extending LRPE into a multi-dimensional space through MD-LRPE operation.

Speed Test.
The current linear complexity models employ multiplicative linear recurrence in sequence modeling and necessitate at least two scans for multi-dimensional data, resulting in processing time denoted by the “2 Scan” in Fig. 1. In contrast, our LightNet requires only a single scan, leading to a processing time denoted by the “1 Scan”. As evident from the figure, the advantage of the “1 Scan” becomes increasingly pronounced with the growth of sequence length.

## 6 Conclusion

In this paper, we have addressed the inefficiency of "multiplicative" linear recurrence in multi-dimensional sequence modeling by introducing a novel "additive" linear recurrence that handles multi-dimensional data within a single scan. We developed LightNet, a new multi-dimensional linear attention model enhanced by two new multi-dimensional linear relative positional encoding methods, MD-TPE and MD-LRPE. Empirical evaluations across tasks like image classification, image generation, bidirectional language modeling, and autoregressive language modeling demonstrate LightNet’s superior performance and versatility. LightNet offers a significant advancement in efficiency and scalability, providing a promising pathway for future research and applications in multi-dimensional sequence modeling.

Limitations.
Our empirical evaluation of LightNet is conducted on a smaller scale compared to other large-scale models. Potential negative social consequences include the misuse of brain models for inappropriate purposes or applications, necessitating prohibition through appropriate regulations.

## 7 Acknowledgments

This research was supported in part by the National Natural Science Foundation of China (62271410), the National Key R&D Program of China (NO.2022ZD0160100), and the Fundamental Research Funds for the Central Universities. Yuxin Mao is sponsored by the Innovation Foundation for Doctor Dissertation of Northwestern Polytechnical University (CX2024014).

## Appendix A Appendix

### A.1 Proof of Eq 2

Note that

$$ $\begin{gathered}A_{t}{y_{t}}=A_{t}a_{t}{y_{t-1}}+A_{t}{x_{t}}=A_{t-1}{y_{t-1}} +A_{t}x_{t},\\ {A_{t}}{y_{t}}-{A_{t-1}}{y_{t-1}}={A_{t}}{x_{t}},\\ \ldots,\\ {A_{2}}{y_{2}}-{A_{1}}{y_{1}}={A_{2}}{x_{2}}.\end{gathered}$ $$

By summing up, we can obtain:

$$ ${A_{t}}{y_{t}}-{A_{1}}{y_{1}}=\sum_{s=2}^{t}{A_{s}}c{x_{s}},\\ {y_{t}}{A_{t}}=\sum_{s=1}^{t}{A_{s}}{x_{s}},\\ y_{t}=\sum_{s=1}^{t}\frac{A_{s}}{A_{t}}x_{s}.$ $$

### A.2 More experiments

In this section, we provide additional experimental results. In Table 7, we show the performance of LightNet under the Commonsense Reasoning Tasks. In Table 8, we present the effects of LightNet on image generation tasks across various sizes.

**Table 7: Performance Comparison on Commonsense Reasoning Tasks. PS, T, HS, WG stand for parameter size (billion), tokens (billion), HellaSwag, and WinoGrande, respectively.**
| Model | P | T | PIQA | HS | WG | ARC-e | ARC-c | OBQA | AVG |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| GPT-Neo | 0.13 | 300 | 63.06 | 30.40 | 50.43 | 43.73 | 23.12 | 26.20 | 39.49 |
| OPT | 0.16 | 300 | 62.95 | 31.35 | 50.43 | 43.52 | 22.70 | 28.00 | 39.83 |
| Pythia | 0.16 | 300 | 61.32 | 30.16 | 51.93 | 43.18 | 23.12 | 26.80 | 39.42 |
| RWKV-4 | 0.17 | - | 65.07 | 32.26 | 50.83 | 47.47 | 24.15 | 29.60 | 41.56 |
| HGRN | 0.15 | 100 | 65.02 | 33.33 | 50.20 | 46.68 | 23.81 | 28.60 | 41.27 |
| LightNet | 0.15 | 100 | 63.49 | 30.55 | 50.83 | 46.04 | 24.40 | 27.40 | 40.45 |
| OPT | 0.35 | 300 | 64.58 | 36.69 | 52.49 | 44.02 | 23.89 | 28.20 | 41.65 |
| Pythia | 0.40 | 300 | 67.08 | 40.52 | 53.59 | 51.81 | 24.15 | 29.40 | 44.43 |
| BLOOM | 0.56 | 350 | 64.09 | 36.97 | 52.80 | 47.35 | 23.98 | 28.20 | 42.23 |
| RWKV-4 | 0.43 | - | 67.52 | 40.90 | 51.14 | 52.86 | 25.17 | 32.40 | 45.00 |
| HGRN | 0.35 | 100 | 66.70 | 38.12 | 51.70 | 49.20 | 25.26 | 30.60 | 43.60 |
| LightNet | 0.39 | 100 | 66.87 | 38.82 | 50.75 | 51.39 | 25.17 | 28.20 | 43.53 |
| GPT-Neo | 1.3 | 300 | 71.11 | 48.93 | 54.93 | 56.19 | 25.85 | 33.60 | 48.44 |
| OPT | 1.3 | 300 | 71.71 | 53.70 | 59.35 | 57.24 | 29.69 | 33.20 | 50.82 |
| Pythia | 1.4 | 300 | 70.67 | 47.18 | 53.51 | 56.99 | 26.88 | 31.40 | 47.77 |
| BLOOM | 1.1 | 350 | 67.14 | 42.98 | 54.93 | 51.47 | 25.68 | 29.40 | 45.27 |
| RWKV-4 | 1.5 | - | 72.36 | 52.48 | 54.62 | 60.48 | 29.44 | 34.00 | 50.56 |
| HGRN | 1.0 | 100 | 70.89 | 48.02 | 51.62 | 55.64 | 27.90 | 31.60 | 47.61 |
| LightNet | 1.0 | 100 | 69.75 | 44.68 | 51.85 | 55.35 | 27.65 | 32.80 | 47.01 |
| OPT | 2.7 | 300 | 73.83 | 60.60 | 61.01 | 60.77 | 31.31 | 35.2.0 | 53.79 |
| Pythia | 2.8 | 300 | 74.10 | 59.31 | 59.91 | 64.14 | 33.02 | 35.60 | 54.35 |
| BLOOM | 3.0 | 350 | 70.57 | 54.53 | 58.48 | 59.43 | 30.38 | 32.20 | 50.93 |
| RWKV-4 | 3.0 | - | 72.42 | 58.75 | 57.30 | 62.92 | 35.15 | 36.20 | 53.79 |
| LightNet | 2.9 | 100 | 73.56 | 55.47 | 57.77 | 61.57 | 32.94 | 33.60 | 52.49 |

**Table 8: Performance Metrics Across Different LightNet Configurations**
| Model | 50K | 100K | 150K | 200K | 250K | 300K | 350K | 400K |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LightNet-S/8 | 192.79 | 172.23 | 161.23 | 154.34 | 150.25 | 147.40 | 145.27 | 143.31 |
| LightNet-S/4 | 167.33 | 132.89 | 118.77 | 110.88 | 105.15 | 101.25 | 97.56 | 94.90 |
| LightNet-S/2 | 145.66 | 119.20 | 104.90 | 94.45 | 87.18 | 82.41 | 78.63 | 75.61 |
| DiT-S/2 | - | - | - | - | - | - | - | 67.16 |
| LightNet-B/8 | 170.79 | 146.43 | 134.63 | 127.31 | 122.18 | 118.50 | 115.40 | 113.02 |
| LightNet-B/4 | 126.37 | 93.86 | 81.44 | 74.11 | 68.80 | 65.09 | 62.34 | 59.81 |
| LightNet-B/2 | 104.19 | 74.27 | 59.60 | 51.22 | 45.70 | 41.65 | 38.60 | 36.45 |
| DiT-B/2 | - | - | - | - | - | - | - | 42.76 |
| LightNet-L/8 | 157.76 | 130.29 | 116.06 | 107.50 | 101.10 | 96.47 | 92.79 | 89.51 |
| LightNet-L/4 | 104.18 | 77.02 | 64.55 | 56.16 | 49.99 | 45.58 | 41.91 | 37.54 |
| LightNet-L/2 | 84.38 | 48.98 | 35.32 | 28.05 | 23.75 | 21.06 | 18.94 | 17.42 |
| DiT-L/2 | - | - | - | - | - | - | - | 24.37 |
| LightNet-XL/8 | 158.75 | 129.23 | 114.72 | 105.75 | 99.35 | 94.53 | 90.66 | 87.22 |
| LightNet-XL/4 | 101.39 | 70.84 | 56.75 | 48.04 | 42.04 | 37.43 | 34.16 | 31.51 |
| LightNet-XL/2 | 79.22 | 45.46 | 31.61 | 25.55 | 21.37 | 18.74 | 16.84 | 15.52 |
| DiT-XL/2 | - | - | - | - | - | - | - | 19.20 |

### A.3 Configurations

In this section, we provide training configurations for all experiments. The configuration for Bidirectional Language Modeling is the same as (), while the configurations for the other experiments are as shown in Table 9, 10, 11, 12.
We use Pytorch () and A100 for training.

**Table 9: Comprehensive Configurations of the Model and Training Procedures for LightNet Experiments “Total batch size” means $\mathrm{batch_per_gpu}\times\mathrm{update_freq}\times\mathrm{num_gpus}$; “ALM” stands for Autoregressive Language Model; “IM” stands for Image Modeling, “IG“ stands for image generation.**
|  | ALM | IM | IG |
| --- | --- | --- | --- |
| Dataset | WikiText-103 | ImageNet-1k | ImageNet-1k |
| Tokenizer method | BPE | - | - |
| Src Vocab size | 50265 | - | - |
| Sequence length | 512 | - | - |
| Total batch size | 128 | 2048 | 256 |
| Number of updates/epochs | 50k updates | 300 epochs | 80 epochs |
| Warmup steps/epochs | 4k steps | 20 epochs | - |
| Peak learning rate | 5e-4 | 5e-4 | 1e-4 |
| Learning rate scheduler | Inverse sqrt | Cosine | - |
| Optimizer | Adam | Adamw | Adamw |
| Adam $\epsilon$ | 1e-8 | 1e-8 | 1e-8 |
| Adam $(\beta_{1},\beta_{2})$ | (0.9, 0.999) | (0.9, 0.98) | (0.9, 0.98) |
| Weight decay | 0.1 | 0.1 for Base, else 0.05 | 0 |
| Gradient clipping | - | 5.0 | - |
| GPUS | 4 | 8 | 8 |

**Table 10: Configurations for LLM**
| Params(B) | Layers | Hidden Dim | L.R. | Batch Size Per GPU | SeqLen | GPUs |
| --- | --- | --- | --- | --- | --- | --- |
| 0.15 | 15 | 768 | 3.00E-04 | 26 | 2048 | 8 |
| 0.385 | 26 | 1024 | 3.00E-04 | 15 | 2048 | 8 |
| 1.0 | 18 | 2048 | 3.00E-04 | 10 | 2048 | 16 |
| 2.9 | 36 | 2560 | 3.00E-04 | 36 | 2048 | 48 |

**Table 11: Model Configurations for Image Generation task.**
| Model | Layers | Hidden Dim | Heads | Params |
| --- | --- | --- | --- | --- |
| LightNet-S | 18 | 384 | 6 | 33M |
| LightNet-B | 18 | 768 | 6 | 131M |
| LightNet-L | 36 | 1024 | 16 | 470M |
| LightNet-XL | 42 | 1152 | 16 | 680M |

**Table 12: Model Configurations for Image Classification task.**
| Model | Layers | Hidden size | Heads | Params |
| --- | --- | --- | --- | --- |
| LightNet-T | 12 | 192 | 6 | 6.0M |
| LightNet-S | 12 | 384 | 16 | 22.6M |
| LightNet-B | 12 | 768 | 16 | 87.7M |

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 5
presentation: 4
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?
soundness: 4
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.
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.
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

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 1
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).
soundness: 2
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: Please see my comments above.
soundness: 2
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.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 3
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.
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.
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.

Meta-review: 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).

Decision: Reject