Logit‑KL Flow Matching: Non‑Autoregressive Text Generation via Sampling‑Hybrid Inference
Authors: Egor Sevriugov
Nikita Dragunov
Anton Razzhigaev
Andrey Kuznetsov
Ivan Oseledets
Venue: ICLR 2026 Poster (2026)
Keywords: Flow Matching
NAR text generation
Primary area: generative models

Abstract: Non-autoregressive (NAR) language models offer notable efficiency in text generation by circumventing the sequential bottleneck of autoregressive decoding. However, accurately modeling dependencies in discrete sequences remains challenging in this paradigm. In this work, we advance the field of NAR generation by applying conditional flow matching (CFM) methods grounded in geometrically principled interpolation, specifically leveraging Kullback-Leibler (KL) divergence geodesics, which correspond to linear interpolation in logit space. We rigorously establish that maximizing conditional likelihood in this setting precisely recovers the flow matching velocity field, supplying the theoretical justification for this approach in sequence modeling. To address practical performance gaps of \emph{basic} inference, we propose a novel empirical \emph{sampling} strategy that iteratively denoises and re-noises, along with a \emph{hybrid} scheme that integrates our \emph{sampling} method with \emph{basic} procedure. Across unconditional and conditional text and code infilling, the approach improves perplexity and downstream metrics over prior NAR baselines under matched settings.

Published as a conference paper at ICLR 2026
LOGIT-KL FLOW MATCHING: NON-AUTOREGRESSIVE TEXT GENERATION VIA SAMPLING-HYBRIE INFER- ENCE
Egor Sevriugov
Nikita Dragunov
LigandPro, Moscow, Russia Applied Al institute, Moscow, Russia egor sevri sevriugov@gma com
FusionBrain Lab, AXXX, Moscow, Russia
Anton Razzhigaey
Andrey Kuznetsov
FusionBrain Lab, AXXX, Moscow, Russia
FusionBrain Lab, AXXX, Moscow, Russia
Ivan Oseledets
AXXX, Moscow, Russia Applied AI institute, Moscow, Russia
ABSTRACT
Non-autoregressive (NAR) language models offer notable efficiency in text gen- eration by circumventing the sequential bottleneck of autoregressive decoding. However, accurately modeling dependencies in discrete sequences remains chal- lenging in this paradigm. In this work, we advance the field of NAR generation by applying conditional flow matching (CFM) methods grounded in geometrically principled interpolation, specifically leveraging Kullback-leebllr (KI.) divergence geodesics, which correspond linear interpolation in logit space. We rigorously establish that maximizing conditional likelihood in this setting precisely recovers the flow matching velocity field, supplying the theoretical justification for this ap- proach in sequence modeling. To address practical performance gaps of basic in- ference, we propose novel empirical sampling strategy that iteratively denoises and re-noises, along with hybrid scheme that integrates our sampling method with basic procedure. Across unconditional and conditional text and code infill- ing, the approach improves perplexity and downstream metrics over prior NAR baselines under matched settings.
INTRODUCTION
Non-autoreeressive (NAR) language models have emerged as efficient alternatives traditional autoregressive models in NLP by generating all tokens simultaneously. However, capturing complex dependencies in discrete textual data remains challenging without sequential modeling.
We investigate conditional flow matching (CFM) methods for text generation. building on recent advances such as Discrete Flow Matching (DFM) Gat al. (2024), Dirichlet Flow Matching Stärk et al. (2024), and Fisher-Flow Davis et al. (024)). which represent tokens as one-hot vectors in a 1-dimensional simplex. These methods interpolate sequence of distributions Pt from an initial Po to data distribution P1: for text, the latter is sampled as discrete sequences in the simplex. Prior work identifies issues with naive linear interpolation simplex space Stärk et al. (2024). We propose instead using KL-geodesics, equivalent to linear interpolation in logit space, to better capture the underlying geometry.
Our CFM framework leverages this interpolation, training with denoiser maximizing the condi- tional likelihood Po (z) Iz),, enabling tractable approximation of the joint distribution. While theoretical guarantees previously existed only for single-token predictions, show that maximiz-

---

