Constrained Diffusion with Trust Sampling

Authors: william  huang
yifeng  jiang
tom  van wouwe
c karen liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:37:43

Abstract: Diffusion models have demonstrated significant promise in various generative tasks; however, they often struggle to satisfy challenging constraints. Our approach addresses this limitation by rethinking training-free loss-guided diffusion from an optimization perspective. We formulate a series of constrained optimizations throughout the inference process of a diffusion model. In each optimization, we allow the sample to take multiple steps along the gradient of the proxy constraint function until we can no longer trust the proxy, according to the variance at each diffusion level. Additionally, we estimate the state manifold of diffusion model to allow for early termination when the sample starts to wander away from the state manifold at each diffusion step. Trust sampling effectively balances between following the unconditional diffusion model and adhering to the loss guidance, enabling more flexible and accurate constrained generation.  We demonstrate the efficacy of our method through extensive experiments on complex tasks, and in drastically different domains of images and 3D motion generation, showing significant improvements over existing methods in terms of generation quality.  Our implementation is available at https://github.com/will-s-h/trust-sampling.

Main paper: 
 introduction: 
 Diffusion models are a class of generative models that have been highly successful at modeling complex domains, ranging from the generations of images [22, 13] and videos [24] , to 3D geometries [32, 54, 4] and 3D human motion [52, 50] , outperforming other deep generative models, such as GANs and VAEs [50, 13, 23] . Originally for unconditional generation, Diffusion models soon became used for cross-domain conditioned generation, such as text-conditioned image generations [43, 39] , and generating human movements from audio [4] .For more fine-grained conditional generation where the samples need to precisely follow specified constraints, such as generating images following a certain contour, high-level controls like text prompts become insufficient. Guided diffusion has recently emerged to be a powerful paradigm on a variety of such constraints. One category of guided diffusion uses a separately trained classifier (as in classifier guidance [13] ) or the score of a conditional diffusion model in classifier-free guidance [21] . For new constraints, the classifier or the conditional diffusion model must be retrained [60, 42, 57] .Alternatively, one can use the gradient of a loss function representing a constraint as guidance to achieve conditional diffusion [47, 11] . This flexible paradigm allows various constraints to be applied on a pre-trained diffusion model without compute cost on extra training. On this front, since the seminal works of Chung et al. [11] and Ho et al. [24] , a number of techniques have been proposed to improve the quality of loss-guided diffusion, such as better step size design [58] , multi-point MCMC Figure 1 : Trust Sampling can be applied to complex constraint problems in drastically different domains.approximation [47] , and incorporation of measurement models [46] . Several challenges remain for the current paradigm when trying to apply loss-guided diffusion for challenging constraints. For one, performance drops significantly when using a smaller budget of inference computation with fewer neural function evaluations (NFEs) [11, 58] . The methods are also sensitive to initialization, where previous evaluations often times take the best of a few generated samples for each constraint input.In light of these challenges in training-free guided diffusion, we introduce Trust Sampling, a novel method that strays from the traditional approach of alternating between diffusion steps and loss-guided gradient steps in favor of a more general approach, considering each timestep as an independent optimization problem. Trust Sampling allows for multiple gradient steps on a proxy constraint function at each diffusion step, while scheduling the termination of the optimization when the proxy cannot be trusted anymore. Additionally, Trust Sampling estimates the state manifold of the diffusion model to allow for early termination, if the predicted noise magnitude of the sample exceeds the expected one in each diffusion step. Our framework is flexible, efficient, and performs well, achieving higher quality across widely different domains (e.g. human motion and images). We demonstrate the generality of Trust Sampling across a large number of image tasks (super-resolution, box inpainting, Gaussian deblurring) and motion tasks (root trajectory tracking, hand-foot trajectory tracking, obstacle avoidance, etc.). When compared to existing methods, we find that Trust Sampling satisfies constraints better and achieves higher fidelity. 
