[ICML 2026] On the Collapse of Generative Paths: A Criterion and Correction for Diffusion Steering

Authors: Ziseok Lee, Minyeong Hwang, Wooyeol Lee, Sanghyun Jo, Jihyung Ko, Young Bin Park, Jae-Mun Choi, Eunho Yang, Kyungsu Kim

🚧 Work In Progress (Camera-Ready Release) We are currently finalizing the codebase for official release. The complete repository will be fully available by May 31, AoE (Anywhere on Earth).

Abstract

Inference-time steering enables pretrained diffusion/flow models to be adapted to new tasks without retraining. A widely used approach is the ratio-of-densities method, which defines a time-indexed target path by reweighting probability-density trajectories from multiple models. This construction, however, harbors a critical failure mode: Marginal Path Collapse, where intermediate densities become non-normalizable even though endpoints remain valid.

This collapse arises systematically when composing heterogeneous models trained on different noise schedules or datasets. In this work, we provide a novel and complete solution:

1. Path Existence Criterion: We derive a criterion that predicts exactly when collapse occurs based on noise schedules and exponents.
2. Adaptive path Correction with Exponents (ACE): We introduce ACE, which extends Feynman-Kac steering to time-varying exponents to guarantee a valid probability path.

Repository Structure

Once fully released, this repository will be organized into three main sections corresponding to the experiments in our paper:

• 2d_synthetic/: Scripts for the Synthetic 2D benchmark demonstrating the path existence criterion, path collapse, and how ACE repairs collapsed paths.
• molecule/: Implementation of ACE for drug design tasks presented in the paper: scaffold decoration and fragment linking.
• image/: Experiments applying ACE to compositional image generation.

Detailed README.md files and instructions will be provided within each subdirectory.

Installation

We provide a shared environment for all experiments.

Setup

# Clone the repository
git clone https://github.com/ziseoklee/ACE.git
cd ACE

# Create a virtual environment
python3 -m venv ace_env
source ace_env/bin/activate  # On Windows use `ace_env\Scripts\activate`

# Install dependencies
pip install -r requirements.txt

Prerequisites (Toy Experiments) To run the pretraining code which will save checkpoints under /PretrainedToyModels (Takes ~30 mins on a single RTX A6000):

python ace_lib/train_toy_models.py

Citation

If you find this code or our paper useful for your research, please cite:

@inproceedings{
lee2026on,
title={On the Collapse of Generative Paths: A Criterion and Correction for Diffusion Steering},
author={Ziseok Lee and Minyeong Hwang and Wooyeol Lee and Jihyung Ko and Young Bin Park and Jae-Mun Choi and Eunho Yang and Kyungsu Kim},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=emv2qsi3TG}
}