End-to-End Differentiable Diffusion Conditioning for Physics-Informed Optimization

AI summary

Overview

Research area: Machine learning for scientific design — specifically offline model-based optimization (MBO) with conditional diffusion models, applied to physics-constrained engineering problems.

Technical level: Advanced. The paper assumes familiarity with denoising diffusion, DDIM sampling, classifier-free guidance, and surrogate-based optimization.

Scope: The paper introduces Differentiable Diffusion Conditioning (D2C), a method that turns a frozen conditional diffusion model into a gradient-based optimizer by backpropagating through the deterministic reverse sampling chain to update the conditioning signal, and evaluates it on two real-world physical design tasks plus four Design-Bench tasks.

What This Paper Is About

Many scientific design problems — shaping a laser pulse for fusion experiments, or scheduling workloads to cut a data center's carbon footprint — involve expensive evaluations, high-dimensional design spaces, and hard physical constraints that cannot be violated. In offline model-based optimization, a method gets only a fixed dataset of previously evaluated designs and no access to the true objective at test time. The paper's goal is to let a pretrained generative model keep refining its output against precise physics-based objectives at test time, rather than sampling once from a fixed target and stopping there.

Key Contributions

1. End-to-end conditioning optimization through diffusion. The authors exploit the fact that deterministic DDIM sampling (with zero stochastic noise) makes the full reverse chain differentiable with respect to the conditioning input. This lets gradients flow from a composite physics objective back to the conditioning variables, making the conditioning signal — not the latent noise — the optimization variable. The paper reports that optimizing conditioning outperforms optimizing the latent noise (Appendix D.7).

2. A frozen proposer–evaluator architecture. A conditional diffusion model (the proposer) generates candidate designs, and a separately trained forward surrogate (the evaluator) scores them. Both stay frozen at test time; only the conditioning is updated, which the authors argue limits out-of-distribution surrogate exploitation.

3. Physics-informed multi-objective optimization. Heterogeneous objectives — a surrogate objective term, constraint penalties (integral conservation, amplitude bounds, spectral stability), and a smoothness term — are folded directly into the differentiable sampling chain, with loss weights drawn from a symmetric Dirichlet distribution (α_k = 1) so each trajectory emphasizes a different trade-off.

4. Empirical results on real-world and benchmark tasks. D2C reports the best results on all three real-world settings and the best average rank (1.75) across four Design-Bench tasks.

Main Findings

• ICF Standard (2,000 randomly sampled pulses): D2C improves neutron yield by 17.98% ± 1.54, versus the best baseline GTG at 5.43% ± 0.74.

• ICF Top (1,000 highest-yield designs, the hardest cases): D2C reaches 10.06% ± 0.88, more than two and a half times the best baseline, Brute Force, at 3.93% ± 0.52.

• Data center scheduling: D2C reduces carbon emissions by 8.98% ± 0.75, versus the best baseline GP-qEI at 5.93% ± 0.82 and the best generative baseline GTG at 5.50% ± 0.68.

