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

Paper: 2606.12400 Authors: Xingjian Diao, Wenbo Li, Yashas Malur Saidutta, Avinash Amballa, Lazar Valkov, Srinivas Chappidi Categories: cs.CL, cs.IR

The Gap

Current context distillation methods like Doc-to-LoRA generate a single LoRA adapter per document, amortizing compression into one forward pass. While efficient, this monolithic adapter treats all queries identically: it internalizes the entire document without distinguishing which parts are relevant to a given question. This causes irrelevant-query interference (the adapter’s parameters are polluted by unrelated content) and limits compositional recall (facts are entangled in one weight matrix, making multi-hop reasoning difficult). Moreover, scaling to very long documents forces a single adapter to encode too much, degrading performance.

The paper identifies a specific boundary: parametric memory for long-document QA is bottlenecked by the lack of decomposition and query-conditional composition. The logical flow:

Problem: Long-input attention is quadratic cost
  |
  v
Assumption: Compressing context into model params can avoid quadratic cost
  |
  v
Prior method: Doc-to-LoRA (single adapter per doc) is simple but suffers interference and poor recall
  |
  v
Hypothesis: Decomposing doc into typed atoms + composing query-specific adapter improves both
  |
  v
Method: Doc2Atom - compile atoms to micro-LoRAs with provenance keys, route queries to relevant atoms
  |
  v
Evidence: Outperforms Doc-to-LoRA on 6 QA benchmarks with lower memory cost
  |
  v
Conclusion: Compositional parametric memory is viable and superior to monolithic adaptation

The Increment

One sentence: Before Doc2Atom, parametric memory for documents meant a single, entangled adapter; after Doc2Atom, we can decompose a document into independent, query-addressable memory atoms that are compiled and composed dynamically, achieving better recall and efficiency.

Core Mechanism

Doc2Atom works in two phases: compilation and inference.

Compilation phase: Given a document D, a base LLM (frozen) processes it through a compressor module that outputs a sequence of “atoms.” Each atom is a triple: (type, content embedding, provenance key). The types are semantic categories (e.g., entity, relation, numeric fact, temporal fact) learned automatically. For each atom, a tiny micro-LoRA adapter is generated via a learned projection from the content embedding. Simultaneously, a provenance key (a vector summary) is produced for retrieval.

Inference phase: A query comes in. A lightweight router network computes attention between the query representation and all provenance keys, selecting a subset of relevant atoms (e.g., top-k). The corresponding micro-LoRA adapters are linearly combined (summed or averaged) into a single query-specific LoRA adapter. This adapter is then injected into the frozen base LLM, which processes the query and generates the answer.

The entire system is trained end-to-end with a multi-objective distillation loss: (1) reconstruction loss compressing the document into atoms, (2) retrieval loss making provenance keys match relevant queries, and (3) generation loss for final answer quality.

[ASCII diagram of method internals]

Compilation Phase:
Document D
    |
    v
[Compressor] --> Atom List: { (type, embedding, key), ... }
    |                    |
    v                    v
[Micro-LoRA Gen]    [Provenance Key Gen]
    |                    |
    v                    v
Micro-LoRA set       Key set
    \                    /
     \                  /
Inference Phase:
Query Q ----> [Router]
                   |
                   v
          [Attention: Q x Keys] --> selected keys (top-k)
                   |
                   v
          [Composer: sum selected micro-LoRA weights]
                   |
                   v
          Query-specific LoRA adapter --> [Frozen LLM] --> Answer

Structural metaphor: A modular reference library.

Imagine a large reference library where each book (document) is initially treated as one monolithic tome. A researcher (query) would have to scan the entire book, even if only a specific chapter is relevant, causing interference from unrelated chapters and making it hard to combine facts across different books. Doc2Atom’s insight is to break each book into individual “knowledge cards” (atoms), each containing a single fact, with a subject label (type) and a classification number (provenance key). When a researcher arrives with a question, a card catalog (router) quickly points to the relevant cards. The librarian then takes those cards and assembles a custom booklet (query-specific LoRA) containing only the needed facts. The researcher uses this booklet instead of the entire books. This avoids irrelevant content, enables combining facts from multiple books easily, and scales to huge libraries because each card is independent.

Key Concepts

