PaperlayerView exact revision on arXiv

Source-linked AI summary

Interdomain Attention: Beyond Token-Level Key-Value Memory

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

arXiv:2605.24330v2cs.LGstat.ML

TL;DR

Softmax attention has high computational and memory costs, while common hybrid architectures stack attention and recurrent-memory layers. IA instead combines attention and recurrent memory in one layer by compressing key features and values into recurrent states. Across language and genomic evaluations, it improves on S4D and is competitive with or better than the compared attention and genomic models.

  • Problem

    Softmax attention has O(N^2) computation and O(N) KV-cache memory, while common hybrids stack attention and recurrent-memory layers, complicating attribution and often requiring extensive architecture search.

    Source: §1 · pp. 1–2

  • Method

    IA uses a kernel-regression view of attention to compress key features and values into recurrent memory and implicitly reconstructs them for attention-like readout.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    Across 125M–1.3B language models at matched state size, IA improves validation perplexity over S4D; at 1.3B it also beats softmax attention on validation perplexity and common-sense reasoning, while Caduceus-based IA is competitive or better on genomic evaluations.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    Ablations identify the attention-like, query-dependent readout as the primary source of IA’s gains over SSM-based architectures.

    Source: §1 · p. 2

    IA training throughput is 2.8 times slower than PyTorch Flash SDPA, primarily because its GPU kernel is less optimized than FlashAttention2.

    Source: §6 · p. 10

Abstract

from arXiv · show

Softmax attention-based transformers and recurrent memory-based models (e.g., state space models (SSMs)) form the foundation of modern sequence models. Hybrid architectures aim to combine the expressivity from attention and the computational efficiency from recurrent memory, yet so far common practices simply stack both attention and recurrent layers. We propose Interdomain Attention (IA) as a novel hybrid attention architecture that mixes attention and recurrent memory in the same layer. Motivated by a kernel regression view of softmax attention, IA compresses values and kernel features of keys into a recurrent memory parameterized by an SSM. Then IA approximates dot-product attention by implicitly reconstructing the attention matrix and values from memory states; importantly this construction maintains SSM's computational efficiency advantage. We demonstrate IA's efficacy by designing and pre-training IA network architectures, where ablation studies identify attention-like readout as the primary source of gain over SSMs. Across 125M-1.3B models pre-trained on text data, IA improves over an S4D baseline at matched memory size. It also outperforms softmax attention-based Transformer in perplexity and accuracy on common-sense reasoning tasks. Based on Caduceus backbone, IA matches or outperforms the original SSM-based architecture in both perplexity and downstream tasks on genomic data.

1 INTRODUCTION

Interdomain Attention (IA) combines recurrent memory and dot-product attention within one layer, addressing the cost of softmax attention and the recall limitations of recurrent models. Across language and genomic evaluations, IA performs competitively or better than the compared architectures, with ablations identifying its attention-like readout as a key source of gains.

1 INTRODUCTION

  • Softmax attention has O(N^2) computation and O(N) KV-cache memory, while SSMs use fixed-size recurrent state and O(N) total computation but perform worse on precise-recall tasks.Source: §1 · p. 1
  • Hybrid models commonly stack recurrent-memory and attention layers, complicating attribution of each layer type’s importance and often requiring extensive architecture search.Source: §1 · p. 2
  • IA interprets softmax attention as kernel regression, compresses key features and values into recurrent memory, then implicitly reconstructs them for dot-product attention.The implicit attention step retains recurrent memory’s O(1) computational cost per token.Source: §1 · p. 2
  • Across 125M–1.3B language models, IA improves validation perplexity over softmax attention and S4D at matched state size; at 1.3B it beats softmax attention on seven of eight reasoning tasks and their average.On 14 recall tasks, 1.3B IA also has higher mean log-likelihood than S4D; in genomics, Caduceus-IA closely matches Caduceus on HG38, while Caduceus-IA-PS outperforms Caduceus-PS on 14 of 18 downstream tasks.Source: §1 · p. 2
  • Ablations in language and genomics identify IA’s attention-like, query-dependent readout as key to improved performance over SSM-based architectures.Source: §1 · p. 2