Published as a conference paper at ICLR 2026
Simplex
F5 Po
Pl
Denoleer natwork
Denoiser network
iuterpclate(so.kk)
3p no
X
x.V
Fash m log(sts) + Manestnnaa Faty
Simplex
Denoiser network
plati
Interpolnte(Xp...
Sertrantll
Fb NM Jk
(a) Training Phase
(b) Basic inference
(c) Sampling inference
Figure I: Overview of the Proposed Approach: Training: Sample Zo Po (uniform distribution on simplex), 31 ~ Pt (target distribution represented by samples); interpolate to obtain It. The denoiser network predicts log po (2: ( trained via log-probability maximization. Inference: For basic inference, numerically solve an ODE with vector field: (xx==p((r)==))) [v[2,, t (x)) using Euler method with N steps and a step size of h = 1/N. Alternatively, in sampling inference, interpolate between xo po and #1 p(x (mn)) each step.
ing this conditional likelihood still exactly recovers the flow matching velocity field in logit space for sequence modeling, lending theoretical support to our approach.
Standard inference procedures with this framework yield suboptimal results, so we introduce a novel sampling strategy: given state Ty, we sample 2i from p(æ, xx)) and re-noise Fith, iterating this process. Despite the lack of full theoretical analysis, this method yields stronger empirical results. We further propose hybrid inference scheme blending our basic and sampling strategies, yielding improved performance tasks such as text generation, conditional question answering, and code infilling (see Figure
Our contributions are:
Using KL-geodesic (logit-space linear) interpolation for flow matching discrete se- quences.
Theoretical analysis showing conditional likelihood maximization exactly recovers the flow matching velocity field for logit-space interpolation.
A novel sampling and hybrid inference strategy with strong empirical results.
Empirical improvements: at least 27% lower perplexity for unconditional generation (Fine- Webe, and at least 17%, 26% BLEU boosts for conditional tasks (Lamini Instruction, WMT 14 de-en); plus 56% and 14% gains in @ass and Pass @ 10 for code infilling where 10% of the code lines were omitted. Prior methods are trained and evaluated under the same setup to ensure fair comparison.
BACKGROUND
2023 constructs a deterministic transport from a simple base distri- bution Po (e.g.. (V0, I)) to an unknown data distribution P1 given by samples. It introduces a time- dependent density p(x, () and velocity field v(x, 1) for E [0, 1] that satisfy the mass-conservation (Liouville) equation
Flow matching Lipman et al.
t, - V (p(x, t)
(1)
with boundary conditions p(:.0) = Po and p(:, = Pr. Once is known, samples are generated by integrating the characteristic ODE
dx. = v(ze,t).
Tt=o ~ Pu.
(2)
dt
2

---

Published as a conference paper at ICLR 2026
Rufoww
Kuflow
Fihver Fiee
Figure 2: Qualitative and quantitative comparison of three distinct classes of geodesics on the prob- ability simplex: Linear, Fisher-Rao, and KL-Flow. The first two panels juxtapose the trajectories of the Linear and Fisher Rasher-Rao interpolations against the KL-Flow interpolation. The rightmost panels depict the temporal evolution of the divergence KL(zi]]ar) for two different simplex dimensions, [V/ 10 and | 10 000.
Table I: Perplexity (lower is better) obtained by 150M-parameter language models trained on the Fine Web corpus under the Linear, Fisher-Raa Rao, KL.-Flow geodesics.
Geodesic
Llama
GPT 3
GPT
Linear
1344
15418
13881
Fisher-Rao
192
298
379
KL-Flow
41
53
62
and taking Tz=1 as a draw from P1-
Learning the velocity field (conditional flow matching). Because u unknown, approxi- mated with neural network wa, t) using interpolation between initial and target data samples. Let Te (= x1) be any interpolation with no = zo and 71 = I1 draw t ~ 20, 1], xo ~ Po. and # ~ P1. d
and set 2; = 7x(00, x.). Define vector field v(a,, & Io, x1)
ae 7 (1r, æ1). The conditional
flow matching objective is
(CF = Ee, [I| vo (7r(20, x1 m1), v ((7 an ) vo. x |12]
At inference, we integrate the ODE with from t=0 to f=1 to obtain samples.
3 CONDITIONAL FLOW MATCHING FOR DISCRETE SEQUENCES
In language modelling, the terminal random variable is one-hot vector (a vertex of the (V-m)-si Following SStäkk al. (2024), we take the initial distribution Po be the uni- form (Dirichlet(1)) measure the simplex, so Zg Po is strictly positive probabilistic token mixture. central design choice flow matching is the interpolation between Tg and Z1-
KL geodesic on the simplex. While linear interpolation in probability space is possible, its draw- backs for discrete data have been documented (Stärk et al. 2024); Risher-Rao geodesics have also been proposed (Davis et al. 2024). We instead use the geodesic induced by the Kullback-Leibler (KL) divergence- the canonical inoomatioonthhorcee discrepancy on the simplex.
Definition 3.1 (KL geodesic). For E [0, 1]. the KL-geodesic joining Do and xı is
36 x
=
1-1 C. 0,3 Lii
C.
(4)
Σ
where C, normalizes onto the simplex.
is linear in logits, 4 = (1 t) log xo + tog 21 with x, Softmax(4).. Moreover, KL-geodesics preserve usable learning signal: shown in Fig. 2 (L[(r1|/2) decays substantially more slowly along the KL path- especially for large vocabularies Fisher-Rao paths collapse KT. (1)1) near too early (for i close to 0), effectively turning the
111=1000000whreee Linear and
3

---

Published as a conference paper at ICLR 2026
transport into one-shot step and depriving the model of informative gradients most of the time horizon. Empirically. Table [hooss that training with Linear or Fisher-Ran Rao objectives yields markedly worse perplexity, consistent with this geometric analysis.
Logit parameterization Write lo = log zo and 4 = log æi, and define the logit-linear interpo- lation (1 1)40 111, so that 2ं (l.) Because log is undefined zero, we use standard | smoothed target for the one one-hot 21.
x (1 3)
B (0,
V
where ò, is the canonical basis vector of the realized token and 1 is the all-ones vector. Equivalently, in logit space we could write linear ODE:
dlc
di 11 lo:
(5)
so the KL path is straight line in logits whose image under Softmax remains intrinsic to the simplex.
3.1 DENOISING OBJECTIVE
Single-token case. Consider the special case in which the sequence length equals one. The en- tire imput is then represented by single vector whose dimensionality matches the vocabulary size. As introduced in Definition 3.1 the KL-geodesic reproduces the conditional flow- www-mathing objec-
tive equation
(6) where H2 = Softmax(4)) denotes the intermediate point obtained by applying the softmax map to the logit vector 4 = - t) lo + tlg The quantities lo and are, respectively, the logits generating the initial state øy and the target state 21 after projection onto the probability simplex. Both the conditional vector field v(z,, Zo.
LCFM (0) = E0, | t) - (4 - 40)l
1 lo and its learnable counterpart 0o(x,, t) admit the
following reparametrisation in terms of 4e:
4 4
0g (g t) -4
4
4
o0( t)
1
Substituting the identities in equation into the loss equation 6 transforms the original objective into a denoising style regression problem in which the model must recover the clean target logit I = x1 from the corrupted observation dz:
(CF B (1 t) 111
(8)
Proposition 3.2. Let (CFM(() be defined as in equation g function
For every E (0,1) and every x the
0g (Ir, 1) = Eerxp(erlze) 4
the (almost surely) unique minimiser of the loss equation 8
Corollary 3.3. Suppose approximate the true conditional p(æ1 x²) with parametric model po(21 2z). Then an estimate of the vector field compatible with equation
u(2,, t) 1-t (
(Ens-ps(ay (mn) 4-44) 14).
(10)
The subscript 0 is omitted in v(x,, t) to emphasise that learning proceeds through the conditional density (1( 2t), rather than through direct parametrisation of the vector field itself.
Sequences of length S We now extend the analysis from the single-token setting to sequences that contain exactly tokens. As a prior over sequences we assume S independent Dirichlet distri- butions, each defined on the (V 1) simmpeee associated with the vocabulary of size V. In contrast, supported on the vertices of the Cartesian product of sim-
the "clean" or target distribution plices. Following the prescriptions in SSärk al. (2024); Gat et al. (2024). we interpolate each token independently along the KL-geodesic. Consequently, the logit representation becomes an S X V matrix 4 whose k-th row 2(x) corresponds token k.
4

---

Published as a conference paper at ICLR 2026
Fixing an index l {l...., S) and specialising Equation equation 9
the present context yields
(Ie, t) = Buimp(an(ar) (K)
(11)
where (14) denotes the logits that would generate the clean token 21
Proposition 3.4. For the KL- Kl-geodesic described above, the expression in equation 17 factorises individual tokens, and the optimal vector field for the k-th coordinate be writlen
e (Zr,t) = E 1 ~pl (x.) (R)
(12)
where p(x)")
| (T) is the marginal conditional distribution associated with the k-th token.
Consequently, under the KL-geodesic, computing the optimal velocity field reduces to evaluating the exact marginal posteriors p(zz) (k) 22a) for each token à independently. In practice approx- imate these posteriors with a parametric model po (#1 xz). We draw #1 pi (from the data distribution) and (0, 1), set Zy ~ Po. and form 2; Softmax - t) log my + log 21). The model outputs tosen-wise conditionals pa( xt); for which minimize the sequence-level
NLL:
C = E a ((
Σ k=1
log po (1
(k) |
(13)
22) can be obtained by adapting Trans- former architecture: the standard causal attention is replaced with bidirectional attention so that the representation of each token has access the entire sequence Wts and additional conditioning mechanism is introduced to incorporate the continuous time variable
A practical realisation of the conditional model (=] (K)
4 INFERENCE: ITERATIVE SAMPLING SCHEME
We present three complementary inference procedures under the Kl.ggodesic interpolation intro- duced earlier: deterministic KL-flow integrator, a stochastic iterative sampler, and hybrid routine that combines both. Unless stated otherwise, logits evolve along the logit-linear path
4 = (1 - t)4 lo + th It Softmax(4.).
4.1 DETERMINISTIC INFERENCE VIA KL-FLOW
Within classical flow matching, samples are generated by numerically integrating the ODE associ- ated with the KL-geodesic For the interpolation in Definition 3.1, the logit vector obeys the linear ODE
dlc 4
4
(14)
di 1-1
Algorithm D imppleennts an explicit scheme (Euler with step size h = 1/N) that advances / from 1. In experiments we refer this baseline as KL-flow (basic).
4.2 STOCHASTIC INFERENCE BY DIRECT SIMULATION
An alternative is to simulate the one-step transport induced by small time increment h > 0. Conditioning on the current iterate T,, the next iterate admits the Markov factorization
p( xx) / p( x1) p(#1 x.)) dx1-
(15)
The exact posterior p(r, 1T)) is intractable at the sequence level. The optimization of objective from equation 13 gives the product of tokenwise marginals produced by the denoiser:
S
Po (0 xx) = II
po xre)-
k=

---

Published as a conference paper at ICLR 2026
Table 2: Summary of inference methods.
Method
Description
Update rule
Limitations
KI.-Flow (basic)
Deterministic
1 (0(1))))))
Higher perplexity (lower text quality);
integration of learned KL-flow vector field on the simplex.
4+4: = le+ + 4-44
344:: So+tmax((+as))
KL-Flow (sampling)
Stochastic sampling along the flow using the factorised conditional.
31 ~ po(ar xx)
Assumes p(a (x)) II. p(x) 22); low entropy (reduced diversity).
2o ~ p(r))
Te+ir = interpolate(2o, x1)
KL-Flow (hybrid)
Combination of basic and sampling schemes with a switching time
Basic update for t S t",
Requires tuning t*
sampling update for >> 6*
Because the KL-geodesic interpolation also factorizes across tokens we obtain the tractable kernel
pa(e xz)
II
po xx) II
d ( pd (K) xx) dz (K)
(16)
k=1
k=1
Iterating these kernels defines an implicit model distribution over terminal states,
po(21) = p(ra) (o zo) pp) x1-1)-
(17)
This construction underpins the sampling routine summarized below; see Algorithm
Corollary 4.1 (Iterative sampler). To draw Iı ~ po(x1). initialize zy ~ Po and iterate for t
0, h 2h,....
(i) For each token k = 1,..., S. sample (hhh pa(x1 (k) xz)).
( For each token k, advance along the KL KL-geodesic by sampling Tith pl (k)
(k)
)
(ii) Set t- + h and repeat while
This KL-flow (sampling) procedure requires one forward pass of the denoiser model po(2) xx)
per iteration and thus matches the complexity of the ODE solver.
4.3 LIMITATIONS AND HYBRID SOLVER
The denoiser trained with the sequence- level NLL equation[1 furnishes only token wise marginals
po(x;" (k) Iz).. Treating these conditionally independent yields p(z1 xx) I Pa [= (k) 2;). This surrogate is exact at 1 but may degrade as t decreases due to emerging inter-token de- pendencies. To balance the stability of early- time deterministic transport,
adopt KL-flow (hybrid) procedure: integrate the ODE of Algorithm from = 0 up to a threshold t*, then switch the sampler of Algorithm 2 for the remaining horizon. Empirically, this combination improves perplexity/entropy trade-off's relative to either component alone. A concise overview of all inference schemes is provided Table 2 and a more detailed analysis is given in Appendix
[from
RELATED WORK
Non-autoregressive text generation methods can be divided into those operating continuous la- tent spaces Li et al. (2023): Ye et al.] y Gong et al. (2022): (Strudel et al. (2022) and those working directly with discrete token representations, as considered in this work. Among the latter, Campbell et al. (2024) proposed Discrete Flow Models, which combine Continuous-Time Markov Chains and normalising flows to model both discrete and continuous variables. achieving of the-art results protein generation. Gat et al. (024) introduced Discrete Flow Matching, defining sample paths between distributions via learned posterior approximations such as probability denois- Stärk al. 2024) extended this line by proposing Dirichlet Flow Matching, limiting paths to

---

Published as a conference paper at ICLR 2026
Dirichlet mixtures for tractable density calculations. Davis et al. (2024) developed Fisher-Flow, util- ising the Fisher-Raas Riemannian metric to transport mass between categorical distributions along hypersphere geodesics. Alternatively, [ou et al. (2024) presented
ceffuuiion-basee approach, gen- eralising score matching to discrete spaces for the construction of discrete diffusion models. These advances collectively demonstrate the strength of flow matching and diffusion methods for discrete generative modelling (see Appendix G for further discussion).
6 EXPERIMENTS
We evaluated KL-Flow on diverse text generation tasks, spanning unconditional language modeling, conditional sequence generation, and code infilling. All models except GPT-2 used bidirectional Transformer backbone (adapted from modded-NanoGGin
with continuous time embeddings as in DiT Peebles & Xie (2023) and logit interpolation fixed 3 0.01; top-k sampling (k 1) was used for sampling inference scheme (see Appendix E). We employed two model sizes: a 150M- parameter configuration for Storiso and 1.5B- B5B-arraeeer setup for other data domains, follow- ing the architectural and hyperparameter details of the original repository. Further hyperparameters and ablation results are provided in Appendix [F] KL-Flow was compared with DFM Gat et al. 2024), GPT-2 Jordan et al.] (2024), and SEDD Lou et al. (2024). All models taken for comparison were trained from scratch in the same setup and on the same data subset as our proposed KL-Flow model to force comparison validity. All training conducted 4 NVIDIA H100 GPUs (80GB
each).
6.1 DATASETS
Unconditional generation. The TinyStories dataset Eldan & Li (023) consists of synthetically generated short narratives authored by GPT-3.5 and GPT-4. All models were trained on B tokens with the maximum sequence length capped at 512.
To verify the scalability of KL-Flow, we further considered B tokens sampled from the Fine- Fibe dataset M-A-P et al.] (2024), which contains deduplicated and quality-filtered English web documents. Each training instance was truncated or padded to a uniform length of 024 to- kens. Models trained on this source served as the initialization traiirrainnn) for all subsequent conditional-genetti experiments.
Conditional text generation was evaluated on two sequences-to-sequeene datasets. (i) The Lamini Instruction benchmark Wu et al. (2023). (ii) The WMT14 German-Englin translation dataset BBjar jet al. (2014). In both cases the concatenation of the prompt and the ground-truth response was restricted to 512 tokens. Total training exposure was fixed at 4 B tokens.
For the code infilling task we curated an open-source Python corpus Only files comprising fewer than 024 tokens were retained. During training. for each example a uniformly random proportion between 10% % and 90 % of the lines was masked, and the model was instructed to reconstruct the clided span. Generalization was quantified on the MBPP benchmark Austin et al. (2021b).
To ensure the validity of comparisons, all baseline models were trained on the identical data subsets, using the same dataset shuffles and number of tokens to train
6.2 EVALUATION TECHNIQUES
The quality of unconditional text generation was evaluated using generative cerplexitymeesuuee by scoring generated samples with large language models (GPT-2 Radford al. (2019). GPT-3 Brown jt al. 2020)). and Llama-2 Touvron et al. (2023)) ) and diversity was assessed via empirical entropy (values above 5 indicated substantial lexical variety). For the Tiny Stories dataset, we additionally reported grammar, creativity. consistency, and plot coherence, as in Eldan & Li (2023). When scoring with external LMs (GPT-2/3, Llama-2), we use their tokenizers for perplexxiy evaluation.
500k

---

Published as a conference paper at ICLR 2026
Table 3: Comparison of unconditional text generation models trained on the Tiny Stories dataset. The results of the best-performing models are indicated in bold, while the instances where our approach matches or exceeds the performance of alternative Non-Autoregressive (NAR) methods are highlighted in blue.
Method
Grammar
Creativity f Consistency f
Plot 1
Perplexity +
GPT2
5.3
4.9
4.9
15.4
DFM
3.5
5.7
3.6
3.5
20.8
SEDD
4.2
6.1
4.0
3.8
20.7
KL-Flow
4.4
6.1
4.0
3.7
19.0
Table 4: Generative perplexity on unconditional text generation compared to prior work. Models were trained on FineFineWWl dataset. The best results are highlighted bold.
Method
Data NFE
Llama
GPT 3
GPT
9.2
15.8
31.4
GPT
1024
48.7
84.9
97.2
DFM
256/512/1024
150.6/107.3/75.0
312.8/198.9//225
381.4/245.8/1577
SEDD
256/512/1024
70.8/57.7/47.6
123.8/95.7774.8
145.8/1 14.2/90.2
KL-flow
256/512/1024
61.0/47.1/35.1
101.7/75.8/54.11
117.3/88. 1/62.9
(150M)
KL--flow
256/512/1024
51.5/41.7/32.7
81.1/63.7/48.4
96.6/76.2/565
(1.5B)
Ssqueence-t--sqqeene outputs measured using ROUGE-L (longest common subsequence over- lap) Lin (2004). BERTScore (semantic similarity via contextual embeddings) Zhang et al. (2020),
and BLEU (clipee n-gram precision with brevity penalty, s 4) Papineni al. (2002).
Code infilling was evaluated by Passēk k (fraction of synthesized functions passing all unit tests out of & samples) and Cempless@ (fraction of code snippets compiling/executing without syntax errors), for & E {1,10},
6.3 UNCONDITIONAL LANGUAGE MODELING
The experimental evaluation of the proposed framework was carried out with the KL-Flow (hybrid) inference strategy that was introduced Section 4 The numerical evidence summarised in Table B) demonstrates that KL-Flow consistently surpasses all alternative non-autte aon-autoregressive baselines across the majority of metrics, although the traditional autoregressive GPT-2 model retains an overall lead on this relatively simple dataset. In contrast, the FineFineWeh dataset imposes a significantly higher level of linguistic and semantic difficulty. aablerreport. perplexity values measured for a range of verified that
numbers of function evaluations (NFE). Before analysing comparative performance, every model under consideration preserves sufficient output variability by computing the empirical entropy of produced token distributions; all entropy scores exceeded the threshold of 5, thereby firming generation diversity. When the NFE parameter is kept at its default value 1024, KL-Flow in the intermediate 50M configuration already establishes clear advantage and flow-based non-autoregressive competitors. Reducing the computational budget by a factor of two (NFE equal 512) does not alter this observation: KL-Flow maintains comfortable margin. Even under an aggressive four-fold reduction to 256 evaluations, the model preserves performance that is comparable or superior GPT-2, underscoring the method's capacity for substantial gen- eration acceleration without sacrificing linguistic plausibility. Scaling the architecture from 150M to 1.5B parameters further accentuates these gains. In the larger setting, KL-Flow attains the best perplexities across all three reference language models (Llama 2, GPT-3, and GPT-2) and for every NFE level examined.
both diffusion-based
6.4 CONDITIONAL LANGUAGE MODELING
The empirical evaluation of the conditional generation framework was carried out on two com- plementary benchmarks, namely the Lamini Instruction and the WMT14 Gennah-Fnglish trans-
8

---

Published as a conference paper at ICLR 2026
Table 5:
Evaluation of conditional text generation on test set compared to prior works. The best
results are highlighted in bold.
Dataset
Method
BLEU Score
ROUGE-L
BERT Score
Top-5
Avg
Top-5
Avg
Top-5
Avg
GPT 2
7.8
3.1
28.9
18.2
63.8
56.4
DFM
8.1
3.6
30.0
19.2
61.6
53.6
Lamini Instruction
SEDD
5.4
2.1
25.9
15.8
61.0
53.7
KL-flow
9.5
4.3
34.5
23.9
67.9
61.1
(hybrid)
KL-fflow
7.7
4.1
31.3
21.5
66.6
60.1
(sampling)
GPT
19.7
9.8
48.3
36.7
78.1
71.0
DFM
21.3
11.2
50.0
38.8
77.1
69.6
WMT14 De-En
SEDD
14.6
6.5
44.9
34.5
74.2
68.2
KI.-flow
23.8
13.7
53.5
44.7
82.1
77.7
(hybrid)
KL-flow
27.0
18.1
56.9
49.4
84.5
81.2
(sampling)
def move_ntn (test_stri: def move_ (teat_atr): def www.ttttt_ttttt def (test_str):
dig
test_str, uid, dig test_ str if cle.itdigit
Convers given string ele (trst_str))s ele.indigtt!):
For ele strs
for ele in test str:
if ildiisdigtt
ele.indigi: dig 1- ele else:
dig 1- ele else:
dig 1- ele else:
dig 1- ele
else:
ele
1- ele
ele
ele
dig
dig
return (zes)
Ires)
return (zes)
(zes)
(a) KL-flow
(b) DFM
GPT
(d) SEDD
An illustrative example of code infilling. The highlighted lines were generated by the model. Lines highlighted in green indicate correct infilling, while those highlighted in red denote incorrect infilling.
Figure 3:
lation datasets. Performance was quantified through the standard metrics BLEU, ROUGE-L, and BERTScore; the corresponding results, reported in Table 5 include both the maximum value ob- tained among the top 5 decoded responses and the mean over this candidates, thereby providing simultaneous insight into peak quality and output stability. Inspection of the numerical results re- veals that the KL-Flow consistently surpasses all prior works. When the conditional distribution admits multiple plausible continuations, as in the Lamini Instruction scenario, the hybrid inference strategy achieves the highest scores across all metrics. By contrast, in the lower-entropy setting of deterministic machine translation, the purely sampling based variant exhibits a clear advantage.
6.5
CODE INFILLING
The code-infilling problem requires a model to reconstruct those program lines that have been re- moved, using both the surrounding source context and the natural-language task description. In the present study the network must generate a replacement of arbitrary length, up to tokens. During training and evaluation we conceal a randomly chosen fraction of the original lines; this fraction is drawn uniformly between 10% and 90% of code lines. Figure 3 illustrates infilling example. For completeness adapted GPT-2 baseline the same setting. Each masked line is replaced by the specified token and the transformer is trained autoregressively so that, after producing the unmasked part of the program, it appends the content of every hidden line in order.
Table 6 summarises the outcomes for three representative masking regimes: 10%, 50%. and 90% of the code removed. Across all regimes Kl-FFlw model with hybrid inference scheme surpasses

