Inter-Trajectory Importance Sampling Improves Diffusion Samplers
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.