background: 
 Diffusion models. There are several equivalent formulations for diffusion models used in literature. Here, we briefly offer background on the denoising diffusion probabilistic model (DDPM) [22] formulation. Beginning from the data distribution x 0 ∼ p(x), we can use a variance schedule β 1 , . . . , β T to produce latent variables x 1 , . . . , x T through the forward diffusion process q(x t |x 0 ) = N ( √ α t x 0 , (1 -α t )I), where α t := t s=1 (1 -β s ). In turn, a de-noising model ϵ θ can be trained by minimizing the following loss function, which is a re-weighting of the variational lower bound [22] :EQUATIONwhere x 0 ∼ p(x), t ∼ Unif{1, . . . , T }, ϵ ∼ N (0, I), and x t = √ α t x 0 + √ 1 -α t ϵ. The diffusion model can then be sampled in the reverse process, via DDPM [22] or the denoising diffusion implicit model (DDIM) formulation [45] :EQUATIONwhere z ∼ N (0, I) in both DDIM and DDPM, whereasσ t = (1 -α t-1 )/(1 -α t ) 1 -α t /α t-1is fixed in DDPM and can be chosen freely in DDIM. x0 (x t ) denotes the predicted x 0 at timestep t, and can be written asEQUATIONNotably, DDPM and DDIM sampling can also be thought of a special case of gradient-based MCMC sampling (or a probability flow, in cases of DDIM without noise), where the goal is to refine the starting sample at each level x t towards maximizing the likelihood x t-1 ∼ p(x t-1 ). In the case of DDPM/DDIM, instead of taking multiple MCMC steps [47] following the score function, only one step is taken at each level.Training-free Guided Diffusion. One important application of Diffusion models is controlled (guided) generation. Instead of sampling from the unconditional data distribution p(x), the goal is to sample from p(x|y), where y is the usually under-specified guidance signal. For example, an animator may wish to use an unconditional diffusion model of human motion p(x) to generate motions with a constraint y that the character's right hand reaches to a specific location. Previous works [13, 11] transform the maximization of p(x t |y) at each diffusion level t with Bayes' rule:EQUATIONwhere we note that ∇ xt log p(y) = 0. Existing algorithms therefore alternate between following the score function of the trained Diffusion ∇ xt log p(x t ), and following the guidance gradient ∇ xt log p(y|x t ). However, directly optimizing p(y|x t ) is generally intractable [11] , as can be seen by the following probability factorization:EQUATIONwhere we used the fact that given x 0 , y is conditionally independent from x t . In general approximating p(x 0 |x t ) requires many denoising iterations of the Diffusion model, which is impractical when needing to alternate with optimizing p(x t ).To address this difficulty, previous works [11, 58] approximate p(y|x t ) with p(y|x 0 (x t )). Their observation is that in many practical applications, practitioners do have access to a closed-form differentiable function L(x 0 , y) that can measure how good a clean (predicted or ground-truth) sample matches the desired condition y. For example, L can simply be the mean-squared error between target and actual positions of the right hand in our aforementioned animation application. Technically, by defining L(x 0 , y) such that p(y|x 0 ) ∝ exp(-L), p(y|x 0 (x t )) can be maximized by following the gradient direction -∇ xt L(x 0 , y).Such frameworks open the door for highly flexible guided diffusion. Using the same unconditional model trained for p(x), we can now plug in various different L for different y during inference time, without having to train additional networks for each possible new y. 
