PaperlayerView exact revision on arXiv

Source-linked AI summary

Graph Diffusion Transformers are In-Context Molecular Designers

Gang Liu, Jie Chen, Yihan Zhu, Michael Sun, Tengfei Luo, Nitesh V Chawla, Meng Jiang

arXiv:2510.08744v1cs.LGcs.AI

TL;DR

Molecular design has many assay and material-property tasks but sparse labels, while language-model in-context learning has shown limited success for molecular design. DemoDiff conditions graph diffusion on molecule–score demonstrations and uses motif-level tokenization for pretraining; across 33 design tasks, it matches or surpasses much larger LLMs and ranks higher than specialized methods.

  • Problem

    Molecular design spans millions of assays and property types, many with few labeled examples, and prior in-context learning has been less effective than specialized molecular methods.

    Source: §1 · p. 1

  • Method

    DemoDiff uses molecule–score demonstrations as task contexts for a denoising graph Transformer and employs Node Pair Encoding for motif-level molecular representation.

    Source: §1 · pp. 1–2, §3.1 · p. 4

  • Results

    Across 33 design tasks, DemoDiff matches or surpasses LLMs 100–1000× larger and achieves rank 3.6 versus 5.25–10.20 for ten specialized models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors present DemoDiff as a promising molecular foundation model for in-context design, with potential to scale further using larger models.

    Source: §6 · p. 9

    The in-context learning framing assumes demonstrations reflect a latent concept θ from a family of concepts.

    Source: §2.1 · p. 2

Abstract

from arXiv · show

In-context learning allows large models to adapt to new tasks from a few demonstrations, but it has shown limited success in molecular design. Existing databases such as ChEMBL contain molecular properties spanning millions of biological assays, yet labeled data for each property remain scarce. To address this limitation, we introduce demonstration-conditioned diffusion models (DemoDiff), which define task contexts using a small set of molecule-score examples instead of text descriptions. These demonstrations guide a denoising Transformer to generate molecules aligned with target properties. For scalable pretraining, we develop a new molecular tokenizer with Node Pair Encoding that represents molecules at the motif level, requiring 5.5$\times$ fewer nodes. We curate a dataset containing millions of context tasks from multiple sources covering both drugs and materials, and pretrain a 0.7-billion-parameter model on it. Across 33 design tasks in six categories, DemoDiff matches or surpasses language models 100-1000$\times$ larger and achieves an average rank of 3.63 compared to 5.25-10.20 for domain-specific approaches. These results position DemoDiff as a molecular foundation model for in-context molecular design. Our code is available at https://github.com/liugangcode/DemoDiff.

1 INTRODUCTION

Molecular design spans many assay and material-property tasks, but examples are sparse and language-model in-context learning has been less effective than specialized methods. DemoDiff uses molecule–score demonstrations to condition graph diffusion, with motif-level tokenization and broad pretraining supporting its reported cross-task performance.

1 INTRODUCTION

  • Molecular design spans millions of biological assays and material properties, many with only a few labels, while prior molecular ICL has been less effective than specialized methods.Some specialized methods rely on extensive Oracle calls or large labeled datasets.Source: §1 · p. 1
  • DemoDiff defines task contexts with molecule–score demonstrations rather than a large property embedding, and conditions a denoising Transformer on them to generate target-aligned molecules.The demonstrations include positive, medium, and negative examples to represent task concepts more completely.Source: §1 · pp. 1–2
  • Node Pair Encoding represents molecules at the motif level, reducing average node counts by 5.5×, while the pretraining dataset combines drug and materials tasks.The model has 0.7B parameters and was pretrained on over 1.6 million tasks from 155K unique properties.Source: §1 · p. 2
  • A 0.7B-parameter DemoDiff matches or surpasses LLMs 100–1000× larger and ranks 3.6 against 5.25–10.20 for ten specialized models.The evaluation covers 33 design tasks across six categories.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries frame in-context learning as inference over latent task concepts from demonstrations, then describe graph diffusion as a way to model molecular structures under property conditions.

