Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching

We thank the reviewer for supporting our paper. Note that we have added additional experiments and figures to provide more insight into our work. See https://sites.google.com/view/adjointsamplingrebuttal. Below we answer the reviewer’s questions in detail.

It seems the training procedure can be conducted in an off-policy manner.

We wish to clarify that our method is not actually off-policy, in the sense that we cannot sample from an arbitrary SDE, which is what is often referred to as off-policy by other methods. That being said, there are many downsides to an off-policy method – such as requiring full trajectories for optimization, and our method is actually much more scalable due to requiring only $(X_t, X_1)$ samples.

Are there any other possibilities to improve the sample efficiency by using several off-policy training schemes?

The existing off-policy methods come with their own drawbacks, such as requiring full trajectories and not using the gradient of the energy for the training loss. We added additional experiments using the log-variance loss [1], but it does not perform well. We find that such off-policy methods only work well when the gradient of the energy is used directly as part of the drift parameterization [2], which is something we need to avoid as we move on to utilizing more computationally costly energy functions.

Is there any reason why the samples are uniformly sampled from the buffer?

This is a good suggestion. We did not specifically look into prioritized replay buffer, but it is possible to use importance weights to design a prioritized replay buffer where higher weights are sampled more frequently [3]. We note that such an approach is difficult to implement for our proposed (amortized) benchmark, which aims at solving 24,000+ sampling sub-problems using a single model, as the priority over samples needs to be computed for each sub-problem separately. In order to train on all sub-problems and be able to generalize, we cannot sample each sub-problem too many times as it incurs a high computational cost so we just used a simple uniform buffer.

It will also be nice to compare the time complexity or performance of naive adjoint sampling (e.g., without reciprocal adjoint matching).

We have included additional ablations without the Reciprocal projection (see Tables 1 & 2 of the link), and also a runtime comparison between different methods (Figure 2 of the link). We note that the naive Adjoint Matching in terms of computational cost is on-par with Discrete Adjoint (PIS) in terms of runtime, while the ablation “Adjoint Sampling w/o Reciprocal” is on-par with Adjoint Sampling in terms of runtime. However, without the Reciprocal projection, the performance deteriorates because it does not as efficiently search the sample space. With Reciprocal AM, the samples $X_t$ are uncorrelated, whereas without it, the samples across time are from the same trajectory (see Figure 5 of the link). We hypothesize that the improved performance when using Reciprocal AM (in terms of Energy W2 for Table 1 and Recall for Table 2) is due to the ability to see more diverse X_t samples during training.

[1] “Improved sampling via learned diffusions”

[2] “No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers”

[3] “Sequential Controlled Langevin Diffusions”

We understand and agree with the reviewer’s concerns regarding additional baselines and the lack of clarity around the proposed benchmark.

To respond, we (i) have additional experiments, (ii) expand the discussion to related off-policy methods, and (iii) clarify why the proposed benchmark is much harder than existing ones. See: https://sites.google.com/view/adjointsamplingrebuttal

the log-variance loss has similar benefits as adjoint sampling

There are problems with the log-variance loss that inhibits scalability. Firstly, it requires full trajectories of length N, and N network evals, per gradient. Partial trajectories can only be used if the time-marginals are either learned [1] or prescribed [2]. In contrast, Adjoint Samping only uses pairs of (X_t, X_1) samples, and one network eval, for each gradient.

Secondly, the log-variance loss doesn’t make use of the energy gradient information. For log-variance to work in high-dimensions, the gradient of the energy is used directly as part of the drift parameterization and it can fail without this [3], but we avoid this since our focus is on computationally expensive energy functions (such as the proposed benchmark). In our tests, the log-variance loss fails to learn on LJ because the potentials have extremely large values.

other sampling methods

We added multiple baselines, including log-variance, DDS, and ablations. Additionally, we have incorporated the suggested reference [4] to estimate ESS and ELBO metrics. Please see Tables 1 & 2 in the above link.

more established benchmarks

We believe the DW and LJ problems are very common and have been used by many prior works.We mainly focus on our proposed benchmark.

how high-dimensional are the proposed benchmark problems?

We agree this was not emphasized well and hope to remedy this. The number of atoms range from 3-50 (median 39), sample dimension is 3x the number of atoms, while the conditioning dimension is quadratic (bonds). See Figs 3 and 4.

This is an amortized setting for finding the distribution of conformers for over 24,000 molecules (i.e., 24,000 conditional sampling problems). The benchmark is designed to test for generalization to unseen molecules, whereas most existing sampling benchmarks (such as DW, LJ, alanine dipeptide) only test performance on a single & cheap energy function.

