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.

Humans do not need to remember every word they read when recalling something from context. Instead, they focus more on the more informative context and pay less attention to other irrelevant words. We believe that a language model could achieve similar results, delivering both strong performance and lower latency. We design a model that automatically decides whether to apply expensive but strong nonlinear softmax attention operations or cheaper but weaker linear attention operations for each token. The model will determine which operation to apply on its own and does not necessarily need to follow any pre-defined fixed rule. This decision process can be learned end-to-end with minimal computational overhead. Experiments showed that our token-level hybrid attention model, Neural Attention Search Linear (NAtS-L), achieves better retrieval performance with lower computational cost in long-context scenarios.