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

2026-06-10

AI Summary

This work addresses the high computational cost of traditional attention mechanisms in long-document understanding and the limitations of existing context distillation methods, which rely on a single monolithic adapter and suffer from query interference, weak compositional recall, and poor scalability. The authors propose a composable parameterized memory framework that decomposes documents into semantically typed knowledge atoms, each compiled into an independent micro-LoRA adapter paired with a retrieval key. During inference, a lightweight query router dynamically assembles relevant atoms to construct a query-specific adapter, which is then injected into a frozen base model. This approach enables, for the first time, end-to-end compilation of documents into structured, composable memory atoms, significantly outperforming Doc-to-LoRA baselines across six question-answering benchmarks while reducing memory overhead and improving inference efficiency, accuracy, and scalability.

Abstract

Long input sequences are central to document understanding and multi-step reasoning in Large Language Models, 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.

Problem

• long-context reasoning
• context distillation
• parameter-efficient fine-tuning
• compositional memory
• query interference

Innovation

• compositional memory
• knowledge atoms
• micro-LoRA
• context distillation
• query routing