Doc2Atom: Compositional Memory for LLMs

• Doc2Atom is a compositional parametric memory framework that decomposes long documents into semantically typed knowledge atoms.
• It compiles each atomic unit into micro-LoRA adapters using an end-to-end multi-objective distillation framework for efficient and selective query-based memory injection.
• Doc2Atom enhances long-context processing by reducing GPU memory usage by 44% while mitigating irrelevant-query interference compared to monolithic document-level adapters.

Searching arXiv for the specified paper and the cited prior work to ground the article in the current literature. arXiv query: (Diao et al., 10 Jun 2026) Doc-to-Atom, abbreviated Doc2Atom, is a compositional parametric memory framework for long-context inference in LLMs. It is designed for settings in which long input sequences are central to document understanding and multi-step reasoning, but where the quadratic cost of attention makes inference memory-intensive and slow. The framework decomposes each document into semantically typed knowledge atoms, compiles each atom into an independent micro-LoRA adapter and a provenance retrieval key, and at inference time uses a lightweight query router to assemble a query-specific adapter that is injected into a frozen base model. The system is trained end-to-end through a multi-objective distillation framework and is presented in "Doc-to-Atom: Learning to Compile and Compose Memory Atoms" (Diao et al., 10 Jun 2026).

1. Position within context distillation

Doc2Atom is situated within the line of work on context distillation, in which long contexts are "written" into model parameters so that a frozen LLM can answer as if it had seen the full text at inference time. Standard context distillation requires expensive per-prompt fine-tuning. Doc-to-LoRA amortizes this procedure through a hypernetwork that maps any document $D$ to one low-rank adapter $\Delta W_D$ in a single forward pass, after which $\Delta W_D$ is applied to all queries about $D$ (Diao et al., 10 Jun 2026).

The central claim of Doc2Atom is that the monolithic structure of a single document-level adapter creates three specific limitations: irrelevant-query interference, limited compositional recall, and poor scalability. In the formulation given for Doc-to-LoRA, $\Delta W_D$ perturbs hidden states even for queries unrelated to $D$; a single low-rank adapter cannot flexibly surface distinct facts; and $\Delta W_D$ has fixed rank $r$ regardless of $|D|$, so longer documents cannot be fully internalized. Doc2Atom addresses these constraints by replacing the monolithic adapter with a set of composable atomic memories.

This design reframes document internalization as selective composition rather than single-shot compression. A plausible implication is that the framework treats long-document reasoning less as a one-time summarization problem and more as a retrieval-and-assembly problem in parameter space.

2. Knowledge atoms and semantic factorization

The foundational object in Doc2Atom is the knowledge atom. Given an input document $D$, the system first uses an LLM-driven annotation to factor $\Delta W_D$0 into

$\Delta W_D$1

where each atom $\Delta W_D$2 is a minimal, self-contained semantic unit of information. The paper specifies five atom types: fact_claim, entity_attribute, event_relation, process_step, and evidence_fragment (Diao et al., 10 Jun 2026).

Formally, the decomposition is written as

$\Delta W_D$3

where $\Delta W_D$4 is the source span, and $\Delta W_D$5 is the text of that span plus any necessary qualifiers. Each atom is also annotated with retrieval_text_i, a normalized anchor for retrieval, together with is_answer_bearing_i \in \{0,1\}, abstraction_level_i \in \{\text{abstract}, \text{evidence}, \text{hybrid}\}, conflict_group_i, confidence_i \in [0,1], and relations to other atoms.

These annotations give the atom a dual status: it is both a semantic unit for memory compilation and a typed retrieval object with provenance. The explicit inclusion of confidence, conflict grouping, and answer-bearing labels indicates that the atomization process is not merely segmentation. It is a structured representation intended to support routing, abstention, and controlled composition.

3. Memory compilation into micro-LoRA adapters

After decomposition, Doc2Atom compiles each atom independently. The framework shares the first $\Delta W_D$6 layers of the frozen base LLM as an AtomEncoder $\Delta W_D$7. For an atom $\Delta W_D$8,

$\Delta W_D$9

A MemoryCompiler consisting of an MLP trunk and four parallel heads then processes this embedding. With

$\Delta W_D$0

the provenance key head produces

$\Delta W_D$1

For memory injection, the micro-LoRA heads generate adapter factors for each target module $\Delta W_D$2 in q_proj, v_proj, o_proj, and down_proj, and for each designated layer $\Delta W_D$3:

$\Delta W_D$4

These are produced by two linear heads,

$\Delta W_D$5

$\Delta W_D$6

with $\Delta W_D$7 initialized to zero so that initially $\Delta W_D$8. The atomic weight update is then

$\Delta W_D$9

The framework also includes optional micro-KV prototypes for each atom, intended for prefix injection, and optional per-$D$0 logits that yield a sigmoid gate $D$1 (Diao et al., 10 Jun 2026).

The key architectural property is independence at compile time. Each atom yields its own retrieval key and low-rank perturbation, allowing subsequent query-conditioned composition. This differs materially from a document-level adapter because the atomic updates are designed to remain separable until routing.

4. Query routing, adapter composition, and end-to-end objectives

At inference time, Doc2Atom uses the same encoder to represent the query. Given query $D$2,

$D$3

and the query key is

$D$4

The router scores all atoms by cosine similarity plus metadata bias:

$D$5

where $D$6 is a learned scalar bias for atom $D$7 based on semantic type, confidence, answer-bearing flag, conflict group, and related metadata. Retrieval proceeds in two stages: first, the top-$D$8 atoms by maximum inner-product search; second, an optional reranking of those $D$9 atoms with a frozen cross-encoder to produce refined scores $\Delta W_D$0. The final top-$\Delta W_D$1 selected atoms $\Delta W_D$2 receive normalized routing weights

