|
Enhancing Diffusion-Based Sampling with Molecular Collective Variables |
Soundness: 4: excellent
Presentation: 3: good
Contribution: 3: good
Rating: 6: marginally above the acceptance threshold
Confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked. |
This paper proposes WT-ASBS (Well-Tempered Adjoint Schrödinger Bridge Sampler), a method that combines a neural diffusion-based sampler (ASBS) with a well-tempered bias potential defined in collective-variable (CV) space. The goal is to accelerate exploration in molecular systems by gradually flattening the free-energy landscape while retaining thermodynamic consistency through reweighting.
The authors claim that this hybrid design achieves faster and broader sampling compared to standard ASBS and conventional enhanced sampling methods such as WTMetaD. Empirical results on alanine dipeptide, alanine tetrapeptide, and chemical reaction benchmarks suggest improved exploration and free-energy reconstruction.
**Conceptually innovative hybridization.**
The integration of Schrödinger bridge–based diffusion sampling with a well-tempered metadynamics bias is novel. It bridges probabilistic transport and enhanced sampling, offering a unified view of energy-driven exploration in molecular systems.
**Clear motivation from sampling efficiency.**
The work addresses a concrete challenge in molecular generative modeling — slow mode discovery and poor coverage of rare events — and proposes a pragmatic biasing mechanism that is both simple and theoretically interpretable.
**Transparent discussion of limitations.**
The paper openly discusses the non-amortized bias, potential generalization issues, and differences from traditional WTMetaD, which reflects scientific maturity and honesty.
**1. Conceptual tension between learned sampler and non-learnable bias.**
Although the method is presented as a neural diffusion sampler, the core bias $V_{WT}$ is manually accumulated (non-neural). This design choice blurs whether the sampler truly learns the target distribution or merely follows a hand-crafted MetaD bias.
**2. Lack of fair and direct comparison with ASBS.**
Since WT-ASBS is an ASBS variant augmented with a well-tempered bias, the most meaningful benchmark is ASBS itself under identical conditions. However, ASBS only appears in Fig. 3c (explored states) and is not evaluated for PMF accuracy, making the incremental benefit of WT unclear.
**3. # of force-calls inconsistency between figures.**
Figure 3.c shows up to $\approx$2 M evaluations, while Figure 3.d extends to $\approx$40 M. Because WT-ASBS reaches chemical accuracy after $\approx$4–5 M evaluations, truncating Fig. 3c prevents fair comparison at later stages.
**4. Limited diffusion-based baselines.**
Comparisons are restricted to WTMetaD, which mainly highlights conceptual differences rather than validating the WT mechanism across diffusion methods. Adding at least one diffusion-based baseline (annealed, energy-guided, or score-guided) would strengthen the evaluation.
**5. Scalability and consistency.**
WT-ASBS excels on alanine dipeptide (Figure 2.f) but underperforms WTMetaD on tetrapeptide (Figure 3.d), suggesting potential sensitivity to CV dimensionality or generalizability to complex system.
**6. Lack of statistical robustness.**
The reported results appear to be based on a single simulation run for WT-ASBS. Given the stochastic nature of both diffusion and metadynamics sampling, multiple independent runs (with mean and variance) are crucial to assess convergence stability and reproducibility. Without them, it is difficult to judge whether the reported trajectories and PMF curves reflect consistent behavior or a favorable random seed.
(Q1) If $V_{WT}(\xi)$ is not parameterized by a neural network, how does the model ensure that the learned sampler contributes beyond a standard MetaD bias? Could the authors clarify whether WT-ASBS should be interpreted as (a) a neural sampler guided by an analytic bias, or (b) a learnable bias parameterized through the neural transport model?
(Q2) In Figure 8, the reweighted and bias-derived PMFs are nearly identical, suggesting that the sampling in CV space is almost uniform. If so, what ensures that the sampled structures remain physically valid? Could the authors provide representative molecular structures sampled from different regions of the CV space (e.g., high- vs low-free-energy regions) to illustrate what kind of configurations the model actually produces? This would clarify whether the method truly explores meaningful conformations or simply performs uniform random exploration over CVs.
(Q3) Could you extend Figure 3.c to the same energy-evaluation range ($\approx$10–40 M) as Figure 3.d to show late-stage exploration and better quantify the WT bias’s long-term effects? Please include an ASBS PMF accuracy curve for Ala2 system.
(Q4) Could you summarize cost–benefit metrics (energy evaluations, wall-clock time, ESS, PMF error) across ASBS, WT-ASBS, and WTMetaD for standardized comparison?
(Q5) ASBS’s slow exploration (Figure 3.c) might result from strong pre-training confinement near the initial mode (marked * in Figure 3.b).
Have you explored reducing pretraining strength or fine-tuning its duration to encourage broader exploration?
(Q6) When computing the free-energy difference, which states or basins were used? Were they predefined minima or clusters in CV space?
Clarifying this would help interpret the reported PMF results.
(Q7) Are all reported curves from single runs, or averaged over multiple simulations? If single, could you comment on the run-to-run variability? Showing variance bands or repeated trials would help establish statistical reliability. |
Fully AI-generated |
|
Enhancing Diffusion-Based Sampling with Molecular Collective Variables |
Soundness: 3: good
Presentation: 3: good
Contribution: 3: good
Rating: 6: marginally above the acceptance threshold
Confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. |
In this work, the authors present a collective-variable-based approach to enhance diffusion-based samplers. By incorporating a biasing force that is gradually deposited, similar to metadynamics, the approach improves the exploration ability of diffusion-based samplers. This addresses a consistent issue with diffusion-based samplers being susceptible to mode collapse.
The specific diffusion-based sampler enhanced using collective variables (CVs) is the Adjoint Schrödinger Bridge Sampler (ASBS), although the CV-based approach does not appear to be necessarily specific to this instance of diffusion-based sampler. The biasing along the CV is approached similarly to most biased enhanced-sampling methods: an additional bias potential is added to the base potential energy. This bias potential is then updated using the same rule as in well-tempered MetaD. This results in a two-step optimization process: an inner step, in which ASBS is performed (resulting in mode collapse), and an outer step, in which the bias potential is updated. In addition, a few practical improvements are implemented, such as the use of a replay buffer and a warm start.
The presented method is evaluated on a number of interesting systems. First, the evaluation considers two peptides, Alanine Dipeptide and Alanine Tetrapeptide, and, following this, the authors study reactive pathways in the form of a simpler SN2 reaction and a more complex reaction with multiple products.
The authors present an interesting and seemingly successful approach to alleviate the well-known issue of mode collapse in diffusion-based samplers within the domain of molecular conformations. The focus on ASBS as the diffusion-based sampler to enhance, and the choice of the well-tempered bias update, also seems the most sensible approach. Furthermore, the paper is well written and does an excellent job highlighting the core considerations that go into an enhanced-sampling method.
The experimental evaluation is relatively in-depth and focuses on two interesting problems by considering both simple peptides and reactions. While quantitative results are limited, the presented results suggest a significant improvement over standard WTMetaD.
All in all, the paper presents an excellent approach to enhance diffusion-based samplers to make them more appropriate for the equilibrium sampling of a molecular system. As such, I vote to accept the paper for publication.
- Looking at Figure 2, I’m surprised to see roughly uniform sampling from the second outer-loop iteration onwards (assuming that this is what the dots in panel c represent). I would have expected the samples to be located around the two smaller modes at this point, or still located at the original modes. Only at the later iterations would I expect to see roughly uniform sampling. Notably, this is also directly in contrast with the intuition presented in Figure 1. An author response to this inconsistency is required for me to increase my score beyond a borderline accept.
- There are a number of important components to the training setup that would be good to see studied more in-depth, such as the use of a replay buffer and the warm start. Most importantly, I am interested in seeing how these components influence the number of energy evaluations needed.
- Due to the reliance on collective variables, it is difficult to envision the applicability of the presented approach outside domains where CVs are known and provide a clear boundary to the sampling domain. For example, in areas such as the study of protein folding dynamics, it is hard to define CVs that sufficiently restrict the important sampling domain.
- As the work could potentially spark new interest in enhanced sampling, it would be good to see a more in-depth discussion of work in this area, such as umbrella sampling and adaptive biasing force methods. This could potentially be combined with the discussion of the requirements in Section 2.1, whose current purpose is not entirely clear. Additionally, it would be good to see additional works discussed, such as [1–4] (not a complete list), all of which use a similar biasing force to that proposed here in the context of enhanced sampling (but are more focused on transition path sampling).
[1] Seong, Kiyoung, et al. “Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers.” arXiv preprint arXiv:2405.19961 (2024).
[2] Singh, Aditya N., Avishek Das, and David T. Limmer. “Variational path sampling of rare dynamical events.” Annual Review of Physical Chemistry 76 (2025).
[3] Holdijk, Lars, et al. “Stochastic optimal control for collective-variable-free sampling of molecular transition paths.” Advances in Neural Information Processing Systems 36 (2023): 79540–79556.
[4] Du, Yuanqi, et al. “Doob’s Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling.” Advances in Neural Information Processing Systems 37 (2024): 65791–65822.
See weaknesses. |
Fully human-written |
|
Enhancing Diffusion-Based Sampling with Molecular Collective Variables |
Soundness: 4: excellent
Presentation: 4: excellent
Contribution: 3: good
Rating: 6: marginally above the acceptance threshold
Confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. |
This paper introduces WT-ASBS, a diffusion-based generative sampler enhanced by biasing along CVs, inspired by WTMetaD.
The method replaces the local MD step in WTMetaD with a diffusion-based neural sampler (ASBS) that learns to generate molecular configurations directly from the potential energy function. A repulsive potential is iteratively constructed in CV space to encourage exploration of new regions and flatten free-energy barriers, with subsequent reweighting to recover the unbiased Boltzmann ensemble. Experiments demonstrate that WT-ASBS improves sampling efficiency and mode discovery for molecular systems.
1. The idea of integrating diffusion-based neural samplers with enhanced-sampling techniques from molecular simulation is elegant and promising. While both ASBS and WTMetaD exist independently, their combination is novel and creates a bridge between machine-learning-based and physics-based sampling paradigms.
2. The methodology is technically well-founded, with a clear derivation and an explicit algorithmic description (Algorithm 1). The paper includes a convergence proposition showing that the bias potential approaches the well-tempered limit, ensuring theoretical consistency.
3. The paper is clearly written, logically organized, and visually well-supported by figures. The workflow in Fig. 1 and the schematic diagrams for peptide systems help the reader follow the training and bias-updating procedure.
4. From an application perspective, the work demonstrates that diffusion-based samplers can efficiently handle physically meaningful molecular systems. The results on peptides and reactive landscapes show large efficiency gains over WTMetaD. This points toward real practical potential in molecular modeling.
1. The method mainly replaces the MD propagation in WTMetaD with ASBS without introducing new learning objectives or architectures. The innovation lies more in application and system integration than in core ML algorithmic development. For ICLR, where the focus is typically on advancing machine-learning methodology, the contribution might appear not so significant. The work may find a more natural home in computational chemistry or physics-oriented venues.
2. Algorithm 1 involves two nested loops (outer bias update over k, inner ASBS optimization over l). In addition, an iid sampling procedure is involved. However, the paper does not specify when or how to stop the iterations in practice. ASBS itself is self-consistent and lacks a single scalar loss whose decrease guarantees convergence. The absence of heuristic or empirical convergence criteria (e.g., stabilization of bias potential or free-energy estimates) makes reproducibility difficult.
3. All experiments use well-known CVs such as torsional angles or bond distances. For realistic, high-dimensional systems, CV discovery is non-trivial. It would strengthen the paper to evaluate the method with ML-discovered CVs (e.g., TICA, SPIB, or FMRC). Without this, the method’s applicability to problems where good CVs are not known remains uncertain.
4. Boltzmann Generators and related approaches can guarantee asymptotically correct estimates via MCMC or importance-sampling refinement. WT-ASBS claims reweighting via the bias but does not discuss whether such reweighting ensures unbiased expectations in the presence of neural-sampler approximation error.
5. In Fig. 3d, the free-energy mean-absolute-error of WT-ASBS is slightly larger than that of WTMetaD. It would be valuable to analyze the reason to understand when diffusion-based sampling may underperform.
6. The method’s feasibility for explicit-solvent systems remains unclear.
1. How should practitioners determine that the two-loop WT-ASBS training has converged?
2. Can the authors clarify whether the proposed reweighting scheme guarantees unbiased Boltzmann expectations, or whether the residual approximation in the learned diffusion process introduces systematic bias? Could additional refinement (e.g., short MCMC runs or importance sampling) ensure asymptotic correctness similar to Boltzmann Generators?
3. Have the authors considered coupling WT-ASBS with ML-based CV discovery?
4. WT-ASBS achieves broader exploration but sometimes slightly worse free-energy accuracy (Fig. 3d). Can the authors comment on this?
5. Can the authors explore the computational feasibility for systems with explicit solvent? |
Fully AI-generated |
|
Enhancing Diffusion-Based Sampling with Molecular Collective Variables |
Soundness: 3: good
Presentation: 3: good
Contribution: 3: good
Rating: 8: accept, good paper
Confidence: 2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked. |
This paper introduces the Well-Tempered Adjoint Schrödinger Bridge Sampler (WT-ASBS), a novel method that significantly enhances the exploration capabilities of energy-based diffusion models for molecular systems. The core innovation is integrating an adaptive bias, inspired by well-tempered metadynamics (WTMetaD), into the Adjoint Schrödinger Bridge Sampler (ASBS) training loop. This bias is deposited along pre-defined Collective Variables (CVs), which are low-dimensional projections of molecular coordinates. The bias effectively increases the sampling temperature in the CV space, mitigating the mode collapse pitfall often seen in standard diffusion samplers. Crucially, the method retains the ability to recover the correct Boltzmann ensemble through importance reweighting. The authors demonstrate WT-ASBS's efficacy on conformational sampling benchmarks (alanine dipeptide and tetrapeptide) and, notably, achieve the first demonstration of reactive sampling using a diffusion-based model on $S_N2$ and post-transition-state bifurcation reactions, showing significant efficiency gains over traditional WTMetaD in terms of wall-clock time and energy evaluations.
The reviewer here acknowledges that I am familiar with MCMC methods and general energy-based sampling, but is less knowledgeable about the specifics of modern molecular sciences applications.
- The work presents a highly original and timely unification of two distinct fields: enhanced sampling (metadynamics) and generative diffusion modeling. By incorporating the well-tempered biasing mechanism directly into the iterative proportional fitting (IPF) scheme of the ASBS, the authors solve the critical problem of mode collapse for diffusion models applied to complex, high-dimensional, and multi-modal free energy landscapes. The successful application of reactive sampling, which is extremely challenging for standard diffusion models, is a significant first.
- The paper is technically sound and well-executed.
- The authors provide a convergence guarantee (Proposition 3.1), showing that WT-ASBS provably approaches the desired well-tempered target distribution.
- The experiments are convincing. On the alanine tetrapeptide, WT-ASBS discovers all eight metastable states much faster than WTMetaD (Figure 3c). The ability to accurately resolve complex free energy landscapes (PMFs) for both conformational and reactive systems, with close agreement between PMF from bias and PMF from reweighting, strongly supports the method's correctness and efficiency.
- For the complex reactive systems using universal ML interatomic potentials (uMLIPs), WT-ASBS achieves convergence with significantly fewer energy evaluations and less wall-clock time compared to WTMetaD (Figure 5).
- The paper is well-structured and clearly written, making complex concepts accessible. The methodological details, especially the two-time-scale algorithm and the practical implementation with a replay buffer (Algorithm 2), are laid out logically. Figure 1 clearly illustrates the overall scheme, and the experimental figures (e.g., Figure 2) effectively show the evolution of the PMF during training.
The effectiveness of WT-ASBS hinges on the accurate selection of several critical hyperparameters inherited from the enhanced sampling domain, such as the initial Gaussian height $h$, the Gaussian width $\sigma$, and the bias factor $\gamma$.
It seems that the paper lacks a systematic ablation study to investigate the sensitivity of the method to these parameters (e.g., how the convergence speed or final accuracy changes with different $\gamma$ values). Given that WT-ASBS generates uncorrelated samples, the optimal choice for parameters like $h$ might deviate significantly from WTMetaD conventions. It requires clearer guidelines for practitioners.
See weakness |
Fully AI-generated |