Doc-to-Atom: Teaching AI to Decompose Documents into Composable Memory Atoms

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 (arXiv:2606.12400), addresses the quadratic cost and 'lost in the middle' problems LLMs face with long documents. Building on Doc-to-LoRA (D2L) context distillation—which compresses an entire document into a single LoRA adapter—Doc2Atom instead decomposes documents into semantically typed, self-contained 'memory atoms' via an LLM-driven annotation pipeline, then compiles each atom into micro-LoRA factors, provenance keys, and metadata using a memory compiler. At query time, a two-stage router (cosine retrieval plus a BAAI/bge-reranker-base cross-encoder) selects the most relevant atoms, which are weight-combined into a query-specific adapter and written to the final four decoder layers. Tested on Gemma-2-2B-It and Qwen3-4B-Instruct across six QA benchmarks (SQuAD, DROP, ROPES, 2WikiMultiHopQA, QASPER, LongBench), Doc2Atom improves overall F1 to 37.99 (+8.58 over D2L-atom), reaches up to 97.5% F1 on irrelevant-query rejection versus D2L's sub-40%, cuts inference memory from 34.62 GB to 5.02 GB on Qwen3, and achieves up to 4.4x gains on LongBench zero-shot tasks.

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

A Chinese forum post (zhichai.net) reviews the Doc2Atom paper (Diao et al., Samsung AI Center & Dartmouth College, arXiv:2606.12400), which rethinks how LLMs internalize long documents.

Key points

Problem: Transformer attention scales quadratically with input length, and models tend to "lose the middle" of long contexts—making naive long-document QA slow, memory-hungry, and unreliable.
Prior art: Context distillation and Doc-to-LoRA (D2L, Sakana AI) compress a document into a single LoRA adapter via a hypernetwork. But one monolithic adapter per document causes (1) interference on irrelevant queries (D2L F1 mostly below 40% on rejection tests), (2) poor compositional recall from a flat low-rank representation, and (3) fixed capacity regardless of document length.
Core idea: Instead of one adapter per document, decompose each document into many "atoms"—independent, composable, semantically typed knowledge units, each with semantic type, self-contained content, provenance, retrieval key, confidence score, conflict-group tags, and typed inter-atom relations (9 types).

Pipeline

Offline — Atomization: An LLM (e.g., GPT-4) splits documents at semantic boundaries into non-overlapping, verifiable atoms (substring-checked against the source to eliminate fabricated citations). "Irrelevant probe generation" creates surface-similar but unanswerable questions to teach calibrated abstention.
Offline — Memory Compiler: A shared encoder (first 4 layers of the base model) maps atoms to 256-dim embeddings; parallel heads output provenance keys, rank-8 micro-LoRA factors (scaling 16, B initialized to zero), optional KV prototypes, and sparse masks.
Online — Two-stage router: Cosine retrieval narrows hundreds of atoms to 32 candidates; a frozen BAAI/bge-reranker-base plus a 4-dim metadata bias selects K atoms (K ramps 1→8 during training).
Online — Memory Assembler: Selected atoms' LoRA factors are weight-summed (A_c = Σ w_i × A_i, B_c = Σ w_i × B_i), gated sparsely, and injected only into the last 4 decoder layers.

Training

Four-stage curriculum (K = 1, 4, 8, 8) with 7–9 losses: language modeling, KL distillation (τ=1.5), routing supervision, irrelevant-query suppression, knowledge protection, sparsity regularization, composition consistency, conflict resolution, and confidence alignment.

Results (Gemma-2-2B-It and Qwen3-4B-Instruct; benchmarks: SQuAD, DROP, ROPES, 2WikiMultiHopQA, QASPER, LongBench)

Overall F1: 37.99 on Gemma (+8.58 over D2L-atom); 35.72 on Qwen3 (+6.42).
LongBench zero-shot: up to 4.4× D2L raw (NarrativeQA 4.64% → 20.58%); MFQA-en 7.14% → 26.00%; MuSiQue 6.95% → 20.86%.
Irrelevant-query rejection F1: typically 85–97% (peak 97.50%) vs. D2L below 40%.
Memory: 34.62 GB → 5.02 GB (−85%) on Qwen3; −44% on Gemma. Trade-off: longer offline compilation (2.10s vs 0.31s).
Ablation: atom memory + full context beats context alone (+2.26 overall F1; MuSiQue +9.11), suggesting atoms are a structured complement for multi-hop reasoning, though gains vanish on simple fact lookup (TriviaQA).

Limitations and future work

Not all datasets benefit; ~3% of samples fail annotation; generated probes occasionally become answerable; curriculum tuning is complex; training used 8× A100 80GB. Future directions include end-to-end learnable atomization, dynamic atom granularity, cross-document atom sharing, incremental updates, and multimodal atoms.

References

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

*Structured English summary of a Chinese-language forum review; all facts, figures, and citations are preserved from the original post.*