Conditional Diffusion Guidance under Hard Constraint: A Stochastic Analysis Approach

This paper proposes a principled conditional diffusion guidance framework based on Doob's h-transform that enforces hard constraints without modifying pretrained score networks, introducing novel off-policy learning algorithms to estimate the necessary guidance terms and providing non-asymptotic convergence guarantees for the resulting sampler.

The Gist

Here is an explanation of the paper "Conditional Diffusion Guidance under Hard Constraint" using simple language, analogies, and metaphors.

The Big Picture: The "Magic Paintbrush" with a Strict Rulebook

Imagine you have a Magic Paintbrush (a Diffusion Model) that has spent years learning to paint beautiful landscapes. It knows how to paint mountains, rivers, and forests perfectly. This is the "pretrained model."

Now, imagine a client comes to you and says: "I love your landscapes, but I have a strict rule: Every single painting you make for me must contain a golden bridge. If even one painting doesn't have a golden bridge, it's useless to me."

This is the problem the paper solves.

• Soft Guidance (The Old Way): Most AI tools try to satisfy this by saying, "Okay, I'll try really hard to paint a bridge. I'll add a 'reward' if I paint a bridge." But sometimes, the AI gets lazy or confused and paints a bridge that looks like a tree, or forgets it entirely. It's a "soft" promise.
• Hard Constraint (The New Way): The client says, "No excuses. 100% of the time, the bridge must be there." This is a "hard constraint."

The authors of this paper invented a new mathematical "guide" that forces the Magic Paintbrush to obey this rule guaranteed, without needing to retrain the whole brush from scratch.

---

The Core Problem: Why is this so hard?

Usually, if you want the AI to paint a specific thing (like a golden bridge), you might try Rejection Sampling.

• The Analogy: Imagine you ask the AI to paint 1,000 landscapes. You check them one by one. If a painting doesn't have a golden bridge, you throw it in the trash and ask for another.
• The Problem: If "golden bridges" are rare (maybe only 1 in 10,000 paintings has one), you have to throw away 9,999 paintings just to get one good one. This is incredibly slow and expensive.

The authors wanted a way to change the AI's brain so that it only paints landscapes with bridges, without throwing anything away.

The Solution: The "Doob's H-Transform" (The GPS)

The paper uses a concept from advanced math called Doob's h-transform. Let's translate that into a GPS analogy.

1. The Original Path: The AI usually wanders randomly through a forest of possibilities to find a landscape.
2. The Goal: The client wants the AI to end up at a specific destination (The Golden Bridge).
3. The Magic: The authors realized that if you know the probability of reaching the destination from anywhere in the forest, you can draw a GPS line that pulls the AI toward the bridge at every single step.

They call this the $h$-function. It's like a "Hope Meter."

• If the AI is in a spot where it's easy to get to the bridge, the Hope Meter is high.
• If the AI is in a dead end, the Hope Meter is low.

The new algorithm adds a "drift" (a gentle push) to the AI's movement. If the AI starts drifting away from the bridge, the GPS pushes it back. If it's heading toward the bridge, the GPS lets it glide. Crucially, this happens at every single step of the painting process, not just at the end.

How They Taught the AI (The Two Algorithms)

The tricky part is: The AI doesn't know the "Hope Meter" ($h$) or the "GPS direction" ($\nabla \log h$) yet. It has to learn it. The authors proposed two clever ways to teach the AI using only the paintings it already made (without needing new data).

1. The "Martingale Loss" (The Consistency Check)

• The Analogy: Imagine you are betting on a horse race. You know the horse will win (the constraint).
• The Logic: If you are at the start of the race, your "probability of winning" is low. As the race goes on, if the horse is running well, your "probability of winning" goes up.
• The Trick: The authors realized that if you track this "probability of winning" as the AI paints, it should behave like a fair game (a martingale). It shouldn't jump up and down randomly; it should move smoothly toward 100% certainty at the end.
• The Algorithm (CDG-ML): They trained a small neural network to predict this "probability" so that it behaves like a fair game. If the prediction is wobbly, they correct it.

2. The "Martingale-Covariation Loss" (The Speedometer)

• The Problem: Knowing the "probability" (the Hope Meter) isn't enough. You need to know the direction and speed to push the AI. It's hard to get the direction right just by looking at the probability.
• The Analogy: Imagine you are driving a car. You know your destination is 10 miles away. But you also need to know how fast you are moving toward it.
• The Trick: The authors used a concept called Quadratic Variation. In math, this is like looking at the "jitter" or "wiggles" of the path.
• The Algorithm (CDG-MCL): They realized that the "wiggles" in the AI's path contain hidden information about the direction it needs to go. They taught a second network to read these wiggles to figure out exactly how to steer the AI toward the constraint.