2.1 IN-CONTEXT LEARNING WITH DEMONSTRATIONS

  • In-context learning treats demonstrations as evidence of a latent concept θ, which a model marginalizes over to generate an output for a query.The paper applies this framing to molecular design, where molecule–score pairs represent task concepts and the query is a target score.Source: §2.1 · pp. 2–3

2.2 MOLECULAR DESIGN WITH GRAPH DIFFUSION TRANSFORMERS

  • Molecules are discrete graphs of atoms and bonds, represented in Graph DiTs by tokens combining atom features with their bond connections.Discrete diffusion corrupts these graph features in a forward process and progressively removes noise in the reverse process.Source: §2.2 · p. 3
  • Graph DiTs condition molecular denoising on property values, while this work explores using demonstrations as the conditions for in-context learning.The model is trained to predict the original graph from its noisy state and property conditions.Source: §2.2 · p. 3

3 LEARNING DIFFUSION MODEL WITH DEMONSTRATIONS

DemoDiff combines motif-level graph tokenization with demonstration-conditioned diffusion to generate molecules for new task contexts. The section describes its tokenizer, context-based inference, consistency scoring, and large-scale pretraining setup.

3 LEARNING DIFFUSION MODEL WITH DEMONSTRATIONS

  • DemoDiff’s reverse diffusion process starts from noise and denoises molecules conditioned on molecule–score demonstrations represented at the motif level.A tokenizer bridges motif and atom representations.Source: §3 · p. 3, Figure 2 · p. 4

3.1 MOLECULAR GRAPH TOKENIZATION WITH NODE PAIR ENCODING

  • Node Pair Encoding builds a motif vocabulary by repeatedly merging adjacent substructures and selecting the most frequent candidate motif in the dataset.The resulting motif graph uses directed edges with bond-type and attachment information, enabling reconstruction of the atom-level graph.Source: §3.1 · pp. 3–4
  • Constrained NPE incorporates frequent rings as complete motifs to avoid reconstruction ambiguity, while rare rings retain atom-level representations.The method integrates ring constraints during vocabulary initialization and motif merging.Source: §3.1 · p. 4
  • Motif tokenization compresses molecules by an average ratio of 5.446 ± 2.569, reducing the median representation from 30 atoms to 5 motifs.A 38-atom example is represented using four motifs.Source: §3.1 · p. 4

3.2 IN-CONTEXT LEARNING WITH GRAPH DIFFUSION TRANSFORMERS

  • DemoDiff replaces property conditions with a context of molecule–score pairs and a query score, using the demonstrations to infer a latent task concept during diffusion.The model can assess generated molecules with a fingerprint-based consistency score against positive, medium, and negative demonstration groups.Source: §3.2 · pp. 4–5

3.3 MODEL DESIGN AND PRETRAINING

  • The model encodes scaled property scores as positional signals and uses motif-token limits to fit the target and demonstrations into context.Because context molecules are structurally disjoint, their edge connectivity marks context boundaries without explicit delimiter tokens.Source: §3.3 · p. 5
  • Pretraining combines ChEMBL assays and polymer properties into about 1.6 million tasks over 155K unique properties, then trains a 0.7B-parameter model.Tasks assign the target a score of 1 and divide up to 15 context examples per group into positive, medium, and negative score ranges.Source: §3.3 · p. 5

4 EXPERIMENT

The experiments compare DemoDiff with molecular-design baselines across 33 tasks and examine its performance across task types, model sizes, demonstration contexts, and case studies. Results show strong performance overall, with clearer advantages on property-driven tasks and evidence that context and consistency choices affect generation.

4 EXPERIMENT

  • The evaluation covers 33 downstream tasks across six categories and compares DemoDiff with 13 baselines using ten valid, unique, and novel generated molecules per task.The reported harmonic mean combines averaged oracle scores and structural diversity.Source: §4 · p. 6