• Knowledge Atom: A minimal, self-contained unit of factual knowledge extracted from a document. Each atom is typed (e.g., “Person-BirthYear”, “Company-Location”) and has a learned embedding that captures its content. Atoms are the currency of memory. Think of them as factoid triplets, but derived end-to-end without manual annotation.

• Micro-LoRA Adapter: A tiny adapter (e.g., rank 1 or 2) generated from a single atom. Unlike a full LoRA for the whole document, a micro-LoRA is so small that it only modifies a few weight dimensions in the base model. But combined, a set of micro-LoRAs can steer the model to recall specific facts without interfering with others. The key is that micro-LoRAs are additive and composable by simple sum.

• Query Router with Provenance Keys: The router is a lightweight neural module that maps a query to attention weights over provenance keys (vectors representing each atom). The provenance keys are learned to be discriminative: atoms that help answer similar queries have similar keys. The router selects top-k keys, effectively performing a sparse retrieval from the parametric memory pool. This ensures that only relevant facts are incorporated.

Framework Shift

Before (Doc-to-LoRA monolithic approach):        After (Doc2Atom compositional approach):

 Document                                          Document
    |                                                 |
    v                                                 v
[Compressor] --> one LoRA adapter                 [Atomizer] --> many micro-LoRAs + keys
    |                                                 |
    v                                                 v
Frozen LLM uses same adapter for ALL queries     Frozen LLM uses query-specific composed adapter
    |                                                 |
    v                                                 v
Answer (interference from irrelevant parts)      Answer (only relevant atoms contribute)

One sentence: From a one-size-fits-all parametric adapter to a modular, query-addressable memory bank, the core shift is decomposition and dynamic composition — treating parametric memory not as a monolithic blob but as a collection of independent facts that can be retrieved and assembled.

Expert Assessment

Problem choice: The gap is real. The quadratic cost of attention is a well-known bottleneck, and parametric memory approaches like Doc-to-LoRA are promising but indeed suffer from interference in multi-topic documents. This paper targets a genuine limitation with a principled decomposition. It sits at the intersection of efficient inference and memory augmentation, a hot area.

Method maturity: The idea is clever — using micro-LoRAs as modular memory and training them end-to-end with a router. But it borrows heavily from retrieval-augmented generation (RAG) concepts (dense retrieval over documents) and applies them to parametric memory. The novelty is in the marriage of micro-LoRA composition with learned provenance keys. The complexity is moderate; the main challenge is training stability (multi-objective optimization) which they seem to handle. A simpler alternative might be to use a mixture-of-experts style routing over adapter banks, but that would be less efficient.

Experimental integrity: Baselines are reasonable (Doc-to-LoRA, full-context, etc.). Results on 6 QA benchmarks show clear improvement, but I would have liked to see ablation studies on the importance of typing, number of atoms, and router mechanism. The memory cost is reported lower — good. However, I notice that all benchmarks are extractive or short-form QA; multi-step reasoning (e.g., HotpotQA) would be stronger evidence for compositional recall. The paper claims multi-step but only uses simple QA sets like SQuAD, TriviaQA. That’s a gap in evidence. Also, no evaluation on very long docs (e.g., 100k tokens) — they mention scalability but test on moderate length. So, numbers are solid but scope is limited.

Writing quality: Generally clear. The method section is well-structured. The weakest part is the related work — it’s a bit thin and doesn’t position deeply against retrieval-based methods (which also do composition). If they rewrote the discussion to compare more explicitly with RAG (retrieval over external databases vs parametric atoms), the paper would gain context. The abstract promises “compositional recall” but the experiments don’t clearly isolate that ability.

Verdict: weak accept — a solid incremental step that fixes a real issue in parametric memory, but the evidence for multi-step compositional reasoning is weaker than claimed, and the novelty is more in engineering than conceptual breakthrough.

Takeaways

• Use micro-adapters as modular memory units: Instead of one LoRA per task or document, consider extreme decomposition. This idea could transfer to continual learning (each new fact = a micro-adapter) or multi-task learning (composing adapters per instance).

• Learned provenance keys for retrieval within a parametric memory: This is a general technique: any parametric memory (e.g., hypernetworks, adapters) can be equipped with a key-query routing mechanism for conditional use.

• Multi-objective distillation for end-to-end training of memory components: The combination of reconstruction, retrieval, and generation losses is a recipe that can be reused in other memory-augmented models.