trust sampling: formulating guided diffusion as optimization: 
 Our work revisits training-free guided Diffusion from the perspective of optimization. Previous works decouple the two terms p(x t ) and p(y|x t ) in Eq. 4 -they use the unconditional Diffusion model to optimize for p(x t ), and then took one gradient step of log p(y|x t ) for the constraint (guidance) term. As our experiment results will demonstrate, single gradient steps for constraints can lead to less optimal samples. With previous works mitigating this issue by carefully selecting the step sizes of ∇ xt log p(y|x t ) [58] or by better approximating p(y|x t ) [47] , we explore a new direction which leads to a robust practical algorithm across multiple domains and various constraint diffusion tasks.To start with, following the gradients of log p(y|x t ) indicates we can reformulate the constraint diffusion problem as an optimization problem:EQUATIONwhere we replace x t with x ′ to signify that the state variable x ′ can deviate from the Diffusionpredicted x t during this optimization. It is important to note that, first, for optimizing p(x t ), we still follow standard diffusion inference given its widespread empirical success in multiple domains, and second, we constrain x ′ to stay in the distribution of all possible x t at diffusion level t, as to not create a train-test discrepancy for the base diffusion model. This optimization formulation opens the door for more more flexibility in algorithm design, as we are no longer limited to taking only one gradient step.Our key improvement of this work is the use of iterative gradient-based optimization to solve Eq. 6. While only taking one single gradient step proves to be sub-optimal, as we will demonstrate in this section, optimizing until the objective saturates is also not ideal. To see this, recall that p(y|x t ) is generally not tractable, while p(y|x 0 (x t )) is. As we replace the optimization objective p(y|x t ) with the proxy p(y|x 0 (x t )), it is crucial to terminate in time before the proxy becomes a poor approximation of the true objective. Formally this relaxation can be written as:EQUATIONwhere readers are reminded that minimizing L(x 0 , y) is equivalent to maximizing p(y|x 0 ), and d is a relaxation threshold newly introduced. To reason about the gap d between true and proxy objectives, note that:p(y|x ′ ) = x0 p(x 0 |x ′ )p(y|x 0 )dx 0 = E x0∼p(x0|x ′ ) [f (x 0 )], p(y|x 0 ) = f (x 0 ) = f (E x0∼p(x0|x ′ ) [x 0 ]),where we use f (•) as a shorthand for exp(-L(• ; y)). While a similar but tedious analysis exist for general multivariant x, for the purpose of practical algorithm design, looking at the special case where x is a scalar random variable is more intuitive for understanding. Let f ′′ denote the curvature of f (x), and a = inf f ′′ and b = sup f ′′ denote the range of the curvature assuming x has a finite span, we now have f -1 2 ax 2 and 1 2 bx 2 -f both as convex functions. Applying Jensen's inequality to both functions:EQUATION)After rearranging this gives:EQUATIONThis indicates, rather intuitively, that the approximation error increases with the variance of x. As a result, we can trust the proxy optimization min x ′ L(x 0 (x ′ ), y) more when the variance of x is smaller and the proxy becomes less reliable when the variance is large. Since it is intractable to estimate the true value of the gap d during the course of optimization, we opt to design a trust schedule of maximally allowed gradient iterations that is correlated to the variance of x at each diffusion iteration t. In our experiments, we will demonstrate that simple schedules, such as a constant function g trust (t) = c, or a linear function g trust (t) = m • t + c, work surprisingly well for the diverse set of tasks we attempted.The previous section reformulates constrained guided diffusion as a gradient-based optimization, with our proposed algorithm designed to timely terminate the iterations based on the trustworthiness of the proxy objective. Equations 6 and 7 additionally require us to characterize the space that a forward sample x t can possibly visit at diffusion level t, so that we can ensure, during inference time, that x ′ will not leave the state manifold where the base model is trained on. In practice, the robustness of diffusion models can produce valid samples even if the input is slightly outside of the state manifold, allowing the constraint x ′ ∼ p(x t ) to be relaxed. However, stepping outside of the state manifolds might require more unnecessary "corrective" steps, affecting the run-time performance. To speed up the computation during inference time, we describe a method for early termination of the optimization when the sample leaves the estimated boundary of the state manifold at each diffusion step.We leverage the boundaries of a Diffusion model's intermediate state manifolds M t,δ , which we define per diffusion timestep t as the manifold on which a diffusion model has likely seen training data from with probability ≥ 1 -δ:EQUATIONGiven sufficiently small δ and a sufficiently well-trained diffusion model, the idea is that any x t ∈ M t,δ will converge to some point x 0 in the original data distribution p(x 0 ). As such, the optimization problem from Eq. 7 becomes:EQUATIONBy definition, M t,δ is a larger manifold when t is larger, meaning it gradually shrinks to true data manifold during diffusion inference. Nevertheless, M t,δ would be challenging to compute in closedform given the unknown true data distribution p(x 0 ). Our observation is that in all formulations of Diffusion models, we do have access to the model's predicted noise ϵ. For a particular x ′ , the ideal value for ϵ θ (x ′ , t) isEQUATIONIf x ′ is within the state manifold boundary, the integrandx ′ - √ αtx0 √ 1-αtfor each sample of x 0 should correspond to a multivariant Gaussian N (0, I). This implies that we can estimate the boundary ofM t,δ with ||ϵ θ (x ′ , t)||. When ||ϵ θ (x ′ , t)|| is far away from zero, x ′ - √ αtx0 √1-αt is unlikely to be sampled from N (0, I). Consequently, x ′ is likely to be outside of the state manifold at the current diffusion step. In practice, we set such a threshold ϵ max by observing the approximate average ||ϵ θ (x t , t)|| across several unconstrained samples running the base Diffusion model.Comparing with standard DDIM sampling, our Trust Sampling algorithm (Algorithm 1) takes multiple gradient steps of constraint guidance up to a maximum of J t , which denotes a max-iteration according to the trust schedule, g trust (t). We experiment with different linear schedules, detailed in the Experiments section, to show the positive impacts of our algorithm on the quality of generated data. The inner optimization loop will also be terminated by the condition when the magnitude of the predicted noise ϵ being larger than ϵ max . Following Yang et al. [58] , we normalize the gradient for numerical stability. w is a constant step size which we keep either as 0.5 or 1.0 for each specific task.Algorithm 1: Trust Sampling with DDIM Require: x T ∼ N (0, I), T , observation y, trust schedule g trust (t), norm upper bound ϵ max , guidance weight w for t = T, . . . , 1 doµ θ ← √ α t-1 x0 (x t ) + 1 -α t-1 -σ 2 t ϵ θ (x t , t) x * t-1 , j ← µ θ , 0 J t ← g trust (t) while j < J t and ||ϵ θ (x * t-1 , t)|| < ϵ max do x * t-1 ← x * t-1 -w∇ x * t-1 L( x0 (x * t-1 ), y)/||∇ x * t-1 L( x0 (x * t-1 ), y)|| j ← j + 1 end ϵ t ∼ N (0, I) x t-1 ← x * t-1 + σ t ϵ t endAdapting Inequality Constraints. We use the mean-squared value over all constraint violations to compose L in the case of equality constraints. However, we need to make an to handle inequality constraints. In the case of an inequality constraint c i (x) > a, we choose formulate L i = max(0, ac i (x)). We then compose L as the mean-squared value over all L i .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Justification: We provide short proofs for all newly introduced theoretical results in Section 3. For theoretical results shown by previous works, we cite them in Sections 2, 3, and 4.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
