Linear Attention Architecture Explained: State Update Mechanisms and Efficient Chunked Computation

1. Background: Why Rethink Linear Attention

Transformer architectures perform well in natural language processing and multimodal tasks, but their core softmax attention mechanism faces significant computational and memory pressure in long-sequence scenarios. Traditional attention requires maintaining an attention matrix whose size scales quadratically with sequence length, so training cost and inference latency rise quickly as input length grows.

To mitigate this problem, researchers have turned to linear attention. The core idea is to transform attention computation from explicit query-key matrix multiplication into a recursively updatable state representation, decoupling per-step computational complexity from sequence length. In simple terms, the model no longer stores all historical key-value pairs; instead, it compresses historical information into a fixed-size state matrix.

However, early linear attention methods still fall short in expressiveness, forgetting control, and numerical stability. How to maintain efficient computation while giving the state matrix more flexible capabilities for updating and evicting information has become a key open question in this direction.

2. Overall Architecture: Hybrid Attention Units Working with a Linear Attention Module

The architecture discussed here adopts a hybrid design that combines conventional multi-head latent attention units with a linear attention module at a fixed ratio. The ratio is typically set at three to one: most computational resources remain on the traditional attention path, while the linear attention module serves as a complement, handling long-range information compression and efficient inference.

The advantage of this hybrid approach is that the traditional attention path retains strong local modeling and content-addressing capability, while the linear attention path reduces the model’s overall sensitivity to sequence length. Combined, the two paths give the model a more balanced configuration between expressiveness and computational efficiency.

In practice, the linear attention module is not a simple drop-in replacement for an existing layer. Instead, it is embedded as an independent attention branch. Outputs from the different branches are merged or weighted along the feature dimension, allowing the model to automatically adjust its reliance on each type of information according to the training objective.

3. State Update Mechanisms: From Pure Linear Attention to Gated State Matrices

To understand the design of this module, it helps to first review several typical linear attention state update formulas.

The most basic way to maintain state in linear attention is:

S_t = S_{t-1} + v_t k_t^T

where S_t is the state matrix at the current time step, and k_t and v_t are the key and value vectors of the current input. This formula simply accumulates each new key-value pair into the historical state. It is easy to compute, but historical information can only be diluted by continuously adding new content—there is no active forgetting mechanism.

DeltaNet builds on this by introducing an update strength control term:

S_t = (I - β_t k_t k_t^T) S_{t-1} + β_t v_t k_t^T

Here β_t can be understood as a scalar coefficient that controls how strongly the current input overwrites existing content along the current key direction in the state matrix. When β_t is large, old states associated with the current key are replaced more aggressively; when β_t is small, more of the old state is preserved.

Gated DeltaNet goes a step further by adding a forgetting gate α_t:

S_t = α_t (I - β_t k_t k_t^T) S_{t-1} + β_t v_t k_t^T

α_t controls the overall retention ratio of the previous state matrix. If α_t is close to 1, more historical information is retained; if α_t is significantly smaller than 1, the model can decay old information quickly. This design addresses the limitation of the default DeltaNet, which needs many steps to fully overwrite old states, and allows the model to actively choose to forget based on the current context.

4. Diagonal Gating and Low-Rank Correction: Finer-Grained Forgetting

Another technical route applies gating on a per-dimension basis to the state matrix. For example:

S_t = S_{t-1} Diag(α_t) + v_t k_t^T

Here α_t is no longer a scalar but a vector, and Diag(α_t) places α_t on the main diagonal of a diagonal matrix. As a result, each row of the state matrix has an independent decay rate. Compared with globally uniform forgetting, per-dimension gating preserves information across different feature channels at a finer granularity.

On top of this, a low-rank correction term can be added:

S_t = S_{t-1} (Diag(α_t) + a_t b_t^T) + v_t k_t^T

where a_t and b_t are low-rank vectors, and a_t b_t^T forms a low-rank matrix. This term can be viewed as adding cross-dimensional interaction beyond per-dimension decay, so that state updates are not confined to the diagonal direction but can also perform structural adjustment within a low-dimensional subspace.

The module discussed here combines “key-direction overwrite control + per-dimension gating,” with its state update summarized as:

S_t = (I - β_t k_t k_t^T) Diag(a_t) S_{t-1} + β_t v_t k_t^T

Compared with the standard gated DeltaNet, forgetting here does not use a global scalar α_t. Instead, a vector a_t acts on each dimension of the state matrix. With this design, the model can precisely adjust the replacement strength of old values associated with the current key through β_t, while applying differentiated decay to different state dimensions through Diag(a_t). Compared with uniform gating, this approach comes closer to letting the model “remember what should be remembered and forget what should be forgotten.”

