Handling unchanged-noise-level corner cases in Diffusion Forcing sampling

Determine the correct procedure for handling the corner case in the Diffusion Forcing sampling algorithm when the scheduling matrix K prescribes that a token’s noise level k_t remains unchanged during a sampling step. Specifically, decide between (i) copying the existing value without updating or (ii) performing a backward diffusion followed by a forward diffusion to resample the token at the same noise level.

Background

Diffusion Forcing sampling proceeds by denoising a sequence according to a scheduling matrix K that sets per-token noise levels across steps. Because K is fully flexible, a corner case can arise in which the prescribed noise level for a given token k_t remains the same between successive sampling steps.

In this situation, the authors describe two plausible update strategies: copying the token’s current value without change, or resampling by applying a backward diffusion step followed by a forward diffusion step to return to the same noise level. The paper notes this design decision is an open question and, in their implementation, they choose the resampling option to support Monte Carlo Tree Guidance.

References

In our sampling algorithm, due to the flexibility of the scheduling matrix \mathcal{K}, there are corner cases when k_t is required to stay at its same noise level during a sampling step. The core question of this corner case is whether we should updatek_t at all. One option is just copying over the old value. The other option is to run a backward diffusion followed by a forward diffusion back to its old noise level to resample under the diffusion process. While we conclude this can be an open question, we prefer the later approach, resampling, and use it in Monte Carlo Guidance to generate multiple samples.

— Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion  (2407.01392 - Chen et al., 2024) in Appendix, Implementation Details of Sampling with Guidance, Corner case of sampling noise