PaperlayerView exact revision on arXiv

Source-linked AI summary

CoMPASS: Collaborative Molecular Property Prediction via Adaptive Small-Large Model Synergy

Wentao Li, Jiangjie Qiu, Yijun Li, Leyi Zhao, Xiaonan Wang

arXiv:2608.30674v1cs.LGcs.AI

TL;DR

Molecular property predictors must combine calibrated numerical outputs with chemical reasoning, but graph models can be locally uncertain while LLMs are unreliable as standalone predictors. CoMPASS uses retrieval-grounded LLM reasoning to make bounded, agreement-aware corrections to a GAT anchor, improving it across eight benchmarks while limiting intervention when confidence is high.

  • Problem

    Molecular property prediction requires calibrated outputs and chemical reasoning, while graph predictors can be unreliable when local evidence is sparse or conflicting.

  • Method

    CoMPASS retrieves relevant molecules, supplies attention-grounded evidence to an LLM, and converts its proposal into a bounded correction of a GAT anchor through agreement-aware gating.

  • Results

    CoMPASS improves its paired GAT anchor across six classification and two regression benchmarks, with gains concentrated where local evidence reveals correctable uncertainty.

  • Takeaways & Limitations

    LLMs are most useful as evidence-grounded, agreement-filtered corrections to supervised graph predictors rather than as replacements for them.

  • Takeaways & Limitations

    CoMPASS depends on retrieval quality, can over-correct with misleading neighborhoods, adds LLM inference costs, and is evaluated on fixed benchmark splits.

Abstract

from arXiv · show

Accurate molecular property prediction requires both statistical reliability and chemical reasoning. Graph neural networks can be calibrated directly on labeled assays but remain limited by the coverage of their training data. Large language models (LLMs) can compare molecular evidence and articulate chemical rationales, yet are unreliable as standalone quantitative predictors. The central challenge is therefore to determine when an LLM should influence a calibrated model and by how much. Here we present CoMPASS, a retrieval-calibrated framework for small-large model collaboration. CoMPASS retains a graph attention network (GAT) as the predictive anchor, retrieves locally relevant training molecules, provides attention-grounded evidence to an LLM, and converts its proposal into a bounded correction through an agreement-aware gate. Across six classification and two regression benchmarks, CoMPASS improves the GAT anchor in regions of correctable uncertainty while limiting LLM intervention in high-confidence regimes. Ablations show that the gains arise from validation-calibrated retrieval and bounded fusion rather than prompting alone. These results suggest that generative reasoning should augment calibrated prediction through evidence-grounded, controlled corrections rather than direct output replacement. Code is available at https://github.com/littlepeachs/CoMPASS.

Introduction

CoMPASS addresses unreliable molecular property predictions by combining a calibrated GAT anchor with retrieval-grounded LLM reasoning. It treats LLM input as a bounded, agreement-aware correction rather than replacing the predictor’s output.

  • Motivation: Molecular property prediction must produce reliable numerical outputs for screening, prioritization, and risk assessment.The introduction frames reliability as necessary for downstream decisions in drug discovery and toxicology.
  • Motivation: GAT reliability is limited by training-distribution coverage, especially for rare scaffolds, sparse activity regions, decision boundaries, and noisy assays.Compact graph models may remain confident when local evidence is weak or conflicting.
  • Contributions: CoMPASS frames small–large collaboration as evidence-grounded correction of a calibrated graph predictor rather than direct output replacement.The framework is presented as retrieval-calibrated collaboration with controlled LLM intervention.
  • Method: CoMPASS keeps the GAT as the predictive anchor while retrieval provides locally relevant molecules and a validation-calibrated numerical estimate.The retrieved evidence is also supplied to the LLM alongside the target molecule, GAT prediction, and attention evidence.
  • Method: The LLM proposes an evidence-conditioned correction, and an agreement-aware bounded gate controls its influence on the final prediction.The gate limits intervention in high-confidence regimes while allowing correction in uncertain regions.

Method

