AtomComposer: Compositional Assembly Systems

• AtomComposer is a compositional system that decomposes complex objects into smaller 'atoms' to facilitate targeted construction across diverse domains.
• It is applied in chemical-space discovery, achieving up to an order of magnitude more valid 3D isomers than traditional reinforcement-learning baselines, and in assembling memory components for LLMs.
• The approach uses dynamic, query-conditioned routing and on-the-fly selection of micro-adapters or molecular fragments, ensuring scalable, efficient inference and robust performance.

Searching arXiv for papers mentioning "AtomComposer" and the cited IDs to ground the article. AtomComposer is a name used in recent arXiv literature for systems that decompose a complex object into smaller “atoms” and then construct, select, or compose them for a downstream objective. In “AtomComposer: Discovering Chemical Space from First Principles with Reinforcement Learning,” it denotes “a self-guided agent that autonomously constructs valid 3D isomers under stoichiometric constraints and is trained exclusively online using reinforcement learning” (Hastrup et al., 27 May 2026). In “Doc-to-Atom: Learning to Compile and Compose Memory Atoms,” it denotes the query-conditioned assembly mechanism that aggregates retrieved micro-LoRA adapters into a single adapter for a frozen base model (Diao et al., 10 Jun 2026). Related technical descriptions also apply the label to an end-to-end, first-principles blueprint built around atom-in-molecule fragments (“amons”), active learning, and kernel-ridge regression (Huang et al., 2017).

1. Nomenclature and scope

The term is not attached to a single standardized architecture. In the material considered here, it refers to distinct systems that share a compositional logic but operate in different technical domains.

Context	Paper	Role of “AtomComposer”
Chemical-space discovery	(Hastrup et al., 27 May 2026)	“a self-guided agent” for valid 3D isomer construction
Parametric memory for LLMs	(Diao et al., 10 Jun 2026)	query-conditioned aggregation of retrieved micro-LoRA atoms
Amon-based quantum ML blueprint	(Huang et al., 2017)	end-to-end first-principles pipeline using amons selected on-the-fly

This multiplicity matters because the same name can otherwise suggest a single method family. The literature here instead supports a narrower statement: AtomComposer is a recurring label for compositional assembly mechanisms, but the underlying objects being composed differ sharply—3D molecular structures in one case, low-rank parametric memories in another, and atom-in-molecule fragments in a third.

2. AtomComposer as a chemical-space discovery agent

In the 2026 paper “AtomComposer: Discovering Chemical Space from First Principles with Reinforcement Learning,” the stated problem is the discovery of “novel stable molecules without training data,” described as “a grand scientific challenge” (Hastrup et al., 27 May 2026). The paper contrasts current molecular generative models, which are said to be “trained on large, pre-curated datasets,” with a different paradigm: “autonomous, generalized agents capable of mapping vast, unknown chemical spaces without any pretraining.”

The central claim is that AtomComposer is “a self-guided agent that autonomously constructs valid 3D isomers under stoichiometric constraints and is trained exclusively online using reinforcement learning” (Hastrup et al., 27 May 2026). This formulation places validity, 3D structure, stoichiometric feasibility, and online RL at the center of the method. The same abstract further states that, “unlike existing approaches that generally overfit to a specific chemical formula,” the method establishes “a multi-composition training scheme that enables a broad generalization across diverse chemistry, guided by energy- and validity-based rewards.”

The reported comparative result is that the agent “can discover up to an order of magnitude more valid isomers on unseen test formulas than existing single-composition reinforcement-learning baselines trained with per-step energy rewards” (Hastrup et al., 27 May 2026). A plausible implication is that the method is positioned not merely as an optimizer over one formula, but as a generalized explorer of chemical configuration space across compositions. The paper’s summary statement is correspondingly programmatic: the results “fulfill the promise of online reinforcement learning as a powerful paradigm for scalable, from-scratch exploration of chemical configuration space.”

3. AtomComposer in Doc-to-Atom: atomization and memory compilation

