Adjoint Schrödinger Bridge Sampler

We thank the reviewer for raising all valuable comments and the comprehensive summary. We are excited that the reviewer identified the novelty of technical contributions and connection to the Iterative Proportional Fitting algorithm, appreciated the well-written presentation, and acknowledged our strong empirical results. We believe ASBS takes a significant step toward scalable diffusion samplers with flexible model design. We try our best to resolve all raised concerns in the responses below.

1. Additional experiments on mode searching

• We thank the reviewer for raising the comments. Following the reviewer’s suggestion, we conduct additional experiments on the 40-mode GMM in 2D from [1]. Specifically, we instantiate ASBS with a uni-variance Gaussian source distribution centered at zero, effectively assuming that users have no prior knowledge of the target modes, as the initial distribution does not coincide with any target modes. We also run a vanilla Langevin baseline for 1 million steps, starting from the same source distribution.

• The table below presents the count of identified modes out of a total of 40. Notably, ASBS is able to identify almost all modes. In contrast, the vanilla Langevin baseline appears to suffer from a slow mixing rate, recovering less than half of the total modes even after 1 million steps. We highlight this distinction as an advantage of constructing diffusion samplers from the stochastic control and Schrödinger Bridge frameworks, which allows theoretical convergence to target distribution within a finite horizon. Finally, we believe that with proper tuning of the ASBS noise schedule, its performance can be further enhanced.

  	# of identified modes out of 40
  Vanilla Langevin	15/40
  ASBS	38/40

• We kindly note that, in accordance with this year’s NeurIPS policy, we are unable to provide any external link to images or updating the manuscript. However, we believe these are valuable discussions and will include the figures in the revision. We thank the reviewer again for bringing up the topic.

2. Clarification on nested optimization (Eqs 14 & 15)

• Regarding the sensitivity of the interplay between Adjoint Matching (Eq. 14) and Corrector Matching (Eq. 15), we observe in practice that both optimizations exhibit stable convergence due to their formulation as simple regression (i.e., matching) objectives. While exact convergence to the minima of Eqs. 14 and 15 is required for global convergence (see Thm 3.2), we found that early termination with a fixed number of iterations did not cause training to diverge or become unstable in any of the tasks considered.

• In practice, the Corrector Matching (Eq 15) often requires significantly fewer iterations to converge compared to the Adjoint Matching (Eq 14). Specifically, the number of iterations for CM ($M_\text{crt}$) are typically 2 to 10 times fewer than for AM ($M_\text{adj}$). We refer the reviewer to Table 5 in Appendix for the exact values of $M_\text{crt}$ and $M_\text{adj}$ on each task.

[1] Flow Annealed Importance Sampling Bootstrap
[2] Non-Equilibrium Transport Sampler

We thank the reviewer for raising all valuable comments and questions. We are excited that the reviewer identified the novelty of technical contributions in connecting SB with SOC, appreciated the scalability of our SB-based diffusion samplers, and the efficiency evidenced by empirical results on real-world sampling problems. We believe ASBS takes a significant step toward scalable diffusion samplers with flexible model design. We try our best to resolve all raised concerns in the responses below.

1. Additional experiments on ALDP

• We first thank the reviewer for raising the comments. The experiments for ALDP (Table 3 & Fig 5) were conducted in internal coordinates mainly to maintain consistency with prior SOC-based diffusion samplers, specifically PIS [2], to which our ASBS sampler also belongs. That being said, our ASBS can indeed be run on the full Cartesian coordinates for ALDP.

• In the table below, we report preliminary results for ASBS (using harmonic source) trained from scratch in full Cartesian coordinates, without relying on any explicit target or MCMC samples. Similar to Table 3, we report the KL divergences for the 1D marginal across five torsion angles, $(\phi, \psi, \gamma_1, \gamma_2, \gamma_3)$, as well as the Wasserstein-2 on joint distribution of $(\phi, \psi)$. The latter serves as a quantitative measure of the Ramachandran plot. We kindly note that, in accordance with this year’s NeurIPS policy, we are unable to provide any external link to images or updating the manuscript.

  	$\phi$	$\psi$	$\gamma_1$	$\gamma_2$	$\gamma_3$	$(\phi,\psi)$
  ASBS (Cartesian)	0.78	0.54	0.20	0.11	0.11	1.30

