Scaling-Aware Adapter for Structure-Grounded LLM Reasoning

TL;DR

Cuttlefish tackles the problem of faithful, scalable structure-grounded reasoning in all-atom LLMs by introducing Scaling-Aware Patching to adapt the query-token budget to structural complexity and a Geometry Grounding Adapter to inject geometry-aware tokens into the LLM. The architecture uses an SE(3)-equivariant EGNN encoder, an instruction-conditioned anchor-patching scheme, and cross-attention-based grounding to produce modality tokens that are fused into the LLM, mitigating geometry hallucinations. Trained on GEO-AT with multimodal all-atom data, and evaluated against general LLMs and modality-specific baselines across molecules, proteins, and nucleic acids, Cuttlefish achieves consistent, across-the-board gains in structure-grounded reasoning and token efficiency. The work provides a unified, scalable framework for atom-level reasoning and releases GEO-AT and code to spur further development in all-atom grounded LLMs.

Abstract

Large language models (LLMs) are enabling reasoning over biomolecular structures, yet existing methods remain modality-specific and typically compress structural inputs through sequence-based tokenization or fixed-length query connectors. Such architectures either omit the geometric groundings requisite for mitigating structural hallucinations or impose inflexible modality fusion bottlenecks that concurrently over-compress and suboptimally allocate structural tokens, thereby impeding the realization of generalized all-atom reasoning. We introduce Cuttlefish, a unified all-atom LLM that grounds language reasoning in geometric cues while scaling modality tokens with structural complexity. First, Scaling-Aware Patching leverages an instruction-conditioned gating mechanism to generate variable-size patches over structural graphs, adaptively scaling the query token budget with structural complexity to mitigate fixed-length connector bottlenecks. Second, Geometry Grounding Adapter refines these adaptive tokens via cross-attention to modality embeddings and injects the resulting modality tokens into the LLM, exposing explicit geometric cues to reduce structural hallucination. Experiments across diverse all-atom benchmarks demonstrate that Cuttlefish achieves superior performance in heterogeneous structure-grounded reasoning. Code is available at the project repository.

Figures (15)

• Figure 1: Mol-Llama performance on the Mol-Instructions captioning task, evaluated across five molecule length bins with 6 metrics (left y-axis, detailed in App \ref{['app: metrics']}) plotted as curves with dashed overall averages, and the background bars indicate the proportion of samples in each length bin (right y-axis).
• Figure 2: Architecture of Cuttlefish. The framework accepts all-atom inputs (spatial graph: atom features, coordinates, and spatial relations) processed by EGNN for modality embeddings. The model incorporates Scaling-Aware Patching through an instruction-conditioned gate and soft patch-growing mechanism. Then Geometry Grounding Adapter utilizes cross-attention to enrich adaptive tokens with granular geometric features derived from modality embeddings, subsequently projecting these modality tokens into the LLM's embedding space.
• Figure 3: Schematic of all-atom encoder pretraining. Masked reconstruction on all-atom spatial graphs with multi-task heads.
• Figure 4: Loss (all atom vs. single modality) across training stages: encoder (a, training; b, evaluation) with objective-wise losses, then modality alignment (c) and LLM adaptation (d). Dashed lines denote evaluation loss. X-axis: global steps (scaled by 0.1).
• Figure 5: Evaluation on DNA-Chat. (a) Mean Matthews correlation coefficient (MCC) over 18 DNA tasks, comparing Cuttlefish with representative baselines (x-axis). (b) Per-task MCC on the same 18 tasks, overlaying Cuttlefish and top-2 baselines.

Theorems & Definitions (12)

• Lemma 1.1: Differentiability of $g_{\text{anc}}$
• proof
• Proposition 1.2: Piecewise gradient correctness through discrete selection
• proof
• Lemma 1.3: Softmax Jacobian
• Lemma 1.4: Gradients of distances and anchor bias
• Lemma 1.5: Pooling gradients
• proof
• Theorem 1.6: Gradient existence and correctness
• proof