Hybrid Associative Memories

Leon Lufkin; Tomás Figliolia; Beren Millidge; Kamesh Krishnamurthy

cs.LGMar 20, 2026

Hybrid Associative Memories

Authors: Leon Lufkin, Tomás Figliolia, Beren Millidge, Kamesh Krishnamurthy

Organizations: 1Zyphra

Abstract

Recurrent neural networks (RNNs) and self-attention are both widely used sequence-mixing layers that maintain an internal memory. However, this memory is constructed using two orthogonal mechanisms: RNNs compress the entire past into a fixed-size state, whereas self-attention's state stores every past time step growing its state (the KV cache) linearly with the sequence length. This results in orthogonal strengths and weaknesses. Self-attention layers excel at retrieving information in the context but have large memory and computational costs, while RNNs are more efficient but degrade over longer contexts and underperform for precise recall tasks. Prior work combining these mechanisms has focused primarily on naively interleaving them to reduce computational cost without regard to their complementary mechanisms. We propose the Hybrid Associative Memory (HAM) layer, which combines self-attention and RNNs while leveraging their individual strengths: the RNN compresses the entire sequence, while attention supplements it only with information that is difficult for the RNN to predict, which is hence the most valuable information to explicitly store. HAM layers enable data-dependent growth of the KV cache, which can be precisely controlled by the user with a single, continuous threshold. We find that this fine-grained control of the KV cache growth rate has a smooth trade-off with loss and performance. Empirically, we show that our hybrid architecture offers strong, competitive performance relative to RNNs and Transformers even at substantially lower KV-cache usage.

Figures & tables

Panel of Figure 1 — Figure 1: Hybrid Associative Memory that maintains both an RNN state $S_{t}$ and KV cache which stores surprising tokens. The two memories work together in a complementary fashion.