• Our preliminary results imply that ASBS can generate samples with reasonable structures, with Ramachandran plots that match those of prior SOC-based diffusion samplers in internal coordinates (see Table 3). Notably, ASBS achieves this without using any explicit target samples during training, relying solely on the energy function of ALDP. This contrasts with previous attempts on ALDP in Cartesian coordinates—such as [1,3] mentioned by the reviewer and Reviewer jVAy—which often require explicit samples from or near the target distribution. We emphasize this distinction as a key advantage of adjoint-based diffusion samplers, which we further detail in the second section.

• Finally, while we believe that the current results could be further improved with additional tuning and/or pretraining, we also highlight an intriguing application of ASBS: using model samples from [1,3] as its source distribution. This effectively transforms ASBS to a “fine-tuning” method, offering an alternative to importance sampling at inference time, which can be costly in terms of energy evaluation. We highlight this flexibility in choosing source distribution, a feature that is otherwise infeasible in prior adjoint-based methods, particularly Adjoint Sampling.

2. Discussion on Adjoint Sampling (AS)

• While our ASBS is closely related to AS, as noted by the reviewer, we emphasize that the motivation behind ASBS differs significantly from that of AS. Practically, AS can only be used for sampling from a distribution, while ASBS can be used to transport from any (e.g. a similar-but-incorrect) distribution. Despite this additional capability, we found ASBS also outperforms AS in pure sampling due to the less restrictive requirements. Theoretically, ASBS is designed to solve the Schrödinger Bridge (SB) problem (Eq 3) whereas AS solves a specific form of SOC problem (Eq 7). The theoretical connection between these two problems only became clear with the presentation of Thm 3.1. This is also acknowledged by the reviewer, i.e., "Elegant SOC formulation of SB problem”.

• Given the practical success of AS in terms of both performance and energy efficiency, along with our theoretical results in Thm 3.1 that generalizes AS to a broader, non-memoryless SOC problem, it is natural to pursue a new learning method that retains the algorithmic advantage of AS while addressing additional complexity from non-memorylessness and arbitrary source distributions. As shown in Sec 3.3, this can be achieved by incorporating an additional corrector matching objective (Eq 15). Since Eq 15 employs a simple matching objective, the overall ASBS method (Alg 1) remains as scalable as AS (Fig 4).

• Such a straightforward algorithmic modification from AS to ASBS, although theoretically complex, has led to significant performance gains, evidenced by all of our experiments (Tables 2,3,4). Furthermore, it does not impede convergence to global solution (Thm 3.2). To summarize, we believe ASBS complements AS with additional flexibility in selecting informative sources and utilizing richer model classes, leading to significant performance gains without compromising algorithmic scalability. Together, these advancements enable ASBS to achieve strong performance and open new avenues for future applications.

3. Clarifications on questions

• LV-PIS & LV-DDS in Table 1. We thank the reviewer for raising this question. Our intention was to convey that LV directly minimizes the variance of the importance weights (IWs), and thus, methods in this category involve the computation of IWs. We acknowledge that they do not perform importance sampling, which may have been incorrectly inferred from the current Table 1. We will clarify this in Table 1 in the revision.

• Pretraining with IPF. As the reviewer conjectured, the global convergence of IPF requires initialization from either $u=0$ or $h=0$. Initializing IPF with a pretrained $u$—which corresponds to the first stage of Iterative Markovian Fitting (IMF)—has been theoretically explored recently [4] and shown to maintain global convergence in Gaussian setups. However, in more complex scenarios, we still observe stable local convergence. Alternatively, we can initialize ASBS by constructing its source distribution using samples from pretrained models, thereby leveraging ASBS's flexibility in selecting source distributions without compromising global convergence.

• Log normalizer of ASBS. The importance weights in Eq 78 can be computed by re-parametrizing $u$ and $h$ with two potential networks, specifically $u = \nabla v$ and $h = \nabla \bar h$ as shown in Eq 80, and then calculating the importance weights using Eqs 78 & 81. Alternatively, we can train an additional backward drift $v_t := \nabla \log \hat \varphi_t$ with standard SB matching [5], and then use Eq 79 to compute the ratio $\log \frac{\hat \varphi_1}{\hat \varphi_0}$. While both methods are tractable, they introduce additional complexity. Exploring alternative approaches to compute the weights would be a promising direction for future research (L295).

• Buffer size in Table 5. This’s a typo. The correct buffer size $|\mathcal{B}|$ for conformer generation should be 64000, which is consistent with the order of magnitude used for other tasks. We thank the reviewer for the meticulous reading!