---

Published as a conference paper at ICLR 2026
Table 6: Quantitative comparison of several code-infilling approaches on the MBPP benchmark.
For each masking ratio the two quality indicators Pass @ k and Compiles@ k are reported for & E {1, 10}. The highest value every column appears in bold.
Infilling 10%
Infilling 50%
Infilling 90%
Method
Pass
Compiles
Pass
Compiles@
Pass
Compiles@@
10
10
1
10
10
10
10
GPT-2
8.8
20.1
54.2
92.8
0.7
3.4
27.5
67.6
0.1
0.6
15.7
56.2
DFM
11.1
25.5
39.7
88.8
2.6
8.0
15.7
59.3
0.1
1.1
7.0
33.2
SEDD
9.2
22.1
51.7
93.7
1.8
6.6
30.3
77.9
0.1
0.3
16.8
60.2
KL-Flow
17.4
29.2
73.7
92.0
4.4
11.2
58.1
87.4
0.2
1.7
60.4
90.8
prior approaches in both functional correctness and syntactic validity, Detailed curves covering the entire masking spectrum appear in Appendix [
7 CONCLUSIONS AND FUTURE WORK
In this work, we propose using Kullback-Leibllrr (KL) divergence geodesics-equine to linear interpolation in logit space-as a principled approach to flow matching in discrete sequence model- ing. Our theoretical analysis shows that the likelihood maximizer precisely matches the exact flow matching velocity, establishing strong foundation for our method. We also introduce new em- pirical sampling algorithm which, despite limited theoretical guarantees, consistently outperforms baselines in conditional text modeling on benchmarks such as WMT14 de-en translation and code infilling. Additionally, dures, achieving strong results in unconditional and conditional generation tasks, including Lamini Instruction dataset. Our findings show that larger models further improve performance, though cur- rent progress is limited by computational resources. Therefore, future work should focus on scaling model size and training unlock further gains.
hybrid inference approach combines both basic and sampling proce-
ETHICS STATEMENT
This work does not involve human subjects, personally identifiable information, or any sensitive data. All datasets used are publicly available and widely used in prior research. We are not aware of any ethical issues or potential negative societal impacts related the methods results presented in this paper.
REPRODUCIBILITY STATEMENT
We have made significant efforts to ensure the reproducibility of our work. Theoretical claims are supported with formal derivations and proofs provided in Sections 3 and Appendix [A] The main inference scheme is described in Section and further detailed in Algorithms T and 2 the Appendix. Model architectures, dataset descriptions, training procedures, and hyperparameters are provided in Section 6 and Appendix F An anonymous implementation of our method, including training and sampling scripts, also provided in the supplementary submission.
10 ACKNOWLEDDGMENTS
The work was supported by the grant for research centers in field AI provided by the Ministry of Economic Development of the Russian Federation in accordance with the agreement 600000C113925P4F0002 and the agreement Nv139-10-200 10-205-033
REFERENCES
Jacob Austin, Daniel Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2021a. URL httpt:tppp.... /tttts://arxi.... wwtp:/iwwwogoo///oo
10

---

Published as a conference paper at ICLR 2026
2107. 030000
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoe Le, et al. Program synthesis with large language models, 2021b. URL /ettps:/aaren
org.orggg//.oo.. abs /tabs/2100 2200/200.... 07732)
Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow. Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ales Tamchyna. Findings of the 2014 workshop on statistical machine transla- tion. In Proceedings of the Ninth Workshop on Statistical Machine Translation, pp. 12-58, Baltimore, Maryland, USA, 2014. Association for Computational Linguistics. URL https: //aclantholoqy orq //////11111 -
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Rad- ford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL
https://arxivs http:///axxxv...raaaboooo httgg//bb//0/00// 14165
Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnad Doucet. A continuous time framework for discrete denoising models, 2022. URL
https://arxiv www.wwwwwwww..raaaoaaa.aaaoeee htgg/abs/2205..
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design, 2024. URL URRLpttps://arxct
org/ abs / 2202244022220 04997
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman. Maskgit: Masked generative image transformer. I Proceedings of the IEEF/CVF Conference on Computer Vision and Patlern Recognition (CVPR), pp. 11305-11315, 2022.
Oscar Davis, Samuel Kessler, Mircea Petrache, Ismail Ilkan Ceylan, Michael M. Bronstein, and Avishek Joey Bose. Fisher flow matching for generative modeling over discrete data, 2024. URL
wwwpiiiwwrrrr////o atcp///bs///433 14664
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, and Jonas Adler. Continuous diffusion for categorical data, 2022. URL
hittss://artie org. abs 2211 15089
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english?, 2023. URL httpt:tttt///rttttt /atpss:///rxvv../ org/ abs 2305 2305.07759
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching, 2024.
15595
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. In Proceedings of the 2019 Conference on Empirical Methods Natural Language Processing (EMNLP), 2019.
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models, 2022. URL 2210. 08933
Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. Ssd-Im: Semi-autoregresssvve simplex-based diffusion language model for text generation and modular control. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (ACL), 2022.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. In Advances Neural Information Processing Systems (NeurIPS). 2021.
11

---

Published as a conference paper at ICLR 2026
Keller Jordan, Jeremy Bernstein, Brendan Rappazzo, fembear, Boza Vlado, You Jiacheng, Franz Cesista, Braden Koszarsky, and Grad62304977. modded-nanogpt Speedrunning the nanogpt baseline.
2024.
Junyi Li, Wayne Xin Zhao, Jian-Yun Nie, and Ji-Rong Wen. Glyphdiffusion: Text generation as image generation, 2023. URL httttt://axxiii.. org. abaaaba//2222
Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion- Im improves controllable text generation, 2022. URL
14217
Chin-Yeur Lin. Rouge:
A package for automatic evaluation of summaries. Text Summarization Branches Out: Proceedings of the ACI.-04 Workshop, pp. 74-81, Barcelona, Spain, 2004, Asso- ciation for Computational Linguisties. URL
antho logy- .rrgy.orgg///.. wwtp////g///0000000000 00100004-113
Zhenghao Lin, Yeyun Gong, Yelong Shen, Tong Wu, Zhihao Fan, Chen Lin, Nan Duan, and Weizhu Chen. Text generation with diffusion language models: A pre-training approach with continuous paragraph denoise, 2023.
11685
Yaron Lipman, Ricky T. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In Proceedings of the Eleventh International Conference on Learning Representations (ICLR), 2023. URL https: https:://ppenreview. festtfoorumii.
PqUMRDCJ19E]
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete difflssio modeling by estimating the ratios of the data distribution, 2024. URL https: http:::::/aaxxi..o http///axiv.oraab////o
16834
Justin Lovelace, Varsha Kishore, Chao gang Wan, Eliot Shekhtman, and Kilian Q Weinberger. Latent diffusion for language generation, 2022. URL htttp:://arxiv.oort wttp:/r///aaba///////
09462
M-A-P, Ge Zhang, Xinrun Du, Zhimiao Yu, Zili Wang. Zekun Wang, Shuyue Guo, Tianyu Zheng, Kang Zhu, Jerry Liu, Shawn Yue, Binbin Liu, Zhongyuan Peng, Yifan Yao, Jack Yang. Ziming Li, Bingni Zhang, Minghao Liu, Tianyu Liu, Yang Gao, Wenhu Chen, Xiaohuan Zhou, Qian Liu, Taifeng Wang, and Wenhao Huang. Finefineweb: A comprehensive study on fine-grained co/
domain web corpus. Dataset on Hugging Face, 2024. URL datasets/n a aataa-appFFiiie ineWeb Version 1..1.0.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics. pp. 31 3-18,1 Philadelphia. Pennsylvania, USA, 2002. Association for Computational Linguisties. URL hntpt::tt///ett Tlttps::/acaatthoooogg wwwpattooooyyoooo/EOo.o prg/PO2 1040
William S. Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 4172-4182, 2023.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl Technical Report, 2019. URL http:/ttts//ccn.pppeeeeee
pdf
Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step- unrolled denoising autoencoders for text generation, 2022. URL
06749
Yang Song, Jascha Sohl-Dicksteni Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Proceedings of the International Conference Learning Representations (ICLR), 2021. URL htttp:://
Robin Strudel, Corentin Tallec, Florent Altché, Yilun Du, Yaroslav Ganin, Arthur Mensch, Will Grathwohl, Nikolay Savinov, Sander Dieleman, Laurent Sifre, and Rémi Leblond. Self- conditioned embedding diffusion for text generation, 2022. URL
abs/2211. 04236
12

