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

Feiyang Fu; Hehe Fan

arxiv: 2606.24140 · v2 · pith:H3RH3TILnew · submitted 2026-06-23 · 💻 cs.LG

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

Feiyang Fu , Hehe Fan This is my paper

Pith reviewed 2026-06-26 00:41 UTC · model grok-4.3

classification 💻 cs.LG

keywords discrete flow matchingtime reparameterizationcumulative intensity extrapolationtau-leapingsampling efficiencygenerative modelingdiscrete state spacestext-to-image generation

0 comments

The pith

The TR-CIE sampler improves sampling quality in discrete flow matching under limited function evaluations by reparameterizing time and extrapolating cumulative intensities without extra model calls.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Discrete flow matching models data on discrete spaces through continuous-time Markov chain dynamics, yet standard tau-leaping discretizations lose quality when the number of function evaluations is kept low. The paper introduces the TR-CIE sampler whose first component rescales the time grid according to the noise schedule so that the growth term is absorbed under factorized rate parameterizations. Its second component reuses the prior model output as a history term to extrapolate the cumulative intensity on the resulting non-uniform grid. A local error bound and convergence result are derived, and the method keeps exactly one network evaluation per step. Experiments on synthetic tasks, text generation, and text-to-image generation show higher sample quality at the same computational budget.

Core claim

Under standard factorized DFM rate parameterizations, schedule-based time reparameterization absorbs the schedule-dependent growth term and mitigates stiffness near the terminal stage, while a cumulative-intensity extrapolation rule that reuses cached outputs improves stepwise approximations on the non-uniform grid; the resulting sampler requires one NFE per step, introduces no additional model evaluations relative to tau-leaping, and is accompanied by a bound on local approximation error together with convergence results.

What carries the argument

The schedule-based time reparameterization combined with the cumulative-intensity extrapolation updating rule that reuses previous model outputs as a history term.

If this is right

The sampler requires exactly one neural-network evaluation per step.
A theoretical bound is obtained on the local approximation error of the cumulative intensities.
Convergence of the overall sampling process is established.
Higher sample quality is reported on synthetic, text, and text-to-image tasks when the number of function evaluations is restricted.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The same reparameterization-plus-extrapolation pattern could be tested on other discrete diffusion or flow models that share factorized rate forms.
The stiffness reduction near the end of sampling may allow practitioners to drop the total step count further while preserving quality.
Combining the extrapolation history term with existing higher-order integrators could yield additional efficiency gains.

Load-bearing premise

The method assumes standard factorized DFM rate parameterizations so that the time-reparameterization transformation absorbs the schedule-dependent growth term and mitigates stiffness.

What would settle it

On the text-to-image or text-generation benchmarks, run TR-CIE and the baseline tau-leaping sampler at identical limited NFE budgets and observe no gain or a loss in sample quality metrics.

Figures

Figures reproduced from arXiv: 2606.24140 by Feiyang Fu, Hehe Fan.

**Figure 1.** Figure 1: Comparison on the synthetic countdown task. We compare our TR-CIE sampler with other samplers across various NFE. The error rate (lower is better) indicates that our method generally achieves better performance. We also observe that TRCIE sampler significantly reduces errors especially in the low-NFE regime. 6. Conclusion This paper addresses the challenge of efficient sampling in DFM. Standard approximat… view at source ↗

