Graph DiT

A polymer for gas separation has to be synthesizable, permeable to oxygen, and not permeable to nitrogen — three targets that must hold in the same molecule. Conditional graph diffusion models could be steered by one property at a time, so the working practice was to generate a batch per property and look for a structure near the top of every list. Across 100 test polymers, more than half had no structure shared by all three condition-specific candidate sets, and the median rank of the shared structure was 30 out of 30.

Graph DiT, from the University of Notre Dame group of Gang Liu, Jiaxin Xu, Tengfei Luo and Meng Jiang, makes the property targets part of the denoiser instead. A condition encoder converts a mixed set of numerical and categorical property values into one representation, and that representation replaces the mean and variance of every hidden layer through adaptive layer normalization, so the conditions shape the reverse diffusion trajectory directly rather than nudging it from outside. No separate property predictor is trained to guide sampling — guidance is predictor-free — which makes a new design brief a new set of inputs at sampling time, not a new training run. Posted in January 2024, the work was presented as an Oral at NeurIPS 2024, and its name follows Peebles and Xie's Diffusion Transformer, the vision architecture whose AdaLN conditioning it adapts to graphs.

It is the earliest of three graph diffusion transformers from the same lab: Llamole later embeds a pretrained Graph DiT as a frozen graph decoder behind Llama, Qwen and Mistral, and DemoDiff carries the line into in-context molecular design.

#Key Features

• Multi-property conditioning without a predictor: Numerical and categorical targets are encoded jointly and injected into the denoiser, so no auxiliary property model is needed to steer sampling.
• Clustering-based numerical encoding: Continuous property values pass through learnable centroids and a soft assignment rather than a linear projection, which controlled gas permeability better than direct or interval encodings in ablations.
• Graph-dependent noise model: Atoms and bonds are packed into one token matrix and corrupted together through a transition matrix built from atom–bond co-occurrence frequencies, rather than receiving independent noise as in earlier graph diffusion.
• Valid graphs without rule-based repair: Generation stays above 0.8 chemical validity with no rule checking, where several diffusion baselines fall below 0.6 once their post-processing is removed.
• Reusable as a frozen generator: A pretrained Graph DiT plugs unchanged into Llamole as its graph decoder behind three different language-model backbones.

#Technical Details

The denoiser is a stack of Transformer layers with AdaLN in place of layer norm and a final MLP predicting atom and bond type probabilities; the released default configuration is 6 layers, hidden size 1152, 16 attention heads, and discrete diffusion over 500 cosine-schedule steps with a guidance scale of 2, trainable on a single A6000 GPU. Training data spans a gas-separation polymer database with oxygen, nitrogen and carbon dioxide permeability conditions (553 molecules for the three-gas task, 609 for the O₂/N₂ design study) and class-balanced MoleculeNet subsets for drug design — BACE (1,332), BBBP (872) and HIV (2,372) — each pairing a categorical activity label with numerical synthetic-accessibility and complexity scores. Molecules are limited to 50 heavy atoms, and a separate model is trained per dataset. Evaluation generates 10,000 molecules per task and scores nine metrics against random-forest oracles: mean absolute error on the polymer permeability conditions averaged 17.8% lower than the strongest baseline, and accuracy on the categorical drug-design conditions exceeded 0.9 where baselines stayed under 0.6.

#Applications

The model suits inverse design where the brief is a set of numbers and labels rather than a scaffold: small molecules conditioned on HIV replication inhibition, blood–brain-barrier permeability or β-secretase 1 inhibition alongside a synthesizability budget, and polymer membranes conditioned on gas permeabilities. In the paper's design study it generated 1,000 polymers under O₂/N₂ conditions, improving average condition error by 9.9% over baselines; four polymer scientists then blind-ranked 25 candidates pooled across all methods, and the top three were all Graph DiT's. The torch-molecule package exposes it through a scikit-learn-style generator class.

#Impact

Graph DiT established that predictor-free guidance transfers from image diffusion to discrete molecular graphs, and that conditioning through layer statistics beats bolting a predictor onto a generator. Its clearest downstream evidence is architectural reuse: the same generator became the chemistry half of Llamole. The limits sit on the artifact rather than the idea — no trained checkpoint is released for the paper's own tasks, so the repository ships datasets and splits and expects training from scratch on a few hundred to a few thousand molecules, and every reported property value is scored by a random-forest oracle rather than measured, with the only human evaluation on the polymer side.