Adaptive Path Correction with Exponents (ACE)

• The paper introduces ACE, a novel framework that adaptively adjusts exponent schedules to maintain valid probability flows in compositional diffusion tasks.
• ACE employs a weighted Feynman–Kac sampler and bump function corrections to robustly counteract Marginal Path Collapse without needing model retraining.
• Empirical evaluations show significant error reduction and improved sample quality in synthetic and molecular generation benchmarks using ACE.

Adaptive Path Correction with Exponents (ACE) is a principled framework for enabling robust inference-time steering of pretrained diffusion and flow models via the adaptive combination of multiple model trajectories, each modulated by time-varying exponents. ACE addresses the fundamental challenge of Marginal Path Collapse in ratio-of-densities steering, ensuring the existence of valid intermediate probability paths for tasks including molecular generation and compositional image synthesis. It operates without retraining and leverages only black-box access to pretrained model scores or velocities, transforming ratio-based composition from an unstable heuristic into a reliable algorithmic tool (Lee et al., 11 Dec 2025).

1. Ratio-of-Densities Steering and Marginal Path Collapse

Ratio-of-densities steering forms the backbone of various compositional generative modeling tasks. Given $n$ pretrained diffusion models $\{q^{(i)}_t\}_{t\in[0,1]}$, potentially on different coordinates or subspaces, one constructs an unnormalized trajectory

$$h_t(x) = \prod_{i=1}^n \bigl(\tilde q^{(i)}_t(x)\bigr)^{\gamma_i(t)},$$

with $\tilde q^{(i)}_t(x) = q^{(i)}_t(\pi_i(x))$ for a suitable projection $\pi_i$ into each expert’s domain, and exponent schedule $\gamma_i(t)\in\mathbb{R}$ that may be positive (numerator) or negative (denominator). If $h_t\in L^1$, then $p^*_t(x) = h_t(x)/Z_t$ with $Z_t = \int h_t(x)\,dx$ defines a valid probability flow between start and target $t=0,1$.

Marginal Path Collapse occurs when, despite $h_0,h_1\in L^1$, there exists $t^*\in(0,1)$ with $\int h_{t^*}(x)\,dx=\infty$, making $p^*_{t^*}$ non-normalizable. For example, in a one-dimensional Gaussian setting, the ratio construction can be integrable at endpoints and divergent at intermediate times due to negative effective variance, a phenomenon systematically triggered when numerator variances contract more slowly than denominator variances.

2. Path-Existence Criterion

A closed-form criterion predicts precisely when Marginal Path Collapse will occur in composition tasks involving compactly supported densities. Assuming each expert $i$ follows a noise schedule $\alpha^{(i)}_t,\beta^{(i)}_t$ and operates on coordinates $I_i$, define

$$C_k(t) = \sum_{i: k\in I_i} \frac{\gamma_i(t)}{(\alpha^{(i)}_t)^2}, \qquad C(t) = \min_k C_k(t),$$

where $k$ indexes coordinates. The integrability of the path is guaranteed by

$$\Bigl\{ h_t \Bigr\} \text{ is integrable for all } t\in [0,1] \quad \Longleftrightarrow \quad C_k(t)>0\ \forall\, k,\, t\in [0,1).$$

Violation of this condition for any $k$ and $t$ signals an impending Marginal Path Collapse. This criterion holds under the requirement that the final densities $q^{(i)}_1$ possess compact support.

3. Construction of Adaptive Path Correction with Exponents (ACE)

ACE introduces a two-pronged solution: (a) adaptive adjustment of exponents to guarantee $C_k(t)>0$ everywhere, and (b) a weighted SDE/ODE sampler for unbiased path tracking.

3.1. Adaptive Exponents Using Bump Functions

Let $\gamma_i(0),\gamma_i(1)$ be the desired endpoint exponents and assume $C_k(0)>0$ and $\lim_{t\to 1^-}C_k(t)>0$. There exists an index $j$ and a bump function $b(t) = t(1-t)$ such that modifying the exponent trajectory via

$$\tilde\gamma_j(t) = \gamma_j(t) + B\, b(t), \quad \tilde\gamma_i(t) = \gamma_i(t)\ (i\ne j)$$

for some analytically determined $B>0$ ensures $C_k(t)>0$ for all $k$ and $t\in [0,1)$. This approach suffices to correct the path without altering the boundary behavior, as the bump vanishes at $t=0$ and $t=1$.

3.2. Weighted Feynman–Kac Sampler

