SSD: Shell-Guided Spherical Diffusion for Molecular Geometry Generation

Chuang, Yun-Yen; Gu, Chen-Sheng; Hsu, Hung-Min; Lin, Kevin; Chang, Ray-I

How it works

Initialize. Corrupt. Generate.

SSD on the left, a standard Gaussian prior on the right, at the same scale. Drag to rotate.

SSD · shell prior

Gaussian prior

v_rad v_rep v_score

Toy 14-atom molecule for intuition, not model output. Both panes use the same scale.

\[ d\mathbf{x}_t^{(i)} = \big[\,{\color{#ffc14d}\mathbf{v}_{\mathrm{rad}}^{(i)}} + {\color{#b48cff}\mathbf{v}_{\mathrm{rep}}^{(i)}} + {\color{#3ee6c1}\mathbf{v}_{\mathrm{score}}^{(i)}}\,\big]\,dt + \sigma_t\, d\bar{\mathbf{w}}_t^{(i)} \]

Radial attraction · same speed α_t for every atom

Short-range repulsion · keeps atoms ≥ d_min apart

SE(3) correction · the backbone's own score network

Training

Still trained, but it converges faster

The backbone is trained as usual. Only its noisy inputs change: they are drawn in closed form around SSD's straight-line forward path instead of Gaussian noise. The loss form, optimizer, and training cost stay the same. Training curves converge faster and higher (Fig. 2a–d).

Sampling

A shell-aware hybrid sampler

Generation starts from a fresh shell sample and follows v_rad + v_rep + v_score. The atom-to-point pairing π is unknown at inference, so v_rad pulls toward the centre. This is a hybrid sampler, not an exact time reversal of the forward process.

Results

One change. Five backbones. Consistent gains.

backbones improved: GeoDiff, SubGDiff, EDM, MCF, Flow

extra learnable parameters in the diffusion backbones, with unchanged training cost

COV-R on Drugs molecules with ≥ 50 atoms

Backbone + SSD MCF · Drugs gain is not statistically significant (Table 3)

Full tables, ablation & trajectory statistics

Each backbone is evaluated under its own protocol with identical training and sampling budgets. Mean values are shown. "–" means the metric is not reported in the original paper. SSD entries are single runs (seed 42); see Tables 1 and 3 in the paper for medians and three-seed statistics. The MCF gain on Drugs (+0.76 COV-R) is sign-consistent but not statistically significant. SSD-Flow also swaps SemlaFlow's encoder for a lighter EDM-style one (Appendix M.1).

Conditional generation · SubGDiff protocol

Model	Data	COV-R ↑	COV-P ↑	MAT-R ↓	MAT-P ↓
GeoDiff	QM9	80.36	53.66	0.2820	0.6673
SSD-GeoDiff	QM9	92.00	61.40	0.2425	0.4188
GeoDiff	Drugs	64.12	52.79	–	–
SSD-GeoDiff	Drugs	91.69	70.23	0.8463	1.1543
SubGDiff	QM9	90.91	50.16	0.2460	0.6114
SSD-SubGDiff	QM9	93.20	63.50	0.2380	0.4100
SubGDiff	Drugs	76.16	–	1.0003	–
SSD-SubGDiff	Drugs	92.50	71.10	0.8300	1.1400

Conditional generation · MCF protocol

Model	Data	COV-R ↑	COV-P ↑	MAT-R ↓	MAT-P ↓
MCF	QM9	95.00	93.70	0.1030	0.1190
SSD-MCF	QM9	96.15	94.61	0.0993	0.1087
MCF	Drugs	84.70	66.80	0.3900	0.6180
SSD-MCF	Drugs	85.37	68.82	0.3792	0.6089

Unconditional refinement · SemlaFlow protocol (%)

Model	Data	AtomStab ↑	MolStab ↑	Valid ↑	Unique ↑
EDM	QM9	98.7	82.0	91.9	98.9
SSD-EDM	QM9	99.3	99.8	100.0	100.0
EDM	Drugs	81.3	–	–	–
SSD-EDM	Drugs	90.5	70.1	98.5	99.2
SemlaFlow	QM9	99.9	99.7	99.4	95.4
SSD-Flow	QM9	99.9	99.9	100.0	100.0
SemlaFlow	Drugs	99.8	97.3	93.9	100.0
SSD-Flow	Drugs	99.8	97.6	100.0	100.0

Ablation · the shell and the dynamics are both required (QM9, SubGDiff protocol)

Variant	Init. / Prior	Dynamics	COV-R ↑	COV-P ↑	MAT-R ↓	MAT-P ↓
(1) Gaussian baseline	Gaussian	v_SubGD	90.91	50.16	0.2460	0.6114
(2) Gaussian + rep./corr.	Gaussian	v_rep + v_score	90.95	50.20	0.2455	0.6098
(3) Shell(r_Gauss)	Shell(r_Gauss)	None	90.92	50.18	0.2462	0.6107
(4) Shell(r_Gauss) + SSD dyn.	Shell(r_Gauss)	v_rad + v_rep + v_score	91.50	51.00	0.2443	0.6020
(5) SSD (ours)	Shell(r_SSD)	v_rad + v_rep + v_score	93.20	63.50	0.2380	0.4100

Trajectory stability (QM9)

Method	Regime	Radial dev. ↓	Path excess ↓
SubGDiff	SDE	0.4534	246.91
SSD-SubGDiff	SDE	0.1781	1.71
SemlaFlow	ODE	0.2145	1.29
SSD-Flow	ODE	0.1096	1.11

SSD vs Gaussian-based diffusion — **SSD vs. Gaussian diffusion.** A direct path from a chemically scaled shell, versus a long walk from a dimension-scaled prior.

Convergence, sampling-step robustness and long-range drift on QM9 — **Faster convergence, robust to sampling steps, stable under long-range drift** (QM9). Orange = SSD-SubGDiff, blue = SubGDiff, red = SSD-Flow, green = SemlaFlow.

Qualitative case studies of generated conformations — **Case studies.** On long-drift and large (> 21 atoms) molecules, SSD stays chemically plausible where SubGDiff distorts.

Abstract

Diffusion models for 3D molecular geometry almost always use an isotropic Gaussian prior. Its Frobenius-norm scale σ_T√3n is set by the noise level and atom count, not by chemistry. The result is a scale mismatch that grows with n and produces high spatial entropy and unstable early trajectories. We introduce Shell-guided Spherical Diffusion (SSD), a model-agnostic framework that replaces the Gaussian prior with a chemically scaled spherical-shell initialization. It adds coordinated radial attraction, short-range repulsion, and an SE(3)-equivariant correction field to both the forward and reverse processes. Initialization and dynamics have to be designed together: neither a shell alone nor radial fields alone reproduces SSD's stability or accuracy. We evaluate SSD on five coordinate-space backbones (GeoDiff, SubGDiff, EDM, SemlaFlow-style flow matching, and MCF), each under its own evaluation protocol. With identical training and sampling budgets, SSD consistently improves both quality and diversity. It lifts weaker backbones such as GeoDiff and EDM to match or surpass stronger diffusion-, VAE-, and flow-based baselines. SSD is therefore a plug-in geometric enhancement for coordinate-space molecular generation models that changes neither their architectures, loss functions, nor training pipelines.

Shell-Guided Spherical Diffusion