Furthermore, each molecule has a number of conformers (representing low energy regions), ranging from a handful to a few hundred (Figs 6,7 in submission). The benchmark is specifically designed to test if a sampling method can find ALL the conformers for every molecule. As such, the metrics of interest are precision and recall against a set of highly-diverse ground truth samples computed using density functional theory (DFT) that took days to obtain even on a large cluster.

Finally, the energy function for this benchmark is a large graph neural network. As such, care must be taken in order to not incur computational cost by evaluating the energy function too many times. See Fig 2 for a runtime plot.

To our knowledge, this is the most difficult sampling benchmark so far, which we believe can aid in directing future research on sampling methods. Performing well on this benchmark has direct consequences in advancing computational chemistry and drug discovery.

other sampling method on the novel benchmark

It is VERY difficult to get reasonable performance on this benchmark. The only other method that we are aware of that can train with only pairs of (X_t, X_1) samples is iDEM, which we now include (see Table 2 of the link). iDEM is biased when few MC samples are used; it typically uses 512 MC samples (== energy evaluations) per gradient, which is prohibitive (see Figure 2). We also added an ablation for the Reciprocal projection (see Table 1 & 2 of the link); without it, the model performs worse, which we hypothesize is due to the lack of exploration (see Figure 5 in the link).

[...] Alanine Dipeptide, which is a common problem in the sampling community?

We did not specifically test the alanine dipeptide setup (which is 22 atoms and uses a classical energy function). The molecules in our benchmark are much larger, and our energy function is much more expensive as it is a GNN trained to approximate forces from DFT.

Does the proposed objective avoid mode-collapse such as e.g. relative entropy minimization? Our method is related to the reverse KL rather than the forward KL (i.e., relative entropy). Mode collapse is inherently difficult as there is no way to know where modes are without fully exploring the search space. Our method relies on the choice of base process to determine what region to search. Finally, we again note that the new benchmark is specifically designed to test mode coverage as it requires finding all minima.

[1] https://arxiv.org/abs/2310.02679

[2] https://arxiv.org/abs/2412.07081

[3] https://arxiv.org/abs/2502.06685

[4] https://arxiv.org/abs/2406.07423

We thank the reviewer for being candid and providing us the opportunity to substantiate our claims, which we believe we can do. The reviewer is concerned with our claims regarding (i) the efficiency and analysis of Adjoint Sampling and (ii) our proposed large-scale sampling benchmark. We agree that we did not emphasize these contributions enough, and hope to clarify them in the following answers.

Please find additional ablations, experiments, and runtime figures here: https://sites.google.com/view/adjointsamplingrebuttal

Analysis and comparison to Adjoint Matching

Efficiency from a simple base processes

Firstly, note that the original Adjoint Matching (AM) method is designed for general base process, and has the same computational cost as the path integral sampler, requiring full trajectory simulation and solving the lean adjoint state. One of our first key observations is that the lean adjoint state can be solved in closed-form when the base process is simple.

Ablation and analysis of the Reciprocal projection

Furthermore, we note that the Reciprocal projection further enhances optimality, allowing us to explore the search space much faster (see Figure 5 of above link) as it allows training on uncorrelated X_t samples. We performed the requested ablation where we removed Reciprocal projection (see “Adjoint Sampling w/o Reciprocal Projection” in Tables 1 & 2). We see that Adjoint Sampling with Reciprocal projection performs better in all settings. On the LJ classical force fields, it avoids high energy regions much better (better Energy W2 metric), and on our amortized conformer generation task it covers the target distribution much better (better recall and precision metrics). We hypothesize this is because the Reciprocal projection lets the model see far more diverse trajectories across different molecular conditionings. In the amortized benchmark, the number of samples per molecule is extremely small, so having a strong learning signal for each energy evaluation becomes very important. Moreover, we know that the Reciprocal projection is theoretically grounded and preserves the optimal solution as a unique fixed point.

Runtime plot

We have included a run-time breakdown in Figure 2 of the link. It shows that Adjoint Sampling is computationally efficient in terms of time-spent evaluating the energy and sampling model. This is what enables us to scale up to both larger architectures and energy models, enabling us to sample from molecular energy foundation models. Other methods will either require full trajectory optimization (such as PIS, log-variance) or require an intractable number of energy evaluations (such as iDEM).

Large-scale experiments for sampling

As mentioned above, Adjoint Sampling is specifically designed for sampling from unnormalized distributions where a simple base process can be used. The text-to-image finetuning problem is a more general problem statement where a pre-trained generative model is used as the base process. Since our methodological contributions rely on the choice of a simple base process, we restrict ourselves to the pure sampling setting, where only an unnormalized density is given.

