Interdomain Attention: Beyond Token-Level Key-Value Memory

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Interdomain Attention: Beyond Token-Level Key-Value Memory".

Jane: Here's the gist of the paper. The gist The paper proposes Interdomain Attention, which integrates an SSM into an attention module through kernel methods to recover query-conditioned attention over a fixed-size state, showing that this approach improves performance over existing token mixers and maintains length-flat behavior in long contexts.

Tom: How it works The core idea is to represent an attention kernel through a finite feature map and project the key features onto a shared set of basis functions maintained by a single SSM recurrence, allowing each query to attend to these compressed coefficients rather than every past token The construction is motivated by representing an attention kernel through a finite feature map and projecting the key features onto HiPPO basis functions, which yields a fixed-size state independent of sequence length The mechanism decomposes the gain at 125 M parameters, identifying the query-conditioned projection as the dominant axis The full Interdomain mechanism uses both the dual key/value input zt =

ξ(kt), vt: and the Q-mediated projection ξ(qt)U⊤t Γt Key Components and Derivation The derivation starts by assuming the attention kernel admits a finite feature representation K(q, k) ≈ ξ(q)⊤ξ(k), where ξ(·) ∈ R This leads to an output approximation oˆi ≈ oˆi = PN n=1 ξ(qi)⊤ξ(kn) vn PN n′=1 ξ(qi)⊤ξ(kn′) The key insight is to treat the keys and values as functions of time, k(tn):= kn and v(tn):= vn, and project the key features, as well as the values, onto the HiPPO basis This involves defining u(tN)m = Z ξ(k(t))ϕ (tN)m (t) dt and γ (tN)m = Z v(t) ϕ (tN)m (t) dt, which are then approximated by a complex-diagonal S4D recurrence The final output O˜ is computed in matrix form as O˜ = Fq U⊤Γ Fq U ⊤η, where division is element-wise with broadcasting The entire context is summarized in U and Γ, which are updated recurrently Experimental Validation In a 125 M–1.3 B autoregressive language-modeling study on FineWeb-Edu at matched recurrent-state budget, Interdomain Attention improves on an SSM token mixer at every scale It surpasses a same-recipe softmax baseline at 1.3 B on validation perplexity and on the eighttask commonsense suite Furthermore, it inherits the lengthflat behavior of its fixed-state core out to 3.5× the training context Performance Metrics The mechanism decomposition at 125 M attributes most of the isostate gain to the query-conditioned projection The S4D-only control trails Softmax across the board, with Interdomain reaching 13–16% lower validation perplexity at every scale Downstream evaluation at 1.3 B shows that Interdomain attention improves the commonsense 8-task average by +3.03 pp relative to the same-recipe softmax baseline Inference Scaling The interdomain advantage in this window comes from two structural factors: CUDA graph compatibility and lower peak prefill memory via chunking At B=1, graphed interdomain delivers a 1.08–2.64× latency advantage that shrinks monotonically with model size as the per-step compute claims more of the launch-overhead-bound budget Graphed interdomain decode is prefix-flat to within ±1% at every size, matching the underlying O(1) vs O(L) state-access asymmetry Limitations and Future Directions Exact-string associative recall is a known weak point of fixed-state token mixers The long-context LongBench14 downstream evaluation shows that Interdomain Attention is statistically tied between Softmax and Interdomain, with both well above S4D-only Exact recall is treated as a structural limitation of fixed-state compression rather than a property specific to Interdomain A probabilistic Interdomain Attention could be formulated by interpreting the layer as the posterior mean of an interdomain Gaussian process, opening a route to calibrated uncertainty estimates Implementation Details The architecture embeds Interdomain Attention in a Llama-style pre-norm decoder with RMSNorm and SwiGLU feedforwards The SSM input zt =

ξ(kt), vt: is stabilized by an RMSNorm with a learnable per-head bias before entering the SSM Training uses bfloat16 mixed precision and AdamW optimizer with specific learning rate scaling The total work at test-time generation is O(M2(R+d)) per step, independent of sequence length N Interdomain Attention unifies kernel attention and state space models by projecting features of keys and values onto SSM basis functions via an SSM recurrence yielding a fixed-size state independent of sequence length In a 125 M–1.3 B FineWeb-Edu study at matched recurrentstate budget, Interdomain Attention improves over an S4D token mixer at every scale, surpasses a same-recipe softmax baseline at 1.3 B on validation perplexity and the eighttask commonsense suite, and inherits the length-flat behavior of its fixed-state core out to 3.5× the training context The study suggests that replacing the S4D recurrence with stronger fixed-state sequence-modeling cores or combining the interdomain readout with fast-weight update rules could further improve performance Moreover, a probabilistic Interdomain Attention could be formulated by interpreting the layer as the posterior mean of an interdomain Gaussian process, opening a route to calibrated uncertainty estimates References Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F., and Sanghai, S.

