PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.25100v1cs.CL

TL;DR

Large language models need external knowledge, but current approaches can risk forgetting, add inference retrieval overhead, or require costly token-level memory supervision with shallow fusion. DFM predicts chunk-level memory and fuses it into intermediate decoder states with a frozen backbone; it reports improved results across QA tasks and domain adaptation.

  • Problem

    Knowledge injection must address missing, poorly retained, or outdated facts while preserving pretrained capabilities and keeping acquisition and inference costs manageable; existing parametric memory also faces costly token-level supervision and shallow output-level fusion.

    Source: §1 · pp. 1–2

  • Method

    DFM predicts memory from preceding chunk context using offline retrieved chunk embeddings as supervision, then jointly trains the predictor and gated fusion modules with the backbone frozen.

    Source: §1 · p. 2, §3.2 · pp. 3–4

  • Results

    Across six QA tasks, DFM beats the strongest baseline by 12.71% on Mistral-7B-v0.3 and 11.19% on Qwen3-8B-Base.

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

  • Takeaways & Limitations

    The results support shared supervision and learned internal fusion as a way to acquire external knowledge efficiently and use memory effectively.

    Source: §6 · p. 11

    Memory generation for next-token prediction uses only input available by the chunk boundary, with the preceding chunk complete before its pooled query is used.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

Large language models need effective ways to acquire and use external knowledge beyond what they retain from pretraining. Continued pretraining risks catastrophic forgetting, while retrieval at inference time can add substantial overhead. External parametric memory avoids online retrieval, but constructing its training data requires extensive offline search, and output-level interpolation provides only shallow fusion. We introduce Deep Fusion Memory (DFM), a parametric memory architecture that reduces this construction workload and incorporates learned memory into intermediate hidden states. A causal memory predictor learns to generate memory from the preceding context, using retrieved chunk embeddings from the corpus as supervision. The predicted memory is shared across tokens within each chunk and fused into their hidden states through gated cross-attention. Memory prediction and fusion are trained end to end with the backbone frozen, and inference requires no explicit retrieval. Compared with existing parametric memory approaches, which often construct token-level memory targets, DFM requires $32\times$ fewer retrieval queries and $2048\times$ fewer retrieved items to construct its training data. Across six QA tasks, DFM performs better over the strongest baseline by 12.71% on Mistral-7B-v0.3 and 11.19% on Qwen3-8B-Base. Beyond QA, domain adaptation yields average gains of 6.51 and 4.38 points over Qwen3-4B-Base and Qwen3-8B-Base, respectively.

1. Introduction

The introduction frames knowledge injection as a balance between preserving pretrained capabilities, incorporating corpus knowledge, and controlling acquisition and inference costs. It presents DFM as a chunk-level parametric memory with deep fusion and reports gains in training-data efficiency and downstream performance.

Motivation

  • Knowledge injection must address missing, poorly retained, or outdated facts while preserving existing capabilities and managing acquisition and inference costs.Continued pretraining and LoRA can risk catastrophic forgetting, whereas retrieval methods can add inference latency.Source: §1 · p. 1
  • Earlier parametric-memory approaches incur offline costs from token-level retrieval targets and use output-probability interpolation rather than decoder hidden-state fusion.DFM is introduced to reduce target-construction workload and incorporate memory into intermediate states.Source: §1 · p. 2

Approach

  • DFM shares chunk-level memory supervision across token predictions and injects learned corpus knowledge into intermediate hidden states while keeping the backbone frozen.The memory predictor and fusion module are trained jointly, and inference requires no explicit retrieval.Source: §1 · p. 2

Results

  • Across six QA tasks, DFM beats the strongest baseline by 12.71% on Mistral-7B-v0.3 and 11.19% on Qwen3-8B-Base.It also uses 32× fewer retrieval queries and 2048× fewer retrieved items to construct training data than token-level-target parametric memory approaches.Source: §1 · p. 2
  • Across code, mathematics, and law, DFM reports average gains of 6.51 and 4.38 points over Qwen3-4B-Base and Qwen3-8B-Base, respectively.The introduction also reports that fusion placement and gate initialization matter, while tested chunk sizes and memory-slot counts produce comparatively small changes in average QA performance.Source: §1 · p. 2

