Scaling-Aware Adapter for Structure-Grounded LLM Reasoning

AI summary

Overview

• Research area: Multimodal large language models that reason over scientific structures — small molecules, proteins, DNA, and RNA — using all-atom 3D geometry.
• Technical level: Advanced (the paper assumes familiarity with LLM adapters, graph neural networks, cross-attention, and SE(3) equivariance), though the core ideas are explainable in plain terms.
• Scope: The paper introduces Cuttlefish, a unified structure-aware LLM whose connector scales its token budget with structural complexity and injects explicit geometric cues, evaluated across molecule, protein, DNA, and RNA benchmarks.

What This Paper Is About

Existing structure-aware LLMs are usually built for one modality only, and they typically convert 3D structural inputs into a fixed number of tokens using sequence tokenization or a fixed-length query connector. A fixed token budget over-compresses large structures (hurting performance on big molecules) while over-allocating capacity to small ones, and sequence-only inputs carry no geometry, which lets models produce structural claims that the input never supported. The paper's goal is a single connector that adaptively sizes its token budget to the structural complexity of each input and injects verifiable geometric evidence into the language model to suppress structural hallucination.

Key Contributions

1. Scaling-Aware Patching — an instruction-conditioned anchoring and patch-growing mechanism that produces a variable-size set of structural query tokens, so the modality token count grows with structural complexity instead of being pinned to a fixed length.
2. Geometry Grounding Adapter — a component that refines the adaptive patch tokens through cross-attention against full modality embeddings and injects the resulting modality tokens into the LLM embedding space, exposing explicit geometric cues.
3. Cuttlefish — a unified multimodal LLM that reasons over heterogeneous all-atom modalities (molecule, protein, DNA, RNA) and, per the paper, achieves strong performance across all-atom understanding benchmarks.
4. GEO-AT — described as the first all-atom instruction dataset with atom-level coordinates and language annotations spanning molecules, proteins, and nucleic acids, used for both training and evaluation and released as a community resource.

Main Findings

