💬 NLP

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

This paper introduces Neural Attention Search Linear (NAtS-L), a framework that dynamically assigns either efficient linear attention or expressive softmax attention to individual tokens within the same layer, thereby achieving a strong balance between computational efficiency and model expressivity for long-context scenarios.

Original authors: Difan Deng, Andreas Bentzen Winje, Lukas Fehring, Marius Lindauer

Published 2026-06-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Difan Deng, Andreas Bentzen Winje, Lukas Fehring, Marius Lindauer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a story, but you have a massive library of notes in front of you. Some notes are tiny, fleeting thoughts (like "the sky is blue"), while others are crucial plot points you'll need to remember for the very last page (like "the hero has a secret twin").

The Problem: The "Everything" vs. The "Too Much"
Current AI models (called Transformers) are like a librarian who insists on reading every single note in the library every time they write a new sentence.

The Good: They remember everything perfectly.
The Bad: As the library grows, this becomes incredibly slow and expensive. It's like trying to find a specific book by reading the entire library's catalog from start to finish every single time you ask a question. This is the "quadratic complexity" bottleneck mentioned in the paper.

On the other hand, there are "Linear" models. These are like a librarian who only keeps a single summary card of the library.

The Good: They are super fast and cheap to run, no matter how big the library gets.
The Bad: Because they only keep one summary card, they often forget the specific, important details needed for long stories. They lose the "long-term memory."

The Solution: The "Smart Hybrid" Librarian (NAtS-L)
The authors of this paper propose a new system called NAtS-L (Neural Attention Search Linear). Instead of choosing between "read everything" or "keep a summary," they built a librarian that decides on the fly which notes to read carefully and which to just glance at.

Here is how it works, using a simple analogy:

1. The "Chunking" Strategy

Imagine the librarian doesn't look at the notes one by one. Instead, they grab a stack of 64 notes (called a "chunk") at a time.

2. The "Traffic Cop" (The Search)

Before processing the stack, a smart "Traffic Cop" (the Neural Attention Search) looks at the notes and asks: "Do these notes contain a crucial plot point, or are they just filler?"

If the notes are crucial (Long-term memory): The Traffic Cop flags them for the "Slow, Careful Reader" (Softmax Attention). This reader will cross-reference these notes with everything else to ensure the hero's secret twin isn't forgotten.
If the notes are just filler (Short-term memory): The Traffic Cop flags them for the "Fast, Summary Reader" (Linear Attention). This reader just quickly summarizes the stack and moves on, saving massive amounts of time.

3. The Magic of "Same Layer, Different Jobs"

Previous attempts at mixing these two styles were like having one floor of the library dedicated to "Slow Readers" and another floor for "Fast Readers." This was rigid.

NAtS-L is different. In every single stack of notes, the system dynamically decides: "Note #1 needs the Slow Reader, but Note #2, #3, and #4 can be handled by the Fast Reader."

It's like a team of workers where some people are doing detailed, slow work on specific items, while others are quickly packing up the rest, all happening at the exact same time.

Why This Matters (According to the Paper)

The paper claims this approach gets the best of both worlds:

Speed: It's much faster than reading every note because it skips the heavy lifting for boring parts.
Memory: It's much smarter than just keeping a summary because it knows exactly when to stop and pay close attention to important details.

The Results:
When the authors tested this on tasks like:

Finding a needle in a haystack: (Can the model find a specific fact hidden in a huge text?)
Reading long stories: (Can it remember the beginning of a story by the time it reaches the end?)

NAtS-L performed better than models that only use the "Slow Reader" or only use the "Fast Reader." It managed to remember long-term details without slowing down the computer as much as the old "read everything" method.

In a Nutshell:
NAtS-L is a smart AI that learns to ignore the boring stuff to save energy, but zooms in on the important stuff so it doesn't forget the plot. It doesn't force the computer to do the hard work for every single word; it lets the computer choose the right tool for the job, word by word.

Technical Summary: Neural Attention Search Linear (NAtS-L)

Problem Statement

The quadratic computational complexity ( $O(L^2)$ ) of standard softmax-based self-attention in Transformers has become a significant bottleneck for long-context scenarios. While linear attention models offer a promising alternative by compressing past Key-Value (KV) values into a fixed-size hidden state to achieve linear complexity ( $O(L)$ ), they often suffer from limited expressivity due to the constraints of this fixed state size. Conversely, standard Transformers preserve all KV context information, enabling better long-term retrieval but incurring high memory and computational costs.