4.1 PERFORMANCE ON DIVERSE MOLECULAR DESIGNS TASKS

  • DemoDiff achieves an average rank of 3.62, outperforming GraphGA (rank 5.25) across six task categories while producing accurate scores and diverse molecules.Its designs score closer to the query and have better structural diversity than LLM-based ICL designs, which often produce structurally similar molecules.Source: §4.1 · pp. 6–7
  • DemoDiff scores 0.67–0.79 on drug and material design, versus around 0.44–0.56 on rediscovery and structure-constrained tasks.It still ranks highly on structure-constrained tasks; the authors attribute its stronger property-driven results to broader solution ranges in chemical-space exploration.Source: §4.1 · p. 7

4.2 ABLATION STUDIES AND PERFORMANCE ANALYSIS

  • At 739M parameters, scaling benefits become more evident; at 311M, performance improves in most tasks except drug design.The tested sizes are 78.7M, 311M, and 739M parameters; Table 2 reports top-10 harmonic means across six categories.Source: §4.2 · p. 7, Table 2 · p. 7
  • Longer contexts and diverse positive, medium, and negative demonstrations support better ICL performance; positive examples alone are insufficient.Fewer positive examples can still yield reasonable results, and the case study examines inference from only negative examples.Source: §4.2 · p. 7
  • Consistency-based confidence filtering improves performance by 0.8% to 27.5%, although fingerprint consistency correlates only moderately with oracle scores on explicit structure-constrained tasks.For drug MPO and materials design, high fingerprint consistency with positive examples does not always indicate high target scores; filtering may still remove false positives.Source: §4.2 · pp. 7–8, Figure 5 · p. 8

4.3 CASE STUDIES

  • With only negative demonstrations (scores < 0.5), DemoDiff can generate desirable candidates for a target score of 1 across structure-constrained, drug MPO, and target-based design.The candidates align with the task concept while remaining structurally distinct from negative examples.Source: §4.3 · p. 8, Figure 6 · p. 8
  • In Albuterol rediscovery, the diffusion trajectory’s similarity score rises from 0.22 initially to 0.74 as successive steps refine the molecule toward the target structure.The Albuterol Oracle measures similarity to the ground truth.Source: §4.3 · p. 8, Figure 7 · p. 9

5 RELATED WORK

Prior molecular optimization methods include search algorithms and graph diffusion, but many require costly oracle calls or task-specific training. Research on in-context learning also examines how model, data, and learning mechanisms shape its performance.

5 RELATED WORK

  • Molecular optimization benchmarks may require about 10,000 oracle calls, while trained diffusion models are limited to specific tasks and need hundreds of labeled molecules.LLM-based chemical foundation models have focused on other molecular tasks or relied on fine-tuning within a restricted scope of design tasks.Source: §5 · p. 9
  • In-context learning research spans model, data, and learning-mechanism perspectives, including implicit Bayesian inference and how task distributions shape learning strategies.The cited work describes a balance between storing common task information in weights and developing in-context learning from rare tasks.Source: §5 · p. 9

6 CONCLUSION

DemoDiff combines demonstration-conditioned molecular diffusion with motif-level graph modeling, and the paper presents it as a promising molecular foundation model. The appendices describe its graph-token losses, diffusion transitions, consistency scoring, and pretraining context construction.

6 CONCLUSION

  • DemoDiff matches or outperforms much larger LLMs and ranks higher than domain- and task-specific methods, supporting its promise as a molecular foundation model.The conclusion also identifies larger models, broader datasets, and greater compute as potential routes for further scaling.Source: §6 · p. 9

A.1 GRAPH-LEVEL TOKENS AND PRETRAINING LOSS

  • Graph-token features comprise motif types, bond types, and attachment specifications, with pretraining loss decomposed into motif, bond, and attachment losses.The feature dimension aligns tokens using dense edges; null bonds are included in bond loss, while null-edge attachments are excluded from attachment loss.Source: §A.1 · p. 12

