Doc-to-Atom: Learning to Compile and Compose Memory Atoms

Abstract: Long input sequences are central to document understanding and multi-step reasoning in LLMs, yet the quadratic cost of attention makes inference both memory-intensive and slow. Context distillation mitigates this by compressing contextual information into model parameters, and recent work such as Doc-to-LoRA amortizes context distillation into a single forward pass that generates one LoRA adapter per document. However, producing a single monolithic adapter for all queries leads to irrelevant-query interference, limited compositional recall, and poor scalability to long-document reasoning. To address these challenges, we propose Doc-to-Atom (Doc2Atom), a compositional parametric memory framework that decomposes each document into semantically typed knowledge atoms. Each atom is compiled into an independent micro-LoRA adapter and a provenance retrieval key. At inference time, a lightweight query router selects and assembles only the relevant atoms into a query-specific adapter, which is then injected into a frozen base model. The entire system is trained end-to-end through a multi-objective distillation framework. Experiments on six diverse QA benchmarks demonstrate that Doc2Atom outperforms Doc-to-LoRA baselines while reducing the memory cost of document internalization.

• The paper presents a novel approach where document content is decomposed into fine-grained knowledge atoms to overcome interference and scalability limitations of monolithic adapters.
• It details an atomization pipeline using explicit semantic annotations to compile micro-LoRA factors and sparse gate masks, enabling selective query-conditioned memory assembly.
• Empirical results demonstrate significant improvements in parameterized accuracy and GPU efficiency, with up to 8.6 F1 gains and 44–85% reduced memory usage compared to baseline methods.

Doc-to-Atom: Compositional Parametric Memory via Knowledge Atoms

Introduction and Motivation

Context distillation techniques in LLMs aim to eliminate expensive attention over long contexts but typically collapse the representation of a document into a monolithic adapter, as in Doc-to-LoRA (D2L). This monolithic paradigm results in (i) irrelevant-query interference—where the same adapter perturbs outputs for unrelated queries, (ii) reduced compositional recall—since evidence for a query is not dynamically selected, and (iii) poor scalability—because fixed-size adapters are forced to compress ever-larger documents. To address these limitations, "Doc-to-Atom: Learning to Compile and Compose Memory Atoms" (2606.12400) introduces a compositional parametric memory framework, Doc2Atom, which internalizes documents as fine-grained, semantically typed “knowledge atoms” that can be selectively composed into query-specific memory adapters.

Figure 1: Comparison—Doc-to-LoRA compresses documents into single adapters, while Doc2Atom decomposes documents into composable semantic atoms.

Methodology

Data Atomization Pipeline

Doc2Atom begins with explicit document decomposition, using LLM-assisted, XML-driven annotation to extract non-overlapping, semantically-typed knowledge atoms (e.g., fact claims, entity attributes, event relations). Each atom includes content, source span, retrieval anchor, type, answer-bearing flag, abstraction level, conflict group, confidence, and inter-atom relations (see details in Appendix). This granular decomposition is a fundamental shift from previous token chunking schemes and is directly optimized for minimal, independently retrievable and composable units.

Figure 2: Raw documents are atomized into a semantically annotated library; questions are mapped to relevant atoms for precise supervision.

Atom-level decomposition is augmented by synthetically generated irrelevant probes—questions surface-similar to in-domain queries but unsupported by the document—to explicitly supervise refusal behavior. Each query is aligned to gold, supporting, and distractor atoms, resulting in richly structured training instances for robust learning.

Compositional Memory Compilation and Query Routing

Doc2Atom’s core engine consists of:

1. Offline Compilation: Each annotated atom is encoded via a shared text encoder (re-using the early layers of the frozen base LLM), then passed through a memory compiler to produce:
   • A provenance key for retrieval
   • Micro-LoRA factors (per-atom, per-layer/module)
   • (Optional) micro-KV prototypes for memory injection
   • Sparse gate masks to restrict write locations
2. Online Query-Conditioned Assembly: At inference, a query is encoded and routed via a two-stage process—cosine MIPS over keys, followed by reranking using a cross-encoder and a metadata-informed bias. The top-$K$ relevant atoms are composed by weighted summation of their micro-LoRA factors, masked for sparsity and injected only into designated layers of the base LLM.

   Figure 3: Memory compilation and query-conditioned inference in Doc2Atom—only query-relevant atoms are dynamically assembled and injected.

This design enables selective perturbation: only queries supported by specific atoms trigger substantial modifications to the base model’s behavior, while unrelated queries yield null adapters, eliminating irrelevant-query interference.

Multi-Objective Training

Doc2Atom uses a comprehensive multi-objective distillation framework that includes:

• Cross-entropy and KL-divergence distillation to match the teacher’s full-context behavior
• Binary and listwise losses to supervise router selection against gold, supporting, and distractor atom masks
• Irrelevant-query suppression and gating sparsity constraints
• Knowledge protection terms to minimize base-model drift
• Composition consistency loss to guarantee multi-atom adapters do not degrade with increased evidence

A staged curriculum progressively increases the number of composed atoms and loss complexity, efficiently balancing learning of all objectives.

Empirical Results

Doc2Atom is evaluated on two base LLMs (Gemma-2-2B-It and Qwen3-4B-Instruct) across six in-domain QA tasks and eight LongBench zero-shot generalization tasks. In all settings, Doc2Atom achieves the highest parameterized accuracy (up to $8.6$ F1 improvement over D2L baselines on Gemma, $6.4$ on Qwen) and demonstrates robust cross-dataset gains enabled by compositional structure.

Strong refusal is achieved: when queries are out-of-domain, Doc2Atom returns "unanswerable" or base-model outputs with F1 above $85$ on most datasets, far exceeding D2L which typically falls below $40$. Efficiency analysis reveals Doc2Atom uses $44$-$85\%$ less GPU memory than D2L during document compilation, with the only trade-off being modestly higher offline compilation latency. Crucially, the top-$K$ atom composition bounds memory use independent of document length.

Discussion and Implications

Doc2Atom demonstrates that parametrizing knowledge as combinatorial, semantically-typed atoms rather than monolithic adapters yields significant gains in selective recall, compositionality, and interference avoidance. This architecture aligns with cognitive models of knowledge retrieval and supports practical extensions:

• Scalability: Offline document compilation and atom-level memory enable knowledge accumulation without ballooning runtime cost.
• Robust Retrieval: Explicit atom-level annotations and distractor supervision mean compositional memory routes only answers genuinely supported by evidence.
• Knowledge Lifecycling: Documents can be recompiled, atoms updated or invalidated, and memory libraries managed modularly.

  Figure 4: The XML-driven pipeline for atom decomposition, source-span validation, and robust question-to-atom annotation.

Potential future work includes scaling Doc2Atom to larger LLMs, adapting the atom taxonomy dynamically, and integrating with broader multimodal knowledge sources. While Doc2Atom substantially reduces the gap with in-context learning—even surpassing it on multi-hop and aggregation tasks when combined with document access—a full match to ICL has not yet been attained.

Conclusion

Doc2Atom (2606.12400) introduces compositional parametric memory at a granularity and selectivity not previously achieved in context distillation architectures. By decomposing documents into semantic atoms and dynamically composing query-specific adapters, it robustly internalizes new knowledge, achieves reliable refusal on out-of-domain queries, and scales efficiently with document size. This work operationalizes a practical and theoretically compelling direction for LLM memory: fine-grained, compositional, and query-conditioned parametric knowledge. Going forward, Doc2Atom provides a concrete basis for scalable, modular, and updatable LLM memory management.