$\Delta W_D$3

(Diao et al., 10 Jun 2026).

Composition is performed separately for each layer $\Delta W_D$4 and module $\Delta W_D$5:

$\Delta W_D$6

With element-wise gating,

$\Delta W_D$7

where

$\Delta W_D$8

the composed update is

$\Delta W_D$9

This query-specific $D$0 is then injected into the frozen base model.

Training is end-to-end and multi-objective. If $D$1 denotes logits from feeding $D$2 into the frozen LLM, and $D$3 denotes logits from feeding only $D$4 plus the composed $D$5, the optimization includes:

• language modeling,

$D$6

• distillation,

$D$7

• routing supervision with gold mask $D$8 and higher $D$9 for distractors,
• irrelevant-query suppression,

$\Delta W_D$0

• knowledge protection,

$\Delta W_D$1

• sparse regularization with target density $\Delta W_D$2,

$\Delta W_D$3

• composition consistency using a gold-atom set,

$\Delta W_D$4

The full loss is

$\Delta W_D$5

and optimization uses AdamW with cosine-decay, warm-up, and gradient clipping. A four-stage curriculum gradually increases $\Delta W_D$6 and activates losses in dependency order: single-atom, multi-atom, robustness, and full.

5. Implementation profile and efficiency characteristics

The reported implementation details are given for a Gemma-2-2B-It run. The key hyperparameters are atom embedding dimension $\Delta W_D$7, LoRA rank $\Delta W_D$8, and $\Delta W_D$9. The encoder uses $r$0 layers, and memory is injected in the last $r$1 decoder layers, specifically $r$2–$r$3. The router uses $r$4, final $r$5, and routing temperature $r$6. Training uses batch size $r$7, learning rate $r$8, and $r$9 (Diao et al., 10 Jun 2026).

For per-document internalization, the paper reports the following empirical compile-step resource usage on Gemma-2-2B-It:

Method	GPU memory	Compile time
Doc-to-LoRA	$|D|$0	$|D|$1
Doc2Atom	$|D|$2	$|D|$3

The stated scaling property is that memory scales with $|D|$4 rather than $|D|$5, so Doc2Atom reduces peak GPU memory by $|D|$6 at the cost of higher offline compile latency. This trade-off is important for interpreting the system correctly: it is not presented as an unconditional speedup. Rather, it exchanges slower offline compilation for lower peak memory during document internalization and more selective inference behavior.

6. Benchmarks, comparative results, and interpretive distinctions

The empirical evaluation compares five systems: Base+ctx, defined as a frozen LLM with full document in context and treated as an ICL upper bound; D2L_ckpt, the official Doc-to-LoRA checkpoint; D2L_raw, Doc-to-LoRA retrained on the reported datasets; D2L_atom, Doc-to-LoRA trained on atomized documents; and Doc2Atom (Diao et al., 10 Jun 2026).

The reported overall F1 scores on Gemma-2-2B-It and Qwen3-4B-Instruct are:

Model	Gemma-2-2B-It F1	Qwen3-4B F1
Base+ctx	48.79	52.22
D2L_ckpt	37.93	18.70
D2L_raw	28.55	28.85
D2L_atom	29.41	29.30
Doc2Atom	37.99	35.72

Relative to the best D2L_atom baseline, Doc2Atom improves by $|D|$7 F1 on Gemma and $|D|$8 F1 on Qwen3. On the five in-domain QA sets—2WikiMultiHopQA, DROP, QASPER, ROPES, and SQuAD—Doc2Atom is described as competitive with full ICL and consistently above all D2L variants. In zero-shot generalization to eight LongBench subsets, it narrows the gap to ICL by up to $|D|$9 over D2L_raw and reliably refuses irrelevant queries, with refusal F1 $D$0, whereas D2L fails to abstain.

Two distinctions are especially important for interpretation. First, atomization alone does not reproduce the reported gains: the D2L_atom baseline remains substantially below Doc2Atom on both backbone families. This suggests that the principal effect comes from query-conditioned retrieval and composition of micro-adapters, not merely from preprocessing the document into smaller units. Second, the refusal result directly addresses the irrelevant-query interference identified in the motivation. The framework is explicitly trained with an irrelevant-query suppression objective, so abstention is part of the design rather than an incidental by-product.

7. Conceptual significance and scope

Taken together, Doc2Atom advances parametric context distillation by introducing fine-grained, semantically typed atoms and by transforming document internalization into a compositional memory system that is selective, scalable, and trainable end-to-end (Diao et al., 10 Jun 2026). Its internal representation couples semantic decomposition, provenance-aware retrieval, and low-rank parameter injection, making it a hybrid of structured memory compilation and adapter-based conditioning.

Within the terms reported for the method, Doc2Atom is best understood neither as ordinary retrieval-augmented generation nor as a standard single-adapter distillation system. It does not simply retrieve source text into the prompt, because inference proceeds by composing atom-level parameter updates into $D$1. Nor is it equivalent to a monolithic document adapter, because only a small subset of atoms is selected and assembled for a given query. A plausible implication is that the framework occupies an intermediate point between non-parametric retrieval and full parametric internalization, with provenance keys and atom composition providing the operative mechanism.

The reported limitations are also integral to its scope. The compile step is slower offline than Doc-to-LoRA, and the framework introduces a more elaborate decomposition, routing, and supervision pipeline. At the same time, the results indicate that this additional structure is associated with reduced peak GPU memory during internalization, stronger performance relative to Doc-to-LoRA baselines, and markedly better handling of irrelevant queries.