• Hallucination comparison (Table 1, 200 molecules and proteins from GEO-AT): Adding structural input reduces hallucination rate. Mol-Llama drops from 0.28 (0.91 HPM) sequence-only to 0.12 (0.59) with structure; ProtChatGPT from 0.34 (0.99) to 0.10 (0.55); Prot2Chat from 0.29 (2.30) to 0.06 (0.23); 3D-MoLM from 0.59 (2.23) to 0.23 (1.15). Answer rate changes vary, e.g. 3D-MoLM falls from 0.89 to 0.83, while ProtChatGPT stays at 0.99.
• Gains over general LLMs on GEO-AT (Table 2, average METEOR / BERTScore): Cuttlefish + Llama-3.1-8B-I reaches 0.428 / 0.864 and Cuttlefish + Qwen3-8B reaches 0.428 / 0.876, versus best general sequence-only baselines such as Llama-3.1-8B-Instruct at 0.186 / 0.694 and Mistral-3-8B-Reasoning-2512 at 0.176 / 0.744. Improvements hold across every tested backbone (Qwen2.5-7B, Llama-3.1-8B-I, Mistral-3-8B-I, GLM-4-9B, Qwen3-8B, R1-7B, Mistral3-8B-R).
• Modality-specific comparisons: On Mol-Instructions molecule tasks, Cuttlefish reports Captioning ROUGE-L 0.766 / METEOR 0.715, Design FTS 0.422, Forward FTS 0.792, Retro FTS 0.747, Reagent FTS 0.509, OpenQA BERTScore 0.884, T/F accuracy 0.660, MC accuracy 0.890, ER F1 0.780, and IE F1 0.270, versus Mol-LLama-3.1 at 0.759 / 0.707, 0.392, 0.774, 0.708, 0.411, 0.812, 0.600, 0.880, 0.700, and 0.200 respectively.
• Proteins: Cuttlefish reports IE F1 0.273, PF ROUGE-L 0.486, FD 0.520, CA 0.551, and DM 0.495, compared with Mol-Instructions (0.224 / 0.430 / 0.440 / 0.520 / 0.460) and ProtLLM (0.176 / 0.450 / 0.435 / 0.463 / 0.384).
• DNA and RNA: On 18 DNA tasks (mean MCC), Cuttlefish achieves top-1 average performance, slightly exceeding ChatNT, with limited gains attributed to substantial structural repetition in DNA. On RNA-QA, Cuttlefish reports ROUGE-1/2/L of 0.494 / 0.371 / 0.512 on Abstract Information Summary and 0.519 / 0.432 / 0.471 on Divide and Conquer, compared with RNA-GPT at 0.503 / 0.367 / 0.475 and 0.479 / 0.269 / 0.441.
• Scaling stability: On molecular captioning and protein functional description stratified by atom-count bins, Cuttlefish maintains stable quality across bins and shows its largest relative gains in the high-size regime, in contrast to the size-induced degradation the paper attributes to fixed-capacity connectors.
• Token efficiency: The paper reports that, compared with ChatNT — which the paper says scales token count linearly — Cuttlefish maintains a substantially lower structural token budget while still scaling with entity size.
• Ablations (Table 6, average drop from full model): Single-modality encoder is the mildest change (2.17% METEOR, 2.69% BERTScore); residue-level patching drops 10.49 / 10.53; removing instruction-guided gating drops 14.01 / 9.89; direct projection instead of fusion blocks drops 20.96 / 12.38; a Q-Former variant with the same maximum token length drops 43.15 / 11.31; sequence-only baseline drops 54.32 / 16.83; sequence-only with enhanced tokenizer drops 55.89 / 19.46; removing soft patch growth drops 66.83 / 22.57.
• Matched-budget comparison (Table 7, BERTScore): Q-Former at fixed 256 tokens yields 0.778 / 0.743 / 0.658 (molecule / protein / DNA&RNA) with a 0.726 average, at 512 yields 0.744 average, at 1024 yields 0.768 average, and degrades to 0.760 at 2048. Cuttlefish matched to a 256 budget reports 0.842 / 0.659 / 0.693 for a 0.731 average. The table is truncated in the provided content.
• Training behaviour: The encoder converges stably and quickly on three objectives, which the authors say indicates a large amount of data is not needed to learn spatial relations in all-atom entities. Mixed all-atom modality training converges better than single-modality training, which the authors attribute to cross-modality diversity.
• Not reported in the provided content: Results for the structure-availability analysis, backbone size ablation, extended hallucination analysis, coordinate-noise robustness, and data contamination analysis are referenced as Sections 4.8–4.12 but their outcomes are not included in the supplied text.

Methodology in Plain English

The system has three stages plus a training recipe.

Encoding geometry. Atomic coordinates and features for each all-atom entity are passed through an SE(3)-equivariant graph neural network (EGNN), which produces one embedding per atom while staying invariant to rotations and translations. This encoder is pretrained by masking atoms and edges and asking it to predict element identity, pairwise distances, and injected noise on direction vectors, combining three losses (type, distance, and direction) with weighting coefficients.

Choosing how many tokens to use (Scaling-Aware Patching). Instead of always producing the same number of structural tokens, the model scores every atom for its relevance to the specific instruction using an instruction-conditioned gate. It normalizes those scores, sorts them, and keeps adding atoms until their cumulative probability mass passes a threshold, capped at a maximum. Those atoms become anchors. Each anchor then "grows" a patch over nearby atoms using a soft assignment weight that blends spatial distance with the anchor's relevance score, so highly relevant anchors capture wider regions. Atoms within a patch are combined by membership-weighted pooling to form patch tokens. Because the number of anchors depends on the input, large structures get more tokens and small ones get fewer.

Grounding and injecting into the LLM (Geometry Grounding Adapter). Patch tokens become queries that go through fusion blocks performing self-attention then cross-attention against the full atom-level embeddings, retrieving the fine geometric detail that pooling abstracted away. The refined features are projected into the LLM's embedding dimension and inserted at placeholder positions in the instruction sequence, with attention and label masks updated to match.

