On the Collapse of Generative Paths: A Criterion and Correction for Diffusion Steering

This paper identifies "Marginal Path Collapse" as a critical failure mode in inference-time steering of diffusion models caused by mismatched noise schedules or negative exponents, and proposes the Adaptive Path Correction with Exponents (ACE) framework to mathematically guarantee path existence and significantly improve performance in complex compositional tasks like drug design and image generation.

The Gist

Imagine you are trying to bake the perfect cake. You have three different expert chefs, each with their own unique recipe and style:

1. Chef A is great at making the cake base (the structure).
2. Chef B is a master at adding the right flavor (the specific taste).
3. Chef C is an expert at decorating the top (the final look).

In the world of AI image and molecule generation, we often want to combine these "experts" to create something new without training a brand-new chef from scratch. We try to mix their instructions together.

The Problem: The "Collapse"
The paper identifies a hidden disaster that happens when you try to mix these chefs, especially if they were trained using different "timers" or "noise schedules" (think of this as them working at different speeds or using different measuring cups).

When you try to combine their instructions, the math sometimes breaks in the middle of the process. The paper calls this "Marginal Path Collapse."

Here is a simple analogy: Imagine the chefs are trying to guide a ball from a starting point (pure noise) to a finish line (the perfect cake).

• The Goal: The ball should roll smoothly along a clear path.
• The Collapse: Because the chefs are using different rules, the path suddenly disappears or turns into a bottomless pit in the middle of the journey. The ball falls off the edge. The AI tries to keep rolling, but it's now rolling through a "ghost" path that doesn't actually exist. It might still move, but it ends up at the wrong destination, or it creates a broken, nonsensical result (like a molecule that falls apart or an image with weird artifacts).

The paper notes that this isn't a rare glitch; it happens very often when combining different types of AI models, especially in complex tasks like designing new medicines.

The Solution: ACE (Adaptive Path Correction with Exponents)
The authors propose a fix called ACE. Think of ACE as a smart traffic controller that watches the chefs in real-time.

1. The Check-Up (The Criterion): Before the ball starts rolling, ACE checks the math to see if the path is safe. It asks, "Is there a solid road ahead, or is there a cliff?"
2. The Adjustment (The Correction): If the path looks shaky or about to collapse, ACE doesn't just let the ball fall. It gently nudges the chefs' instructions. It slightly changes how much weight it gives to each chef's advice at every single moment of the journey.
   • Analogy: Imagine the chefs are shouting directions. If Chef A is shouting too loudly and causing the path to wobble, ACE turns Chef A's volume down just a tiny bit for a second, then turns it back up. It dynamically adjusts the "volume knobs" (exponents) so the path stays solid and safe all the way to the finish line.

Why It Matters
The paper shows that without this traffic controller, the AI often fails when trying to combine different experts, especially when you ask for high-quality results (high "guidance").

• In Drug Design: The authors tested this on a task called "scaffold decoration," where they try to build a new drug molecule that fits a specific protein pocket. Without ACE, the AI often produced broken molecules or failed to connect the pieces. With ACE, it successfully built stable, valid molecules that fit the pocket perfectly.
• In Image Generation: They also tested it on creating images with specific objects in specific spots. Even when the path didn't completely collapse, ACE made the images sharper and more accurate by keeping the "ball" on the tightest, most direct path.

The Bottom Line
This paper provides a mathematical safety net. It tells us exactly when combining AI models will break the process and gives a tool (ACE) to fix it on the fly. It turns a risky, heuristic guess into a reliable, guaranteed method for mixing different AI experts to solve complex problems.

Technical Summary

Technical Summary: On the Collapse of Generative Paths

1. Problem Statement: Marginal Path Collapse

The paper identifies a fundamental failure mode in inference-time steering of generative models, specifically when composing heterogeneous experts via ratio-of-densities constructions. While standard steering methods (e.g., Classifier-Free Guidance, Feynman-Kac Correctors) assume that the intermediate density defined by the product of expert marginals remains normalizable, the authors demonstrate that this assumption often fails when experts are trained with mismatched noise schedules or operate on different data dimensions.

This failure is termed Marginal Path Collapse (MPC). It occurs when the intermediate density $h_t(x) = \prod_i q_i(x)^{\gamma_i(t)}$ becomes non-integrable (i.e., the normalizing constant $Z_t = \int h_t(x) dx$ diverges to infinity), even if the initial ($t=0$) and final ($t=1$) endpoints are valid.

• Mechanism: MPC arises from a mismatch in tail contraction rates. If the variances of numerator terms shrink "slower" than those of denominator terms during the diffusion trajectory, the combined density can become explosive (non-normalizable) at intermediate timesteps.
• Consequence: When collapse occurs, the score function of the intended target distribution becomes mathematically undefined. While numerical solvers may continue to run, they effectively simulate an unintended path, leading to terminal distributions that diverge significantly from the target. The authors show this is not an edge case but a prevalent issue in scientific applications like drug design, where heterogeneous experts (e.g., de-novo, conformer, and pocket-conditioned models) must be combined.

2. Methodology

The proposed framework consists of two main components: a diagnostic criterion and a corrective sampling algorithm.

A. Path Existence Criterion (PEC)

