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

AI summary

Overview

Research area: Generative modeling — specifically inference-time steering (composition) of pretrained diffusion and flow-matching models, with applications to structure-based drug design and compositional image generation.

Technical level: Advanced. The paper is built on stochastic interpolants, stochastic differential equations (SDEs), Feynman–Kac correctors, and a Gaussian-to-compactly-supported density analysis. The mathematics is dense, though the core failure mode it identifies is intuitive.

Scope (one sentence): The paper diagnoses a failure mode ("Marginal Path Collapse") that occurs when heterogeneous pretrained experts are composed by ratio-of-densities steering, proves a criterion that certifies when the composed path is mathematically valid, and proposes a method (ACE) that repairs invalid paths by making the steering exponents time-varying.

What This Paper Is About

Diffusion and flow models are often adapted to new tasks at inference time by multiplying or dividing their probability densities by pretrained "experts" with fixed exponents — the same trick behind classifier-free guidance and product-of-experts. The paper shows that when these experts come from different sources (different noise schedules, different data dimensions, different supports), the combined intermediate distribution can become non-normalizable — it "collapses" — even though the starting and ending distributions are both perfectly valid. The goal is to give practitioners a cheap test that predicts this collapse before sampling, and a correction that guarantees a valid path while keeping the intended endpoints.

Key Contributions

1. Identification of Marginal Path Collapse (MPC): The authors name and demonstrate a failure mode where the intermediate density defined by a ratio-of-densities becomes non-integrable, so the score is undefined and the sampler silently transports a different density path than intended. They show it can occur with widely used noise schedules (linear, cosine, DDPM, sigmoid, polynomial combinations).

2. A sharp sufficient Path Existence Criterion (PEC): For experts with Gaussian priors and compactly supported targets, they give a coordinate-wise criterion C_k(t) = Σ_{i: k ∈ I_i} γ_i(t) / (α_t^(i))² > 0 (with C(t) = min_k C_k(t)), where γ_i(t) are exponents and α_t^(i) are noise schedules. The paper reports that C(t) > 0 and C(t) < 0 are the sharpest possible universal conditions, with the boundary case C(t) = 0 handled separately (it can yield either existence or collapse).

3. Adaptive Path Correction with Exponents (ACE): The authors generalize Feynman–Kac steering to time-varying exponents γ_i(t). Given schedules and endpoint exponents, they show there always exist differentiable exponent functions that match the endpoints and keep C(t) > 0 for all t ∈ [0, t_end], and they derive the associated importance-weighted sampling dynamics (which reduces to the constant-exponent Feynman–Kac corrector when γ̇_i(t) = 0).

4. Concentration control and empirical validation: They show the (1−ε)-quantile radius of the intermediate distribution scales roughly as R_t(ε) ≈ O(1/√C(t)), so ACE keeps radii bounded (acting as variance reduction), and they validate ACE on a synthetic checkerboard benchmark, flexible-pose scaffold decoration, and COCO-MIG compositional image generation.

Main Findings

• Collapse can occur with valid endpoints: In a ratio of Gaussians with schedules q_t^(1) = N(0, ((1−t)² + ½t²)I), q_t^(2) = N(0, ((1−t)² + 7t²)I), and q_t^(3) = q_t^(4) = N(0, ((3/2)(1−t)² + t²)I), the endpoints are integrable, yet at t = 0.5 the ratio satisfies h_t(x) ≥ C·exp(+0.01‖x‖²), which is not integrable on R^d.

• The intuitive cause: collapse happens when the variances of numerator terms shrink more slowly than those of denominator terms, producing a temporary imbalance where the combined density explodes at infinity instead of decaying.

• Standard schedules can trigger it: Panel (a) of the schedule figure shows that many heterogeneous three-expert compositions h_t = q_t^(1) q_t^(2) / q_t^(3) built from linear, cosine, DDPM, and related schedules enter a region with C(t) < 0.

• Synthetic checkerboard results: For a 2D checkerboard synthesized from two 1D priors plus a 2D constraint, standard steering (NR, no resampling) is biased; Feynman–Kac correctors (FKC) fail "catastrophically" because C(t) goes negative; ACE keeps C(t) > 0 and recovers the ground truth. The abstract reports ACE reduces error by 4× on this benchmark.