Training. First, modality alignment tuning freezes both the EGNN encoder and the LLM and trains only the connector on the GEO-AT all-atom instruction corpus, so the geometry-to-text interface is learned without disturbing pretrained language ability. Second, LLM adaptation tuning unfreezes the LLM and continues end-to-end with a smaller learning rate so the model can adapt to the new modality injection distribution. The paper contrasts this with Q-Former-style pipelines that need heavy contrastive pretraining to fit a fixed set of learnable queries, arguing that its structural queries are dynamically created and carry explicit geometric meaning, so alignment can be achieved by instruction-supervised end-to-end tuning alone.

Evaluation. The model is compared on Mol-Instructions (about 700K samples), DNA-Chat (a 7.8M-sample corpus over 18 tasks built from the Nucleotide Transformer benchmark), RNA-QA (400K+ literature-grounded QA pairs derived from RNAcentral), and the authors' own GEO-AT. Text-output tasks use METEOR and BERTScore, generation tasks use validity and RDKit fingerprint similarity, closed-form QA uses accuracy, DNA uses Matthews correlation coefficient, and RNA summarization uses ROUGE-1/2/L.

Why This Matters

The paper argues that structure-aware LLMs cannot be built by squeezing all-atom geometry through a fixed-size bottleneck, and that fluent language output without geometric grounding correlates with structural hallucination. If the reported trends hold, this points toward adapters that allocate capacity in proportion to structural complexity and that keep a verifiable geometric channel open between the structure encoder and the language model.

Real-world applications:

• Drug discovery and medicinal chemistry — the paper evaluates molecular design, forward prediction, retrosynthesis, and reagent prediction, all tasks that sit inside computational chemistry pipelines.
• Protein engineering and annotation — protein function, functional description, catalytic activity, and domain or motif prediction are evaluated directly, with structures retrieved from PDB or generated by AlphaFold2 when missing.
• Nucleic acid research — DNA task classification and RNA literature summarization and functional annotation suggest use in genomics and transcriptomics workflows.
• Scientific assistants that must not fabricate structure — the hallucination test on 200 molecules and proteins from GEO-AT is motivated by dialogue settings where a model invents geometry it was never given.

Industry relevance: The results are measured across seven different open-source LLM backbones, which suggests the connector is designed as a drop-in component rather than a bespoke model, and the code is released at github.com/zihao-jing/Cuttlefish. Genomic, pharmaceutical, and biotechnology organizations that want to condition existing LLMs on atomic-level structure would be the natural adopters.

Future Directions

• Repair the weak spots. Cuttlefish matched to a 256-token budget trails Q-Former on protein and DNA&RNA BERTScore in the matched-budget comparison (0.659 and 0.693 versus 0.743 and 0.658 respectively by modality), so closing that gap at low budgets is an open problem.
• Exploit modalities where geometry helps less. The paper notes only limited gains from atom-level inputs on DNA, attributing this to structural repetition. Determining when 3D grounding pays off versus when sequence information suffices is unresolved.
• Understand the scaling budget. How the mass threshold and maximum anchor count should be set, and whether the observed token efficiency can be pushed further without losing the high-size-bin gains, is not settled by the reported ablations.
• Complete the referenced analyses. Structure availability, backbone size, extended hallucination, coordinate-noise robustness, and data contamination studies are referenced as Sections 4.8–4.12 but their findings are not in the available content, leaving those robustness questions open.

Target Audience

This paper is most valuable to researchers and engineers working on multimodal LLMs, scientific foundation models, or geometric deep learning, particularly those building connectors between graph or 3D structure encoders and language models. It is also relevant to computational chemists and biologists who want LLMs grounded in atomic structure, and to practitioners who need to choose between fixed-length Q-Former-style connectors and adaptive token budgets. Readers without a background in graph neural networks or adapter architectures will find the methods dense, though the problem framing and the hallucination and scaling results are accessible.

Authors’ abstract

Large language models (LLMs) are enabling reasoning over 2D and 3D 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 grounding 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 multimodal 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 interdisciplinary all-atom benchmarks demonstrate that Cuttlefish achieves superior performance in heterogeneous structure-grounded reasoning. Code: github.com/zihao-jing/Cuttlefish.