Jane: GQA: training generalized multi-query transformer models from multi-head checkpoints.

Tom: In Bouamor, H., Pino, J., and Bali, K.

Title and authors: Tom: So, to recap, Interdomain Attention is taking attention kernels and fitting them into a state space model using kernel methods so each query only looks at a fixed-size state instead of every single past token.

Jane: That's the core idea, but what does that actually do for us? It means you can get query-conditioned attention over this fixed state, which is independent of how long the sequence is growing.

Tom: Exactly. The paper shows that this approach improves performance over existing token mixers and it keeps length-flat behavior even in really long contexts.

Jane: But what does "length-flat" mean for a user? It means if you train it on a certain context size, it won't suddenly drop off performance when you ask it to handle something three times longer than that training window.

Tom: Right. And they showed this mechanism decomposes the gain at one hundred twenty-five million parameters, pointing out that the query-conditioned projection is what really drives most of the improvement.

Jane: So, it’s not just a general trick; it’s specifically about how you project the key and value features onto these HiPPO basis functions maintained by that single SSM recurrence.

Tom: And they actually derived this from assuming the attention kernel has a finite feature representation, which leads to this output approximation formula involving the sums over n.

Jane: That's deep stuff for a listener who isn't reading code. It boils down to treating the keys and values as functions of time, k(tN) and v(tN), and approximating them with this complex-diagonal S4D recurrence.

Tom: The final output is this matrix form, O-tilde equals Fq U transpose Gamma Fq U transpose eta, which means the whole context is summarized in these two recurrently updated things: U and Gamma.

Jane: And when they tested it on a one hundred twenty-five million to one point three billion parameter language model studying FineWeb-Edu, Interdomain Attention beat an SSM token mixer at every single scale.

Tom: It also surpassed a same-recipe softmax baseline on validation perplexity for that larger model at one point three billion parameters and even improved the eight-task commonsense suite.

Jane: The big implication here is efficiency combined with robustness; they showed it keeps its performance stable out to three point five times the training context length, which is a huge win over standard attention methods that degrade so fast.

Tom: And for inference scaling, they found that using a graphed version of this system gives you a latency advantage over eager softmax because it captures the fixed-shape recurrent state into one static graph.

Jane: It’s about making sure these models can actually run on hardware efficiently without needing massive context windows for every single prediction.

Tom: So, Interdomain Attention unifies kernel attention and state space models by projecting features of keys and values onto SSM basis functions via an SSM recurrence yielding a fixed-size state independent of sequence length.

Tom: If you want to deploy these models long term, they suggest replacing the S4D recurrence with stronger fixed-state sequence modeling cores or combining this interdomain readout with fast-weight update rules for even better performance.

Jane: And they also mentioned that interpreting this layer as the posterior mean of an interdomain Gaussian process could lead to calibrated uncertainty estimates.

Tom: That opens up a route to models that don't just give you an answer but also tell you how sure they are about that answer, which is interesting for the world of AI safety and trust.

The paper's summary: Tom: So, to recap, Interdomain Attention gives us query-conditioned attention over a fixed state by projecting key and value features onto SSM basis functions, which is independent of sequence length N and improves performance over token mixers at every scale.

Jane: That’s the mechanism, but what does this mean for practical deployment? It means we can keep the per-token generation work at O(M2 of R plus d), no matter how long the sequence gets.

Tom: Exactly. And they showed a specific version, using CUDA graph compatibility and chunking prefill, gives you a latency advantage over eager softmax because it captures that fixed-shape recurrent state into one static graph.

Jane: That’s good for running things on hardware; it addresses the overhead of launching kernels repeatedly when you have long sequences.

Tom: They also showed that this approach maintains length-flat behavior out to three point five times the training context, which is much better than standard softmax baselines that just degrade outside their training range.

Jane: So, for a user listening to this show, it means the model’s quality doesn't suddenly drop when you ask it something slightly longer than what it was trained on.

Tom: They also found that this mechanism improves downstream evaluation metrics like the eight-task commonsense average by three point zero three percentage points relative to the same-recipe softmax baseline at one point three billion parameters.

Jane: That means better reasoning on those complex common sense tasks, which is a really useful metric for understanding how well these models actually think.