2. Related Work

Related work spans model adaptation, retrieval-augmented generation, and parametric memory. DFM is positioned as reducing token-level supervision costs while fusing predicted memory inside decoder hidden states without online search.

Model Adaptation

  • CPT and LoRA adapt models using additional corpora or parameter-efficient updates, but such adaptation can cause catastrophic forgetting.DFM instead keeps the backbone frozen and learns an added memory module from retrieval-derived targets.Source: §2 · pp. 2–3

Retrieval-Augmented Language Models

  • RAG, RETRO, kNN-LM, and xRAG expose corpus information through retrieval-based interfaces that retain an external retrieval dependency at inference.DFM predicts continuous memory parametrically and uses it inside the decoder without online search.Source: §2 · p. 3

Parametric Memory

  • MLP Memory, Memory Decoder, Memory Decoder at Scale, and MemSFT avoid explicit inference-time retrieval but use token-level targets and shallow output-probability fusion.DFM shares target memory across tokens within each chunk and fuses predictions into hidden states through gated cross-attention.Source: §2 · p. 3

3. Deep Fusion Memory

DFM predicts memory from preceding chunk context and reuses that memory for token predictions in the next chunk. It trains a predictor and decoder fusion modules jointly with a frozen backbone, then supplies predicted memory at inference without constructing target memory.

3. Deep Fusion Memory

  • DFM combines a causal chunk-level memory predictor with cross-attention modules inserted into selected decoder layers.The added modules are trained jointly while the backbone remains frozen; at inference, the predictor supplies memory from the input sequence.Source: §3 · p. 3

3.1. Training data construction

  • For each chunk after the first, the preceding chunk is used to retrieve top-K corpus neighbors whose fixed dense embeddings form the offline target memory.That target supervises a memory prediction shared across token predictions in the current chunk; the first chunk has no preceding-chunk query and is excluded from memory supervision.Source: §3.1 · p. 3

3.2. Training

  • The causal predictor mean-pools preceding-chunk token embeddings and uses their prefix to generate K memory vectors for the next chunk.A causal Transformer with K separate output heads produces slots in retrieval-rank order; the first input embedding is a masked placeholder.Source: §3.2 · p. 4
  • A shared projector maps predicted memory into decoder hidden space, where selected layers use per-head gated cross-attention alongside causal self-attention.Attention selects information from shared memory for each token, gates control each head’s contribution, and other decoder layers retain their original computation.Source: §3.2 · pp. 4–5
  • End-to-end training uses next-token cross-entropy and an MSE memory-prediction loss, with the backbone, retrieval encoder, and target memory fixed.The language model consumes predicted memory during training; gradients from language modeling pass through fusion to the predictor, while MSE directly supervises predicted slots.Source: §3.2 · p. 5

3.3. Inference

  • At inference, predicted memory comes from input embeddings without datastore search, retrieved text, or an external embedding encoder.For QA, memory is computed from the visible prompt and remains fixed during answer generation; selected layers attend to it at prompt-chunk positions.Source: §3.3 · p. 5

4. Experiments

The experiments evaluate DFM on six knowledge-intensive QA tasks and domain adaptation, comparing performance and resource costs with existing methods. DFM leads the six-task QA averages on both backbones, improves domain-adaptation averages over the base models, and reduces target-construction workload.

4.1. Experimental Setup

  • The QA evaluation uses six benchmarks on Wikipedia data, while domain adaptation tests Code, Math, and Law using separate domain datasets.QA uses approximately 4B tokens of English Wikipedia; domain evaluation covers HumanEval, HumanEval+, GSM8K, and LawBench.Source: §4.1 · p. 6