Within the Doc-to-Atom framework, AtomComposer is the assembly component of a broader “compositional parametric memory framework” designed for long-document reasoning in LLMs (Diao et al., 10 Jun 2026). The starting point is document atomization. An input document $D$ is first split by “an LLM-driven annotator” into $N$ semantically typed atoms, $\mathcal{A}=\{a_1,\dots,a_N\}$, where each atom is “a minimal, self-contained fact, attribute, event, etc., with associated metadata (answer-bearing flag, confidence, conflict group, etc.).”

Memory compilation is performed offline. A shared text encoder, described as “the first $L_{\rm enc}$ layers of the frozen base LLM,” maps each atom to an embedding

$e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$

A small MLP “memory compiler” or hypernetwork then maps $e_i$ to a provenance key $k_i\in\mathbb{R}^{d_k}$, micro-LoRA factors $\{A_{m,i}^{(\ell)},B_{m,i}^{(\ell)}\}$ for each target module $m$ and layer $\ell$, and optionally a sparse mask $N$0 and micro-KV prototypes: $N$1 All such artifacts are stored in a “Memory Bank.”

The compilation step defines the atoms as independent parametric memories rather than text chunks. This is crucial to the framework’s contrast with monolithic context distillation methods. The abstract states that producing “a single monolithic adapter for all queries leads to irrelevant-query interference, limited compositional recall, and poor scalability to long-document reasoning,” whereas the atomized design is intended to isolate and later recombine document-local knowledge only when needed (Diao et al., 10 Jun 2026).

4. Query-conditioned assembly and routing

AtomComposer proper is the online query-conditioned assembly stage. A query $N$2 is encoded by the same shared encoder to $N$3 and projected to a normalized query key

$N$4

A two-stage router then retrieves relevant atoms. Stage 1 performs cosine MIPS over the stored provenance keys $N$5; an optional cross-encoder reranker and “learned metadata biases” refine the scores and produce soft routing weights $N$6 (Diao et al., 10 Jun 2026).

The similarity and routing equations are specified as

$N$7

$N$8

Here $N$9 is a learned metadata bias, with the description explicitly naming “answer-bearing, confidence, type, conflict” metadata.

Given the selected atoms, AtomComposer aggregates the retrieved micro-LoRA factors into a single query-specific adapter: $\mathcal{A}=\{a_1,\dots,a_N\}$0 If sparse masks are used, the aggregate gate is

$\mathcal{A}=\{a_1,\dots,a_N\}$1

followed by element-wise gating of the composed factors. The resulting adapter $\mathcal{A}=\{a_1,\dots,a_N\}$2, and optionally micro-KVs, is “injected into the frozen base model to answer $\mathcal{A}=\{a_1,\dots,a_N\}$3” (Diao et al., 10 Jun 2026).

The routing system is explicitly two-stage. The technical description states: fast MIPS retrieves top $\mathcal{A}=\{a_1,\dots,a_N\}$4 candidates “(e.g. 32)” through an approximate nearest-neighbor index; a Stage-2 cross-encoder may rerank; the final selection keeps top $\mathcal{A}=\{a_1,\dots,a_N\}$5 atoms “(e.g. 8)” and softmaxes their scores into routing weights. For irrelevant queries, “if all $\mathcal{A}=\{a_1,\dots,a_N\}$6 are very low, the softmax yields near-zero weights; the adapter effectively becomes null and the base model falls back to plain inference.”

5. Training objectives, ablations, and trade-offs

Doc-to-Atom trains the atomized memory system “end-to-end through a multi-objective distillation framework” (Diao et al., 10 Jun 2026). The loss decomposition includes language-model cross-entropy $\mathcal{A}=\{a_1,\dots,a_N\}$7, a knowledge-distillation KL term $\mathcal{A}=\{a_1,\dots,a_N\}$8, routing supervision $\mathcal{A}=\{a_1,\dots,a_N\}$9, irrelevant-query suppression $L_{\rm enc}$0, knowledge protection $L_{\rm enc}$1, sparse gate regularization $L_{\rm enc}$2, and composition consistency $L_{\rm enc}$3. The total loss is written as

$L_{\rm enc}$4

with “curriculum-scheduled weights $L_{\rm enc}$5.”