• Quantitative checkerboard metrics (5 seeds, Table 2): NR gives W1 = 0.89 ± 0.02, W2 = 1.18 ± 0.02, MMD (RBF) = 0.092 ± 0.004. FKC gives W1 = 1.37 ± 1.09, W2 = 1.59 ± 1.16, MMD = 0.419 ± 0.579. ACE with B2 = 1.5 and B1 = 10 gives W1 = 0.20 ± 0.04, W2 = 0.29 ± 0.05, MMD = 0.012 ± 0.003 — the best values in the table. ACE at B1 = 0 (0.78 ± 0.15) and at B1 = 20, 30, 40, 50 performs progressively worse than at B1 = 10, so the bump magnitude is not monotone in benefit.

• The tested schedules were deliberately mismatched: the checkerboard experiments used α_t^(1) = cos(π/2 · t), α_t^(2) = DDPM, α_t^(3) = 1 − t, which are known to induce collapse. The paper states the failure is not specific to this choice and reports other schedule combinations (DDPM, sigmoid, linear, cosine, polynomial) show the same patterns.

• Scaffold decoration: the constant-exponent path from the Bayesian decomposition violates the PEC for guidance scales ω ≥ 1.1, causing collapse. ACE remains PEC-valid across all tested ω values.

• Scaffold decoration results vs. steering baselines (Table 3): ACE reaches OSR of 0.71 (ω = 1.1), 0.65 (1.2), 0.68 (1.3), 0.75 (1.4), all PEC-valid; FKC reaches 0.52, 0.43, 0.49, 0.40 (all PEC-invalid); NR reaches 0.46, 0.46, 0.42, 0.47 (all PEC-invalid). ACE's best top-25% Vina score is −8.70 (ω = 1.4) and its average Vina is −7.10 (ω = 1.4) versus −6.77 for the reference molecule; ACE's top-25% Vina at ω = 1.1 is −8.30 versus −8.0 for the reference.

• Scaffold decoration vs. task-specific baselines (Table 4): ACE achieves OSR 0.71 (ω = 1.1) and 0.75 (ω = 1.4); ACE-lite achieves 0.67 and 0.66; AutoFragDiff achieves 0.36 and Delete achieves 0.47 (both marked as requiring a reference ligand pose). The abstract states ACE exceeds specialized monolithic baselines in optimization success rates while maintaining competitive drug-likeness.

• Compositional image generation: on COCO-MIG, ACE improves attribute success rates by +9.6%p over constant-exponent baselines.

• Why time-varying exponents help: the quantile-radius bound R_t(ε) ≈ O(1/√C(t)) means a vanishing C(t) → 0⁺ implies R_t → ∞, producing overly diffuse intermediates that destabilize particle weights. ACE keeps C(t) positive and bounded away from zero, shrinking the effective radius.

• Resampling matters: in the trajectory visualizations, ACE assigns weights such that invalid (out-of-distribution) samples are discarded at each resampling step, whereas the no-resampling heuristic (e.g., CFG) has no corrective mechanism and leaves invalid samples in the batch.

Methodology in Plain English

The setup. Each "expert" is a pretrained model that generates data along its own probability path, defined by its own noise schedule — how quickly noise is added or removed over time. To combine experts, the standard recipe multiplies their densities together and raises each to a fixed power (exponent). The authors instead allow those exponents to vary over time.

Step 1 — Diagnose. Compute C(t) from the noise schedules and exponents. If it goes negative or near zero, fixed-exponent steering is invalid or unstable at that time.

Step 2 — Repair. Add a "bump function" to the exponent of one positively-weighted expert (or a covering subset of them when supports are partial). The bump is zero at the start and end of sampling, so the intended endpoints are preserved, and strictly positive in between, so C(t) stays above a positive floor δ. The specific shape used is a sum of a quadratic term B1·t(1−t) and a linear term B2·min(t, τ(1−t)), with positive constants chosen so the criterion holds at every sampler timestep.

