Post-training diffusion models: a unified view from stochastic analysis

Abstract

Post-training has attracted significant attention as a way to adapt pre-trained generative AI models to align with human preferences, structural constraints, or downstream task requirements. A wide range of approaches have been proposed, including reinforcement learning with human feedback (RLHF), stochastic control–based formulations, and classifier-guided or conditioning-based methods for diffusion models. These methods are typically formulated under different assumptions and settings, with limited unifying theoretical understanding.

In this talk, we present a unified framework for post-training diffusion model problems, addressing both soft constraints and hard constraints, through the lens of stochastic analysis. We show that, a broad class of post-training tasks can be rigorously formalized as an add-on mechanism to a fixed pre-trained diffusion model, without requiring full retraining. This perspective clarifies the connections between existing methods and reveals their common mathematical structure.

Building on this framework, we discuss the design of provably efficient post-training algorithms with theoretical guarantees grounded in martingale analysis, Malliavin calculus, and stochastic control. We illustrate the performance of the algorithms on applications including stress testing in financial systems and scheduling problems in queueing networks.

The talk is based on a joint work with Wenpin Tang (Columbia) and Zhengyi Guo (Columbia), as well as a joint work with Yinbin Han (Stanford) and Meisam Razaviyayn (USC).

Bio

Renyuan Xu is currently an assistant professor in the Department of Management Science and Engineering at Stanford University. Her research interests include financial engineering, stochastic analysis, stochastic controls and games, and machine learning theory.