DemoDiff: In-Context Molecular Design

• The paper introduces DemoDiff, a demonstration-conditioned diffusion model that leverages molecule-score pairs to design molecules via in-context learning.
• It employs a Graph Diffusion Transformer architecture with motif-level representations to efficiently generate molecules with desired properties.
• Extensive evaluations on drug, docking, and polymer design tasks show that DemoDiff outperforms traditional methods and scales to many design properties.

Searching arXiv for DemoDiff and closely related molecular design papers to ground the article. arxiv_search(query="DemoDiff demonstration-conditioned diffusion molecular design Graph Diffusion Transformer", max_results=10, sort_by="relevance") Searching arXiv for the specific DemoDiff paper by title and related graph diffusion molecular design work. arxiv_search(query="Graph Diffusion Transformers are In-Context Molecular Designers", max_results=10, sort_by="relevance") DemoDiff is a demonstration-conditioned diffusion model for in-context molecular design that defines tasks through a small context of molecule-score examples rather than textual descriptions or fixed property identifiers (Liu et al., 9 Oct 2025). Introduced as a Graph Diffusion Transformer, it generates molecules from a conditional distribution $p_\theta(X \mid \mathcal{C}, Q)$, where $\mathcal{C}$ is a set of demonstration pairs and $Q$ is a query score, typically set to 1 to denote the ideal target (Liu et al., 9 Oct 2025). The method is framed for the regime in which chemical and materials databases contain very large numbers of assays and properties, while each individual assay or property typically has only a small number of labeled molecules.

1. Problem setting and conceptual basis

Inverse molecular design seeks molecules that satisfy a desired property profile such as bioactivity, ADMET constraints, or materials properties. In DemoDiff, the central difficulty is that existing resources such as ChEMBL contain millions of assays or properties, but each individual property often has only a small number of labeled examples, making per-task training impractical (Liu et al., 9 Oct 2025). The method therefore adopts an in-context learning formulation: instead of training a separate model for each property, a single pretrained model receives a small set of molecule-score pairs that define a new task and must generate molecules consistent with that task.

The demonstrations are explicit task descriptors. A context is written as $\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$, where $X_{ij}$ are molecules and $Y_{ij} \in [0,1]$ are normalized scalar scores indicating relative desirability; DemoDiff then samples from

$p_\theta(X \mid \mathcal{C}, Q).$

The scores are interpreted as relative distances to the target score 1, and contexts are organized into positive, medium, and negative examples using score ranges $[0.75,1]$, $(0.5,0.75]$, and $[0,0.5]$, respectively (Liu et al., 9 Oct 2025).

This formulation is explicitly contrasted with three alternatives. Text-conditioned generative models are limited because many assays and materials properties have no meaningful natural-language description beyond an identifier and a numeric field. Property-prediction-plus-optimization pipelines require substantial labeled data per task and often rely on $\mathcal{C}$0–$\mathcal{C}$1 oracle calls. Conditional diffusion models with fixed property vectors do not scale naturally to $\mathcal{C}$2–$\mathcal{C}$3 properties and generalize poorly to unseen properties because conditioning is tied to indices or embeddings rather than example-based descriptions (Liu et al., 9 Oct 2025). A plausible implication is that DemoDiff treats heterogeneous design tasks—bioactivity, docking, and polymer selectivity—as instances of a single demonstration-conditioned generation problem.

2. Diffusion formulation and Graph Diffusion Transformer architecture

DemoDiff is built on Graph Diffusion Transformers and uses discrete diffusion over motif-level molecular graphs rather than atom-level strings or text sequences (Liu et al., 9 Oct 2025). A molecule is represented as a graph $\mathcal{C}$4, but the model operates on a motif graph $\mathcal{C}$5, where $\mathcal{C}$6 is a set of disjoint motifs and $\mathcal{C}$7 is a set of directed motif-to-motif edges with bond types and attachment specifications. The denoising objective follows a conditional diffusion loss of the form

$\mathcal{C}$8

with factorization over motif types, bond types, and attachment variables (Liu et al., 9 Oct 2025).

The forward diffusion process is discrete and defined through transition matrices. DemoDiff specifies a joint graph transition matrix $\mathcal{C}$9 for nodes and edges, with separate motif, bond, and co-occurrence transitions, and uses a cosine noise schedule

$Q$0

Motif and bond types are diffused, whereas attachment specifications are predicted during denoising rather than diffused directly (Liu et al., 9 Oct 2025).