Why is this cool? Both methods are Off-Policy.

• On-Policy (Old way): The AI tries to learn while it's driving the car. If it crashes, it has to start over. It's dangerous and unstable.
• Off-Policy (New way): The AI learns by watching recordings of its old, unguided drives. It learns the rules of the road without ever crashing. This makes it much safer and faster.

The Results: Stress Testing and Rare Events

The paper tested this on two real-world scenarios:

1. Finance (Stress Testing):

   • Scenario: "What happens to a stock portfolio if the market crashes?"
   • The Issue: Market crashes are rare. Normal AI models rarely generate them because they are trained on "normal" days.
   • The Fix: The authors forced the AI to generate "Crash Scenarios" (Hard Constraint).
   • Result: The AI could generate realistic-looking market crashes that followed all the complex rules of how stocks move together. This helps banks prepare for disasters.

2. Supply Chain (Hospital Queues):

   • Scenario: "What happens to a hospital if a flu season hits and everyone gets sick at once?"
   • The Issue: Normal simulations assume average traffic.
   • The Fix: They forced the AI to simulate a "Flu Season" where patients arrive faster and recover slower.
   • Result: The simulation showed exactly which hospital wards would get overwhelmed, allowing managers to add beds before the crisis happens.

Summary: Why This Matters

• Safety First: In critical fields (medicine, finance, engineering), you can't have "maybe" constraints. You need guarantees. This paper provides a mathematical guarantee that the AI will obey the rules.
• Efficiency: It stops the AI from wasting time generating bad samples that get thrown away.
• Smart Learning: It teaches the AI how to follow strict rules by analyzing its own past behavior, rather than forcing it to learn from scratch.

In a nutshell: The authors built a GPS for AI that guarantees the AI never misses its destination, even if that destination is a rare, difficult, or dangerous event. They did this by teaching the AI to read the "wiggles" in its own path to find the right way forward.

Technical Summary

Here is a detailed technical summary of the paper "Conditional Diffusion Guidance under Hard Constraint: A Stochastic Analysis Approach."

1. Problem Statement

The paper addresses the challenge of conditional generation in diffusion models under hard constraints.

• Context: In safety-critical applications (e.g., autonomous driving, medical diagnosis) and rare-event simulation (e.g., financial stress testing, system overloads), generated samples must satisfy specific constraints with probability one (almost surely).
• Limitation of Existing Methods: Current guidance techniques (soft guidance, classifier guidance, RLHF-based methods) typically optimize for reward or penalty terms. While computationally convenient, they offer no guarantee that the final sample satisfies the constraint, often resulting in violations or requiring inefficient rejection sampling (which has an expected cost of $O(1/\rho)$ for rare events with probability $\rho$).
• Goal: Develop a principled framework to modify the sampling dynamics of a pretrained diffusion model to enforce hard constraints without retraining the underlying score network, ensuring the terminal output satisfies the constraint $Y_T \in S$ with probability 1.

2. Methodology

The authors propose a framework rooted in stochastic analysis, specifically leveraging Doob's $h$-transform, martingale theory, and quadratic variation.

A. Theoretical Foundation: Doob's $h$-Transform

The core idea is to change the measure of the diffusion process to condition it on the event $\{Y_T \in S\}$.

• Let $h(t, y) = P(Y_T \in S \mid Y_t = y)$ be the conditional probability of satisfying the constraint given the current state.
• The guided dynamics $Y^S_t$ are derived by adding a drift correction term to the pretrained SDE:
  $$dY^S_t = \left( \bar{f}(t, Y^S_t) + s_\theta(t, Y^S_t) + g(t)^2 \nabla \log h(t, Y^S_t) \right) dt + g(t) dB_t$$
• Here, $s_\theta$ is the pretrained score function, and $\nabla \log h$ is the guidance term. Crucially, this approach does not modify $s_\theta$; it only augments the drift.

B. Learning the Guidance Function ($h$ and $\nabla h$)

Since $h(t, y)$ is unknown, the authors propose two novel off-policy learning algorithms that estimate $h$ and its gradient using only trajectories from the pretrained model (no access to the true data distribution or on-policy sampling required).