---

Published as a conference paper at ICLR 2026
Hannes Stärk, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design, 2024. URL
https://arxiv org/ htt//aas//aa 2402. 05841
Hugo Touvron, Louis Martin, Kevin R. Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, Daniel M. Bikel, Lukas Blecher, Cristian Cantón Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernan- des, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, An- thony S. Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Ma- dian Khabsa, Isabel M. Kloumann, Alexei V. Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao. Xavier Martinet, Todor Mi- haylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, R. Subramanian, Xia Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zhengxu Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
2307. 09288
Minghao Wu, Abdul Waheed, Chiyu Zhang, Muhammad Abdul-Mageed, and Alham Fikri Aji. Lamini-Im: diverse herd of distilled models from large-scale instructions, 2023. URL https:
// 7/arxiv. org/ abs/ 2304 14402
Jiasheng Ye, Zaixiang Zheng, Yu Bao, Lihua Qian, and Mingxuan Wang. Dinoiser: Diffused con- ditional sequence learning by manipulating noises, 2023. URL https://aaax.ii.. org. abs/
2332..0022
Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. Bertscore: Evalu- ating text generation with bert. In Proceedings of the International Conference on Learning Rep- resentations (ICLR), 2020. /ttttttts///eeereeeeet htwp:eewwetttttorouoeseseeeeeeee
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang. Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du, Chen Yang, Yushuo Chen, Zheng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang, Zikang Liu, Peiyu Liu, Jianyun Nie, and rong Wen. A survey of large language models, 2023. URL http:::://ardiv.orr/bb////333 18223
Alon Ziv, Itai Gat, Gael Le Lan, Tal Remez, Felix Kreuk, Alexandre Défossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. Masked audio generation using a single non-autoregressive transformer, 2024. URL https: arxiv. org abs /20.. 04577
13