Step 3 — Sample. Because the exponents now change over time, the sampling dynamics acquire an extra weight term Σ_i γ̇_i(t) log q̃_t^(i)(X_t). The authors derive the weighted SDE (and auxiliary dynamics via Itô's formula) that tracks this term, implement it as a particle system with importance weights, and resample whenever the effective sample size drops below a threshold. Choosing the vector field v*_t = Σ_i γ_i(t) ṽ_t^(i) avoids costly divergence computations and was used in the scaffold decoration experiment for numerical stability and efficiency.

The headline intuition: the collapsed path is still numerically solvable — the sampler will happily run — but it is solving the wrong problem. ACE restores the missing probability path before sampling it, which is why the authors describe it as a "repair pipeline" rather than a sampler.

Why This Matters

Impact on research. The paper reframes an implicit assumption in the growing literature on diffusion steering — that the composed ratio path is normalizable at all sampling times — as an assumption that quietly breaks in exactly the settings science cares about. It supplies a cheap checklist (evaluate C(t) on the sampler's timesteps) plus a constructive fix, which should make heterogeneous composition of independently trained models testable rather than trial-and-error. It also connects a binary validity question to a continuous quality question: even a path that technically exists can be too diffuse to sample well.

Real-world applications.

• Structure-based drug design: composing de-novo, conformer, and pocket-conditioned experts to decorate a scaffold while allowing its 3D pose to flex within the binding pocket.
• Linker generation and protein–protein glue design, which the authors identify as sharing the same heterogeneous-conditioning structure (a local constraint on part of the object plus a global constraint on the whole).
• Compositional image generation, where multiple attributes must hold simultaneously (demonstrated on COCO-MIG).
• Any pipeline that mixes independently released checkpoints trained under different schedules or on different dimensionalities, where retraining a single monolithic model would be prohibitively expensive.

Industry relevance. The pitch is modular reuse: instead of paying to retrain a unified model, teams can compose third-party experts they already have. That only works if the composition is reliable, so a validity criterion and an automatic correction have direct engineering value. The paper also reports that forcing schedule alignment across experts can be suboptimal in the scaffold-decoration setting, which matters because it argues heterogeneity is often the right choice rather than a messy accident to be avoided.

Future Directions

• Beyond the boundary case. The paper notes that C(t) = 0 can yield either path existence or collapse, and gives a 1D example showing both outcomes. A general characterization of the boundary remains open.
• Beyond compactly supported targets. The tightest results assume Gaussian priors and compactly supported target distributions; extending the sharp criterion to general non-Gaussian targets (only a Gaussian mixture case is analyzed) is a natural next step.
• Scaling and cost. The paper does not report the computational overhead of tracking auxiliary log-densities and performing resampling relative to the baselines, nor wall-clock or memory comparisons; those trade-offs would determine practical adoption.
• Choice of bump. The checkerboard results are non-monotone in the bump coefficient (B1 = 10 beat B1 = 0, 20, 30, 40, 50), so how to select bump parameters automatically — rather than by hand or by satisfying a lower bound δ — is an obvious open question.
• Broader evaluation. Only two application domains are demonstrated. The full details of the image-generation study are referenced to an appendix, and the provided content is truncated before the complete experimental section, so a fuller picture of behavior across other modalities is not available here.

Target Audience

Researchers and practitioners working on diffusion models, flow matching, and inference-time guidance who need to compose independently trained models — especially in scientific domains like structure-based drug design where experts naturally differ in noise schedule, dimensionality, and support. It is also relevant to applied engineers building modular generative pipelines who want a pre-flight check before trusting a composition, and to theorists interested in probability-path validity and Feynman–Kac formulations. The paper is mathematically demanding, but the diagnosis step (compute C(t), check for negative values) is actionable without following the full derivations.

Note on completeness: the supplied content is truncated mid-sentence in the experiments section. Details referenced only to appendices — such as the complete COCO-MIG results beyond the reported +9.6%p, dataset size figures, full hyperparameter studies, and the full schedule surveys in Tables E.7–E.9 — are not reported in the available text and are therefore not included above.

Authors’ abstract

Inference-time steering adapts pretrained diffusion and flow models to new tasks without retraining, often utilizing ratio-of-densities constructions that reweight time-indexed marginals with fixed exponents. We identify Marginal Path Collapse, a failure mode in which the intermediate density defined by such compositions becomes non-normalizable despite valid endpoints. This collapse can arise when composing heterogeneous experts trained with mismatched noise schedules (and/or negative exponents / partial supports). To address this, we provide (i) a sharp sufficient Path Existence Criterion that characterizes when the composed intermediate densities are mathematically well-defined, and (ii) Adaptive Path Correction with Exponents (ACE), which generalizes Feynman-Kac steering to support time-varying exponents. Our analysis reveals that ACE controls the quantile radius of the intermediate distributions, providing a theoretical mechanism for path stabilization observed in experiments. On flexible-pose scaffold decoration, a drug design task composed of de-novo, conformer, and protein-conditioned experts, ACE prevents collapse and significantly outperforms constant-exponent baselines. Furthermore, ACE improves attribute success rates in compositional image generation, establishing it as a general framework for compositional sampling. Project Page: https://ziseoklee.github.io/projects/ACE/