Deep Fusion Memory: A Chunk-Level Parametric Memory for Large Language Models
Rubin Wei, Siyuan Sun, Hao Sun, Yuliang Liu, Xiang Shi, Jiarui Wang, Junming Zhang, Zhouhan Lin
Deep Fusion Memory (DFM) integrates external knowledge into LLM hidden states via chunk-level parametric memory, eliminating online retrieval.
How can we integrate external knowledge into LLMs efficiently without the high inference latency of retrieval or the catastrophic forgetting of continued pretraining?
Large language models struggle to incorporate external knowledge without either suffering from catastrophic forgetting during retraining or incurring high latency from online retrieval. Deep Fusion Memory (DFM) solves this by training a causal predictor to generate memory vectors from preceding context, which are then injected into intermediate decoder layers via gated cross-attention. This architecture achieves significant gains on knowledge-intensive QA and domain adaptation tasks while requiring 32× fewer retrieval queries and 2048× fewer retrieved items for training data construction than token-level parametric baselines.
Paper Primer
Existing parametric memory methods often rely on token-level targets, which are computationally expensive to construct and provide only shallow fusion by interpolating output probabilities. DFM shifts to chunk-level supervision: it predicts continuous memory vectors for 32-token chunks and fuses them directly into the decoder's hidden states using gated cross-attention. This allows the model to process external knowledge alongside the evolving context in intermediate layers, rather than treating memory as a final-stage adjustment.
DFM significantly outperforms existing parametric and non-parametric baselines on knowledge-intensive QA.
Across six QA tasks, DFM improved the mean score by 12.71% on Mistral-7B-v0.3 and 11.19% on Qwen3-8B-Base compared to the strongest baseline. Absolute gains of 3.07 and 2.66 points on the six-task mean for Mistral and Qwen3, respectively.
Why is "deep fusion" into intermediate hidden states superior to the shallow interpolation used in earlier parametric memory models?
Integrating memory into intermediate layers allows subsequent decoder layers to further process and combine the retrieved knowledge with the evolving context, rather than simply adjusting the final output probability distribution.
Does DFM require an external datastore or retrieval engine during inference?
No. While the model is trained using offline retrieval to generate target memory, the inference process relies solely on the learned memory predictor to supply memory vectors directly from the input sequence.
DFM demonstrates that chunk-level parametric memory is a viable, efficient alternative to online retrieval, enabling models to acquire and use external knowledge without the latency of explicit search or the instability of continued pretraining.
Motivation and Problem Framing
LLMs need external knowledge, but existing solutions trade forgetting for latency; DFM tackles this gap.
Large language models excel at language tasks but still miss facts that are not encoded in their parameters. Continuing pretraining can refresh knowledge but often erases previously learned capabilities—a phenomenon known as catastrophic forgetting. On the other hand, retrieval‑augmented generation (RAG) fetches external documents at inference time, which restores factuality but introduces noticeable latency.
When a model is further trained on new data, it can overwrite weights that previously encoded other knowledge, causing the model to lose performance on earlier tasks.
RAG augments a language model’s generation by first retrieving relevant passages from an external corpus and then conditioning the decoder on those passages.
Deep Fusion Memory (DFM) addresses the trade‑off by learning a chunk‑level memory predictor that shares supervision across tokens, and by fusing the predicted memory into hidden states via gated cross‑attention. This design eliminates the need for online retrieval while allowing later decoder layers to refine the injected knowledge.
The core trade‑off is between static pretraining that risks forgetting and dynamic retrieval that incurs latency; DFM shifts the balance toward static, learned memory without sacrificing factual coverage.
The Deep Fusion Memory Architecture
Core mechanism that predicts chunk‑level memory and fuses it via gated cross‑attention.
The naive approach of attaching a full‑size external datastore to every decoder step incurs prohibitive latency; the trick is to predict a compact, chunk‑level memory from the already‑available context and fuse it directly inside the model.
Predict a small set of memory vectors for each chunk of the prompt, then let selected decoder layers attend to them alongside self‑attention, with a per‑head gate deciding how much memory to use.
How does Deep Fusion Memory differ from a standard retrieval‑augmented generation pipeline?
Standard RAG performs a live lookup at inference time, pulling external documents for every token. Deep Fusion Memory replaces that lookup with a predictor that generates a fixed set of $K$ memory vectors from the prompt itself, eliminating the external datastore and latency while still providing chunk‑level knowledge.
The predictor compresses the embeddings of the previous chunk into $K$ vectors, one per memory slot, using a causal Transformer that can only see earlier chunks.
Linear projection $W_p$ (identity for simplicity) maps $C_1$ to the predictor input.
The causal Transformer processes $C_1$ and outputs two head vectors: $v_{1,1} = (0.6,0.2,0.1)$ and $v_{1,2} = (0.1,0.7,0.3)$.
These vectors form the predicted memory $V_1 = \begin{bmatrix}v_{1,1}\\ v_{1,2}\end{bmatrix}$ for the next chunk $X_1$.
The shared projector $P_\psi$ (a layer‑norm followed by a linear map) transforms $V_1$ into decoder hidden space $M_1$.
The predictor creates a fixed‑size representation of the preceding chunk without any external retrieval, enabling fast inference while still providing contextual knowledge.
Why does the predictor use $K$ separate linear heads instead of a single output vector?
Each head corresponds to one slot in the retrieval‑rank order, preserving the ordering of the offline target memory. This alignment lets the cross‑attention attend to the same ranked neighbors the predictor was trained to mimic.
For each token, the model queries the $K$ predicted memory slots, computes a softmax over them, and then mixes the results with a per‑head gate that decides how much of the memory signal to inject.
How is Memory Cross‑Attention different from standard multi‑head attention?
Standard attention queries the same token sequence for keys and values. Memory Cross‑Attention queries a fixed set of $K$ predicted memory vectors that come from the preceding chunk, and a learned gate per head modulates each head’s contribution, allowing the model to blend external knowledge with self‑attention selectively.
**Figure 1 |** Overview of Deep Fusion Memory. (a) Memory Prediction. The memory predictor maps pooled chunk embeddings to predicted memory, supervised by offline target memory. A shared projector maps predicted memory to the decoder’s hidden space. (b) LLM Backbone. Selected decoder layers incorporate memory for next-token prediction. (c) Memory Fusion. Memory predicted from the preceding context supports token predictions in the next chunk. Per-head gates control the memory cross-attention contribution alongside self-attention. Inference requires no retrieval.
Partition each training sequence into non‑overlapping chunks of $C$ tokens.
For each chunk $X_j$ ( $j\ge1$ ), retrieve the top‑$K$ neighbor embeddings from a frozen datastore to form target memory $V_j$.
Mean‑pool token embeddings of the preceding chunk $X_{j-1}$ to obtain $C_j$.
Feed the sequence $C_{\le j}$ into the causal predictor to generate predicted memory $\hat{V}_j$.
Project $\hat{V}_j$ with $P_\psi$ and fuse it via gated cross‑attention in the selected decoder layers.
Compute the joint loss $L = L_{\text{LM}} + \beta L_{\text{MSE}}$ and back‑propagate through all added modules while keeping the backbone frozen.
Inference with Deep Fusion Memory – no external retrieval.
Empirical Evaluation
DFM delivers strong QA performance while cutting training cost.
We now present the quantitative results that demonstrate DFM’s advantage over prior baselines.
DFM cuts training time and boosts QA performance compared to MLP Memory.
Training finishes in $24$ h versus $64$ h, and the average QA score rises from $23.80$ to $26.46$ ($\Delta$ +3.07 points).
It learns token‑level k‑nearest‑neighbor distributions and interpolates them with the model’s output probabilities.
How does MLP Memory differ from DFM’s chunk‑level approach?
MLP Memory predicts memory for each individual token, which forces a retrieval for every token; DFM predicts memory for 32‑token chunks, reducing both the number of queries and the parameter budget.
Tasks that require factual information not encoded in the model’s parameters, such as open‑domain question answering.
Why can DFM improve Knowledge‑Intensive QA without online retrieval?
DFM trains a chunk‑level memory predictor on Wikipedia, so at inference time the decoder receives injected factual embeddings directly, eliminating the need for a separate retrieval step.
**Figure 2.** Resource comparison with token-level MLP Memory. (a) Added parameters. (b) Training cost: MLP Memory and DFM take approximately 64 and 24 hours, respectively, on 16 A800 GPUs, a 62.5% reduction. (c) Retrieval workload: retrieval queries and total retrieved items (queries $\times K$) for a 4B-token corpus. Reference bars in (b,c) equal 100%; arrows show percentage reductions.
DFM achieves superior performance with lower training costs than MLP Memory.
Ablations and Sensitivity Analysis
Key ablations confirm DFM’s robustness and the conclusion highlights its overall gains.
Recall that DFM injects chunk‑level memory into the decoder, trading off retrieval latency for a compact, learned knowledge store.
Chunk size 32 tokens yields the highest six‑task average (27.21), only 0.14 points below the 64‑token setting.
Average scores: 26.51 (size 16), 27.21 (size 32), 27.07 (size 64).
Four intermediate‑to‑late fusion layers achieve up to 27.54 average, a 6.43‑point gain over fusing only the final layer.
Scores range 27.09–27.54 for 2–8 layers; final‑layer‑only fusion scores 21.11.
Using only two memory slots (Top‑2) attains 26.88 average, just 0.33 points shy of the Top‑16 setting (27.21).
Average scores: Top‑2 = 26.88, Top‑4 = 27.04, Top‑8 = 27.09, Top‑16 = 27.21.
Removing the token‑wise gate drops the six‑task mean from 27.21 to 26.83, a 0.38‑point decline.
Mean with gate = 27.21; mean without gate = 26.83.
Initializing the gate bias at 0 yields the best mean (27.90), 0.69 points above the default bias $-2$.
Bias $0$ row in Table 3 shows Avg = 27.90; default bias $-2$ row shows Avg ≈ 27.21.
**Figure 3.** DFM ablations with predicted memory on Mistral-7B-v0.3. Points show single-run six-task QA means. Filled circles and dashed lines mark the main setting (chunk 32, layers [9, 16, 24, 31], Top-16). In (b), the diamond and triangle use four fusion layers at [28, 29, 30, 31] and [2, 10, 21, 31], respectively.
**Figure 4** | Task-level score changes relative to the main DFM configuration on Mistral-7B-v0.3. Blue and brown indicate increases and decreases, respectively. Avg. gives the change in the unweighted six-task mean. Fusion labels denote the number of fusion layers, while Last four and Start layer 2 use [28, 29, 30, 31] and [2, 10, 21, 31], respectively. All differences are computed before rounding.
Overall, DFM delivers a compact, learned memory that improves QA across diverse tasks without incurring the latency of external retrieval, confirming the paper’s central claim.
Extended Results and Case Studies
Limits and case studies assess DFM’s efficiency and qualitative retrieval strengths.
This limits section quantifies DFM’s training‑time and parameter efficiency and then illustrates how its retrieval‑augmented generation improves answer correctness on qualitative examples.
**Figure 5.** Six selected HotpotQA answer comparisons using Mistral-7B-v0.3. Columns show Base (left), MLP Memory (middle), and DFM (right). Check and cross symbols indicate correct and incorrect normalized exact-match answers.
**Table.** Examples of retrieval cases showing the preceding query chunk, retrieved chunk text, and ground-truth continuation.
Qualitative retrieval cases demonstrate DFM's ability to ground generation.
Experimental Details
Appendix details for QA experiments, data, training, baselines, and evaluation.
Data construction uses Wikipedia December 2018 section and infobox text. Articles are normalized, split 99.8/0.1/0.1 % (seed 42), deduplicated, EOS‑appended, and packed into $2{,}048$‑token sequences; incomplete tails are dropped. Tokenization yields $4.70\!B$, $4.29\!B$, and $4.08\!B$ training tokens for Mistral, Qwen3, and OLMo3 respectively, which are then chunked into non‑overlapping $32$‑token pieces.
Training freezes the backbone and jointly optimizes the chunk‑level memory predictor, shared projector, and gated fusion modules. The predictor comprises $28$ layers, hidden width $1{,}024$, FFN width $3{,}072$, $16$ query heads, and $8$ KV heads, outputting $16$ continuous memory slots clipped to $[-1,1]$. Memory cross‑attention is inserted at layers $[9,16,24,31]$ with gate weights initialized to $0$ and biases to $-2$.
Baselines include RAG (retrieves top‑1 Wikipedia document), xRAG (projects a retrieved embedding into a single prepended token), RETRO (chunk‑level retrieval with cross‑attention), CPT (full‑model fine‑tuning), LoRA (low‑rank adapters matching DFM’s parameter budget), and MLP Memory (token‑level k‑NN prediction interpolated with backbone logits). For MLP Memory we sweep interpolation weight $\lambda\in\{0,0.05,\dots,1\}$ and report the per‑split optimum.
Evaluation follows Table 4, covering six QA benchmarks with metrics $F1$ or substring match on a $0\!-\!100$ scale. DFM runs zero‑shot greedy decoding without online retrieval or chain‑of‑thought prompts; generation stops at a newline, token limit, or punctuation stop strings for WebQA. Queries are formed from the current visible prompt chunk, and memory predictions remain static throughout generation.
**Table 4.** Evaluation datasets and metrics for question answering. Max. tokens is the maximum number of newly generated tokens.