2 BACKGROUND

The background presents softmax attention as kernel regression over key-value pairs and SSMs as recurrent mechanisms for compressing input history. Figure 1 contrasts standard token-level attention with IA’s readout from compressed interdomain states.

2 BACKGROUND

  • Standard dot-product attention can be viewed as kernel regression over key-value pairs, with each query receiving the model’s prediction.The cited prior work relates this formulation to the Nadaraya–Watson estimator.Source: §2 · p. 2
  • Figure 1 contrasts standard attention’s Nq × N score matrix with IA’s Nq × M interdomain attention matrix, using recurrently compressed key and value projections.IA maps queries and keys into a shared feature space and reads out from states U and Γ.Source: Figure 1 · p. 3, §2 · p. 3
  • SSMs maintain a compressed representation of input history through recurrence, with HiPPO deriving transition matrices that maintain projections onto time-varying orthogonal polynomial bases.Source: §2 · pp. 2–3

3 INTERDOMAIN ATTENTION

Interdomain Attention (IA) combines recurrent memory and dot-product attention within one layer by projecting key features and values into a fixed-size recurrent state. Its practical design uses SSM recurrences and chunkwise execution, with optional refinements to query construction, memory updates, and value paths.

3.1 IMPLICIT DOT-PRODUCT ATTENTION VIA INTERDOMAIN KERNEL FEATURES

  • IA approximates the attention kernel with finite features and projects key features and values onto HiPPO basis functions for recurrent memory.The key coefficients reconstruct key features implicitly, while the coefficients are updated online with the HiPPO ODE.Source: §3.1 · pp. 3–4
  • IA projects key features and values onto a shared temporal basis, then uses query features to retrieve from those projections through an interdomain attention matrix.The recurrent states U, Γ, and η have sizes independent of sequence length; practical readout drops the softmax denominator.Source: §3 · p. 3, §3.1 · p. 4

3.2 A PRACTICAL IMPLEMENTATION OF INTERDOMAIN ATTENTION

  • IA uses a complex diagonal S4D recurrence for function projections, with chunkwise execution combining parallel work within chunks and sequential state passing across chunks.IA has O(N) training computation; at inference it stores recurrent state of size O(M(R + d)) independent of N, with O(M(R + d) + M^2) work per decoding step.Source: §3.2 · pp. 4–5

3.3 OPTIONAL ARCHITECTURAL ENHANCEMENTS

  • Optional enhancements include nonlinear query projection, selective writing and memory decay, and value-path and output-gating refinements.The nonlinear query adds a residual MLP; selective writing uses input-dependent gating to control information stored in memory.Source: §3.3 · p. 5

4 RELATED WORK

Prior work reduces attention cost through implementation changes or alternative mechanisms, and recurrent memory and linear attention offer distinct ways to summarize history. IA is presented as a hybrid operator that combines recurrent compression and content-based retrieval within one layer.

Efficiency in attention mechanisms

  • FlashAttention retains quadratic arithmetic while reducing GPU memory transfers, whereas Performer approximates the softmax kernel with positive random features for linear time and space.Flash Linear Attention uses chunkwise parallel computation and kernel fusion; IA also adopts chunkwise execution.Source: §4 · p. 6

SSMs and Linear Attention

  • SSMs and linear attention maintain fixed-size recurrent states, but IA represents keys and values with shared-basis projection coefficients and combines them at readout.Unlike linear attention’s accumulated key–value products, IA combines the separate key and value representations during retrieval.Source: §4 · p. 6

Hybrid architectures

  • Many hybrid architectures stack recurrent or convolutional layers with attention, while IA integrates recurrent compression and content-based retrieval within one layer.The related-work examples include Jamba, Griffin, Qwen3.8-Next, Kimi K3, and Evo 2; H3 combines SSMs and multiplicative interactions within a layer.Source: §4 · p. 6