related work: 
 Our work is most closely related to zero-shot guided Diffusion methods for general loss functions. The seminal works of [11] and [24] introduced a method that alternates between taking one denoising step of the unconditional base diffusion model to maximize data distribution and taking one constraint gradient step to guide the model for conditional sample generation. This approach effectively balances data fidelity and conditional alignment. DSG [58] enhanced [11] by normalizing gradients in the constraint guidance term and implementing a step size schedule inspired by Spherical Gaussians.LGD-MC [47] addressed the inherent approximation errors in DPS by using multiple samples instead of a single point, which provided a better approximation of the guidance loss. Manifold Constrained Gradient (MCG) [10] and Manifold Preserving Guided Diffusion (MPGD) [19] use projections on the constraint gradient and predicted de-noised sample respectively to leverage the manifold hypothesis for better constraint following. In contrast, our work explores improving this paradigm using iterative gradient-based optimization.Various methods have been developed specifically for guided diffusion of image restoration. RED-Diff [36] extends the principles of Regularization by Denoising (RED) for image noise removal [40] to a stochastic setting, offering a variational perspective on solving inverse problems with diffusion models. Techniques such as [10, 27, 55, 14, 49] assume linear distortion models and utilize the measurement operator matrix to improve guidance for image restoration. To handle non-linear distortion models, approaches like [41] and [56] have been proposed. These methods can accommodate complex distortion but require specialized initialization schemes, which limits their general applicability. In contrast, our approach initializes from the standard unit Gaussian, ensuring broader applicability in general tasks. Similarly, ΠGDM [46] addresses inverse problems for image restoration with diffusion, but it is confined to certain loss types, while RePaint [31] enhances diffusion-based image in-painting by repeating crucial diffusion steps to improve fidelity.Recent advancements like [59] and [6] tackle guided diffusion tasks based on conditional models, with conditions including textual information. FreeDOM [59] additionally adopts an energy-based framework and generalizes the repeating strategy found in RePaint [31] with a novel time travel strategy. DiffPIR [62] balances the data prior term from the unconditional diffusion model with the constraint term from measurement loss to improve image restoration tasks.Other methods adopt additional training for controlled diffusion. Ambient Diffusion Posterior Sampling [1] builds upon DPS [11] by training the base model on linearly corrupted data. [48] learns a score function for the noise distribution, specifically targeting structured noise in images. ControlNet [60] and OmniControl [57] train additional Diffusion branches to process input constraints and conditions, achieving notable results in image or motion domains. DreamBooth [42] fine-tunes a base diffusion model to place subjects in different backgrounds using a few images, demonstrating versatility in content generation. Other notable related works include [15] , which focuses on composing multiple diffusion models. The proposed MCMC framework replaces simple gradient addition with a more robust iterative optimization process, similar to our framework for solving guided diffusion. D-PNP [17] reformulates diffusion as a prior for various guidance tasks but has been observed to struggle with more complex diffusion models, such as those trained on ImageNet [12] .Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We provide all parameters with regard to trust sampling at each interval in Section 5 and Appendix A, and we fully provide the sampling algorithm needed to reproduce our results. 
