Molecular Property Prediction (MPP) is a fundamental problem in drug discovery that has recently attracted growing attention. Large Language Models (LLMs), known for their impressive proficiency across domains, show promise as generalist models for MPP. However, their current performance remains below the threshold needed for practical adoption. To bridge this gap, we propose TreeKD for distilling the knowledge of tree-based specialist models into LLMs to complement the internal knowledge of LLMs and improve their predictive accuracy. For each property, we train a specialist decision tree using features derived from 40K functional groups in the input molecules. Then, the predictive rule learned by the decision tree, which encodes its knowledge, is verbalized and incorporated into the prompts for training LLMs. In addition, by replacing a single decision tree with a Random Forest, we introduce a test-time scaling technique called rule-consistency, which aggregates predictions generated from different prompts constructed with different rules. An extensive evaluation with two LLMs, Gemma-2-2B and Granite-3.3-2B, on the TDC benchmark with 22 prediction tasks shows that our method substantially enhances the performance of LLMs, advancing the development of generalist models for MPP.
Figures & tables
Figure 1: An example illustrates that the LLM sees a molecule as a whole, while a decision tree operates on each dimension of functional groups.
Figure 2: An overview of our proposed method, consisting of 4 stages: ① Extracting FGs, ② Building Specialist Models, ③ Distilling Knowledge to LLMs, and ④ Test-Time Scaling.
Figure 3: An example of rule verbalization. The decision tree is turned into a hierarchy of if-then conditions, where indentation indicates the depth of nodes.
Type
Property
Metric
Specialist Models
TxGemma-2B
Gemma-2-2B
MapLight
Decision Tree
Naive
TreeKD
w/ self-consistency
w/ rule-consistency
Reg
Caco-2 Permeability
MAE (↓)
0.269
0.413
0.629
0.547
0.476 ( + 0.071 )
0.551 0.008 ( - 0.075 )
0.528 0.005 ( - 0.052 )
Lipophilicity
MAE (↓)
0.547
0.848
0.598
0.672
0.579 ( + 0.093 )
0.634 0.009 ( - 0.055 )
0.600 0.006 ( - 0.021 )
Solubility
MAE (↓)
0.793
1.771
0.950
2.921
0.921 ( + 2.000 )
0.982 0.019 ( - 0.061 )
0.868 0.025 ( + 0.053 )
PPBR
MAE (↓)
7.654
11.492
9.986
7.556
7.503 ( + 0.053 )
8.380 0.168 ( - 0.877 )
7.459 0.106 ( + 0.044 )
VDss
Spearman (↑)
0.722
0.414
0.387
0.522
0.647 ( + 0.125 )
0.602 0.013 ( - 0.045 )
0.675 0.007 ( + 0.028 )
Table 1: The results on the TDC benchmark. ( green ) and ( red ) values for TreeKD indicate performance gaps relative to naive training; boldface denotes better results in a direct comparison with TxGemma2-2B. In the 2 rightmost columns, ( green ) and ( red ) values indicate performance gaps relative to TreeKD for test-time scaling techniques; underlining denotes better results between them. The subscripts report standard deviations across 3 runs.
Figure 4: An example of a complete prompt on Caco-2 Permeability. All numerical values are rounded to 6 decimal places.
Type
Property
Gemma-2-2B
Naive
Naive
Naive
w/ FGs
w/ FGs
w/ a predictive rule
Reg
Caco-2 Permeability
0.547
+ 0.055
+ 0.016
Lipophilicity
0.672
- 0.005
+ 0.098
Solubility
2.921
+ 2.012
- 0.012
Table 2: The performance gains on the TDC benchmark with evolved prompt variants. ( green ) and ( red ) values indicate performance gaps relative to the preceding one.
Type
Property
Specialist Models
Gemma-2-2B
Decision Tree
TreeKD
d = 3
d = 6
d = 3
d = 6
Reg
Caco-2 Permeability
0.566
0 0.413 ( + 0.153 )
0.442
0.476 ( - 0.034 )
Lipophilicity
0.913
0 0.848 ( + 0.065 )
0.636
0.579 ( + 0.057 )
Solubility
1.863
0 1.771 ( + 0.092 )
0.949
0.921 ( + 0.028 )
PPBR
11.734
11.492 ( + 0.242 )
7.827
7.503 ( + 0.324 )
Table 3: The results of TreeKD on the TDC benchmark using maximum depths ( d ) of 3 and 6. ( green ) and ( red ) values indicate performance gaps when using d=3 relative to using d=6 .
Figure 5: The distribution of 10498 molecules in the test set of 13 classification tasks from the TDC benchmark, according to whether they are correctly predicted by the decision tree and the LLMs.
Figure 6: The distribution of 55 pairs of molecules exhibiting activity cliffs that the decision tree fails to distinguish, according to whether both molecules in each pair are correctly predicted by the LLMs.
Figure 7: A few examples of pairs of molecules exhibiting activity cliffs on Ames Mutagenicity (differences between them are highlighted in red ), showing whether both molecules in each pair are correctly predicted.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Type
Property
Dataset Size
Dataset Split
Reg
Caco-2 Permeability
906
Scaffold
Lipophilicity
4200
Scaffold
Solubility
9982
Scaffold
PPBR
1797
Scaffold
VDss
1130
Scaffold
Half Life
667
Scaffold
Appendix
Table 4: Detailed information and statistics of datasets.
Type
Property
Gemma-2-2B
N = 50
N = 100
self-*
rule-*
self-*
rule-*
Reg
Caco-2 Permeability
0.551
0.528
0.537 ( + 0.014 )
0.520 ( + 0.008 )
Lipophilicity
0.634
0.600
0.627 ( + 0.007 )
0.595 ( + 0.005 )
Solubility
0.982
0.868
0.978 ( + 0.004 )
0.869 ( - 0.001 )
PPBR
8.380
7.459
7.913 ( + 0.467 )
7.329 ( + 0.130 )
Appendix
Table 5: The results of TreeKD on the TDC benchmark with test-time scaling techniques using N=50 and N=100 . Note that * stands for consistency. underlining denotes better results between techniques when using the same N . ( green ) and ( red ) values indicate performance gaps when using N=100 relative to using N=50 .
Type
Property
Random Forest
N = 50
N = 100
Reg
Caco-2 Permeability
0.355
0 0.354 ( + 0.001 )
Lipophilicity
0.842
0 0.841 ( + 0.001 )
Solubility
1.702
0 1.703 ( - 0.001 )
PPBR
11.043
11.123 ( - 0.080 )
VDss
0.489
0 0.491 ( + 0.002 )
Appendix
Table 6: The results of Random Forest on the TDC benchmark using N=50 and N=100 . ( green ) and ( red ) values indicate performance gaps when using N=100 relative to using N=50 .
Type
Property
Metric
Gemma-2-2B
Naive
TreeKD
Reg
Caco-2 Permeability
MAE (↓)
000 5.847
000 0.607 ( 000 + 5.240 )
Lipophilicity
MAE (↓)
000 1.687
000 0.981 ( 000 + 0.706 )
Solubility
MAE (↓)
00 42.173
000 1.937 ( 00 + 40.236 )
PPBR
MAE (↓)
00 58.217
00 10.364 ( 00 + 47.853 )
VDss
MAE (↓)
0 796.001
00 11.765 ( 0 + 784.236 )
Appendix
Table 7: The results on the TDC benchmark (on regression tasks) in the zero-shot setting. ( green ) and ( red ) values for TreeKD indicate performance gaps relative to naive inference.
Type
Property
Metric
Gemma-2-2B
Naive
TreeKD
Reg
Caco-2 Permeability
MAE (↓)
000 2.803
000 0.593 ( 000 + 2.210 )
Lipophilicity
MAE (↓)
000 1.214
000 0.957 ( 000 + 0.257 )
Solubility
MAE (↓)
000 2.567
000 1.930 ( 000 + 0.637 )
PPBR
MAE (↓)
00 13.330
00 10.031 ( 000 + 3.299 )
VDss
MAE (↓)
000 4.653
000 2.278 ( 000 + 2.375 )
Appendix
Table 8: The results on the TDC benchmark (on regression tasks) in the 3-shot setting. ( green ) and ( red ) values for TreeKD indicate performance gaps relative to naive inference.
Type
Property
Gemma-2-2B
Naive
Naive
Naive
w/ a predictive rule
w/ FGs
w/ a predictive rule
Reg
Caco-2 Permeability
0.547
+ 0.014
+ 0.057
Lipophilicity
0.672
- 0.026
+ 0.119
Solubility
2.921
- 0.084
+ 2.084
Appendix
Table 9: The performance gains on the TDC benchmark with evolved prompt variants. ( green ) and ( red ) values indicate performance gaps relative to the preceding one.
Figure 8: The inference times (in seconds) across different batch sizes.
Large language models (LLMs) are widely applied across chemical tasks, such as molecular property prediction, which underpins drug discovery. Molecular LLMs represent a molecule through several modalities, notably a 1D SMILES sequence or a 2D molecular graph. Both encode molecular information implicitly, so the contribution of individual substructures remains opaque. Retrieval and augmentation methods add context, but from external sources. However, the cues chemists reason over are the internal substructures that drive a property up or down. We propose MR-MoL, a multi-granular rationale-guided molecular LLM that supplies this evidence directly. A fine-tuned GNN scores each substructure through masking, and the most influential ones are serialized as a ranked, direction-tagged rationale that the LLM reads alongside the SMILES sequence and molecular graph. The rationale spans three levels of granularity: Murcko scaffolds with their side chains, BRICS fragments, and functional groups. This is, to our knowledge, the first method to expose GNN-derived attributions to an LLM as evidence for property prediction. On eight MoleculeNet tasks, MR-MoL achieves the best overall results among generalist models and narrows the gap to specialist models tuned for each task. Five diagnostics further confirm that the model reads the rationale rather than merely benefiting from its presence. Its direction, rank, and substructure each shape the prediction, and its attributions reproduce known structure-property relationships.
Large language models (LLMs) have shown promise for molecular property prediction, but their ability to reason over chemical structures remains limited, as molecular representations such as SMILES differ substantially from the natural language on which LLMs are primarily trained. To bridge this semantic and chemical knowledge gap, we propose MolE-RAG, a training-free, molecule-centric retrieval-augmented generation framework for LLM-based molecular property prediction. MolE-RAG augments each prediction with three complementary sources of inference-time context: retrieved chemistry literature, molecule-specific information including compound synonyms, identifiers, functional group annotations, and physicochemical descriptors, and structurally similar molecules retrieved from the training set. We evaluate MolE-RAG across nine molecular property prediction tasks using proprietary, chemistry-specialized, and open-source LLMs. Across general-purpose LLMs, MolE-RAG improves ROC-AUC by up to 28 percentage points on classification tasks and reduces regression RMSE by up to 67% relative to a SMILES-only baseline. We further find that the utility of each context source varies across models and tasks, with different models benefiting most from textual retrieval, molecular context, or structural retrieval. These results suggest that molecule-centric retrieval can improve LLM-based molecular property prediction without model fine-tuning while providing a flexible framework for integrating heterogeneous chemical knowledge at inference time.
Joey Chan, Wonbin Kweon, Ashley Shin +4
University of Illinois Urbana-Champaign · University of California, San Diego
Large language models (LLMs) are increasingly employed to predict molecular properties. However, prediction error alone cannot distinguish prediction from retrieval of published values. We audit 22 frontier models on 12 molecular regression datasets in a zero-shot setting, assessed against a molecule-blind reference derived from each dataset's labels. Significant retrieval is concentrated on 5 datasets, with isolated flagged LLMs elsewhere. Increasing the reasoning setting raises the number of flagged model--dataset combinations from 47 to 89 of 264. An in-context blinding experiment reduces retrieval but leaves a quarter of the combinations flagged. Blinding changes model rankings and increases errors. Because blinding also removes chemically interpretable structure, the error increase can only be partially attributed to reduced retrieval.
Matthias Busch, Marius Tacke, Sviatlana V. Lamaka +4
Institute for Artificial Intelligence and Simulation in Mechanics, Hamburg University of Technology, Eißendorfer Straße, 21073 Hamburg, Germany · Institute of Material Systems Modeling, Helmholtz-Zentrum Hereon, Max-Planck-Straße, 21502 Geesthacht, Germany · Institute of Surface Science, Helmholtz-Zentrum Hereon, Max-Planck-Straße, 21502 Geesthacht, Germany +3