4.2. Question Answering Performance

  • 12.71% on Mistral-7B-v0.3 and 11.19% on Qwen3-8B-Base are DFM’s relative improvements over the strongest baseline in six-task QA averages.DFM averages 27.21 versus xRAG’s 24.14 on Mistral and 26.46 versus MLP Memory’s 23.80 on Qwen3.Source: §4.2 · p. 6, Table 1 · p. 7
  • DFM improves over Base on all four multi-hop QA benchmarks for both backbones, while explicit retrieval remains stronger on some benchmarks.On Mistral, DFM scores 25.00 on 2Wiki versus RAG’s 31.02; on Qwen3, DFM scores 26.46 on HotpotQA versus xRAG’s 27.74.Source: §4.2 · p. 7

4.3. Domain Adaptation

  • DFM’s domain-adaptation averages are 64.69 on Qwen3-4B-Base and 67.52 on Qwen3-8B-Base, gains of 6.51 and 4.38 points over Base.It improves all four domain benchmarks over Base on both backbones, though its advantage over adaptation methods varies by benchmark.Source: §4.3 · p. 8, Table 2 · p. 8

4.4. Efficiency Analysis

  • On 16 A800 GPUs, DFM training takes approximately 24 hours versus 64 hours for MLP Memory, a 62.5% reduction; DFM adds 0.75B parameters versus 1.5B.Both configurations use the same corpus and number of epochs.Source: §4.4 · pp. 8–9, Figure 2 · p. 9
  • For a 4B-token corpus, chunk-level target construction reduces retrieval queries by 96.88% and retrieved items by 99.95% versus token-level construction.This corresponds to 125M versus 4B queries and 2B versus 4.096T retrieved items.Source: §4.4 · p. 9, Figure 2 · p. 9

5. Ablations and Analysis

The ablations examine chunk size, fusion placement and depth, memory-slot count, and gating. The reported results favor a 32-token chunk setting and intermediate-to-late fusion, with similar average QA performance possible using fewer slots.

5.1. Ablation Studies

  • 27.21 is the six-task QA mean with 32-token chunks, the best of the tested sizes; 64-token chunks halve construction queries with a 0.14-point drop.The 16-token setting doubles retrieval queries without improving the mean.Source: §5.1 · p. 9, Figure 3 · p. 10
  • Two to eight fusion layers starting at layer 9 achieve 27.09–27.54, versus 21.11 with final-layer-only fusion and 22.07 with the final four layers; intermediate-and-late access is more effective.The four-layer [2, 10, 21, 31] configuration scores 26.32, 0.89 points below the default.Source: §5.1 · pp. 9–10
  • Top-2 remains within 0.33 points of Top-16’s 27.21 six-task mean, while using substantially fewer memory slots.Top-16 exceeds Top-8 by only 0.12 points.Source: §5.1 · p. 10, Figure 3 · p. 10

5.2. Memory Gating

  • Removing the gate lowers the six-task mean from 27.21 to 26.83, while zero gate-bias initialization yields the highest mean, 27.90.The retained default bias is -2, which initially scales down the memory branch while allowing the gate to adapt during training.Source: §5.2 · p. 10, Table 3 · p. 11

6. Conclusion

DFM combines chunk-level supervision with gated memory fusion inside the decoder, training added modules while the backbone remains frozen. It reduces target-construction retrieval and improves QA averages over the strongest baseline on both tested backbones.

6. Conclusion

  • DFM shares target memory across tokens within each chunk and uses gated cross-attention to integrate predicted memory into intermediate decoder states.Memory prediction and fusion are trained end to end with the backbone frozen, and inference requires no explicit retrieval.Source: §6 · p. 11
  • DFM improves six-task QA averages over the strongest baseline by 12.71% on Mistral-7B-v0.3 and 11.19% on Qwen3-8B-Base.Beyond QA, it improves performance over the base models across code, mathematics, and law.Source: §6 · p. 11

A. Experimental Details

The QA experiments specify Wikipedia-based memory construction, frozen-backbone training, and a six-benchmark evaluation protocol. Domain adaptation applies the same DFM components to code, math, and law data, with domain-specific training and evaluation settings.