experiments: 
 We evaluate our method on two drastically different domains: images and 3D human motion. In both domains, we compare against recent zero-shot guided diffusion algorithms for solving general constraint diffusion: DPS [11] , DSG [58] , and LGD-MC [47] .Tasks. We evaluate our method on three challenging image restoration problems: Super-resolution, Box Inpainting, and Gaussian Deblurring. These common linear inverse problems are standard across DPS [11] , DSG [58] , and LGD-MC [47] ; we note that in this paper, we do not inject noise into the initial observations. Each of these image restoration problems can be thought of as a constraint satisfaction problem, where the constraint is that the generated picture appear the same as the source image upon applying the particular distortion. Distortion for these problems, respectively, was performed via (i) bicubic downsampling by 4×, (ii) randomly masking a 128 × 128 square region (sampled uniformly within a 16 pixel margin of each side), and (iii) Gaussian blur kernel of size 61 × 61 with standard deviation 3.0. We experimented on two datasets: FFHQ 256 × 256 [26] and ImageNet 256 × 256 [12] on 100 validation images each given our limited compute access. For a fair comparison between methods, we used the same pretrained unconditional diffusion models across methods for FFHQ [11] and ImageNet [13] following previous works. Quantitative evaluation of images is performed with two widely used metrics for image perception: Fréchet Inception Distance (FID) [20] and Learned Perceptual Image Patch Similarity (LPIPS) [61] .Results. Quantitative evaluation results can be seen in Tables 1 and 2 . Our method outperforms diffusion model baselines by a significant margin across all three image tasks on both FID and LPIPS and on both FFHQ and ImageNet. Qualitative results can be seen in Fig. 2 . In super-resolution, Trust Sampling shows an ability to adhere to the original down-sampled image better, even recovering text much better. In box inpainting, Trust Sampling fills in the box with realistic output; for example, the eyes in the human faces generated on the right of Fig. 2 are much more natural. Unconditional Motion Diffusion Model. For all tasks we use the same unconditional diffusion model, which we trained on the AMASS [33, 2, 28, 16, 8, 7, 30, 37, 34, 29, 53, 25, 51, 44, 3] dataset excluding the following datasets that are used for testing: danceDB [5] , HUMAN4D [9] and Weizmann [35] . The architecture is an adapted version of the EDGE motion model [52] , where we removed the branches handling conditions.Metrics. We evaluate DPS, DPS+DSG, LGD-MC, and Trust on several constrained motion generation tasks. We train an autoencoder and use the encoder as a feature extractor for motion clips, to allow for calculation of motion realism and diversity metrics [38] . We use the following metrics to evaluate performance:• FID: We extract features using the aforementioned encoder and calculate FID between different methods vs. ground-truth, as in Action2Motion [18] . • Diversity: We extract features and calculate the diversity metric as in Action2Motion [18] for the generated and ground truth motions. A result is claimed better than others if its score is closer to the score of the ground truth. • Constraint Violation: A task-dependent metric that describes how well the generated motion adheres to the provided constraints.Tasks. We first evaluate on two tasks where we have ground truth motions from the test dataset: root trajectory tracking and right hand & left foot trajectory tracking. Here the diffusion model should be guided to generate natural human movements that closely follow specified root motions or hand/foot motions. Note that the generations do not need to match the ground-truth motions due to under-specification of the constraints; we are only using them for generating the control constraints which are guaranteed to be physically feasible for human movements. Specifically, we randomly select a total of 1000 slices from the mentioned three test sets, and we extract their root motion and right hand and left ankle motion as constraint signals for the respective tasks. Note also that the observation mapping, from full motion states to the constraint signals, is highly non-linear in the hand/foot tracking task. This is because the full motion state of Diffusion only uses local joint rotations, but the hand/foot trajectory is defined in the global Cartesian space (see EDGE [52] for more details).Results. Our method strikes the best balance to matching the constraint without sacrificing realism nor diversity. DPS has the best FID score closely followed by ours. However, this comes at a large cost for DPS that violates the constraints. DSG satisfies constraints slightly better than our method, but it sacrifices both diversity and realism significantly. Our method outperforms DPS and DSG on the Diversity score for both root tracking and right hand & left foot tracking. While LGD-MC balances fidelity and constraint following better, it still has worse fidelity than Trust and struggles with harder tasks such as right hand & left foot tracking. Note that for these tasks the constraint metric is the root-mean-square tracking error in meter. The difference between ours, DSG, and LGD-MC are hardly noticeable when performing visual comparison between the generated motions, especially for root tracking. Visualizations that support these observations are in Appendix D, Fig. 6 , but are best viewed in the supplementary video.More Challenging Tasks. We further experimented our method with more difficult tasks such as sparse spatio-temporal constraints, inequality and highly non-linear constraints, and compositing multiple constraints. This is in drastic contrast to the image tasks, where a single "dense" (closer to being fully-specified) constraint must be satisfied. We designed the following tasks with additional two composite constraints on each-a translation constraint on the initial and the final frames.• Obstacle Avoidance: We add an inequality constraint to avoid penetration between any joint and three pseudo-randomly placed obstacle spheres. • Jump: We add an inequality constraint at the middle frame, to impose that all joints have a vertical position that is higher than a selected value between 0.6 m and 1.0 m. • Angular Momentum: We add an inequality constraint to impose different minimum values for the average angular momentum around a horizontal axis. This serves as a way to control dynamicism of a motion. Angular momentum is approximated as: 4 i=1 v i × p i . with v i , p i the relative velocity and position of an end effector (wrists and ankles) with respect to the root.As quantitative metrics would not be informative in these tasks (for example, it is not reasonable to compute the distributional distance between ground-truth test set and jumping motions), we focus on qualitative demonstrations. We consistently found that for easier inequality constraints (e.g. lower jumping heights) all methods could match the constraints. Howver, our method was more robust when constraints became harder, while DSG sacrificed physical realism and DPS violated the constraints. See Fig. 6 , and supplemental videos for more details on these observations.To examine the influence of Trust sampling, we performed ablations on the same three image tasks on FFHQ. In addition to FID and LPIPS, we look at the number of neural function evaluations (NFEs) as an implementation-agnostic metric of efficiency. In our case, NFEs is the number of times a pass through the pretrained model occurs.Trust Scheduling. We decouple just the trust schedule and do not use state manifold estimates for this part. The results (Table 4 ) show that our method is not sensitive to scheduling parameters as all schedules still outperform the DPS and DSG baselines on all three image tasks by significant margins. Within the different schedules, we see that linear schedules with non-zero slope (i.e. non-constant schedules) typically outperform constant schedules. This aligns with our notion of trust, as earlier diffusion steps tend to be noiser and therefore the proxy constraint function is less trustworthy, so it is less productive to take gradient steps at earlier times. Although linear trust schedule is better than constant schedules, the results indicate the best slope is dependent of the task and NFEs.Fewer NFEs. Table 4 also shows when decreasing NFEs from 1000 (same as baselines) to 600, the performance of our method barely drops and are still significantly better than baselines. To control the desired number of NFEs (1000 or 600 in our experiments), we choose a few combinations of the slope m and the offset c of the trust schedule g trust (t) = m • t + c, such that T t=1 g trust (t) equals the desired number of NFEs, where T is the total number of diffusion iterations.Manifold Boundary Estimates. We examined the effect of using manifold boundary estimates on the image tasks on FFHQ and ImageNet. We compare the effect of manifold boundary estimates when added to the trust schedule, as compared to only trust scheduling. Table 5 shows the results of using manifold boundary estimates. The use of manifold boundary reduces the needed NFEs by 10-20% without any substantial loss in quality, resulting in better compute efficiency. This performance boost is evidently robust across image task, dataset, and NFEs. Table 5 also shows that if instead of adopting manifold boundary, we want to achieve the same NFE save by tuning the start and end points of the linear schedule, model quality can suffer. Table 6 shows the effect of varying ϵ max . We observe that ϵ max generally has an acceptable range (e.g. 440-442 for FFHQ Super Resolution (4×)), within which performance varies only slightly. For the motion tasks we did not find a significant effect when introducing manifold boundary estimates. 6 : ϵ max ablations. Metrics calculated on Super Resolution (4×) on 100 validation images of FFHQ 256 × 256. To isolate purely the effect of ϵ max while keeping the number of NFEs comparable, constant linear schedules were chosen so that the number of NFEs was close to 1,000. Bold: best.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: To our best knowledge, we have provided enough detail to fully recreate our experiments. Our code will be open-sourced at the time of publication, at the GitHub provided in the abstract.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 We introduce trust sampling, a novel and effective method for guided diffusion, addressing the current limitations of meeting challenging constraints. By framing each diffusion step as an independent optimization problem with principled trust schedules, our approach ensures higher fidelity across diverse tasks. Extensive experiments in image super-resolution, inpainting, deblurring, and various human motion control tasks demonstrate the superior generation quality achieved by our method.Our findings indicate that trust sampling not only enhances performance but also offers a flexible and generalizable framework for future advancements in constrained diffusion-based modeling. To further improve generation quality, future research should adopt a holistic approach by incorporating additional concepts from traditional numerical optimization into this framework, beyond just the termination criterion. This includes techniques such as step size line search and fast approximation of higher-order derivatives. Moreover, automating the setting of heuristic parameters, which are currently manually adjusted for each base diffusion model, would be beneficial. Justification: We touched on limitations of our work in Section 6.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Our paper focuses on an inference technique on top of existing pretrained models. We specify all the parameters we used to test our new sampling method in Section 5.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: All metrics we provide are calculated on sets of 100 images or 100+ motion examples.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We have verified that all datasets used (FFHQ, ImageNet, AMASS) are licensed to allow non-commercial scientific research. We believe there to be a small, but non-zero, risk for the usage of our method for deceptive interactions such as deepfake images, but because our method does not really optimize for recreating particular people in other contexts, we believe the risk is low. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Diffusion is a core technique that enables many of the applications within Generative AI, many of which have positive benefits on society. Being able to precisely control the generations of diffusion models can improve the safety and accuracy of those applications. Although we do not see specific negative impacts of our work, there is always potential for misuse of our technology. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.Answer: [NA] Justification: We do not release any new assets along with this paper. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: • The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: All reviewers are positive  on the paper after rebuttal.
