🤖 machine learning

Interdomain Attention: Beyond Token-Level Key-Value Memory

The paper proposes Interdomain Attention, a novel architecture that integrates State Space Models (SSMs) into attention mechanisms via kernel methods to achieve query-conditioned attention over a fixed-size state, thereby combining the scalability and length-robustness of SSMs with the performance benefits of content-based matching.

Original authors: Naoki Kiyohara, Harrison Bo Hua Zhu, Riccardo El Hassanin, Zhuo Sun, Wenlong Chen, Samir Bhatt, Yingzhen Li

Published 2026-05-26

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Naoki Kiyohara, Harrison Bo Hua Zhu, Riccardo El Hassanin, Zhuo Sun, Wenlong Chen, Samir Bhatt, Yingzhen Li

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

The Big Problem: The "Too Much Stuff" Dilemma

Imagine you are trying to write a story, but you have a terrible memory.

The Old Way (Standard Transformers): To write the next sentence, you have to look back at every single word you've written so far. If you write a 100-page book, your brain has to hold all 100 pages in your head at once to find the right connections. This is incredibly slow and requires a massive amount of mental energy (computing power). As the story gets longer, your brain gets overwhelmed.
The Alternative Way (State Space Models/SSMs): To save energy, you summarize the whole story into a single, tiny note card. You only look at this note card to write the next sentence. This is super fast and efficient, no matter how long the story gets. However, because you only have a tiny note card, you often miss specific details. You might forget the name of a character mentioned three chapters ago.

The Goal: The authors wanted to build a system that has the speed and efficiency of the tiny note card but the memory and detail of looking at the whole book.

The Solution: "Interdomain Attention"

The authors created a new method called Interdomain Attention. Think of it as a smart librarian who manages a library in a very specific way.

1. The "Compressed Library" (The SSM Core)

Instead of keeping every single book (token) on a shelf, the librarian uses a special machine (an SSM) to compress the entire history of the story into a fixed-size set of "summary coefficients."

Analogy: Imagine you have a 1,000-page novel. Instead of keeping all 1,000 pages, you distill the story into 64 specific "themes" or "vibes" (like "the hero's journey," "the villain's motive," "the setting's mood"). These 64 themes are your "state." No matter if the story is 10 pages or 10,000 pages, you only ever need to hold these 64 themes in your hand.

2. The "Smart Query" (The Attention Part)

When you want to write the next sentence, you don't just look at the 64 themes blindly. You ask a specific question (a "query").

The Magic Trick: The system takes your question and translates it into the same "language" as those 64 themes. It then checks: "Which of these 64 themes is most relevant to my current question?"
The Result: You get the best of both worlds. You are looking at a compressed summary (fast!), but you are selecting the right parts of that summary based on what you are currently thinking about (smart!).

How It Works (The "Kitchen" Analogy)

Imagine you are a chef making a soup (the output).

Standard Attention: You have a giant pot of ingredients (the whole history). To add salt, you have to taste every single ingredient in the pot to decide how much salt to add. This takes forever as the pot gets bigger.
Standard SSM: You have a tiny spice rack with only 64 jars. You just grab a jar and add it. It's fast, but you can't taste the specific ingredients anymore.
Interdomain Attention:
1. You have a machine that constantly updates a fixed-size spice rack (the 64 themes) based on everything you've cooked so far.
2. When you want to add salt, you don't taste the whole pot. Instead, you hold a special tasting spoon (the query feature map) that instantly tells you: "Based on the current flavor, you need to mix 30% of the 'Tomato' jar and 70% of the 'Herb' jar."
3. You mix those specific amounts from your fixed spice rack. It's fast because the rack is small, but it's accurate because your spoon knows exactly how to mix them for the current moment.

What the Paper Actually Found

The authors tested this new "Interdomain Attention" on language models ranging from small (125 million parameters) to large (1.3 billion parameters). Here are their key claims:

It Beats the "Tiny Note Card" (SSM): In every test, Interdomain Attention was better at writing text than the standard SSM method. It understood context better while staying fast.
It Beats the "Giant Pot" (Standard Attention) at Scale: At the largest size they tested (1.3 billion parameters), Interdomain Attention actually wrote better text (lower "perplexity" and better common sense scores) than the standard method that looks at every word.
It Doesn't Forget Long Stories: The biggest win is length. Standard methods get confused and make mistakes when the story gets longer than they were trained on. Interdomain Attention stayed calm and consistent even when the story was 3.5 times longer than what it was trained on. It didn't get slower or more confused; it just kept working.
The Secret Sauce: They did a "mechanism decomposition" (a fancy way of taking the machine apart to see which part did the work). They found that the query-conditioned projection (the "smart spoon" that translates your question into the summary language) was the main reason for the success. Without that specific step, the system just acted like a standard SSM and performed poorly.

