Inter-Trajectory Importance Sampling Improves Diffusion Samplers

Lukas Gruber, Sandeep Suresh Cranganore, Sepp Hochreiter, Sebastian Lehner

Keywords: Amortized sampling, Diffusion samplers, Mode collapse, Importance sampling, Generative modeling

TL;DR: We introduce ITIS, a multi-sample objective that couples diffusion trajectories via path-measure importance weighting to suppress mode collapse and achieve state-of-the-art performance with significantly fewer function evaluations.

Abstract: Sampling from unnormalized distributions is a fundamental challenge in computational sciences. Diffusion samplers have emerged as a promising approach to this problem, yet mode collapse remains one of their most persistent and practically limiting failure modes, arising from the mode-seeking nature of the reverse KL divergence. We introduce Inter-Trajectory Importance Sampling (ITIS), a multi-sample training objective that couples the loss contributions of concurrently generated trajectories through importance weighting. By reweighting the loss contribution of each trajectory step using the path measures of concurrently generated trajectories, ITIS renders the loss sensitive to inter-trajectory diversity while leaving the underlying reverse KL objective intact. We provide theoretical grounding via an exit-time analysis characterizing ITIS' suppression of mode collapse, and demonstrate that ITIS serves as a drop-in replacement loss for established diffusion samplers with significant and consistent performance gains across standard sampling benchmarks. Notably, these improvements can be achieved using considerably fewer diffusion steps and energy function evaluations.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 227