• Symmetry for ALDP. In practice, we utilize the full internal coordinate transform from FAB [6], which converts Cartesian coordinates into bond lengths, bond angles, and dihedral angles, effectively removing 6 degrees of freedom (3 for translation and another 3 for rotation). Non-augular coordinates are further normalized using samples with minimal energies. We refer the reviewer to Appendix F.1 in [6] for further details and will include these discussions in the revised Appendix.

• Corrector matching (CM) for Gaussian prior. We first clarify that CM is used for all ASBS models, including those with Gaussian priors. This is because the base processes $p^\text{base}$ for all ASBS models are not memoryless—with $f :=0$ and $\sigma_t$ set according to Table 5— hence necessitating the use of CM, as discussed in Sec 3.2. As the reviewer conjectured, there is indeed a memoryless base process (specifically, the one considered in DDS) for Gaussian prior, which eliminates the need for CM. This corresponds exactly to the leftmost plot in Fig 1, where significant noise is injected into sampling to ensure the conditional independence of the memoryless process (Eq 6). We have also reported its performance (i.e., DDS) on synthetic energies (Table 2) and ALDP (Table 3), where ASBS outperformed DDS in both cases. We will include these clarifications in the revision.

• SCLD. We thank the reviewer for catching the error. We did use subtrajectory splitting (with the default value of 4) for SCLD, so it does not degenerate to CMCD with LV. In practice, we find that using subtrajectories yields better results. This will be corrected in the revision, and we appreciate the reviewer’s meticulous reading.

• Harmonic vs Gaussian prior in synthetic energies. In the below table, we present a comparison of ASBS using harmonic and Gaussian priors. The results are consistent with those in Table 4, showing slight improvements with harmonic priors in both cases.

  	DW4 (W2)	DW4 (EW2)	LJ13 (W2)	LJ13 (EW2)	LJ55 (W2)	LJ55 (EW2)
  Harmonic	0.38	0.19	1.59	1.28	4.00	27.69
  Gaussian	0.42	0.20	1.60	1.82	4.09	38.70

[1] Temperature-Annealed BG
[2] Path Integral Sampler
[3] Scalable Equilibrium Sampling with Sequential BG
[4] Diffusion & Adversarial SB via IPMF
[5] Diffusion SB Matching
[6] Flow Annealed Importance Sampling Bootstrap

We thank the reviewer for raising all valuable comments. We are excited that the reviewer identified the novelty of technical contributions in relaxing the memoryless condition, appreciated the well-written presentation, and acknowledged our strong empirical results. We believe ASBS takes a significant step toward scalable diffusion samplers with flexible model design. We try our best to resolve all raised concerns in the responses below.

1. Additional experiments on ALDP

• We first thank the reviewer for raising the comments. The experiments for ALDP (Table 3 & Fig 5) were conducted in internal coordinates mainly to maintain consistency with prior SOC-based diffusion samplers, specifically PIS [4], to which our ASBS sampler also belongs. That being said, our ASBS can indeed be run on the full Cartesian coordinates for ALDP.

• In the table below, we report preliminary results for ASBS (using harmonic source) trained from scratch in full Cartesian coordinates, without relying on any explicit target or MCMC samples. Similar to Table 3, we report the KL divergences for the 1D marginal across five torsion angles, $(\phi, \psi, \gamma_1, \gamma_2, \gamma_3)$, as well as the Wasserstein-2 on joint distribution of $(\phi, \psi)$. The latter serves as a quantitative measure of the Ramachandran plot. We kindly note that, in accordance with this year’s NeurIPS policy, we are unable to provide any external link to images or updating the manuscript.

  	$\phi$	$\psi$	$\gamma_1$	$\gamma_2$	$\gamma_3$	$(\phi,\psi)$
  ASBS (Cartesian)	0.78	0.54	0.20	0.11	0.11	1.30

• Our preliminary results imply that ASBS can generate samples with reasonable structures, with Ramachandran plots that match those of prior SOC-based diffusion samplers in internal coordinates (see Table 3). Notably, ASBS achieves this without using any explicit target samples during training, relying solely on the energy function of ALDP. This contrasts with previous attempts on ALDP in Cartesian coordinates—such as [1,3] mentioned by the reviewer and Reviewer 5V2q—which often require explicit samples from or near the target distribution. We emphasize this distinction as a key advantage of adjoint-based diffusion samplers, which we further detail in the second section.

• We would like to highlight that the amortized conformer generation task (Table 4) already includes larger molecules, with 50-65 atoms in the test sets of SPICE and GEOM-DRUG, respectively, whereas the suggested alanine tripeptide consists of 33 atoms. Given ASBS's strong performance in amortized conformer generation and its capability to generate the Boltzmann distribution of ALDP in Cartesian spaces, we believe ASBS has great potential for tackling larger molecules. Scaling ASBS for amortized Boltzmann generators would be a promising direction worthy of its own in-depth investigation.

