ArXiv: 2510.07318

🎯 Pitch

Augmenting a 3B Transformer with a tiny recurrent head that compresses old context into a fixed-size state slashes FLOPs by 40.5% and memory by 74% on 128k-token inputs, yet somehow boosts long-context benchmark scores beyond the original full-attention model. The trick is to keep a giant 32k sliding window for exact recent memory and only switch on the cheap recurrent compressor for the older tokens, so short-context behavior stays identical.


1. Executive Summary

This paper introduces Artificial Hippocampus Networks (AHNs), a memory framework that transforms lossless attention KV cache outside a sliding window into a fixed-size compressed long-term memory, inspired by the Multi-Store Model from cognitive science. Evaluated on long-context benchmarks LV-Eval and InfiniteBench using Qwen2.5-Instruct models (3B, 7B, 14B) augmented with RNN-like AHN instantiations (Mamba2, DeltaNet, GatedDeltaNet), AHNs reduce FLOPs by 40.5% and memory cache by 74.0% at 128k sequence length while improving LV-Eval average score from 4.41 to 5.88—matching or exceeding full-attention performance with substantially lower cost. The framework achieves this efficiency by activating AHNs only when sequence length exceeds a large 32k sliding window, establishing that hybrid lossless-compressed memory architectures can substitute for full attention on extra-long sequences without sacrificing performance on short-context tasks where AHNs remain inactive.

2. Context and Motivation

The Core Problem: The Inherent Tension Between Memory Fidelity and Computational Efficiency

The central challenge this paper addresses is a fundamental architectural trade-off that has shaped the entire trajectory of sequence modeling research: lossless memory provides perfect recall but incurs growing computational cost, while compressed memory offers constant efficiency but inevitably loses information. This tension becomes acute when processing extremely long sequences — precisely the regime where modern LLMs are increasingly expected to operate.

To understand why this matters, we must first appreciate what "memory" means in neural sequence models. Every model that processes a sequence of tokens must somehow retain information about past tokens to inform future predictions. How this retention is implemented determines both the model's effectiveness (can it recall specific facts from 50,000 tokens ago?) and its efficiency (how much GPU memory and computation does each new token require?). The two dominant paradigms — attention mechanisms and recurrent neural networks — represent opposite poles of this fidelity-efficiency spectrum.

Attention-based Transformers maintain what is essentially a perfect, lossless memory. For every token processed, the model generates a key vector and a value vector that are appended to an ever-growing key-value (KV) cache. When generating the next token, the attention mechanism computes relevance scores between the current query and every key in the cache, then aggregates the corresponding values. This means the model has direct, unmediated access to the exact representation of every token it has ever seen. The benefits are profound: Transformers can learn to retrieve specific information across arbitrary distances, they enable in-context learning where the model conditions on provided examples, and they support chain-of-thought reasoning that builds on earlier steps without degradation. These capabilities have driven the Transformer's dominance across virtually all sequence modeling tasks.

However, the lossless nature of the KV cache is, as the paper puts it, "a double-edged sword." The memory cache size grows linearly with sequence length — storing two vectors per token, per attention head, per layer. With models having dozens of layers and dozens of heads, each with hundreds of dimensions, this cache balloons rapidly. For a 7B parameter model at 128k context length, the KV cache alone consumes approximately 14.7 GB of GPU memory (Table 9). Simultaneously, the computational cost of computing attention scores between the current query and all cached keys grows quadratically with sequence length. At 128k tokens, a 7B model's attention mechanism performs roughly 3.23×10153.23 \times 10^{15} floating-point operations for the token mixer alone — dominating the total model computation (Section 2.3, Table 1, Table 9).

This creates a practical crisis for long-context applications: the hardware requirements become prohibitive, latency increases to unacceptable levels, and the throughput of serving systems plummets. As the paper demonstrates in Figure 3d, CUDA memory usage grows linearly under FlashAttention for standard Qwen2.5 models, while the AHN-augmented version maintains nearly constant memory. The quadratic FLOPs growth is equally stark (Figure 3a), with the cost difference between the full-attention and AHN models diverging rapidly beyond the 32k window threshold.

RNN-like models offer the opposite trade-off. They maintain a single fixed-size hidden state — essentially a vector (or matrix) of predetermined dimension — that is updated at each step by combining the new input with the previous hidden state. Since the state size is constant regardless of how many tokens have been processed, both the memory cost per token and the update computation per token remain constant (O(1)\mathcal{O}(1) relative to sequence length). This makes RNNs extraordinarily efficient for long sequences: they can theoretically process sequences of arbitrary length without growing memory or per-step computation.

The cost, of course, is information loss. A fixed-size vector must represent arbitrary amounts of history — compressing 10,000 tokens into the same 4096-dimensional representation that previously held 100 tokens. The paper explicitly cites work by Wen et al. (2025) showing that "RNNs are not transformers (yet): The key bottleneck on in-context retrieval" — the fundamental limitation being that fixed-size compressed memory struggles with "tasks that require precise long-range information recall" (Section 1). This is not a minor degradation; it represents a fundamental capacity ceiling. An RNN simply cannot store the exact text of a specific paragraph from 50 pages ago in the way that a Transformer's KV cache can.

Why This Problem Matters Now

This trade-off has become urgent for several converging reasons that the paper makes clear through both explicit discussion and implicit framing:

1. The explosion of long-context use cases. Real-world applications increasingly demand processing of entire documents, multi-turn conversations spanning thousands of exchanges, codebases with millions of lines, and legal contracts or scientific papers that run hundreds of pages. Benchmarks like LV-Eval and InfiniteBench specifically test 128k-token contexts, and models are being marketed with 1M-token context windows. At these scales, the quadratic complexity of attention is not just a theoretical concern — it becomes the binding constraint on what is practically deployable.

2. The economic reality of serving LLMs. Inference costs are dominated by memory bandwidth and compute requirements that scale with context length. The paper's concrete numbers tell the story: for Qwen2.5-3B at 128k context, full attention consumes 3.29×10153.29 \times 10^{15} model FLOPs with a 9.44 GB memory cache, while AHN-augmented versions reduce this to 1.95×10151.95 \times 10^{15} FLOPs (59.4% of original) and 2.45 GB cache (26.0% of original) — a 40.5% FLOP reduction and 74.0% cache reduction (Figure 1b, Table 9). These are not marginal improvements; they represent the difference between a model that requires expensive high-memory GPUs and one that runs on commodity hardware.

3. The gap between capability and efficiency. Modern LLMs are trained with finite context windows (e.g., 32k for Qwen2.5) but are expected to generalize to much longer sequences at inference. When sequences exceed the training context length, two problems arise: perplexity degrades sharply (as shown in Figure 3c for Qwen2.5-3B on a 57k-token PG19 passage), and the attention mechanism's quadratic cost becomes unbearable precisely when the model's native capacity is already strained. The paper notes that Qwen2.5-3B-Instruct's perplexity "rises sharply once the 32k token context window is exceeded" (Section 3.2), while AHN-augmented models maintain consistently low perplexity — addressing both the quality and efficiency challenges simultaneously.

Where Prior Approaches Fall Short

The paper situates itself within a rich landscape of solutions that have been proposed for the long-context efficiency problem, identifying specific limitations in each category. Understanding these limitations provides the motivation for AHN's design choices.

Sliding window attention (SWA): dropping distant tokens entirely. The simplest approach is to restrict attention to a fixed window of recent tokens, discarding everything beyond it. This is the baseline method used in the original Transformer paper (Vaswani et al., 2017) and is a natural first approximation: if you can only attend to WW tokens, both memory and computation become bounded by WW rather than the full sequence length LL. However, as the paper notes, "this method discards KV pairs outside the window, thereby losing long-range context" (Section 4.2). The model has zero access to information from earlier in the sequence — including crucial context like document introductions, task instructions, or reference facts that were stated many tokens ago.

The paper demonstrates this limitation quantitatively. In Table 2, sliding window attention with attention sinks (a refinement that keeps a few initial "sink" tokens for stability) achieves only 4.59 average score on LV-Eval for Qwen2.5-3B, compared to 4.41 for full attention. While it's competitive on some tasks, it fundamentally cannot answer questions that require retrieving information from beyond the window — a hard ceiling on capability. The attention sink mechanism (Xiao et al., 2024) helps stabilize the attention distribution by keeping initial tokens that serve as "sinks" for attention mass, but it doesn't solve the information loss problem for content between the sinks and the window.

Sparse attention: selective but still lossy. Sparse Transformer variants (Child et al., 2019) retain KV pairs at specific pattern positions — every kkth token, dilated patterns, or combinations of local and global attention — rather than keeping all tokens or just a contiguous window. This captures some long-range dependencies while reducing the number of cached tokens. The paper acknowledges this approach but notes that it "still drops portions of the KV cache, potentially missing important information" (Section 4.2). The fundamental issue is that pattern-based selection is content-agnostic: a token is kept or discarded based on its position, not its semantic importance. A crucial number or entity reference at the "wrong" position gets dropped just as readily as a preposition.