Table 1: Model comparison at 800M scale. All HAM models use 50% KV cache (

\rho_{KV}=0.5

Table 2:

T_{\mathrm{KV}}/T

sweep results at 800M scale. ∗ GDN baseline (no KV).

Figure 2: Language modeling loss(top) and long-context accuracy on RULER (bottom) as a function of the KV-cache usage in HAM. The loss and the accuracy show a smooth relation with the amount of KV-cache HAM uses – a user-controlled parameter. Notably, the HAM with the learned router with

50\%

KV-cache usage shows strong performance significantly outperforming the GDN hybrid with full self-attention layers and identical KV-cache usage

Figure 3: Trajectory of the prediction error for the toy NIAH single 1 sequence. Obtained from the 10th layer of an 800M HAM model using a cosine-similiarity routing score trained on 50B tokens from Long Data Collections.

Figure 4: Routing scores $e_{t}$ averaged over 1,000 sequences from Institutional Books for selected layers of a HAM model with 0.5 KV usage. The average scores show a consistent downward trend along the sequence indicating that HAM is picking up long-range structure in the sequences and using its context to make better predictions towards the end of a sequence. The massive deviations around the mean trend indicate the large degree of sequence-to-sequence variability in the routing scores.

Figure 5: Routing scores and gate ( $\alpha_{t}$ ) for examples sequences and layers. The vertical lines indicates time-points where $\alpha_{t}<0.05$ i.e the state $S_{t}$ is effectively reset. In addition to the heterogeneity of the routing scores, the plot illustrates how frequently GDN resets the RNN state along the sequence.

Figure 6: KV cache growth: the top panels show the mean (thick line) and some example sequence plots for the KV-cache usage for two layers on

1000

sequences from the Institutional Books dataset. The KV-cache usage for a globally set KV-cache target can be quite heterogenous across layers (bottom left), and HAM’s method for tuning the KV-cache threshold quickly achieves the desired target after a period of freezing the threshold (bottom right, orange line); however, the different layers can converge to very different KV-cache usage (thin blue lines)

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure C.1 : Trajectory of the routing for the toy NIAH single 1 sequence used in Section 6.1 . Again, we observe a large spike in the routing score when the needle is encountered, . Obtained from the 10th layer of an 800M HAM model using an end-to-end learnable router trained on 50B tokens from Long Data Collections.

Figure D.1 : KV cache growth for the learnable router

	HAM / GDN	Transformer
$d_{\mathrm{QK}}$	$\tfrac{5}{7}d$	$d$
$d_{\mathrm{V}}$	$\tfrac{15}{14}d$	$d$
$d_{\mathrm{RNN\text{-}QK}}^{\mathrm{head}}$	$256$	—
$d_{\mathrm{RNN\text{-}V}}^{\mathrm{head}}$	$384$	—
$d_{\mathrm{KV\text{-}QK}}^{\mathrm{head}}$	$128$	$128$
$d_{\mathrm{KV\text{-}V}}^{\mathrm{head}}$	$192$	$128$

Appendix

Table E.1 : Architectural relations used for simplifying parameter and FLOPs expressions.

Model	Training FLOPs (zFLOPs)	Relative
HAM ( $T_{\mathrm{KV}}^{*}/T=0.5$ )	0.3511	ref.
GDN	0.2467	$-$ 29.7%
Transformer	0.4592	+30.8%
GDN-GSA ( $k=2$ )	0.3429	$-$ 2.3%

Appendix

Table E.2 : Empirical training FLOPs for each model.

Model	Parameters
HAM	$\frac{48075}{401408}d^{3}+\frac{72591}{200704}d^{2}+\frac{448061}{7}d$
GDN	$\frac{375}{3136}d^{3}+\frac{33}{112}d^{2}+\frac{7168703}{112}d$
Transformer	$\frac{23}{240}d^{3}+\frac{23}{960}d^{2}+64001d$
GDN-GSA	$\left(\frac{375k-27}{3136k}\right)d^{3}+\left(\frac{33k-30}{112k}\right)d^{2}+\left(\frac{7168703k-591}{112k}\right)d$

Appendix

Table E.3 : Total model parameters as a function of

d

Model	Forward FLOPs
HAM	$T\left[\frac{375}{1568}d^{3}+\frac{99417}{3136}d^{2}+\frac{28713903}{224}d\right]+TT_{\mathrm{KV}}\left[\frac{75}{3136}d^{2}+\frac{15}{56}d\right]$
GDN	$T\left[\frac{12015}{50176}d^{3}+\frac{4689327}{401408}d^{2}+128004d\right]$
Transformer	$T\left[\frac{23}{120}d^{3}+\left(\frac{23}{90}+\frac{989}{40960}T\right)d^{2}+128004d\right]$
GDN-GSA	$T\left[\left(\frac{12015k-879}{50176k}\right)d^{3}+\left(\frac{10836T+4689327k-4572591}{401408k}\right)d^{2}+128004d\right]$

Appendix

Table E.4 : Total forward FLOPs as a function of

d

Model	Forward Memory (bytes)
HAM	$\frac{48075}{200704}d^{3}+\frac{809871}{100352}d^{2}+\frac{75}{1568}T_{\mathrm{KV}}d^{2}+\frac{896122}{7}d$
GDN	$\frac{375}{1568}d^{3}+\frac{3111}{392}d^{2}+\frac{7168703}{56}d$
Transformer	$\frac{23}{120}d^{3}+\frac{23}{480}(1+T)d^{2}+128002d$
GDN-GSA	$\left(\frac{375k-27}{1568k}\right)d^{3}+\left(\frac{12444k-12360+75T}{1568k}\right)d^{2}+\left(\frac{7168703k-591}{56k}\right)d$

Appendix

Table E.5 : Total forward memory as a function of

d

(bytes, BF16).

Model	Forward FLOPs (per token)	Forward Memory (bytes)
HAM	$2P+(130+0.1T_{KV})P^{\frac{2}{3}}-8.5\text{K}T_{KV}$	$2P+T_{KV}(0.208P^{\frac{2}{3}}-13P^{\frac{1}{3}}-11.4\text{K})$
GDN	$2P+46P^{\frac{2}{3}}$	$2P+30P^{\frac{2}{3}}$
Transformer	$2P+T(0.115P^{\frac{2}{3}}-11\text{K})$	$2P+T(0.23P^{\frac{2}{3}}-21\text{K})$
GDN-GSA ( $k=2$ )	$2P+T(0.06P^{\frac{2}{3}}-4P^{\frac{1}{3}}-3.3\text{K})$	$2P+T(0.107P^{\frac{2}{3}}-7P^{\frac{1}{3}}-5.9\text{K})$

Appendix

Table E.6 : Approximate scaling for longer contexts. All expressions are accurate to within

{\sim}2\%

of the full expressions in Table E.6 .

Figure E.1 : Forward FLOPs for GDN, Transformer, and GDN-GSA relative to HAM, for

T_{\mathrm{KV}}/T=0.10

(left),

T_{\mathrm{KV}}/T=0.25

(center), and

T_{\mathrm{KV}}/T=0.5

(right).

Figure E.2 : Forward memory for GDN, Transformer, and GDN-GSA relative to HAM, for

T_{\mathrm{KV}}/T=0.10

(left),

T_{\mathrm{KV}}/T=0.25

(center), and

T_{\mathrm{KV}}/T=0.5

(right).

Pre-norm
Source of weights	# parameters
Pre-norm	$d_{\mathrm{hidden}}$
Subtotal	$d_{\mathrm{hidden}}$
Q, K, V projections (shared)
$W_{\mathrm{Q}},W_{\mathrm{K}}$	$2\cdot d_{\mathrm{hidden}}\cdot d_{\mathrm{QK}}$
$W_{\mathrm{V}}$	$d_{\mathrm{hidden}}\cdot d_{\mathrm{V}}$

Appendix

Table E.7 : Per-layer parameter count for HAM.

Pre-norm
Source of weights	# parameters
Pre-norm	$d_{\text{hidden}}$
Subtotal	$d_{\text{hidden}}$
Q, K, V projections
$W_{\text{Q}}$	$d_{\text{hidden}}\cdot d_{\text{hidden}}$
$W_{\text{K}}$	$d_{\text{hidden}}\cdot d_{\mathrm{KV}}$

Appendix

Table E.8 : Per-layer parameter count for Transformer.

Pre-norm
Source of weights	# parameters
Pre-norm	$d_{\text{hidden}}$
Subtotal	$d_{\text{hidden}}$
Q, K, V projections
$W_{\text{Q}},W_{\text{K}}$	$2\cdot d_{\text{hidden}}\cdot d_{\mathrm{QK}}$
$W_{\mathrm{V}}$	$d_{\text{hidden}}\cdot d_{\mathrm{V}}$

Appendix

Table E.9 : Per-layer parameter count for Gated DeltaNet.

FFN projections
Source of weights	# parameters
Pre-norm	$d_{\text{hidden}}$
Gate projection ( $W_{\text{gate}}$ )	$d_{\text{hidden}}\cdot d_{\text{int}}$
Up projection ( $W_{\text{up}}$ )	$d_{\text{hidden}}\cdot d_{\text{int}}$
Down projection ( $W_{\text{down}}$ )	$d_{\text{int}}\cdot d_{\text{hidden}}$

Appendix

Table E.10 : Per-layer parameter count for FFN (SwiGLU).

RNN scalars
Operation	FLOPs
$\beta$ (from $b$ proj)	$3Td_{\mathrm{hidden}}h_{\mathrm{RNN}}$
$g$ (from $a$ proj & $A_{log}$ )	$5Td_{\mathrm{hidden}}h_{\mathrm{RNN}}$
Subtotal	$8Td_{\mathrm{hidden}}h_{\mathrm{RNN}}$
Chunked delta rule
$g$ cumsum	$h_{\mathrm{RNN}}T$

Appendix

Table E.11 : RNN block (GDN) FLOPs.

Rotary positional encoding
Operation	FLOPs
RoPE on Q, K	$4Td_{\mathrm{QK}}$
Subtotal	$4Td_{\mathrm{QK}}$
Attention
$QK^{\top}$ (causal)	$TT_{\mathrm{KV}}d_{\mathrm{QK}}$
Softmax (causal)	$2TT_{\mathrm{KV}}h_{\mathrm{KV}}$

Appendix

Table E.12 : KV block (attention) FLOPs.

Pre-norm
Operation	FLOPs
Pre-norm	$4Td_{\mathrm{hidden}}$
Subtotal	$4Td_{\mathrm{hidden}}$
Q, K, V projections (shared)
$W_{\mathrm{Q}},W_{\mathrm{K}}$	$4Td_{\mathrm{hidden}}d_{\mathrm{QK}}$
$W_{\mathrm{V}}$	$2Td_{\mathrm{hidden}}d_{\mathrm{V}}$

Appendix

Table E.13 : Full HAM layer FLOPs. References

\mathcal{F}_{\mathrm{RNN}}

from Table E.11 and

\mathcal{F}_{\mathrm{KV}}

from Table E.12 .

Pre-norm
Operation	FLOPs
Pre-norm	$4Td_{\mathrm{hidden}}$
Subtotal	$4Td_{\mathrm{hidden}}$
Q, K, V projections
$W_{Q}$	$2Td_{\mathrm{hidden}}^{2}$
$W_{K}$	$2Td_{\mathrm{hidden}}d_{\mathrm{KV}}$

Appendix

Table E.14 : Transformer attention layer FLOPs (excluding FFN).

Pre-norm
Operation	FLOPs
Pre-norm	$4Td_{\mathrm{hidden}}$
Subtotal	$4Td_{\mathrm{hidden}}$
Q, K, V projections
$W_{\text{Q}},W_{\text{K}}$	$4Td_{\text{hidden}}d_{\mathrm{QK}}$
$W_{\mathrm{V}}$	$2Td_{\text{hidden}}d_{\mathrm{V}}$

Appendix

Table E.15 : Pure Gated DeltaNet layer FLOPs.

Operation	FLOPs
FFN (SwiGLU)
Gate projection ( $W_{\text{gate}}$ )	$2Td_{\text{hidden}}d_{\text{int}}$
Up projection ( $W_{\text{up}}$ )	$2Td_{\text{hidden}}d_{\text{int}}$
SiLU activation	$3Td_{\text{int}}$
Gate $\cdot$ Up	$Td_{\text{int}}$
Down projection ( $W_{\text{down}}$ )	$2Td_{\text{int}}d_{\text{hidden}}$

Appendix

Table E.16 : FFN (SwiGLU) FLOPs.

Associative Memory

Associative memory studies how systems store patterns or associations and retrieve them from partial, noisy, or context-dependent cues.

See all 120 papers

Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs
Oct 5, 2026Hyunji Lee, Joykirat Singh, Zaid Khan +5Long-Context Language Modeling Recurrent Transformers
UNC Chapel Hill · The University of Texas at Austin · Mila +1
Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing
Jul 8, 2026Tommaso Cerruti, Tim Rieder, George Rowlands +2Linear Attention Delta-Rule Linear Attention
1ETH Zurich, D-INFK · 2ETH AI Center, ETH Zurich
Variational Linear Attention: Stable Associative Memory for Long-Context Transformers
May 11, 2026Vishal Pandey, Gopal SinghEfficient Attention Linear Attention
Independent Researcher London, UK · Metriqual Athens, GR

Oct 5, 2026J/K move · Enter open · S save

Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs

Hyunji Lee, Joykirat Singh, Zaid Khan, Justin Chih-Yao Chen+4

UNC Chapel Hill · The University of Texas at Austin · Mila+1

Long-Context Language Modeling Recurrent Transformers Hybrid Attention

Open paper AskPDF arXiv

Recurrent-attention hybrid language models (LMs), which interleave attention and recurrent layers, are increasingly used to combine the efficiency of the recurrent layers with the strong performance of attention layers. Prior work suggests that attention and recurrent layers offer complementary pathways to use past information: attention supports precise memory recall from earlier tokens, while recurrent layers support consolidation of disparate information over long contexts. However, we observe that simply having access to both pathways does not mean that hybrid LMs are effectively using them. We find that they rely substantially more on attention than on the recurrent state. Standard supervised fine-tuning improves overall performance but does not improve how the two memory pathways are coordinated: the model becomes more reliant on information propagated by attention layers, while its use of information propagated by recurrent layers remains limited. To encourage better coordination between the two memory pathways, we add an auxiliary loss that limits attention's access to earlier context while the recurrent state propagates through the full sequence. This objective encourages the model to retain and use information through the recurrent pathway alongside attention. It improves overall performance, with particularly strong gains on tasks involving longer contexts or requiring information aggregation, consistent with the strengths of recurrent layers observed in analysis. Crucially, this imbalance and the benefit of our auxiliary loss generalize: they apply to multiple recurrent-attention LMs in question-answering and agentic tasks, as well as to attention-based LMs that combine different forms of memory. Together, our findings show that simply providing multiple memory pathways does not ensure their effective use, and that targeted supervision is needed to better coordinate them.

Abstract

Associative Memory

Explore similar work

Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs

Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing

Variational Linear Attention: Stable Associative Memory for Long-Context Transformers