A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matching

2026-06-23

AI Summary

This work addresses the challenge in discrete flow matching (DFM) where existing sampling methods struggle to balance efficiency and generation quality under a limited number of function evaluations (NFE). The authors propose TR-CIE, a novel sampler that uniquely integrates time reparameterization with cumulative intensity extrapolation. Time reparameterization alleviates terminal stiffness through adaptive noise scheduling, while cumulative intensity extrapolation enhances sampling accuracy on non-uniform time grids by reusing historical model outputs—without requiring additional model evaluations. TR-CIE significantly improves generation quality at low NFE and is supported by theoretical guarantees on local error bounds and convergence. Empirical results demonstrate its consistent superiority over state-of-the-art samplers across synthetic data, text generation, and text-to-image tasks.

Abstract

Discrete flow matching (DFM) provides a principled framework for generative modeling on discrete state spaces via continuous-time Markov chain dynamics. In practice, sampling for DFM commonly employs discretizations such as $τ$-leaping, yet efficient sampling methods under a limited number of function evaluations (NFE) remain less studied. To address this gap, we propose the Time-Reparameterized Cumulative Intensity Extrapolation (TR-CIE) sampler, which aims to improve sampling quality when function evaluations are restricted. TR-CIE consists of two components. First, a schedule-based time reparameterization rescales the time grid according to the noise schedule. Under standard factorized DFM rate parameterizations, this transformation of variables absorbs the schedule-dependent growth term and mitigates stiffness near the terminal sampling stage. Second, we introduce a cumulative-intensity extrapolation updating rule. By reusing cached model outputs from the previous step as a history term, this improves the approximation of stepwise cumulative intensities on the resulting non-uniform time grid. We provide a theoretical analysis that bounds the local approximation error of cumulative intensities and establishes convergence results. The resulting sampler requires one NFE per step and introduces no additional model evaluations compared to the standard $τ$-leaping sampler. Extensive experiments on synthetic tasks, text generation, and text-to-image benchmarks demonstrate that our method improves sampling quality under limited NFE.

Problem

• Discrete Flow Matching
• sampling efficiency
• limited function evaluations
• continuous-time Markov chain
• generative modeling

Innovation

• Discrete Flow Matching
• Time Reparameterization
• Cumulative Intensity Extrapolation
• Stiffness Mitigation
• Efficient Sampling