The empirical description emphasizes three properties. First, on QA performance, “Doc-to-Atom overall F1: 37.99 vs Doc-to-LoRA raw (flat adapter) 29.41 (+8.6 pts)” for a “Gemma-2-2B-It base model as example.” It also reports “Zero-shot LongBench suite F1 gains up to 6× over monolithic adapters.” Second, on irrelevant-query refusal, “Monolithic Doc-to-LoRA: 1–15 % on long-context tasks” versus “AtomComposer: 85–97 % (near-perfect refusal).” Third, on memory cost and compile efficiency, the reported figures are “Monolithic Doc-to-LoRA compile GPU memory: 24 GB (Gemma) / 34.6 GB (Qwen),” against “Doc-to-Atom compile additional memory: 13.4 GB (Gemma, –44 %) / 5.0 GB (Qwen, –85 %),” with “Compile latency (amortized offline): 2.1 s vs 0.3 s per document (acceptable since it runs offline)” (Diao et al., 10 Jun 2026).

A common misunderstanding would be to attribute the gains simply to improved chunking. The atomization-only ablation argues against that interpretation: “retraining the flat Doc-to-LoRA on atomized docs (vs raw) yields a small +0.8 F1 gain, confirming that the biggest win comes from dynamic composition rather than just better chunking.” The system-level summary in the same description is therefore precise: AtomComposer “replaces a single low-rank adapter per document with a pool of fine-grained micro-adapters (‘atoms’) that can be routed and sparsely composed on a per-query basis.”

6. Relation to amon-based fragment selection and kernel methods

A separate technical blueprint applies the AtomComposer label to an “end-to-end, first-principles” pipeline derived from Huang and von Lilienfeld’s AML methodology (Huang et al., 2017). In that setting, the atomic unit is not a document fact or a micro-adapter, but an atom-in-molecule fragment, or “amon.” A query molecule is represented by atomic numbers $L_{\rm enc}$6, Cartesian coordinates $L_{\rm enc}$7, and a bond-order/connectivity graph $L_{\rm enc}$8. An amon $L_{\rm enc}$9 is defined as “a small, connected subgraph of $e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$0 whose atoms and bonds preserve the local hybridization and bond-order of the parent.”

The blueprint states that “every heavy-atom environment in the query is contained in at least one selected amon.” Its local representation is many-body and atomic: $e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$1 where $e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$2 is the one-body term, $e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$3 is a discretized “Spectrum of London” two-body distribution, and $e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$4 is a discretized three-body “Axilrod–Teller–Muto” angular distribution. The description identifies AML’s default as the “aSLATM” many-body spectrum.

Amon selection is performed “on-the-fly” and in increasing fragment size. The pseudocode description enumerates connected subgraphs up to a maximum size, checks preservation of hybridizations and bond-orders, forms hydrogen-saturated candidate fragments, locally minimizes them, and retains fragments whose relaxed connectivity remains unchanged. The accompanying comments stress three points: fragments are generated “in increasing size,” subgraph isomorphism ensures “exact local motifs,” and “geometry-relaxation + connectivity check filters out unstable or chemically reorganizing fragments” (Huang et al., 2017).

The predictive model is kernel-ridge regression with atomic kernels. The atomic kernel is

$e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$5

the global kernel is the double sum

$e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$6

and the KRR weights are

$e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$7

Prediction for a query molecule is then

$e_i = \mathrm{Encode}(a_i)\in\mathbb{R}^d.$8

The scaling claims are specific. For “11 000 diverse C₉HₓOₓNₓ queries,” “chemical accuracy (MAE≲1 kcal/mol) reached with ≈50 amons of size ≤7 heavy atoms.” For biomolecules including “chignolin, 77 atoms” and “ubiquitin, 602 atoms,” the description states that energy, polarizability, atomic Mulliken-charges, NMR shifts, core-level shifts, and forces “all converge within DFT- or experiment-level accuracy after training on ∼100–500 amons (max heavy-atom-count ∼7–11).” It further reports that quantum chemistry for ubiquitin “took ∼4 CPU-years of a DFT code,” whereas AML training on amons “took ∼CPU-hours on a laptop,” with prediction in minutes (Huang et al., 2017).

Taken together, these uses of the term indicate a shared compositional motif: a large target object is decomposed into smaller atoms, and downstream inference depends on selective recombination rather than monolithic processing. This suggests a family resemblance at the level of systems design rather than a single canonical AtomComposer method.