Doc-to-Atom: Compiling Documents into Composable Memory Atoms for LLMs

Forum topic · 小凯 · 2026-06-11

Summary

Doc-to-Atom (Doc2Atom), a paper by Xingjian Diao et al. from Samsung Electronics AI Center and Dartmouth College, addresses a key limitation of context-distillation approaches like Doc-to-LoRA (D2L), which compress an entire document into a single monolithic LoRA adapter. Instead of one adapter per document, Doc2Atom decomposes documents into typed, self-contained "memory atoms" — minimal knowledge units with provenance, retrieval keys, confidence scores, conflict markers, and inter-atom relations. An offline pipeline uses an LLM to atomize documents and compile each atom into micro-LoRA factors, KV prototypes, and sparse masks. At inference, a two-stage router retrieves and reranks atoms, then assembles a query-specific composite adapter. Trained with a four-stage curriculum and up to nine loss terms, Doc2Atom outperforms D2L across six QA benchmarks: 37.99 vs 29.41 overall F1 on Gemma-2-2B-It, up to 6x gains on LongBench zero-shot tasks, over 85% F1 (peak 97.50%) on irrelevant-query rejection versus below 40% for D2L, and an 85% memory reduction (34.62 GB to 5.02 GB) on Qwen3-4B-Instruct. Ablations show atom memory complements in-context text, especially for multi-hop reasoning.

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

> "If you can't explain it to a six-year-old, you don't really understand it." — Richard Feynman

This post is a Chinese forum's in-depth explainer of the paper *Doc-to-Atom: Learning to Compile and Compose Memory Atoms* (Diao et al., Samsung Electronics AI Center & Dartmouth College, arXiv:2606.12400). Below is a structured English rendition of its full analysis.

Key points

The problem: Transformer attention scales quadratically with input length, and models exhibit "lost in the middle" behavior, making long-document QA slow, memory-hungry, and unreliable.
Prior approach (Doc-to-LoRA, Sakana AI): Context distillation via a hypernetwork that generates a *single* LoRA adapter per document in one forward pass. Limitations: interference from irrelevant queries (F1 below 40% on rejection tests), limited compositional recall from a flat low-rank representation, and poor scaling to long documents under a fixed rank budget.
Doc2Atom idea: Decompose each document into multiple independent, composable, semantically typed "atoms" — like dismantling a Lego castle into labeled bricks instead of stuffing it into one bag.

How it works

Atoms carry: semantic type (fact, entity attribute, event relation, procedure, evidence), self-contained content, provenance (exact source span), retrieval key, confidence score, conflict-group marker, and 9 relation types to other atoms.

Offline stage — atomization and compilation:

An LLM (e.g., GPT-4) splits documents along semantic boundaries into non-overlapping atoms; each must pass substring verification against the source, eliminating fabricated citations.
"Irrelevant probe" generation creates superficially similar but unanswerable questions, teaching the router fine-grained answerability beyond lexical overlap.
A memory compiler encodes each atom (shared encoder, first 4 layers of the base model, 256-dim embeddings) and produces: provenance keys, rank-8 micro-LoRA factors (scale 16, B initialized to zero), optional micro-KV prototypes, and sparse masks over layers/modules.

Online stage — routing and assembly:

A two-stage router: cosine retrieval over provenance keys selects 32 candidates; a frozen cross-encoder (BAAI/bge-reranker-base) plus a 4-dim learned metadata bias (confidence, direct answer, conflict group, type) reranks to select K atoms (K grows from 1 to 8 during training).
Selected atoms are weight-summed into a composite adapter (A_c = Σ w_i A_i, B_c = Σ w_i B_i) and gated by sparse masks, with KV prefixes concatenated; writes are confined to the last 4 decoder layers.
Irrelevant queries yield near-zero routing weights, falling back to base model behavior.

Training: Four-stage curriculum (1 → 4 → 8 atoms) over 7–9 loss terms: LM loss, KL-distillation from the teacher with full context (τ = 1.5), routing supervision, irrelevant-query suppression, knowledge protection, sparsity regularization, composition consistency, conflict resolution, and confidence alignment.

Results

Overall: 37.99 F1 / 37.88 ROUGE-L on Gemma-2-2B-It vs 29.41/29.32 for atomized D2L (+8.58/+8.56); 35.72/35.58 on Qwen3-4B-Instruct (+6.42/+6.41).
LongBench zero-shot: up to ~6x over D2L raw — MFQA-en 26.00% vs 7.14%, MuSiQue 20.86% vs 6.95%, NarrativeQA 20.58% vs 4.64%, QASPER 29.85% vs 12.29%.
Irrelevant-query rejection: 85–97% F1 (peak 97.50%) vs D2L's <40%; wins on every benchmark even with Qwen3's tendency to answer unsupported questions.
Memory: 5.02 GB vs 34.62 GB on Qwen3-4B (−85%); 13.43 GB vs 24.02 GB on Gemma (−44%). Offline compilation is slower (2.10 s vs 0.31 s) but is a one-time cost.
Ablation: Doc2Atom + in-context text beats context alone (48.79 → 51.05 F1 overall; MuSiQue +9.11, MFQA-en +7.14, DROP +6.52), showing atom memory is a *structured complement*, most valuable for multi-hop and aggregation tasks. Simple factoid tasks (TriviaQA) see no gain.

Limitations and outlook

Not all datasets benefit (simple fact lookup doesn't need atomization); ~3% of samples fail full annotation and need fallbacks; irrelevant probes are occasionally accidentally answerable; the 4-stage curriculum with 9 losses is complex to tune; training requires 8× A100 80GB GPUs.
Future directions: end-to-end learnable atomization, dynamic atom granularity, cross-document atom sharing, incremental atom updates, and multimodal atoms.

Closing analogy

The author frames D2L as alchemy — everything thrown into one pot — while Doc2Atom is chemistry: decompose into labeled elements, then recompose on demand. The goal is AI that knows what it knows — and knows when to say "I don't know."

References

1. Diao, X., Li, W., Saidutta, Y. M., Amballa, A., Valkov, L., & Chappidi, S. (2026). *Doc-to-Atom: Learning to Compile and Compose Memory Atoms*. arXiv:2606.12400. 2. Hu, E. J., et al. (2022). *LoRA: Low-Rank Adaptation of Large Language Models*. ICLR. 3. Liu, N. F., et al. (2024). *Lost in the Middle: How Language Models Use Long Contexts*. TACL. 4. Charakorn, R., et al. (2026). *Doc-to-LoRA: Context Distillation with LoRA Adapters*. Sakana AI, ICML. 5. Ha, D., Dai, A., & Le, Q. V. (2016). *HyperNetworks*. ICLR. 6. Yang, A., et al. (2025). *Qwen3 Technical Report*. arXiv:2505.09388.

*Original Chinese explainer by “Xiaokai”, written in a Feynman-style explanatory voice.*