5 EXPERIMENTS

Experiments evaluate IA on language and genomic tasks, compare it with attention and recurrent-memory baselines, and ablate architectural choices. IA performs strongly across the reported evaluations, with results varying by task and configuration.

5.1 LANGUAGE MODELING

  • At 1.3B on FineWeb-Edu, IA achieves validation perplexity 7.98 versus 8.62 for Softmax and 9.13 for S4D-only, and exceeds Softmax by +3.03% on eight-task commonsense average accuracy.Source: §5.1 · p. 7
  • At contexts up to 14K, IA maintains nearly constant perplexity, while Softmax perplexity rises from 14.21 to 62.06; IA’s decoding speedup ranges from 1.11× at batch 1 to 0.69–0.94× at batches 8–64.The context-length evaluation is 3.5× the 4K training context; the speed comparison uses CUDA-graph capture.Source: §5.1 · p. 7, Figure 2 · p. 7

5.2 GENOMIC SEQUENCE MODELING

  • Caduceus-IA replaces Mamba layers while preserving bidirectional and reverse-complement handling; after 10,000 steps at context length 1024, its validation perplexity is 2.614 versus 2.612 for native Caduceus.A separate 7.7M-parameter configuration is pretrained at context length 131,072 for 50,000 updates; downstream evaluations use models pretrained at 1K or 131K contexts.Source: §5.2 · p. 8
  • On GenomicBenchmarks, Caduceus-IA-Ph leads mean accuracy on three of eight tasks and shares the lead on two; IA variants beat corresponding Caduceus baselines on 14 and 12 of 18 Nucleotide Transformer tasks.Caduceus-IA-Ph leads eight H3/H4 tasks, while NT-v2 leads all promoter and splice-site tasks.Source: §5.2 · p. 8, Figure 3 · p. 8
  • For variant-effect prediction, Caduceus-IA-PS reports higher mean AUROC than HyenaDNA and NT-v2 across distance groups, but its baseline comparisons vary and Caduceus-IA-Ph underperforms Caduceus-Ph in every group.Against its respective baseline, Caduceus-IA-PS matches 0.679 within 30 kb, scores 0.620 versus 0.612 at 30–100 kb, and scores 0.574 versus 0.586 beyond 100 kb.Source: §5.2 · pp. 8–9

5.3 ANALYSIS OF THE INTERDOMAIN ATTENTION ARCHITECTURE

  • The attention-like, query-conditioned readout is the largest contributor to IA’s improved performance; final configurations reduce validation perplexity from 16.595 to 16.008 in language modeling and from 2.737 to 2.614 in genomics.Language gains come from most tested enhancements except value skip and selective writing; genomic gains come mainly from Q/K convolution, value convolution with SiLU, and residual query MLP.Source: §5.3 · p. 9, Figure 4 · p. 9, Table 3 · p. 9

6 CONCLUSION AND FUTURE WORK

The conclusion presents IA as a recurrent-memory approach to approximating dot-product attention, reports competitive language and genomic results, and identifies slower training throughput as a limitation.

6 CONCLUSION AND FUTURE WORK

  • IA projects key features and values into recurrent memory to approximate dot-product attention, improving over matched-state S4D from 125M to 1.3B and outperforming softmax attention at 1.3B on validation perplexity and commonsense reasoning.Caduceus-based IA architectures show competitive or improved performance; ablations identify attention-like readout as the primary performance source.Source: §6 · p. 10
  • IA training throughput is 2.8-times slower than PyTorch Flash SDPA, primarily because its GPU kernel is less well optimized than FlashAttention2.Source: §6 · p. 10
  • Future work includes applying interdomain projections to Mamba-3 or Gated DeltaNet, optimizing GPU kernels, and formalizing a unified probabilistic framework.Source: §6 · p. 10

REPRODUCIBILITY STATEMENT