A.2 TRANSITION MATRICES IN DIFFUSION MODELS

  • The motif-level transition matrix combines node-to-node, edge-to-node, edge-to-edge, and node-to-edge transitions, while diffusion perturbs bond types but predicts attachments directly.Transitions use training-data marginal and co-occurrence distributions with a cosine noise schedule.Source: §A.2 · pp. 12–13

A.3 DETAILS ON CONSISTENCY SCORE

  • Selecting the top 100 of 1,000 generated molecules by consistency score improves empirical results across task categories before Oracle evaluation.The score favors molecules whose structural similarities follow the positive, medium, and negative demonstration ordering; Table 3 reports average Top-10 harmonic scores.Source: §A.3 · p. 13, Table 3 · p. 13

B DETAILS ON PRETRAINING

  • The pretraining set contains 1,084,566 molecules or polymers, 155,150 unique assays or properties, and 1,639,515 tasks, each with a target score fixed at 1.Up to 45 demonstrations are grouped into positive, medium, and negative examples, with half the context window allocated to positive examples and one quarter each to the other groups.Source: B DETAILS ON PRETRAINING · p. 13

B.1 PRETRAINING DATASET

The pretraining dataset builds molecule-score contexts from biological assays and polymer properties. Both sources form positive, medium, and negative demonstrations relative to a target molecule.

ChEMBL dataset

  • ChEMBL activity contexts use pChEMBL values, selecting targets with pChEMBL > 6 and grouping same-assay molecules by normalized activity distance.The positive, medium, and negative groups correspond to activity distances [0, 0.25), [0.25, 0.5), and [0.5, 1.0], with up to 15 molecules sampled per group (Zdrazil et al., 2024).Source: §B.1 · p. 14

Polymeric materials

  • Polymer contexts span thermal, electronic, structural, and transport properties from PolyInfo, MSA, and Kuenneth et al. (2021).Values are min–max normalized, with logarithmic transformation for ranges exceeding 1000 and non-negative shifts when needed; up to 15 candidates per context group are sampled by distance.Source: §B.1 · p. 14

B.2 TOKENIZER PREPARATION

Node Pair Encoding (NPE) builds motif vocabularies by merging neighboring molecular motifs, reducing graph node counts relative to atom-level representations. The selected vocabulary size balances node-count reduction against further gains.

B.2 TOKENIZER PREPARATION

  • NPE initializes atom- and ring-based motifs, repeatedly merges structurally compatible neighbors, and adds the most frequent merged motif to the vocabulary.The procedure treats rings as units under its structural constraints and stops when the vocabulary reaches its maximum size.Source: §B.2 · pp. 14–15
  • When K ≥ 3000, mean and maximum node counts no longer significantly decrease and the median remains unchanged, motivating K = 3000 with K_ring = 300.The analysis varies motif vocabulary size and reports mean, maximum, and median node counts.Source: Figure 8 · p. 14, §B.2 · p. 15
  • The compression ratio ranges from 1 to 40, with median and mean around 5.5, indicating roughly fivefold graph-size reduction.The atom- and motif-level node-count distributions are heavy-tailed; Figure 10 defines compression ratio as uncompressed graph size divided by compressed graph size.Source: §B.2 · p. 15
  • For molecules with 150 to 200 atoms, motif representations reduce node counts by up to a factor of 15, and compression ratios increase mildly with molecule size.This comparison relates atom-level and motif-level representations.Source: §B.2 · p. 15

B.3 MODEL PRETRAINING

DemoDiff was pretrained as a 0.7B-parameter model, with training monitored through reconstruction and generation metrics. The authors identify larger models, more data, and additional compute as possible routes to improvement.

B.3 MODEL PRETRAINING

  • Training stopped after reconstruction accuracy exceeded 0.99 for every component, due to resource limits; generated molecules reached about 0.83 validity and 0.69 MACCS similarity.The 0.7B-parameter model trained for 550 epochs over 49 days on 2–4 H100 GPUs, or about 146 GPU days.Source: §B.3 · p. 15
  • The reported improvements in loss, reconstruction accuracy, validity, and similarity suggest that larger models, more data, and additional compute could further improve DemoDiff.The training trends are described as continuing to show incremental gains near the end of training.Source: §B.3 · pp. 15–16