5. Parameter Behavior: Separating Update Strength from Forgetting Speed

In this structure, β_t and a_t carry distinct control responsibilities.

β_t mainly controls how strongly the current input updates content along the “same key direction” within the existing state. Since k_t k_t^T forms a projection matrix along the current key direction, a larger β_t means the current value vector v_t replaces the old state more strongly. When outdated information needs to be overwritten quickly, the model can output a larger β_t; when historical context should be preserved, it can output a smaller β_t.

Diag(a_t), in contrast, mainly controls the forgetting speed of each dimension of the state matrix. The closer a component of a_t is to 1, the more completely the corresponding dimension’s historical state is retained; the closer to 0, the faster the decay. The advantage of this separation is that the model can make more flexible combined decisions between “replacing old key-values” and “global forgetting.” In some scenarios, even though the current input overlaps little with historical key directions, the model may still want to decay certain state dimensions overall—this is where per-dimension gating comes into play.

Compared with schemes that rely only on a scalar forgetting gate, this structure delivers higher expressiveness without significantly increasing computational complexity. For long-sequence tasks, this added flexibility helps improve the quality of long-term dependency modeling.

6. Hardware-Oriented Chunked Computation: State Passing and Local Attention Expansion

Although recursive state updates reduce theoretical complexity, computing them step by step is inefficient on parallel hardware such as GPUs. To balance parallelism with state expressiveness, a chunked computation strategy is typically adopted.

The core idea is to divide the input sequence into multiple consecutive chunks, use parallelizable local attention computation within each chunk, and pass historical information between chunks through the state matrix. The cross-chunk state update can be expressed as:

S_{[t+1]} = Diag(γ_{[t]}^C) S_{[t]} + (Γ_{[t]}^{i→C} ⊙ K_{[t]})^T (U_{[t]} − W_{[t]} S_{[t]})

In this formula, S{[t]} denotes the state at the start of chunk t, while K{[t]}, U{[t]}, and W{[t]} correspond to the keys, values, and intermediate variables within the chunk. Γ and γ are preset or computed gating factors, and ⊙ denotes element-wise multiplication. The formula first passes the historical state to the next chunk after diagonal decay, then adds the correction contributed by the current chunk to the state matrix.

The intra-chunk output can be written as:

O_{[t]} = (Γ_{[t]}^{i→C} ⊙ Q_{[t]}) S_{[t]}
         + Tril((Γ_{[t]}^{i→C} ⊙ Q_{[t]}) (K_{[t]} / Γ_{[t]}^{i→C})^T) (U_{[t]} − W_{[t]} S_{[t]})

where Q_{[t]} is the query matrix within the chunk, and Tril denotes lower-triangular mask computation. The first term reflects the contribution of the historical state to the current chunk’s output, while the second term represents local interactions between time steps within the chunk. Since the second term only involves the intra-chunk length, it can be efficiently parallelized through matrix multiplication.

This chunking strategy breaks long-sequence computation into multiple small chunks, preserving cross-chunk state recursion while leveraging intra-chunk parallelism to improve hardware utilization. Chunk size is typically adjusted based on the target hardware characteristics and sequence length to balance computational efficiency and memory overhead.

7. Numerical Stability and Precision in Engineering Implementation

When implementing linear attention state updates, numerical stability is a non-negligible concern. Repeated multiplication of gating terms, projection matrices, and the state matrix can lead to numerical overflow or precision loss, especially with long sequences and large state dimensions.

Common practices include constraining the gating factors—for example, bounding β_t within a certain range through an activation function, and parameterizing a_t appropriately to avoid extreme values. Scaling intermediate variables during chunked computation also helps mitigate accumulated error. In addition, in half-precision training scenarios, state matrix updates usually need to be restored to higher precision at specific steps to reduce rounding errors introduced by low-bit-width representations.

These engineering details do not change the model’s theoretical expressiveness, but they have a significant impact on final training stability and inference accuracy. In real deployments, targeted tuning is required based on the target hardware and model scale.

8. Summary

The linear attention module discussed in this article combines key-direction overwrite control with per-dimension gated decay in its state update mechanism, allowing the model to decide more flexibly what historical information to retain and what to replace. Compared with traditional linear attention or single-gate schemes, this structure offers finer-grained adjustment between expressiveness and computational efficiency.

At the same time, the chunked computation strategy allows recursive state passing and intra-chunk parallel computation to work together, alleviating the problem that purely recursive implementations struggle to fully utilize hardware performance. For developers focused on long-sequence modeling efficiency, understanding these state update mechanisms and chunked algorithms helps inform better technical choices in model design or inference optimization. Further practice can explore directions such as gating parameter initialization, chunk size configuration, and hybrid branch ratios.