R e GLA : Refining Gated Linear Attention (2025)

Abstract: Recent advancements in Large Language Models (LLMs) have set themselves apart with their exceptional performance in complex language modelling tasks. However, these models are also known for their significant computational and storage requirements, primarily due to the quadratic computation complexity of softmax attention. To mitigate this issue, linear attention has been designed to reduce the quadratic space-time complexity that is inherent in standard transformers. In this work, we embarked on a comprehensive exploration of three key components that substantially impact the performance of the Gated Linear Attention module: feature maps, normalization, and the gating mechanism. We developed a feature mapping function to address some crucial issues that previous suggestions overlooked. Then we offered further rationale for the integration of normalization layers to stabilize the training process. Moreover, we explored the saturation phenomenon of the gating mechanism and augmented it with a refining module. We conducted extensive experiments and showed our architecture outperforms previous Gated Linear Attention mechanisms in extensive tasks including training from scratch and post-linearization with continual pre-training.

## introduction in the rapidly evolving field of natural language processing (nlp), transformer models have emerged as a groundbreaking innovation. these models have demonstrated unparalleled success across a wide array of tasks, revolutionizing our approach to understanding and generating natural language. they have proven their mettle in analyzing intricate documents, executing professional writing, and performing sophisticated reasoning tasks, thereby setting new benchmarks in the realm † corresponding author. *work conducted while at huawei noah's ark lab of nlp @xcite @xcite . the cornerstone of these transformer models is the softmax attention mechanism. this mechanism, an extension inspired by the attention mechanism employed in recurrent neural network (rnn) systems, has played a pivotal role in the success of transformer models @xcite . the softmax attention has outperformed rnn models in terms of parallelizability and the stability of gradient propagation over time, making it a preferred choice for many nlp tasks. however, the softmax attention mechanism is not without its challenges. it requires substantial computational resources and high memory usage, which can be a significant hurdle in practical applications. as the length of the input increases, the required computation grows quadratically. this growth restricts the context window size and complicates the deployment of these models in real-world scenarios @xcite . in addition to the issue of computational complexity, several studies have highlighted the limited length extrapolation capability of self-attention-based models @xcite . specifically, transformer models tend to underperform during inference if the sequence length of the test data exceeds that of the training data. as an orderinvariant encoding mechanism, the self-attentionbased encoder heavily depends on position embeddings (pes) to model input orders. however, these studies reveal that the inability of transformers to handle long sequences can be attributed to the limited length generalization ability of these position embedding methods @xcite @xcite . this finding underscores the need to explore alternative architectures to address the challenges associated with long-sequence processing. numerous studies have been conducted with the aim of mitigating this drawback by introducing linear attention operator @xcite @xcite @xcite @xcite @xcite . unfortunately, existing linear attention mechanisms frequently struggle to match the modeling quality of softmax attention. some work introduce gating mechanisms to improve the performance of linear attention @xcite @xcite . in this work, we delve into the different components of the gated linear attention mechanism with the goal of optimizing the training process while ensuring rapid inference. our contribution can be summarized as follows: first, we find that previous suggestions overlook some crucial aspects. we address the instability issue of feature mapping functions by proposing a normalized exponential solution. additionally, we introduce a variance reduction scaling factor to enhance its performance. then we revisit the normalization layer, emphasizing its role in stabilizing the training process. finally, we investigate the saturation phenomenon of the gating mechanism and enhance it with a refining module. by integrating our findings, we propose a novel architecture that outperforms previous gated linear attention mechanisms across various tasks. ## background we first briefly revisit the linear attention. our method is grounded on these works by analyzing the essential components of them. ## experiments in this section, we evaluate our method with other linear attention and the conventional transformer. this comparison spans autoregressive language modeling training from scratch and finetuning pretrained language models after replacing its softmax attention with linear variants. to justify our design choices for regla, we conduct a comprehensive ablation study and efficiency analysis. ## analysis and discussion in this section, we delve into a comprehensive discussion of our regla method. this includes an evaluation of the effectiveness of the gating mechanism, an analysis of speed and memory usage and an ablation study to understand the impact of each component. all of these aspects are examined in a controlled manner to ensure the reliability of our findings. in addition to the aforementioned evaluations, we also conducted a detailed analysis of our refining gate mechanism. as depicted in figure 4 , we examined the distributions of the forget gate activations for both the gated linear attention (gla) and our refining gated linear attention (regla) methods, both before and after the training process. ## related work there has been a great surge of research to design efficient variants of softmax attention or propose other alternatives for sequence modeling directly. the efficient variants broadly include two categories: sparsified attention and linear attention. sparsified attention @xcite @xcite computes attention maps with pre-defined or learnable masks. for instance, slide window attention (swa) @xcite limits each query input only attend to a certain number of preceding tokens. another efficient variant is linear attention @xcite @xcite @xcite , which exchanges the computation order by decomposing the softmax function with randomized or learnable feature functions. alternatively, @xcite propose modeling sequential data with state-space models (ssms) and show surprisingly good performance on a benchmark for comparing transformers over long sequence data @xcite . the h3 model @xcite expanded ssms with gated connections and a conventional local convolution layer. they also show ssm can work in tandem with attention mechanism in a hybrid manner. @xcite propose to substitute the ssm layer with a global convolution parameterized by mlp. @xcite incorporates data-dependent gating to ssms and show comparable performance as transformer-based language models. @xcite develops rwkv architecture which absorbs insights from rnn and attention-free transformer @xcite . the retnet model @xcite and transformer-llm @xcite apply a decay factor to the current hidden state before incorporating the current input information and achieving impressive improvements. yang et al. (2024a) and yang et al. (2024b) develop chunkwise forms of gla and deltanet respectively to parallelize the computation of gated linear recurrence models and provide a triton-based library to accelerate the training speed of linear attention model @xcite . another interesting line of work dedicated to substituting the softmax attention in a pre-trained model with linear attention and performing continual training to bridge their performance gap. @xcite take a pre-trained sa transformer, swap the sa modules with linear attention, and continue training the entire model on the same task. mao (2022) adopts the same procedure by optimizing it with the fast decay rules and removing the relu function in the feature maps, namely, a simple identity map. @xcite ; @xcite linearized existing large pretrained transformers into recurrent neural networks (rnns) with a modest continual pre-training budget to recover their performance. wang et al. improve hybrid models by applying knowledge distillation @xcite from pre-trained transformers to mamba, enhancing efficiency and inference speed. ## conclusion in this study, we conduct an in-depth examination of three pivotal components that significantly influence the performance of the gated linear attention mechanism: feature maps, normalization, and the gating mechanism. we posit the unstable issue of commonly used feature mapping functions and develop stable exponential functions. apart from that, we also provide a corresponding variance reduction scaling factor to further improve its performance. then we revisit the normalization layer and give additional justification for the incorporation of normalization layers to stabilize the training process. furthermore, we explore the saturation phenomenon of the gating mechanism and enhance it with a refining mechanism. by integrating our findings, we propose a novel architecture that surpasses the performance of previous gated linear attention mechanisms in extensive tasks.