• Design-Bench headline: D2C places first on TF Bind 8 (0.973 ± 0.011) and TF Bind 10 (0.684 ± 0.023), second on D'Kitty (0.974 ± 0.009, behind GTG's 0.976 ± 0.009), and third on Ant (0.876 ± 0.020, behind DDOM's 0.925 ± 0.019).

• Average rank: D2C's average rank of 1.75 across the four Design-Bench tasks beats GTG (3.50), DDOM (4.38), and all other baselines listed.

• Beating the dataset's best example: On TF Bind 8, TF Bind 10, D'Kitty, and Ant, D2C's scores exceed the best design present in the offline dataset (0.439, 0.467, 0.884, and 0.565 respectively) under the evaluation protocol used.

• Conditioning beats latent-noise optimization: With proposer, surrogate, and evaluation protocol held fixed, optimizing the conditioning outperforms optimizing the latent noise on both real-world tasks.

• Cost is the trade-off: Backpropagating through the DDIM chain at each iteration makes D2C roughly 8–17× slower than single-pass methods (Table 16).

• Constraint handling is validity-first: Only candidates passing domain-specific hard constraints count as improvements; if none pass, the baseline design is returned and contributes 0% gain. The paper states that method quality should therefore be read jointly through achieved gain and the rate at which valid improving solutions are found.

Methodology in Plain English

The setup has three moving parts. First, a conditional diffusion model is trained offline on a dataset of designs paired with their measured performance metrics. It learns to generate a design (say, a 256-step laser power profile) when told a set of target metrics. Second, a surrogate model is trained on the same data to predict those performance metrics from a design. Third, at test time, both models are frozen and only the input targets are adjusted.

The trick is how the targets are adjusted. Standard diffusion sampling is random, but the deterministic DDIM variant is a fixed chain of arithmetic steps from a starting noise vector to a final design. Because every step is a differentiable operation on the conditioning input, you can push a gradient backwards through the entire chain: score the produced design with the frozen surrogate and the physics penalties, compute the loss, and propagate the gradient all the way back to the target metrics. Gradient descent then nudges those targets, the design changes, and the loop repeats.

To avoid getting stuck in one mode, the method draws N_z = 20 independent noise vectors for real-world tasks, each held fixed during optimization, and within each, samples N_λ = 5 loss-weight vectors from a Dirichlet distribution — yielding 100 candidate trajectories per sample. For Design-Bench, 100 latent trials are used with a single objective weight. Every candidate is then checked against hard physical constraints, and the best valid one is selected. Classifier-free guidance is used both to bias generation toward the conditioning and, per the authors, to provide a stable regime for gradient updates on the targets.

Why This Matters

Impact on research. The paper identifies a specific gap: inverse and trajectory-based MBO methods treat the conditioning target as a static input fixed before generation, so they cannot refine a candidate against a downstream objective. D2C shows that the deterministic diffusion sampler itself can serve as a differentiable bridge, and that optimizing the conditioning is a better use of that bridge than optimizing the latent noise. It also demonstrates this on genuinely constrained physical problems where violating a constraint is not merely a scoring penalty but a safety failure.

Real-world applications:

• Inertial confinement fusion: shaping 256-step laser power profiles to raise neutron yield while conserving pulse energy, staying temporally smooth enough to manufacture, and keeping frequency content physically plausible.
• Sustainable data centers: building 24-hour workload power profiles at 10-minute resolution that cut carbon emissions while delivering exactly the demanded compute and respecting inlet temperature limits.
• Biomolecular design: the TF Bind 8 and TF Bind 10 tasks involve designing DNA sequences for transcription factor binding.
• Robot morphology: the Ant and D'Kitty tasks optimize physical robot body parameters.

Industry relevance. The method requires no retraining or fine-tuning of the generative model — it is a test-time procedure layered on top of existing frozen assets. That is attractive for organizations that have already invested in trained diffusion models and surrogates and want to extract more optimization value from them without a new training pipeline. The 8–17× inference cost is the explicit caveat for latency-sensitive deployment.

Future Directions

• Calibrating the surrogate. The authors note the gradient signal is bounded by the accuracy of the frozen surrogate, and suggest incorporating uncertainty quantification so that optimization does not converge on high-scoring surrogate designs that fail on the true objective.
• Cutting the compute cost. Backpropagating through the full T-step DDIM chain dominates runtime. The ablations reportedly show that the number of DDIM steps and optimization steps can be reduced with minimal quality loss, but the method remains unsuitable for latency-critical use.
• Better constraint satisfaction rates. Because invalid candidates are discarded and the baseline is kept when none pass, improving how often a valid improving design is found is a direct lever on reported gains.
• Broadening beyond the two studied domains. The method's loss terms, constraint definitions, and tuning are domain-specific; how much of that transfers automatically to other physics-constrained design problems is left open.

Target Audience

Researchers and practitioners working on offline model-based optimization, diffusion-based generative design, or surrogate-assisted engineering optimization. It will be most valuable to readers already comfortable with diffusion sampling internals who want to apply generative models to constrained scientific design — particularly in fusion energy, energy-efficient computing, and other settings where evaluations are expensive and constraint violations are unacceptable. Readers new to diffusion models will find the preliminaries section helpful but should expect a demanding read.

Authors’ abstract

Optimizing high-dimensional physical systems under expensive evaluations and complex constraints is a central challenge in science and engineering. In offline model-based optimization (MBO), methods must improve designs using only a fixed dataset and no test-time access to the true objective. However, existing methods typically suffer from conservatism, unconstrained surrogate exploitation, or function as static inference-time samplers that cannot refine candidates against precise physical objectives. In this paper, we propose Differentiable Diffusion Conditioning (D2C) for offline MBO, which leverages the fact that deterministic sampling makes the reverse denoising chain differentiable with respect to the conditioning signal. D2C turns a frozen conditional diffusion model into an end-to-end gradient-based optimizer by directly optimizing conditioning variables at test time via gradients from frozen surrogate objectives and physics-informed penalties backpropagated through the reverse chain. We implement this idea in a composed proposer-evaluator framework and evaluate it on laser pulse optimization for inertial confinement fusion, sustainable data-center workload scheduling, and four Design-Bench tasks. Our method achieves the strongest results on the two real-world tasks and the best average rank across four Design-Bench tasks.