A.1. Question Answering

  • QA training uses 32-token chunks, top-16 retrieved target embeddings, and a frozen backbone with jointly trained memory modules.The Wikipedia datastore contains training chunks only; target memory is retrieved for the next chunk, and the first chunk is masked from supervision.Source: §A.1 · p. 14
  • The QA comparison includes retrieval, continued-pretraining, LoRA, and token-level parametric-memory baselines, alongside DFM.LoRA's added parameter budget is matched to DFM; MLP Memory uses task-specific interpolation weights selected on each evaluation split.Source: §A.1 · p. 14
  • QA is evaluated on six benchmarks using zero-shot greedy decoding without online retrieval, gold documents, or chain-of-thought demonstrations.Memory queries come from the visible prompt and remain fixed during generation; scores use task-specific metrics on a 0–100 scale, with an unweighted six-task mean.Source: §A.1 · pp. 14–15, Table 4 · p. 15

A.2. Domain Adaptation

  • Domain adaptation uses 100M training tokens each for Code and Math and 123.5M for Law, with domain-specific data preparation.Code and Math use document-level 98/1/1 splits; Law uses DISC-Law-SFT and preserves its Reference field when present.Source: §A.2 · p. 15
  • Domain-adaptation DFM freezes the backbone, trains for one epoch, and uses the same memory modules and gate initialization as the QA experiments.Sequence length is 8K only for Code on Qwen3-4B-Base and 2K for other backbone–domain pairs; CPT and LoRA use the same data and sequence lengths.Source: §A.2 · p. 15
  • Domain evaluation covers HumanEval, HumanEval+, GSM8K, and LawBench, with the main-table average equally weighting the four benchmarks.LawBench uses 500 examples per task across 20 tasks in zero-shot evaluation, and its score is the equal-weight macro average of those task metrics.Source: §A.2 · p. 15

B. Additional Experimental Results

The additional experiments report QA results for OLMo-3-1025-7B and provide task-level LawBench results for two Qwen3 model sizes. DFM's OLMo six-task mean exceeds the strongest listed baseline.

B.1. Question Answering Performance

  • DFM reaches 25.49 on OLMo-3-1025-7B, 1.86 points above RETRO’s 23.63 six-task mean.DFM improves four tasks over Base, including 22.60 points on 2Wiki and 12.12 on Bamboogle, while WebQA and HotpotQA decrease by 1.28 and 0.63 points.Source: §B.1 · p. 16, Table 5 · p. 16

B.2. Domain Adaptation

  • Tables 7 and 8 report zero-shot results for 20 LawBench tasks on Qwen3-4B-Base and Qwen3-8B-Base, with an equal-weight macro average.The tables compare Base, RAG, CPT, LoRA, and DFM; Table 6 groups the task definitions into legal knowledge memorization, understanding, and application.Source: §B.2 · p. 16, Table 6 · p. 17, Table 7 · p. 18

C. Ablation Studies

The ablations change one factor at a time from the Mistral-7B-v0.3 configuration and use a common six-task QA protocol.

C. Ablation Studies

  • All ablations use predicted memory, a frozen backbone, one training epoch, and the same six-task QA protocol; results are single runs.Source: C. Ablation Studies · p. 20

C.1. Chunk Size, Fusion Placement, and Memory Slots

The ablations vary chunk size, fusion placement, memory-slot count, learning rate, projector sharing, predictor depth, and prediction loss. Their reported results show both task-level trade-offs and stability in several average scores.

C.1. Chunk Size, Fusion Placement, and Memory Slots

  • With 16-token chunks, 2Wiki and MuSiQue improve while the other four task scores decrease; 64-token chunks improve 2Wiki by 1.03 points.For a fixed token budget, target-construction query counts scale inversely with chunk size; elapsed construction time is not measured.Source: §C.1 · p. 20
  • 27.54 is the highest mean among the tested fusion-depth settings, attained by two layers; the four-layer reference is retained for consistency.The sweep changes layer count and placement together; the [28, 29, 30, 31] four-layer setting scores 22.07, while [2, 10, 21, 31] scores 26.32.Source: §C.1 · pp. 20–21
  • Across Top-2 to Top-16, the six-task mean varies by only 0.33 points, although Top-8 is 3.07 points lower on Bamboogle than Top-16.Top-8 exceeds Top-16 by 0.07–0.80 points on five tasks; Figure 4 shows task-level changes relative to the main configuration.Source: Figure 4 · p. 20, §C.1 · p. 21