---

Published as a conference paper at ICLR 2026
A PROOFS OF PROPOSITIONS
Proposition A.1. Let (f be defined as in equation For every t E (0, 1) and every Ia the function
bg (Iz, t) = Berrp(s 4
(18)
the (almost surely) unique minimiser of the loss equation 8
Proof. Fix arbitrary pair (Ix, l). Since equation is quadratic in (0 l), its minimiser is ob- tained by differentiating the integrand with respect to the candidate value and equating the derivative to zero. Concretely,
(1,, p(ax)) / p( 11p(0
Io, 21) p( 21) day dr1-
(19)
Assuming an independent coupling (( = po(#o))1 (11) and carrying out the integral with respect zo yields
1
/ 1 p(aa x1) (1
bg (r,, 1) = p(x))
dri.
(20)
By Bayes' theorem, p(z1 x1) p(a.(2)) pi(1) Substituting this identity into equation 20 imme-
p(a,))
diately furnishes equation 9 completing the argument.
Proposition A.2. For the KL-geodesic described above, the expression in equation 77 factorises over individual tokens, and the optimal vector field for the k-th coordinate can be writlen as
((k) (21,t) = E 1 ~p(= (a.) (R)
(21)
where p(z xx)) is the marginal conditional distribution associated with the k-th token.
Proof. Because ((t)) is deterministic function of 11(h) alone, may integrate out all remaining coordinates to obtain
(or, t) f 1/4 p(a1 ax) dr f 11th p(x]{1) a dxr" dzr (M)
which coincides with equation 12 While, in principle, the geodesic interpolation could introduce
dependencies among tokens through the joint kernel p(æ, Ho, æ1), empirical findings reported in Stärk al. (2024); Gat al. (2024) indicate that treating the coordinates independently suffices for practical purposes. Hence, the optimal vector field for each token depends solely on its own marginal posterior.
B FEW-SHOT TEXT GENERATION
We evaluate the capability of the considered non-autoregressive (NAR) models few-shot text generation task and compare them to the proposed KL-Flow model. The quantitative results in Table that KL-Flow consistently achieves substantially lower perplexity than the baseline NAR methods (DFM and SEDD) across all numbers of refinement iterations (4, 8, and 16). All models are trained on the Fine Fine Web dataset with sequence length of 1024.
In addition to perplexity, we measure the diversity of generated text using token-level entropy. We observe that KL-Flow tends to produce slightly less entropic (less variable) text than the baselines. This reduction in entropy is most pronounced at 4 refinement steps, where the entropy of KL-Flow markedly lower than that DFM and SEDD. For and steps, the entropy partially recovers and approaches that of the basclines, while preserving the perplexity gains. Overall, these results suggest that the proposed KL-Flow methodology well-suited for few-shot text generation, offering strong improvements in perplexity; however, for very small numbers of refinement steps, additional tuning may be beneficial to mitigate entropy reduction and better preserve output diversity.
To complement the quantitative evaluation, we additionally report unconditional generations. Ta- ble [ shows representative samples produced by the NAR baselines and the proposed KL-Flow model. The generated samples truncated the first 300 characters for clarity of presentation.
14

---

Published as a conference paper at ICLR 2026
Table Few-shot text generation evaluation for number iterations equal 4, 8, 16. Following NAR methods considered: DFM, SEDD, KL-Flow. The evaluation performed with models trained on Fine Fine Web dataset with sequence length 1024.
Perplexity
Entropy
Method
4
8
16
4
8
16
DFM
1017.4
687.7
451.9
5.5
5.5
5.5
SEDD
839.8
561.2
321.0
5.5
5.5
5.5
KL-Flow
76.6
179.4
99.8
3,6
5.1
5.1
Table 8: Unconditional text generation examples produced by non-autoregresssve baselines (DFM, SEDD) and the proposed KL-Flow model.
Model
Generated text
Refinement steps 8
DFM
bill the age, important two growers with fatty foods or vegetables and equipment kicked have to be rachest and powered healthy GAFO depending upon plate/ conceived order, et, This Newsletter se
size.b outl to cook, Luphem. departure. press- step- to comply with monthly priorities of b
SEDD
for exclusive content through what you are seeing daily The LiveIt numbers images 422 picture maximum the initial Cments menu applying (see ReveFast* latest information) to purchase different videos are special or not AutnRstan SAC MOI Engineer Handbook page I, в These goals have enough that
KL-Flow (ours)
15 hectares and the village plots average 660 per tree.the Polsripini area is close to the sierra where the agricultural terrain spans 15 000 metres from 2.5 hectares. what Makes You Growers ? Farming, Buddha's ear clove is highly appreciated for more resicky. Excellent butter ta
Refinement steps 16
DFM
excitementS farmers quclve/pro'N contributers west-told had much difficulty in Iccal, production and pride surprised farmersFish production areajpfuisemakers profuse like Hos According to aniseed, adding a similar crop can help alleviate (GwGamm),, afford fertiliser, and were placing in com
SEDD
men to risk upwards territeries from throwing fresh land deposits elements; days you can. paylead unit be 1. p.ddf, obtained at the Space International Eg at 3:27s appropriate sustainment capacity 3 months of mission the International Space Station).
servicing00
super-fast launch you can change
KL-Flow (ours)
-like plants that can lead to somo starvation. Sustainabiliss is when they are forced to live in an unwanted direction through the internal parts of the plant. "realizing that plants behavior in a way that one could imagine, it not evident that plants also refer to ot
inspired by their
C ALGORITHMS OF INFERENCE SCHEMES
Here, present algorithms for basic and sampling inference schemes, Algorithms I and
15

---

Published as a conference paper at ICLR 2026
Algorithm Inference scheme (basic)
1: Input: Initial distribution Pu; denoiser model po(#1/2)) parameter N (number of iterations); parameter (time step size, default 1/N).
2: Set = 0
3: Sample 34 ~ Po
4: for 1 to N do
5: Compute = pe (a(p() (21 (z))
Compute smoothed target logits 1 = w log (1 B + +) + (1 - w) log 2
6:
7:
Compute l +- 4 1-1 (11 - 12)
8:
Update 24 + Softmax(4)
9: Update t t +h h
10:
end for
11: Return Zt.
Algorithm 2 Inference scheme (sampling)
1: Input: Initial distribution Pu; denoiser model pe(x1 (x)); parameter N (number of iterations); parameter (time step size, default 1/N).
2: Set 0
3: Sample 24 ~ Po
4: for 1 to N do
(k) 5: Sample #1 ~ pe(o)) (xx)) for & E S]
6: Sample za po
7: Compute 4+h = - log(zo) + (t + h) log(z1)
8: Update 24 Softmax(1e41)
9: Update <- +
end for
11:
Return Zt.
D ADDITIONAL CODE INFILLING EXPERIMENT
In this section we present full comparison of code infilling task for an arbitrary amount of masked lines. The results were summarized in Figure 4 For most cases the KL-Flow outperforms other approaches across all considered metrics. The most noticeable advantage could be seen in Com- piles@1 metric, where for any portion of missed code lines the difference from closest competitor above 10%.
E COMPARISON OF INFERENCE SCHEMES WITH ANALYSIS OF TOP- SAMPLING EFFECTS
Table Quantitative comparison of inference schemes in terms of perplexity and entropy.
Method
Perplexity
Entropy
KL-FFlow (basic)
154.2
5.6
KL-Flow (sampling)
3.8
1.9
KL-Flow (hybrid)
41.4
5.2
In this section analyse the inference procedures introduced in Section 4 and study how their hyperparameters affect performance. Our main practical proposal is the sampling inference scheme, which repeatedly denoises and re-noises the current state. Its derivation relies the factorisation assumption
p(r1 xz) ~ Il
((= (0)
xx),
(22)
16

---

Published as a conference paper at ICLR 2026
0.175
GPT-2
0.7
DFM
SEDD
0.6
0.125
KL-Flow
Compiles@)
0.5
GPT-2
Pass@11
0.100
DFM
SEDD
0.075
KL-Flow
0.3
0.050
0.2
0.025
0.1
0.000
Percentage of removed lines
Percentage of removed lines
0.30
GPT-2
0.9
DFM
0.25
SEDD
0.8
KL-Flow
Compiles@10
Pass@10
0.7
0.15
0.6
0.10
0.5
GPT-2
DFM
0.05
0.4
SEDD
KL-Flow
0.00
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Percentage of removed lines
Percentage of removed lines
Figure 4: Comparison of the performance of prior models against Kl.-Flow on the code infilling task
where ( denotes the i-th token. This approximation is exact at and clearly fails at O. Our goal is therefore to understand for which times & the factorised approximation is accurate enough, and identify threshold such that for t > t" the difference between p(z) xx) and
p(=)) 2x) is negligible.
As shown in Corollary 3.3) in the single-token case the solution of the Now-matching problem coin- cides with the conditional model Po(x1 zz). In this setting we can derive the exact conditional in closed form.
Proposition E.1. Consider the KL geodesic the simplex with uniform prior zu and sequence of length one. The exact solution p(æ1 22) is given (up to a normalising constant) by
log,po(<1 xx)) llgg((z) V log Σ exp(Lo) C.
(23)
where C. is normalisation constant, V the vocabulary size, and
4 t
(24)
1-t
a (V V) matrix whose rows contain the logits of the preimages No associated with each simplex veriex of 1- The summation Σ is taken over the last dimension of Lo- Throughoul, capital letters for matrices whose first dimension indexes vertices and whose second dimension indexes simplex coordinates.
Proof. By Bayes' rule,
1og (4)) log log p(x1) log p(=; x1)-
(25)
17

---

Published as a conference paper at ICLR 2026
The marginal p(z)) does not depend æ1 and can be absorbed into the normalisation constant. The remaining term can be written using the thange-of-variabee formula:
dXg dx,
logp(r; (11) = log
(26)
dXe is three- direee-imensionll tensor whose first index enumerates vertices and whose last two dimensions correspond to the Jacobian with respect to Zt. The determinant is taken over the last two dimensions, resulting in a vector over vertices.
where da,
From the KL-geodesic interpolation (equation 3.1
we obtain the of preimages Xo of shape
(V, V) as
4
LL!
Xg = Softmax
(27)
1-t
with 4 log st. The Jacobian of the Softmax map with respect its logits is
Softmax(z) diag(x) .r
(28)
This matrix has one zero eigenvalue because Softmax is invariant under adding a constant to all logits. Consequently, its determinant is the product of the non-zero eigenvalues only. Writing A diag(x) xr and using the characteristic polynomial
det(A- - AT) xA))),
(29)
the product of the non-zero eigenvalues is g(0), which can be obtained
d Yx
(0) det(A - Ar) X=o
(30)
Since A - A1 = diag(x) AI - xx a diagonal matrix plus a rank-one update, its determinant
admits the closed form
det det (diag(z) - AI Zr ) = II (xt 2) (1- x. diag -(x 2)
(31)
Differentiating yields
(0)) = V
II 3
(32)
up multiplicative constant that is absorbed into normalisation. Therefore,
14 - tL
log p(x x1) = C + Σ log Softmax
(33)
1 1-t
where the summation is over the last dimension and
C' = log
Σ log
Ht
(34)
1 1
collects all terms independent of æ1.
Using the identity
log Softmax(() log Σ exp(1),
(35)
we obtain
logp(#; (m)) = Σ -
th log
Σ exp
4 4-t111 1-t
(36)
again up to an additive constant. The first term does not affect the relative probabilities over Ti Σ 4 is constant and Σ L1 contributes equally to every vertex. Hence the dependence on #1 arises entirely through
-V Σ exp(Lo),
(37)
with Lo 4-th; which completes the expression 1-1
log (x xx)) (ogp(z)) V log Σ exp((Lo C.
(38)
18

---

Published as a conference paper at ICLR 2026
p(xi) = k(x)
10
p(x)) U
t= 0.28
p(x, (x))
KL(p(x [x)]
2-
0.0 0.2
0.4
0.6
0.8
1.0
Figure 5: KL divergence between the exact conditional p(21 2x) and the KL-geodesic-only ap- proximation p(r1 œx) as function of time for a vocabulary of size V = 50k. We report results for two priors
vertices, p(z1) = and p(x1) U. The vertical line indicates the threshold = 0.28, beyond which the approximation error becomes negligible and the dynamics are effec- tively governed by the KL--geddsic term.
Proposition E. shows that, in the single-token setting under KL-geodesic with uniform prior, the exact posterior decomposes into two contributions: a vertex term log p(x1) capturing the prior probability of the token, and KL-geodesic term -V log Σ exp(Lo) capturing how likely it to reach a given vertex from the current state zt. The relative strength of these two terms varies with time t.
Corollary E.2. In the setting of Proposition E.1
0 the posterior reduces to the vertex term,
logpo(2; (x) xg)
logp(x1) + +C,
(39)
whereas for L + 1 the KL-geodesic contribution grows the prior term p(x1)
V I through Lo 4--L1, and dominates I-t
For large vocabularies (e.g., V ~ 50k) the balance between these terms quickly shifts in favour of the KL-geodesic component as increases. This behaviour is illustrated in Figure 5 which reports the KL divergence between the full conditional
p( xx) Soflmax(log((z1)) V log exp((Lo))
(40)
and the approximation that retains only the KL-geodesic term,
p(z) xx) = Softmax( - V log Σ eep((Lo)
(41)
For clarity we consider two extreme priors: a point mass p(r1) = å and the uniform distribution p(z1) U. The vertical line t" 0.28 marks the threshold at which p(x1 xx) p(œ1 xx), indicating that for t > t" the dynamics are largely governed by the KL-geodesic term and become effectively insensitive the prior p(x1).
This observation is crucial for extending the analysis to multi-token sequences (sequence length 1). For two tokens,
log (2
(1) (x] (2 xx)) log (x)) x2) + log (x)) (1)))
(2)
(42)
19

---

Published as a conference paper at ICLR 2026
tp exty
Encose Eroway traahald
Perkotos
Perste
-1-028
107
200
300
2.0
2.5
3.5
5.0
5.5
Entropy
(a) Variation across different top-k values.
(b) Selection of the optimal switching time t".
Figure Effect of inference hyperparameters on generation quality. (a) Influence of the top-k parameter in the sampling inference scheme. Performance is measured by perplexity under Llama 2 and by token-level entropy; the horizontal line marks the empirical entropy threshold associated with diverse text generation. (b) Dependence of perplexity and entropy the switching time t* in the hybrid inference scheme, which transitions from basic to sampling updates. The optimal trade-off attained at t* 0.28.
where superscripts denote token indices. The vertex term now also encodes inter-token dependencies (1)). The discrepancy between the exact posterior p(#1 se)) and the tokenwise factorisation p(x xx.) is entirely due to these dependencies. The single-token analysis and Figure 5 together suggest that for t 2 0.28 the KL-geodesic term dominates sufficiently to suppress the effect of inter-token correlations, making the factorisation good approximation:
through the conditional p(=)?) | =1,,
p( x4) ~
I[
xr)) for 1 12
(43)
This justifies the use of the sampling inference scheme in the aate-time regime.
Empirical effect of and top-k. We now study the impact of in practice by measuring per- plexity (using Llama 2 as the scorer) and token-level entropy. Two main factors are varied.
First, sweep the threshold that controls the relative share of basic versus sampling steps in the hybrid inference routine. The results in Figure 6(b) indicate that t* 0.28 yields the best compromise between low perplexity and high entropy. At this setting, the entropy remains above the diversity threshold of 5. Moving t* away from this optimum leads to marked degradation in either perplexity or entropy, harming quality or diversity respectively.
Second, we examine the role of top-k sampling during the sampling phase of inference; see Fig- ure ([... Increasing l initially improves entropy, but large values of eventually deteriorate text quality reflected by perplexity. In practice we adopt k = 1, which already achieves sufficiently diverse outputs (entropy 5) while maintaining strong perplexity.
Finally, Table summarises performance across the three inference schemes: hasic, sampling. and hybrid. The sampling-only variant suffers from low entropy, consistent with the discussion Sec- tion 3.2, paragraph Approximation of the conditional p(z1 xz). Conversely, the basic scheme produces comparatively high-entropy but tow-quality text, as indicated by large perplexity. The hy- brid method, which combines both regimes and leverages the threshold t*, delivers the best overall trade-off.
F OPTIMAL TRAINING CONFIGURATION
In this section, we discuss several critical aspects and technical strategies for addressing the Flow Matching (FM) problem. The foundational code and architecture employed for training were derived
20

---

Published as a conference paper at ICLR 2026
Ir Adam 0.00018
Ir Adam = 0.000036
Loss
Ir Adam = 0.00036
Ir Adam 0.0036, no Muon optimizer
Ir Adam = 0.0036
4
0
25k
50k
75k
100k
125k
150k
175k
200k
Iteration
Figure 7: Comparison of the impact of learning rate values training GPT-like model for the Flow Matching problem. The base implementation utilizes the Muon optimizer for certain model parameters, while the tag "no Muon optimizer" indicates that the Muon optimizer has been replaced with the Adam optimizer.
from an open-source GitHub repository featuring efficient implementation of the GPT-2 model, designed for standard language modeling tasks. However, our investigation revealed that the initially suggested optimal configuration is not truly optimal for the FM problem.
key factor influencing convergence is the selection of an appropriate learning rate. Figure we present comparison of various learning rate values, alongside an assessment of how the integration of the Muon optimizer- proposse that the standard learning rate of (Ir 0.0036) is not optimal. learning rate reduced by factor of ten significantly accelerates convergence and mitigates the risk of stagnation during the initial determined that the ratio of learning rates between the Adam
the original repository- affectt model performance. We found
phases of training. Furthermore, optimizer and the Muon optimier yields optimal results. Additionally, the application of the Muon optimizer for specific model parameters enhances convergence, even when employing non-optimal learning rate.
Another critical consideration is the method of incorporating temporal information into the model architecture. We identified three primary strategies for this purpose:
Time Token: Transform the time value into an embedding vector and incorporate as separate token within the sequence.
21

---

Published as a conference paper at ICLR 2026
7.0
Time token
Layer norm Standard
6.5
6.0
5.5
Loss
5.0
4.5
4.0
3.5
25k
50k
75k
100k
125k
150k
175k
200k
Iteration
Figure 8: Comparison of various strategies for time insertion within model architecture.
Layer Normalization: Employ method akin that used in the DiT architecture, where the time embedding is utilized to adjust the mean and standard deviation of the data within the layer normalization module.
Standard Addition: Simply append the time embedding to each token embedding.
Our findings, as presented in Figure 8 indicate that the Layer Normalization strategy is the most effective approach, as provides better convergence and achieves a lower loss value after 200k training steps.
G RELATED WORK FULL DISCUSSION
this section, we review the literature on modeling discrete sequences. The authors Campbeil et al. (2024) present Discrete Flow Models (DFMs) that combine discrete and continuous data using Continuous Time Markov Chains, improving traditional diffusion methods for protein co-design and achieving state-of-the ahe-art results in protein structure generation.
Additionally, Song et al. (202 propose a stochastic differential equation (SDE) for transforming complex data distributions using neural networks for accurate estimation. The work by Camp- bell et al. (2022) introduces a continuous time framework for denoising diffusion models of discrete data, resulting in high- performance samplers that surpass traditional methods.
Research by Gat et al. (202)) introduces Discrete Flow Matching, focusing on generating high- dimensional discrete data, such language, while enhancing generative perplexity. Meanwhile,
22

---

Published as a conference paper at ICLR 2026
Ghazvininejad al. (2019) masked language modeling to predict target words based imput text, and Austin et al. (2021a) improve multinomial diffusion models. Finally, Hoogeboom et al. 2021 provide extensions for categorical data, demonstrating high efficacy in text modeling and image segmentation.
Recent advancements have focused applying continuous space diffusion methods to discrete datasets Dieleman et al. (2022); : Li et al. (2022); Han et al. (2022). Notable contributions from Lin et al. (2023) improve diffusion flow modeling, while new Continuous Flow Matching techniques introduced by Lovelace et al. (2022) and Stärk et al. (2024).
Autoregressive models have been crucial in natural language processing Zhao et al. (2023). exem- plified by the GPT-2 model Radford et al. (2019), which showcased the potential of autoregres- sive approaches in generating coherent text. Research highlights the effectiveness of autoregressive methods in addressing complex linguistic challenges.
Masked generative modeling has emerged promising area, utilizing techniques to generate con- tent by obscuring parts of imput data Ghazvininejad et al. (2019). Studies by [Savinov et al. (2022) refined traditional masking methods, leading to innovations like MaskGIT, which employs advanced techniques for high-resolution image synthesis Chang et al. (2022). Furthermore, ZZiv et al. (2024) demonstrated the effectiveness of text-to-music model, showing that the MaskGIT M sig- nificantly improves the quality of generated outputs.
23


Reviews:
[{"review_id": "8YapN2rGE5", "reviewer": "Reviewer_VSxC", "summary": "This paper presents a non-autoregressive text-generation framework that performs conditional flow matching on the probability simplex using KL-divergence geodesics. The authors show (in their way) that maximizing the token-level conditional likelihood exactly recovers the flow-matching velocity field, providing a principled training objective.", "questions": "No, see weakness.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The mathematical descriptions in this paper are relatively accurate.\n2. The paper is well structured.", "weaknesses": "1. The used base model in the experiments seems obsolete, which makes it unclear whether the proposed method still works for the SOTA models nowadays.\n2. Lines 226 ~ 228 seem confusing. A bidirectional attention is used to model sequence-level NLL, but the condition variable $x_t$ represents a single token, not a total sequence. I am not convinced of this modelling approximation.\n3. Line 309 says all models used in the experiments are bidirectional backbones. How do you apply this to text generation, where you do not have access to any future token information?\n4. I see some autoregressive models have been used in the experiments (GPT-2), so you replaced the causal attention in the model with a bidirectional attention? If so, how do you initialize your model weights? Is it the same with the pre-trained weights?\n5. The experiment setting does not include baseline introductions, which makes the reader very hard to get familiar with the relevant work.\n6. The proposed method in this paper has two inference methods, namely, basic inference and sampling inference. However, there is no formal section in the paper to systematically compare these two inference methods and discuss the corresponding advantages and disadvantages. The experiment section also missed this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MH7PZYlbxF", "reviewer": "Reviewer_qQ8c", "summary": "This paper proposes a novel non-autoregressive text generation framework that uses a Kullback-Leibler (KL) divergence geodesic for interpolation, which is shown to be equivalent to linear interpolation in logit space. The objective function is to minimize the negative log-likelihood of the target distribution at the sequence level. Experiments demonstrate that the proposed method outperforms other non-autoregressive (NAR) baselines such as DFM and SEDD.", "questions": "In the experiments, the authors state that all models use a bidirectional transformer backbone. Was the GPT-2 baseline trained in an autoregressive manner, and was its causal attention mechanism replaced with bidirectional attention?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed method achieves strong performance against other NAR baselines across various tasks.\n\nThe discussion of the inference process is insightful.", "weaknesses": "It is unclear why the deterministic inference process performs poorly, given that the \"Logit-KL Flow Matching\" objective recovers the velocity field.\n\nThe efficiency of the proposed method is not thoroughly discussed, particularly in comparison to methods that are trained with an MSE loss (in training) and solve ODEs using numerical techniques (in inference).\n\nPerplexity is measured using samples from GPT-2, GPT-3, and Llama-2, which may introduce bias from these reference models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E74OgT9M3R", "reviewer": "Reviewer_eLVC", "summary": "The paper proposes a non-autoregressive (NAR) text generator built on conditional flow matching (CFM) in logit space. Instead of interpolating token probabilities linearly on the simplex or along Fisher–Rao geodesics, it interpolates logits between a simple Dirichlet-like start and the target one-hot token, which the authors argue is the KL geodesic; they show that, under this path, maximizing the conditional likelihood ( $\\log p_\\theta(x_1 \\mid x_t, t)$ ) exactly recovers the desired flow velocity field. On top of that, they introduce a hybrid inference scheme that runs a deterministic ODE-style update in the early time steps and switches to an iterative sampling / re-noising procedure in later steps to fix token-level errors. On several text, conditional, and code-infilling benchmarks, the method outperforms earlier discrete / Dirichlet / Fisher flow-matching baselines and comes close to, but does not fully match, similarly sized AR models.", "questions": "1. How does this behave under true few-step regimes (e.g. 8–16 NFE) against step-consistent discrete flows like FS-DFM or consistency-trained DFM? Right now the advantage is shown largely when you afford hundreds of steps or a tailored hybrid schedule; but the practical motivation for NAR is low latency. A head-to-head with step-budget-aware models is missing. Can the KL-geodesic path still deliver better gradients than Dirichlet / Fisher when you compress the time discretization that aggressively?\n2. What exactly is the uniqueness claim over existing KL-geodesic / logit-FM papers? Several public works from late 2024 onward already stated (i) “logit-space interpolation is the KL geodesic,” (ii) “maximizing ($p_\\theta(x_1 \\mid x_t,t)$) gives you the right field,” and (iii) “basic deterministic inference is weak; add an iterative sampling-and-noise scheme; hybridize.” If this paper’s contribution is a stronger sequence-level derivation or better scaling to open-domain data, please spell out the technical gap (e.g. a specific lemma about token-wise optimality under KL-paths, or a complexity advantage in hybrid inference) that is not present in Sevriugov & Oseledets (2024) or the concurrent KL-geodesic variants. Right now it reads more like a consolidation than a breakthrough.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Clear geometric diagnosis + concrete fix. The paper identifies a real failure mode of earlier probability-space paths — linear on the simplex, Fisher-Rao sphere, even some Dirichlet settings — namely that $KL((x_{\\text{data}}|x_t))$ collapses too quickly so mid-time supervision vanishes on large vocabularies. The logit-space (KL-geodesic) interpolation directly targets this and shows improved calibrations against those baselines. This is well aligned with prior observations in Dirichlet Flow Matching that “naïve linear FM on the simplex is pathological.”\n2. Bridging “train a denoiser” and “learn the flow field” for sequences, not just single tokens. Earlier CFM/DFM papers had versions of “conditional likelihood recovers the field,” but mostly in single-site or weaker sequence assumptions; this paper pushes the argument specifically for logit-KL paths and uses it to justify a very practical objective (just NLL on corrupted sequences). That reduces the gap between elegant flow theory and what people actually train.\n3. Inference is engineered rather than hand-waved. Many discrete flow papers stop at “we have the field, integrate it”; here they run a 3-way comparison (deterministic / stochastic / hybrid) and show that pure ODE is insufficient and pure sampling collapses entropy, while a staged hybrid fixes both. That’s a useful empirical lesson for the whole discrete-flow community.", "weaknesses": "1. Novelty margin over very close contemporaries is thin. At least two 2024–2025 papers already investigated conditional text generation via KL-geodesic / logit-space flow matching and even proposed almost the same empirical “sampling + noise re-injection + hybrid” recipe to fix the underperforming basic sampler. The descriptions in Sevriugov & Oseledets (2024) and its 2025 extensions match this work’s geometric choice and sampling intuition almost line-for-line. If the contribution here is meant to be “we prove the conditional-likelihood = exact field under this path and scale it to larger datasets,” the paper needs to separate itself much more crisply from those concurrent KL-geodesic efforts, especially since they also claim better results over discrete FM. Right now the delta looks incremental.\n2. Key equivalence rests on a factorized / per-position view that is not obviously valid in early timesteps. The derivation leans on the idea that the optimal vector field at time (t) can be written as an expectation of target logits under ($p(x_1 \\mid x_t)$) token-wise. But in NAR text, ($p(x_1 \\mid x_t)$) is usually not well factorized when (t) is small: tense, agreement, long-range topic constraints all couple positions. The paper’s fix is “do deterministic updates early, sampling later,” which is an empirical workaround, not a proof that the factorization is OK. So a central theoretical selling point (“likelihood = flow”) is relying on a data-distribution property that’s weakest exactly where the model needs guidance most. That’s a structural, not cosmetic, gap.\n3. Evaluation uses LM-perplexity proxies and medium-scale AR baselines, so the true competitiveness is unclear. Measuring perplexity by scoring NAR outputs with an external LM is standard for discrete flows, but it is a proxy; it is known to favor models that mimic the scorer’s style rather than models that are truly diverse or controllable. And the main AR point of comparison is GPT-2-class models, not the 2025-era instruction-tuned or code-tuned LLMs that NAR methods would actually have to replace. In other words, the paper shows “better than prior discrete flows” but not “this can plausibly replace competitive AR models under identical training budgets.” That’s a material, not rhetorical, limitation.\n4. Hybrid inference is admitted to be heuristic and under-analyzed. The whole motivation of the paper is “basic ODE flow isn’t good enough on text,” which is fair; but the proposed fix (early deterministic, late sampling + noise) is only justified by curves. There is no stability analysis, no guarantee of staying close to the learned KL-geodesic, and no complexity comparison with recent few-step DFM / consistency-trained flows that do target a fixed number of steps directly. A reviewer can reasonably ask why we shouldn’t just adopt FS-DFM-style step-budget-aware training or consistency distillation to get the same effect with a cleaner theory.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AwaObmaEbM", "reviewer": "Reviewer_5rVW", "summary": "This paper proposes Logit-KL Flow Matching (KL-Flow), a non-autoregressive (NAR) text generation framework that performs flow matching in logit space rather than probability space.\nBy interpreting the KL divergence geodesic as a linear path in logits, the authors provide a theoretically grounded interpolation scheme for discrete sequence modeling. They further introduce an iterative sampling–hybrid inference procedure combining deterministic ODE integration and stochastic denoising steps.\nEmpirically, KL-Flow shows consistent gains over prior discrete flow and diffusion baselines (DFM, Dirichlet Flow, Fisher Flow, SEDD) on unconditional, conditional, and code infilling benchmarks.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper rigorously connects conditional likelihood maximization to flow velocity recovery in logit space.\n2. The token-wise conditional likelihood formulation is both simple and tractable, providing a bridge between discrete diffusion and flow models.\n3. The hybrid inference scheme demonstrates consistent improvements across several datasets (TinyStories, FineWeb, WMT14, MBPP), showing competitive results for both language and code tasks.", "weaknesses": "1. While the paper covers some discrete flow and diffusion models, it does not include strong diffusion-based text generation baselines such as MDLM (Masked Diffusion Language Model) and conditional NAR transformers such as Tracformer (https://arxiv.org/pdf/2502.07616?)\n2. Most experiments focus on small to mid-scale models (≤1.5B parameters); Would scaling the model or using stronger backbones (e.g., Llama-2 or Mistral) improve performance beyond the 1.5B-parameter setting?\n3. Although empirically effective, the hybrid scheme’s transition point t* is largely heuristic. Could the authors provide intuition or ablation results for how the choice of the deterministic–sampling switch point t* influences diversity and perplexity across datasets?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)