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
cs.CL
2026-10-18
DFM writes predicted chunk memory into a frozen LLM and skips retrieval at inference. On Mistral-7B the six-QA mean rises from 24.14 to 27.21, 12.71% over the best baseline.
Continued pretraining and LoRA fold an extra corpus into parameters and keep inference self-contained, but the new data crowds out older skills. RAG and RETRO read retrieved evidence while generating, so the index can change, and every request pays for search plus the extra context. Parametric memory moves the search offline. MLP Memory and Memory Decoder train an add-on to predict a retrieval-derived next-token distribution, with a separate retrieval at every token. A 4B-token Wikipedia then takes 4 billion queries, and the fusion stops at the output probabilities, so intermediate hidden states never see the memory.
Deep Fusion Memory (DFM) shares one supervision target across a chunk, 32 tokens in the main setup, and writes the predicted memory into intermediate decoder states. The backbone stays frozen. Inference does not search.
Training sequences are split into non-overlapping chunks. From the second chunk on, frozen Qwen3-Embedding-0.6B encodes the previous chunk and retrieves the top 16 neighboring chunks. Their fixed embeddings, stacked in rank order, are the target memory. The first chunk has no predecessor, so it receives no memory loss and a zero memory update.
The predictor is a causal Transformer with K linear heads, one head per memory slot. Its input is the mean-pooled token embedding of the previous chunk. Causal masking lets the memory for chunk j see only earlier chunks, the same context next-token prediction is allowed to use. A shared projector runs LayerNorm and two linear layers with GELU in between, mapping the K vectors into the backbone hidden size. Selected decoder layers share that projection.
Fusion is inserted at layers 9, 16, 24, and 31. The token state is the query; the K slots are keys and values. Each head has a sigmoid gate, a scalar in (0, 1) computed from the current state and multiplied onto that head. The gate bias starts at -2. Because σ(-2) ≈ 0.119, the memory branch opens small and is less likely to shove the frozen backbone off its pretrained behavior. The branch runs beside self-attention, both add into the residual, and the original FFN follows.
The loss is next-token cross-entropy plus memory MSE, with β = 1. The backbone, the retrieval encoder, and the stored targets stay fixed. At inference the predictor reads input embeddings and does not search. For QA, memory is computed once from the prompt, held fixed during the answer, and read only at prompt positions.
QA training is one epoch on English Wikipedia from December 2018, about 4B tokens. The tasks are NQ, WebQA, HotpotQA, 2Wiki, MuSiQue, and Bamboogle. The reported average is the unweighted mean of six. Both QA backbones add 0.75B parameters and train with AdamW at 1e-5. LoRA's added budget matches DFM, and CPT, LoRA, and DFM each train for one epoch. RAG prepends the top-1 passage from the same embedding model. xRAG compresses one document vector into a single prefix token.
| Backbone | Strongest baseline | Baseline avg | DFM avg | Relative gain |
| Mistral-7B-v0.3 | xRAG | 24.14 | 27.21 | 12.71% |
| Qwen3-8B-Base | MLP Memory | 23.80 | 26.46 | 11.19% |
On Mistral every task rises, and the mean moves from 13.63 to 27.21. NQ moves from 20.30 to 33.57, 6.18 above xRAG at 27.39. WebQA moves from 29.28 to 42.47, 2.26 above LoRA at 40.21. HotpotQA reaches 29.53, only 0.44 above RAG at 29.09. Bamboogle moves from 3.91 to 22.65. 2Wiki moves from 5.87 to 25.00, and RAG still scores 31.02. CPT averages 13.48, LoRA 15.47, MLP Memory 22.09.
On Qwen3-8B the mean moves from 22.71 to 26.46, up 2.66 absolute. WebQA falls from 41.78 to 39.47, while MLP Memory scores 43.26 there. NQ rises from 17.60 to 27.46, above RAG at 24.42. 2Wiki rises from 24.83 to 32.38. MuSiQue and Bamboogle are also best under DFM. HotpotQA lands at 26.46, below xRAG at 27.74.
Code, math, and law each get a separate model. Across HumanEval, HumanEval+, GSM8K, and LawBench, Qwen3-4B-Base moves from 58.18 to 64.69 (+6.51) and Qwen3-8B-Base from 63.14 to 67.52 (+4.38). On the 4B model, HumanEval moves from 62.04 to 69.51, HumanEval+ from 57.32 to 62.80, and GSM8K from 83.40 to 86.05, all three above CPT and LoRA. LawBench on 4B moves from 29.96 to 40.41, still under CPT at 41.66. On 8B it moves from 38.45 to 46.61, above CPT at 45.09 and LoRA at 45.73. LoRA still leads 8B HumanEval+ and GSM8K, 65.85 and 87.19 against DFM's 64.02 and 86.43. RAG drops LawBench to 24.35 on 4B and 26.89 on 8B.
Against token-level MLP Memory on the same 4B-token corpus, queries fall from 4 billion to 125 million, 32× fewer. MLP Memory keeps 1024 neighbors per query and DFM keeps 16, so retrieved items fall from 4.096T to 2 billion, 2048× fewer. On 16 A800 GPUs, training takes about 24 hours versus about 64, 62.5% less time. Added parameters are 0.75B versus 1.5B in that MLP Memory setup, and 0.6B on Qwen3-4B.
Ablations on Mistral are single runs. Chunk lengths 16, 32, and 64 score 26.51, 27.21, and 27.07. Doubling the chunk, which halves the queries, costs 0.14. Slot counts 2, 4, 8, and 16 score 26.88, 27.04, 27.09, and 27.21. Fusing only the last layer scores 21.11, and the last four layers score 22.07. Two to eight fusion layers starting at layer 9 land between 27.09 and 27.54. Four layers at positions 2, 10, 21, and 31 score 26.32, 0.89 under the default. No gate scores 26.83. Initializing the gate bias at 0 scores 27.90, above the default of -2.
Inference does not touch a vector store. The forward pass is ordinary decoding. Updating the knowledge still means rebuilding offline targets and retraining the memory module.
On Qwen3, CPT averages 15.13 and LoRA 14.66, both under the base mean of 22.71, while DFM reaches 26.46. With the backbone frozen, continued training does not rewrite the original weights.
Explicit retrieval still scores higher on some multi-hop sets, and LoRA or CPT still leads some code and math metrics. The gains are a few points. The fitting setup is a frozen backbone, no online retrieval, and memory written into intermediate layers.
WebQA on Qwen3 goes down. On Mistral 2Wiki, RAG scores 31.02 and DFM 25.00, a gap of 6.02. A continuous vector is not a substitute for source text placed in the context.
The 2048× figure multiplies two choices. Chunks of 32 tokens account for the 32× drop in queries. Cutting neighbors from 1024 to 16 accounts for another 64×, and that K is a hyperparameter. Top-8 averages 27.09, only 0.12 under top-16, and it is higher on five tasks. Most of top-16's extra points come from Bamboogle. A top-1024 denominator inflates the item count.
RETRO sits on top of the base model: 13.73 versus 13.63 on Mistral, 22.69 versus 22.71 on Qwen3. The paper does not explain that flat result. In the gate table, Bamboogle is 14.94 at bias -6, 22.65 at the default -2, 19.35 with no gate, and 19.69 at bias 0. The six-task mean moves with this small set, and every ablation is one run. Bias 0 averages 27.90, above the chosen default of -2, which is kept so the memory branch disturbs the backbone less early in training.
For QA, memory is computed once from the prompt, then left unchanged, and only prompt positions read it. Recall while generating is untested. The corpus stops at December 2018 Wikipedia, and the age gap to current QA sets is not discussed. Domain runs leave out MLP Memory, scale stops at 8B, and the code is listed as not yet released.