Paper deep dive

Integration Matters: Rollout-Based Training for Constrained Diffusion Models

Xiaoxuan Liang, Saeid Naderiparizi, Berend Zwartsenberg, Frank Wood

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 92

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%

Last extracted: 7/18/2026, 2:33:54 AM

Summary

The paper introduces IMCD (Integration Matters for Constrained Diffusion), a fine-tuning framework that aligns diffusion model training with sampling by incorporating constraint guidance from online rollouts. It addresses training-time misalignment and sampling-time distribution shift by differentiating through a fixed noise schedule, utilizing a learned scaling network for adaptive guidance, and retaining EDM denoising loss to preserve data fidelity. Experiments on bouncing ball and traffic scene trajectory prediction tasks demonstrate improved constraint satisfaction while maintaining competitive sampling quality.

Entities (10)

IMCD · method · 98%Online Rollout · technique · 94%Training-time Misalignment · problem · 93%Constrained Diffusion Models · model-type · 92%Learned Scaling Network · component · 91%EDM Denoising Loss · loss-function · 89%Sampling-time Correction · approach · 88%Bouncing Ball Task · evaluation-task · 87%Traffic Scene Trajectory Prediction · evaluation-task · 86%LoRA Adapters · technique · 85%

Relation Signals (9)

IMCD → addresses → Training-time Misalignment

confidence 95% · We present IMCD ... that addresses the training-time misalignment.

IMCD → uses → Online Rollout

confidence 94% · We propose a fine-tuning framework that incorporates constraint guidance obtained through online rollout into the training process

IMCD → improves → Constraint Satisfaction

confidence 92% · Experiments across multiple tasks show that our method improves constraint satisfaction while maintaining competitive sampling quality compared to prior methods.

IMCD → incorporates → Learned Scaling Network

confidence 91% · We instead propose to learn a scaling function that adapts the correction dynamically to both the current state and diffusion time.

Training-time Optimization → suffersfrom → Training-time Misalignment

confidence 90% · Training-time optimization approaches optimize on states induced by the training distribution, which can differ substantially from those encountered during sampling.

IMCD → retains → EDM Denoising Loss

confidence 89% · To preserve distributional fidelity, we retain the EDM denoising loss as a regularizer.

Sampling-time Correction → causes → Distribution Shift

confidence 88% · Sampling-time correction methods instead modify the sampling process at inference, introducing distribution shift and requiring expensive tuning

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:Constrained generative models aim to produce samples that satisfy complex feasibility constraints while remaining faithful to the data distribution. Existing constrained generation methods typically enforce constraints either through training-time optimization or sampling-time correction. Training-time optimization approaches optimize on states induced by the training distribution, which can differ substantially from those encountered during sampling. Sampling-time correction methods instead modify the sampling process at inference, introducing distribution shift and requiring expensive tuning, particularly for few-step sampling. We propose a fine-tuning framework that incorporates constraint guidance obtained through online rollout into the training process, which aligns training with sampling by differentiating through the fixed noise schedule used to numerically integrate the denoising process. This exposes the model to violations that arise along the denoising trajectory and aligns diffusion learning with the sampling process. Experiments across multiple tasks show that our method improves constraint satisfaction while maintaining competitive sampling quality compared to prior methods.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

91,061 characters extracted from source content.

Expand or collapse full text