CoMPASS combines a GAT anchor with attention-grounded retrieval, an LLM proposal, and a deterministic bounded gate. Retrieval is calibrated on validation data, while uncertainty, local conflict, attention consistency, and directional agreement control the final correction.

  • GAT anchor: CoMPASS decomposes prediction into four stages, beginning with a GAT that produces the supervised prediction, confidence, attention evidence, and graph representation.Classification uses prediction sharpness for confidence, whereas regression derives confidence from Monte Carlo dropout uncertainty.
  • GAT anchor: Attention highlights graph-model-important atoms or functional groups and serves as a communication channel to the LLM.The method uses attention as evidence rather than only for visualization.
  • Retrieval calibration: The retriever selects K calibration molecules from the training split using Morgan overlap, graph-embedding cosine similarity, and attention overlap.The query is excluded by SMILES match, and validation/test molecules are never inserted into retrieval.
  • Retrieval calibration: Similarity-weighted retrieved labels form a local estimate, and λ is selected on validation data to create a retrieval-calibrated numerical correction.The classification objective is negative ROC-AUC and the regression objective is RMSE, making retrieval more than prompt context.
  • Bounded fusion: The LLM receives GAT predictions, confidence, attention-highlighted groups, retrieved examples, conflict, and consistency, then returns a numeric proposal, correction direction, optional gate, and rationale.The final prediction interpolates between the calibrated base and LLM proposal using a deterministic bounded gate based on uncertainty, retrieval conflict, attention consistency, and directional agreement; agreement permits the largest cap, while disagreement is most conservative.
  • Bounded fusion: The agreement-dependent gate uses normalized uncertainty and assigns the largest cap to agreeing retrieval and LLM shifts, a smaller cap to weak agreement, and the most conservative cap to disagreement.Classification uncertainty is clipped from entropy, while standardized regression uncertainty is clipped from dropout uncertainty.

Experimental Setup

CoMPASS is evaluated on eight molecular-property benchmarks covering activity, pharmacokinetics, toxicity, and physicochemical regression, using fixed DeepChem splits and task-appropriate metrics. The main configuration uses Llama-3-8B with validation-selected retrieval settings, three-seed averaging, and specified GAT and LLM inference protocols.

  • Benchmarks and metrics: Eight benchmarks cover target activity, pharmacokinetics, toxicity, and physicochemical regression, with six classification tasks using ROC-AUC and ESOL and Lipo using original-scale RMSE.The benchmarks are BACE, BBBP, ClinTox, HIV, Tox21, ESOL, Lipo, and CYP450.
  • Data splits and selection: All runs use fixed DeepChem train/validation/test splits, while retrieval weight λ⋆ and retrieval size are selected on validation data and test labels are reserved for final evaluation.The evaluation protocol separates validation-based configuration selection from final test assessment.
  • Main configuration: The main configuration uses Llama-3-8B, evaluates K ∈{4, 6, 8}, and averages results across three independent seeds.GAT training runs for up to 30 epochs with early stopping; inference uses batch size 32, learning rate 10−4, weight decay 10−5, Morgan radius 2 with 2048 bits, and LLM temperature 0.7.

Results

CoMPASS improves the paired GAT anchor most when retrieved local evidence exposes correctable uncertainty, while bounded gating preserves performance where the anchor has limited headroom. Ablations and confidence-binned analyses attribute these gains to calibrated retrieval, attention-grounded reasoning, and agreement-aware correction rather than unrestricted LLM replacement.

  • Main results: CoMPASS improves the paired GAT across the reported benchmarks, with larger gains on BACE, BBBP, ClinTox, and HIV and smaller gains on CYP450 and Tox21.The results associate larger gains with retrievable boundary or calibration mistakes and smaller gains with already-strong anchors whose gates mainly limit harm.
  • Main results: Bootstrap intervals are positive for BACE, BBBP, ClinTox, HIV, ESOL, and Lipo, while Tox21 is positive but small and CYP450 has low headroom.The pattern supports improvement when local evidence reveals correctable GAT uncertainty and conservative gating when headroom is limited.
  • Retrieval analysis: Retrieval size is validation-calibrated: BACE, ClinTox, and Tox21 peak at K = 6, whereas BBBP, HIV, ESOL, and Lipo peak at K = 8.More neighbors can strengthen coherent local neighborhoods but dilute signal when property variation is sharp; K controls whether the LLM receives focused counterfactuals or noisier summaries.
  • Ablation analysis: Ablations show that LLM-only correction without calibrated retrieval is weak, while retrieval, attention evidence, and agreement-filtered correction provide the useful collaborative signal.Retrieval-only variants can improve regression and some classification tasks, but prompt evidence and LLM rationales become useful when capped by agreement and uncertainty.
  • Interface robustness: The collaboration interface is not tied to one LLM or graph architecture, with different LLMs competitive across endpoint types and gains extending to an attention-bearing graph transformer.Llama-3-8B is strongest on BACE and BBBP, chemistry-oriented LLMs are competitive on regression, and the transformer effect is largest on BACE.
  • Confidence analysis: Low-confidence GAT predictions benefit most from retrieved evidence and LLM interpretation, whereas high-confidence predictions require capping because over-correction risk is greater.The confidence-binned pattern explains why an unrestricted LLM mixture is undesirable.

Case and Failure Analysis