The reproducibility statement says the manuscript enumerates algorithm and experiment details and promises to release model weights and code if the paper is accepted.

REPRODUCIBILITY STATEMENT

ETHICS STATEMENT

The ethics statement describes the work as using public datasets without human subjects or new data collection, while noting a possible dual-use risk for biological sequence modeling.

ETHICS STATEMENT

  • The work uses publicly available datasets under their licenses, involves no human-subjects research or new data collection, and could in principle contribute to harmful biological sequence design.The statement also notes potential benefits for studying gene regulation and variant effects.Source: ETHICS STATEMENT · p. 10

AI USE STATEMENT

The authors describe AI assistance for implementation, stabilization, and presentation while retaining responsibility for the paper’s core idea and final content. The appendices specify the model architectures, chunkwise SSM kernels, and language-model training setup.

AI USE STATEMENT

  • The authors developed IA’s core idea and derivations without AI assistance, verified AI-assisted outputs against implementations and experiment records, and take responsibility for the final paper.AI tools assisted with stabilization refinements, ablations, implementation, training settings, and editing and presentation tasks.Source: AI USE STATEMENT · p. 10

A.1 ARCHITECTURE DETAILS

  • Language models use a Llama-style pre-norm decoder, while genomic models replace Caduceus Mamba layers with IA; Caduceus-IA-PS and -Ph use distinct reverse-complement handling.Language models also use RMSNorm, SwiGLU, RoPE on queries and keys, untied embeddings, and no dropout.Source: §A.1 · p. 20

A.2 SSM KERNELS

  • Chunkwise S4D kernels parallelize computation within chunks and pass recurrent states across chunk boundaries; the default chunk length is 64.A second parallel pass recomputes each chunk from its corrected initial state.Source: §A.2 · p. 20

A.3 LANGUAGE MODELING TRAINING

  • Softmax, S4D-only, and IA share dimensions at each language-model scale, including context length 4096, state dimension M=64, and feature dimension R=dh.The models use the same backbone dimensions, while IA and S4D-only have less than 1% more parameters than Softmax, largely due to the full complex readout matrix.Source: §A.3 · p. 20
  • Language-model training uses AdamW with weight decay 0.1, gradient clipping 1.0, linear warmup, and cosine decay; runs use bfloat16 mixed precision on 4 to 32 GH200 GPUs.The peak learning rate varies by model scale, and each update uses 524,288 tokens in packed sequences.Source: §A.3 · pp. 20–21

B.1 ABLATION SETTINGS AND PARAMETER COUNTS

The appendices specify ablation configurations and baselines, then report that IA improves language-model validation perplexity over S4D-only and Softmax across the tested scales. They also describe downstream and long-context evaluations.

B.1 ABLATION SETTINGS AND PARAMETER COUNTS

  • Ablations use four language-model and genomic configurations, shared validation batches, and specified short-convolution settings to compare incremental IA components.The table describes each row as adding a component to the preceding row; within each domain, configurations share recurrent-state size and identically initialized shared parameters.Source: §B.1 · p. 22, Table 6 · p. 22

B.2 NATIVE CADUCEUS BASELINE SETTINGS

  • The native Caduceus-PS baseline has four layers, width 256, state dimension 16, and 1,934,592 parameters, and is trained with the stated HG38 protocol.Its original initialization and BF16 residual accumulation differ from Caduceus-IA’s FP32 residual accumulation.Source: §B.2 · p. 22

B.3 CADUCEUS-IA MECHANISM ABLATION ON BASE IA

  • The Caduceus-IA mechanism ablation compares selective and non-selective variants on Base IA using final validation NLL and parameter counts.Base IA omits the residual query MLP, value convolution, value skip connection, and output gate.Source: §B.3 · p. 23, Table 7 · p. 23

