Neural Attention Search Linear: Towards Adaptive Token-Level Hybrid Attention Models

TL;DR

The paper tackles the inefficiency of softmax attention in long-context transformers by introducing NAtS-L, which adaptively assigns tokens to either linear or softmax attention within the same layer. It implements a chunk-wise framework with an Attention Score Layer to decide per chunk whether to use linear or softmax operations, and uses a DeltaNet/GDN-based token mixer to manage historical memory efficiently. Gradient-based optimization jointly learns operation choices and model weights, achieving reduced computational cost while preserving long-range retrieval capabilities; experiments show strong long-context performance and competitive latency against full softmax transformers. This token-level hybrid approach demonstrates robust long-context modeling with practical efficiency gains, particularly for retrieval-heavy and extended-context tasks.

Abstract

The quadratic computational complexity of softmax transformers has become a bottleneck in long-context scenarios. In contrast, linear attention model families provide a promising direction towards a more efficient sequential model. These linear attention models compress past KV values into a single hidden state, thereby efficiently reducing complexity during both training and inference. However, their expressivity remains limited by the size of their hidden state. Previous work proposed interleaving softmax and linear attention layers to reduce computational complexity while preserving expressivity. Nevertheless, the efficiency of these models remains bottlenecked by their softmax attention layers. In this paper, we propose Neural Attention Search Linear (NAtS-L), a framework that applies both linear attention and softmax attention operations within the same layer on different tokens. NAtS-L automatically determines whether a token can be handled by a linear attention model, i.e., tokens that have only short-term impact and can be encoded into fixed-size hidden states, or require softmax attention, i.e., tokens that contain information related to long-term retrieval and need to be preserved for future queries. By searching for optimal Gated DeltaNet and softmax attention combinations across tokens, we show that NAtS-L provides a strong yet efficient token-level hybrid architecture.

Figures (8)

• Figure 1: A comparison between different attention modules. In \ref{['fig:SoftmaxAttnMap']}, softmax attention needs to compute each $\mathbf{Q}\mathbf{K}$-pair. In \ref{['fig:LinearAttnMap']}, linear attention models perform chunk-wise computation and therefore only maintain a fixed hidden state (the color darkness indicates hidden states from different time steps). \ref{['fig:MixAttnMap']}, NAtS-L automatically determines if a model belongs to softmax or linear attention and merge their output to achieve both performance and efficiency.
• Figure 2: The NAtS-L Architecture. The projection layers for $\mathbf{q},\mathbf{k},\mathbf{v}$ are a linear layer followed by a short conv and SiLU activation function. The scores for each operation are computed by a mean pooling layer followed by a linear layer (Equation \ref{['eq:natsscore']}). The model then selects the operations with the highest score for each chunk and feeds them to different attention operations. The outputs from each attention model are then normed and weighted summed (Equation \ref{['eq:outnorm']}), where the weights $w$ for each attention output are mapped from the $\mathbf{q}$ matrix.
• Figure 3: Per-token perplexity on different datasets, all the models are trained with 4096 tokens (the black vertical line)
• Figure 4: Inference Latency with different input context length
• Figure 5: Fraction of softmax attention tokens within each layer for NAtS-L Hybrid in different tasks.