They find the paper well motivated and the method original albeit simple.
The rebuttal addressed most concerns

Reviews: review: paper_summary: paper_summary: This paper tackles the task of sampling from diffusion models with additional inference-time constraints. In this setting, synthesis needs to simultaneously follow the diffusion model-defined generative prior as well as a constraint objective. The paper proposes two techniques to achieve this in a robust fashion. First, trust schedules define the number of optimization steps with respect to the constraint that are carried out between each diffusion model generative denoising step. Moreover, state manifold boundaries prevent the model from falling off the diffusion model-defined data manifold during optimization. These techniques enable more robust and stable constraint-guided sampling of diffusion models. The proposed methods are validated on constrained image modeling tasks (superresolution, deblurring, inpainting) as well as several human motion synthesis benchmarks, and favourable results compared to baselines are achieved.
strength_weakness: strength_weakness: I think there are two main weaknesses:

First, both the trust schedules and also the approach to estimate state manifold boundaries are merely well-motivated heuristics and they both require hyperparameter tuning (how many steps in the trust schedule, $\epsilon_{max}$). For the trust schedules, there is an elaborate derivation that relates approximation errors to state variances (Eq. (9)). But in the end, none of this is tractable and it just motivates the use of a step-dependent optimization schedule, which needs to be chosen fully heuristically. And this schedule also depends on the optimization step size, which needs to be chosen manually, too. The state manifold boundary constraint follows a similar $\epsilon_{max}$ heuristic. In the end, the paper proposes some useful and well-motivated optimization heuristics, but not an overly novel and rigorous framework for constrained diffusion model sampling. This makes the paper somewhat less original.

Second, as the authors pointed out in their related work section, there are many previous papers in the area. However, the authors only compare to DPS as well as DPS+DSG, although there should be further applicable baselines. Most importantly, the authors argue that LGD-MC has no code available. However, I believe reimplementing the LGD-MC method should be quite easy, and a comparison to this work would be very appropriate here. Moreover, maybe LGD-MC could even be combined with, or enhanced by the techniques proposed here. This possibility is not discussed.

**Conclusion:** In conclusion, even though I see some weaknesses and evaluation could be expanded, the proposed heuristics seem useful and the paper is of overall satisfactory quality. Hence, despite my concerns, I am carefully leaning towards suggesting acceptance.
questions: questions: I have one question: In line 173, the authors write that they replace max with ReLU to have a smooth gradient. What is meant by this? ReLUs do not have a smooth gradient, but it is a step function.

