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

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.

💡 Research Summary

The paper addresses the well‑known quadratic time and linear memory bottlenecks of softmax‑based transformers when processing very long sequences. While linear‑attention families (e.g., Performer, FAVOR+, and recent kernel‑based methods) reduce the computational cost to linear time by compressing past key‑value (KV) pairs into a fixed‑size hidden state, they suffer from limited expressivity because a single hidden vector cannot capture all long‑range information. Prior hybrid approaches interleave softmax and linear layers, but the softmax layers remain a performance bottleneck.

Neural Attention Search Linear (NAtS‑L) proposes a token‑level hybrid architecture that decides, for each token (or chunk of tokens), whether to process it with a softmax attention operation (preserving full KV information for long‑term retrieval) or with a linear attention operation (compressing it into a hidden state for short‑term impact). The decision is made automatically during training by an “Attention Score Layer”. This layer first mean‑pools the representations of a chunk and then applies a linear projection to produce two scores: one for the linear branch and one for the softmax branch. The higher‑scoring branch is selected for that chunk.

Formally, the input sequence X∈ℝ^{L×d} is split into T = L/C chunks of size C. For each chunk t, the score vector s_t = W_score·Mean(X