The authors derive a rigorous, sharp sufficient condition to certify whether a composed path exists. For a set of experts with noise schedules $\alpha^{(i)}_t$ and exponents $\gamma_i(t)$, the criterion $C(t)$ is defined coordinate-wise as:
$$C_k(t) := \sum_{i: k \in I_i} \frac{\gamma_i(t)}{(\alpha^{(i)}_t)^2}$$
where $I_i$ represents the coordinates acted upon by expert $i$.

• Condition: The path exists (is integrable) for all $t \in [0, 1)$ if and only if $C_k(t) > 0$ for all coordinates $k$.
• Implication: If $C_k(t) < 0$ for any coordinate, the path collapses. The paper proves that for Gaussian-to-compactly-supported interpolants, this condition is both necessary and sufficient.

B. Adaptive Path Correction with Exponents (ACE)

To resolve MPC, the authors introduce ACE, a framework that generalizes Feynman-Kac steering to support time-varying exponents.

1. Exponent Correction: Instead of using fixed exponents $\gamma_i$, ACE dynamically adjusts them to $\tilde{\gamma}_i(t)$ using a "bump function" protocol. This modification preserves the boundary conditions ($\tilde{\gamma}_i(0) = \gamma_i(0)$ and $\tilde{\gamma}_i(1) = \gamma_i(1)$) while ensuring $C_k(t) > 0$ throughout the trajectory.
2. Sampling Dynamics: The correction introduces a time-dependency ($\dot{\gamma}_i(t) \neq 0$) that requires an update to the standard Feynman-Kac sampling dynamics. The authors derive a weighted Stochastic Differential Equation (SDE) where particle weights evolve to account for the changing exponents:
   $$d \log w_t = \left( F(\dots) + \sum_i \dot{\gamma}_i(t) \log \tilde{q}^{(i)}_t(X_t) \right) dt$$
   This allows the sampler to track the corrected probability path unbiasedly.
3. Stabilization: Theoretically, ACE acts as a variance reduction mechanism. By maintaining $C(t)$ positive and bounded away from zero, it controls the quantile radius of intermediate distributions, preventing the "explosive" variance expansion associated with near-collapse regimes.

3. Key Contributions

1. Identification of MPC: The paper formally defines Marginal Path Collapse as a critical failure mode in heterogeneous model composition, explaining why standard constant-exponent steering fails in these settings.
2. Path Existence Criterion (PEC): A sharp, analytically tractable condition ($C(t) > 0$) that diagnoses the validity of a composed generative path based solely on noise schedules and exponents.
3. ACE Framework: A general correction method that guarantees path existence by adaptively adjusting exponents. It extends Feynman-Kac theory to time-varying constraints, providing a theoretical mechanism for path stabilization.
4. Empirical Validation: The method is validated on synthetic benchmarks and complex scientific tasks, demonstrating that it prevents collapse and significantly outperforms existing baselines.

4. Experimental Results

Synthetic Benchmarks

On a 2D checkerboard dataset composed of heterogeneous experts with mismatched schedules:

• Baselines: Standard Heuristics (NR) and Feynman-Kac Correctors (FKC) failed catastrophically when the path existence criterion was violated, producing high distributional error (Wasserstein distance increased by ~4x compared to ACE).
• ACE: Successfully eliminated collapse, recovering the ground truth distribution with significantly lower error.

Flexible-Pose Scaffold Decoration (Drug Design)

This task involves composing three heterogeneous experts: a de-novo (DN) model, a conformer (CONF) model, and a structure-based drug design (SBDD) model.

• Performance: ACE enabled stable composition at high guidance scales ($\omega \ge 1.4$), where baselines (NR, FKC) suffered from path collapse, resulting in fragmented molecules and poor docking scores.
• Metrics: ACE achieved an Optimization Success Rate (OSR) of 0.75 at $\omega=1.4$, significantly outperforming specialized monolithic baselines (e.g., Delete, AutoFragDiff) and FKC (OSR ~0.40).
• Quality: ACE generated chemically valid, connected molecules with superior Vina scores (average -7.10 kcal/mol) and drug-likeness (QED) compared to baselines.

Compositional Image Generation

Even in homogeneous settings where path collapse does not occur, ACE improved attribute success rates by +9.6% over constant-exponent baselines on the COCO-MIG benchmark, demonstrating that time-varying exponents can sharpen intermediate distributions and improve sample quality beyond mere validity repair.

5. Significance and Claims

The paper claims to establish a theoretically grounded foundation for the modular composition of generative models.

• From Heuristic to Guarantee: It transforms ratio-of-densities steering from an unstable heuristic into a provably valid methodology. By providing a diagnostic tool (PEC) and a repair mechanism (ACE), it enables the reliable use of heterogeneous experts in high-stakes scientific domains like drug discovery.
• Generalizability: The framework is not limited to specific architectures but applies to any stochastic interpolant (diffusion or flow matching) where experts can be embedded in a common space.
• Necessity of Heterogeneity: The authors argue that forcing schedule alignment (homogenization) is often suboptimal for scientific tasks; therefore, a method capable of handling inherent heterogeneity (like ACE) is essential for advancing AI in science.

The work concludes that ensuring mathematical validity (normalizability) is a prerequisite for building safe and effective AI tools, particularly when combining specialized models for complex, multi-constraint tasks.