javid48/readout-guidance-qm9-cv

Readout Guidance — QM9 heat capacity ($C_v$)

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

# Sample 4096 molecules at the recommended operating point
python code/sample_cv.py \
  --readout_ckpt head/best.pt \
  --edm_ckpt  models/pretrained_models/EDMsecond/generative_model_ema.npy \
  --edm_args  models/pretrained_models/EDMsecond/args.pickle \
  --num_samples 4096 --batch_size 256 \
  --guidance_scale 12 --t_max_norm 0.7 --eta 1.0 --run_eval

# Retrain the head
python code/train_cv.py --tag raw --target_mode raw \
  --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.

Reproducing the floor

python code/sample_cv.py --readout_ckpt head/best.pt --guidance_scale 0 --eta 1.0 --run_eval
python code/steer_diag.py --runs samples/FINAL_raw_s16_t0.7_eta1 samples/FINAL_FLOOR_s0_eta1
python code/prop_stats.py     # analytic floors for Cv and alpha

Citation

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.