Conditioning is implemented by packing the noisy target graph and all demonstration molecules into a single token sequence. The model therefore performs joint processing of the target and demonstrations, with graph connectivity serving implicitly to distinguish disjoint molecular components; no explicit delimiter tokens are required (Liu et al., 9 Oct 2025). Timestep $Q$1 is encoded and added to token embeddings, while demonstration scores $Q$2 and query score $Q$3 are embedded using Rotary Position Embedding, treating scores as continuous positions along a one-dimensional axis from negative to positive examples (Liu et al., 9 Oct 2025). This architecture allows attention both within each graph and across graphs, so the Transformer can transfer task information from demonstrations to the target denoising trajectory.

The paper interprets this behavior through the lens of in-context learning as implicit Bayesian inference: $Q$4 where $Q$5 denotes a latent task concept. In DemoDiff, the latent concept is not provided symbolically; it is inferred from the pattern of demonstration graphs and scores (Liu et al., 9 Oct 2025).

3. Node Pair Encoding and motif-level graph representation

A central enabling component is Node Pair Encoding, a molecular tokenizer that converts atom-level graphs into motif graphs (Liu et al., 9 Oct 2025). The tokenizer starts from a motif vocabulary containing 118 atom types plus a polymerization point “$Q$6”, which guarantees that any molecule can still be represented at atom level in the worst case. It then iteratively merges frequently co-occurring neighboring motifs, with explicit handling of ring structures through a constrained merging procedure. The reported vocabulary choice is $Q$7 motifs with $Q$8 (Liu et al., 9 Oct 2025).

The representation gain is substantial. The abstract states that Node Pair Encoding requires 5.5$Q$9 fewer nodes, and the detailed analysis reports an average compression ratio of $\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$0, with the median node count moving from approximately 30 atoms to approximately 5 motifs (Liu et al., 9 Oct 2025). The pretraining dataset exhibits compression ratios ranging from 1 to 40, with a median of approximately 5.5. Increasing the vocabulary size beyond 3000 yields diminishing returns in node count reduction (Liu et al., 9 Oct 2025).

The motif graph itself is exact in the sense that the tokenizer provides both $\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$1 and $\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$2, mapping atom-level molecules to motif graphs and back (Liu et al., 9 Oct 2025). Each graph-level token contains motif identity together with outgoing bond and attachment information. This tokenization is significant because the context window is defined in motif tokens rather than atoms; shorter sequences permit larger demonstration sets within a fixed Transformer budget. This suggests that NPE is not merely a compression heuristic but a structural precondition for graph-based in-context learning at practical context lengths.

4. Pretraining corpus, task construction, and inference procedure

DemoDiff is pretrained at substantial scale. The reported corpus contains 1,084,566 molecules or polymers, 155,150 unique assays or properties, and 1,639,515 tasks, derived from ChEMBL v35 and multiple polymer datasets (Liu et al., 9 Oct 2025). Task construction is based on anchor molecules or anchor polymers, with candidate examples ranked by normalized distance in property space and partitioned into positive, medium, and negative groups; up to 15 demonstrations are sampled per group, producing at most 45 demonstrations per pretraining task before truncation to the motif-token context length (Liu et al., 9 Oct 2025).

Three parameter scales are reported: 78.7M, 311M, and 739M. The largest model, DemoDiff-0.7B, uses 24 Transformer layers, hidden size 1280, 16 attention heads, and MLP ratio 4, and is pretrained for 550 epochs, requiring approximately 49 days on 2–4 H100 GPUs, or about 146 H100 GPU-days (Liu et al., 9 Oct 2025). Training continues until reconstruction accuracy of motifs, bonds, and attachments exceeds 0.99; reported curves also show generation validity reaching approximately 0.83 and structural similarity to ground truth approximately 0.69 (Liu et al., 9 Oct 2025).

Inference on a new task follows a fixed sequence. Molecule-score demonstrations are normalized to $\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$3, encoded into motif graphs, and packed into the context window, with roughly half positive, one quarter medium, and one quarter negative when available (Liu et al., 9 Oct 2025). The query score is set to $\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$4. Diffusion begins from a noisy motif graph drawn from the stationary distribution and iteratively denoises to a candidate molecule, which is then decoded back to atom level.

Candidate selection is refined by a context consistency score that compares a generated molecule’s similarity to the positive, medium, and negative context groups. Using groupwise Tanimoto similarities, DemoDiff defines a difference-based score

$\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$5

where each margin measures whether the generated molecule is more similar to positives than to medium examples, and more similar to medium than to negatives (Liu et al., 9 Oct 2025). In practice, the model generates 1000 candidates, retains the top 100 by consistency score, and then evaluates those with the oracle. The reported effect is an improvement in top-10 harmonic scores of up to 27.5% on some categories (Liu et al., 9 Oct 2025).

5. Evaluation protocol, empirical performance, and ablations

