Researchers propose the Time-Reparameterized Cumulative Intensity Extrapolation (TR-CIE) sampler to improve sampling quality in discrete flow matching when function evaluations are restricted. The method combines schedule-based time reparameterization with a cumulative-intensity extrapolation updating rule to mitigate stiffness and improve approximation accuracy.

  • TR-CIE rescales the time grid according to the noise schedule, absorbing schedule-dependent growth terms under standard factorized DFM rate parameterizations.
  • The sampler uses cached model outputs from previous steps as a history term to approximate stepwise cumulative intensities on non-uniform time grids.
  • It requires only one function evaluation per step and introduces no additional model evaluations compared to standard τ-leaping.
  • Theoretical analysis bounds the local approximation error of cumulative intensities and establishes convergence results.
  • Experiments on synthetic tasks, text generation, and text-to-image benchmarks demonstrate improved sampling quality under limited NFE.

This approach addresses the lack of efficient sampling methods for discrete flow matching by providing a theoretically grounded technique that enhances performance without increasing computational overhead.