Integration Matters: Rollout-Based Training for Constrained Diffusion Models Xiaoxuan Liang1 2 Saeid Naderiparizi1 2 Berend Zwartsenberg2 Frank Wood1 2 3 1University of British Columbia 2Inverted AI 3Alberta Machine Intelligence Institute liang51@cs.ubc.ca Abstract Constrained generative models aim to produce samples that satisfy complex feasibility constraints while remaining faithful to the data distribution. Existing constrained generation methods typically enforce constraints either through training-time optimization or sampling-time correction. Training-time optimization approaches optimize on states induced by the training distribution, which can differ substantially from those encountered during sampling. Sampling-time correction methods instead modify the sampling process at inference, introducing distribution shift and requiring expensive tuning, particularly for few-step sampling. We propose a fine-tuning framework that incorporates constraint guidance obtained through online rollout into the training process, which aligns training with sampling by differentiating through the fixed noise schedule used to numerically integrate the denoising process. This exposes the model to violations that arise along the denoising trajectory and aligns diffusion learning with the sampling process. Experiments across multiple tasks show that our method improves constraint satisfaction while maintaining competitive sampling quality compared to prior methods. 1 Introduction In many real-world applications, generative tasks are defined not only by the underlying data distribution but also by complex feasibility constraints that govern the validity of the sample space. In robotics systems, generated states should remain geometrically feasible and avoid collisions with the environment and other agents [40, 6]. In autonomous driving scenarios, predicted trajectories should remain on the road [47, 54]. Diffusion models [20, 44, 23] provide a powerful framework for high-quality generation [39, 21, 17], but standard denoising objectives do not explicitly enforce feasibility constraints. As a result, generated samples may violate constraints that are critical for deployment. Specifically, standard diffusion models are trained with denoising objectives evaluated at independently sampled noise levels, without directly enforcing constraints on the final sample. Existing constrained diffusion models address this gap through constraint-aware fine-tuning [2, 32, 26] or post-training guidance [11, 9, 19, 50, 30, 15, 49, 8, 25]. Fine-tuning methods incorporate constraint guidance during training, but often optimize constraint losses on forward noising states, which can differ from the denoising states encountered during sampling. Sampling-time correction methods instead apply guidance only at inference, forcing corrections on states the model was not trained to handle [10, 48]. We tackle constrained generation with a fine-tuning approach that aligns training with sampling. Instead of enforcing constraints only through sampling-time corrections, we optimize the model through online denoising rollouts that match the inference procedure. Constraint violations are measured on the terminal generated samples, and the resulting objective trains the model on the states it encounters during sampling. To preserve distributional fidelity, we retain the EDM denoising loss as a regularizer. Our contributions are: • We introduce a fine-tuning framework that optimizes constraint satisfaction through the reverse denoising trajectory. • We provide formal analysis of the training-time alignment induced by our training objective and show that the terminal denoising state converges towards the feasible region. • We validate our method on two constrained generation tasks: bouncing ball and traffic scene trajectory prediction experiments. 2 Background Diffusion model generation Diffusion models [42, 43, 20, 44] generate samples by gradually corrupting data with noise and then learning to reverse this process through denoising. Formally, the forward process perturbs data 0∼p0x_0 p_0 into a sequence of increasingly noisy states tx_t, while the reverse process reconstructs clean samples from noise. The forward process is defined by the Ito stochastic differential equation (SDE): dt=f(t;t)dt+g(t)dt, dx_t=f(x_t;t)\,dt+g(t)\,dw_t, (1) where f(⋅;⋅)f(·;·) and g(⋅)g(·) are drift and diffusion coefficients, and tw_t is a standard Wiener process. The corresponding reverse-time SDE is given by [1]: dt=[f(t;t)−g2(t)∇log⁡pt(t)]dt+g(t)d¯t. dx_t=[f(x_t;t)-g^2(t)∇ p_t(x_t)]\,dt+g(t)\,d w_t. (2) where ~t w_t denotes a Wiener process evolving backward in time. Since the score function ∇tlog⁡pt(t) _x_t p_t(x_t) is intractable, diffusion models learn a neural approximation sθ(t;t)s_θ(x_t;t) trained via denoising score matching (DSM) objective [46], which minimizes ℒDSM(θ)=t,0,t[λ(t)∥sθ(t;t)−∇tlog⁡pt(t)∥2], L_DSM(θ)=E_t,x_0,x_t [λ(t) s_θ(x_t;t)- _x_t p_t(x_t) ^2 ], (3) where t∼p(t∣0)x_t p(x_t x_0) is sampled from the forward process. For the Gaussian perturbation process used in EDM [23], consider the diffusion-only SDE dt=g(t)dt. dx_t=g(t)\,dw_t. (4) It has marginals p(t∣0)∼(0,σ2(t))p(x_t x_0) N(x_0,σ^2(t) I), σ2(t)=∫0tg2(s)sσ^2(t)= _0^tg^2(s)\,ds, where σ(t)σ(t) denotes noise level at time t. Choosing g(t)=2tg(t)= 2t leads to t∼(0,t2)x_t N(x_0,t^2 I). Thus the noise level satisfies σ(t)=tσ(t)=t, and we use t to denote the noise scale throughout. This equivalence is standard [44, 23]. Through Tweedie’s formula [13], score estimation can be recast as a denoising problem. Given a denoiser network Dθ(t;t)D_θ(x_t;t) that approximates [0∣t]E [x_0 x_t ], the score can be estimated as ∇tlog⁡pt(t)=[0∣t]−t2≈Dθ(t;t)−t2. _x_t p_t(x_t)= E [x_0 x_t ]-x_tt^2≈ D_θ(x_t;t)-x_tt^2. (5) Constrained diffusion generation In many applications, generated samples are required to satisfy problem-specific constraints that define a feasible set Ω⊆ℝd R^d. These constraints often do not admit an explicit analytic boundary and are only implicitly specified through task-specified functions. Let ℓΩ() (x) denote a constraint violation function such that, for a sample ∈ℝdx∈R^d, ℓΩ()=0 (x)=0 if ∈Ωx∈ and ℓΩ()>0 (x)>0 otherwise. We assume that the data distribution is supported on the feasible set, supp(p0)⊆Ωsupp(p_0) . Constrained diffusion generation seeks to produce samples through a reverse process that remains faithful to p0p_0 while satisfying the feasibility condition of lying in Ω . However, learning such constraints from data alone is challenging, as training data only consists of feasible examples and lacks explicit information about infeasible regions [37]. Prior work [31] further shows that increasing dataset size alone does not eliminate constraint violations, leaving the model’s behavior outside the data support and instead determined by inductive biases. Consequently, maximum-likelihood training alone cannot achieve high constraint satisfaction, motivating the need for additional signals that explicitly penalize infeasible generations. 3 Limitations of existing approaches to constrained diffusion 3.1 Training-time misalignment in constraint fine-tuning Prior constrained diffusion methods [2, 26] derive their training signals from constraint loss evaluated at Dθ(t;t)D_θ(x_t;t), where t∼q(t∣0)x_t q(x_t x_0). At low noise levels, tx_t remains close to the data manifold and is typically near-feasible. Since the pretrained denoiser trained on feasible data tends to recover specific training samples [4], Dθ(t;t)D_θ(x_t;t) also tends to satisfy the constraints. This makes ℓΩ(Dθ(t;t))≈0 (D_θ(x_t;t))≈ 0 and causes the correction signal to vanish. In contrast, sampling starts from Gaussian noise and follows discrete denoising steps. High-noise samples often introduce violations under uninformative guidance. These violations can persist until low noise, where correction becomes reliable. However, low-noise infeasible states produced by reverse rollouts are rarely encountered during training. This leads to a misalignment: the model is trained mostly on near-feasible states, yet must correct infeasible rollout states at inference. This is illustrated in Figure˜1. We evaluate the constraint loss of the one-step lookahead Dθ(t;t)D_θ(x_t;t) at matched noise levels, where tx_t is obtained either from the forward process or reverse sampling using the EDM model later evaluated in Section˜5.1. While both exhibit similar behavior at high noise levels, they diverge significantly at low noise. Forward states remain near-feasible and yield near-zero constraint loss, whereas reverse states incur persistent violations. This indicates that training on forward states underestimates the violations encountered during sampling. Prior work has studied the discrepancy between forward sampling states tx_t and reverse sampling states ~t x_t [35, 24]. We examine its effect by comparing constraint losses on the denoised lookahead Dθ(t;t)D_θ(x_t;t) and Dθ(~t;t)D_θ( x_t;t). This shows that constraint signals remain weak on forward sampled states, but become nontrivial along the reverse denoising trajectory. Figure 1: Infraction loss of the one-step lookahead prediction as a function of intermediate state tx_t and noise level t, comparing forward process samples and reverse-time sampling states. Shaded regions denote ±0.2± 0.2 standard deviations. (a) Generated sample (b) Ground truth sample Figure 2: Generated versus ground-truth bouncing ball trajectories. 3.2 Sampling-time correction: distribution shift and few-step degradation Another common approach enforces constraints only at sampling time, using either projection [8, 25] or gradient-based updates [19] dt=[f(t;t)−g2(t)(sθ(t;t)−γ(t)∇ℓΩ(t;t))]dt+g(t)d¯t, dx_t= [f(x_t;t)-g^2(t)(s_θ(x_t;t)-γ(t)∇ (x_t;t)) ]\,dt+g(t)\,d w_t, (6) where γ(t)γ(t) controls the guidance strength. These methods correct violations online during denoising and can reduce the infeasibility rate in practice, but they are highly sensitive to the guidance schedule [9]. Weak guidance may fail to remove violations, while strong guidance can push intermediate state tx_t away from the data manifold [11], which places samples in regions where the learned score is unreliable. This further widens the misalignment between forward sampled and reverse sampling states. Moreover, because sampling-time correction relies on incremental corrections across the denoising trajectory, it becomes sensitive to sampling budget: many steps allow small corrections to accumulate, while few steps require stronger updates that can distort or destabilize sampling [41]. Figure˜2 illustrates this behavior using trajectories generated by a sampling-time correction method called MPGD without projection [19] under 50-step sampling setting later used in Section˜5.1. The full-scene and zoomed views use hollow and solid circles denoting start and end positions. With gray balls shown for context, balls 1, 3, and 4 are highlighted because of their jittering motion, which produces sharp trajectory changes that avoid collision but distort the motion. In contrast, the ground truth rollout remains smooth and physically consistent. Thus, this method enforces feasibility without preserving realistic motion. These observations reveal two limitations. Training-time optimization suffers from misalignment between the states encountered from training and sampling distributions. Sampling-time correction introduces distribution shift and becomes fragile under few-step sampling. Together, these limitations motivate training on reverse-time rollout states and learning constraint-aware updates that enforce feasibility. 4 Methodology We present IMCD (Integration Matters for Constrained Diffusion), a rollout-based fine-tuning framework for constrained generative modeling that addresses the training-time misalignment. IMCD evaluates constraint losses at terminal rollout samples and backpropagates through the full denoising trajectory for model updates, aligning optimization with inference. During each rollout step, it computes constraint-aware corrections using denoised lookahead estimates and applies them through a learnable scaling mechanism, enabling adaptive guidance across noise levels. To preserve data fidelity, IMCD retains the EDM loss and updates only lightweight trainable components. We assume access to a pretrained denoiser denoted by D, where the pretrained model is parameterized by fixed model weights θ. Since θ remains frozen throughout fine-tuning, we omit it from the notation. The fine-tuned denoiser network is initialized from θ and denoted by DϕD_φ, where ϕφ collects the lightweight trainable components introduced during fine-tuning. The score functions induced by these denoisers are denoted by s(t;t)s(x_t;t) and sϕ(t;t)s_φ(x_t;t), respectively. 4.1 Constraint Guidance Mechanism A key design choice in constrained diffusion is how to construct a constraint-aware score estimator. Generalizing the definitions in [32], we consider estimators of the form sϕΩ(t;t)=sϕ(t;t)−γ(t;t)GΩ(t;t), s_φ (x_t;t)=s_φ(x_t;t)-γ(x_t;t)G (x_t;t), (7) where GΩG produces a direction that reduces violations of the constraint Ω and γ controls its strength. We call GΩG the guidance function and γ the scaling function. Table˜3 summarizes the definition of these functions in our method and in prior work. Guidance function We follow Liang et al. [26] for the choice of guidance function GΩ(t;t):=∂1ℓΩ(D(t;t);t)G (x_t;t):= _1 (D(x_t;t);t) where ∂1 _1 denotes the partial derivative with respect to the first argument. Effectively, this computes the gradient of the constraint loss function at the one-step denoising prediction by a pretrained denoiser D. Liang et al. [26] show that this leads to a more stable model (both during training and sampling) than computing the gradients at the noisy state tx_t. It is important to note that ∂1ℓΩ(D(t;t);t) _1 (D(x_t;t);t) differs from ∇tℓΩ(D(t;t);t) _x_t (D(x_t;t);t). The latter corresponds more directly to the diffusion formulation, since it fully accounts for the dependence of the denoised prediction on the noisy latent tx_t by backpropagating through the denoiser D. In contrast, ∂1ℓΩ(D(t;t);t) _1 (D(x_t;t);t) ignores this dependency and does not backpropagate gradients through D, making it significantly more computationally efficient in practice. Scaling function The scaling γ used in prior work [33, 26] is independent of the sample tx_t and is typically selected via grid search. In practice, this tuning procedure can be computationally expensive, as the optimal value of γ depends on the sampling regime, noise level, and constraint strength. We instead propose to learn a scaling function that adapts the correction dynamically to both the current state tx_t and diffusion time t. This eliminates the need for costly hyperparameter sweeps. This parameterization is also motivated by the exact gradient ∇tℓΩ(D(t;t);t) _x_t (D(x_t;t);t). In particular, the discrepancy between ∂1ℓΩ(D(t;t);t) _1 (D(x_t;t);t) and the exact gradient induced by backpropagation through the denoiser can be interpreted as a state- and time-dependent rescaling of the guidance direction. More precisely, ∇tℓΩ(D(t;t))=[∂D(t;t)∂t]⊤∂1ℓΩ(D(t;t);t). _x_t (D(x_t;t))= [ ∂ D(x_t;t)∂x_t ] _1 (D(x_t;t);t). (8) Prior work [13, 5] shows that the Jacobian scales with a posterior covariance term of the form ∂D(t;t)/∂t≈cov[0∣t]/t2∂ D(x_t;t)/∂x_t [x_0 x_t ]/t^2. The covariance term captures state-dependent uncertainty, while the denominator amplifies corrections as the sample approaches the data manifold. Definition 4.1. (Learned scaling network) Given a trainable score sϕs_φ, a constraint violation function ℓΩ(t;t) (x_t;t), and a learned scaling network γϕ:ℝd×[0,T]→ℝ+ _φ:R^d×[0,T]→R^+, we define: γϕ(t;t)=αϕ(t;t)⋅tβϕ(t;t), _φ(x_t;t)= _φ(x_t;t)· t _φ(x_t;t), (9) where αϕ(t;t)≥0 _φ(x_t;t)≥ 0, βϕ<−1 _φ<-1. The constrained score is then sϕΩ=sϕ(t;t)−γϕ(t;t)∂1ℓΩ(D(t;t);t), s_φ =s_φ(x_t;t)- _φ(x_t;t) _1 (D(x_t;t);t), (10) This parameterization keeps the simple power-law guidance form but makes it adaptive to both the current state tx_t and noise level t. Proposition 4.2. (Convergence property) The constrained score 10 corresponds to the score of pΩ(t;t)∝p(t;t)exp⁡(−γϕ(t;t)ℓΩ(D(t;t))). p (x_t;t) p(x_t;t) (- _φ(x_t;t) (D(x_t;t))). (11) Furthermore, γϕ(t;t)→∞ _φ(x_t;t)→∞ as t→0t→ 0 for any tx_t with αϕ(t;t)>0 _φ(x_t;t)>0, such that limt→0pΩ(t;t)∝limt→0p(t;t)Ω(t), _t→ 0p (x_t;t) _t→ 0p(x_t;t) 1_ (x_t), (12) preserving stronger corrections near the data manifold while allowing the magnitude to adapt to the current state and noise level. The discussion above considers a single constraint gradient. For tasks with multiple constraints, we combine their gradients into a single compatible update direction inspired by PCGrad [51], then apply the scaling network to this direction to learn a global correction magnitude. We provide the combination procedure and scaling-network architecture in Appendix Appendix˜C. 4.2 Rollout-based training Now we address the training-time misalignment discussed in Section˜3.1. We construct training states by rolling out the reverse-time process from pure Gaussian noise. The training-time rollout shares the same noise sampling schedule as inference such that the intermediate states are generated by the fine-tuned network DϕD_φ. We use (t;t)(x_t;t) for continuous-time notation and adopt step-indexed notation (i,ti)(x_i,t_i) for the discretized reverse-time rollout. Let the noise schedule be T=tN>tN−1>⋯>t0=tmin>0T=t_N>t_N-1>…>t_0=t_ >0. Starting from ~N∼(0,tN2) x_N N(0,t_N^2 I), we simulate denoising steps backward along this schedule and denote the resulting reverse rollout trajectory by ~ii=1N\ x_i\_i=1^N. At each rollout step i, we compute the denoised state guidance ∂1ℓΩ(D(~i;ti);ti) _1 (D( x_i;t_i);t_i) using the pretrained denoiser. The fine-tuned denoiser DϕD_φ is conditioned on this guidance through a lightweight embedding, and the scaling network γϕ(~i;ti) _φ( x_i;t_i) defined in Equation˜9 determines its strength at the current rollout state. We also introduce LoRA adapters [22] within DϕD_φ to provide adaptation capacity while keeping the pretrained backbone weights fixed. This gives the constraint-aware denoiser: DϕΩ(~i;ti)=Dϕ(~i;ti,∂1ℓΩ(D(~i;ti);ti))+γϕ(~i;ti)∂1ℓΩ(D(~i;ti);ti). D_φ ( x_i;t_i)=D_φ ( x_i;t_i, _1 (D( x_i;t_i);t_i) )+ _φ( x_i;t_i) _1 (D( x_i;t_i);t_i). (13) We use this constraint-aware denoiser DϕΩD_φ at each rollout step. Thus the model receives the constraint signal at its input while also applying an output-level correction with a learned, state-dependent strength. Training on the resulting rollout states ~i\ x_i\ under the same discretization schedule as inference exposes the model to off-manifold states produced by its own sampling process. This alignment directly targets the training-time misalignment identified in Section˜3.1, enabling the model to correct constraint violations during generation. 4.3 Training objective and sampling procedure We optimize the fine-tuning parameters ϕφ with an EDM denoising term and a terminal rollout constraint term, where κ balances constraint optimization against data fidelity: ℒ(ϕ)=ℒEDM(ϕ)+κℒrollout(ϕ), L(φ)=L_EDM(φ)+κL_rollout(φ), (14) EDM denoising loss To preserve data-consistent behavior, we retain the EDM denoising objective for the constraint-augmented denoiser DϕΩD_φ : ℒEDM(ϕ)=t,0,t[λ(t)∥DϕΩ(t;t)−0∥2], L_EDM(φ)=E_t,x_0,x_t [λ(t) D_φ (x_t;t)-x_0 ^2 ], (15) where t∼q(t∣0)x_t q(x_t x_0) is sampled from the forward process independently of the rollout, and λ(t)λ(t) is the EDM loss weighting. This objective regularizes the model toward data-consistent denoising while allowing deviations required to satisfy the constraints. Rollout constraint loss For constraint optimization, constraint loss is applied at the terminal rollout: ℒrollout(ϕ)=p(~T),pϕ(~0∣~T)[ℓΩ(~0)]. L_rollout(φ)=E_p( x_T),p_φ( x_0 x_T) [ ( x_0) ]. (16) During optimization, the sampled states ~i\ x_i\ are not updated by gradients, and optimization is only performed with respect to ϕφ through the fine-tuned denoiser evaluations. For efficiency, we adopt the EDM discretization [23], which achieves high sampling quality with fewer denoising steps. In addition, we employ gradient checkpointing [7] to control memory when propagating through rollout trajectories. Since supervision is applied only at the terminal state, the constraint signal is sparse and becomes sparser as the infraction rate decreases. We therefore adaptively weight the rollout loss: κ=sg(ℒEDM¯)/sg(ℒrollout¯+ϵ),ϵ=10−5 κ=sg( L_EDM)/sg( L_rollout+ε), ε=10^-5 (17) where the overline denotes the batch mean, sg(⋅)sg(·) stops gradients through the scaling factors. The adaptive stop-gradient weight keeps the rollout loss on the same scale as the EDM loss, maintaining effective constraint supervision without introducing additional gradient pathways. We summarize the full training procedure in Algorithm˜1. At inference time, sampling follows the same reverse-time rollout, with the same noise discretization schedule used during training. The rollout starts from a Gaussian initialization at the highest noise level and iteratively denoises to obtain the final sample. The constraint-aware denoiser DϕΩD_φ is used for sampling, while the pretrained denoiser D provides the denoised estimate for constraint gradient computation. Terminal loss evaluation and parameter updates are omitted. 4.4 Theoretical analysis We formalize the training-time alignment by viewing sampling as a composition of differentiable denoising maps, showing that rollout backpropagation optimizes constraint violations under the same sampling distribution at inference. Theorem 4.3. (Training-time alignment) Let tii=0N\t_i\_i=0^N be the fixed discretization schedule. Define each denoising step as a deterministic function: i−1=fϕ(i)(i;ti,ti−1,ϵi),ϵi∼(0,), x_i-1=f_φ^(i)(x_i;t_i,t_i-1, _i), _i N(0, I), (18) and fϕf_φ is differentiable in ix_i and ϕφ. For an initial noise sample Nx_N, the terminal sample is obtained by composing the update maps: Fϕ=fϕ(1)∘fϕ(2)∘⋯∘fϕ(N),s.t.0ϕ=Fϕ(N)is the terminal state. F_φ=f_φ^(1) f_φ^(2) ·s f_φ^(N), s.t. x_0^φ=F_φ(x_N)\ is the terminal state. (19) Then the gradient of unscaled rollout objective 16 ℒrollout(ϕ)L_rollout(φ) with respect to ϕφ is: ∇ϕℒrollout(ϕ)=p(N)[∑i=1N∇0ϕℓΩ(0ϕ)∂0ϕ∂i−1ϕ∂fϕ(i)(i)∂ϕ]. _φL_rollout(φ)=E_p(x_N) [ _i=1^N _x_0^φ (x_0^φ) ∂x_0^φ∂x_i-1^φ ∂ f_φ^(i)(x_i)∂φ ]. (20) This concludes that backpropagating ℓΩ through the rollout computes the exact gradient of ℒrollout(ϕ)L_rollout(φ), directly minimizing constraint violations under the model’s own sampling distribution. Remark 4.4. (Full training objective) The full objective in Equation˜14 augments the rollout loss with ℒEDM(ϕ)L_EDM(φ). Since ℒEDM(ϕ)L_EDM(φ) is evaluated on clean data samples independently of the rollout trajectory, its gradient is simply an additive term outside the rollout computation graph. Thus, Theorem˜4.3 still applies to ℒ(ϕ)L(φ), with ∇ϕℒEDM(ϕ) _φL_EDM(φ) added independently. 5 Experiments We evaluate whether IMCD improves constraint enforcement while preserving generative fidelity under the same experimental conditions as prior work. To enable direct comparison, we adopt the experimental setup of [26] and focus on isolating the effect of IMCD. Specifically, we assess whether the proposed design reduces infeasible generation through rollout-based training and improves stability across noise levels. 5.1 Bouncing balls We begin with a synthetic bouncing ball experiment, where multiple balls move inside a closed box with elastic ball-ball and ball-wall collisions. This setting enables controlled evaluation of both physical consistency and constraint satisfaction. Following [26], we generate simulated trajectories with randomized initial positions and velocities using a physics-based simulator [16]. The dataset contains 100,000 scenarios, each with 10 balls over 100 timesteps. The goal is to model the trajectory distribution while enforcing two constraints: (i) boundary constraints, keeping all balls inside the box, (i) overlap constraints, preventing overlap between balls. We compare against a diverse set of baselines spanning unconstrained, sampling-time correction and fine-tuning approaches. EDM [23] serves as the unconstrained pretrained diffusion model. MBM [32] and MBM++ [26] incorporate constraint gradients into the score function with scaling and guidance functions in Table˜3 to enforce constraint satisfaction at the clean sample. MPGD [19] is a sampling-time correction method that we adapt with the same diverging scaling schedule as MBM++, omitting decoder projection. PIDM [2] adds denoised-state constraint losses to the DSM objective. Since the original formulation does not specify whether pretraining is used, we evaluate both training-from-scratch and fine-tuning variants, denoted PIDM and PIDM (FT). We also include rollout-based fine-tuning baselines, Adjoint Matching (AM) [12] and DPOK [14], which optimize rollout-level objectives while regularizing deviation from the unconstrained pretrained model. All methods use the same sampling configuration with 50 reverse-time steps. We evaluate constraint satisfaction using overlap and boundary violation rates, and assess distributional fidelity using reweighted ELBO (r-ELBO). To capture physical motion consistency, we introduce three additional metrics. We measure maximum frame-to-frame displacement (F2F) to quantify temporal inconsistency and detect jittering artifacts, maximum contact distance (MCD) to capture worst-case proximity during ball-ball and ball-wall interactions, and maximum energy deviation (MED) to measure violations of kinetic energy conservation over time. We also compute these metrics on ground-truth trajectories for reference. Full definitions of these metrics are provided in the appendix Section˜D.1. Table˜4 reports results on the bouncing ball experiment. Baselines reveal a trade-off between feasibility, physical realism, and fidelity. MPGD enforces constraints perfectly, but its high F2F and MED indicate abrupt frame-to-frame corrections and jittery motion, as visualized in Figure˜2. MBM and MBM++ reduce violations but remain unstable or lose physical plausibility under few-step sampling. PIDM and PIDM (FT) preserve plausible dynamics but fail to enforce hard constraints, while DPOK and AM either degrade fidelity or only partially satisfy constraints. Our method avoids these trade-offs: it achieves near-zero boundary and overlap violations while keeping F2F and MED close to the EDM baseline, indicating smooth constraint-aware rollouts rather than abrupt corrections. Its competitive fidelity shows that retaining EDM denoising preserves the training data distribution, yielding the strongest overall balance across metrics. To better understand these gains, we further provide ablation studies on the effect of rollout-based training and the trainable components in Section˜D.2. Table 1: Bouncing ball experiment results. We report constraint violation rates, distributional fidelity (r-ELBO), and physical consistency metrics (F2F, MCD, MED). Ground truth values are provided for reference. Constraint (%) ↓ Fidelity ↑ Physical plausibility (×10−1× 10^-1) ↓ Method Boundary rate Overlap rate r-ELBO (×10−2× 10^-2) F2F MCD MED EDM [23] 6.01±0.026.01± 0.02 39.82±0.2139.82± 0.21 −21.7±0.1-21.7± 0.1 3.1±0.03.1± 0.0 1.0±0.01.0± 0.0 0.4±0.00.4± 0.0 MPGD [19] 0.00±0.000.00± 0.00 0.00±0.000.00± 0.00 −27.0±0.1-27.0± 0.1 8.6±0.48.6± 0.4 0.8±0.00.8± 0.0 29.7±5.429.7± 5.4 MBM [32] 0.36±0.000.36± 0.00 0.16±0.000.16± 0.00 −22.2±0.1-22.2± 0.1 59.7±0.459.7± 0.4 1.7±0.01.7± 0.0 510.9±11.5510.9± 11.5 MBM++ [26] 0.00±0.000.00± 0.00 0.00±0.000.00± 0.00 −22.8±0.1-22.8± 0.1 10.3±0.610.3± 0.6 0.8±0.00.8± 0.0 41.1±5.941.1± 5.9 PIDM [2] 11.07±0.1511.07± 0.15 40.91±0.1940.91± 0.19 −26.9±0.1-26.9± 0.1 2.9±0.02.9± 0.0 1.1±0.01.1± 0.0 0.3±0.00.3± 0.0 PIDM (FT) [2] 6.40±0.046.40± 0.04 29.15±0.1029.15± 0.10 −21.6±0.1-21.6± 0.1 3.0±0.03.0± 0.0 0.8±0.00.8± 0.0 0.3±0.00.3± 0.0 AM [12] 0.49±0.010.49± 0.01 4.64±0.064.64± 0.06 −39.4±0.2-39.4± 0.2 3.4±0.03.4± 0.0 3.8±0.03.8± 0.0 0.5±0.00.5± 0.0 DPOK [14] 2.04±0.052.04± 0.05 1.69±0.031.69± 0.03 −53.5±0.1-53.5± 0.1 2.8±0.02.8± 0.0 4.3±0.04.3± 0.0 0.3±0.00.3± 0.0 IMCD (Ours) 0.01±0.000.01± 0.00 0.01±0.000.01± 0.00 −22.9±0.0-22.9± 0.0 3.4±0.13.4± 0.1 1.2±0.01.2± 0.0 0.6±0.00.6± 0.0 Ground truth 0.00±0.000.00± 0.00 0.00±0.000.00± 0.00 – 2.8±0.02.8± 0.0 0.4±0.00.4± 0.0 0.3±0.00.3± 0.0 5.2 Traffic scene trajectory prediction We next evaluate real-world traffic scene trajectory prediction on the Interaction dataset [52], which contains vehicle trajectories from 11 traffic scenarios. Given a short history of observed states, the model predicts future positions and headings for all agents in the scene. This setting tests multi-agent driving behavior modeling under road geometry and vehicle interactions, including staying within drivable areas and avoiding collisions. Following prior work [27, 26], we observe one second of motion and predict a three-second future horizon. We compare against the same classes of baselines as in the bouncing ball experiment, excluding MBM and AM, and additionally include CriticSMC [27]. Our bouncing ball results show that MBM underperforms MBM++, and we found that AM converges slowly and incurs high computational cost in this large-scale trajectory prediction setting. We use DJINN [34], an EDM-based diffusion model for joint trajectory prediction, as the pretrained backbone, fine-tuning it with the corresponding objective for baselines that require training and for our method. We report our main variant in the main text and defer additional rollout/LoRA variants to Section˜E.2. For evaluation, we generate six trajectory samples per scene and report metrics covering constraint satisfaction, trajectory match, physical plausibility and diversity. Constraint satisfaction is measured by offroad and collision rates. Trajectory match is evaluated using ego minADE6/minFDE6, where scene metrics minSADE6/minSFDE6 are computed over all vehicles. Physical plausibility is measured by maximum scene frame-to-frame displacement (SF2F) and maximum scene lateral velocity (SMLV), which capture abrupt motion and unrealistic lateral shift. Diversity is measured by maximum ego final distance (MFD6) across six samples. Table˜2 reports results on the INTERACTION DR_DEU_Merging_MT scenario. Existing methods remain limited in different ways: EDM and PIDM (FT) better preserve diversity and physical plausibility but retain substantial violations, while MPGD and MBM++ reduce violations but still incur non-negligible offroad errors. In contrast, our IMCD achieves the strongest feasibility, substantially reducing offroad violations and eliminating collisions while maintaining good trajectory match. Its physical plausibility metrics remain close to EDM, indicating that constraint satisfaction is achieved without severe motion distortion. Overall, IMCD gives the best balance between feasibility and quality under few-step generation. Additional experiments and visualizations on different locations are reported in Section˜E.5. Table 2: Traffic scene trajectory prediction results on the INTERACTION DR_DEU_Merging_MT scenario. Ground truth values are provided for reference. Constraint (%) ↓ Min traj. match ↓ Phys. plaus. ↓ Diversity ↑ Method Offroad Collision ADE6 FDE6 SADE6 SFDE6 SF2F6 SMLV6 MFD6 EDM [23] 9.039.03 0.420.42 0.200.20 0.500.50 0.270.27 0.770.77 1.111.11 0.410.41 2.752.75 MPGD [19] 3.783.78 0.290.29 0.190.19 0.460.46 0.230.23 0.670.67 1.041.04 0.340.34 1.921.92 MBM++ [26] 3.733.73 0.290.29 0.170.17 0.450.45 0.220.22 0.660.66 1.011.01 0.310.31 1.881.88 PIDM (FT) [2] 8.528.52 0.340.34 0.180.18 0.500.50 0.240.24 0.760.76 0.960.96 0.260.26 2.672.67 DPOK [14] 5.455.45 0.030.03 3.453.45 9.949.94 3.343.34 9.739.73 0.880.88 0.400.40 1.901.90 CriticSMC [27] – 1.031.03 0.350.35 – – – – – 2.372.37 IMCD (Ours) 0.300.30 0.000.00 0.180.18 0.450.45 0.230.23 0.660.66 1.071.07 0.430.43 1.911.91 Ground truth 0.000.00 0.000.00 – – – – 0.830.83 0.020.02 – 6 Related Work 6.1 Diffusion models for constrained generation Constrained diffusion has been approached by modifying either the forward or reverse process. Liu and Wu [29] employ Doob’s h-transform [38] to construct a drift that confines the diffusion process to the constraint domain at training time. Reflected diffusion model [30] and constrained diffusion framework proposed by Fishman et al. [15] alter the forward noising process by reflected Brownian motion or log-barrier metrics to confine samples to feasible regions. Projected diffusion [8] instead enforces constraints at inference time by projecting noisy samples onto the constraint set at each denoising step, with extensions to multi-robot motion planning [25]. However, repeated projection of noisy intermediate states can distort the denoising trajectory. More importantly, these methods all require constraints to be analytically formulated upfront, which makes them less applicable when feasibility is observed through violations but not available in closed form, as in our setting. 6.2 Reward-driven fine-tuning and controlled diffusion Fine-tuning diffusion models with reward or energy objectives is closely related to our training formulation. We defer a detailed comparison with AM and DPOK to Appendix˜A. A related line of work views diffusion sampling as stochastic optimal control, where methods such as AM, Adjoint sampling (AS) [18] and Adjoint Schrodinger bridge Sampler (ASBS) [28] learn reverse-time controls that drive samples toward reward- or energy-defined targets. ASBS further extends AS to arbitrary source distributions. However, both methods are designed to optimize reward or energy objective directly, rather than leverage a pretrained model as the base distribution. In contrast, our goal is constrained generation, which fine-tunes a pretrained denoiser through its sampling rollout to improve feasibility while preserving the learned data distribution. 6.3 Diffusion models for autonomous driving Several recent works have explored diffusion models for autonomous driving planning. DJINN [34] jointly diffuses trajectories of all agents for realistic scenario generation, but is primarily designed for simulation rather than planning, and does not focus on constraint satisfaction. Diffusion-ES [47] combines gradient-free evolutionary search with a diffusion trajectory model to optimize non-differentiable test-time reward objectives, but operates on the ego vehicle only without modeling interactions with other agents. Diffusion Planner [54] jointly models ego planning and agent prediction with a unified diffusion architecture, and applies classifier guidance at inference time to achieve constraint satisfaction, but may encounter the sampling-time correction limitation discussed in Section˜3.2. 7 Discussion and future work Our method IMCD provides a rollout-based fine-tuning framework for constrained diffusion generation that addresses training-time misalignment. By minimizing constraint violations on terminal rollout samples, the model learns corrections under its own sampling dynamics rather than relying on post hoc guidance. Both theory and experiments show that our method improves constraint satisfaction while preserving distributional fidelity across bouncing ball simulation and autonomous driving tasks. While effective, the current inference procedure still involves constraint-gradient computation during the reverse rollout, adding modest computational overhead. A promising direction is to improve efficiency through distillation, where a standalone denoiser learns to approximate the constraint-aware updates directly. Another important direction is to extend rollout-based fine-tuning to higher-dimensional visual domains and broader application settings where feasibility is specified by task-specific violation signals. Acknowledgment We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), the Alberta Machine Intelligence Institute (Amii) through the Canada CIFAR AI Chairs Program, Inverted AI, Mitacs, and Google. This research was enabled in part by technical support and computational resources provided by the Digital Research Alliance of Canada (alliancecan.ca), the Advanced Research Computing at the University of British Columbia (arc.ubc.ca), and Amazon Web Services. References [1] B. D. Anderson (1982) Reverse-time diffusion equation models. Stochastic Processes and their Applications 12 (3), p. 313–326. Cited by: §2. [2] J. Bastek, W. Sun, and D. M. Kochmann (2024) Physics-informed diffusion models. arXiv preprint arXiv:2403.14404. Cited by: §E.5, Table 7, Table 8, §1, §3.1, §5.1, Table 1, Table 1, Table 2. [3] I. Beltagy, M. E. Peters, and A. Cohan (2020) Longformer: the long-document transformer. arXiv preprint arXiv:2004.05150. Cited by: §D.5. [4] G. Biroli, T. Bonnaire, V. De Bortoli, and M. Mézard (2024) Dynamical regimes of diffusion models. Nature Communications 15 (1), p. 9957. Cited by: §3.1. [5] B. Boys, M. Girolami, J. Pidstrigach, S. Reich, A. Mosca, and O. D. Akyildiz (2023) Tweedie moment projected diffusions for inverse problems. arXiv preprint arXiv:2310.06721. Cited by: §4.1. [6] J. Carvalho, A. T. Le, M. Baierl, D. Koert, and J. Peters (2023) Motion planning diffusion: learning and planning of robot motions with diffusion models. In 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), p. 1916–1923. Cited by: §1. [7] T. Chen, B. Xu, C. Zhang, and C. Guestrin (2016) Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174. Cited by: §4.3. [8] J. K. Christopher, S. Baek, and F. Fioretto (2024) Projected generative diffusion models for constraint satisfaction. arXiv preprint arXiv:2402.03559. Cited by: §1, §3.2, §6.1. [9] H. Chung, J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye (2022) Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687. Cited by: §1, §3.2. [10] H. Chung, J. Kim, G. Y. Park, H. Nam, and J. C. Ye (2024) Cfg++: manifold-constrained classifier free guidance for diffusion models. arXiv preprint arXiv:2406.08070. Cited by: §1. [11] P. Dhariwal and A. Nichol (2021) Diffusion models beat gans on image synthesis. Advances in neural information processing systems 34, p. 8780–8794. Cited by: §1, §3.2. [12] C. Domingo-Enrich, M. Drozdzal, B. Karrer, and R. T. Chen (2024) Adjoint matching: fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.08861. Cited by: Appendix A, §5.1, Table 1. [13] B. Efron (2011) Tweedie’s formula and selection bias. Journal of the American Statistical Association 106 (496), p. 1602–1614. Cited by: §2, §4.1. [14] Y. Fan, O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, P. Abbeel, M. Ghavamzadeh, K. Lee, and K. Lee (2023) Dpok: reinforcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Systems 36, p. 79858–79885. Cited by: Appendix A, §E.5, Table 7, Table 8, §5.1, Table 1, Table 2. [15] N. Fishman, L. Klarner, V. De Bortoli, E. Mathieu, and M. J. Hutchinson (2023) Diffusion models for constrained domains. Transactions on Machine Learning Research. Cited by: §1, §6.1. [16] Z. Gan, C. Li, R. Henao, D. E. Carlson, and L. Carin (2015) Deep temporal sigmoid belief networks for sequence modeling. Advances in Neural Information Processing Systems 28. Cited by: §5.1. [17] W. Harvey, S. Naderiparizi, V. Masrani, C. Weilbach, and F. Wood (2022) Flexible diffusion modeling of long videos. Advances in neural information processing systems 35, p. 27953–27965. Cited by: §1. [18] A. Havens, B. K. Miller, B. Yan, C. Domingo-Enrich, A. Sriram, B. Wood, D. Levine, B. Hu, B. Amos, B. Karrer, et al. (2025) Adjoint sampling: highly scalable diffusion samplers via adjoint matching. arXiv preprint arXiv:2504.11713. Cited by: §6.2. [19] Y. He, N. Murata, C. Lai, Y. Takida, T. Uesaka, D. Kim, W. Liao, Y. Mitsufuji, J. Z. Kolter, R. Salakhutdinov, et al. (2024) Manifold preserving guided diffusion. In The Twelfth International Conference on Learning Representations, Cited by: §D.3, §E.5, Table 7, Table 8, §1, §3.2, §3.2, §5.1, Table 1, Table 2. [20] J. Ho, A. Jain, and P. Abbeel (2020) Denoising diffusion probabilistic models. Advances in neural information processing systems 33, p. 6840–6851. Cited by: §1, §2. [21] J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet (2022) Video diffusion models. Advances in neural information processing systems 35, p. 8633–8646. Cited by: §1. [22] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al. (2022) Lora: low-rank adaptation of large language models.. Iclr 1 (2), p. 3. Cited by: Appendix A, §4.2. [23] T. Karras, M. Aittala, T. Aila, and S. Laine (2022) Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems 35, p. 26565–26577. Cited by: §E.5, Table 7, Table 8, §1, §2, §2, §4.3, §5.1, Table 1, Table 2. [24] M. Li, T. Qu, R. Yao, W. Sun, and M. Moens (2023) Alleviating exposure bias in diffusion models through sampling with shifted time steps. arXiv preprint arXiv:2305.15583. Cited by: §3.1. [25] J. Liang, J. K. Christopher, S. Koenig, and F. Fioretto (2025) Simultaneous multi-robot motion planning with projected diffusion models. arXiv preprint arXiv:2502.03607. Cited by: §1, §3.2, §6.1. [26] X. Liang, S. Naderiparizi, Y. Liu, B. Zwartsenberg, and F. Wood (2026) Improved constrained generation by bridging pretrained generative models. arXiv preprint arXiv:2603.06742. Cited by: Table 3, §D.5, §E.4, §E.5, Table 7, Table 8, §1, §3.1, §4.1, §4.1, §5.1, §5.1, §5.2, Table 1, Table 2, §5. [27] V. Lioutas, J. W. Lavington, J. Sefas, M. Niedoba, Y. Liu, B. Zwartsenberg, S. Dabiri, F. Wood, and A. Scibior (2023) Critic sequential monte carlo. In The Eleventh International Conference on Learning Representations, Cited by: §E.5, Table 7, Table 8, §5.2, §5.2, Table 2. [28] G. Liu, J. Choi, Y. Chen, B. K. Miller, and R. T. Chen (2025) Adjoint schr\ " odinger bridge sampler. arXiv preprint arXiv:2506.22565. Cited by: §6.2. [29] X. Liu and L. Wu (2023) Learning diffusion bridges on constrained domains. In international conference on learning representations (ICLR), Cited by: §6.1. [30] A. Lou and S. Ermon (2023) Reflected diffusion models. In International Conference on Machine Learning, p. 22675–22701. Cited by: §1, §6.1. [31] S. Naderiparizi, X. Liang, S. Cohan, B. Zwartsenberg, and F. Wood (2024) Don’t be so negative! score-based generative modeling with oracle-assisted guidance. In Forty-first International Conference on Machine Learning, Cited by: §2. [32] S. Naderiparizi, X. Liang, B. Zwartsenberg, and F. Wood (2025) Constrained generative modeling with manually bridged diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 39, p. 19607–19615. Cited by: Table 3, §1, §4.1, §5.1, Table 1. [33] S. Naderiparizi, A. Scibior, A. Munk, M. Ghadiri, A. G. Baydin, B. J. Gram-Hansen, C. A. S. De Witt, R. Zinkov, P. Torr, T. Rainforth, et al. (2022) Amortized rejection sampling in universal probabilistic programming. In International Conference on Artificial Intelligence and Statistics, p. 8392–8412. Cited by: §4.1. [34] M. Niedoba, J. Lavington, Y. Liu, V. Lioutas, J. Sefas, X. Liang, D. Green, S. Dabiri, B. Zwartsenberg, A. Scibior, et al. (2024) A diffusion-model of joint interactive navigation. Advances in Neural Information Processing Systems 36. Cited by: §E.4, §5.2, §6.3. [35] M. Ning, E. Sangineto, A. Porrello, S. Calderara, and R. Cucchiara (2023) Input perturbation reduces exposure bias in diffusion models. arXiv preprint arXiv:2301.11706. Cited by: §3.1. [36] E. Perez, F. Strub, H. De Vries, V. Dumoulin, and A. Courville (2018) Film: visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, Vol. 32. Cited by: Appendix C. [37] L. Regenwetter, G. Giannone, A. Srivastava, D. Gutfreund, and F. Ahmed (2024) Constraining generative models for engineering design with negative data. Transactions on Machine Learning Research. Cited by: §2. [38] L. C. Rogers and D. Williams (2000) Diffusions, markov processes, and martingales: volume 1, foundations. Cambridge university press. Cited by: §6.1. [39] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer (2022) High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 10684–10695. Cited by: §1. [40] J. Schulman, Y. Duan, J. Ho, A. Lee, I. Awwal, H. Bradlow, J. Pan, S. Patil, K. Goldberg, and P. Abbeel (2014) Motion planning with sequential convex optimization and convex collision checking. The International Journal of Robotics Research 33 (9), p. 1251–1270. Cited by: §1. [41] Y. Shen, X. Jiang, Y. Yang, Y. Wang, D. Han, and D. Li (2024) Understanding and improving training-free loss-based diffusion guidance. Advances in Neural Information Processing Systems 37, p. 108974–109002. Cited by: §3.2. [42] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli (2015) Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, p. 2256–2265. Cited by: §2. [43] Y. Song and S. Ermon (2019) Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems 32. Cited by: §2. [44] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2021) Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, Cited by: §1, §2, §2. [45] A. Vaswani (2017) Attention is all you need. Advances in Neural Information Processing Systems. Cited by: §D.5. [46] P. Vincent (2011) A connection between score matching and denoising autoencoders. Neural computation 23 (7), p. 1661–1674. Cited by: §2. [47] B. Yang, H. Su, N. Gkanatsios, T. Ke, A. Jain, J. Schneider, and K. Fragkiadaki (2024) Diffusion-es: gradient-free planning with diffusion for autonomous driving and zero-shot instruction following. arXiv preprint arXiv:2402.06559. Cited by: §1, §6.3. [48] L. Yang, S. Ding, Y. Cai, J. Yu, J. Wang, and Y. Shi (2024) Guidance with spherical gaussian constraint for conditional diffusion. arXiv preprint arXiv:2402.03201. Cited by: §1. [49] H. Ye, H. Lin, J. Han, M. Xu, S. Liu, Y. Liang, J. Ma, J. Zou, and S. Ermon (2024) Tfg: unified training-free guidance for diffusion models. Advances in Neural Information Processing Systems 37, p. 22370–22417. Cited by: §1. [50] J. Yu, Y. Wang, C. Zhao, B. Ghanem, and J. Zhang (2023) Freedom: training-free energy-guided conditional diffusion model. In Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 23174–23184. Cited by: §1. [51] T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn (2020) Gradient surgery for multi-task learning. Advances in neural information processing systems 33, p. 5824–5836. Cited by: Appendix C, §4.1. [52] W. Zhan, L. Sun, D. Wang, H. Shi, A. Clausse, M. Naumann, J. Kummerle, H. Konigshof, C. Stiller, A. de La Fortelle, et al. (2019) Interaction dataset: an international, adversarial and cooperative motion dataset in interactive driving scenarios with semantic maps. arXiv preprint arXiv:1910.03088. Cited by: §5.2. [53] M. Zhang, Z. Cai, L. Pan, F. Hong, X. Guo, L. Yang, and Z. Liu (2022) Motiondiffuse: text-driven human motion generation with diffusion model. arXiv preprint arXiv:2208.15001. Cited by: §D.5. [54] Y. Zheng, R. Liang, K. Zheng, J. Zheng, L. Mao, J. Li, W. Gu, R. Ai, S. E. Li, X. Zhan, et al. (2025) Diffusion-based planning for autonomous driving with flexible guidance. arXiv preprint arXiv:2501.15564. Cited by: §1, §6.3. Appendix A Algorithm Algorithm 1 Rollout-Based Constrained Fine-Tuning 1:Pretrained denoiser D, fine-tuned denoiser DϕD_φ initialized from D, scaling network γϕ _φ, constraint loss ℓΩ , noise schedule tii=0N\t_i\_i=0^N, training data D, division stabilizer ϵ=10−5ε=10^-5. 2:while not converged do 3: // Rollout 4: Sample T∼(0,tN2)x_T N(0,t_N^2 I) 5: for i=T,…,1i=T,…,1 do 6: Compute constraint-aware fine-tuned denoiser DϕΩ(~i;ti)D_φ ( x_i;t_i) with Equation˜13. 7: Update the rollout state to ~i−1 x_i-1. 8: end for 9: // Losses 10: ℒrollout←ℓΩ(~0)L_rollout← ( x_0) ⊳ Terminal state only 11: Sample 0∼x_0 D, t∼q(t∣0)x_t q(x_t x_0) 12: ℒEDM←λ(t)∥DϕΩ(t;t)−0∥2L_EDM←λ(t) D_φ (x_t;t)-x_0 ^2 13: // Dynamic Scaling and Update 14: κ←sg(ℒEDM¯)sg(ℒrollout¯)+ϵκ← sg( L_EDM)sg( L_rollout)+ε 15: ℒ←ℒEDM+κℒrolloutL←L_EDM+κL_rollout 16: Update ϕ←ϕ−η∇ϕℒφ←φ-η _φL ⊳ Gradients flow only through ϕφ 17:end while Rollout-based fine-tuning and stochastic optimal control Our rollout-based training admits a natural stochastic optimal control (SOC) interpretation: the reverse denoising process induced by the pretrained denoiser D(t;t)D(x_t;t) and constraint-aware fine-tuned denoiser DϕΩD_φ correspond to the base and controlled stochastic process respectively, and the rollout constraint loss acts as a terminal cost. This connects our approach to Adjoint Matching (AM) [12], which provides a formal SOC grounding for reward fine-tuning of generative models via rollout-based optimization. It is also related to DPOK [14], which frames the denoising chain as a Markov decision process (MDP) and applies policy gradient with KL regularization. Both AM and DPOK approaches regularize the fine-tuned model by minimizing its KL divergence from the pretrained model, treating the pretrained model as the reference. AM additionally incurs significant computational overhead beyond rollout simulation, as each gradient update requires integrating the lean adjoint ODE backward through the trajectory. However, in our setting, the pretrained denoiser can produce samples that violate the constraints, so enforcing proximity to it can perpetuate these errors. Instead, we rely on initialization from the pretrained model, anchor to the data distribution directly via the denoising objective, and combine it with rollout-time constraint losses to drive the sampler toward the feasible manifold, preserving consistency with the data rather than with the pretrained model’s potentially infeasible outputs. Gradient embeddings Following MBM++, the constraint gradient is embedded with a lightweight MLP and added once to tx_t as input to Dϕ(t;t)D_φ(x_t;t). Since this modifies the model input distribution seen by the pretrained backbone, we introduce LoRA adapters [22] in the attention layers of DϕD_φ to provide the necessary adaptation capacity. Table 3: A summary table of scaling and guidance functions for various methods. MBM [32] MBM++ [26] IMCD (ours) γ(t;t)γ(x_t;t) c/t2c/t^2 c/t2c/t^2 αϕ(t;t)⋅tβϕ(t;t) _φ(x_t;t)· t _φ(x_t;t) GΩ(t;t)G (x_t;t) ∂1ℓΩ(t;t) _1 (x_t;t) ∂1ℓΩ(D(t;t);t) _1 (D(x_t;t);t) ∂1ℓΩ(D(t;t);t) _1 (D(x_t;t);t) Appendix B Derivation of proofs B.1 Proof of Proposition˜4.2 Proof. Given the formulation of pΩ(t;t)p (x_t;t), take gradient ∇tlog⁡pΩ(t;t) _x_t p (x_t;t) and treat γt(t;t) _t(x_t;t) as constant with respect to tx_t. This yields the score decomposition in Definition˜4.1. For the convergence of denoising terminal state to the constrained region, since βϕ<−1 _φ<-1, tβϕ→∞t _φ→∞ as t→0t→ 0. Given αϕ>0 _φ>0, we have γϕ(t;t)→∞ _φ(x_t;t)→∞ as t→0t→ 0. Then limt→0exp⁡(−γϕ(t;t)ℓΩ(t;t))=1ifℓΩ(t;t)=00ifℓΩ(t;t)>0 _t→ 0 (- _φ(x_t;t) (x_t;t))= \ array[]c1&if\ (x_t;t)=0\\ 0&if\ (x_t;t)>0 array . (23) Therefore, limt→0pΩ(t;t)∝limt→0p(t;t)Ω(t), _t→ 0p (x_t;t) _t→ 0p(x_t;t) 1_ (x_t), (24) meaning that as t→0t→ 0, the distribution concentrates on the constrained region Ω , ensuring that the terminal denoised sample 0∼p(0,0)x_0 p(x_0,0) lies within Ω . ∎ B.2 Proof of Theorem˜4.3 Proof. Since each ϵi _i is independent of ϕφ, the trajectory ii=0N\x_i\_i=0^N is a deterministic differentiable function of ϕφ. Because each fϕ(i)f_φ^(i) is differentiable, then the composed map FϕF_φ is differentiable. By chain rule, ∇ϕℓΩ(Fϕ(N))=∇0ϕℓΩ(0ϕ)∂Fϕ(N)∂ϕ. _φ (F_φ(x_N))= _x_0^φ (x_0^φ) ∂ F_φ(x_N)∂φ. (25) Then the gradient of ℒrollout(ϕ)L_rollout(φ) with respect to ϕφ is ∇ϕℒrollout(ϕ) _φL_rollout(φ) =∇ϕp(N)[ℓΩ(Fϕ(N))]=p(N)[∇ϕℓΩ(Fϕ(N))] = _φE_p(x_N) [ (F_φ(x_N)) ]=E_p(x_N) [ _φ (F_φ(x_N)) ] (26) =p(N)[∇0ϕℓΩ(0ϕ)∂Fϕ(N)∂ϕ] =E_p(x_N) [ _x_0^φ (x_0^φ) ∂ F_φ(x_N)∂φ ] (27) by interchanging differentiation with expectation. Now we expand ∂Fϕ(N)∂ϕ ∂ F_φ(x_N)∂φ by applying the chain rule through the individual step fϕ(i)f_φ^(i): ∂Fϕ(N)∂ϕ=∑i=1N∂0ϕ∂iϕ⋅∂iϕ∂ϕ, ∂ F_φ(x_N)∂φ= _i=1^N ∂x_0^φ∂x_i^φ· ∂x_i^φ∂φ, (28) where ∂0ϕ∂iϕ ∂x_0^φ∂x_i^φ is the Jacobian of the terminal state 0ϕx_0^φ with respect to the intermediate state iϕx_i^φ. The second term ∂iϕ∂ϕ ∂x_i^φ _φ, where iϕ=fϕ(i+1)(i+1)x_i^φ=f_φ^(i+1)(x_i+1), can be further decomposed to: ∂iϕ∂ϕ=∂fϕ(i+1)(i+1)∂i+1ϕ⋅∂i+1ϕ∂ϕ+∂fϕ(i+1)(i+1)∂ϕ. ∂x_i^φ∂φ= ∂ f_φ^(i+1)(x_i+1)∂x_i+1^φ· ∂x_i+1^φ∂φ+ ∂ f_φ^(i+1)(x_i+1)∂φ. (29) Given this recursive pattern with base case ∂N∂ϕ=0 ∂x_N∂φ=0 as Nx_N is independent of ϕφ, we have ∂N−1ϕ∂ϕ ∂x_N-1^φ∂φ =∂fϕ(N)(N)∂Nϕ⋅∂Nϕ∂ϕ+∂fϕ(N)(N)∂ϕ=∂fϕ(N)(N)∂ϕ, = ∂ f_φ^(N)(x_N)∂x_N^φ· ∂x_N^φ∂φ+ ∂ f_φ^(N)(x_N)∂φ= ∂ f_φ^(N)(x_N)∂φ, (30) ∂N−2ϕ∂ϕ ∂x_N-2^φ∂φ =∂fϕ(N−1)(N−1)∂N−1ϕ⋅∂N−1ϕ∂ϕ+∂fϕ(N−1)(N−1)∂ϕ = ∂ f_φ^(N-1)(x_N-1)∂x_N-1^φ· ∂x_N-1^φ∂φ+ ∂ f_φ^(N-1)(x_N-1)∂φ (31) =∂fϕ(N−1)(N−1)∂N−1ϕ⋅∂fϕ(N)(N)∂ϕ+∂fϕ(N−1)(N−1)∂ϕ = ∂ f_φ^(N-1)(x_N-1)∂x_N-1^φ· ∂ f_φ^(N)(x_N)∂φ+ ∂ f_φ^(N-1)(x_N-1)∂φ (32) =∂N−2ϕ∂N−1ϕ⋅∂fϕ(N)(N)∂ϕ+∂fϕ(N−1)(N−1)∂ϕ. = ∂x_N-2^φ∂x_N-1^φ· ∂ f_φ^(N)(x_N)∂φ+ ∂ f_φ^(N-1)(x_N-1)∂φ. (33) This gives a general form for i∈0,1,⋯,N−1i∈\0,1,·s,N-1\: ∂iϕ∂ϕ ∂x_i^φ∂φ =∑j=i+1N∂iϕ∂j−1ϕ⋅∂fϕ(j)(j)∂ϕ. = _j=i+1^N ∂x_i^φ∂x_j-1^φ· ∂ f_φ^(j)(x_j)∂φ. (34) Then ∂Fϕ(N)∂ϕ=∂0ϕ∂ϕ=∑i=1N∂0ϕ∂i−1ϕ⋅∂fϕ(i)(i)∂ϕ. ∂ F_φ(x_N)∂φ= ∂x_0^φ∂φ= _i=1^N ∂x_0^φ∂x_i-1^φ· ∂ f_φ^(i)(x_i)∂φ. (35) Therefore, ∇ϕℒrollout(ϕ)=p(N)[∑i=1N∇0ϕℓΩ(0ϕ)∂0ϕ∂i−1ϕ∂fϕ(i)(i)∂ϕ]. _φL_rollout(φ)=E_p(x_N) [ _i=1^N _x_0^φ (x_0^φ) ∂x_0^φ∂x_i-1^φ ∂ f_φ^(i)(x_i)∂φ ]. (36) ∎ Corollary B.1. (Unbiased gradient estimator) Let Nx_N be a single initial noise sample and let iϕi=0N−1\x_i^φ\_i=0^N-1 be the trajectory by rolling out fϕ(i)i=1N\f_φ^(i)\_i=1^N. Then the single sample estimator h(ϕ)=∑i=1N∇0ϕℓΩ(0ϕ)∂0ϕ∂i−1ϕ∂fϕ(i)(i)∂ϕ h(φ)= _i=1^N _x_0^φ (x_0^φ) ∂x_0^φ∂x_i-1^φ ∂ f_φ^(i)(x_i)∂φ (37) is an unbiased estimator of ∇ϕℒrollout(ϕ) _φL_rollout(φ). Proof. This is an immediate conclusion from Theorem˜4.3. Since h(ϕ)h(φ) is a fixed realization, this result is obtained by taking expectation over p(N)p(x_N) and interchanging expectation and differentiation, which is, ∇ϕℒrollout(ϕ)=p(N)[h(ϕ)]. _φL_rollout(φ)=E_p(x_N) [h(φ) ]. (38) ∎ Appendix C Multiple constraint gradients and scaling network architecture To handle potential conflicts between multiple gradients, we apply a symmetric gradient projection inspired by PCGrad [51]. Let jj=1J\g_j\_j=1^J denote the gradients associated with J constraint losses. For each pair of constraints (j,k)(g_j,g_k), if j⊤k<0g_j g_k<0, we remove the conflicting component from both gradients: k′ g_k ←k−k⊤j∥j∥2+ϵj, ←g_k- g_k g_j g_j ^2+ε\,g_j, (39) j′ g_j ←j−j⊤k∥k∥2+ϵk, ←g_j- g_j g_k g_k ^2+ε\,g_k, (40) where ϵ=10−5ε=10^-5. If j⊤k≥0g_j g_k≥ 0, the gradients are left unchanged. The corrected gradients are then combined to form the final constraint direction. Therefore, the scaling network learns a global correction magnitude rather than per-constraint weights while mitigating destructive interference between constraints. In both our experiments, J=2J=2, so this reduces to the symmetric two gradient updates described above. The scaling network γϕ _φ is implemented as three MLPs with linear layers and SiLU activations. The first MLP processes the noisy state tx_t to a feature embedding. The second MLP maps the noise level t to a time embedding, which is split into scale and shift components for FiLM-style [36] conditioning of the feature embedding. The third MLP maps the conditioned features to αϕ _φ and βϕ _φ. Together, these outputs determine the correction magnitude through Equation˜9. Experiment-specific parameterizations of αϕ _φ and βϕ _φ are provided in the corresponding experiment detail sections. Appendix D Bouncing balls experiment details D.1 Definition of constraint losses and evaluation metrics We provide the full definitions of the constraint losses and evaluation metrics used in the bouncing balls experiment. Each scene contains B=10B=10 balls in a 10×1010× 10 square box, with ball radius r=0.5r=0.5. Let τ(b)∈ℝ2x_τ^(b)∈R^2 denote the center position of ball b at frame τ, for τ=1,⋯,τ=1,·s,T and =100T=100. Constraint losses We first define the boundary and overlap losses used to measure constraint violations. The boundary loss measures ball-wall penetration. Since the center of each ball remains in [r,10−r]2[r,10-r]^2, we define ℓboundary(τ(b))=max⁡r−τ,1(b),τ,1(b)−(10−r),r−τ,2(b),τ,2(b)−(10−r),0. _boundary(x_τ^(b))= \r-x_τ,1^(b),x_τ,1^(b)-(10-r),r-x_τ,2^(b),x_τ,2^(b)-(10-r),0\. (41) The overlap loss measures overlap between pairs of balls: ℓoverlap(τ(b),τ(b′)=max(2r−∥τ(b)−τ(b′)∥,0). _overlap(x_τ^(b),x_τ^(b )= (2r- x_τ^(b)-x_τ^(b ) ,0). (42) Reported metrics The boundary rate is the fraction of frames in which at least one ball has positive boundary loss: Boundary rate=1∑t=1[maxb⁡ℓboundary(τ(b))>0]. rate= 1T _t=1^T 1 [ _b _boundary (x_τ^(b) )>0 ]. (43) The overlap rate is the fraction of frames in which at least one pair of balls has positive overlap loss: Overlap rate=1∑t=1[maxb<b′⁡ℓoverlap(τ(b),τ(b′))>0]. rate= 1T _t=1^T 1 [ _b<b _overlap (x_τ^(b),x_τ^(b ) )>0 ]. (44) The maximum frame-to-frame displacement (F2F) measures abrupt motion or jitter: F2F=maxb,t<∥τ+1(b)−τ(b)∥2. 2F= _b,t<T x_τ+1^(b)-x_τ^(b) _2. (45) The maximum contact distance (MCD) measures whether sharp direction changes are supported by nearby contacts. We estimate velocity directly from consecutive positions, which is more reliable than using predicted velocity in generated samples. Specifically, we define τ(b)=τ+1(b)−τ(b). v_τ^(b)=x_τ+1^(b)-x_τ^(b). (46) We then mark (b,τ)(b,τ) as a candidate bounce if the angle between τ+1(b)v_τ+1^(b) and τ(b)v_τ^(b) exceeds π6 π6. Formally, ℬ=(b,τ):arccos⁡(τ−1(b)⊤τ(b)∥τ−1(b)∥∥τ(b)∥+ϵ)>π6, B= \(b,τ): ( v_τ-1^(b) v_τ^(b) v_τ-1^(b) v_τ^(b) +ε )> π6 \, (47) where τ=2,⋯,−1τ=2,·s,T-1. For each candidate bounce (b,τ)(b,τ), we compute the signed distance to the nearest possible contact object over a local temporal window τ=τ−1,τ,τ+1. W_τ=\τ-1,τ,τ+1\. (48) This accounts for small discretization offsets between the detected direction change and the actual contact frame. The signed wall distance is dwall(b,s)=min⁡s,1(b)−r,10−r−s,1(b),s,2(b)−r,10−r−s,2(b), d_wall(b,s)= \x_s,1^(b)-r,10-r-x_s,1^(b),x_s,2^(b)-r,10-r-x_s,2^(b)\, (49) and the signed ball-ball distance is dball(b,s)=minb′≠b⁡(∥s(b)−s(b′)∥−2r). d_ball(b,s)= _b ≠ b ( x_s^(b)-x_s^(b ) -2r ). (50) We collect all candidate signed contact distances in the local window: (b,τ)=dwall(b,s):s∈τ∪dball(b,s):s∈τ. D(b,τ)=\d_wall(b,s):s∈W_τ\∪\d_ball(b,s):s∈W_τ\. (51) MCD is then defined as MCD=max(b,τ)∈ℬ⁡mind∈(b,τ)⁡|d|. = _(b,τ)∈B _d∈D(b,τ)|d|. (52) Finally, the maximum energy deviation (MED) measures frame-wise inconsistency in system kinetic energy. Assuming unit mass and constant time step, we define the kinetic energy proxy Eτ=∑b=1B12∥τ+1(b)−τ(b)∥22. E_τ= _b=1^B 12 x_τ+1^(b)-x_τ^(b) _2^2. (53) The maximum energy deviation is: MED=maxt<−1⁡|Eτ+1−Eτ|. = _t<T-1|E_τ+1-E_τ|. (54) All maximum-based metrics are computed within each generated trajectory and then averaged over evaluation batches. We report the mean and standard deviation over four independent evaluation runs. D.2 Ablation study We further conduct ablations to isolate the contribution of each component in our framework. FP training replaces rollout-based training with standard forward process training, which introduces a mismatch between training and sampling dynamics. It keeps the learned scaling network and gradient embeddings trainable. Rollout-only keeps rollout-based training and the same trainable scaling network and gradient embeddings, but removes LoRA adaptations. Finally, IMCD combines rollout-based training approach, the learned scaling network, gradient embeddings, and LoRA adaptation. Table 4: Ablation study of bouncing ball experiment Constraint (%) ↓ Fidelity ↑ Physical plausibility (×10−1× 10^-1) ↓ Method Boundary rate Overlap rate r-ELBO (×10−2× 10^-2) F2F MCD MED FP training 0.78±0.030.78± 0.03 2.83±0.062.83± 0.06 −21.7±0.1-21.7± 0.1 4.1±0.04.1± 0.0 0.5±0.00.5± 0.0 1.3±0.11.3± 0.1 Rollout-only 0.07±0.010.07± 0.01 0.01±0.000.01± 0.00 −23.4±0.0-23.4± 0.0 5.4±0.05.4± 0.0 0.9±0.00.9± 0.0 1.9±0.11.9± 0.1 IMCD 0.01±0.000.01± 0.00 0.01±0.000.01± 0.00 −22.9±0.0-22.9± 0.0 3.4±0.03.4± 0.0 1.2±0.01.2± 0.0 0.6±0.00.6± 0.0 Figure 3: Scaling network output versus noise level σ along denoising trajectory for bouncing ball experiment, Means are taken over batches; shaded bands show ±1± 1 standard error. Rather than performing a discrete ablation on the scaling network, we analyze its learned behavior across noise levels. As shown in Figure˜3, both learned components are noise dependent: αϕ _φ is not a constant multiplier, and βϕ _φ varies smoothly along the sampling trajectory. To compare with prior schedules of the form ct−2ct^-2 in Table˜3, we plot t2γϕ(t;t)t^2 _φ(x_t;t). This quantity would be constant under such a ct−2ct^-2 schedule, but instead changes substantially across noise levels. This suggests that the scaling network learns a nontrivial modulation beyond a fixed power law. Moreover, γϕ(t;t) _φ(x_t;t) grows rapidly as the noise level decreases, consistent with stronger constraint enforcement near the data manifold. D.3 Effect of sampling steps on sampling-time correction MPGD without projection [19] performs constraint correction only during sampling, without updating the pretrained model. We evaluate it on the bouncing ball task with 50,100,150,20050,100,150,200 denoising steps, where 5050 matches the main method setting. This ablation tests how sensitive sampling-time correction is to the number of steps for incremental updates along the denoising trajectory. Table 5: Effect of sampling steps on MPGD without projection in the bouncing ball task Constraint (%) ↓ Physical plausibility (×10−1× 10^-1) ↓ Steps Boundary rate Overlap rate F2F MCD MED 50 0.0±0.00.0± 0.0 0.0±0.00.0± 0.0 8.6±0.48.6± 0.4 0.8±0.00.8± 0.0 29.7±5.429.7± 5.4 100 0.0±0.00.0± 0.0 0.0±0.00.0± 0.0 3.8±0.23.8± 0.2 0.6±0.00.6± 0.0 3.2±0.93.2± 0.9 150 0.0±0.00.0± 0.0 0.0±0.00.0± 0.0 3.2±0.03.2± 0.0 0.6±0.00.6± 0.0 0.6±0.10.6± 0.1 200 0.0±0.00.0± 0.0 0.0±0.00.0± 0.0 3.1±0.03.1± 0.0 0.6±0.00.6± 0.0 0.4±0.10.4± 0.1 Table˜5 shows that MPGD achieves zero boundary and collision violations across all various sampling steps, but its physical plausibility depends strongly on the number of denoising steps. With 5050 steps, matching the main evaluation setting, MPGD exhibits large frame-to-frame changes (F2F) and high energy deviation (MED), indicating that feasibility is obtained through many incremental corrections. This supports the claim that sampling-time correction methods can enforce constraints, but may rely on many incremental updates to avoid physically infeasible trajectories. D.4 Scaling network parameterization The third MLP outputs αϕ _φ and a raw exponent parameter βϕraw _φ^raw. We use a softplus output for αϕ _φ to ensure a nonnegative magnitude, and map the raw exponent to the effective component as 2sigmoid(βϕraw)−42sigmoid( _φ^raw)-4. This offset incorporates the known scaling of the denoised-state guidance term: the denoiser Jacobian contributes a factor cov[0∣t]/t2cov [x_0 x_t ]/t^2, and the score parameterization ([0∣t]−t)/t2(E [x_0 x_t ]-x_t)/t^2 contributes an additional t−2t^-2 factor. Thus the t−4t^-4 component is fixed analytically, while the network learns the remaining state- and noise-dependent covariance modulation. This parameterization preserves the desired divergence property as noise level approaches zero without requiring a manually tuned schedule. D.5 Training details We follow the prior work [26] for the experimental setup, using the same transformer-based diffusion backbone [45, 53, 3] and gradient embedding MLP. We use the same log-linear noise schedule for sampling with σmin=3×10−5 _ =3× 10^-5 and σmax=80 _ =80 but use 50 sampling steps. The standard EDM model is trained from scratch for 950950k iterations with learning rate 3×10−43× 10^-4 on a Tesla V100 GPU, taking approximately 7070 hours. Our proposed fine-tuned model additionally uses the scaling network described in Appendix˜C, and includes LoRA adapters. Our model and other baseline models are either initialized from scratch or from the pretrained EDM checkpoint, as specified in their method descriptions, and are trained with learning rate 3×10−53× 10^-5 for approximately 7272 hours on an NVIDIA L40S GPU. The number of training iterations varies across methods due to differences in computational cost. All models use the Adam optimizer; our method and DPOK use batch size 1616, while other fine-tuning models use batch size 3232. Appendix E Traffic scene trajectory prediction experiment details E.1 Definition of constraint losses and evaluation metrics We provide the definition of the constraint losses and evaluation metrics used in the traffic scene prediction experiment. Each scene contains a varying number of agents. Given the first 10 frames observed, the model jointly predicts all agents over a future prediction horizon =30T=30 frames. We denote the predicted center position of the oriented bounding box for agent a at frame τ by τ(a)∈ℝ2x_τ^(a)∈R^2, and denote the corresponding bounding box by τ(a)O_τ^(a). In our implementation, the first agent in the predicted scene tensor is the ego agent. Constraint losses The offroad loss measures how far a vehicle extends outside the drivable region. Given the drivable mesh ℳM, we compute the four corners of each vehicle box. For a corner τ,jax_τ,j^a, where j=1,⋯,4j=\1,·s,4\, we define its distance to drivable mesh as dmesh(τ,j(a))=min∈ℳ∥τ,j(a)−∥22. d_mesh (c_τ,j^(a) )= _m∈M c_τ,j^(a)-m _2^2. (55) The offroad loss is the maximum distance among corners that lie outside the drivable region: ℓoffroad(τ(a))=maxj⁡[τ,j(a)∈ℳ]dmesh(τ,j(a)). _offroad (O_τ^(a) )= _j 1 [c_τ,j^(a)∈M ]d_mesh (c_τ,j^(a) ). (56) The collision loss measures the overlap area between pairs of vehicles. For two agents a and a′a , we define ℓcollision(τ(a),τ(a′))=Area(τ(a),τ(a′)). _collision (O_τ^(a),O_τ^(a ) )=Area (O_τ^(a),O_τ^(a ) ). (57) Reported metrics For each scene, we generate K=6K=6 predicted futures. We use k∈1,⋯Kk∈\1,·s K\ to index generated samples and a=0a=0 to denote the ego agent. The offroad rate is the fraction of scenes in which the ego vehicle goes offroad at least once during the future prediction horizon Offroad rate=1K∑k=1K[maxτ=1,⋯,⁡ℓoffroad(τk,(0))>0]. rate= 1K _k=1^K 1 [ _τ=1,·s,T _offroad(O_τ^k,(0))>0 ]. (58) The collision rate is the fraction of scenes in which the ego vehicle collides at least once with another vehicle during the future prediction Collision rate=1K∑k=1K[maxτ=1,⋯,⁡ℓcollision(τk,(0))>0]. rate= 1K _k=1^K 1 [ _τ=1,·s,T _collision(O_τ^k,(0))>0 ]. (59) The scene maximum frame-to-frame displacement (SF2F) measures abrupt motion or jitter, same as in the bouncing ball experiment: SF2F6=maxk,a,τ<∥τ+1k,(a)−τk,(a)∥2. 2F_6= _k,a,τ<T x_τ+1^k,(a)-x_τ^k,(a) _2. (60) Both offroad and collision rates are then averaged over evaluation batches. The scene maximum lateral velocity (SMLV) measures sideways motion relative to each vehicle heading. We estimate velocity from consecutive predicted positions: τk,(a)=τ+1k,(a)−τk,(a). v_τ^k,(a)=x_τ+1^k,(a)-x_τ^k,(a). (61) Given heading hτk,(a)h_τ^k,(a), the lateral velocity is lat,τk,(a)=−τ,xk,(a)sin⁡hτk,(a)+τk,(a)cos⁡hτk,(a). v_lat,τ^k,(a)=-v_τ,x^k,(a) h_τ^k,(a)+v_τ^k,(a) h_τ^k,(a). (62) We define maximum lateral velocity as SMLV6=maxk,a,τ<⁡∥lat,τk,(a)∥. _6= _k,a,τ<T v_lat,τ^k,(a) . (63) The maximum final distance (MFD) measures endpoint diversity among generated ego trajectories. We compute the maximum pairwise distance between final predicted ego position across K generated samples, MFD6=maxk,k′⁡∥k,(0)−k′,(0)∥. _6= _k,k x_T^k,(0)-x_T^k ,(0) . (64) SF2F6, SMLV6, MFD6 are computed within K=6K=6 generated samples, and then average over evaluation batches. Figure 4: Scaling network output versus noise level σ along denoising trajectory for traffic scene trajectory prediction experiment. Means are taken over batches; shaded bands show ±1± 1 standard error. E.2 Method variants We report two additional variants to our IMCD in Table˜6: Rollout-only trains the scaling network and gradient embedding without LoRA adaptation; and IMCD (staged) uses the same trainable components and training objective as IMCD, but initializes from the trained Rollout-only checkpoint before enabling LoRA adapters for a second fine-tuning stage. IMCD is the variant reported in Table˜2. As shown in Table˜6, IMCD (staged) further reduces the offroad rate, but requires additional training through the initial rollout-only stage. E.3 Effect of scaling network and its parameterization As shown in Figure˜4, αϕ(t;t) _φ(x_t;t) and βϕ(t;t) _φ(x_t;t) both change with the noise level, so the learned scaling does not collapse to a constant coefficient. Unlike the bouncing balls case, t2γϕ(t;t)t^2 _φ(x_t;t) peaks at an intermediate noise level, suggesting that the model places the strongest relative correction in the middle of the trajectory. The final scale γϕ(t;t) _φ(x_t;t) still grows rapidly as t decreases, preserving strong constraint enforcement near the data manifold. For parameterization, the third MLP outputs αϕ _φ and a raw exponent parameter βϕraw _φ^raw as before. We apply a softplus to αϕ _φ to keep the correction magnitude nonnegative, and map the raw exponent to the effective exponent as 2sigmoid(βϕraw)−32sigmoid( _φ^raw)-3. This effective exponent, rather than the raw MLP output, is the βϕ _φ used in the scaling function. This retains the t−2t^-2 factor induced by the score parameterization, while using the remaining t−1t^-1 factor as a milder base scaling. Compared with the −4-4 offset used in bouncing balls, this reduces the strength of low-noise guidance. We conjecture that this milder scaling is better suited to trajectory prediction because it uses few sampling steps, where overly aggressive low-noise guidance can lead to abrupt late-stage correction. The bounded learned exponent keeps the divergence property as t goes to zero. E.4 Training details We follow prior work [26] for the experimental setup, using the same transformer-based diffusion backbone [34] and gradient embedding MLP as the bouncing ball experiment. Our fine-tuned models additionally use the scaling network described in Appendix˜C and LoRA adapters. We use the same log-linear sampling noise schedule with σmax=80 _ =80 and σmin=2×10−4 _ =2× 10^-4, and sample with 2020 steps. The standard diffusion model is trained from scratch with learning rate 3×10−43× 10^-4 on a Tesla V100 GPU, taking approximately 168 hours for about 140 epochs. Our fine-tuned models, and other fine-tuned baselines are trained with learning rate 3×10−43× 10^-4 for approximately 7272 hours on two NVIDIA L40S GPUs, except for PIDM, which uses learning rate 3×10−53× 10^-5, and IMCD (staged). The IMCD (staged) variant follows the staged procedure described above: the Rollout-only model is trained for three days, and the LoRA adapters are then trained for another three days. The number of training iterations varies across methods due to differences in computational cost. All models use the Adam optimizer with batch size 1616. E.5 More experimental results Table 6: Traffic scene trajectory prediction results on the INTERACTION DR_DEU_Merging_MT scenario. Constraint (%) ↓ Min traj. match ↓ Phys. plaus. ↓ Diversity ↑ Method Offroad Collision ADE6 FDE6 SADE6 SFDE6 SF2F6 SMLV6 MFD6 Rollout-only 0.490.49 0.020.02 0.180.18 0.450.45 0.230.23 0.660.66 1.101.10 0.490.49 1.921.92 IMCD 0.300.30 0.000.00 0.180.18 0.450.45 0.230.23 0.660.66 1.071.07 0.430.43 1.911.91 IMCD (staged) 0.040.04 0.000.00 0.190.19 0.470.47 0.240.24 0.690.69 1.021.02 0.370.37 1.961.96 Table 7: Traffic scene trajectory prediction results on the INTERACTION DR_DEU_Roundabout_OF scenario. Constraint (%) ↓ Min traj. match ↓ Phys. plaus. ↓ Diversity ↑ Method Offroad Collision ADE6 FDE6 SADE6 SFDE6 SF2F6 SMLV6 MFD6 EDM [23] 23.9223.92 1.701.70 0.350.35 1.111.11 0.500.50 1.801.80 1.151.15 0.490.49 6.006.00 MPGD [19] 11.9611.96 1.421.42 0.330.33 1.041.04 0.450.45 1.571.57 1.241.24 0.650.65 3.983.98 MBM++ [26] 12.4212.42 1.401.40 0.320.32 1.051.05 0.440.44 1.591.59 1.211.21 0.610.61 3.913.91 PIDM (FT) [2] 14.1214.12 1.351.35 0.340.34 1.241.24 0.490.49 1.931.93 0.960.96 0.330.33 5.005.00 DPOK [14] 3.733.73 0.110.11 2.992.99 8.888.88 2.962.96 9.949.94 0.970.97 0.560.56 2.422.42 CriticSMC [27] – 0.080.08 0.450.45 – – – – – 3.433.43 Rollout-only 0.680.68 0.030.03 0.370.37 1.061.06 0.510.51 1.631.63 2.072.07 1.351.35 4.074.07 IMCD 1.221.22 0.010.01 0.370.37 1.071.07 0.510.51 1.661.66 1.691.69 1.061.06 4.024.02 IMCD (staged) 0.370.37 0.000.00 0.370.37 1.081.08 0.510.51 1.661.66 1.681.68 1.051.05 3.953.95 Ground truth 0.000.00 0.000.00 – – – – 0.840.84 0.060.06 – Table 8: Traffic scene trajectory prediction results on the INTERACTION DR_USA_Intersection_MA Constraint (%) ↓ Min traj. match ↓ Phys. plaus. ↓ Diversity ↑ Method Offroad Collision ADE6 FDE6 SADE6 SFDE6 F2F6 MLV6 MFD6 EDM [23] 1.361.36 1.541.54 0.240.24 0.660.66 0.400.40 1.371.37 1.291.29 0.460.46 4.444.44 MPGD [19] 0.490.49 0.830.83 0.220.22 0.640.64 0.350.35 1.221.22 1.241.24 0.420.42 3.173.17 MBM++ [26] 0.420.42 0.830.83 0.200.20 0.630.63 0.330.33 1.211.21 1.221.22 0.360.36 3.113.11 PIDM (FT) [2] 1.151.15 1.201.20 0.210.21 0.670.67 0.360.36 1.391.39 1.201.20 0.280.28 4.264.26 DPOK [14] 0.370.37 0.040.04 2.072.07 6.286.28 2.102.10 6.606.60 0.960.96 0.520.52 1.771.77 CriticSMC [27] – 0.090.09 0.450.45 – – – – – 2.872.87 Rollout-only 0.030.03 0.000.00 0.230.23 0.640.64 0.350.35 1.221.22 1.331.33 0.530.53 3.173.17 IMCD 0.050.05 0.000.00 0.220.22 0.650.65 0.350.35 1.231.23 1.261.26 0.460.46 3.173.17 IMCD (staged) 0.010.01 0.000.00 0.230.23 0.670.67 0.360.36 1.261.26 1.241.24 0.450.45 3.173.17 Ground truth 0.000.00 0.000.00 – – – – 1.111.11 0.050.05 – Traffic scene trajectory prediction results on the INTERACTION DR_USA_Roundabout_FT Constraint (%) ↓ Min traj. match ↓ Phys. plaus. ↓ Diversity ↑ Method Offroad Collision ADE6 FDE6 SADE6 SFDE6 F2F6 MLV6 MFD6 EDM [23] 9.329.32 1.371.37 0.250.25 0.720.72 0.410.41 1.441.44 1.231.23 0.470.47 4.264.26 MPGD [19] 2.642.64 0.930.93 0.230.23 0.680.68 0.350.35 1.241.24 1.211.21 0.500.50 2.972.97 MBM++ [26] 2.492.49 0.920.92 0.210.21 0.680.68 0.340.34 1.231.23 1.181.18 0.460.46 2.922.92 PIDM (FT) [2] 5.795.79 0.860.86 0.230.23 0.830.83 0.390.39 1.531.53 1.021.02 0.310.31 3.873.87 DPOK [14] 1.951.95 0.060.06 2.552.55 7.847.84 2.652.65 8.258.25 1.041.04 0.580.58 1.951.95 CriticSMC [27] – 0.070.07 0.440.44 – – – – – 2.972.97 Rollout-only 0.280.28 0.110.11 0.240.24 0.700.70 0.360.36 1.261.26 1.601.60 1.011.01 3.013.01 IMCD 0.250.25 0.060.06 0.240.24 0.710.71 0.370.37 1.281.28 1.371.37 0.750.75 2.942.94 IMCD (staged) 0.160.16 0.050.05 0.240.24 0.720.72 0.360.36 1.291.29 1.351.35 0.770.77 3.003.00 Ground truth 0.000.00 0.000.00 – – – – 0.950.95 0.060.06 – In the rendered visualizations below, we show the final predicted frame so that the full trajectory is visible. Gray dot trajectories indicate ground-truth agent motion, and orange trajectories denote model predictions. Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged) Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged) Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged) Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged) Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged) Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged) Standard Diffusion MPGD w/o projection MBM++ PIDM DPOK Rollout-only IMCD IMCD (staged)

Instant research discovery