Segment-level recurrence: fixed-capacity FIFO memory. Transformer-XL (Dai et al., 2019) introduced the idea of caching hidden states from previous segments and attending to them as a form of extended memory. This is a FIFO (First-In, First-Out) buffer: when the buffer is full, the oldest segment is evicted to make room for the newest. The Compressive Transformer (Rae et al., 2020) — which the paper implements as a baseline — extends this by compressing older segments into a secondary, compressed FIFO memory using pooling operations (max or average pooling in the paper's implementation). However, as the paper points out, this approach "still discards memory once the slots are full" (Section 4.2). The FIFO structure means that compression is a one-time operation applied uniformly to entire segments, and compressed memories are still eventually evicted. There is no mechanism for selective retention — the model cannot decide that a particular piece of information is important and should be preserved indefinitely.

The paper's experimental results confirm the limitations of this approach. Compressive Transformers with max or average pooling (CT-Max, CT-Average in Table 2) achieve LV-Eval scores of 4.12 and 4.47 respectively for Qwen2.5-3B, underperforming both sliding window baselines and AHN variants. The fixed, uniform compression strategy — applying the same pooling operation regardless of content — cannot adapt to the varying importance of different tokens. A mathematical formula from 30 pages ago deserves different treatment than a filler word, but CT cannot distinguish them.

KV cache selection and eviction: after-the-fact triage. A significant body of work has focused on post-hoc selection: after generating the full KV cache, decide which entries to keep and which to discard based on heuristics or learned importance scores. Methods like H2O (Zhang et al., 2023), Scissorhands (Liu et al., 2023), and SnapKV (Li et al., 2024) compute attention-based importance scores for cached KV pairs and prune the less important ones. The paper acknowledges these methods (Section 4.2) but their fundamental limitation is that they make retention decisions after all tokens have been processed, meaning the full memory and computation cost has already been paid during prefill. They reduce memory for the decoding phase but don't address the quadratic prefill cost. Moreover, importance-based pruning still loses some tokens entirely — there's no attempt to preserve partial information from "unimportant" tokens through compression.

Hybrid attention-RNN architectures: the missing comprehensive framework. The paper positions recent work on interleaving attention and RNN layers as the closest intellectual precursor to AHNs. Jamba (Lieber et al., 2024) alternates Transformer layers with Mamba state-space model layers. Infini-attention (Munkhdalai et al., 2024) performs chunk-wise attention and maintains a compressive memory updated chunk-by-chunk. However, the paper identifies a critical gap in these approaches: they typically use small attention windows (e.g., 64 tokens in LoLCATs and HQLT, 2048 in Infini-attention). This means the RNN components are activated for most of the sequence, even on short contexts where attention alone would be both efficient and more capable. The paper argues this is unnecessary: "Since quadratic attention remains efficient for short and medium sequences, the quadratic-complexity bottleneck only appears when sequences become extra long" (Appendix F). By activating RNN compression only for tokens outside a large (32k) window, AHN preserves the full power of attention where it's cheap and effective, while deploying compression selectively where it's needed.

How This Paper Positions Itself Relative to Existing Work

The paper's positioning is distinctive and careful, establishing AHNs as both a conceptual framework and a practical implementation that addresses gaps left by prior work along several dimensions:

1. A cognitive-science-inspired memory framework rather than just an architectural tweak. The paper explicitly roots its design in the Multi-Store Model (MSM) of memory from cognitive psychology (Atkinson and Shiffrin, 1968). In MSM, short-term (or working) memory has limited capacity — roughly 7±2 items in the classic formulation — but maintains information with high fidelity. The hippocampus then consolidates short-term memories into long-term cortical representations that are compressed, semantically organized, and can persist indefinitely. The paper's mapping is direct: the sliding window attention is "lossless short-term memory," and the AHN module plays the role of the hippocampus, "recurrently compressing the out-of-window context into a fixed-size state as the long-term compressed memory" (Section 1).

This framing does more than provide a biological metaphor. It motivates specific design choices: compression should be continual (happening at each step, not in chunks), recurrent (building on previous compressed states, not starting fresh), and learnable (the hippocampus doesn't use a fixed pooling function — it learns what to consolidate). The paper contrasts this with simple heuristics like max-pooling used in Compressive Transformers, which lack the ability to learn compression strategies from data.

2. A large-window philosophy that preserves attention's strengths. The paper's most distinctive architectural position is its insistence on a large (32k tokens) sliding window for attention. This is not an arbitrary choice — it reflects a clear argument about when attention's efficiency problems actually matter. By keeping a generous window, the model preserves perfect recall over a substantial recent context that covers most practical tasks. The AHN only activates when sequences exceed this window, meaning the model is "exactly as a standard full-attention model" for short-context tasks (Appendix H). This directly contrasts with approaches like LoLCATs (Zhang et al., 2025) that use tiny attention windows (64 tokens) and rely heavily on the RNN component, requiring additional effort to match attention's short-context performance.

The paper demonstrates the importance of this choice through its ablation on inference window size (Figure 4): performance steadily improves as the window grows from 1k to 16k tokens, plateaus at 32k, and eventually declines beyond 64k due to attention dilution. The 32k default is the sweet spot — large enough to capture most needed context, small enough that the quadratic cost hasn't yet become prohibitive. This is also larger than the pretraining context length of many base models (e.g., Qwen2.5's 32k), meaning the model operates within its native attention capabilities inside the window.

3. An efficient training paradigm that leverages existing models. Rather than training hybrid models from scratch (as in HQLT; Irie et al., 2025) or distilling entire models (as in MiL; Wang et al., 2024), the paper introduces a self-distillation approach where the base LLM's weights are entirely frozen and only the AHN parameters are trained. This is both computationally efficient — training a 7B model's AHNs requires only ~10 hours on 32 A100 GPUs with 1B tokens — and respectful of the substantial investment in pretraining the base model. The self-distillation objective (KL divergence between the full-attention teacher's output distribution and the AHN-augmented student's output distribution) provides a dense training signal that teaches the AHN to preserve the teacher's behavior even when attending over a limited window plus compressed memory.

The paper shows that this training scheme is not just efficient but also effective: replacing self-distillation with standard next-token prediction (CE loss) causes a notable performance drop (Table 4), from 40.59 to 39.59 on LongBench for the 7B model. The authors hypothesize that "CE provides sparse learning signals, and pushes the small AHN modules towards shortcuts in the training data" while self-distillation offers "denser guidance over the teacher's entire output distribution" (Section 3.4).

4. A framework that generalizes across instantiations. By defining AHNs as an abstract recurrent memory update — htW=AHN((ktW,vtW),htW1)h_{t-W} = \text{AHN}((k_{t-W}, v_{t-W}), h_{t-W-1}) — the paper shows that the concept can be instantiated with different RNN architectures. The experiments span Mamba2 (a state-space model), DeltaNet (a fast weight programmer), and GatedDeltaNet (a gated variant), with all three showing consistent improvements over sliding window baselines (Table 2, Table 3). This demonstrates that the framework's benefits come from the architectural decomposition (lossless window + compressed long-term memory) rather than from a specific recurrent architecture. The paper positions this as a contribution of the concept, not just the implementation.

5. Acknowledging the fundamental limitation: compression loses exact recall. Unlike some prior work that overclaims the capabilities of compressed memory, the paper is forthright about its limitations. Table 5 in Appendix B shows that on exact-recall needle-in-a-haystack tasks from RULER, AHN-GDN performs on par with sliding window attention but markedly worse than full attention. The paper states directly: "While AHN-augmented models enable efficient long-context reasoning, they inevitably struggle on tasks that require exact-recall from the compressed memory" and suggests "memory management that preserves critical information in lossless memory while leveraging compression for efficiency" as future work (Appendix B). This intellectual honesty about the inherent trade-off — you can't have both perfect recall and constant memory — actually strengthens the paper's contribution by clearly delineating the boundary conditions under which the approach is appropriate.

3. Technical Approach

3.1 Reader Orientation

This paper presents a system for augmenting existing large language models so they can process extremely long sequences efficiently without retraining the entire model. The core problem it solves is that standard Transformers have memory and computation costs that grow quadratically with sequence length, making very long contexts (e.g., 128,000 tokens) prohibitively expensive. The solution is a hybrid memory architecture: keep a large sliding window of exact, lossless attention memory for recent tokens, and use a small recurrent neural network (the Artificial Hippocampus Network) to continually compress older tokens into a compact, fixed-size memory as they leave the window.

3.2 Big-Picture Architecture (Diagram in Words)

The AHN-augmented model consists of five major components that interact during inference:

  1. Base LLM (e.g., Qwen2.5-Instruct) — A standard Transformer with frozen pretrained weights. It provides the core language modeling capability and the attention mechanism for lossless memory.

  2. Sliding Window Attention — A modified version of the base model's self-attention that restricts the attention receptive field to the most recent WW tokens (default W=32,768W = 32,768). Within this window, attention operates exactly as in a standard Transformer — queries attend to all keys in the window, providing perfect recall of recent context. Attention sink tokens (a small number of initial tokens, default 128) are also preserved for numerical stability.

  3. Artificial Hippocampus Network (AHN) — A small recurrent neural network (instantiated as Mamba2, DeltaNet, or GatedDeltaNet) that compresses the key-value pair for each token as it exits the sliding window. The AHN takes the KV pair (ktW,vtW)(k_{t-W}, v_{t-W}) and the previous compressed memory state htW1h_{t-W-1}, and produces an updated compressed state htWh_{t-W} that summarizes all tokens from positions 1 through tWt-W. This operates at every attention head in every layer.

  4. AHN Output Projection — A gated linear transformation that converts the compressed memory state htWh_{t-W} into an output vector yAHN,ty_{\text{AHN},t} compatible with the attention output. This involves a per-head scalar gate γ(xt)\gamma(x_t) that modulates how much the compressed memory contributes, followed by a per-head output projection matrix WoW_o.

  5. Summation Point — The outputs from sliding window attention and the AHN compressed memory pathway are simply summed to produce the final token-mixing output for the current token: yt=yAHN,t+Attention({(ki,vi)}i=tW+1t,qt)y_t = y_{\text{AHN},t} + \text{Attention}(\{(k_i, v_i)\}_{i=t-W+1}^t, q_t).

Information flows as follows: a token sequence enters the model → QKV projection produces queries, keys, and values for standard attention → if sequence length W\leq W, only standard full attention is used (AHN is inactive) → if sequence length >W> W, the token at position tt attends within the window [tW+1,t][t-W+1, t] via standard attention, while the token at position tWt-W (just exiting the window) feeds its KV pair into the AHN to update the compressed memory → the current query qtq_t reads from both the compressed memory (via the AHN output projection) and the lossless window memory (via attention) → the two outputs are summed and fed into the subsequent MLP layer.

During training, a different flow is used (self-distillation): an input sequence is fed to both a frozen full-attention teacher model and the AHN-augmented student model → the student uses a random sliding window size (sampled from a set during training) → the KL divergence between teacher and student output distributions is computed → gradients flow backward to update only the AHN parameters, with the base model's weights remaining frozen.

3.3 Roadmap for the Deep Dive

  • First, the formal definition of the AHN recurrent update (Equation 3) and the memory integration (Equation 4) — this establishes the core computational abstraction: what the AHN ingests, what it produces, and how it interacts with attention.
  • Second, the AHN-GDN instantiation (Equations 5–7) in full detail — the most expressive AHN variant, which serves as the canonical example for understanding how AHNs actually transform KV pairs into compressed memory and generate outputs.
  • Third, the complexity analysis — a precise accounting of memory, FLOPs, and parameters compared to full attention, establishing exactly where the efficiency gains come from.
  • Fourth, the self-distillation training framework — including the KL divergence objective, the randomized window size training strategy, the parameter freeze vs. train split, and the specific training hyperparameters.
  • Fifth, the AHN-Mamba2 and AHN-DN alternative instantiations (Appendix A, Equations 9–10) — showing the framework's generality across recurrent architectures.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architectural design and empirical validation paper whose core idea is that a large sliding window attention combined with a small recurrent compressor for out-of-window memory achieves the efficiency of linear-attention models while preserving — and sometimes exceeding — the quality of full quadratic attention on long-context tasks.


Formal Definition of the AHN Recurrent Update

The AHN is defined by a single recurrence equation that describes how compressed memory is updated at each step:

htW=AHN((ktW,vtW),htW1)h_{t-W} = \text{AHN}((k_{t-W}, v_{t-W}), h_{t-W-1})

where tt is the current token position (1-indexed), WW is the sliding window size, (ktW,vtW)(k_{t-W}, v_{t-W}) is the key-value pair for the token that has just exited the sliding window (i.e., the token at position tWt-W), htW1h_{t-W-1} is the compressed memory state from the previous step (summarizing all tokens up to position tW1t-W-1), and htWh_{t-W} is the updated compressed memory state now including information from position tWt-W.

What it computes: the AHN is a function that ingests a single KV pair — the exact, lossless representation of a token that is about to be evicted from the sliding window — along with its own previous memory state, and produces an updated memory state that now incorporates that token's information in compressed form. This happens once per token position for every position t>Wt > W. The function AHN(,)\text{AHN}(\cdot, \cdot) is instantiated by specific recurrent architectures (Mamba2, DeltaNet, GatedDeltaNet), but the interface is uniform: two inputs (the KV pair to be compressed and the previous state), one output (the updated state). The compressed memory htWh_{t-W} is depicted as a fixed-size representation — it can be a vector or matrix — that does not grow with sequence length.

Why this form: the recurrence mirrors the classic RNN state update ht=f(xt,ht1)h_t = f(x_t, h_{t-1}) but with a crucial difference: the input is not the current token's embedding (which would be xtWx_{t-W}) but its KV pair (ktW,vtW)(k_{t-W}, v_{t-W}). This is a deliberate design choice. The K and V projections already encode information in a format optimized for attention-based retrieval — the key kk captures what the token can be queried on, and the value vv captures what should be retrieved. By feeding these exact projections into the AHN, the framework ensures that the compressed memory stores information in a representation space that is compatible with how the model's attention mechanism will later try to access it. An alternative would be to compress the token's hidden state or embedding directly, but this would create a representational mismatch between the lossless window memory (which stores KV vectors) and the compressed memory. The AHN approach maintains representational consistency.

The recurrence processes tokens one at a time as they exit the window, not in chunks. This token-wise design differs from Infini-attention's chunk-wise recurrence and allows the sliding window size to be configured arbitrarily at inference time without retraining — a property the paper exploits in its window size ablation (Figure 4) and randomization strategy during training.


Integration with Lossless Memory

Once the compressed memory htWh_{t-W} is updated, the model must use it alongside the lossless window memory to produce the output for the current token at position tt. This is defined as:

yt=f(htW,{(ki,vi)}i=tW+1t,qt)y_t = f(h_{t-W}, \{(k_i, v_i)\}_{i=t-W+1}^t, q_t)

where htWh_{t-W} is the compressed memory from the AHN (summarizing tokens 1 through tWt-W), {(ki,vi)}i=tW+1t\{(k_i, v_i)\}_{i=t-W+1}^t is the set of W key-value pairs currently within the sliding window (tokens tW+1t-W+1 through tt, the most recent WW tokens), and qtq_t is the query vector derived from the current token embedding xtx_t.

What it computes: the function ff combines information from compressed long-term memory, lossless short-term memory, and the current query to produce the token-mixing output yty_t for position tt. In the concrete instantiation (Equation 7), ff is realized as a simple summation: the compressed memory is queried via qtq_t and transformed by an output projection, the lossless memory is accessed via standard causal attention within the window, and the two outputs are added. The summation means both memory pathways contribute additively to the final representation — neither pathway is gated or thresholded at the output level.

Why this form: the additive combination allows the model to use both memory sources for every prediction. If a piece of information is present in the window memory (i.e., it was encountered in the last WW tokens), the attention mechanism can retrieve it with perfect fidelity through the standard attention pathway. If the information is older and only resides in the compressed memory, the AHN pathway provides a compressed-but-accessible representation that can still influence the output. The summation means the model doesn't need to learn which memory source to consult — it can draw from both, with the relative contribution determined by the learned parameters of each pathway (the attention weights for the lossless path, and the gate γ(xt)\gamma(x_t) and projection WoW_o for the compressed path).

A critical practical detail: the paper also supports attention sinks (Xiao et al., 2024), where a small number of initial tokens (default 128) are retained as lossless memory in addition to the sliding window. In this configuration, illustrated in Appendix Figure 6, the lossless memory consists of SS sink tokens + WW window tokens, and the AHN compresses only tokens that exit this combined lossless region. The attention sink mechanism prevents attention distributions from collapsing when operating over very long sequences by providing stable initial tokens that serve as attention "anchors."


AHN-GDN Instantiation: Gated Delta Rule Compression

The paper's most expressive AHN instantiation is based on GatedDeltaNet (Yang et al., 2025), a modern linear recurrent architecture that combines the delta rule for memory updates with input-dependent gating. The AHN-GDN compresses the KV pair (ktW,vtW)(k_{t-W}, v_{t-W}) into the memory state htWh_{t-W} using a gated delta rule update (Equation 5):

htW=α(xtW)(Iβ(xtW)ktWTktW)htW1+β(xtW)ktWTvtWh_{t-W} = \alpha(x_{t-W})(I - \beta(x_{t-W}) k_{t-W}^T k_{t-W}) h_{t-W-1} + \beta(x_{t-W}) k_{t-W}^T v_{t-W}

where:

  • xtWRDx_{t-W} \in \mathbb{R}^D is the token embedding at the position being compressed, with DD being the hidden dimension,
  • ktWRHk_{t-W} \in \mathbb{R}^H is the key vector for that position, with HH being the per-head dimension,
  • vtWRHv_{t-W} \in \mathbb{R}^H is the value vector for that position,
  • htW1RH×Hh_{t-W-1} \in \mathbb{R}^{H \times H} is the compressed memory state from the previous step, a square matrix of head dimension HH,
  • htWRH×Hh_{t-W} \in \mathbb{R}^{H \times H} is the updated compressed memory state,
  • α(xtW)(0,1)\alpha(x_{t-W}) \in (0, 1) is a scalar forget gate (also called the decay factor), computed as α(xtW)=σ(WαTxtW)\alpha(x_{t-W}) = \sigma(W_\alpha^T x_{t-W}) where WαRD×1W_\alpha \in \mathbb{R}^{D \times 1} is a learnable weight vector per head and σ\sigma is the sigmoid function,
  • β(xtW)(0,1)\beta(x_{t-W}) \in (0, 1) is a scalar input gate (also called the learning rate), computed as β(xtW)=σ(WβTxtW)\beta(x_{t-W}) = \sigma(W_\beta^T x_{t-W}) where WβRD×1W_\beta \in \mathbb{R}^{D \times 1} is another learnable weight vector per head,
  • II is the H×HH \times H identity matrix,
  • ktWTktWRH×Hk_{t-W}^T k_{t-W} \in \mathbb{R}^{H \times H} is the outer product of the key vector with itself, producing a rank-1 matrix,
  • ktWTvtWRH×Hk_{t-W}^T v_{t-W} \in \mathbb{R}^{H \times H} is the outer product of the key and value vectors, also rank-1.

What it computes: this equation performs a two-part memory update in each step.

Part 1 — Forgetting: The term α(xtW)(Iβ(xtW)ktWTktW)htW1\alpha(x_{t-W})(I - \beta(x_{t-W}) k_{t-W}^T k_{t-W}) h_{t-W-1} first scales the previous memory htW1h_{t-W-1} by a forget gate α\alpha, then applies a selective erasure via the delta-like correction matrix (IβktWTktW)(I - \beta k_{t-W}^T k_{t-W}). The correction subtracts a rank-1 projection of the memory onto the current key direction ktWk_{t-W}, weighted by the input gate β\beta. In plain language: the model first decays the old memory according to α\alpha (a scalar forget rate between 0 and 1, where smaller α\alpha means faster forgetting), and then further attenuates information in htW1h_{t-W-1} that is associated with the current key ktWk_{t-W} — essentially "making room" for new information about whatever this key represents. The amount of attenuation is controlled by β\beta: when β\beta is near 0, the correction is negligible and the old memory is preserved; when β\beta is near 1, a full projection of the memory onto kk is removed.

Part 2 — Writing: The term β(xtW)ktWTvtW\beta(x_{t-W}) k_{t-W}^T v_{t-W} adds a rank-1 update to the memory: the outer product of the current key and value vectors, scaled by the input gate β\beta. This is the classic Hebbian-like associative memory update: the memory now associates the key ktWk_{t-W} with the value vtWv_{t-W}, so that later queries qtq_t that are similar to ktWk_{t-W} (as measured by the dot product qtktWq_t k_{t-W}) will retrieve the stored value vtWv_{t-W}. The β\beta gate controls the "learning rate" of this association — how strongly the current KV pair gets written into memory.

The overall effect is a content-dependent, gated associative memory: the model decides (through the learned gates α\alpha and β\beta) how much to forget, how much to overwrite existing associations aligned with the current key, and how strongly to store the new key-value association. This is fundamentally different from fixed compression schemes like average or max pooling, which cannot adapt compression strength to token importance.

Why this form: the gated delta rule has two critical properties that make it well-suited for long-context compression over simpler recurrent updates (such as ht=λht1+ktTvth_t = \lambda h_{t-1} + k_t^T v_t with a fixed decay λ\lambda).

First, content-dependent forgetting: the gates α(xtW)\alpha(x_{t-W}) and β(xtW)\beta(x_{t-W}) are functions of the token embedding xtWx_{t-W}, meaning the model can learn to treat different tokens differently. A mathematical formula token might get a high input gate β\beta (strong write) and low forget gate α\alpha (slow decay) so it persists in memory for a long time. A function word or punctuation token might get a low β\beta (weak write) and high decay (fast forgetting), effectively being ignored by the compressed memory. This is the architectural mechanism that produces the gradient-based selectivity shown in Figure 5, where mathematical symbols show small gradients (well-preserved in memory) while pronouns and special tokens show large gradients (poorly represented). The paper confirms this through the probing experiment: "AHN tends to preserve the information of mathematical symbols and numbers while neglecting less critical ones such as pronouns and special tokens" (Section 3.5).

Second, associative removal via the delta rule: the term (IβktWTktW)(I - \beta k_{t-W}^T k_{t-W}) goes beyond simple exponential decay. Standard gated RNNs like the GRU or LSTM can forget information globally (via a forget gate) but cannot selectively remove associations to specific keys. The delta rule correction allows the model to say: "I'm storing new information associated with key kk, so any old information stored under that key should be overwritten." This prevents the memory from accumulating stale or contradictory associations over very long sequences. If a character's name appears with different descriptions across a long document, later descriptions can overwrite earlier ones sharing the same key representation, rather than the memory blindly averaging all descriptions.

AHN output computation: Once the compressed memory htWh_{t-W} is updated, the model must access it from the current query qtq_t. The AHN output pathway is defined by Equation 6:

yAHN,t=γ(xt)qthtWWoy_{\text{AHN},t} = \gamma(x_t) q_t h_{t-W} W_o

where:

  • γ(xt)(0,1)\gamma(x_t) \in (0, 1) is an output gate, computed as γ(xt)=σ(WγTxt)\gamma(x_t) = \sigma(W_\gamma^T x_t) with learnable weight WγRD×1W_\gamma \in \mathbb{R}^{D \times 1} per head,
  • qtR1×Hq_t \in \mathbb{R}^{1 \times H} is the query vector for the current position, treated as a row vector,
  • htWRH×Hh_{t-W} \in \mathbb{R}^{H \times H} is the compressed memory matrix,
  • WoRH×HW_o \in \mathbb{R}^{H \times H} is a per-head learnable output projection matrix, grouped by heads (each head has its own WoW_o),
  • yAHN,tR1×Hy_{\text{AHN},t} \in \mathbb{R}^{1 \times H} is the AHN's contribution to the token-mixing output for this head.

What it computes: the current query qtq_t is multiplied by the memory matrix htWh_{t-W}, producing a row vector qthtWR1×Hq_t h_{t-W} \in \mathbb{R}^{1 \times H}. This is the standard linear attention operation: qtq_t acts as the "address" to read from the associative memory htWh_{t-W}, whose stored content was built from outer products kTvk^T v of past keys and values. If qtq_t is similar to a past key ktWk_{t-W} (in dot product), the product qtktWTvtWq_t k_{t-W}^T v_{t-W} will retrieve a value similar to vtWv_{t-W}. The result is then projected by WoW_o (which can mix information across head dimensions and transform the representation) and gated by the scalar γ(xt)\gamma(x_t). The gate allows the model to learn when to attend to compressed memory versus rely on the window attention — γ(xt)\gamma(x_t) near 1 means "the compressed memory is important here," near 0 means "ignore it."

Why this form: the gate γ\gamma is per-token and content-dependent. This is essential because the utility of long-term compressed memory varies dramatically across tokens. For a token in the middle of a math derivation that depends on definitions stated at the very beginning, γ\gamma should be high so the compressed memory (which stored those definitions) contributes strongly. For a token in a routine connective phrase, γ\gamma should be near zero so the model relies only on immediate window context. The scalar gate (one value per head, not per dimension) is simple but expressive enough for this binary-like decision.

The final output for each head (Equation 7) simply adds the compressed memory pathway and the lossless window attention pathway:

yt=yAHN,t+Attention({(ki,vi)}i=tW+1t,qt)y_t = y_{\text{AHN},t} + \text{Attention}(\{(k_i, v_i)\}_{i=t-W+1}^t, q_t)

where Attention()\text{Attention}(\cdot) is standard causal multi-head attention restricted to the WW tokens in the sliding window, producing an output of the same dimension HH per head. The addition means the model never has to choose between the two memory sources — it can use both simultaneously, with their relative contributions learned rather than hard-coded.

Parameter count for AHN-GDN: for each attention head, the trainable parameters consist of WαRD×1W_\alpha \in \mathbb{R}^{D \times 1} (D scalars), WβRD×1W_\beta \in \mathbb{R}^{D \times 1} (D scalars), WγRD×1W_\gamma \in \mathbb{R}^{D \times 1} (D scalars), and WoRH×HW_o \in \mathbb{R}^{H \times H} (H2H^2 scalars, head dimension squared). With NqN_q attention heads, this gives 3DNq+H2Nq3DN_q + H^2 N_q learnable parameters. All other parameters (QKV projections, MLPs, layer norms, embeddings) remain frozen from the base model. The paper reports this as "approximately 0.4% relative to the frozen base model's parameters" (Section 2.4). For Qwen2.5-3B-Instruct, this is 13.0 million parameters out of 3 billion (Table 9). For Qwen2.5-7B-Instruct, it is 21.3 million (0.3%). For Qwen2.5-14B-Instruct, it is 61.0 million (0.4%).


Complexity Analysis

The paper provides a precise accounting of how AHN alters the computational and memory profile of the attention mechanism. Table 1 in the paper summarizes the complexity for vanilla full attention versus window attention augmented with AHN-GDN across three dimensions: parameter count, memory cache size, and FLOPs. We'll walk through each carefully.

Parameter complexity:

  • Full attention: 2DH(Nq+Nkv)2DH(N_q + N_{kv}) parameters from the Q, K, V, and output projection matrices, where DD is hidden dimension, HH is head dimension, NqN_q is number of query heads, and NkvN_{kv} is number of key-value heads (for grouped-query attention, NkvNqN_{kv} \leq N_q).
  • Window attention + AHN-GDN: the same 2DH(Nq+Nkv)2DH(N_q + N_{kv}) for the base attention projections, plus 3DNq+H2Nq3DN_q + H^2 N_q from the three per-head gate vectors (WαW_\alpha, WβW_\beta, WγW_\gamma, each of size D×1D \times 1) and the per-head output projection WoRH×HW_o \in \mathbb{R}^{H \times H}.

The extra parameters are small because the gate vectors project from the full hidden dimension DD to a scalar (one parameter per input dimension per gate per head), and the output projection is head-dimension-squared rather than hidden-dimension-squared. For Qwen2.5-3B-Instruct with D=2048D=2048, H=128H=128, Nq=32N_q=32, the extra parameters are 3×2048×32+1282×32=196, ⁣608+524, ⁣288=720, ⁣8960.723 \times 2048 \times 32 + 128^2 \times 32 = 196,\!608 + 524,\!288 = 720,\!896 \approx 0.72M, though the actual reported number is 13.0M (Table 9), likely due to biases and implementation details.

Memory cache complexity:

Memory cache refers to the GPU memory required to store the attention's KV cache during autoregressive generation. This is typically the dominant memory consumer for long sequences.

  • Full attention: 2LHNkv2 L H N_{kv} bytes (assuming 2-byte float16), where LL is the current sequence length. This grows linearly with LL, hence O(L)\mathcal{O}(L).
  • Window attention + AHN-GDN: 2WHNkv+H2Nq2 W H N_{kv} + H^2 N_q. The first term is the KV cache for the sliding window (size WW, not LL), and the second term is the AHN's compressed memory state htWh_{t-W} — one H×HH \times H matrix per head per layer. Since WW is fixed (e.g., 32k) and H2NqH^2 N_q is constant, the total is O(W)\mathcal{O}(W), not O(L)\mathcal{O}(L). The compressed memory state size is independent of sequence length.

At 128k tokens, for Qwen2.5-3B-Instruct, the full attention KV cache is 9.44 GB, while the AHN-augmented version uses only 2.45 GB — a 74.0% reduction (Table 9, Figure 1b). The 2.45 GB consists of 2.42 GB for the 32k-token window KV cache (including 128 sink tokens) plus approximately 0.03 GB for the AHN memory states hh. The ratio is 25.6% of the original cache size, which the paper rounds to 26.0%.

Computational complexity (FLOPs):

The paper counts only matrix multiplication FLOPs, omitting softmax, normalization, and element-wise operations (as stated in Table 1 footnote).

  • Full attention: 4LDH(Nq+Nkv)+2HNqL24LDH(N_q + N_{kv}) + 2H N_q L^2. The first term is the cost of QKV projections and output projection (linear in LL), and the second term is the attention score computation QKTQK^T and the weighted sum AVAV (quadratic in LL). At L=128kL=128k, the quadratic term dominates, making the complexity effectively O(L2)\mathcal{O}(L^2).
  • Window attention + AHN-GDN: 4LDH(Nq+Nkv)+2HNqW2+2(LW)×(2WHNq+H2Nq+3DNq+H2Nq)4LDH(N_q + N_{kv}) + 2H N_q W^2 + 2(L - W) \times (2W H N_q + H^2 N_q + 3D N_q + H^2 N_q). This decomposes into three parts:
    1. 4LDH(Nq+Nkv)4LDH(N_q + N_{kv}): same linear projection cost as full attention (all tokens still need QKV projections).
    2. 2HNqW22H N_q W^2: attention score computation within the sliding window (quadratic in WW, not LL).
    3. 2(LW)×(2WHNq+H2Nq+3DNq+H2Nq)2(L - W) \times (2W H N_q + H^2 N_q + 3D N_q + H^2 N_q): the cost of querying and updating the AHN compressed memory for each of the LWL-W tokens beyond the window. The per-token cost includes: reading from memory (2WHNq2W H N_q, multiplying qtq_t by htWh_{t-W}), writing to memory (H2NqH^2 N_q for the outer product updates), gating (3DNq3D N_q for the three gate projections), and output projection (H2NqH^2 N_q). The dominant term here for large LL is the factor (LW)(L-W), making the total O(WL)\mathcal{O}(WL) — effectively O(L)\mathcal{O}(L) for fixed WW.

At 128k tokens, for Qwen2.5-3B-Instruct, full attention's token mixer performs 2.50×10152.50 \times 10^{15} FLOPs, while AHN-augmented attention performs 1.17×10151.17 \times 10^{15} FLOPs — a 53.3% reduction in attention FLOPs (Table 9). The full model FLOPs (including MLPs) go from 3.29×10153.29 \times 10^{15} to 1.95×10151.95 \times 10^{15} — a 40.5% reduction, since MLP FLOPs are identical in both cases. The mixing FLOP ratio (AHN-augmented attention FLOPs divided by full attention FLOPs) is 46.7%, and the model FLOP ratio is 59.4% (Table 2, Table 9).

Why this complexity profile matters: the reduction from O(L2)\mathcal{O}(L^2) to O(WL)\mathcal{O}(WL) changes the asymptotic behavior dramatically. As LL grows beyond 32k, the full attention cost keeps growing quadratically, while the AHN-augmented model's cost grows only linearly — the gap widens indefinitely. Figure 3a visualizes this: at 250k sequence length, the FLOPs gap is enormous (roughly 4× difference). Similarly, Figure 3b shows the memory cache staying flat for the AHN model while climbing linearly for full attention. The practical consequence is that AHN-augmented models can process sequences of essentially arbitrary length with bounded memory and linear-time computation, while full attention models hit hardware limits.

A critical design note: the complexity analysis in Table 1 marks certain terms in gray, noting they "can be further omitted compared to the other terms." For instance, 2WHNq2W H N_q in the AHN per-token cost is negligible compared to H2NqH^2 N_q when HH is large (e.g., H=128H=128, so H2=16, ⁣384H^2 = 16,\!384 dominates over 2W=65, ⁣5362W = 65,\!536 for practical window sizes). This means the per-token AHN cost is dominated by the matrix multiplications involving the H×HH \times H memory state, not by the attention-like readout.


Self-Distillation Training Framework

Instead of training AHN-augmented models from scratch or fine-tuning all parameters, the paper uses a self-distillation approach that freezes the pretrained base model and trains only the AHN parameters to mimic the full-attention teacher. This is a key practical contribution: it enables rapid training on modest compute and preserves the base model's existing capabilities.

Teacher and student definitions:

  • Teacher model: the original open-weight LLM (e.g., Qwen2.5-Instruct) with standard full attention over the entire sequence. Its output probability distribution over the vocabulary for a given input sequence is denoted pp'. The teacher processes the full context with quadratic attention, providing a "gold standard" of what a lossless-memory model would predict. All parameters of the teacher are frozen and never updated.

  • Student model: the same base LLM, but with its attention mechanism modified to operate over only a sliding window of size WW (randomized during training) and augmented with AHNs. Its output probability distribution is denoted pp. The student's window attention can only see the most recent WW tokens directly; all older information must flow through the AHN's compressed memory. Crucially, the student shares all non-AHN parameters with the teacher — QKV projections, MLPs, embeddings, layer norms — and these shared parameters are frozen. Only the AHN parameters (the gate weights and output projections described in Section 2.3) are updated.

Training objective:

The student is trained to minimize the Kullback-Leibler (KL) divergence between the teacher's output distribution and its own:

l=KL(pp)=vVpvlog(pvpv)l = \text{KL}(p' \| p) = \sum_{v \in \mathcal{V}} p'_v \log\left(\frac{p'_v}{p_v}\right)

where V\mathcal{V} is the vocabulary, pvp'_v is the teacher's predicted probability for token vv, and pvp_v is the student's predicted probability for token vv.

What it computes: the KL divergence measures how much information is lost when using the student's distribution pp to approximate the teacher's distribution pp'. For each position in the sequence, the student is penalized if it assigns low probability to tokens that the teacher assigns high probability to (the pvlogpv/pvp'_v \log p'_v / p_v terms), and more severely penalized if the student assigns high probability to tokens the teacher assigns near-zero probability to (since log(1/pv)\log(1/p_v) explodes as pv0p_v \to 0 when pvp'_v is non-negligible). Minimizing KL divergence forces the student's entire output distribution to match the teacher's, not just the argmax prediction.

Why this form (KL divergence rather than cross-entropy): the teacher's distribution pp' captures rich information about token probabilities, not just the correct token. For example, at a position where multiple continuations are plausible, the teacher might assign 0.4 probability to "the," 0.3 to "a," and 0.2 to "some." The student needs to learn to produce a similar spread, not just pick the top-1 token. Cross-entropy loss with ground-truth next-token labels (the standard language modeling objective) would only train the student to predict the single correct token, providing a sparse signal. The paper confirms this empirically: in Table 4, replacing self-distillation (KL loss) with standard next-token prediction (CE loss) on the same training data causes the LongBench average to drop from 40.59 to 39.59 for the Qwen2.5-7B model with AHN-GDN. The authors hypothesize "CE provides sparse learning signals, and pushes the small AHN modules towards shortcuts in the training data. In contrast, self-distillation offers denser guidance over the teacher's entire output distribution" (Section 3.4).

An additional practical benefit: self-distillation does not require training labels. The teacher generates the target distribution on-the-fly for any input sequence, so the training data can be any text corpus — in this paper, the ChatQA2 dataset of diverse long-context tasks, consisting of 1B tokens.

Randomized sliding window size during training:

To ensure the AHN learns a generalizable compression strategy rather than overfitting to a specific window size, the paper randomizes both the window size and the attention sink size for each training example (Section 3.1):

  • Attention sink size is uniformly sampled from the set [0,32,64,128,512,2048,4096][0, 32, 64, 128, 512, 2048, 4096] after removing any candidate larger than half the sequence length of the current example.
  • Total lossless memory size (sinks + sliding window) is uniformly sampled from [32,64,128,256,512,1024,2048,4096,8192][32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] after filtering out values smaller than one-eighth of the sequence length.
  • The sliding window size is then computed as (total lossless size minus sink size) for each example.

This randomization strategy exposes the AHN to many different window configurations during training. Sometimes the window is very small (e.g., 32 tokens), forcing the AHN to compress almost the entire sequence and carry most of the information burden. Other times the window is large (e.g., 8192 tokens), meaning the AHN compresses relatively little and the window attention handles most of the context. By training across this range, the AHN learns to be robust to the window size — it can operate as a heavy compressor (small window) or a light supplement (large window). The paper verifies this in Table 4: training with a fixed window size of 1024 produces a LongBench score of 38.53, compared to 40.59 for randomized window training, confirming that the fixed-size model "often overfit[s] to that specific configuration and fail[s] to generalize to unseen context lengths" (Section 3.4).

Training hyperparameters and data:

  • Training data: ChatQA2 dataset (Xu et al., 2025), an open-source collection of diverse long-context tasks, totaling 1B tokens.
  • Maximum training sequence length: 24k tokens (the model generalizes to 128k at inference — nearly 5.3× longer than training).
  • Optimizer: AdamW (Loshchilov and Hutter, 2019).
  • Learning rate: 1×1041 \times 10^{-4}, warmed up linearly over the first 10% of steps, then cosine decayed to zero.
  • Global batch size: 128 sequences.
  • Total update steps: 740 (one epoch over 1B tokens at batch size 128).
  • Training compute: approximately 10 hours on 32 A100 GPUs for the 7B model (stated in Section 3.1 and Appendix D).
  • Implementation: PyTorch (Paszke et al., 2019), built on LLaMA-Factory (Zheng et al., 2024) and Flash Linear Attention (Yang and Zhang, 2024).

Parameter efficiency of training: only the AHN parameters (approximately 0.4% of the base model) are updated. All QKV projections, MLP weights, embeddings, and layer norms are frozen. This means the optimizer states (momentum and variance buffers in AdamW) only need to be maintained for the tiny AHN parameter set, dramatically reducing GPU memory during training. The paper reports the FLOPs per training step in Table 7: for the 7B model with full attention, one step costs 1.1519×10171.1519 \times 10^{17} FLOPs under the training configuration (batch size 128, seq length 24k, window size 8k). The AHN-augmented model with frozen base weights costs 1.0359×10171.0359 \times 10^{17} FLOPs, essentially identical to sliding window attention alone (1.0334×10171.0334 \times 10^{17}). This confirms that the AHN adds negligible training overhead.

Why this training strategy over alternatives: training hybrid models from scratch (as in some prior work) would require the full pretraining compute budget — prohibitive for most practitioners. Full fine-tuning of the base model would risk catastrophic forgetting of its pretrained capabilities and would be slower. The paper's approach of freezing the base model and only training tiny AHN modules leverages the substantial investment in pretraining while adapting the model for long-context efficiency. The self-distillation objective ensures the AHN learns to approximate the full-attention model's behavior rather than developing its own (potentially different) inductive biases.

Inference-time configuration: at inference, AHNs are configured with a default sliding window size of 32,768 tokens and 128 attention sink tokens. These values are fixed and not randomized. The 32k window is chosen as the sweet spot in the ablation study (Figure 4): performance improves as the window grows from 1k to 16k, plateaus around 32k, and eventually declines beyond 64k due to attention dilution (where the attention distribution becomes overly diffuse across too many keys). The paper describes the 32k default as "substantially larger than those used in prior attention–RNN hybrid methods (e.g., 64 in [41, 112])" (Section 1).

A critical practical point: when sequence length L32, ⁣768L \leq 32,\!768, the AHN never activates. The model "behaves exactly as a standard full-attention model" (Appendix H), so there is zero performance difference from the base model on short-context tasks. The AHN only affects behavior when L>32, ⁣768L > 32,\!768, precisely the regime where full attention becomes quadratically expensive and memory-constrained.


AHN-Mamba2 and AHN-DN Alternative Instantiations

To demonstrate the architectural generality of the AHN framework, the paper instantiates AHNs with two other modern recurrent architectures beyond GatedDeltaNet. All three share the same interface — they take (ktW,vtW)(k_{t-W}, v_{t-W}) and htW1h_{t-W-1} as inputs, produce htWh_{t-W} as output, and use the same output pathway (Equation 6) — but differ in their internal state update mechanics. The full equations are provided in Appendix A.

AHN-Mamba2 (Equation 9):

htW=exp(Δ(xtW)A)htW1+Δ(xtW)ktWTvtWh_{t-W} = \exp(-\Delta(x_{t-W}) A) h_{t-W-1} + \Delta(x_{t-W}) k_{t-W}^T v_{t-W}

where Δ(xtW)R+\Delta(x_{t-W}) \in \mathbb{R}^+ is a scalar step size computed from the token embedding (via a learned projection followed by a softplus activation to ensure positivity), and AR+A \in \mathbb{R}^+ is a fixed scalar decay rate (typically a learnable parameter initialized to some positive value, treated as a per-head constant).

What it computes: this is a state-space model (SSM) discretization of the continuous-time system dhdt=Ah+kTv\frac{dh}{dt} = -A h + k^T v. The first term exp(ΔA)htW1\exp(-\Delta A) h_{t-W-1} decays the previous memory by a factor that depends on both the learned decay rate AA and the input-dependent step size Δ\Delta. A larger Δ\Delta (faster "time step") causes more forgetting of the old memory (since exp(ΔA)\exp(-\Delta A) is smaller when ΔA\Delta A is larger). The second term ΔktWTvtW\Delta k_{t-W}^T v_{t-W} writes the new key-value association into memory, scaled by the step size — larger Δ\Delta means both more forgetting of the past and a stronger write of the current token.

Why this form: the Mamba2 recurrence is a structured state-space model that has been shown to be highly efficient for long-sequence processing when combined with selective (input-dependent) step sizes. Unlike the delta rule in AHN-GDN, which can selectively erase specific key directions via the (IβkTk)(I - \beta k^T k) correction, the Mamba2 forgetting is isotropic — everything in the memory state decays by the same factor exp(ΔA)\exp(-\Delta A). This makes it simpler but potentially less precise: it cannot selectively preserve some information while forgetting other information associated with the same key. However, Mamba2 benefits from highly optimized CUDA kernels and hardware-efficient implementations through the structured state-space duality framework (Dao and Gu, 2024), making it very fast in practice.

AHN-DN (Equation 10):

htW=(Iβ(xtW)ktWTktW)htW1+β(xtW)ktWTvtWh_{t-W} = (I - \beta(x_{t-W}) k_{t-W}^T k_{t-W}) h_{t-W-1} + \beta(x_{t-W}) k_{t-W}^T v_{t-W}

where β(xtW)(0,1)\beta(x_{t-W}) \in (0, 1) is a scalar input gate, computed identically to AHN-GDN as β(xtW)=σ(WβTxtW)\beta(x_{t-W}) = \sigma(W_\beta^T x_{t-W}).

What it computes: this is AHN-GDN without the forget gate α\alpha. The update consists only of the delta-rule correction and the new key-value association, both scaled by the same gate β\beta. When β\beta is near 0, the memory state is preserved unchanged; when β\beta is near 1, the memory state is updated with both the removal of old key-associated information and the writing of new information. There is no separate mechanism for gradual decay — the only way information leaves memory is by being overwritten by the delta rule correction (IβkTk)(I - \beta k^T k).

Why this form: AHN-DN is simpler than AHN-GDN (one gate instead of two), and corresponds to the original DeltaNet architecture (Schlag et al., 2021; Yang et al., 2024). The lack of a separate forget gate means the memory has less flexibility: it cannot globally decay all information (e.g., at a section boundary or topic shift) independently of writing new information. However, it may be easier to train with fewer parameters and a simpler gating structure. In practice, the paper's results (Tables 2 and 3) show that AHN-DN slightly outperforms AHN-GDN on some tasks (e.g., LV-Eval average for 3B: AHN-DN 5.68 vs AHN-GDN 5.88 — here GDN wins; but for 7B, AHN-DN 6.82/16.48 vs AHN-GDN 6.54/16.93 — DN wins on LV-Eval while GDN wins on InfiniteBench), suggesting the optimal instantiation may be task-dependent.

Comparison across instantiations: Table 2 shows that all three AHN variants consistently outperform sliding window baselines. For Qwen2.5-3B-Instruct on LV-Eval and InfiniteBench (averaging the two): AHN-Mamba2 achieves 8.79, AHN-DN achieves 9.61, AHN-GDN achieves 9.56 — all above SWA's 7.53. The differences between AHN variants are small relative to the gap between any AHN variant and the sliding window baseline, confirming that the core AHN framework (the decomposition into lossless window + compressed out-of-window memory) matters more than the specific recurrent architecture. The number of extra parameters is also similar: 0.4% for Mamba2 and DN, 0.4% for GDN on 3B; 0.2% for Mamba2 and DN, 0.3% for GDN on 7B.

Important note: all three AHN variants use the same output mechanism (Equation 6) with the gate γ(xt)\gamma(x_t) and the output projection WoW_o. The differences are entirely in the recurrent state update (comparing Equations 5, 9, and 10). This modularity — where the compression mechanism and the readout mechanism are decoupled — is part of the framework's design and allows for easy experimentation with new recurrent architectures.


Summary of Key Design Choices

  1. Token-wise rather than chunk-wise recurrence: the AHN compresses one KV pair at a time, as each token exits the sliding window, rather than operating on chunks of tokens at once. This enables arbitrary window size configurations without architectural changes, and keeps the recurrent update lightweight (one rank-1 outer product per step).

  2. Compressing KV pairs rather than hidden states: the AHN ingests the exact key and value vectors from the attention projection, ensuring representational compatibility between the compressed memory and the lossless memory. The model doesn't need to learn separate projections for memory compression.

  3. Additive combination of memory pathways: the outputs from compressed memory and sliding window attention are summed, not concatenated or gated at the output level. This lets the model use both sources for every prediction, with relative contributions determined by learned parameters (the gate γ\gamma for the compressed pathway, the attention weights for the lossless pathway).

  4. Large default window size of 32k tokens: the AHN only activates when sequences exceed this window, preserving full attention performance on short and medium sequences where attention is already efficient and effective. This is the paper's most distinctive architectural position compared to prior hybrid methods.

  5. Self-distillation with frozen base model: training only the AHN parameters (approximately 0.4% of total parameters) using KL divergence against the full-attention teacher, with randomized window sizes to ensure generalization. This is computationally efficient (approximately 10 hours on 32 A100 GPUs for 7B models) and preserves the base model's capabilities.

  6. Content-dependent gating: all AHN variants use input-dependent gates (α\alpha, β\beta, γ\gamma) computed from the token embedding, enabling the model to learn token-level decisions about what to store, what to forget, and when to query compressed memory. This is the mechanism that produces the selective preservation shown in Figure 5.

4. Key Insights and Innovations

Innovation 1: The Large-Window Philosophy — AHNs Activate Only When Attention Fails, Not Always

The paper's most consequential conceptual move is not that we can hybridize attention and recurrence — many prior works have intermixed Transformer and RNN layers — but when the recurrence should activate. Prior hybrid approaches (Jamba, Infini-attention, LoLCATs, HQLT) embed RNN components as an always-on part of the architecture, operating alongside attention at every sequence position regardless of length. This is the natural engineering instinct: if recurrence helps on long sequences, activate it everywhere so it's always available.

The AHN framework rejects this instinct with a sharply argued counter-position: quadratic attention is not a problem on short sequences, so recurrence should remain dormant until the sequence length crosses the threshold where attention becomes compute-bound. The paper sets this threshold at 32k tokens — substantially larger than the 64–2048 token windows used in prior hybrid work (Appendix F). Below 32k, the model is literally identical to the base Transformer: the AHN code path never executes, no compressed memory is allocated, and performance matches the full-attention teacher exactly. Above 32k, the AHN activates to compress tokens as they exit the window, converting the O(L2)\mathcal{O}(L^2) attention cost into O(WL)\mathcal{O}(WL) with W=32, ⁣768W = 32,\!768.

This framing constitutes a genuine conceptual shift because it reframes the problem from "how do we make recurrence work well everywhere?" to "how do we make recurrence work well specifically in the regime where attention is too expensive?" The former question leads to architectures where recurrence permeates every layer and must match attention's quality across all tasks — a tall order given RNNs' documented limitations on exact recall (Wen et al., 2025). The latter question leads to a division of labor: attention handles the regime where it's both efficient and effective; recurrence takes over only where attention's cost is prohibitive and some information loss from compression is an acceptable trade-off.

The evidence that this matters is not just philosophical — it appears in the ablation on inference window size (Figure 4). Performance improves steadily as the window grows from 1k to 16k on both LV-Eval (128k subset) and InfiniteBench (128k subset), confirming that more lossless memory is better when you can afford it. The 32k default is the empirical sweet spot, and the decline beyond 64k (attributed to attention dilution) suggests that unbounded lossless memory has its own failure modes, not just efficiency problems. The paper is implicitly arguing that the right architecture is not "attention everywhere, compressed where necessary" but "attention where possible, compressed where necessary" — a subtle but important difference in design philosophy.

This is a fundamental shift rather than an incremental refinement because it changes the optimization objective. Prior hybrid methods optimize for "how well does the RNN approximate attention?" on all tokens. AHN optimizes for "how well does the RNN approximate attention specifically on tokens older than 32k?" — a much easier learning problem because the model has 32k tokens of perfect context to condition on before needing to rely on compressed memory. The 32k window provides rich conditioning information that makes the compression task more tractable: the AHN compresses old tokens into a form that supplements, rather than replaces, the recent context.


Innovation 2: Memory as a Cognitive Framework, Not Just an Implementation Detail

The paper's explicit cognitive-science framing — mapping the Multi-Store Model (MSM) of Atkinson and Shiffrin (1968) onto neural architectures — is not a decorative metaphor. It is a diagnostic tool that clarifies which information each memory store should handle and why the architecture is structured the way it is.

Prior work on memory management in Transformers (KV cache eviction, segment-level recurrence, linear attention) operates in an engineering paradigm: the goal is to reduce memory and compute while minimizing quality degradation. The design space is explored through ablations and benchmarks, but there is rarely a principled account of what kinds of information different memory mechanisms are well-suited to store. This leads to approaches that apply uniform compression or eviction strategies regardless of content — Compressive Transformers pool all old tokens with the same operation, KV cache eviction methods drop tokens based on generic attention-based importance scores.

AHN's mapping to MSM provides a functional decomposition that is absent from these approaches:

  • Short-term memory (the sliding window): high-fidelity, limited capacity, holding exact representations of recent tokens. This corresponds to the psychological concept of working memory, which can hold roughly 7±2 items with high precision but decays rapidly (Miller, 1956; Peterson, 1959).
  • Long-term memory (the AHN state): compressed, semantically organized, unbounded duration, but lacking exact detail. This corresponds to hippocampal consolidation, which transforms episodic memories into cortical representations that capture gist and structure while losing surface-level specifics (Scoville and Milner, 1957; McClelland et al., 1995).

This decomposition explains why the architecture should work the way it does, not just that it does. Short-term memory is for exact recall of recent tokens — the model needs to know the precise wording of the last few paragraphs, maintain syntactic coherence, and track local discourse structure. Long-term memory is for semantic summarization and retrieval of distant information — the model needs to know that a character was introduced 50 pages ago and is now relevant again, but doesn't need to recall the exact phrasing of the introduction. This maps cleanly onto the empirical finding that AHN-augmented models match or exceed full attention on LV-Eval (which tests comprehension of long documents) but underperform on exact-recall needle-in-a-haystack tasks from RULER (Table 5) — exactly the pattern MSM would predict.

The cognitive framing also motivates the continual, learnable nature of the compression: the hippocampus doesn't use a fixed pooling function, and neither should the AHN. The content-dependent gating (α\alpha, β\beta, γ\gamma) is the architectural realization of the idea that some information is more consolidation-worthy than others — a mathematical formula should persist in long-term memory, while a connective phrase should be forgotten. The gradient probing experiment in Figure 5 (Section 3.5) provides direct evidence that this happens: mathematical symbols and numbers show low gradient magnitudes (well-preserved in the AHN's compressed memory), while pronouns and special tokens show high gradients (poorly represented or quickly forgotten). This is exactly the kind of selective consolidation that a cognitive framework predicts.

This is a fundamental intellectual contribution because it provides a vocabulary and conceptual structure for reasoning about memory in neural architectures. Rather than treating all memory as fungible (as in approaches that uniformly compress or evict tokens), the MSM mapping gives principled answers to design questions: What should be stored in lossless memory? Recent, high-precision information. What should be compressed? Old, semantically important information. How should compression work? Continually, with content-dependent selectivity. These answers aren't derived from an ablation study — they follow from the cognitive framework and are validated by experiments, not generated by them.


Innovation 3: Self-Distillation as a Training Paradigm That Preserves Pretrained Capabilities While Teaching Compression

Training hybrid attention-RNN models presents a fundamental tension: you want to teach the recurrent component to compress information effectively, but you don't want to disturb the carefully optimized attention-based capabilities of the pretrained model. Prior work resolves this tension in ways that are either expensive (training from scratch, as in HQLT) or risky (full fine-tuning or distillation of all parameters, as in MiL which trains the entire token mixer). Both approaches require substantial compute and risk catastrophic forgetting of the base model's capabilities.

The paper's self-distillation approach — freezing all base model parameters and training only the AHN parameters (approximately 0.4% of total) using KL divergence against the full-attention teacher — is not just computationally efficient. It represents a training paradigm that solves the tension by teaching the student to approximate the teacher rather than replace it, and by constraining the approximation to only the newly added parameters.

What makes this distinctive is the density of the learning signal. Standard next-token prediction (cross-entropy loss) provides a single scalar loss per position: was the predicted token correct? This is a sparse signal for a tiny module trying to learn complex compression behavior. Self-distillation with KL divergence provides a loss over the entire vocabulary distribution at every position: not just "did you get the right token?" but "how closely does your full probability distribution match the teacher's full distribution?" The paper's ablation (Table 4) confirms this matters substantially: switching from KL to CE drops LongBench performance from 40.59 to 39.59 (Qwen2.5-7B, AHN-GDN) — a non-trivial degradation that the authors attribute to CE's sparsity pushing "the small AHN modules towards shortcuts in the training data" (Section 3.4).

The randomized window size training strategy is a subtler but equally important contribution to this paradigm. By sampling window sizes from a diverse set (32 to 8192 tokens) during training, the AHN learns to operate across a spectrum from "carry almost all information" (tiny window) to "mostly just supplement attention" (large window). This makes the inference-time window configurable without retraining — the ablation in Figure 4 shows the model works well across window sizes from 1k to 96k — and it prevents the AHN from developing brittle, window-specific compression strategies. Training with a fixed window (1024 tokens, Table 4) causes a notable performance drop (38.53 vs. 40.59), confirming that generalization requires randomized exposure.

This is an incremental but practically important innovation. The individual components (KL distillation, frozen base model, randomized training hyperparameters) are not individually novel, but their combination creates a training recipe that is dramatically cheaper than alternatives (approximately 10 hours on 32 A100 GPUs for a 7B model, using only 1B tokens) while producing models that generalize to sequences 5.3× longer than training (from 24k training length to 128k inference length). The low training cost and modest data requirements (1B tokens vs. 15–20B for alternatives like MiL and HQLT) make the approach accessible to practitioners who cannot afford full-scale pretraining.

The significance goes beyond the specific numbers. This training paradigm demonstrates that you can add efficient long-context capabilities to an existing LLM without touching its pretrained weights, treating compression as a learnable add-on rather than a fundamental architectural redesign. This is a deployment-friendly philosophy: the base model's safety properties, factual knowledge, and instruction-following behavior are preserved exactly (since its weights are unchanged), while the AHN modules add long-context efficiency as a transparent extension. For practitioners deploying models in production, the guarantee that short-context behavior is identical to the validated base model (since AHNs never activate below 32k) is a significant risk-reduction property that training-from-scratch or full-fine-tuning approaches cannot offer.


Innovation 4: The Identification of Attention Dilution as a Failure Mode of Unbounded Windows

The paper's ablation on inference window size (Figure 4) reveals a non-obvious empirical phenomenon: as the sliding window grows beyond 64k tokens for LV-Eval (or 96k for InfiniteBench), performance declines rather than continuing to improve or plateau. The paper attributes this to "the attention-dilution effect, where the attention distribution becomes overly diffuse when the number of keys grows very large, weakening the model's ability to focus on relevant information" (Section 3.4).

This is a conceptually important finding because it challenges the implicit assumption underlying most long-context research: that more lossless context is always better, and the only limitation is computational cost. If attention dilution is real, then even with infinite compute, unbounded attention windows would eventually hurt performance — the attention mechanism itself has a sweet spot beyond which adding more keys reduces its ability to focus.

This finding is not the paper's primary contribution, but it has significant implications for architectural design. It provides a principled justification for the hybrid approach beyond mere efficiency: even if you could afford quadratic attention on 128k-token sequences, you might not want to, because the attention mechanism's selectivity degrades. The hybrid architecture solves not just an engineering problem (cost) but a potential quality problem (dilution) by keeping the attention window at a size where it operates well and routing older information through a different mechanism (compressed memory) that doesn't suffer from dilution.

The specific window sizes where dilution occurs — after 64k on LV-Eval and after 96k on InfiniteBench — are task-dependent and likely related to how much relevant information is distributed across the sequence length. A task requiring integration across the entire 128k context might suffer from dilution sooner than one where most relevant information is concentrated. The paper doesn't deeply analyze the causes of dilution (it merits future work), but the observation itself is valuable because it adds a quality ceiling to the standard efficiency argument for long-context architectures.

This is an incremental empirical finding rather than a fundamental theoretical advance, but it is methodologically important because it justifies the paper's core architectural choice on grounds of both efficiency and effectiveness. The large-window-but-not-unbounded design of AHN (32k default) is positioned at the pre-dilution sweet spot, making it not just a cost-saving compromise but potentially an improvement over unbounded full attention for very long sequences. The fact that AHN-augmented models sometimes outperform full attention on long-context benchmarks (e.g., LV-Eval average of 5.88 vs. 4.41 for Qwen2.5-3B-Instruct in Table 2) may be partially attributable to this effect: full attention at 128k may already be in the dilution regime, while the AHN's 32k window + compressed memory avoids it.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses three long-context benchmarks for evaluation. LV-Eval (Yuan et al., 2024) is a challenging benchmark covering both single-hop and multi-hop QA with 11 datasets, featuring confounding facts insertion, keyword/phrase replacement, and keyword-recall-based metrics; the paper evaluates on the 128k-context subsets. InfiniteBench (Zhang et al., 2024) tests models' ability to process and reason over super-long contexts, also using the 128k-length subset. LongBench (Bai et al., 2024) covers diverse tasks across multiple domains and languages; the paper focuses on six tasks with average sequence lengths exceeding 8k tokens (DuReader, HotpotQA, MuSiQue, NarrativeQA, QMSum, TriviaQA). For qualitative analysis, a 57k-token passage from PG19 (Rae et al., 2020) is used to demonstrate perplexity and memory behavior.

  • Base model(s). All experiments use the Qwen2.5-Instruct series (Yang et al., 2024) at three scales: 3B, 7B, and 14B parameters. The paper states these models are chosen because they are open-weight, representative of contemporary LLM capabilities, and their pretrained context length (approximately 32k for Qwen2.5) creates a natural regime where test-time extensions beyond training context are needed. The Qwen2.5 models serve as both the base architecture (with AHNs added) and the teacher for self-distillation (with full attention, frozen weights).

  • Metrics. The primary metric on LV-Eval is task-specific accuracy (%) computed using the keyword-recall-based grading function released with the benchmark; for InfiniteBench, task-specific metrics (e.g., F1 for QA, accuracy for retrieval tasks) are normalized to percentage scores; for LongBench, task-specific accuracy metrics are reported per-task and averaged across the six selected tasks. In the PG19 illustrative example (Section 3.2), perplexity is measured as the log-likelihood loss per token, and CUDA memory is measured in GB as peak GPU memory allocation during inference. For efficiency analysis, FLOPs are counted for matrix multiplications only (excluding softmax, normalization, and element-wise operations, per Table 1 footnote), and memory cache is measured in GB as the size of the KV cache plus AHN hidden states at float16 precision.

  • Baselines. Four baselines are compared in the main results (Tables 2 and 3):

    • Full attention: the unmodified base Qwen2.5-Instruct model with standard quadratic attention over the entire sequence. This serves as the performance upper bound (and efficiency lower bound).
    • Sinks + SWA: sliding window attention with attention sinks (Xiao et al., 2024), where the model keeps SS initial sink tokens (default S=128S = 128) plus a window of WW most recent tokens (default W=32,640W = 32,640 for ultra-long-context experiments, yielding 32,768 lossless tokens total; for LongBench, W=8,064W = 8,064 with S=128S = 128, yielding 8,192 lossless tokens). All tokens outside this combined lossless region are discarded entirely.
    • CT-Max: Compressive Transformer (Rae et al., 2020) with max-pooling as the compression function, compressing out-of-window tokens at a 4× compression rate. The compressed memory is allocated the same total size (in bytes) as the AHN's hidden state for fair comparison.
    • CT-Average: Compressive Transformer with average-pooling as the compression function, otherwise identical to CT-Max.

    Additionally, for the RULER needle-in-a-haystack analysis (Appendix B, Table 5), full attention and Sinks + SWA are compared against AHN-GDN, with all methods using 128 attention sinks and a 32,640-token sliding window where applicable.

  • Generation budget / compute accounting. For inference efficiency comparisons, compute is measured in FLOPs (specifically matrix multiplication FLOPs as defined in Table 1) and memory cache in GB. Rather than fixing a "generation budget," the paper compares methods at the same sequence lengths (128k for ultra-long-context, variable lengths up to the test set maximums for LongBench) and measures the relative FLOPs and memory required by each method's token mixer. Table 1 provides the analytical complexity formulas, and Tables 2, 7, and 9 provide concrete numbers at 128k sequence length. For training efficiency, Table 7 reports FLOPs per training step under a unified setting (AdamW optimizer, next-token prediction, full-model training, batch size 128, sequence length 24k, window size 8k) to compare the theoretical cost of training full attention vs. sliding window vs. AHN-augmented models, though the paper's actual training uses frozen base weights and only trains AHN parameters.

  • Cross-validation / statistical protocol. The paper does not employ cross-validation for model selection or hyperparameter tuning. AHN models are trained on the ChatQA2 dataset (1B tokens) and evaluated directly on the test sets of LV-Eval, InfiniteBench, and LongBench. Training hyperparameters (learning rate 1×1041 \times 10^{-4}, batch size 128, one epoch, randomized window sizes) are fixed based on a single configuration without reported hyperparameter sweeps. The ablation studies in Section 3.4 vary one factor at a time (training objective, window randomization, inference window size) and report the resulting metric directly without cross-validated confidence intervals. This means reported scores should be treated as point estimates without formal statistical significance bounds.

Main Quantitative Results

Ultra-Long-Context Evaluation: LV-Eval and InfiniteBench (128k Sequences)

The headline result appears in Table 2 and Figure 1b: augmenting Qwen2.5-3B-Instruct with AHN-GDN (+0.4% parameters, adding 13.0M parameters to the 3B base) reduces token-mixer FLOPs by 53.3% (from 2.50×10¹⁵ to 1.17×10¹⁵), reduces model FLOPs by 40.5% (from 3.29×10¹⁵ to 1.95×10¹⁵), and reduces memory cache by 74.0% (from 9.44 GB to 2.45 GB), while improving the LV-Eval average score from 4.41 to 5.88 and the InfiniteBench average from 9.52 to 13.24 — outperforming not just the sliding window baseline but also the full-attention model (LV-Eval 4.41, InfiniteBench 9.52).

Breaking this down across model scales and AHN instantiations:

Qwen2.5-3B-Instruct results (Table 2):

  • On LV-Eval (128k subset, average across all 11 tasks): full attention scores 4.41; Sinks + SWA scores 4.59; CT-Max scores 4.12; CT-Average scores 4.47. All three AHN variants exceed these: AHN-Mamba2 at 5.13, AHN-DN at 5.68, AHN-GDN at 5.88. The best AHN variant (GDN) outperforms the best non-AHN baseline (Sinks + SWA at 4.59) by 28.1%, and outperforms full attention by 33.3%.

  • On InfiniteBench (128k subset, average across English QA and Chinese QA): full attention scores 9.52; Sinks + SWA scores 10.47; CT-Max scores 10.00; CT-Average scores 10.81. AHN variants reach 12.44 (Mamba2), 13.51 (DN), and 13.24 (GDN). The best AHN variant (DN) outperforms the best baseline (CT-Average at 10.81) by 25.0%, and outperforms full attention by 41.9%.

  • Looking at individual LV-Eval tasks (Table 6 for complete results): AHN variants consistently outperform sliding window baselines on most datasets. For instance, on factrecall_en, full attention scores 6.88, Sinks + SWA drops to 3.34, while AHN-Mamba2 scores 5.58, AHN-DN scores 9.22, and AHN-GDN scores 12.51 — the best AHN variant nearly doubles full attention performance. On loogle_SD_mixup, full attention scores 0.89, Sinks + SWA scores 4.59, and AHN-GDN reaches 7.21 — an 8.1× improvement over full attention. However, there are some task-specific exceptions: on multifieldqa_en_mixup, full attention scores 0.00, Sinks + SWA scores 0.33, while AHN-Mamba2 and AHN-DN both score 0.00 and AHN-GDN scores 0.19 — the AHN variants occasionally underperform the sliding window baseline on specific tasks.

Qwen2.5-7B-Instruct results (Table 2):

  • On LV-Eval: full attention scores 3.62; Sinks + SWA scores 5.34; CT-Max scores 4.82; CT-Average scores 5.28. AHN-Mamba2 achieves 6.21, AHN-DN achieves 6.82, AHN-GDN achieves 6.54. The best AHN variant (DN at 6.82) outperforms the best baseline (Sinks + SWA at 5.34) by 27.7%, and nearly doubles full attention (3.62 → 6.82, an 88.4% relative improvement). Notably, the full attention 7B model performs worse than the 3B full attention model on LV-Eval (3.62 vs. 4.41 average), which the paper does not discuss — possibly reflecting that the larger model's training was less optimized for this specific benchmark, or that full attention at 128k pushes the 7B model into a degraded regime.

  • On InfiniteBench: full attention scores 13.50; Sinks + SWA scores 13.16; CT-Max scores 13.00; CT-Average scores 13.31. AHN-Mamba2 achieves 14.21, AHN-DN achieves 16.48, AHN-GDN achieves 16.93. The best AHN variant (GDN at 16.93) outperforms full attention (13.50) by 25.4%. Here, the sliding window baselines are competitive with or slightly below full attention, while AHNs provide a clear lift.

Qwen2.5-14B-Instruct results (Table 2):

  • On LV-Eval: full attention scores 4.99; Sinks + SWA scores 5.69; CT-Max scores 5.28; CT-Average scores 5.64. AHN-Mamba2 achieves 6.43, AHN-DN achieves 6.50, AHN-GDN achieves 6.51. The best AHN variant (GDN at 6.51) outperforms the best baseline (Sinks + SWA at 5.69) by 14.4%.

  • On InfiniteBench: full attention scores 12.21; Sinks + SWA scores 12.54; CT-Max scores 11.66; CT-Average scores 11.95. AHN-Mamba2 reaches 15.21, AHN-DN reaches 17.48, AHN-GDN reaches 16.52. The best AHN variant (DN at 17.48) outperforms full attention (12.21) by 43.2% — the largest relative improvement across all scales and benchmarks.

Efficiency analysis (Table 9): All three model scales show consistent efficiency patterns. The memory cache ratio (AHN-augmented divided by full attention) is 25.6–26.0% across all scales — roughly a 4× reduction. The mixing FLOP ratio is 46.7–49.7%, and the model FLOP ratio is 59.4–65.6% — the larger the base model, the smaller the relative FLOP savings because MLP computation (which is unchanged) constitutes a larger fraction of total FLOPs. The extra parameter ratio ranges from 0.2–0.4%, decreasing slightly with model scale due to the base model parameters growing faster than the per-head AHN parameter count.

Illustrative Example: PG19 Perplexity and Memory (Figure 3)

On a 57k-token passage from the first book of the PG19 test set (Section 3.2, Figures 3c and 3d):

  • Perplexity (Figure 3c): The base Qwen2.5-3B-Instruct model shows log perplexity around 2.2–2.3 for tokens within the pretrained 32k context length, then rises sharply once the context window is exceeded. The AHN-GDN augmented model maintains consistently low log perplexity (around 2.2–2.4) across all 57k tokens, with no sharp increase at the 32k boundary. This demonstrates that the compressed memory effectively preserves information needed for language modeling beyond the training context length.

  • GPU memory (Figure 3d): The base model's CUDA memory grows from approximately 6.0 GB at the start of the sequence to slightly above 7.5 GB at 57k tokens, consistent with linear KV cache growth under FlashAttention. The AHN-GDN model's memory remains roughly constant at approximately 6.0–6.1 GB throughout — the growth that would come from storing KV pairs beyond 32k is eliminated by the compression mechanism, and only the fixed-size AHN state and the 32k window KV cache are maintained.

Long-Context Evaluation: LongBench (6 Tasks with Average Length > 8k)

Table 3 reports results on six LongBench tasks where average sequence lengths exceed 8k tokens, using a smaller lossless memory budget than the ultra-long-context experiments (8,192 tokens total: 128 attention sinks + 8,064-token sliding window). This tests AHN under more constrained conditions where compression must handle a larger fraction of the total context.

Qwen2.5-3B-Instruct on LongBench:

  • Sinks + SWA averages 34.31 across the six tasks; CT-Max averages 34.03; CT-Average averages 34.51.
  • AHN-Mamba2 achieves 34.90, AHN-DN achieves 35.89, AHN-GDN achieves 35.55.
  • The best AHN variant (DN at 35.89) outperforms the best baseline (CT-Average at 34.51) by 4.0%, and outperforms Sinks + SWA by 4.6%. The improvement is most pronounced on MuSiQue (Sinks + SWA: 16.55; AHN-DN: 19.78, a 19.5% relative gain) and NarrativeQA (Sinks + SWA: 15.35; AHN-DN: 19.11, a 24.5% relative gain). On TriviaQA, all methods score similarly (84.93–86.17 range), suggesting that task's difficulty is not primarily driven by long-context retrieval.

Qwen2.5-7B-Instruct on LongBench:

  • Sinks + SWA averages 38.52 across the six tasks; CT-Max averages 38.29; CT-Average averages 38.50.
  • AHN-Mamba2 achieves 40.56, AHN-DN achieves 41.04, AHN-GDN achieves 40.59.
  • The best AHN variant (DN at 41.04) outperforms the best baseline (Sinks + SWA at 38.52) by 6.5%. Gains are concentrated on HotpotQA (Sinks + SWA: 51.57; AHN-DN: 54.24, a 5.2% gain) and MuSiQue (22.34 → 29.30, a 31.2% gain). The MuSiQue improvement is particularly striking — the AHN's compressed memory appears to help substantially on multi-hop reasoning that requires integrating information across long contexts.

Qwen2.5-14B-Instruct on LongBench:

  • Sinks + SWA averages 40.65; CT-Max averages 39.72; CT-Average averages 40.82.
  • AHN-Mamba2 achieves 41.34, AHN-DN achieves 41.83, AHN-GDN achieves 41.90.
  • The best AHN variant (GDN at 41.90) outperforms Sinks + SWA by 3.1%. The improvement is smaller in relative terms at the 14B scale, but still consistent across all AHN variants. HotpotQA shows the largest gain (Sinks + SWA: 55.68; AHN-DN: 58.71, a 5.4% gain).

Cross-task pattern: Across all three model scales on LongBench, AHN-DN tends to achieve the highest MuSiQue scores (19.78, 29.30, 32.92 for 3B, 7B, 14B respectively), while AHN-GDN tends to perform best on DuReader. There is no single AHN instantiation that dominates all tasks — the optimal choice appears task-dependent but all three variants consistently outscore the sliding window baselines, confirming that the AHN framework itself (not the specific recurrent architecture) drives the improvement.

Window Size Generalization on LongBench (Appendix C, Figure 7)

The paper tests how AHN-augmented models (Qwen2.5-7B-Instruct, AHN-GDN variant) perform as the lossless memory budget varies from 1,024 to 8,192 tokens (128 attention sinks + varying window sizes from 896 to 8,064). Across all six LongBench tasks:

  • AHN-GDN consistently outperforms both Sinks + SWA and CT-Average at every window size. The performance gap is largest at smaller window sizes (e.g., at 1,024 total lossless memory) and narrows as the window grows — this is expected because a larger window means the compressed memory needs to carry less information.
  • On DuReader, performance improves monotonically with window size for all methods, with AHN-GDN at the top (from roughly 24.5 at 1k to roughly 26.5 at 8k, compared to Sinks + SWA from roughly 23.5 to 25.5).
  • On HotpotQA, AHN-GDN shows strong performance even at 1k window (roughly 52), improving to roughly 54 at 8k, while Sinks + SWA drops from roughly 50 to 48 as the window shrinks — AHN is more robust to constrained lossless memory.
  • On MuSiQue, AHN-GDN shows the widest margin over baselines at all window sizes, demonstrating that the compressed memory is particularly valuable for multi-hop reasoning tasks where key information is distributed across the full sequence.

Needle-in-a-Haystack: RULER Exact-Recall Tasks (Appendix B, Table 5)

On advanced needle-in-a-haystack tasks from the RULER-128k subset, using Qwen2.5-7B-Instruct as the base model with 128 attention sinks and a 32,640-token sliding window:

  • Full attention dominates: on single_1, full attention scores 98.60 vs. 26.80 for both Sinks + SWA and AHN-GDN; on single_2, 97.20 vs. 25.40/25.20. The gap is massive — full attention can retrieve specific facts with near-perfect accuracy, while sliding-window-based methods (with or without AHN compression) cap around 25–28% because the needle may fall outside the window entirely.
  • AHN-GDN performs on par with Sinks + SWA across all NIAH task variants: multikey_1 (27.40 vs. 27.80), multikey_2 (11.40 vs. 10.60), multikey_3 (8.60 vs. 9.00), multivalue (23.45 vs. 22.95), multiquery (23.35 vs. 24.00). The compressed memory does not help on exact-recall tasks — it does not enable retrieval of specific facts that were never stored in the lossless window. The AHN's compression is lossy by design, and for tasks requiring verbatim recall of a specific sentence from 50k tokens ago, the information is effectively lost.

This result is the clearest evidence for the paper's own stated limitation: "they inevitably struggle on tasks that require exact-recall from the compressed memory" (Appendix B). The AHN framework trades off exact recall fidelity for efficiency, and the RULER results quantify the cost of that trade-off — a substantial one for needle-in-a-haystack-style retrieval.

Ablation Studies and Robustness Checks

Training objective: self-distillation (KL divergence) vs. next-token prediction (cross-entropy): Replacing the KL-divergence-based self-distillation objective with standard next-token prediction cross-entropy loss, while keeping all other training settings identical (Qwen2.5-7B-Instruct, AHN-GDN, randomized window size), causes the LongBench average to drop from 40.59 to 39.59. This is reported in Table 4 (note: the table reports 38.53 for a fixed 1024 window with KL loss, 39.59 for randomized window with CE loss, and 40.59 for randomized window with KL loss — the CE vs. KL comparison is between the latter two rows). The paper attributes this degradation to CE providing "sparse learning signals" that push the AHN modules toward "shortcuts in the training data," while self-distillation provides "denser guidance over the teacher's entire output distribution" (Section 3.4). This is a 2.5% relative degradation from switching the training objective, confirming that the choice of distillation objective is material.

Training window size: randomized vs. fixed: Training AHN-GDN on Qwen2.5-7B-Instruct with a fixed window size of 1024 tokens (rather than randomized window sizes sampled from [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]) causes the LongBench average to drop from 40.59 to 38.53 — a 5.1% relative degradation (Table 4, rows 1 vs. 3). The paper interprets this as evidence that fixed-window training causes the model to "overfit to that specific configuration and fail to generalize to unseen context lengths" (Section 3.4). Notably, the degradation from fixing the window (40.59 → 38.53, a drop of 2.06 points) is larger than the degradation from switching to CE loss (40.59 → 39.59, a drop of 1.00 point), suggesting that window randomization is the more impactful training design choice among the two ablated.

Inference window size sweep (Figure 4): Using the Qwen2.5-7B-Instruct AHN-GDN model (trained with randomized windows), the paper evaluates performance at inference window sizes from 1k to 96k (on LV-Eval 128k) and from 1k to 128k (on InfiniteBench 128k), both with 128 fixed attention sinks:

  • On LV-Eval: AHN performance improves from roughly 1.0 at 1k window to roughly 6.5 at 16k, plateaus around 32k (approximately 6.5), and then declines after 64k (dropping to roughly 5.5 at 96k). The sliding window attention (SWA) baseline shows a similar shape but consistently lower: roughly 0.9 at 1k, improving to roughly 5.5 at 16k, plateauing around 5.2–5.4 from 16k to 128k. Full attention (at 128k) scores approximately 3.6, which is below both AHN and SWA at medium-to-large window sizes — suggesting full attention at 128k may already be in the dilution regime.

  • On InfiniteBench: AHN improves from roughly 9.5 at 1k to roughly 17.0 at 16k, plateaus from 16k to 96k around 16.5–17.0, then declines to roughly 16.0 at 128k. SWA improves from roughly 9.0 at 1k to roughly 14.5 at 16k, plateaus around 14.5–15.0, with a slight decline at 128k. Full attention scores approximately 13.5.

  • The paper attributes the post-peak decline to "the attention-dilution effect, where the attention distribution becomes overly diffuse when the number of keys grows very large" (Section 3.4). The 32k default window is chosen as the sweet spot balancing performance and efficiency. Importantly, AHN's performance advantage over SWA is maintained at every window size — even at 1k, where the compressed memory carries most of the information burden, AHN outperforms SWA.

Pure RNN baseline (Appendix E, Table 8): Removing the sliding window attention entirely and relying only on the AHN (RNN) component — effectively converting the model into a pure recurrent architecture — causes catastrophic performance degradation: LV-Eval average drops from 5.34 (Sinks + SWA) to 0.04 (Pure RNN), and InfiniteBench average drops from 13.16 to 1.19. The AHN framework (Sinks + SWA + RNN) achieves 6.54 on LV-Eval and 16.93 on InfiniteBench, confirming that the compressed memory supplements rather than replaces the lossless window — the large sliding window is essential for the framework to function.

Soft demonstration of selective compression (Section 3.5, Figure 5): The paper probes the AHN's learned behavior by visualizing the gradient magnitudes of the self-distillation loss with respect to out-of-window token embeddings. Tokens with small gradient magnitudes have their information well-captured in the compressed memory, while tokens with large gradients are poorly represented. On an 811-token example from the AceMath training data with a 512-token window, mathematical symbols and numbers show predominantly green (low gradient, well-preserved), while pronouns, special tokens (e.g., <|im_start|>), and function words show red (high gradient, poorly preserved). This provides qualitative evidence that the content-dependent gating mechanism (α\alpha, β\beta, γ\gamma) learns to selectively store semantically important tokens in the compressed memory, consistent with the cognitive-science framing of hippocampal consolidation. However, this is a single illustrative example — the paper does not provide quantitative metrics of selectivity across many sequences or tasks.

Module design validation (Table 1 and Table 8): The ablation in Table 8 confirms that the full AHN framework (Sinks + SWA + RNN) is necessary — neither Sinks + SWA alone nor a pure RNN achieves competitive performance. The complexity analysis in Table 1 is validated by the measured numbers in Table 9: for Qwen2.5-3B at 128k, the predicted mixing FLOP ratio of 46.7% (from the formula) matches the measured ratio. The extra parameter count (from the formulas in Section 2.3: 3DNq+H2Nq3DN_q + H^2 N_q) is consistent with the reported 0.2–0.4% extra parameters across model scales.

Training efficiency validation (Table 7): Under the unified training setting (full-model training, AdamW, batch size 128, seq length 24k, window size 8k), one training step costs 0.6348×10170.6348 \times 10^{17} FLOPs for full attention (3B), 0.5405×10170.5405 \times 10^{17} for Sinks + SWA, and 0.5422×10170.5422 \times 10^{17} for AHN-GDN. The AHN overhead over sliding window attention is negligible (0.5422/0.54051.0030.5422/0.5405 \approx 1.003, or 0.3% additional FLOPs per step). For the 7B model: 1.1519×10171.1519 \times 10^{17} (full), 1.0334×10171.0334 \times 10^{17} (SWA), 1.0359×10171.0359 \times 10^{17} (AHN-GDN). For 14B: 2.5396×10172.5396 \times 10^{17} (full), 2.2252×10172.2252 \times 10^{17} (SWA), 2.2319×10172.2319 \times 10^{17} (AHN-GDN). These numbers confirm the analytical complexity claims — AHN adds minimal training overhead over sliding window attention.

Context length generalization from 24k training to 128k inference: The paper trains models with a maximum sequence length of 24k tokens and evaluates at 128k — a 5.3× generalization in sequence length. The strong performance at 128k (Tables 2, 6, Figure 3c, Figure 4) demonstrates that the AHN's learned compression strategy transfers to sequences substantially longer than those seen during training. However, this generalization ability is not directly ablated — there is no experiment comparing, for instance, a model trained at 8k vs. 24k to quantify how training length affects 128k performance.

Critical Assessment

The experiments in this paper collectively demonstrate that AHN-augmented models achieve the paper's primary efficiency claims while matching or exceeding the base model's performance on comprehension-focused long-context benchmarks. However, the mapping between specific experiments and specific claims requires careful examination.

Claim: "AHN-augmented models consistently outperform sliding window baselines and achieve performance comparable or even superior to full-attention models, while substantially reducing computational and memory requirements." (Abstract, Section 1)

This claim is well-supported by the main results in Tables 2 and 3. Across all three model scales (3B, 7B, 14B), all three AHN instantiations (Mamba2, DN, GDN), and both ultra-long-context benchmarks (LV-Eval, InfiniteBench), AHN-augmented models outperform the sliding window with attention sinks baseline. The margin varies — from modest (14B on LV-Eval: 6.51 vs. 5.69 for Sinks + SWA, a 14.4% relative gain) to substantial (7B on LV-Eval: 6.82 vs. 5.34, a 27.7% gain; 3B on InfiniteBench: 13.51 vs. 10.47, a 29.0% gain). The efficiency numbers (Table 9) are unambiguous: 74% memory cache reduction and 40.5% model FLOP reduction are measured, not just predicted from complexity formulas.

However, the claim that AHN achieves performance "superior to full-attention models" requires more nuance than the Abstract suggests. On LV-Eval, full attention scores 4.41 (3B), 3.62 (7B), and 4.99 (14B) — and AHN variants score higher in all cases (5.13–5.88 for 3B, 6.21–6.82 for 7B, 6.43–6.51 for 14B). On InfiniteBench, full attention scores 9.52 (3B), 13.50 (7B), and 12.21 (14B) — and AHN variants score higher (12.44–13.51, 14.21–16.93, 15.21–17.48 respectively). So AHN does outperform full attention on these benchmarks. But the full attention baseline at 128k sequence length may not represent the best possible full-attention performance — the paper itself shows in Figure 3c that Qwen2.5-3B-Instruct's perplexity degrades beyond its 32k pretraining context length. If the base model was never trained on 128k sequences, its full-attention performance at that length may be suboptimal due to out-of-distribution positional encodings or attention patterns, not due to an inherent limitation of the attention mechanism. The AHN might be outperforming an undertrained full-attention model, not full attention in principle. A stronger baseline would be a model that was fine-tuned with full attention at extended context lengths (e.g., via position interpolation or continued pretraining) — this comparison is absent.

Claim: "AHNs activate only when the sequence length exceeds the 32k window, addressing the quadratic-complexity issue of attention that emerges at that scale." (Abstract, Section 2.2)

This claim is validated by design (AHNs are architecturally inactive below the window size), by the complexity analysis (Table 1: when LWL \leq W, the AHN-GDN complexity reverts to the sliding window complexity, which is still quadratic in WW but bounded), and by the PG19 illustrative example (Figure 3c shows no perplexity degradation at the 32k boundary for AHN-augmented models). However, the claim that AHN "addresses" the quadratic-complexity issue is partially undercut by the baseline setup: the sliding window baseline also addresses the quadratic issue by simply discarding out-of-window tokens. The contribution of AHN is not that it makes complexity linear (sliding window already does that) but that it makes complexity linear while preserving information through compression rather than discarding it. The paper sometimes conflates "AHN reduces complexity" with "AHN preserves information that sliding window discards" — both are true, but the efficiency argument applies equally to any sliding-window-based method.

Claim: The self-distillation training method is "efficient" — only 10 hours on 32 A100 GPUs for 7B models, using only 1B tokens. (Section 3.1, Appendix D)

This claim is specific and verifiable. Table 7 provides the FLOPs numbers backing the efficiency: training one step with AHN-GDN costs essentially the same as training with sliding window attention alone (1.0359 vs. 1.0334 × 10¹⁷ FLOPs for 7B), and substantially less than full attention training (1.1519 × 10¹⁷). The 1B token training budget is modest compared to the 15–20B tokens used by prior work (MiL, HQLT) — a genuine advantage. However, the training data (ChatQA2) is itself a curated long-context dataset; training on arbitrary corpora might require different data volumes. More importantly, the paper does not ablate the amount of training data — there is no experiment showing performance at 500M, 1B, or 2B tokens to assess whether 1B is sufficient or saturating. It is possible that even less data would suffice, or that more data would yield further gains; the paper provides no evidence either way.

Claim: Randomized window size training causes the AHN to learn a generalizable compression strategy. (Section 3.4)

The ablation in Table 4 supports this: fixed window (1024) produces 38.53 on LongBench, randomized windows produce 40.59 — a clear improvement. However, the "fixed" setting uses a single window size (1024), which is at the small end of the range. An alternative interpretation is that training with a small fixed window is simply suboptimal for the 8k-window LongBench evaluation, and training with any window size close to 8k might perform better. A more rigorous ablation would compare randomized windows against a fixed window of the same size as the evaluation setting (e.g., 8k) to determine whether randomization helps beyond simply matching the training window to the inference window. The current ablation confounds window size (1024 vs. mean of the random distribution) with randomization itself.

Genuine weaknesses:

  1. Single model family (Qwen2.5-Instruct). All experiments use Qwen2.5-Instruct models at three scales. While this is three distinct checkpoints, they share the same architecture, pretraining data, and training procedure. Whether AHN would work as well on models from other families (Llama, Mistral, Gemma) — which have different attention patterns, positional encodings, and pretraining recipes — is completely untested. The paper acknowledges this implicitly (no claim of universality), but the reader should not assume these results generalize across model families without replication.

  2. Full attention baseline may be suboptimal at 128k. As noted above, Qwen2.5 models were pretrained with a maximum context length of approximately 32k. Deploying them with full attention at 128k without any context-extension technique (e.g., RoPE scaling, position interpolation) means the model is operating with positional encodings it never saw during training. The fact that AHN sometimes outperforms this baseline may reflect AHN's effectiveness, or it may reflect that full attention at 128k is an unfairly weak baseline. A properly context-extended full-attention model would be a more compelling comparison point.

  3. No ablation on compression rate or memory size. The AHN memory state hh has a fixed size of H×HH \times H per head (e.g., 128×128=16, ⁣384128 \times 128 = 16,\!384 elements per head for Qwen2.5). The paper never varies this — there is no experiment showing what happens with a smaller or larger compressed memory. If the memory were halved (e.g., by using a low-rank approximation), would performance degrade gracefully or collapse? If doubled, would there be diminishing returns? This is a central design parameter that receives zero empirical attention.

  4. Limited statistical rigor. With 500 test questions on LV-Eval and a handful of tasks on InfiniteBench and LongBench, the reported average scores are point estimates without confidence intervals. The differences between AHN variants (e.g., AHN-DN 5.68 vs. AHN-GDN 5.88 on 3B LV-Eval) might not be statistically significant. The paper makes no attempt to quantify variance or report standard deviations, making it difficult to assess whether observed differences are reliable or noise.

  5. The gradient visualization is anecdotal, not systematic. Figure 5 shows a single example with 811 tokens and a 512-token window. This is compelling as an illustration but does not constitute evidence that the AHN systematically learns to prioritize mathematically informative tokens. Quantitative metrics — such as the correlation between token-level gradient magnitude and some measure of token importance (e.g., TF-IDF, surprisal, attention entropy) across many sequences — would be needed to support the claim that AHN learns selective compression.

  6. No comparison to other KV cache compression methods. The baselines include Compressive Transformer (a simple pooling-based compressor) and sliding window (which discards tokens), but not more sophisticated KV cache selection or eviction methods like H2O, SnapKV, or Keyformer. These methods also aim to reduce memory while preserving important information, and a comparison would position AHN more clearly in the landscape of practical long-context inference solutions. The paper acknowledges these methods in the related work (Section 4.2) but does not benchmark against them.

  7. Training efficiency numbers are idealized. The claimed "10 hours on 32 A100 GPUs" is for the AHN-GDN variant on the 7B model, using 1B tokens of training data with frozen base weights. This is certainly efficient relative to pretraining, but several practical costs are not accounted for: (a) the need to forward-pass the full teacher model during training (doubling the required memory for activations), (b) the cost of generating or curating the training data (ChatQA2 is pre-existing but may not be available for all domains), (c) the need to train separate AHN modules for each base model variant. Additionally, the training FLOPs in Table 7 assume full-model training for comparison purposes, but the actual training freezes base weights — the reported 10 hours reflects the actual training cost, not the numbers in Table 7 which are for "a unified setting" of full-model training. This creates a slight disconnect between the analytical efficiency numbers and the actual training cost claim.

  8. Missing experiment: how does AHN performance vary with the amount of training data? The paper uses exactly one epoch over 1B tokens. There is no data scaling curve — would 500M tokens suffice? Would 2B tokens help? This is a practical omission since data efficiency is one of the claimed advantages of self-distillation.

Missing experiments that would strengthen the paper:

  • Ablation of the compressed memory size (H×HH \times H per head vs. smaller variants): how much compression is possible before performance collapses?
  • Evaluation on non-Qwen base models: Llama-3, Mistral, or Gemma to test architectural generality.
  • Comparison with context-extended full attention: apply RoPE scaling or position interpolation to the base model so full attention operates in-distribution at 128k, then compare against AHN.
  • Benchmarking against modern KV cache eviction methods (H2O, SnapKV) on the same LV-Eval/InfiniteBench tasks.
  • Training data ablation: performance as a function of tokens trained (250M, 500M, 1B, 2B).
  • Systematic quantification of selective compression: correlation between gradient magnitude and token-level importance metrics across many sequences of different types (math, code, narrative text).

Summary of evidential support for the central thesis: The paper's central thesis — that a large sliding window attention augmented with a small recurrent compressor for out-of-window tokens achieves strong long-context performance with substantially lower cost than full attention — is well-supported by the presented experiments within the studied configuration. The consistency across three model scales, three AHN architectures, and two distinct long-context benchmarks constitutes compelling evidence that the approach works for Qwen2.5-family models on comprehension-oriented long-context tasks. The weaker performance on exact-recall tasks (RULER, Table 5) is honestly reported and correctly interpreted as a fundamental limitation of lossy compression. The efficiency numbers (74% memory reduction, 40.5% FLOP reduction) are measured, not estimated, and represent genuine practical gains.

The primary uncertainty is generality: would these results hold for other model families, other task types (code generation, dialogue, retrieval-augmented generation), and other training regimes? The paper's contributions are solidly demonstrated within their scope, but that scope — a single model family with a specific pretraining recipe, evaluated on comprehension-focused long-context benchmarks — limits how broadly the conclusions can be applied without further validation. The paper is candid about some limitations (exact recall in Appendix B) but less explicit about others (single model family, the potentially weak full-attention baseline at out-of-distribution lengths). A practitioner considering adopting AHN for a different base model or task domain should treat these results as promising but in need of domain-specific validation.

6. Limitations and Trade-offs

Lossy Compression Precludes Exact-Recall Tasks

The assumption or constraint. The AHN framework is built on the premise that older historical information can be compressed into a fixed-size recurrent state without catastrophic information loss for most tasks. However, the paper acknowledges directly in Appendix B that "their fixed-size compressed memory inevitably entails some information loss and may impair performance on tasks that require exact recall." This is not a side effect — it is a fundamental consequence of the architectural choice. The RNN-like AHN module maps all tokens beyond the sliding window into an H×HH \times H matrix of bounded capacity (e.g., 128×128=16, ⁣384128 \times 128 = 16,\!384 elements per head in Qwen2.5), which means that as sequence length grows, the compression ratio grows without bound. At 128k tokens with a 32k window, roughly 96k tokens are compressed into a state that has the same representational capacity regardless of whether it summarizes 1 token or 96,001 tokens.

The consequence. On needle-in-a-haystack (NIAH) tasks from the RULER-128k benchmark (Appendix B, Table 5), AHN-GDN performs on par with sliding window attention — both score approximately 26% on single_1 compared to 98.6% for full attention — but "markedly worse than full attention on exact-recall tasks" (Appendix B). The compressed memory does not enable retrieval of verbatim facts that were never stored in the lossless window. This is a hard failure mode: any downstream task that requires recalling a specific name, date, number, or code snippet from more than 32k tokens ago will fail at rates comparable to simply discarding the out-of-window context entirely. For applications like long-document legal review, codebase debugging across hundreds of files, or fact-verification over book-length texts — where exact retrieval of distant specific details is the primary task — AHN provides essentially zero benefit over a much simpler sliding window baseline.

What evidence exists in the paper. Table 5 provides direct quantitative evidence: across all eight RULER NIAH subtasks (single_1 through multiquery), AHN-GDN's scores are within ±1 percentage point of the Sinks + SWA baseline, and both are dramatically below full attention (e.g., 25.2 vs. 97.2 on single_2). The gradient visualization in Figure 5 provides mechanistic evidence: the AHN learns to preserve mathematical symbols while discarding function words, which is beneficial for semantic comprehension but means that arbitrary token-level detail (a specific name mentioned once on page 3) is likely lost. Section 5 (Discussion and Conclusion) also acknowledges the limitation, describing it as something "inherent to the trade-off of lossy compression." The paper provides no experiments testing intermediate regimes — for instance, how recall degrades as a function of distance beyond the window, or whether certain types of facts (numbers, names, dates) are differentially affected.

Mitigation status. The paper does not attempt to solve this limitation within the current framework. It suggests future work on "memory management that preserves critical information in lossless memory while leveraging compression for efficiency" (Appendix B) — essentially, a hybrid approach where some tokens beyond the window are selectively retained in full fidelity rather than compressed. This is an aspiration, not an implemented solution. No mechanism is proposed for how the model would decide which tokens merit retention versus compression. The limitation is acknowledged honestly but remains unresolved, and any practitioner evaluating AHN for exact-recall-heavy workloads should treat the RULER results as disqualifying for those use cases.


The Full Attention Baseline Is Potentially Undertrained at 128k Context Length

The assumption or constraint. The paper compares AHN-augmented Qwen2.5 models against the same base models operating with full attention over 128k-token sequences. However, Qwen2.5 models were pretrained with a maximum context length of approximately 32k tokens (as referenced in Section 3.2: "the perplexity of standard Qwen models rises sharply once the 32k token context window is exceeded"). Deploying them at 128k with full attention — without any context-extension technique such as RoPE scaling, position interpolation, or continued pretraining on longer sequences — means the model is operating with positional encodings and attention patterns far outside its training distribution. This is not a minor technicality: it means the full-attention baseline may be fundamentally handicapped, and any performance advantage shown by AHN (which the paper often frames as "superior to full attention") may partially reflect the fact that full attention at 128k is an artificially weak comparison point.

The consequence. The paper's headline finding — that AHN-augmented models often outperform full attention on long-context benchmarks (Table 2: e.g., LV-Eval 5.88 for AHN-GDN vs. 4.41 for full attention on 3B; InfiniteBench 16.93 for AHN-GDN vs. 13.50 for full attention on 7B) — cannot be cleanly attributed to the superiority of the hybrid memory architecture. Part of the gain may simply come from the fact that the AHN-augmented model uses a 32k window (within the training distribution) plus compressed memory, while the full-attention model is forced to attend over tokens with out-of-distribution positional encodings. The perplexity evidence in Figure 3c is direct corroboration: Qwen2.5-3B-Instruct's log perplexity rises sharply beyond 32k tokens even without any architectural modification, indicating that the model was not trained to handle such lengths. If the full-attention baseline were properly context-extended (e.g., via NTK-aware RoPE scaling or continued pretraining at 128k), the performance gap might shrink, disappear, or even reverse — and the paper provides no evidence to distinguish these possibilities.

This also muddies the interpretation of the window-size ablation (Figure 4). The fact that full attention scores lower than both AHN and SWA on LV-Eval at 128k is interpreted by the paper as possible evidence of "attention dilution" — but an equally plausible explanation is simply out-of-distribution positional encodings causing degraded attention patterns. The two hypotheses (dilution vs. distribution shift) have different implications for architecture design: dilution implies that even perfectly trained attention would fail at very long lengths, while distribution shift implies that the failure is fixable through training alone. The paper does not attempt to distinguish them.

What evidence exists in the paper. The primary evidence that the baseline may be weak comes from the paper's own observations: Figure 3c shows perplexity degradation for the base model beyond 32k, and the text notes this explicitly ("the perplexity of standard Qwen models rises sharply once the 32k token context window is exceeded"). The fact that Qwen2.5's pretraining context length is approximately 32k is mentioned in passing but never directly connected to the validity of the full-attention baseline. The paper does not report what context length the base models were evaluated at in any prior published work, nor does it apply any context-extension technique to bring the full-attention model into its operating regime.

Mitigation status. The paper does not acknowledge this as a limitation of the experimental design, nor does it attempt to address it through additional baselines. A proper comparison would include: (1) the base model with a context-extension technique applied (e.g., RoPE scaling to 128k), (2) the base model fine-tuned on long-context data with full attention, or (3) a model from a different family that was natively pretrained at 128k+ context (e.g., Llama 3.1 with 128k native context). None of these comparisons are present. This limitation weakens the "AHN outperforms full attention" narrative, though it does not affect the validity of the efficiency comparisons (which are based on FLOP counting, not benchmark scores) or the comparisons against sliding window baselines (which use the same window size and thus the same in-distribution positional range).


Single Model Family with No Cross-Architecture Validation

The assumption or constraint. All experiments in the paper use models from the Qwen2.5-Instruct family (3B, 7B, 14B), which share the same pretraining architecture, training data, tokenizer, and training procedure developed by the Qwen team (Yang et al., 2024). The paper states in Section 3.1 that these models are chosen as representatives of "open-weight" LLMs, but it provides no evidence that the AHN framework's effectiveness transfers to models with different design choices — different positional encodings (RoPE configuration details, learned positional embeddings), different attention implementations (multi-head vs. grouped-query vs. multi-query attention), different activation functions, different normalization schemes, or different pretraining data distributions and tokenizers.

The consequence. The paper's claim that AHNs are a "general memory framework" (Section 1, Section 2.2) that can augment existing LLMs is only validated within a narrow architectural corridor. A practitioner using Llama, Mistral, Gemma, or any non-Qwen model has no empirical evidence that AHN will work — the reported performance numbers and efficiency gains may not transfer. Specific concerns include: (1) Models with different head dimensions HH would produce AHN memory states of different sizes, potentially changing the compression-to-fidelity trade-off. (2) Models with different positional encoding schemes might handle the sliding window boundary differently, affecting the quality of information that the AHN receives at the window's edge. (3) The self-distillation training approach relies on the full-attention teacher providing high-quality probability distributions — if the teacher itself is poorly calibrated or has different output distribution characteristics (e.g., different temperature scaling in the pretraining), the KL distillation signal might be less effective. (4) The paper uses Qwen2.5 Instruct variants, which have undergone instruction tuning — whether the approach works on base (non-instruction-tuned) models is untested, and the self-distillation objective targeting an instruction-tuned teacher's outputs may embed instruction-following biases that wouldn't apply to base model distillation.

What evidence exists in the paper. The paper experiments on three model scales within the same family, which provides evidence of parameter-count generalization (3B → 7B → 14B) but not architectural generalization. The results are consistent across scales: AHN improves over baselines at all three sizes, and the efficiency ratios are similar (memory cache ratio 25.6–26.0%, mixing FLOP ratio 46.7–49.7%). This is reassuring but limited — all three models share the same architecture and pretraining recipe. There is no ablation studying whether performance depends on Qwen2.5-specific properties (e.g., its particular RoPE base frequency, its grouped-query attention ratio, or its SwiGLU activation).

Mitigation status. The paper does not address this limitation. There is no experiment with a non-Qwen model, nor any discussion of potential architectural dependencies. The phrase "general memory framework" in the abstract should be interpreted as "a framework that can in principle be applied to different models" (which is true — nothing in the AHN design is Qwen-specific) rather than "a framework verified to work across diverse model families" (which is not established). A practitioner should budget for a validation experiment on their specific base model before committing to the approach.


Training Cost Amortization and the Hidden Cost of Teacher Forward Passes

The assumption or constraint. The paper emphasizes the efficiency of its self-distillation training procedure: "only ∼10 hours on 32 A100 GPUs to train AHNs to augment 7B model" (Section 3.1), using 1B tokens of training data with all base model parameters frozen. This is framed as a key advantage over prior work that requires training from scratch or fine-tuning all parameters. However, this cost calculation omits a structural expense: every training step requires running the teacher model (the frozen full-attention Qwen2.5) on the entire input sequence to obtain the target probability distribution pp'. The teacher forward pass processes the full sequence with quadratic attention — it is not accelerated by sliding windows or AHN compression. This means the total computation per training step includes both the student forward pass (window attention + AHN, linear complexity) and the teacher forward pass (full attention, quadratic complexity).

The consequence. The actual training FLOPs are substantially higher than the AHN-only numbers in Table 7 might suggest. Table 7 reports 1.0359×10171.0359 \times 10^{17} FLOPs per step for AHN-GDN training (7B model), but this counts only the student's forward pass under the "unified setting" of full-model training — it does not include the teacher's full-attention forward pass, which costs 1.1519×10171.1519 \times 10^{17} FLOPs per step (the "Full attention" row in Table 7). The total training cost per step is therefore approximately 1.0359+1.1519=2.1878×10171.0359 + 1.1519 = 2.1878 \times 10^{17} FLOPs, or roughly 2.1× the cost of training the student alone. The 10-hour figure presumably accounts for this (since the training was actually run, not just estimated), but the paper's presentation of training efficiency — focusing on the parameter count being trained (0.4%) and the low token count (1B) rather than total FLOPs — can mislead readers into thinking the training is 250× cheaper than full training (since 0.4% of parameters are updated), when it is closer to 2× cheaper per step (and overall cheaper primarily because only 1B tokens are used rather than the trillions used in pretraining).

Additionally, the teacher forward pass consumes GPU memory proportional to the full sequence length with full attention. At the training sequence length of 24k tokens, the teacher's KV cache and attention activations must coexist in memory alongside the student's parameters and activations. This means the peak GPU memory during training may be substantially higher than what would be needed for the student alone, potentially constraining the maximum training sequence length or batch size on fixed hardware.

What evidence exists in the paper. Table 7 provides the relevant numbers if interpreted carefully: the "Full attention" row shows the per-step cost of a full-attention forward pass, and the "AHN-GDN" row shows the cost of the AHN-augmented student forward pass under the same setting. The sum is 2.1878×10172.1878 \times 10^{17} FLOPs per step — the paper does not explicitly perform this addition anywhere. The 10-hour figure is stated as an empirical measurement ("Training AHNs for a 7B base model requires only ∼10 hours on 32 A100 GPUs"), so it presumably reflects the actual wall-clock time including teacher forward passes, but the FLOPs accounting in the paper does not make the teacher cost explicit. There is no ablation studying whether a smaller teacher (e.g., a 3B teacher for a 7B student, or a teacher with a shorter context window) could provide a sufficient training signal at lower cost.

Mitigation status. The paper does not discuss this hidden cost. The text in Section 2.4 describes the self-distillation framework as "computationally efficient" and emphasizes that "the base model's weights are frozen during training, and only the AHN parameters are optimized," but never mentions that the teacher model must run a full forward pass on every training example. A practitioner planning to apply this method should be aware that the training cost is roughly double what the student's forward pass alone would suggest, and that GPU memory must accommodate both the student and teacher simultaneously during training. For very long training sequences or larger models, this could become a binding constraint that is not apparent from the paper's efficiency discussion.


Difficulty Estimation Is Absent — No Mechanism for Dynamic Allocation Between Memory Types

The assumption or constraint. The AHN framework operates with a fixed, predetermined division of labor: the sliding window handles the most recent WW tokens with lossless fidelity, and all tokens beyond WW are compressed through the AHN. The window size WW is a global hyperparameter (default 32,768) chosen based on aggregate benchmark performance (Figure 4) and applied uniformly to all sequences, all tokens, and all positions within a sequence. The paper provides no mechanism for dynamic, per-token or per-sequence allocation — deciding at inference time that some tokens merit retention in the lossless window beyond the default limit, or that some parts of the sequence require a larger or smaller window for optimal processing.

The consequence. The fixed window creates a hard binary: tokens at position tWt-W are compressed; tokens at position tW+1t-W+1 are retained in full fidelity. This boundary is content-agnostic, determined purely by position. The RULER results (Table 5) demonstrate the cost of this rigidity: when a needle (a specific fact needed later) happens to fall outside the window, it is compressed and becomes effectively irretrievable, even if the model "knows" it will need that fact. A dynamic system could, in principle, recognize that a particular sentence contains crucial information and selectively retain it in lossless memory (or allocate it a larger share of the compressed memory's representational budget), while aggressively compressing less important tokens to stay within memory constraints. AHN provides no mechanism for this kind of content-aware memory management — the gating mechanisms (α\alpha, β\beta, γ\gamma) operate within the compression process to control how strongly each token is written into memory, but they cannot prevent a token from being compressed or promote it to lossless retention.

Furthermore, the optimal window size varies by task (Figure 4: LV-Eval peaks around 32k and declines after 64k; InfiniteBench remains flat from 16k to 96k), by sequence (a 200k-token document may benefit from a different window than a 40k-token document), and potentially by position within a sequence (the first few paragraphs of a document may be more important to retain than later, routine text). A fixed 32k window applies the same policy everywhere, leaving potential gains unharvested.

What evidence exists in the paper. The evidence is primarily in what the paper does not explore. The window-size ablation (Figure 4) shows that performance varies with WW, confirming that the choice of WW matters, but the paper treats WW as a static hyperparameter to be set once based on benchmark curves. The gradient visualization (Figure 5) shows that the AHN does learn selective compression — mathematical symbols are better preserved than pronouns — which demonstrates that some level of content-aware processing occurs within the compressed memory. However, this selectivity cannot override the hard window boundary: a crucial mathematical formula at position tW1t-W-1 gets compressed; a filler word at position tW+1t-W+1 gets retained losslessly. There is no experiment testing a dynamic or adaptive window size, no mechanism for "promoting" compressed tokens back to lossless status, and no discussion of how difficulty or importance estimation could inform memory allocation.

Mitigation status. The paper does not frame this as a limitation or propose dynamic alternatives. The fixed window is presented as a feature (simplicity, predictability) rather than a constraint. The paper's future work suggestions focus on "stronger recall mechanisms" (Section 5) without specifying dynamic allocation as a direction. This is a significant omission because it represents a clear path to improving the framework: if the window size or the compression budget could be allocated based on estimated token importance, the hard failure on exact-recall tasks might be partially mitigated (by retaining needles in the window) without sacrificing overall efficiency.


The 5.3× Context-Length Generalization Gap Is Empirically Observed but Not Explained or Bounded

The assumption or constraint. The paper trains AHN-augmented models with a maximum sequence length of 24k tokens (Section 3.1) and evaluates them at 128k tokens — a 5.3× increase. The model is expected to generalize its compression strategy to sequences nearly an order of magnitude longer than any it encountered during training. This generalization is observed to work (the 128k results in Table 2 are strong), but the paper provides no analysis of why it works, when it might fail, or how much further the generalization can be pushed. There is no experiment with training at shorter lengths (e.g., 8k or 12k) to establish a scaling relationship between training length and inference length, no experiment testing even longer evaluation lengths (e.g., 256k or 512k) to find where generalization breaks, and no theoretical argument for why the recurrent compression should be length-agnostic.

The consequence. A practitioner cannot predict what will happen if they deploy AHN at 256k tokens after training at 24k. Will performance degrade gracefully or collapse at some threshold? The linear complexity of the AHN update (O(WL)\mathcal{O}(WL) from Table 1) suggests that the computation will remain tractable, but whether the quality of the compressed memory representations remains useful at extreme lengths depends on factors the paper does not analyze: the capacity of the H×HH \times H memory state, the accumulation of compression errors over many timesteps, the behavior of the forget gate α\alpha over very long horizons (does the memory saturate or continue to update meaningfully?), and the distributional shift between the 1–24k token regime seen during training and the 24–128k regime encountered at inference.

The self-distillation training objective further complicates this: the teacher model was itself not trained at 128k (Qwen2.5's pretraining context is ~32k), so the teacher's output distributions at positions beyond 32k may be unreliable — they are out-of-distribution for the teacher as well. The student is being trained to match a teacher whose own behavior at extreme lengths is uncharacterized and potentially degraded. At evaluation time beyond 32k, the student must generalize from a teacher that was providing targets from the 1–24k range (with a possible teacher degradation starting around 24k, though the paper doesn't discuss this). The fact that this works at 128k is impressive but not well-understood, making it difficult to predict reliability at other lengths or for other model families.

What evidence exists in the paper. The primary evidence that generalization works is the 128k benchmark performance itself (Table 2) and the PG19 perplexity results at 57k tokens (Figure 3c) — both post-training lengths. The randomized window size training (where the window varies from 32 to 8,192 tokens during training) provides some evidence that the AHN learns to handle a range of compression ratios, which likely contributes to generalization. But there is no experiment that systematically varies training length and measures generalization: no curve showing performance at 128k as a function of maximum training length (e.g., trained at 8k, 16k, 24k, 32k). There is no evaluation at 256k or longer to find the breaking point. There is no analysis of memory state saturation — for instance, measuring the effective rank of htWh_{t-W} as a function of the number of compressed tokens to see whether the memory's representational capacity is exhausted at some point.

Mitigation status. The paper does not acknowledge this as an open question. The 128k evaluation is presented as a demonstration that the method works at that length, with no discussion of the generalization gap or its limits. The randomized window training is the implicit mitigation (exposing the AHN to varied compression ratios during training), but its effectiveness at enabling length generalization is only validated post-hoc by the 128k results, not through a controlled experiment. A practitioner pushing beyond 128k should treat this as uncharted territory and budget for evaluation at their target lengths rather than assuming the 128k results will transfer.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper instigates a reorientation in how the field thinks about hybrid attention-RNN architectures, shifting the design question from "how do we make recurrence approximate attention everywhere?" to "when does recurrence actually need to operate, and what should its precise division of labor with attention be?" This is a reframing rather than a paradigm shift — the core components (sliding window attention, recurrent state compression, self-distillation) are individually precedented — but the reframing has substantial practical consequences because it changes what we optimize for and what we consider acceptable behavior from the recurrent component.

The key methodological shift is the large-window philosophy: set the sliding window large enough (32k tokens) that attention handles the vast majority of practical sequence lengths in its native, efficient regime, and activate recurrence only for the tail that would otherwise incur quadratic cost. This is the paper's sharpest departure from prior hybrid work. Infini-attention (Munkhdalai et al., 2024) uses 2,048-token chunks. LoLCATs (Zhang et al., 2025) and HQLT (Irie et al., 2025) use 64-token attention windows — smaller than a typical paragraph. In these designs, the RNN component is a co-equal partner to attention, activated for nearly every token in every sequence. The AHN framework demotes recurrence to a specialized overflow handler: it activates only when the sequence length crosses the threshold where attention becomes compute-bound, and it handles only the tokens that have fallen out of the lossless window. This is a fundamentally different relationship between the two memory systems — one of strict hierarchy rather than partnership.

This reframing matters because it relaxes the performance requirements on the recurrent component. Prior hybrid methods must ensure their RNN approximates attention well across all positions and all sequence lengths, including short contexts where attention is both efficient and highly capable — a tall order given RNNs' documented limitations on exact recall (Wen et al., 2025). AHN only needs to perform well for tokens older than 32k, and only for sequences longer than 32k. It operates in a regime where the alternative is not perfect attention but zero attention (the sliding window baseline discards those tokens entirely). This makes the learning problem dramatically easier: the AHN doesn't need to match attention; it just needs to outperform discarding. The paper's strongest evidence for this division of labor is implicit but telling: on RULER exact-recall tasks (Table 5), AHN performs identically to sliding window attention — the compressed memory adds nothing. But on comprehension tasks like LV-Eval and InfiniteBench (Table 2), AHN substantially outperforms sliding window, confirming that the compressed memory provides semantic gist even when it cannot support verbatim retrieval. The division of labor is functional: attention handles exact recall within the window; AHN handles semantic summarization beyond it.

The paper also resolves — or at least provides a framework for resolving — a tension in the long-context literature between efficiency-motivated methods that discard information (sliding windows, KV cache eviction) and fidelity-motivated methods that preserve everything (full attention, external memory). The AHN framework demonstrates that you can have both, but in different regions of the sequence: lossless fidelity for recent tokens, compressed efficiency for distant tokens. This is not a new idea in cognitive science (the Multi-Store Model dates to 1968), but the paper translates it into a concrete, trainable neural architecture with measured efficiency gains. The 74% memory reduction and 40.5% FLOP reduction at 128k (Table 9) put hard numbers on what this decomposition achieves.

A subtler landscape change is the paper's demonstration that self-distillation with frozen base weights is a viable and efficient training paradigm for augmenting existing LLMs with long-context capabilities. Prior hybrid work either trained from scratch (HQLT: 15B tokens) or fine-tuned substantial fractions of the model (MiL: all token mixer parameters, 20B tokens). AHN's approach — freeze everything, train only 0.4% of parameters on 1B tokens — makes long-context augmentation accessible to practitioners who cannot afford pretraining-scale compute. This lowers the barrier to entry for experimenting with hybrid architectures and shifts the cost calculus: you no longer need to commit to a hybrid architecture at pretraining time; you can retrofit it onto an existing model post-hoc. The 10-hour training time on 32 A100 GPUs for a 7B model is a concrete figure that makes the approach legible for production engineering teams.

The paper also implicitly de-emphasizes exact-recall as the primary metric for long-context models. The RULER results (Table 5) are included honestly — AHN fails at needle-in-a-haystack relative to full attention — but the paper's framing and benchmark selection (LV-Eval, InfiniteBench, LongBench) prioritize comprehension and reasoning over retrieval. This reflects a bet that most real-world long-context applications (document QA, summarization, multi-turn dialogue, code understanding) depend more on semantic integration than on verbatim recall of specific sentences. Whether this bet pays off depends on the application, but the paper makes a clear normative claim: if you need perfect retrieval, use full attention or retrieval-augmented generation; if you need efficient comprehension, use AHN. This is a useful clarification of the design space even if it doesn't solve the retrieval problem.

Research directions that become more attractive after this work:

  • Adaptive memory allocation: if the window size can be chosen based on task (as Figure 4 shows), can it also be chosen per-token or per-sequence based on estimated information density? The AHN framework provides the efficiency headroom to make dynamic allocation practical.
  • Verifier-guided compression: the paper shows that AHN learns selective compression (Figure 5), but this is emergent from end-to-end training. Could an explicit importance signal — perhaps from a separately trained verifier, or from the attention weights themselves — provide a stronger training signal for what to preserve?
  • Retrofitting existing models: the self-distillation recipe is general enough that it could be applied to Llama, Mistral, Gemma, or any Transformer with minimal adaptation. A replication study across model families would be immediately valuable and is now feasible given the training efficiency.

Research directions that become less attractive after this work:

  • Uniform compression of all tokens (e.g., full linear attention, pure RNN language models). The paper's pure-RNN ablation (Appendix E, Table 8) shows catastrophic degradation (LV-Eval 0.04, InfiniteBench 1.19), confirming that compressed memory alone cannot support long-context comprehension. The large window is essential — compression-only approaches are dead ends for this class of tasks.
  • Small-window hybrid architectures. The paper's large-window philosophy (32k) substantially outperforms the 64–2048 token windows used in prior work, and the ablation in Figure 4 shows monotonic improvement as window size increases up to 16–32k. Hybrid methods that use tiny attention windows and rely heavily on recurrence are architecturally suboptimal for comprehension tasks — the attention window should be as large as the hardware budget allows.

Follow-Up Research This Work Enables

Adaptive window sizing based on per-token or per-sequence importance. The paper shows that the optimal sliding window size varies by task (Figure 4: LV-Eval peaks at 32k, InfiniteBench at 96k) and that performance degrades past a certain point due to attention dilution. But the window size is fixed globally — every sequence and every position within a sequence receives the same allocation. A natural extension is to make the window size content-dependent: when the model encounters a token that refers back to something stated long ago (signaled by high attention entropy, or by a dedicated "memory-gate" prediction head), it could dynamically expand the lossless window to retain that information, or alternatively, flag certain KV pairs for retention beyond the window boundary. A strong follow-up would train a lightweight gating module (comparable in cost to the existing AHN gates α\alpha, β\beta, γ\gamma) that predicts, at each position, whether the current token should be retained in lossless memory beyond the default window — essentially a learned "do not compress" flag. The evaluation would test whether this recovers some of the exact-recall capability lost to compression (measured on RULER; Table 5) without substantially increasing memory usage. The paper's gradient visualization (Figure 5) already shows that the model can distinguish informative tokens from filler — the question is whether that signal can be operationalized to control memory allocation rather than just compression strength.

Staged memory with multiple compression levels. The current AHN framework has exactly two memory tiers: lossless (the sliding window) and compressed (the AHN state). But cognitive models of memory typically include multiple stages — sensory memory, short-term/working memory, and long-term memory with varying degrees of consolidation. A natural extension is a three-tier architecture: a small fully lossless window (e.g., 4k tokens for local coherence), a larger partially compressed region (using a lighter-weight compression that preserves more detail than the full AHN), and a fully compressed AHN state for everything beyond that. This would create a "memory hierarchy" where fidelity degrades gradually with distance rather than dropping off a cliff at the window boundary. The paper's own results motivate this: the RULER failure (Table 5) shows that the cliff is real — tokens just beyond the window are as irretrievable as tokens 100k positions away. Intermediate compression might allow approximate retrieval (e.g., "the document mentioned a specific number around position 40k") without the full cost of lossless storage. A concrete experiment would implement a two-stage AHN where the first stage uses a larger memory matrix (e.g., 2H×2H2H \times 2H) for tokens 32k–64k and the second stage uses the standard H×HH \times H for tokens beyond 64k, measuring whether RULER performance improves for needles in the 32k–64k range relative to the single-stage AHN.

Cross-model-family replication and the question of architectural dependence. All experiments in this paper use Qwen2.5-Instruct models. The paper claims AHN is a "general memory framework" (Section 2.2), but this claim is untested beyond a single architectural family. A high-priority replication would apply the exact same self-distillation recipe to Llama-3.1 (which natively supports 128k context through RoPE scaling, removing the out-of-distribution positional encoding confound discussed in Section 6), Mistral, and Gemma at comparable parameter counts. The key measurements would be: (1) Does the 0.4% parameter overhead and 40% FLOP reduction replicate across architectures? (2) Does AHN consistently outperform sliding window baselines by similar margins? (3) Do the three AHN instantiations (Mamba2, DN, GDN) maintain the same relative ordering across model families, or is the optimal instantiation model-dependent? A negative result — e.g., that AHN provides no benefit over sliding window on Llama-3.1 because that model's attention patterns are already robust to long contexts — would be equally informative, as it would delineate the boundary conditions for when hybrid architectures add value beyond context-extended full attention.

Training length scaling: where does the 24k→128k generalization break? The paper trains at 24k maximum sequence length and evaluates at 128k — a 5.3× generalization — without any analysis of when this generalization might fail. A targeted stress test would train AHN models at several maximum training lengths (8k, 16k, 24k, 32k, 64k) and evaluate at 128k, 256k, and 512k to map out the generalization frontier. The key question is whether the AHN's recurrent update is truly length-agnostic (in which case training length should not matter beyond some minimum needed to learn the compression strategy) or whether there is a "training-inference length ratio" beyond which performance degrades. The paper's randomized window training (varying the window from 32 to 8,192 tokens) provides some robustness, but if the AHN has only ever seen sequences up to 24k, its forget gate α\alpha may not have learned appropriate decay rates for 100k+ token horizons. A practical outcome of this experiment would be a recommended "training length multiplier" — e.g., "train at 0.25× target inference length for optimal efficiency-quality tradeoff."

Combining AHN with retrieval-augmented generation for exact-recall tasks. The paper's clearest failure mode is exact recall (RULER, Table 5). An obvious mitigation — which the paper does not explore — is to pair the AHN's compressed memory with a sparse retrieval mechanism that indexes specific facts for later lookup. When the model needs to answer a factoid question about something stated 80k tokens ago, it queries the retrieval index rather than relying on the compressed memory. The AHN would handle semantic comprehension (the kind of reasoning tested by LV-Eval and InfiniteBench), while retrieval handles fact-lookup. The architectural question is whether the AHN's own gates could be trained to recognize when a token contains "retrievable" information and write it to both the compressed state and a separate explicit memory store. A concrete experiment would augment the AHN-GDN model with a lightweight retriever (e.g., a sparse bag-of-words index or a dense passage retriever using the AHN's key vectors as embeddings) and measure whether RULER performance recovers toward full-attention levels while keeping memory and FLOPs well below full-attention costs. The training would need to include retrieval-augmented tasks to teach the model when to use retrieval vs. compressed memory — the ChatQA2 dataset already includes RAG-oriented examples and might serve as a starting point.

Theoretical analysis of the effective capacity of the AHN memory state. The AHN state is an H×HH \times H matrix per head (e.g., 128×128128 \times 128 for Qwen2.5), sized identically regardless of how many tokens it has compressed. As the sequence length grows, the compression ratio increases without bound, but the paper provides no analysis of when the memory saturates. A theoretical or semi-empirical analysis could measure the effective rank of htWh_{t-W} as a function of the number of compressed tokens, using singular value decomposition to track how many independent directions in the memory state are actually being used. If the effective rank saturates at, say, 80 even though the matrix is 128×128128 \times 128, that would indicate unused capacity and might motivate a smaller (more efficient) memory state. Conversely, if the effective rank approaches 128 at 128k tokens, that would suggest the memory is near capacity and longer sequences might degrade. This analysis would provide engineering guidance — "for a model with head dimension HH, the AHN memory can effectively represent up to cHc \cdot H tokens of compressed history before saturation" — that the current paper lacks.

Practical Applications and Downstream Use Cases

Cost-efficient long-document processing in enterprise settings. Organizations that need to process long documents — legal contract review (hundreds of pages), financial report analysis (annual filings, earnings call transcripts), medical record summarization (patient histories spanning years) — currently face a choice between using full-attention models with prohibitive GPU memory costs or sliding-window models that degrade on long-range dependencies. The paper's numbers are directly applicable: a 7B model with AHN reduces KV cache memory from 14.7 GB to 3.76 GB at 128k tokens (Table 9) while improving LV-Eval score from 3.62 to 6.82 (Table 2). For a document processing pipeline handling thousands of documents per day, this 74% memory reduction means either deploying on cheaper GPU instances or handling 4× higher throughput on the same hardware. The FLOP reduction (40.5%) also reduces per-document latency and energy costs. Critically, since AHNs are inactive for documents under 32k tokens, the vast majority of documents (which are shorter) incur zero overhead — the model operates as standard full attention. Only the subset of documents exceeding 32k trigger AHN activation, making this a "pay-as-you-go" efficiency model.

On-device or edge deployment of long-context LLMs. The paper's 3B model with AHN achieves a memory cache of 2.45 GB at 128k context (Table 9, vs. 9.44 GB for full attention). While 2.45 GB is still beyond most mobile devices, it represents a 4× reduction that brings long-context inference closer to the edge. For a 1B model (not tested in the paper but architecturally identical), the KV cache with AHN would be proportionally smaller — potentially under 1 GB at 128k. This enables scenarios like on-device email thread summarization, local document QA without cloud upload, or privacy-preserving analysis of long personal chat histories. The training recipe's efficiency (10 GPU-hours for 7B, likely single-digit GPU-hours for smaller models) makes it feasible for device manufacturers or app developers to train AHN modules for their specific base models and domains. The frozen-base-weights property is especially valuable for deployment: the base model's safety alignment and behavioral characteristics are preserved exactly, reducing the validation burden compared to fine-tuning approaches that modify base model weights.

Efficient inference serving for long-context APIs. LLM API providers (e.g., serving Qwen or Llama models at scale) face a sharp cost increase when users submit long prompts. A 128k-token prompt with full attention might require a high-memory GPU instance (A100 80GB or H100) and tie up that instance for the duration of generation. With AHN, the same 7B model's KV cache shrinks from 14.7 GB to 3.76 GB, potentially allowing the model to run on a cheaper instance or enabling higher batch sizes by fitting more concurrent requests on the same GPU. The FLOP reduction (40.5% at the model level; Table 2) similarly reduces per-token generation latency for prompts exceeding 32k. Since the API provider controls the inference stack, they can deploy AHN-augmented models transparently — users submitting prompts under 32k see identical behavior to the base model, while users exceeding 32k automatically benefit from the efficiency improvements. For a provider serving millions of requests daily, a 40% reduction in long-context compute translates directly to infrastructure cost savings. The AHN training cost (approximately 10 GPU-hours per model variant) is negligible amortized over the serving lifetime.

Retrofitting existing fine-tuned models without retraining them. A common deployment pattern is to fine-tune a base LLM on domain-specific data (e.g., a medical Qwen fine-tuned on clinical notes, a code Qwen fine-tuned on repositories). These fine-tuned variants inherit the base model's architecture and most weights, but retraining them from scratch with AHN would require redoing the domain fine-tuning — expensive and potentially lossy. The AHN framework's key property — all base model weights (including domain-fine-tuned weights) are frozen — means the AHN module can be trained on top of the already-fine-tuned model without touching the domain-specific weights. The self-distillation teacher is the domain-fine-tuned model itself (with full attention), and the student is the same model with AHN-augmented attention. The AHN learns to compress long contexts while preserving the teacher's domain-specific output distribution. This enables domain-adapted long-context models at the cost of training only the 0.4% AHN parameters — a few GPU-hours on domain-specific long documents, rather than a full fine-tuning run. For organizations that have invested in domain-specific model customization, this provides a path to long-context efficiency without discarding that investment.