Existing hybrid approaches attempt to mitigate this by interleaving softmax and linear attention layers or using fixed schedules to switch between them. However, these methods often retain fixed structures that may not flexibly capture all required information across varying contexts, and their efficiency remains bottlenecked by the presence of full softmax layers. Furthermore, sparse attention variations often rely on pre-defined human heuristics, potentially ignoring correlations not covered by their sparse maps.

Methodology

The paper proposes Neural Attention Search Linear (NAtS-L), a framework that applies both linear and softmax attention operations within the same layer but on different tokens. Instead of a fixed architecture, NAtS-L automatically determines the optimal attention operation for each token block based on the input context.

Core Mechanisms

Token-Level Hybridization: The input sequence is split into chunks. An Attention Score Layer (consisting of mean-pooling followed by a linear layer) maps each chunk to a set of scores representing different operation types (Softmax Attention vs. Gated DeltaNet, a linear attention variant). The operation with the highest score is assigned to that chunk.
Unified Search Space: The framework unifies non-linear softmax and linear attention operations, allowing them to be processed within a single layer. It learns the importance of each token block, applying softmax attention to tokens with high long-term impact (requiring long-term retrieval) and linear attention to tokens with short-term impact.
Chunk-wise Computation: To maintain hardware efficiency, NAtS-L utilizes a chunk-wise parallel approach. Within a chunk, both linear and softmax correlations are computed. The model maintains hidden states for both operation types. When a new chunk arrives, the model pseudo-updates both states; however, for chunks not selected as the target operation, the hidden states are reset or rolled back to the start of the chunk, ensuring that only the relevant KV values are retained for the selected operation type.
Output Merging: The outputs from the selected softmax and linear operations are normalized individually (using RMS normalization) and then combined via a weighted sum. The weights for this combination are dynamically generated from the query ( $Q$ ) projections, ensuring the model learns to balance the contributions of both attention types.
Gradient-Based Optimization: The selection of attention types is learned end-to-end. Gradients are computed for the attention score layer, allowing the model to jointly optimize the operation selection and the model weights without requiring auxiliary losses or pre-trained models.

Key Contributions

Unified Framework: The authors provide a unified description for non-linear softmax and linear attention, enabling them to be processed and selected within a single layer.
Adaptive Architecture: NAtS-L is proposed as a framework that automatically determines the optimal attention operation for input context information on a token-level (chunk-level) basis, rather than relying on fixed layer-wise ratios.
Efficiency and Expressivity: The paper demonstrates that NAtS-L achieves a strong balance between performance and efficiency, offering better long-context modeling capabilities than pure linear models while maintaining lower computational latency than full softmax Transformers.

Experimental Results

The authors evaluated NAtS-L on academic-scale pre-training tasks using the Fineweb-Edu dataset with model sizes of 380M and 800M parameters.

Language Modeling & Reasoning: On zero-shot commonsense reasoning benchmarks (e.g., LAMBADA, PIQA, Hellaswag), NAtS-L and its hybrid variant performed comparably to or slightly better than pure Gated DeltaNet (GDN), Mamba2, and Transformer baselines.
Retrieval Tasks: On real-world retrieval tasks (e.g., SWDE, SQD, DROP), NAtS-L Hybrid achieved the best average scores, outperforming pure linear models (which struggled with specific tasks like FDA) and pure Transformers (which struggled with SQD). This highlights the model's robustness in preserving necessary context information.
Length Extrapolation: Evaluated on contexts up to 65,536 tokens (PG19, NarrativeQA, CodeParrot), NAtS-L Hybrid maintained perplexity levels even when the input context exceeded the training length (4096 tokens), whereas pure linear hybrids (GDN Hybrid) failed to extrapolate effectively.
Long-Context Benchmarks: On RULER and LongBench, NAtS-L Hybrid achieved state-of-the-art performance among the compared models, particularly at 8k and 16k context lengths.
Inference Latency: NAtS-L Hybrid showed a 5.4x speedup in pre-filling compared to Transformers and a 2.3x speedup in decoding at 128k context length, while being only 1.66x slower than pure GDN models.

Significance and Claims

The paper claims that NAtS-L provides a flexible hybrid architecture that adaptively assigns tokens to either linear or softmax attention based on their specific informational needs. By learning this distribution end-to-end, the model avoids the rigidity of fixed hybrid structures and the inefficiency of full softmax attention.

The authors emphasize that this approach allows the model to:

Preserve long-context related information via softmax attention where necessary.
Reduce overall computational costs via linear attention for tokens with short-term impact.
Achieve a better accuracy-efficiency trade-off without requiring pre-trained models or auxiliary constraints on the ratio of attention types.

The work suggests that a token-level mixture of linear and softmax attention is a viable path toward scalable, long-context language models that do not compromise on retrieval capabilities or inference speed.