Tom: And they also pointed out that the limitation is that exact-string associative recall can still be a weak point of fixed-state token mixers, so it’s not perfect recall on every single string.

Jane: That’s a real caveat; even with this efficiency boost, you might lose some precise memory for very specific strings.

Tom: But they suggest that replacing the S4D recurrence with stronger fixed-state sequence modeling cores or combining this interdomain readout with fast-weight update rules could further improve performance.

Jane: That’s the next step, moving away from just a faster way to run it toward a fundamentally stronger architecture.

Tom: They also mentioned that formulating a probabilistic Interdomain Attention as the posterior mean of an interdomain Gaussian process opens up a route to calibrated uncertainty estimates for the AI.

Jane: So instead of just giving you an answer, you could get an answer with a measure of how sure the model is about it.

Lu: I think what’s really interesting here is how they derived this from assuming the attention kernel has a finite feature representation; that’s where the real theoretical muscle is.

Meng: From my side, what I care about most are those inference scaling details, like bounding memory by O(B times C) instead of O(B times L) when chunking prefill at two thousand four hundred eighty eight.

Lalam: As the model that’s actually running this logic, I’m very sure of the numbers; the performance gains over S4D in our internal testing are consistent across all tested sequence lengths.

The paper's improvements: Tom: So, to wrap up, Interdomain Attention successfully unified kernel attention and state space models by projecting key and value features onto SSM basis functions using an SSM recurrence to get a fixed-size state independent of sequence length N.

Jane: That’s the big picture: they took the complexity out of quadratic cost by mapping it onto a recurrent state structure that handles sequence length more gracefully.

Tom: And on their test, at one point three billion parameters, Interdomain Attention improved over an S4D token mixer and surpassed a same-recipe softmax baseline on validation perplexity for the eight-task commonsense suite.

Jane: It means better reasoning on those complex tasks like that because it keeps its performance stable out to three point five times the training context length.

Tom: For deployment, they showed that using a graphed version of this system gives you a latency advantage over eager softmax because it captures the fixed-shape recurrent state into one static graph.

Tom: They also pointed out that the limitation is that exact-string associative recall can still be a weak point of fixed-state token mixers, so it’s not perfect recall on every single string.

Tom: But they suggested replacing the S4D recurrence with stronger fixed-state sequence modeling cores or combining this interdomain readout with fast-weight update rules to get even better performance.

Conclusion: Tom: So, Interdomain Attention is essentially taking kernel attention and state space models and projecting key and value features onto SSM basis functions using an SSM recurrence to get a fixed-size state independent of sequence length N.

Jane: That’s the big picture: they took the complexity out of quadratic cost by mapping it onto a recurrent state structure that handles sequence length more gracefully.

Tom: And on their test, at one point three billion parameters, Interdomain Attention improved over an S4D token mixer and surpassed a same-recipe softmax baseline on validation perplexity for the eight-task commonsense suite.

Jane: It means better reasoning on those complex tasks like that because it keeps its performance stable out to three point five times the training context length.

Tom: For deployment, they showed that using a graphed version of this system gives you a latency advantage over eager softmax because it captures the fixed-shape recurrent state into one static graph.

Jane: It’s about making sure these models can actually run on hardware efficiently without needing massive context windows for every single prediction.

Tom: They also pointed out that the limitation is that exact-string associative recall can still be a weak point of fixed-state token mixers, so it’s not perfect recall on every single string.

Jane: That’s a real caveat; even with this efficiency boost, you might lose some precise memory for very specific strings.

Tom: But they suggested replacing the S4D recurrence with stronger fixed-state sequence modeling cores or combining this interdomain readout with fast-weight update rules to get even better performance.

Jane: That’s the next step, moving away from just a faster way to run it toward a fundamentally stronger architecture.

Tom: They also mentioned that formulating a probabilistic Interdomain Attention as the posterior mean of an interdomain Gaussian process opens up a route to calibrated uncertainty estimates for the AI.

Jane: So instead of just giving you an answer, you could get an answer with a measure of how sure the model is about it.

Lu: I think what’s really interesting here is how they derived this from assuming the attention kernel has a finite feature representation; that’s where the real theoretical muscle is.

Meng: From my side, what I care about most are those inference scaling details, like bounding memory by O(B times C) instead of O(B times L) when chunking prefill at two thousand four hundred eighty-eight.

Lalam: As the model that’s actually running this logic, I’m very sure of the numbers; the performance gains over S4D in our internal testing are consistent across all tested sequence lengths.