The Trade-off (What It Can't Do)

The paper is honest about one weakness: Exact Recall.

Analogy: If you ask the system, "What was the exact phone number of the character mentioned on page 4?" the standard method (looking at the whole book) is great at finding it. Interdomain Attention, because it relies on a compressed summary, is not as good at retrieving exact strings of text (like phone numbers or specific names) if they aren't part of the main "themes."
However, the authors note this is a limitation of any system that compresses information, not just their new method.

Summary

Interdomain Attention is a new way for AI to remember things. Instead of trying to hold the whole world in its head (slow) or just a tiny note (forgetful), it keeps a smart, compressed summary of the world. When it needs to write, it uses a special translator to pick the exact right pieces of that summary. This makes it fast, efficient, and surprisingly good at handling very long stories without getting confused.

Technical Summary: Interdomain Attention

Problem Statement

Transformers and deep State Space Models (SSMs) represent two opposing design paradigms for sequence modeling. Standard attention mechanisms (Vaswani et al., 2017) utilize a growing key-value (KV) cache, routing queries through content-based matching. While this provides sharp recall, it incurs quadratic computational costs ( $O(N_q N)$ ) and linear memory growth ( $O(N)$ ) with respect to sequence length. Conversely, deep SSMs (e.g., S4, Mamba) compress the entire context into a fixed-size recurrent state, updating in $O(1)$ time per step. However, this compression comes at the cost of losing direct query-key matching; the fixed state is not directly addressable by specific queries, often leading to weaker performance in tasks requiring precise retrieval or content-based attention.

Hybrid architectures have attempted to interleave these approaches, and linear attention variants have sought to reduce attention costs directly. This paper proposes a different route: integrating an SSM directly into the attention mechanism to maintain the efficiency of a fixed-size state while enabling each query to attend to the compressed history.

Methodology: Interdomain Attention

The core contribution is Interdomain Attention, a token mixer that unifies kernel attention and state space models. The method is derived from a kernel-regression view of attention and a HiPPO-style basis projection.

Theoretical Derivation

Kernel Feature Map: The attention kernel $K(q, k)$ is approximated by a finite feature map $\xi(q)^\top \xi(k)$ .
Basis Projection: Instead of storing all keys and values, the method projects the key features $\xi(k)$ and values $v$ onto a shared set of basis functions maintained by a single SSM recurrence (specifically HiPPO-style). This results in a fixed-size state independent of sequence length.
Compressed Coefficients: The integrals of these projections yield three state components:
- $U$ : Basis-projection coefficients for keys (inducing variables).
- $\Gamma$ : Value projections.
- $\eta$ : Normalizing constants.
Query-Conditioned Readout: At each position, the query $q$ is mapped to the feature space $\xi(q)$ . The output is computed by attending to the compressed coefficients rather than individual tokens:
$\tilde{o} = \frac{\sum_m (\xi(q)^\top u_m) \gamma_m}{\sum_m (\xi(q)^\top u_m) \eta_m}$
In matrix form: $\tilde{O} = \frac{F_q U^\top \Gamma}{F_q U^\top \eta}$ , where $F_q$ is the query feature matrix.

Scalable Implementation

The practical implementation is a learned relaxation of the theoretical derivation, designed for stability and efficiency:

Feature Map: Uses a DeltaNet-style neural feature map $\xi(x) = \text{SiLU}(\text{SC}(x)) / \|\text{SiLU}(\text{SC}(x))\|^2$ , where SC is a causal depthwise convolution.
Recurrent Core: Approximates the HiPPO projection using a complex-diagonal S4D recurrence. The SSM input concatenates the key feature $\xi(k)$ and the value $v$ .
Normalization: Applies input RMSNorm with learnable per-head biases to the key and value halves before entering the SSM.
Denominator-Free Readout: To avoid mixing incompatible scales caused by input rescaling, the normalizing term $\eta$ is dropped, resulting in a linear attention readout: $\tilde{O} = F_q U^\top \Gamma$ .
Architecture: Embedded in a Llama-style decoder with pre-normalization, SwiGLU feedforwards, and rotary position embeddings (RoPE).

Key Contributions

Construction: A query-conditioned fixed-state token mixer derived from kernel regression and HiPPO basis projection, with a clear distinction between the ideal theoretical derivation and the scalable implementation.
Mechanism Decomposition: An ablation study at 125M parameters identifies the query-conditioned projection as the dominant factor driving performance gains over standard S4D mixers. The dual key/value input contributes a smaller, secondary gain.
Scaling Study: A comprehensive evaluation on FineWeb-Edu from 125M to 1.3B parameters under a matched recurrent-state budget.

Experimental Results

The study compares Interdomain Attention against an S4D-only token mixer and a standard softmax attention baseline.

Performance vs. S4D: Interdomain Attention improves upon the S4D token mixer at every scale (125M–1.3B), achieving 13–16% lower validation perplexity.
Performance vs. Softmax: At 1.3B parameters, Interdomain Attention surpasses the same-recipe softmax baseline on validation perplexity (7.98 vs. 8.63) and on an eight-task commonsense suite (+3.03 percentage points average accuracy). At 125M, it is roughly tied with softmax.
Length Extrapolation: Leveraging its fixed-state core, Interdomain Attention exhibits "length-flat" behavior. While the softmax baseline's perplexity degrades significantly beyond the training context (4K), Interdomain maintains stable performance up to 14K (3.5× the training context).
Recall Limitations: The paper notes that Interdomain Attention, like other fixed-state models, struggles with exact-string associative recall (e.g., needle-in-a-haystack tasks) compared to softmax attention, identifying this as a structural limitation of fixed-state compression rather than a specific flaw of the Interdomain mechanism.

Significance and Claims

The paper claims that Interdomain Attention successfully bridges the gap between the efficiency of fixed-state SSMs and the content-based retrieval capabilities of attention. By projecting features onto an SSM basis, it allows queries to attend to a compressed history without the quadratic cost or linear memory growth of standard attention.

The primary significance lies in demonstrating that a query-conditioned projection onto a fixed recurrent state can outperform both pure SSMs and, at larger scales, standard softmax attention, while inheriting the desirable length-extrapolation properties of SSMs. The authors suggest this approach offers a viable path for long-context language modeling where memory efficiency and consistent performance across varying context lengths are critical. Future work directions mentioned include integrating stronger fixed-state cores (e.g., Mamba-3) or combining the readout with fast-weight update rules.