The corrected exponent schedule $\{\tilde\gamma_i(t)\}$ enables unbiased tracking of the normalized path via a stochastic differential equation: $$\begin{aligned} dX_t &= \Bigl(v^*_t(X_t) + \frac{\sigma_t^2}{2} s^*_t(X_t)\Bigr)\,dt + \sigma_t\,dW_t, \ d\log w_t &= \Bigl[ \nabla\cdot v^*_t(X_t) + \sum_{i=1}^n \dot{\tilde\gamma}_i(t)\,\log \tilde q^{(i)}_t(X_t) + \sum_{i=1}^n \tilde\gamma_i(t)\, D^{(i)}_t(X_t) \Bigr] dt \end{aligned}$$ where $$s^*_t(x) = \sum_{i=1}^n \tilde\gamma_i(t)\, \tilde s^{(i)}_t(x)$$ collects expert scores, $v^*_t(x)$ is user-chosen drift, and $D^{(i)}_t(x)$ terms encapsulate velocity divergence and drift mismatch corrections.

3.3. Collapse Prevention

By enforcing $C_k(t)>0$ via the bump-corrected exponents, the integrand $h_t$ remains normalizable for all $t$, preserving unbiased inference. The score $\nabla\log p^*_t$ is well-defined, and the Feynman-Kac framework yields stable tracking.

4. Algorithmic Implementation

The ACE method is operationalized as follows (see Alg. 1 in (Lee et al., 11 Dec 2025)):

1. Initialization: Draw $N$ initial particles $X^j_0 \sim p^*_0$, set log-weights $\log w^j_0 = 0$.
2. For each timestep $\Delta t$:

a. Compute $s^*_t(X^j_t)$ and drift $\mu_t = v^*_\phi + \frac12 \sigma_t^2 s^*_t$. b. For each expert $i$, compute $D^{(i)}_t(X^j_t)$. c. Propagate $$X^j_{t+\Delta t} = X^j_t + \mu_t \Delta t + \sigma_t \sqrt{\Delta t} \xi^j$$. d. Update $\log w^j$ by integrating the SDE corrections. e. If the effective sample size ${\rm ESS}<\tau N$, perform weighted resampling.

1. Output: Approximates $p^*_1$ via weighted samples.

This sampler accommodates both SDE and ODE settings (by setting $\sigma_t=0$, omitting Itô corrections) and can be integrated with Sequential Monte Carlo resampling strategies.

5. Empirical Evaluation

5.1. Synthetic 2D Checkerboard

In a synthetic benchmark, the target $p(x,y\mid A,B)\propto p(x\mid A)p(x,y\mid B)/p(x)$ is realized over $[-4,4]^2$ using three heterogeneous diffusion experts. Collapse frequency is substantial: in $100$ random schedules, Marginal Path Collapse occurs in $41\%$ of cases at guidance $w=1.0$, rising to $80\%$ at $w=15$. ACE ($B=30$) eliminates collapse, achieving $W_1\approx0.28$, $W_2\approx0.40$, and $\mathrm{MMD}\approx0.027$, representing a $>5\times$ reduction in error over constant-exponent baselines (FKC: $W_1\approx2.13$).

Method	$W_1$	$W_2$	MMD
NR	0.78	1.07	0.068
FKC	2.13	—	—
ACE ($B=30$)	0.28	0.40	0.027

5.2. Flexible-Pose Scaffold Decoration

In molecular scaffold decoration, three models are composed: DN (de-novo, trained on ZINC), CONF (topology-conditioned conformers), and SBDD (pocket-conditioned generation). Collapse for constant-exponent steering limits attainable guidance ($\omega>1.1$). Applying ACE with $B=30$ to the SBDD exponent restores valid paths for $\omega$ up to $1.4$. On CrossDock-Weak and CrossDock-SBDD datasets, ACE attains 96–100% chemical validity, optimal structure recovery, and substantially improved docking scores ($-5.44$ for ACE at $\omega=1.2$ vs. $-3.40$ for FKC), surpassing specialized task models such as Delete, DiffDec, and AutoFragDiff.

6. Practical Guidelines and Limitations

• Always verify the path-existence criterion $C(t)>0$ on the discrete sampling grid before steering.
• If $C(t)<0$ for some $t$, introduce a bump $B\,t(1-t)$ in positive exponents; $B=30$ suffices for molecular tasks, and increasing $B$ further broadens safety margins, though may amplify model errors.
• For valid paths, moderate time variation in exponents can sharpen distributions and improve sample quality (see compositional generation on COCO-MIG, Prop. 4.1).
• ACE requires no retraining and is agnostic to the specific network architectures, requiring only black-box score/velocity access.
• Current limitations include: error accumulation under many expert compositions, extension to discrete/hybrid spaces, cost of SDE/ODE distillation for faster inference, and application beyond stochastic interpolants to broader transport tasks.

7. Significance and Research Frontiers

ACE furnishes a theoretically complete framework for addressing Marginal Path Collapse in diffusion steering, enabling previously unattainable degrees of compositionality and guidance strength across synthetic, molecular, and image domains. This framework systematically transforms ratio-of-densities steering into a stable, general methodology, promoting broader adoption in controllable generation. Open research questions include managing expert composition error, discrete-variable extensions, efficient inference distillation, and application to broader generative transport problems (Lee et al., 11 Dec 2025).