2. Discussion on the new RDKit baseline

• We thank the reviewer for referring us to the new RDKit baseline [2], which post-processes RDKit samples using a K-means clustering algorithm. Following the reviewer’s suggestion, we applied their method to our GEOM-DRUG test set and found that the recall converges at the same 0.75Å threshold, as reported in Table 4, to approximately 54.4%. This performance is comparable to ASBS with a harmonic prior and relaxation post-processing.

• We note, however, that [2] appears to generate a much larger number of conformers than the evaluation pipeline used in Table 4. That said, we believe that ASBS could similarly benefit from clustering post-processing by generating a comparable number of conformers. Exploring such post-processing techniques, which have proven effective in prior work, would be an interesting direction for future research. We thank the reviewer for highlighting this opportunity.

[1] Scalable Equilibrium Sampling with Sequential BG
[2] Do DL Methods Really Perform Better in Molecular Conformation Generation
[3] Temperature-Annealed BG
[4] Path Integral Sampler

We thank the reviewer for raising all valuable comments. We are excited that the reviewer identified the novelty of our technical contributions for overcoming Dirac limitations from prior methods, appreciated the advantages of SB-based diffusion samplers, and acknowledged our strong and extensive empirical results. We believe ASBS takes a significant step toward scalable diffusion samplers with flexible model design. We try our best to resolve all raised concerns in the responses below.

1. Clarification on Sec 2 (Eqs 5 & 6)

• We understand that current Sec 2 may be compactly written, mainly due to the space constraints at submission. In response to the reviewer’s suggestion, we will provide further clarifications in the revised Sec 2, along with a self-contained, step-by-step derivation in the revised Appendix. We detail these clarification below.

• The goal of Sec 2, as briefly stated in L86-88, is to formally introduce the mathematical formulation of the SOC problem (Eq 4), its optimal distribution (Eq 5), and how prior diffusion samplers have adopted it for Boltzmann sampling under the memoryless condition (Eq 6). This leads to a specialized, “memoryless”, SOC problem (Eq 7), which contrasts with our generalized, “non-memoryless”, SOC problem inspired by Schrödinger Bridge (Eq 9; also see L129-131).

• As the reviewer noticed, Eq 5 is indeed the key to connect SOC problems with Boltzmann sampling. This equation, representing the optimal solution to the original SOC problem (Eq 4), can be derived as follows:

  • By rewriting Eq 4 as KL-regularized optimization, $\min_{p^u} \mathrm{KL}(p^u || p^\text{base}) + \mathbb{E}_{p^u_1}[g(X_1)]$ using Girsanov theorem [1].
  • Due to the KL term, the aforementioned optimization is convex in $p^u$, allowing for an analytic solution. Setting the first-order variation to zero yields $p^\star(X_1 | X_0) \propto p^\text{base}(X_1 | X_0) e^{-g(X_1)}$.
  • Finally, substituting the normalization constant $\int p^\text{base}(X_1 | X_0) e^{-g(X_1)} \mathrm{d} X_1 = e^{-V_0(X_0)}$ leads to Eq 5.

  While Eq 5 is well-known in SOC literature [2,3], we acknowledge that it may not be immediately clear to readers without sufficient background—even with the interpretation provided right below (L93-96). In the revision, we will include these clarifications along with additional citations for Eq 5, including the original AM paper, which—we kindly note that—was already cited in the beginning of Sec 2 (L86).

• Eq 5 implies that, for any SOC problems (Eq 4) to be effective in sampling Boltzmann distribution, the optimal distribution $p^\star$ must satisfy the desired target distribution $\nu$ at terminal time $t=1$. This requires appropriately selecting the base drift $f_t$, noise schedule $\sigma_t$, source distribution $\mu$, and terminal cost $g$, such that $\int p^\star(X_0, X_1) \mathrm{d} X_0 = \nu(X_1)$. In Adjoint Sampling [4], this is achieved by

  • Enforcing the base process $p^\text{base}$—uniquely defined by $(f_t, \sigma_t, \mu)$—to be memoryless (Eq 6). This simplifies the marginalization to $\int p^\star(X_0, X_1) \mathrm{d} X_0 = C \cdot p^\text{base}(X_1)_1 e^{-g(X_1)}$ with some constant $C$, as detailed between L102-123.
  • Finally, equating $p^\text{base}_1(X_1) e^{-g(X_1)} = \nu(X_1)$, which leads to the proposed terminal cost $g := \log p^\text{base}_1 - \log \nu$ in L102.

  We will include these clarifications for how $g$ is set up in the revision, and we thank the reviewer for pointing it out.