C.1 DETAILS ON EXPERIMENTAL SET-UPS

The experiments compare DemoDiff with molecular optimization, conditional-generation, and language-model baselines across 33 tasks in six molecular-design categories. They define task-specific demonstrations, oracle-based evaluation, and a shared performance-and-diversity measure.

C.1 DETAILS ON EXPERIMENTAL SET-UPS

  • 33 tasks across six categories are evaluated against 13 baselines: eight molecular optimization methods, two conditional generators, and three LLMs.Table 4 reports benchmark task statistics; the categories cover drug rediscovery, drug MPO, structure-constrained design, drug design, target-based design, and material design.Source: §C.1 · p. 16, Table 4 · p. 17
  • Task oracles score similarity-based rediscovery and MPO, SMARTS- or formula-constrained structures, ML-scored drug design, docking targets, and ML-estimated polymer gas selectivity.MPO adds constraints such as logP, TPSA, and Bertz; target-based docking scores are mapped to [0, 1].Source: §C.1 · pp. 18–19
  • Each task has up to 450 molecule–score pairs split among positive, medium, and negative groups, with a target score of 1 and an evaluation-only oracle except for optimization methods.Some tasks have fewer examples when positive examples are insufficient.Source: §C.1 · p. 19
  • Molecular optimization is tested with either 100 oracle calls or 10,000 predictor calls, using a random forest trained on all 450 task pairs for predictor-based evaluation.The authors note that laboratory oracle calls may take days to months, making large oracle budgets impractical in real-world settings.Source: §C.1 · p. 19
  • Each method generates 100 valid, unique, novel molecules per task, and performance is the harmonic mean of top-10 oracle score and diversity.DemoDiff first generates 1,000 candidates and selects the top 100 by consistency with the positive-to-negative context ordering.Source: §C.1 · pp. 19–20
  • DemoDiff uses a 150-motif-token context averaging 23 demonstrations, while LLM comparisons use 12 positive, 6 medium, and 6 negative examples.DemoDiff’s average context is half positive, one quarter medium, and one quarter negative; the LLMs include DeepSeek-V3, GPT-4o, and Qwen-Max, with up to hundreds of billions of parameters.Source: §C.1 · p. 19

C.2 ADDITIONAL DISCUSSION OF EXPERIMENTAL RESULTS

The additional results discussion finds that predictor quality, not query count alone, matters for molecular optimization, and that cross-method performance alignment may partly signal task difficulty. Case studies also show motif-level contexts fitting large molecules that exceed the atom-level context window.

C.2 ADDITIONAL DISCUSSION OF EXPERIMENTAL RESULTS

  • All molecular optimization methods score zero on Valsartan SMARTS when limited-data predictors cannot model latent constraints, including multiple SMARTS patterns and physicochemical properties.By contrast, more predictor calls improve performance on target-based tasks where training data are sufficient.Source: §C.2 · p. 20, Table 14 · p. 25
  • Tasks where molecular optimization performs well, such as target-based and material design, are also more tractable for ICL; this alignment may partially reflect task difficulty.The authors note that task difficulty is challenging to quantify formally because it depends on oracle definitions and data quality.Source: §C.2 · p. 20
  • For DRD2 and JNK3, molecular optimization under limited supervision underperforms while DeepSeek-V3, GPT-4o, and DemoDiff achieve strong results.These comparisons are reported in Table 7.Source: §C.2 · p. 21, Table 7 · p. 21

C.3 DETAILS ON CASE STUDIES

  • With 17, 23, and 26 negative demonstrations, three case-study molecules of 460, 631, and 481 atoms exceed a 150-token atom-level window, while motif representations encode the demonstrations within it.The cases cover isomer generation, Osimertinib MPO, and PARP1 target-based design.Source: §C.3 · p. 22, Figure 14 · p. 27, Figure 15 · p. 28
Loading 2510.08744v1…