1. CDG-ML (Martingale Loss):

   • Exploits the fact that $\{h(t, Y_t)\}_{t \geq 0}$ is a martingale under the pretrained dynamics.
   • Objective: Minimize the $L_2$ loss to learn $h$:
     $$\min_{\ell} \mathbb{E} \left[ \int_0^T \left( \ell(t, Y_t) - \mathbb{1}(Y_T \in S) \right)^2 dt \right]$$
   • The guidance term is approximated as $\nabla \log h_\phi \approx \nabla h_\phi / h_\phi$.

2. CDG-MCL (Martingale–Covariation Loss):

   • Addresses the difficulty of accurately estimating $\nabla \log h$ when $h$ is small (common in rare events).
   • Exploits the quadratic covariation identity: $d[h, Y]_t = g(t)^2 \nabla h(t, Y_t) dt$.
   • Objective: Learn the gradient $\nabla h$ directly by minimizing the error between the empirical covariation and the model prediction:
     $$\min_{q} \mathbb{E} \left[ \int_0^T \left( \frac{1}{g(t)^2} \frac{d[h_\phi, Y]_t}{dt} - q(t, Y_t) \right)^2 dt \right]$$
   • The final guidance is $\frac{q_{\psi^*}}{h_{\phi^*}}$.

C. Extensions

• ODE Sampling: The framework extends to probability-flow ODE samplers for higher efficiency.
• Reinforced Conditioning: Introduces a guidance scale $\eta$ (similar to classifier guidance) to trade off between constraint strictness and sample diversity.

3. Key Contributions

1. Principled Hard Constraint Framework: A theoretically grounded method to enforce hard constraints in diffusion models using Doob's $h$-transform, guaranteeing $P(Y_T \in S) = 1$.
2. Novel Off-Policy Learning: Two new algorithms (CDG-ML and CDG-MCL) that learn the conditioning function and its gradient using only pretrained trajectories, avoiding the instability and distribution shift associated with on-policy reinforcement learning or control-based methods.
3. Rigorous Non-Asymptotic Guarantees:
   • Total Variation (TV) Bounds: Proved bounds on the distance between the target conditional distribution and the guided sampler, explicitly decomposing errors from score approximation and guidance estimation.
   • Wasserstein-2 Bounds: Derived geometric error bounds under regularity assumptions, showing stability even for rare events where TV bounds might scale poorly with $1/\rho$.
4. Convergence Analysis: Provided convergence rates for the stochastic optimization procedures used to learn $h$ and $\nabla h$.

4. Experimental Results

The authors validated the framework on three types of tasks:

• Synthetic Examples:

  • Demonstrated on 1D and 2D Gaussian distributions with truncated constraints.
  • Result: Both algorithms successfully generated samples matching the target conditional distribution. CDG-MCL (using covariation) generally achieved lower Wasserstein distances and better fit than CDG-ML, particularly in capturing the gradient structure.

• Financial Stress Testing:

  • Applied to US stock returns (AAPL, AMZN, TSLA, JPM) to simulate rare market crashes (e.g., TSLA dropping >10% in 10 days).
  • Result: The framework generated realistic portfolio return distributions under stress scenarios. CDG-ML allowed for larger guidance scales ($\eta \approx 150$) with stable performance, while CDG-MCL required smaller scales ($\eta \approx 2\text{--}4$). Both methods accurately captured tail risks (quantiles) compared to real out-of-sample data.

• Supply Chain Simulation:

  • Modeled a hospital queueing system (QGym) under seasonal stress (flu season).
  • Result: The soft-guidance framework successfully simulated unstable regimes (increased arrivals, decreased service rates) where hard truncation would fail or be too rigid. It captured the "smoothing" effect of soft constraints, preventing the system from diverging immediately while still reflecting the stress scenario.

5. Significance

• Safety and Reliability: Provides a mathematically rigorous solution for applications where constraint violation is unacceptable (e.g., safety-critical systems), overcoming the probabilistic nature of soft guidance.
• Rare-Event Simulation: Offers a computationally efficient alternative to rejection sampling for generating rare events, which is crucial for risk management in finance and engineering.
• Theoretical Advancement: Bridges the gap between diffusion models and classical stochastic analysis (martingales, $h$-transforms), providing the first non-asymptotic error analysis for conditional diffusion guidance under hard constraints.
• Practical Utility: The off-policy nature of the learning algorithms makes the method lightweight and applicable to existing pretrained models without expensive retraining.