I do not have any additional questions, but here are a few typos:

- line 30: constrained -> constraints
- line 44: it should be "...state manifold of the diffusion model..."
- line 76: I believe the score function misses a $\log$.
- line 182: I believe it should be "LGD-MC".
- line 251: "first first"

One additional comment: In line 65, DDPM and DDIM sampling are described as forms of gradient-based MCMC sampling. I do not think this is generally quite correct. For instance, if one choses a deterministic DDIM setup without noise injection, then this is a regular probability flow without stochasticity, but not MCMC. MCMC typically always has stochasticity with some form of noise injection. Broadly speaking, one can see DDPM and DDIM always as a flow plus a variable MCMC component, but not always as pure MCMC. I would suggest the authors to rephrase this statement.
limitations: limitations: The paper does not have a detailed discussion of its limitations or its societal impact. This did not significantly influence my paper rating, but I would strongly suggest the authors to add a more detailed discussion on both of these aspects to the paper.
paper_summary: paper_summary: This paper proposes a trust sampling scheme which incorporates given constraints loss function as guidance for constrained generation. This approach conducts early stop while detecting a mismatch between the predicted noises magnitude of sample and the noisy level at the current state manifold at each diffusion step. In experiment section, this approached is applied to image generation and human motion tasks both across various applications.
strength_weakness: strength_weakness: 1. It seems that there is no guarantee that the final samples satisfy the constraints since the reverse process will go to the next diffusion noisy level from the early stop if the predicted noise is higher than the threshold while the loss is not 0 yet, which is also reflected in the human motion experiment.
2. I get the idea that early stop is determined by the norm of predicted noise from the section 3.2: if the norm is large, then $\mathbf{x_t}$ does not reside in its supposed manifold. But is the reverse direction valid: if the norm is small than some threshold, than $\mathbf{x}_t$ is on the manifold? Is this necessarily true?
3. To summarize above, I found the explanation from methodology section tends to give intuitions about why this algorithm would work instead of guarantees or theoretical analysis, such as a bound on the distance between the $\mathbf{x}_t$ before and after applying loss gradients multiple times, or an upper limit of the probability of samples being out of constraints, etc..
questions: questions: - In the algorithm, the early termination happens when predicted noise level is higher than the threshold. Shouldn't this be checked at the end of previous step and not applying the gradient from there since the current $\mathbf{x}_{t-1}^*$ already falls off the manifold? Kindly correct me if I am wrong. 
- I am not sure if image generation tasks are suitable for this paper. I do not see the constraints but simply comparisons with other baselines for checking the performance of the algorithm.
-  Maybe an ablation study is required here to see how the performance is affected by $\epsilon_{\max}$ instead of choosing $\epsilon_{\max}$ from the samples generated by base diffusion models since it determines when early stop should be.
- Could be more typos but not limited to: from Line 172, ``However, we need to make an to handle inequality constraints``, which affects the reading.
limitations: limitations: Based on the checklist guidelines, the authors claimed the positive and negative societal impact are discussed, but I only found the ``justification`` below the question in checklist is the only place they talked about the impact and I assume this should be mentioned somewhere in the main text or appendix. Limitation is also claimed being discussed in the section 6, and I assume that's mentioned in the second paragraph.
paper_summary: paper_summary: The paper presents a method to enhance training-free loss-guided diffusion sampling. The key contributions are:

1. Introduction of Trust Sampling: A novel method called Trust Sampling is proposed, which diverges from the traditional approach of alternating between diffusion steps and loss-guided gradient steps. Instead, it treats each timestep as an independent optimization problem, allowing multiple gradient steps on a proxy constraint function at each diffusion step.

2. Early Termination and State Manifold Estimation: The method includes a mechanism to estimate the state manifold of the diffusion model, enabling early termination if the sample starts to deviate significantly from the expected state. This ensures the proxy constraint function remains trustworthy.

3. Optimization Perspective: The paper reformulates training-free guided diffusion as a constrained optimization problem, providing more flexibility and robustness across various domains and tasks. This approach addresses the limitations of previous methods, such as sensitivity to initialization and performance degradation with fewer neural function evaluations.

4. Experimental Validation: The efficacy of Trust Sampling is demonstrated through extensive experiments in different domains, including image and 3D motion generation. The method shows significant improvements in generation quality and constraint satisfaction compared to existing techniques.

5. Generalization Across Tasks: The paper demonstrates the generality of Trust Sampling across various image tasks (e.g., super-resolution, inpainting, gaussian deblurring) and motion tasks (e.g., trajectory tracking, obstacle avoidance), highlighting its versatility and effectiveness.
strength_weakness: strength_weakness: 1. The Trust Sampling algorithm, which takes multiple gradient steps of constraint guidance on each predicted $x_{t}$, is similar to the  corrector stage of PC sampling [1]. 
2. Training-free loss-guided diffusion sampling have been applied not only on inverse problem but also on variable tasks [2-6], such as refined text-to-image and layout-to-image. The authors need to further verify the effectiveness of Trust Sampling on these tasks. 

[1] Score-based Generative Modeling through Stochastic Differential Equations