**Figure 2.** Figure 2: Comparison in terms of FID on the DFM backbone across four sampling schedules. B.3. Runtime comparison We report the runtime comparison under the same hardware and batch-size settings. The experiments are performed on the uniform DFM backbone using a single NVIDIA RTX 4090 GPU with a batch of 64 samples and report the average generation time per sample at NFE = 8, 16, 32, 64, 128 in [PITH_FULL_IMAGE:figur… view at source ↗

**Figure 3.** Figure 3: Comparison of temporal integration and state drift proxies. (Bottom) In the physical t domain, the integration proxy (blue) diverges exponentially as t → 1 due to the schedule-induced stiffness. (Top) In the reparameterized τ domain (mapped back to physical time t for direct comparison in the x-axis), the integration proxy remains bounded and decays near the boundary. The ratio plots (right) show that inte… view at source ↗

**Figure 4.** Figure 4: Qualitative comparison at NFE=8. We observe that TR-CIE better preserves prompt semantics and reduces low-step artifacts, yielding a better sampling quality. (zoom in for best view) 22 [PITH_FULL_IMAGE:figures/full_fig_p022_4.png] view at source ↗

read the original 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 $\tau$-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 $\tau$-leaping sampler. Extensive experiments on synthetic tasks, text generation, and text-to-image benchmarks demonstrate that our method improves sampling quality under limited NFE.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

TR-CIE adds a reparameterized grid plus extrapolation to tau-leaping in DFM but its main gain depends on factorized rates.

read the letter

The main thing to know is that this paper introduces the TR-CIE sampler for discrete flow matching. It uses a schedule-based time reparameterization to create a non-uniform grid that is meant to absorb the growth term and cut stiffness, then adds a cumulative-intensity extrapolation step that reuses prior model outputs to approximate the integrals better. The sampler still costs one NFE per step and claims a local error bound plus convergence.

It does a couple of things cleanly. The design stays inside the existing DFM setup and does not increase model evaluations, which matters for practical sampling. The experiments run on synthetic tasks, text generation, and text-to-image benchmarks, and the theory section at least states bounds on the approximation error. Those are the right ingredients for a sampling-methods paper.

The soft spot is the modeling assumption. The reparameterization is described as absorbing the schedule-dependent term only under standard factorized DFM rate parameterizations. The paper does not discuss how common non-factorized rates are in current DFM work or what happens to the stiffness and the error bound when the absorption does not occur. If that case is frequent, the claimed quality gains under limited NFE become narrower. The abstract also gives no equations, so the actual tightness of the bound and any extra conditions in the proof remain to be checked in the full text.

This is for people already working on discrete flow matching or similar continuous-time discrete generative models. A reader who needs faster sampling on text or image tasks could get a usable method if the assumption holds. The combination of theory and targeted experiments is enough to send it to a serious referee, though the review should focus on the rate-factorization point and the derivations.

Referee Report

1 major / 1 minor

Summary. The manuscript proposes the Time-Reparameterized Cumulative Intensity Extrapolation (TR-CIE) sampler for discrete flow matching (DFM). It consists of a schedule-based time reparameterization that, under standard factorized DFM rate parameterizations, absorbs schedule-dependent growth terms and mitigates stiffness, together with a cumulative-intensity extrapolation rule that reuses prior model outputs to approximate stepwise cumulative intensities on the resulting non-uniform grid. The sampler is claimed to require exactly one NFE per step with no extra model evaluations relative to τ-leaping, supported by a local approximation-error bound and convergence results. Experiments on synthetic tasks, text generation, and text-to-image benchmarks are reported to show improved sample quality under limited NFE.

Significance. If the local-error bound and convergence results are valid and the empirical gains prove robust across rate parameterizations, the method would supply a low-overhead improvement for NFE-constrained sampling in discrete generative models. The explicit reuse of cached outputs without additional evaluations is a concrete practical strength.

major comments (1)

[Abstract] Abstract (paragraph describing the two components of TR-CIE): the central claim that the time-reparameterization 'absorbs the schedule-dependent growth term and mitigates stiffness' is explicitly conditioned on 'standard factorized DFM rate parameterizations.' No analysis is supplied of how frequently non-factorized rates appear in DFM practice, nor of whether the extrapolation step can compensate when the absorption fails; if the assumption does not hold, both the practical improvement under limited NFE and the applicability of the local-error bound are undermined.

minor comments (1)

The abstract states that theoretical bounds and convergence results are provided, yet the manuscript should ensure that the proof sketches or key steps are presented with sufficient detail to allow verification of the local approximation error bound.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for the constructive feedback on the scope of our assumptions. We address the single major comment below.

read point-by-point responses

Referee: [Abstract] Abstract (paragraph describing the two components of TR-CIE): the central claim that the time-reparameterization 'absorbs the schedule-dependent growth term and mitigates stiffness' is explicitly conditioned on 'standard factorized DFM rate parameterizations.' No analysis is supplied of how frequently non-factorized rates appear in DFM practice, nor of whether the extrapolation step can compensate when the absorption fails; if the assumption does not hold, both the practical improvement under limited NFE and the applicability of the local-error bound are undermined.

Authors: We agree that the absorption property and the associated local-error bound are derived under the assumption of standard factorized rate parameterizations, which is stated explicitly in the abstract and method sections. In the DFM literature this factorization (typically across independent dimensions or tokens) is the dominant practical choice because it yields tractable rate matrices and enables the product-form likelihoods used by nearly all published DFM models. The cumulative-intensity extrapolation rule itself does not rely on factorization and remains applicable on any non-uniform grid; however, when the reparameterization fails to absorb the growth term the stiffness-mitigation benefit is reduced. Because the manuscript supplies neither a survey of non-factorized usage nor an explicit statement of this limitation, we will add a short clarifying paragraph in the revised introduction and a footnote in the abstract. This is a partial revision that leaves the core claims and experiments unchanged while making the scope of the guarantees explicit. revision: partial

Circularity Check

0 steps flagged

No significant circularity; derivation builds on external DFM framework

full rationale

The provided abstract and description present TR-CIE as an extension of the existing discrete flow matching (DFM) framework, with time reparameterization applied under explicitly stated 'standard factorized DFM rate parameterizations' and a new extrapolation rule supported by claimed theoretical bounds on local error. No equations, self-citations, or fitted quantities are shown that reduce the sampler's claimed benefits, error bounds, or convergence results to inputs defined by the method itself. The one-NFE-per-step property is stated as a direct consequence of reusing cached outputs without additional model calls, independent of any self-referential construction. This matches the default expectation of a non-circular paper.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

Only the abstract is available; the ledger is therefore minimal and reflects the domain assumptions stated there.

axioms (1)

domain assumption Standard factorized DFM rate parameterizations allow the time-reparameterization to absorb the schedule-dependent growth term.
Explicitly invoked in the abstract as the setting under which the first component of TR-CIE works.

pith-pipeline@v0.9.1-grok · 5776 in / 1266 out tokens · 19456 ms · 2026-06-26T00:41:35.681332+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

14 extracted references · 5 linked inside Pith

[1]

Anderson, D. F. and Mattingly, J. C. A weak trapezoidal method for a class of stochastic differential equations. arXiv preprint arXiv:0906.3475,

Pith/arXiv arXiv
[2]

N., Ryder, J

Billera, L., Nordlinder, H. N., Ryder, J. C., Oresten, A., St˚almarck, A., Bj ¨ork, T. M., and Murrell, B. Branch- ing flows: Discrete, continuous, and manifold flow matching with splits and deletions.arXiv preprint arXiv:2511.09465,

Pith/arXiv arXiv
[3]

Multi- objective-guided discrete flow matching for control- lable biological sequence design.arXiv preprint arXiv:2505.07086,

Chen, T., Zhang, Y ., Tang, S., and Chatterjee, P. Multi- objective-guided discrete flow matching for control- lable biological sequence design.arXiv preprint arXiv:2505.07086,

arXiv
[4]

alpha-flow: A unified framework for continuous-state discrete flow matching models.arXiv preprint arXiv:2504.10283,

Cheng, C., Li, J., Fan, J., and Liu, G. alpha-flow: A unified framework for continuous-state discrete flow matching models.arXiv preprint arXiv:2504.10283,

arXiv
[5]

The llama 3 herd of models.arXiv preprint arXiv:2407.21783,

Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., et al. The llama 3 herd of models.arXiv preprint arXiv:2407.21783,

Pith/arXiv arXiv
[6]

T., Lopez-Paz, D., Ben-Hamu, H., and Gat, I

Lipman, Y ., Havasi, M., Holderrieth, P., Shaul, N., Le, M., Karrer, B., Chen, R. T., Lopez-Paz, D., Ben-Hamu, H., and Gat, I. Flow matching guide and code.arXiv preprint arXiv:2412.06264,

Pith/arXiv arXiv
[7]

Next-omni: Towards any- to-any omnimodal foundation models with discrete flow matching.arXiv preprint arXiv:2510.13721,

Luo, R., Xia, X., Wang, L., Chen, L., Shan, R., Luo, J., Yang, M., and Chua, T.-S. Next-omni: Towards any- to-any omnimodal foundation models with discrete flow matching.arXiv preprint arXiv:2510.13721,

arXiv
[8]

Drax: Speech recognition with discrete flow matching.arXiv preprint arXiv:2510.04162,

Navon, A., Shamsian, A., Glazer, N., Segal-Feldman, Y ., Hetz, G., Keshet, J., and Fetaya, E. Drax: Speech recognition with discrete flow matching.arXiv preprint arXiv:2510.04162,

arXiv
[9]

Y .-C., Wu, S., Song, Z., Reneau, A., and Liu, H

Su, M., Lu, M., Hu, J. Y .-C., Wu, S., Song, Z., Reneau, A., and Liu, H. A theoretical analysis of discrete flow match- ing generative models.arXiv preprint arXiv:2509.22623,

arXiv
[10]

Error analysis of discrete flow with generator matching.arXiv preprint arXiv:2509.21906,

Wan, Z., Ouyang, Y ., Yao, Q., Xie, L., Fang, F., Zha, H., and Cheng, G. Error analysis of discrete flow with generator matching.arXiv preprint arXiv:2509.21906,

Pith/arXiv arXiv
[11]

Mdns: Masked diffusion neural sampler via stochastic optimal control

Zhu, Y ., Guo, W., Choi, J., Liu, G.-H., Chen, Y ., and Tao, M. Mdns: Masked diffusion neural sampler via stochastic optimal control. InConference on Neural Information Processing Systems, 2025a. Zhu, Y ., Wang, X., Lathuili`ere, S., and Kalogeiton, V . Di [m] o: Distilling masked diffusion models into one-step generator. InInternational Conference on Com...

2024
[12]

13 A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matching Table 9.Runtime per sample across various NFEs. Sampler NFE=8 NFE=16 NFE=32 NFE=64 NFE=128 Eulerτ-leaping 123.8 ms 242.1 ms 490.7 ms 988.9 ms 1978.6 ms Tweedieτ-leaping 74.2 ms 158.2 ms 313.1 ms 615.8 ms 1284.0 ms θ-Trapezoidal 97.3 ms 194.1 ms 388.0 ms 775.6 m...

1978
[13]

Pathwise KL Bound We establish a standard change-of-measure identity for jump processes

C.2. Pathwise KL Bound We establish a standard change-of-measure identity for jump processes. This result allows us to upper bound the terminal KL divergence using an integral functional of the intensities and to facilitate the term-by-term error decomposition. 14 A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matching...

1999
[14]

Explicit bounds for these terms are provided in Appendix C.6

N−1yields the global decomposition: KL(pτN ∥qτN )≤ N−1X n=0 En ≤ Eint +E freeze +E var,(36) where the three components correspond to the three non-zero terms derived above: •E int =P n E R τn+1 τn dBreg(¯un,ˆµτ)dτ; •E freeze =P n CΦE R τn+1 τn |˜uτ(Xτ −)−˜uτ(Xτn)|dτ; •E var =P n CΦE R τn+1 τn |˜uτ(Xτn)−¯un|dτ. Explicit bounds for these terms are provided ...

2001

[1] [1]

Anderson, D. F. and Mattingly, J. C. A weak trapezoidal method for a class of stochastic differential equations. arXiv preprint arXiv:0906.3475,

Pith/arXiv arXiv

[2] [2]

N., Ryder, J

Billera, L., Nordlinder, H. N., Ryder, J. C., Oresten, A., St˚almarck, A., Bj ¨ork, T. M., and Murrell, B. Branch- ing flows: Discrete, continuous, and manifold flow matching with splits and deletions.arXiv preprint arXiv:2511.09465,

Pith/arXiv arXiv

[3] [3]

Multi- objective-guided discrete flow matching for control- lable biological sequence design.arXiv preprint arXiv:2505.07086,

Chen, T., Zhang, Y ., Tang, S., and Chatterjee, P. Multi- objective-guided discrete flow matching for control- lable biological sequence design.arXiv preprint arXiv:2505.07086,

arXiv

[4] [4]

alpha-flow: A unified framework for continuous-state discrete flow matching models.arXiv preprint arXiv:2504.10283,

Cheng, C., Li, J., Fan, J., and Liu, G. alpha-flow: A unified framework for continuous-state discrete flow matching models.arXiv preprint arXiv:2504.10283,

arXiv

[5] [5]

The llama 3 herd of models.arXiv preprint arXiv:2407.21783,

Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., et al. The llama 3 herd of models.arXiv preprint arXiv:2407.21783,

Pith/arXiv arXiv

[6] [6]

T., Lopez-Paz, D., Ben-Hamu, H., and Gat, I

Lipman, Y ., Havasi, M., Holderrieth, P., Shaul, N., Le, M., Karrer, B., Chen, R. T., Lopez-Paz, D., Ben-Hamu, H., and Gat, I. Flow matching guide and code.arXiv preprint arXiv:2412.06264,

Pith/arXiv arXiv

[7] [7]

Next-omni: Towards any- to-any omnimodal foundation models with discrete flow matching.arXiv preprint arXiv:2510.13721,

Luo, R., Xia, X., Wang, L., Chen, L., Shan, R., Luo, J., Yang, M., and Chua, T.-S. Next-omni: Towards any- to-any omnimodal foundation models with discrete flow matching.arXiv preprint arXiv:2510.13721,

arXiv

[8] [8]

Drax: Speech recognition with discrete flow matching.arXiv preprint arXiv:2510.04162,

Navon, A., Shamsian, A., Glazer, N., Segal-Feldman, Y ., Hetz, G., Keshet, J., and Fetaya, E. Drax: Speech recognition with discrete flow matching.arXiv preprint arXiv:2510.04162,

arXiv

[9] [9]

Y .-C., Wu, S., Song, Z., Reneau, A., and Liu, H

Su, M., Lu, M., Hu, J. Y .-C., Wu, S., Song, Z., Reneau, A., and Liu, H. A theoretical analysis of discrete flow match- ing generative models.arXiv preprint arXiv:2509.22623,

arXiv

[10] [10]

Error analysis of discrete flow with generator matching.arXiv preprint arXiv:2509.21906,

Wan, Z., Ouyang, Y ., Yao, Q., Xie, L., Fang, F., Zha, H., and Cheng, G. Error analysis of discrete flow with generator matching.arXiv preprint arXiv:2509.21906,

Pith/arXiv arXiv

[11] [11]

Mdns: Masked diffusion neural sampler via stochastic optimal control

Zhu, Y ., Guo, W., Choi, J., Liu, G.-H., Chen, Y ., and Tao, M. Mdns: Masked diffusion neural sampler via stochastic optimal control. InConference on Neural Information Processing Systems, 2025a. Zhu, Y ., Wang, X., Lathuili`ere, S., and Kalogeiton, V . Di [m] o: Distilling masked diffusion models into one-step generator. InInternational Conference on Com...

2024

[12] [12]

13 A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matching Table 9.Runtime per sample across various NFEs. Sampler NFE=8 NFE=16 NFE=32 NFE=64 NFE=128 Eulerτ-leaping 123.8 ms 242.1 ms 490.7 ms 988.9 ms 1978.6 ms Tweedieτ-leaping 74.2 ms 158.2 ms 313.1 ms 615.8 ms 1284.0 ms θ-Trapezoidal 97.3 ms 194.1 ms 388.0 ms 775.6 m...

1978

[13] [13]

Pathwise KL Bound We establish a standard change-of-measure identity for jump processes

C.2. Pathwise KL Bound We establish a standard change-of-measure identity for jump processes. This result allows us to upper bound the terminal KL divergence using an integral functional of the intensities and to facilitate the term-by-term error decomposition. 14 A Time-Reparameterized Cumulative Intensity Extrapolation Sampler for Discrete Flow Matching...

1999

[14] [14]

Explicit bounds for these terms are provided in Appendix C.6

N−1yields the global decomposition: KL(pτN ∥qτN )≤ N−1X n=0 En ≤ Eint +E freeze +E var,(36) where the three components correspond to the three non-zero terms derived above: •E int =P n E R τn+1 τn dBreg(¯un,ˆµτ)dτ; •E freeze =P n CΦE R τn+1 τn |˜uτ(Xτ −)−˜uτ(Xτn)|dτ; •E var =P n CΦE R τn+1 τn |˜uτ(Xτn)−¯un|dτ. Explicit bounds for these terms are provided ...

2001