In the literature of sampling from unnormalized densities, our proposed benchmark is actually THE most difficult benchmark to date, requiring highly scalable algorithms. Prior related works have only experimented with classical (synthetic) force fields such as the Lennard-Jones experiments we have included. In contrast, the conformer generation benchmark uses a large graph neural network as the energy function, and hence each evaluation incurs a runtime cost. Furthermore, this benchmark is aimed at learning conditional sampling models, amortized over 24000+ molecules, with the test metrics being generalization to unseen molecules. This type of amortized benchmarks for sampling is incredibly rare and has not been well-explored. Due to this amortization, expensive methods that require full simulations for each gradient update (such as Adjoint Matching) does not scale well — again, note the runtime plot (Figure 2 in the link), Adjoint Matching in its basic form has a similar computation cost as the Discrete Adjoint (PIS).

Finally, the reward fine-tuning benchmarks are aimed at finding one good image sample per text prompt. Here, our benchmark requires the model to find a set of diverse samples per molecule (ALL local minima), and the metrics of interest are precision and recall against a set of highly-diverse ground truth samples from density functional theory (DFT) that took weeks to obtain even on a large cluster. Performing well on this benchmark has direct consequences in advancing computational chemistry and drug discovery.

We hope this answers the reviewer regarding (i) why Adjoint Sampling works only for the sampling task, and (ii) how our proposed new benchmark is precisely encouraging the sampling community to look into harder and larger-scale problems.

We thank the reviewer for their insightful comments. We’ve incorporated new baselines and metrics, and have produced figures to better illustrate our claims.

Additional figures & results: https://sites.google.com/view/adjointsamplingrebuttal

I’m not sure if the scalability was properly based on convincing experimental setting

We agree we did not emphasize this enough. Please see the link above for a runtime plot. Adjoint Sampling is the only method so far that can work with $(X_t, X_1)$ pairs — no full trajectories — and can scale to computationally expensive energy functions. Further details are provided in the answers of the reviewer’s other questions.

I strongly encourage the authors to include NLL and ESS metrics for the experiments on LJ potentials and DW-4. LJ potentials experiments should be equipped with the sampled energies histograms compared with the ground truth energy histogram.

I would recommend adding other samplers like FAB [1], or DDS [2] as the baselines for experiments

We have added several new baselines, including DDS. We have also included path-based ESS and ELBO estimates, and energy histograms. New results and figures can be found in the link above. For ESS and ELBO, we used the same method as [1,2,3]. However, we were not able to get reasonable estimates for iDEM at this time following the approach of [1] so have decided not to include them for now.

I think that putting more attention on a recent approaches for scaling the training or abilities of diffusion samplers, would be beneficial for this work, e.g., [1] or [2].

We completely agree! We will add more discussions to existing methods.

We note that there may be a conflation of “off-policy” and “scalability”, which are best decoupled. Off-policy methods, such as the log-variance divergence [4] and trajectory balance [5], do not take gradient of the energy function (as they do not differentiate through the model sample) and actually strongly relies on parameterizing the drift using the gradient of the energy function as discussed in [6]. In our setting, the energy function is more expensive than the drift network (see runtime plot), so our experiments do not use the energy function as part of parameterizing the drift. Another major downside is that these methods require the full trajectory, only being able to use sub-trajectories if the time-marginals are either learned [5] or prescribed [7]. In contrast, Adjoint Sampling is an on-policy method (resulting in a more direct update to the current model), directly works with the energy gradient rather than the energy values, and only requires $(X_t, X_1)$ pairs.

In our additional experiments, we have added a log-variance baseline from [3] to the synthetic energy experiments which did not scale beyond DW4. This is because the LJ potentials can have extremely large values which an importance-sampling based method like log-variance does not handle well (again, without taking the gradient of the energy into the parameterization).

It’s unclear to me why lean adjoint matching and adjoint sampling actually produce good gradients, since we are not minimizing the KL divergence anymore with the lean adjoint matching?

Instead of optimizing the KL, Adjoint Matching (AM) directly regresses onto the control in the manner of a consistency loss. AM is still a fixed point iteration method and has the optimal control as the unique solution, and can be interpreted as removing a stochastic term from the KL gradient that has expectation zero at the optimum (as discussed in the original AM paper). In particular, the lean adjoint state does not depend on the learned control and only the base process.

This last property of the lean adjoint state is incredibly important in the sampling setting, and in the paper we use simple base processes so the lean adjoint state has a closed-form solution. This leads to allowing us to use only $(X_t, X_1)$ samples for training, and the Reciprocal projection further enhances optimality. We plan to provide more details about these derivation steps in the main paper.

[1] “Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling”

[2] “Path Integral Sampler: a stochastic control approach for sampling”

[3] “Improved sampling via learned diffusions”

[4] “From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training.”

[5] “Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization”

[6] “No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers”

[7] “Sequential Controlled Langevin Diffusions”