Budgeted Quotient-Residual Guidance for Frozen Pocket-Conditioned Molecular Diffusion
Authors: Xinyu Wang, Jinbo Bi, Minghu Song
Organizations: Department of Computer Science and Engineering, University of Connecticut Storrs, CT 06269, USA · Institute of Health and Medicine, Hefei Comprehensive National Science Center Hefei 230601, China
Pocket-conditioned molecular diffusion updates ambient atom coordinates, but many lead-optimization objectives are expressed on quotient features such as distances, contacts, and anchored substructures. We introduce budgeted quotient-residual guidance (QRG), an inference-time correction that makes these quotient objectives active without retraining the molecular generator. QRG lifts quotient covectors to metric-horizontal ambient directions and delivers them through a trust budget set by the frozen sampler's own step norm: quotient geometry chooses the direction, while sampler motion bounds the scale. We derive the horizontal lift, closed-form sampler-budget update, KL/kinetic interpretation around a frozen reverse step, equivariance conditions, and a product-budget split for budget-capped section and residual controls. Controlled quotient tasks confirm that sampler-relative delivery activates signals that raw local quotient gradients leave dormant. On frozen TargetDiff backbones, official seed-0 CBGBench ligand-generation/editing sweeps show practical quality-runtime gains: Local-QRG improves validity from 0.815 to 0.864 on fragment growing, 0.664 to 0.707 on scaffold hopping, and 0.681 to 0.712 on linker design, while PredNext-QRG improves fragment/scaffold and remains near-neutral on linker. Novelty remains 1.000 and diversity is preserved in the matched multi-seed molecular slice, giving task-dependent improvements without sampler retraining or backbone modification. Overall, QRG provides a lightweight route to quotient-aware inference for frozen molecular samplers with explicit runtime accounting.
Figures & tables
Figure 1: Budgeted quotient-residual guidance. A frozen molecular sampler moves ambient coordinates, while many scientific objectives live on quotient features such as distances, contacts, and fixed anchors. Raw quotient directions can be dormant at sampler scale. QRG keeps the quotient-residual direction but delivers it through a sampler-relative trust budget. The right panel gives the deployed Budgeted QRG step: section and quotient-residual controls are budget-capped separately by the frozen sampler’s own motion. Deployed variants use PredNext or local-radius quotient features; rollout residuals are teachers and diagnostics, not deployment requirements.
Figure 2: Controlled quotient tasks. Dormant local guidance is almost identical to the base sampler, while trust-region delivery makes the same quotient signal active under the same denoising schedule.
Task
Method
Success ↑
Key error ↓
NFE
Gaussian 2D
Base
0.159
0.528
64
Dormant local
0.159
0.529
64
Budgeted residual
0.246
0.451
64
Orbit points
Base
0.117
0.406
72
Dormant local
0.117
0.406
72
Budgeted residual
0.310
0.299
72
Table 1: Toy-task summary. The key error is quotient MMD for Gaussian 2D and invariant template error for the geometry tasks. All rows within a task use the same denoising budget.
Task
Method
Official parts
Validity ↑
Δ val.
sec/sample
Control/base
Novelty ↑
Diversity ↑
Fragment
Base
5/5
0.815
–
16.68
0.000
1.000
0.898
PredNext-QRG
5/5
0.829
+0.014
18.28
0.192
1.000
0.899
Local-QRG
5/5
0.864
+0.049
19.91
0.265
1.000
0.896
Scaffold
Base
5/5
0.664
–
17.99
0.000
1.000
0.913
PredNext-QRG
5/5
0.683
+0.019
20.36
0.194
1.000
0.911
Local-QRG
5/5
0.707
+0.043
20.18
0.268
1.000
0.912
Table 2: Budgeted inference on frozen ligand-generation/editing tasks. Validity, runtime, and control/base ratio are official seed-0 aggregates; novelty and diversity are from the matched multi-seed molecular slice for the same task and method. Fragment and scaffold improve for both deployed variants. Linker design is near-neutral for PredNext-QRG but improves for Local-QRG , showing that the quality–runtime trade-off is task-dependent rather than uniformly monotone.
Figure 3: Official-run norm accounting for the trust-region mechanism. Left: delivered control stays well below the frozen sampler’s own step scale across official seed-0 ligand-generation/editing shards. Right: for Local-QRG , the delivered section and quotient-residual components spend a small bounded budget relative to the base sampler motion. This is a kinetic-budget diagnostic, not a measured path KL: it supports the mechanism claim that QRG delivers quotient information at sampler scale rather than applying an unbounded external optimizer.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Expanded quality–runtime trade-off from the matched multi-seed molecular slice. Fragment has the richest ablation set; scaffold and linker include base, PredNext-QRG , and Local-QRG . This figure complements the official seed-0 aggregate by showing novelty and diversity preservation across matched seeds.
Cheap surrogate
Validity ↑
Unique
Diversity
sec/sample
Sampled pairs, 16
0.888
854
0.896
17.84
Sampled pairs, 128
0.886
856
0.897
18.15
Local radius
0.902
869
0.896
16.88
Appendix
Table 3: Development-slice fragment seed-0 ablation for cheap quotient surrogates. The local-radius surrogate is not introduced as an arbitrary extra branch: on this slice it outperforms sampled-pair residual surrogates at comparable runtime while preserving diversity.
Task
Method
n
Mean Δ
Median Δ
95% bootstrap CI
Wins/losses/ties
Fragment
PredNext-QRG
61
+0.013
+0.010
[−0.006,0.032]
32/27/2
Fragment
Local-QRG
61
+0.049
+0.030
[0.019,0.078]
42/14/5
Scaffold
PredNext-QRG
64
+0.019
+0.040
[−0.029,0.061]
44/17/3
Scaffold
Local-QRG
64
+0.043
+0.060
[−0.016,0.099]
43/19/2
Linker
PredNext-QRG
43
-0.001
+0.020
[−0.037,0.033]
22/19/2
Linker
Local-QRG
43
+0.030
+0.040
[−0.012,0.073]
25/14/4
Appendix
Table 4: Paired official seed-0 target-level validity deltas. Each row compares a guided method to the unguided sampler on the same targets; bootstrap intervals resample targets. These statistics support the aggregate table while showing target heterogeneity.
Task
Method
n
Mean Δ
Median
P10
P25
P75
P90
Min/Max
Fragment
PredNext-QRG
61
+0.013
+0.010
-0.060
-0.020
+0.050
+0.100
-0.160/+0.210
Fragment
Local-QRG
61
+0.049
+0.030
-0.060
+0.000
+0.100
+0.190
-0.280/+0.400
Scaffold
PredNext-QRG
64
+0.019
+0.040
-0.097
-0.010
+0.103
+0.167
-0.860/+0.330
Scaffold
Local-QRG
64
+0.043
+0.060
-0.246
-0.023
+0.190
+0.271
-0.860/+0.520
Linker
PredNext-QRG
43
-0.001
+0.020
-0.158
-0.035
+0.060
+0.130
-0.340/+0.270
Linker
Local-QRG
43
+0.030
+0.040
-0.148
-0.050
+0.130
+0.156
-0.260/+0.440
Appendix
Table 5: Distribution of per-target official seed-0 validity deltas versus the unguided sampler. Each target contributes one target-level validity difference for the same task and method. The quantiles show that aggregate gains coexist with heterogeneous pockets, including a small number of large negative and large positive target-level shifts.
Fragment variant
Seeds
Validity ↑
Runtime
Control/base
Base
3
0.842
1.00×
0.000
Section-only
1
0.896
1.07×
–
Sham-QRG
3
0.888
1.22×
0.255
PredNext-QRG
3
0.896
1.10×
0.193
Local-QRG
3
0.905
1.22×
0.271
Rollout-QRG Teacher
3
0.898
2.15×
0.253
Appendix
Table 6: Fragment ablation frontier. Rollout-QRG Teacher is useful as a diagnostic upper-cost point but is not the deployed method. The comparison separates the core trust-region delivery rule from rollout-heavy or unstructured guidance controls.
Figure 5: Side-chain decoration exposes a task-dependent budget mismatch. Default ligand-generation budgets reduce side-chain validity relative to the unguided sampler. A sparse low-budget rescue nearly closes the gap, but does not beat unguided in the full sweep.
Task / method
Paired targets
Mean Δ Vina ↓
Better targets
Fragment / PredNext-QRG
30
0.061[−0.049,0.174]
14/30
Fragment / Local-QRG
30
−0.101[−0.189,−0.023]
17/30
Scaffold / PredNext-QRG
44
−0.132[−0.234,−0.035]
28/44
Scaffold / Local-QRG
44
−0.285[−0.397,−0.178]
35/44
Linker / PredNext-QRG
42
−0.001[−0.087,0.081]
20/42
Linker / Local-QRG
42
0.038[−0.055,0.130]
17/42
Appendix
Table 7: Corrected raw-pocket docking diagnostic. Values are paired mean differences versus the base sampler with bootstrap intervals. Negative means the guided variant has better raw AutoDock Vina score. The signal is clearest on scaffold hopping and mixed on the other tasks.
Evaluation asset
Completed rows
Generated samples
Role in the paper
Official ligand-generation/editing sweeps
9
50,400
Main validity, runtime, and control-budget evidence for fragment growing, scaffold hopping, and linker design in table 2 .
Official section-only controls
3
16,800
Residual-off controls used to separate section delivery from quotient-residual guidance.
Official side-chain sweep
4
25,600
Out-of-family budget-transfer stress test; not counted as a ligand-generation win claim.
Matched multi-seed molecular slice
40
–
Novelty/diversity preservation, development ablations, rollout-teacher diagnostics, and sham-control comparisons.
Target-balanced docking diagnostic
6
–
Supporting raw-pocket AutoDock Vina diagnostic with paired target comparisons; not a binding-affinity claim.
Official full aggregate total
16
92,800
Completed official seed-0 aggregate rows across all four molecular tasks and four reported method families.
Appendix
Table 8: Completed evaluation inventory. The table records what was completed and how each asset is used in the paper. It is an accounting table, not an additional performance ranking.
Task
Method
Generated
sec/sample
Approx. GPU h
Control/base
Fragment
Base
6,100
16.68
28.3
0.000
Fragment
PredNext-QRG
6,100
18.28
31.0
0.192
Fragment
Local-QRG
6,100
19.91
33.7
0.265
Scaffold
Base
6,400
17.99
32.0
0.000
Scaffold
PredNext-QRG
6,400
20.36
36.2
0.194
Scaffold
Local-QRG
6,400
20.18
35.9
0.268
Appendix
Table 9: Compute accounting for the official seed-0 ligand-generation/editing rows that carry the main molecular claim. Approximate GPU hours are computed from generated samples times aggregate wall-clock seconds/sample divided by 3600, and therefore include denoising, guidance, postprocessing, and file writing as recorded by the runner.
Method family
Guidance schedule
Guidance deliveries/sample
Denoiser-query interpretation
Base
Frozen reverse sampler only
0
No guidance branch and no extra denoiser rollout.
Section-only
Section branch, residual budget zero
500
Tests delivery without quotient-residual correction.
PredNext-QRG
Predicted-next residual every two denoising steps
250
Reuses the sampler’s predicted next state; deployed point avoids rollout-teacher queries.
Local-QRG
Local-radius quotient surrogate with trust-region delivery
500
Uses a cheap local surrogate rather than querying a rollout teacher.
Rollout-QRG Teacher
One shared rollout before residual evaluation
500
Diagnostic high-cost teacher/oracle point, not the deployed method.
Appendix
Table 10: Denoiser-query and guidance-call accounting. The table separates deployed QRG variants from rollout-teacher diagnostics and explains why the reported runtime overhead is modest relative to a rollout-heavy inference scheme.
Permutation invariance is fundamental in molecular point-cloud generation, yet most diffusion models enforce it indirectly via permutation-equivariant networks on an ordered space. We propose to model diffusion directly on the quotient manifold \calX~=\sRd×N/SN, where all atom permutations are identified. We show that the heat kernel on \calX~ admits an explicit expression as a sum of Euclidean heat kernels over permutations, which clarifies how diffusion on the quotient differs from ordered-particle diffusion. Training requires a permutation-symmetrized score involving an intractable sum over SN; we derive an expectation form over a posterior on permutations and approximate it using MCMC in permutation space. We evaluate on unconditional 3D molecule generation on QM9 under the EQGAT-Diff protocol, using SemlaFlow-style backbone and treating all variables continuously. The results demonstrate that quotient-based permutation symmetrization is practical and yields competitive generation quality with improved efficiency.
Gyeonghoon Ko, Juho Lee
Korea Advanced Institute of Science and Technology
Diffusion-based generative models have reformed generative AI, and also enabled new capabilities in the science domain, e.g., fast generation of 3D structures of molecules. In such tasks, there is often a symmetry in the system, identifying elements that can be converted by certain transformations as equivalent. Equivariant diffusion models guarantee a symmetric distribution, but miss the opportunity to make learning easier, while alignment-based simplification attempts fail to preserve the target distribution. In this work, we develop quotient-space diffusion models, a principled generative framework to fully handle and leverage symmetry. By viewing the intrinsic generation process on the quotient space, the exact construction that removes symmetry redundancy, the framework simplifies learning by allowing model output to have an arbitrary intra-equivalence-class movement, while generating the correct symmetric target distribution with guarantee. We instantiate the framework for molecular structure generation which follows SE(3) (rigid-body movement) symmetry. It improves the performance over equivariant diffusion models and outperforms alignment-based methods universally for small molecules and proteins, representing a new framework that surpasses previous symmetry treatments in generative models.
Yixian Xu, Yusong Wang, Shengjie Luo +4
State Key Laboratory of General Artificial Intelligence, Peking University, Beijing, China · State Key Laboratory of Human-Machine Hybrid Augmented Intelligence, Institute of Artificial Intelligence and Robotics, Xi’an Jiaotong University · Zhongguancun Academy, Beijing, China +2
Consider an auto-regressive model that produces outputs x (e.g., answers to questions, molecules) each of which can be summarized by an attribute vector y (e.g., helpfulness vs. harmlessness, or bio-availability vs. lipophilicity). An arbitrary reward function r(y) encodes tradeoffs between these properties. Typically, tilting the model's sampling distribution to increase this reward is done at training time via reinforcement learning. However, if the reward function changes, re-alignment requires re-training. In this paper, we show that a reward weighted classifier-free guidance (RCFG) can act as a policy improvement operator in this setting, approximating tilting the sampling distribution by the Q function. We apply RCFG to molecular generation, demonstrating that it can optimize novel reward functions at test time. Finally, we show that using RCFG as a teacher and distilling into the base policy to serve as a warm start significantly speeds up convergence for standard RL.