C.2. Learning Rate

  • 27.21 versus 27.10: the default 1 × 10−5 learning rate has a 0.11-point higher mean than 2 × 10−5.The default is higher on NQ, WebQA, HotpotQA, and Bamboogle, but lower on 2Wiki and MuSiQue.Source: §C.2 · p. 21, Table 9 · p. 21

C.3. Projector Sharing

  • 27.21 versus 26.88: the shared projector has a slightly higher mean than independent per-layer projectors.Per-layer projectors improve four tasks but reduce NQ by 0.59 points and Bamboogle by 3.95 points; sharing avoids replicating projector parameters.Source: §C.3 · p. 21, Table 10 · p. 21

C.4. Memory Predictor Depth

  • 27.31, 27.21, and 26.47 are the six-task means for 14-, 28-, and 56-layer predictors, respectively.Halving depth changes the mean by 0.10 points; doubling depth lowers it by 0.75, driven mainly by a 4.74-point Bamboogle decrease.Source: §C.4 · pp. 21–22, Table 11 · p. 22

C.5. Memory Prediction Objective

  • 27.21 for raw MSE exceeds cosine loss at 26.87 by 0.34 points on the six-task mean.Cosine loss improves WebQA and 2Wiki but decreases the other four scores, with the largest drop on Bamboogle at 2.23 points.Source: §C.5 · p. 22, Table 12 · p. 22

D. Efficiency Analysis Details

DFM reduces construction and training costs relative to token-level MLP Memory while adding fewer parameters.

D. Efficiency Analysis Details

  • DFM adds approximately 0.75B parameters, approximately 50% fewer than the 1.5B MLP Memory configuration used for comparison.The reference configuration is used throughout the experiments, including parameter and training-time comparisons.Source: D. Efficiency Analysis Details · p. 23
  • 24 hours versus approximately 64 hours: DFM uses 37.5% of MLP Memory’s measured training time, a reduction of approximately 62.5%.Both methods used 16 A800 GPUs, the same training corpus, and the same number of epochs; timings exclude target construction and evaluation.Source: D. Efficiency Analysis Details · p. 23, Table 13 · p. 23
  • 32× fewer retrieval queries and 2048× fewer returned slots are required for chunk-level rather than token-level targets on a common 4B-token corpus.Returned slots count search outputs before MLP Memory removes the self neighbor and aggregates token probabilities; they are not distinct tokens, stored bytes, or FLOPs.Source: D. Efficiency Analysis Details · p. 23, Table 13 · p. 23

E.1. Question Answering

Selected HotpotQA examples show DFM answering cases missed by both baselines, while selected retrieval examples illustrate the supervision targets and their limits.

E.1. Question Answering

  • DFM answers all six selected HotpotQA cases correctly, while Base and MLP Memory give the same incorrect answer in each case.The examples include plant and film comparisons; selection criteria mean they do not establish how often these differences occur or reveal intermediate reasoning.Source: §E.1 · p. 24, Figure 5 · p. 24

E.2. Offline Retrieval Supervision

  • Seven selected retrieval examples pair a preceding 32-token query chunk with retrieved text and its ground-truth next chunk.For Pakenham, retrieved text supplies the network operator, Metro Trains Melbourne; highlighted spans mark continuation-relevant information missing or incomplete in the query chunk.Source: §E.2 · p. 25, Figure 6 · p. 25
  • The selected retrieval records exclude the query itself, exact query-text matches, and chunks from the same source block, but cross-block near duplicates may remain.These examples do not directly measure predicted-memory fidelity or attribute the answer differences to particular retrieved chunks.Source: §E.2 · p. 25

Related papers

Loading 2610.25100v1…