C LANGUAGE-MODELING SCALING AND DOWNSTREAM DETAILS

  • 13–16% lower validation perplexity than S4D-only at every scale; IA is about 1% lower than Softmax at 125 M and 7.4% lower at 1.3 B.The scaling table reports best-of-run validation perplexity on FineWeb-Edu, with parameter counts closely matched across models at each nominal scale.Source: C LANGUAGE-MODELING SCALING AND… · p. 24, Table 8 · p. 24
  • At 1.3 B, IA raises the commonsense average by 3.03 points over Softmax, while lowering BPB by 0.0105 on WikiText-2 and 0.0764 on LAMBADA.Softmax perplexity rises to 22.85 at 8 K and 62.06 at 14 K, while IA stays between 15.56 and 15.77; IA is also more accurate than S4D-only on RULER and all six Based tasks.Source: C LANGUAGE-MODELING SCALING AND… · pp. 24–25

D.1 MECHANISM AND ROPE ABLATIONS AT 125 M

The appendices test IA’s input and readout mechanisms, characterize recurrent-state size, and report recall, genomic, and inference findings. They also caution that two prefill-memory measurements are not directly comparable.

D.1 MECHANISM AND ROPE ABLATIONS AT 125 M

  • 14.3% lower validation perplexity than S4D-only at the same recurrent-state size; removing the query-conditioned readout makes perplexity worse than S4D-only.Replacing the key feature/value input with linear projections leaves perplexity slightly above Softmax, while RoPE has little effect on IA and worsens S4D-only.Source: §D.1 · p. 26, Table 10 · p. 26

D.2 STATE SIZE AT 1.3 B

  • IA and S4D-only have equal recurrent-state size because their inputs have the same width; Softmax’s KV cache instead grows linearly with sequence length.Each IA head has 2(dfeat + dh)M real degrees of freedom, with dfeat = dh.Source: §D.2 · pp. 26–27

D.7 REFERENCE-ANSWER LIKELIHOOD ON RECALL-HEAVY TASKS

  • IA has higher mean log-likelihood than S4D-only in all 14 recall settings; on Phonebook, both have zero accuracy, but IA leads by 0.51 to 2.90 nats per answer.The evaluation scores reference answers under teacher forcing and reports mean log-likelihood per example.Source: Table 15 · p. 28, §D.7 · p. 29

E.2 GENOMIC DOWNSTREAM RESULTS

  • Caduceus-IA-PS matches Caduceus-PS below 30 kb, is slightly better at 30–100 kb, and is slightly worse beyond 100 kb; Caduceus-IA-Ph has lower mean AUROC in all three groups.The task predicts causal-eQTL effects from frozen embeddings and groups variants by distance to the nearest transcription start site.Source: §E.2 · p. 30, Table 19 · p. 32

F.2 PREFILL MEMORY

  • Chunked IA’s prefill-memory column is not directly comparable with Softmax or non-chunked IA because it builds only the decoding state and computes no logits.The other two configurations run a full forward pass that includes vocabulary logits.Source: §F.2 · p. 34, Table 21 · p. 34

F.3 DECODE LATENCY ACROSS MODEL SIZES

  • At every model size and at B=1 and B=8, IA’s CUDA-graph decode latency varies by less than 1% across prefix lengths.The reported latency table covers completed prefix lengths from 512 to 16,384 tokens.Source: §F.3 · p. 34, Table 22 · p. 34

G TRAINING THROUGHPUT ON NVIDIA GH200

Under the stated four-GH200 training workloads, an IA optimizer update takes longer than either Softmax or S4D-only, and Caduceus-IA-PS takes longer than Caduceus-PS.

G TRAINING THROUGHPUT ON NVIDIA GH200

  • 2.82× as long as Softmax and 1.92× as long as S4D-only for an IA update; Caduceus-IA-PS takes 3.25× as long as Caduceus-PS.Timing covers forward and backward passes, gradient accumulation and communication, clipping, and the optimizer step; it excludes data loading and validation.Source: G TRAINING THROUGHPUT ON NVIDIA GH200 · p. 35, Table 23 · p. 35
Loading 2605.24330v2…