arXiv:2603.12344cs.LG2026-03

用决策树规则指导大模型,提升分子属性预测精度

Can Decision Trees Teach Large Language Models? Distilling Verbalized Knowledge for Molecular Property Prediction

  • 将决策树的预测规则转化为自然语言注入提示词
  • 在TDC基准上22个任务中显著提升大模型性能
  • 适合药物研发中需要通用分子预测的场景

分子属性预测(MPP)是药物发现中的基础问题,近年来受到广泛关注。大语言模型(LLMs)虽具备跨领域表现力,但在MPP上的性能仍不足以实际应用。为此,我们提出TreeKD方法,将基于决策树的专业模型知识蒸馏到LLMs中,以补充其内部知识并提升预测准确率。针对每个属性,利用输入分子中40,000个功能基团的特征训练专用决策树;随后将决策树学习到的预测规则进行自然语言化,并融入训练提示词中。此外,通过用随机森林替代单一决策树,引入一种测试时扩展技术——规则一致性,通过聚合不同规则生成的多个提示所对应的预测结果来增强鲁棒性。在TDC基准数据集上,对Gemma-2-2B和Granite-3.3-2B两个LLM进行22个任务的全面评估,结果表明该方法显著提升了LLMs的性能,推动了通用分子预测模型的发展。

原文摘要 · Abstract (English)

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.

分子预测知识蒸馏大模型决策树

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。