ReGLA: Refining Gated Linear Attention

TL;DR

ReGLA introduces a refined gated linear attention mechanism with improved feature mapping, normalization, and gating, leading to superior performance in large language model tasks while reducing computational complexity.

Contribution

This work presents a comprehensive refinement of Gated Linear Attention by enhancing feature maps, normalization, and gating, resulting in better performance and training stability.

Findings

Outperforms previous Gated Linear Attention methods

Effective in training from scratch and continual pre-training

Reduces computational complexity of attention mechanisms

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.