The evaluation spans 33 design tasks in six categories: drug rediscovery, drug multi-objective optimization, structure-constrained design, drug design, target-based docking design, and polymer gas separation (Liu et al., 9 Oct 2025). Performance is assessed on 100 valid, unique, and novel molecules per task, using the average of top-10 oracle scores, internal diversity, and their harmonic mean; methods are then ranked per task and aggregated by average rank and sum of ranks (Liu et al., 9 Oct 2025).

Across all 33 tasks, DemoDiff attains an average rank of 3.63. The reported comparison figures are 5.25 for GraphGA under 100 oracle calls, 7.34 for GraphGA with predictor calls, 6.91 for a conditional Graph DiT baseline, 6.76 for DeepSeek-V3, 6.34 for GPT-4o, and 11.56 for Qwen-Max (Liu et al., 9 Oct 2025). The abstract additionally states that DemoDiff matches or surpasses LLMs 100–1000$\mathcal{C} = \{(X_{ij}, Y_{ij})\}_{j=1}^{L}$6 larger and that its 0.7-billion-parameter model outperforms or matches specialized baselines on average rank (Liu et al., 9 Oct 2025).

The reported category-level strengths are especially clear in drug design, target-based design, and materials design, where DemoDiff achieves harmonic scores of 0.79, 0.78, and 0.67, respectively (Liu et al., 9 Oct 2025). Top-1 results also show it achieving or matching best single-molecule performance in multiple drug-design and polymer-selectivity tasks. The paper positions these outcomes as evidence that demonstration-conditioned graph diffusion can function as a molecular foundation model rather than as a narrow task-specific optimizer (Liu et al., 9 Oct 2025).

Ablation studies probe scaling, context composition, and filtering. Performance improves with model size, with the 739M model showing a clear gain in drug-design harmonic score relative to 78M and 311M models (Liu et al., 9 Oct 2025). Increasing context length improves harmonic scores in Albuterol rediscovery, and mixed contexts containing positive, medium, and negative examples outperform positive-only contexts, indicating that negative and medium examples help define the task boundary (Liu et al., 9 Oct 2025). Even in tasks with only negative demonstrations, DemoDiff still generates high-scoring molecules in structure-constrained design, Osimertinib MPO, and PARP1 docking, which the paper interprets as evidence that the model has learned how negative examples relate to optimal ones for similar tasks (Liu et al., 9 Oct 2025). Diffusion-trajectory analysis in Albuterol rediscovery shows similarity improving from 0.22 at initial noise to 0.74 at the final step (Liu et al., 9 Oct 2025).

6. Scope, limitations, and nomenclature

The explicit term “DemoDiff” refers to the demonstration-conditioned molecular design model introduced in “Graph Diffusion Transformers are In-Context Molecular Designers” (Liu et al., 9 Oct 2025). This point matters because nearby literature contains superficially similar names with different meanings. “DEMO: Disentangled Motion Latent Flow Matching for Fine-Grained Controllable Talking Portrait Synthesis” concerns audio-driven talking portraits and states that “DemoDiff” is not a term used by its authors; its method is DEMO and it uses flow matching rather than the demonstration-conditioned graph diffusion framework described here (Chen et al., 12 Oct 2025). “Restoring Noisy Demonstration for Imitation Learning With Diffusion Models” likewise states that its official method name is DMDR rather than DemoDiff (Chen et al., 16 Oct 2025). A common misconception is therefore to treat “DemoDiff” as a generic label for any diffusion model involving demonstrations; in the arXiv literature represented here, the explicit name denotes the molecular design model.

The paper also identifies several limitations. Pretraining depends on ChEMBL and existing polymer datasets, so performance may degrade on chemistries or properties far outside that distribution, such as exotic inorganic clusters or complex organometallics (Liu et al., 9 Oct 2025). The fixed context window and tokenizer still impose upper bounds on demonstration count and molecule size. Oracle quality remains a bottleneck because docking and learned predictors are imperfect surrogates for experimental behavior. Interpretability is limited, since the task concept is encoded implicitly in Transformer attention rather than in an explicit symbolic program. Pretraining cost is high, at approximately 146 H100 GPU-days for the 0.7B model (Liu et al., 9 Oct 2025).

The paper also notes dual-use concerns. A strong generative model for molecular design could be misused to design harmful chemicals, including toxins or illicit drugs, and responsible deployment should therefore involve task restrictions, safety filters, and human oversight (Liu et al., 9 Oct 2025). Within the research landscape, the broader implication is that in-context learning can emerge in graph diffusion transformers when pretraining exposes the model to sufficiently diverse tasks with Zipf-like frequency structure, rather than being confined to text-only LLMs (Liu et al., 9 Oct 2025).