Readout head, generated molecules and diagnostics for training-free guidance of
heat capacity $C_v$ on QM9, steering a frozen unconditional EDM
(Hoogeboom et al.) with the TFG (Ye et al. 2024) evaluation protocol.
A small MLP head is trained on frozen EGNN activations of the diffusion
backbone to predict $C_v$ at any noise level. At sampling time the head's
gradient w.r.t. the noisy state $z_t$ steers DDIM. The diffusion model is never
fine-tuned.
Results (4096 generated molecules, TFG protocol)
Validity↑
Uniqueness↑
Novelty↑
Mol. stab.↑
Atom stab.↑
MAE↓
Unguided floor ($s=0$, $\eta=1$)
92.09%
99.31%
82.59%
81.23%
98.26%
2.8908
This head, $s=8$, $\eta=1$
88.21%
99.25%
83.68%
74.41%
96.26%
1.3126
This head, $s=12$, $\eta=1$
90.82%
99.35%
83.50%
78.52%
98.01%
1.2422
This head, $s=16$, $\eta=1$
89.40%
99.18%
83.48%
76.46%
97.76%
1.1660
All rows use $t_{\max}^{\text{norm}} = 0.7$ and 100 DDIM steps on the same
frozen EDM; only the guidance scale differs.
Published training-free baselines on the same task (MAE↓): DPS 5.26, LGD 3.77,
UGD 3.02, MPGD 2.86, FreeDoM 2.84, TFG 2.77.
Why the unguided floor matters
QM9's sampling protocol draws the atom count $n$ from $p(n)$ before generation
and then draws the conditioning target from $p(C_v \mid n)$. Because 83.6% of
$C_v$ variance is explained by atom count alone, any generator that merely
emits a valid molecule with the right $n$ scores well without steering $C_v$ at
all. Measured two ways, that floor is 2.77 (analytic) / 2.854–2.891
(this sampler at $s=0$).
Every published baseline sits at that floor. The right diagnostic is therefore
not MAE but whether the generated residual tracks the requested residual,
$r = C_v - \mathbb{E}[C_v \mid n]$:
run
MAE
corr($r_y$, $r_c$)
slope
unguided, $\eta=0$
2.854
+0.004
+0.004
unguided, $\eta=1$
2.933
−0.039
−0.043
$s=12$, $\eta=1$
1.242
+0.835
+0.960
$s=16$, $\eta=1$
1.166
+0.836
+0.985
A slope near 1 means the generated heat-capacity residual moves essentially 1:1
with the requested one.
Contents
head/ Cv readout head (the one used for all reported results)
best.pt state dict + config + per-t val MAE
config.json training hyperparameters
mae_per_t.json val MAE (cal/mol·K) at t = 0, 50, 100, 250, 500, 750, 999
head_residual/ size-factored variant (predicts Cv - E[Cv|n]); see below
samples/ 4096 generated molecules per run + metrics
FINAL_raw_s16_t0.7_eta1/ best MAE
FINAL_raw_s12_t0.7_eta1/ recommended operating point (+ grad_stats.json)
FINAL_FLOOR_s0_eta1/ unguided control
diagnostics/
prop_stats.json floors and atom-count dependence for Cv and alpha
steer_final.json residual correlation / slope per run
sweep_summary.json full (s, t_max, eta, loss, grad_mode) sweep at N=1024
code/ training, sampling and diagnostic scripts
samples/*/molecules.pkl is a list of 4096 tuples
(one_hot [29,5], charges [29,0], x [29,3], node_mask [29,1], target [1]),
where target is the requested $C_v$ in cal/mol·K, thermochemically
subtracted (QM9's subtract_thermo convention, as used by EDM/TFG).
Head
3 taps from EGNN blocks 2 / 5 / 8 (hidden_nf=256) → Linear(256→128) per tap,
masked sum-pool over atoms → concat → LayerNorm-SiLU-Linear fuse → FiLM on a
sinusoidal timestep embedding → 4-layer MLP → scalar. 379,649 parameters.
Trained on the QM9 second-half split (the same data the EDM saw), L1 loss on the
mean/MAD-normalized label, uniform $t$, AdamW lr 2e-3, batch 512, 40 epochs,
cosine schedule with warmup.
Validation MAE (cal/mol·K; $C_v$ MAD = 4.89):
t
0
50
100
250
500
750
999
MAE
0.272
0.201
0.263
0.896
1.797
1.939
1.975
At $t=999$ the state is pure noise, so 1.975 is what the atom count alone buys —
the gap down to 0.272 at $t=0$ is the size-independent signal that guidance uses.
head_residual/ regresses $(C_v - \mathbb{E}[C_v \mid n]) / \text{MAD}_{\text{resid}}$
instead, with sum+mean pooling and an explicit atom-count embedding, so the
head's output is only the steerable part. It steers harder per unit $s$ but
degrades molecule quality faster; the plain head reaches better MAE at matched
stability, so it is the one used for the reported numbers.
Usage
bash
1# Sample 4096 molecules at the recommended operating point2python code/sample_cv.py \3 --readout_ckpt head/best.pt \4 --edm_ckpt models/pretrained_models/EDMsecond/generative_model_ema.npy \5 --edm_args models/pretrained_models/EDMsecond/args.pickle \6 --num_samples 4096 --batch_size 256\7 --guidance_scale 12 --t_max_norm 0.7 --eta 1.0 --run_eval
89# Retrain the head10python code/train_cv.py --tag raw --target_mode raw \11 --epochs 40 --batch_size 512 --lr 2e-3
Requires the TFG codebase on PYTHONPATH (for tasks.networks.egnn,
tasks.networks.qm9, evaluations.molecule) plus the EDM second-half
checkpoint and the evaluate_Cv oracle predictor from the TFG release.
Note: TFG's QM9 loader downloads from springernature.figshare.com, which now
returns HTTP 202 with an empty body; ndownloader.figshare.com/files/3195389
(data) and .../3195404 (uncharacterized list) still work.
Backbone and protocol: Hoogeboom et al., Equivariant Diffusion for Molecule
Generation in 3D (EDM); Ye et al. 2024, TFG: Unified Training-Free Guidance;
Satorras et al., E(n) Equivariant Graph Neural Networks.