CoMPASS succeeds when retrieved analogues provide useful, attention-grounded evidence for bounded corrections, especially near classification boundaries or along local regression scales. Its main limitation is misleading retrieval, motivating conflict detection and adaptive correction caps rather than simply increasing LLM influence.

  • Successful corrections: Successful classification corrections move near-threshold GAT scores across the 0.5 decision boundary using directional signals from retrieved analogues.The gate permits enough movement to change the decision.
  • Successful corrections: Successful regression corrections move estimates partway toward a local physicochemical property scale rather than replacing predictions with invented labels.The LLM interprets structured evidence instead of directly inventing a label.
  • Failure analysis: Misleading retrieval can pull a correct classification across the threshold or move a regression estimate farther from its label.These failures expose retrieval as the method’s main limitation.
  • Failure analysis: The nonzero gate floor preserves large-model participation, while the cap prevents uncontrolled override.The floor avoids making the large model decorative, and future gains are expected from better conflict detection and more adaptive caps.

Operational Interpretation

CoMPASS operationally preserves task-specific GAT calibration while applying a scalar, bounded LLM correction to retrieval-calibrated task vectors. Its evidence trace and calibrated retrieval-and-gating design distinguish correction mechanisms and constrain LLM influence to supported cases.

  • Multi-task endpoints: For multi-task endpoints, CoMPASS retains the GAT task vector, applies a scalar LLM shift to retrieval-calibrated labels, and preserves the benchmark’s original masks.This prevents the LLM from hallucinating missing-task labels and keeps task-specific calibration with the supervised graph model.
  • Audit trail: Each prediction records molecular evidence, retrieval details, the LLM proposal, gate value, final prediction, and rationale for operational review.The trace distinguishes retrieval-driven correction, disagreement-suppressed proposals, and another proposal outcome with the same final metric effect but different meaning.
  • Why not prompt-only RAG?: Unlike prompt-only molecular RAG, CoMPASS gives retrieval a numerical role before consulting the LLM and uses validation-fitted blending to calibrate local-label movement.This addresses cases where chemically plausible retrieved analogues may nevertheless be statistically misleading.
  • Confidence-binned behavior: Mean correction benefit is largest in lower-confidence GAT bins and shrinks as the GAT becomes more confident.This confidence-binned behavior supports limiting LLM intervention in high-confidence regimes.
  • Why not prompt-only RAG?: Prompting and retrieval help only when calibrated evidence bounds the final prediction and the gate controls how much of the LLM interpretation survives.The framework therefore treats the LLM as a controlled correction rather than an unrestricted replacement for the graph anchor.

Related Work

Prior molecular property prediction work has used graph and fingerprint representations, with modern approaches learning task-specific molecular representations. CoMPASS extends related small–large model collaboration by making the graph model primary and using attention evidence, retrieval calibration, and uncertainty-aware gating to mediate bounded LLM assistance.

  • Molecular representation learning: Molecular property prediction has long relied on graph and fingerprint representations, while modern GNNs learn task-specific molecular representations.The passage cites foundational graph, fingerprint, and GNN approaches, alongside molecular pretraining and contrastive learning.
  • Related molecular RAG: CoMPASS differs from LLM-driven molecular RAG by treating the graph model as the primary predictor and retrieval as both prompt evidence and numeric calibration.This positions retrieval as serving both the language-model interaction and the predictive calibration process.
  • Small–large collaboration: CoMPASS instantiates small–large model collaboration for molecular prediction through attention evidence, retrieval calibration, and uncertainty-aware gating.The framework formalizes collaboration by mediating the large model’s role through these three mechanisms.

Conclusion

CoMPASS supports a restrained role for LLMs: they should interpret retrieved analogues and graph-attention evidence rather than replace the supervised graph predictor. Across eight benchmarks, this interface improves the paired GAT anchor, with gains driven by retrieval, validation-fitted local correction, and agreement-filtered bounded proposals.

  • Conclusion: CoMPASS uses the supervised graph predictor as an anchor and invokes the LLM to interpret retrieved analogues and graph-attention evidence.The LLM is intended to help when the small model may be locally miscalibrated, not to replace it.
  • Conclusion: Across eight benchmarks, CoMPASS improves the paired GAT anchor while exposing where headroom is small.
  • Conclusion: Ablations indicate that retrieval and validation-fitted local correction carry much of the signal, while agreement-filtered bounded proposals make the LLM useful.

Limitations

CoMPASS is intended for offline screening, triage, and mechanistic analysis because LLM inference limits scalability, while retrieval errors and fixed benchmark splits constrain reliability and generalizability.

  • LLM inference makes CoMPASS more appropriate for offline screening, triage, and mechanistic analysis than billion-scale virtual-library scoring.The framework adds an inference cost that limits its suitability for very large-scale virtual screening.
  • Retrieval quality is critical: meaningful analogues expose local GAT errors, whereas misleading neighborhoods can induce over-correction.The method’s correction behavior depends on whether retrieved molecules provide relevant local evidence.
  • Results from fixed benchmark splits, rather than repeated resampling protocols, support conclusions about the retrieval-plus-gating interface but limit broader interpretation.The reported evaluation design leaves robustness across alternative resampling procedures untested.
Loading 2608.30674v1…