[2] Counting Guidance for High Fidelity Text-to-Image Synthesis

[3] Fine-grained Text-to-Image Synthesis with Semantic Refinement

[4] Attend-and-Excite: Attention-Based Semantic Guidance for Text-to-Image Diffusion Models

[5] BoxDiff: Text-to-Image Synthesis with Training-Free Box-Constrained Diffusion

[6] LoRA-Composer: Leveraging Low-Rank Adaptation for Multi-Concept Customization in Training-Free Diffusion Models
questions: questions: Please see weaknesses.
limitations: limitations: The authors have adequately described the limitations and potential negative societal impact of their work.
paper_summary: paper_summary: This paper addresses this limitation by rethinking diffusion without training loss guidance from an optimization perspective. They formulate a series of constrained optimizations throughout the inference process of the diffusion model. In each optimization, they allow the sample to take multiple steps along the gradient of the surrogate constraint function. The termination conditions come from two aspects: one is the accuracy of the approximate surrogate, and the other is the estimation of the manifold.
strength_weakness: strength_weakness: 1. First of all, I have doubts about the theory of this article. Equation 12 in the article seems to be written incorrectly. As far as I know, $\epsilon_\theta(\mathbf{x}',t)=\int(\mathbf{x}'-\sqrt{\alpha_t}\mathbf{x_0})/(\sqrt{1-\alpha_t})p(\mathbf{x}_0|\mathbf{x}')\mathrm{d}\mathbf{x}_0$, it should be conditional expectation here.

2. The motivation for this part is also not very clear. Why does the above integral correspond to a multivariate Gaussian (Line 156)?

3. The number of test samples used in the image experiment is a bit small. Why didn't the author use a data set similar to DPS?

4. In addition, the proposed method still seems to require more than 500 NFE, similar to DPS, and I am a little worried about its efficiency and the need for improvement.

5. DPS itself requires careful tuning of hyperparameters, such as the guidance rate. The proposed method introduces new hyperparameters $m·t+c$ and $\epsilon_max$, which raises scalability concerns.
questions: questions: see Weaknesses.
limitations: limitations: The authors discuss the limitations of their work.

Comments: title: Response to author
Reply to Rebuttal
comment: We thank the reviewer for their continued effort in helping us improve this manuscript. We are pleased that our efforts to address the previous concerns have been acknowledged.

Regarding the new issue raised, we apologize for the omission of the related work [1]. However, after carefully reviewing we believe that our paper's contributions remain distinct. The “key idea” of [1] is that “the information bottleneck in autoencoders naturally incorporates the manifold hypothesis.” While this approach is relevant for many modern image diffusion networks, which often utilize autoencoders for dimensionality reduction, it does not apply as readily to other domains, such as 3D human motion we experimented with in this paper. Indeed, [1] is only tested on image models.

Our method, like DPS, DSG, and LGD-MC, is designed to be domain-agnostic. While [1] proposes a weaker variant that could be applied across domains—by changing \nabla_x_t to \nabla_x_0 in the standard loss guidance equation and using the gradient to update x_0 instead of x_t (note that the loss L is always computed on x0 regardless)—this is a commonly known implementation choice among researchers in the field, for example mentioned in the first two paragraphs of Section 4.2 of https://arxiv.org/pdf/2305.12577. We also experimented with this variant early in our project but did not observe any noticeable improvements for our test problems, possibly also due to the potential issues outlined in the aforementioned paper.

We want to assure all reviewers that this omission was an honest mistake and not a deliberate oversight. We made every effort to include relevant papers by utilizing common online tools. However, when we reviewed the papers citing the seminal work DPS and the state-of-the-art paper LGD-MC, Google Scholar did not list [1]. Only now we could realize that this may have been due to a technical issue with Google Scholar not indexing the paper correctly: https://scholar.google.com/scholar?hl=en&as_sdt=2005&sciodt=0%2C5&cites=696239910969416231&scipsc=1&q=manifold+preserving&oq=

In summary, we believe that the added value of our manuscript to the research community remains intact, with or without [1]. Our work provides new theoretical insights, introduces a novel method which can potentially be combined with [1] for image domain problems, and presents a set of new experimental results.

We will include a more thorough discussion and comparison with this paper in the next revision.
Thanks to the author for the great effort put into the rebuttal, which has alleviated some of my concerns. 

However, I found some work that also constrains DPS [1] during the rebuttal, which seems to be highly relevant to this paper, as it aim to constrain diffusion on a specified manifold. The author seems to have ignored these developments, which may need further discussion.         

[1] MANIFOLD PRESERVING GUIDED DIFFUSION, ICLR 2024.
Hello! With the rebuttal discussion period ending soon, we would like to kindly ask if our rebuttal addresses your questions, and if you have additional questions or concerns. Thank you for your time and effort on this review!
After going through the rebuttals, I decide to increase the score from 4 to 5. Thank you!
I acknowledge having read the authors' rebuttal. My overall assessment of the paper remains unchanged, and I continue to support my current rating.
I would like to acknowledge that I have read the authors' rebuttal. Overall, my impression of the paper remains and I maintain my current, positive-leaning rating.