• Finally, for notational simplicity, we omitted the time specification of distribution—such as $p^\text{base}(X_0)$ and $p^\text{base}(X_1)$ in Eq 6—whenever it can be clearly inferred from the variables $X_0$ and $X_1$. As the reviewer correctly pointed, this is not mathematically rigorous, and we will correct these inconsistencies in the revision.

2. Clarification on the Adjoint Matching (Eq 14) and Corrector Matching (Eq 15)

• Given our proposed SOC problem (Eq 9), which accommodates non-memoryless conditions and arbitrary source distributions, Sec 3.2 identifies two key mathematical objects, namely the optimal control $u^\star$ in Eq 12 and the corrector $\nabla \log \hat \varphi_1$ in Eq 13. These two mathematical objects characterize the optimality condition of our SOC problem (Eq 9) and, equivalently, the Schrödinger Bridge problem (Eq 3). In other words, the diffusion sampler that solves Eqs 9 & 3 must satisfy the optimality conditions in Eqs 12 & 13.

• Building on the reasoning above, our algorithmic goal is to design a learning method that finds a diffusion sampler satisfying Eqs 12 & 13, which correspond to two simple matching-based objectives (RHS of Eqs 12 & 13). However, these two matching objectives cannot be implemented naively due to their interdependency: Solving the matching objective in Eq 12 for the optimal control $u^\star$ requires knowledge of the corrector $\nabla \log \hat \varphi_1$. Likewise, solving the matching objective in Eq 13 for $\nabla \log \hat \varphi_1$ requires knowledge of $u^\star$.

• In Sec 3.3, we address the interdependency between the control $u$ and the corrector $h$ through an alternating training scheme. Specifically,

  • Starting from an initial guess $(u, h)$ of the control and corrector, we first update $u$ by solving Eq 12 with $h \approx \nabla \log \hat \varphi_1$.
  • Then, we update $h$ by solving Eq 13 with $u \approx u^\star$.
  • This alternate update between $u$ and $h$ continues until $u$ converges to a fixed point (see Alg 1).

  This effectively creates a sequence of update, $(u^{(0)}, h^{(0)})$ → … $(u^{(k)}, h^{(k)})$ → …, that may be thought of as running coordinate descent between the control $u$ and the corrector $h$. Intuitively, at each stage $k$, we first find the control $u^{(k)}$ that best aligns with the corrector from previous stage, $h^{(k-1)}$, then update the corrector $h^{(k)}$ accordingly to reflect the “memorylessness” of the current control $u^{(k)}$. We formalize these update rules in Eqs 14 & 15, and their dynamics are visualized on a 2D toy example in Fig 2.

• Whether the aforementioned update sequence converges—and if so, whether it converges to the global solution—is a highly non-trivial question. Nevertheless, we provide a positive answer in Thm 3.2. The proof, detailed in Sec 4, essentially shows that the update rules in Eqs 14 & 15 are contractive maps that progressively bring $u^{(k)}$ closer to the optimal control $u^\star$. Formally, the KL divergence $\mathrm{KL}(p^u_1 || \nu)$ is non-increasing as $k$ grows [5].

Overall, we acknowledge that Sec 3.3 is compactly written and will include these discussions in the revision. We thank the reviewer again for the suggestion.

[1] Applied SDEs (Särkkä & Solin, 2019)
[2] Adjoint Matching
[3] Solving high-dim HJB PDEs using neural networks
[4] Adjoint Sampling
[5] Diffusion SB with App. to Score-Based Generative Modeling

We thank the reviewer again for all valuable comments. In our rebuttal, we have clarified all raised concerns by providing self-contained, step-by-step derivations of Eqs (5), (6), (14), and (15). These clarifications will be included in the revised Sec 2 and Appendix. We appreciate the reviewer’s feedback in enhancing the clarity of our presentation—a key feature that we highly value and has been positively recognized by Reviewers jVAy & NUeM.

As we approach the end of the discussion period (Aug 8th 11:59pm AOE), we would like to ask whether our responses have covered the questions raised in your initial reviews. We are happy to provide any further clarification and discussion. If our replies adequately address your concerns, we kindly ask the reviewer to reconsider their initial rating, so that it better reflects the discussion at the current stage.