Discrete Adjoint Matching
Authors: Oswin So
Brian Karrer
Chuchu Fan
Ricky T. Q. Chen
Guan-Horng Liu
Venue: ICLR 2026 Poster (2026)
Keywords: Discrete Diffusion Model
Fine Tuning
Continuous-Time Markov Chain
Adjoint Matching
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Discrete Adjoint Matching for fine-tuning discrete generative models on reasoning tasks.

Abstract: Computation methods for solving entropy-regularized reward optimization—a class of problems widely used for fine-tuning generative models—have advanced rapidly. Among those, Adjoint Matching (AM, Domingo-Enrich et al., 2025) has proven highly effective in continuous state spaces with differentiable rewards. Transferring these practical successes to _discrete_ generative modeling, however, remains particularly challenging and largely unexplored, mainly due to the drastic shift in generative model classes to discrete state spaces, which are nowhere differentiable. In this work, we propose Discrete Adjoint Matching (DAM)—a discrete variant of AM for fine-tuning discrete generative models characterized by Continuous-Time Markov Chains, such as diffusion-based large language models. The core of DAM is the introduction of _discrete adjoint_—an estimator of the optimal solution to the original problem but formulated on discrete domains—from which standard matching frameworks can be applied. This is derived via a purely statistical standpoint, in contrast to the control-theoretic viewpoint in AM, thereby opening up new algorithmic opportunities for general adjoint-based estimators. We showcase DAM’s effectiveness on synthetic and mathematical reasoning tasks.

Published as a conference paper at ICLR 2026
DISCRETE ADJOINT MATCHING
Oswin Solit* Brian Karrer
Chuchu Fan
Ricky T. Q. Chen 2.*
Guan-Horng Liu
MIT
2FAIR at Meta
Work done at FAIR, Meta
ABSTRACT
Computation methods for solving retropy-regularized reward optimization class of problems widely used for fine-tuning generative models-hawe advanced rapidly. Among those. Adjoint Matching (AM, Domingo-Ennin Enrich et al., 2025) has proven highly effective in continuous state spaces with differentiable rewards. Transferring these practical successes to discrete generative modeling, however, remains particularly challenging and largely unexplored, mainly due the dras- tic shift in generative model classes to discrete state spaces, which are nowhere differentiable. In this work, we propose Discrete Adjoint Matching (DAM) discrete variant of AM for fine-tuning discrete generative models characterized by Continuous- Time Markov Chains, such as diffusion- large language models. The core of DAM is the introduction of discrete adjoint-an estimator of the op- timal solution to the original problem but formulated on discrete domains from which standard matching frameworks can be applied. This is derived via a purely statistical standpoint, in contrast to the control-theoretic viewpoint AM, thereby opening up algorithmic opportunities for general adjoint-based based estimators. We showcase DAM's effectiveness on synthetic and mathematical reasoning tasks.
INTRODUCTION
We consider entropy-regularize optimization in state space X' over a finite horizon t [0. 1]: min E [g(X) + DxL(p" bose (X))
where the objective consists of a terminal loss 9(-) incurred at and Kullback-Leible (KL) divergence (-Il)- that regularizes the model distribution pd over some stochastic processes X
(X1)ee(o.1) parametrized by u, to the base distribution p = pr parametrized by u hase
Optimization problems of the form (I) are ubiquitous at the intersection of reinforcement learn- ing (Levine, 2018), (), stochastic control (Theodorou et al., 2010), and probabilistic inference (Blei al., 2017). the context of generative modeling, the primary motivation for studying (1) stems
from the fact that wwhen the base distribution satisfies the memoryless condition hise (x) = y base (X1)² the optimal solution to (1) known analytically by
(X) = pd (X) X buse (X) 9(X,)
(2)
Such an analytic expression is particularly useful for post-training generative models (Ziegler al., 2019; Ouyang et al., 2022; Fan et al., 2023b), where we wish to fine-tune pretrained models given some reward function rax). Specifically, by considering p base as the pretrained model and setting the terminal loss negative reward, i.e., g(x) = -r(x), the optimization problem in (1) can be best understood as variational inference, seeking a posterior close to the prior, pretrained, distribution while being tilted exponentially toward high-reward regions, i.e., p² px(1) hase
haseee (X1)e
Computational methods for solving (1) has advanced rapidly in recent years. Among those, methods based on Adjoint Matching (AM, Domingo-Enrich et al., 2025) have shown highly effective in continuous state spaces with differentiable terminal losses, e.g., X Ie and Vg Rd - IR. achieving strong results in applications such as image fine-tuning (Domingo-Enrich et al., 2025), molecular generation (Havens et al., 2025), and Boltzmann distribution sampler (Liu et al., 2025a).
"These authors are core contributors to the project. See Author Contributions for details.
'The memoryless condition readily holds for all base models considered. See Appendix A.1 for details.

---

Published as a conference paper at ICLR 2026
Concretely, when p" is the distribution induced by a class of stochastic differential equations (SDEs) parametrized by dx2,, = of (texddtttttlle with some diffusion or, Domingo-Emiin -nrrcch et al. (2025)
observed that any optimal solution must satisfy fxxed-point equation, presented compactly by
f.
ug (2) = ur
buse (x)) (XIxxXx
(x)
Vul Vubuue (X+) ä. (X))7 x1(X) = Vx((1) (3)
(x)
where u base defines the base distribution phuse = p=d and Vu Rdd - Rddd denotes its Jacobian.
The core of Adjoint Matching the introduction of the (lean) adjoint ãe (X) E Rd-an estimator of the optimal solution wi (1z) that can be computed efficiently through backward integration from the terminal condition āj (X) = (X),,,, which is exactly the gradient of the terminal loss.
In this work, we focus specifically discrete generative modeling, heated research area that has gained sigmificant attention across both industry and academia (Touvron et al., 2023; Achiam et al., 2023; Kasneci et al., 2023), fueled by the transformative success of large language models (LLMs) in language understanding (Wei et al., 2022; Bai et al., 2023), code generation (Chen et al., 2021; Team al., 2023), and mathematical reasoning (Shao al., 2024; Nie et al., 2025). In all cases, the state spaces are represented by the product of a discrete finite set, d'diserete = (1,... M MJN. That is, each sample x = (x* , ,... ...) E d corresponds to sequence of N discrete variables, where a vocabulary list of size M.
each variable r' is the index of an element- also called
Having to consider a discrete state space necessitates family of discrete generative models. No- tably, most of which can be formulated as the Continuous- Markov Chain (CTMC, Anderson, 2012) a stochastic jump process on finite set, uniquely characterized by its transition rate,
p(X++Ar At y/ X, x) 5x(1)
( x) = lim Ar-0
At
which quantifies the infinitesimal transition probability from x to any y o x at time I. We denote (y) as the indicator function, equaling one if y and zero if otherwise. CTMC was initially introduced by discrete diffusion models (Austin et al., 2021: Campbell et al., 2022) and later shown to generalize masked and auto-regressive models (Sahoo al., 2024; Shi al., 2024).
Building on the practical success of AM in continuous domains, one question arises immediately:
Does variant of Adjoint Matching exist for discrete generative modeling?
More concretely, we are interested in whether a fixed-point equation analogous to (3) exists when the optimization problem in (1) is formulated on the discrete state space X' = disscette, where p" and prsse corresponds respectively to the distributions of fine-tuned and pretrained CTMC models. Such a discrete extension, if any, remains particularly challenging and has therefore been largely -rimmarrll due the non-differentiability of discrete state spaces and the drastic change in the class of distributions considered. Both prevent naive adaptation of (3) in discrete domains, as ((x) is no longer differentiable and huse now represents the rate of the base CTMC rather than the drift of an SDE. It is perhaps even questionable whether a discrete variant of AM exists at all.
To this end, we propose Discrete Adjoint Matching (DAM), discrete analogy of Adjoint Match- ing that solves the same optimization problem in (1) but for CTMC models on discrete state spaces. Unlike original A.M, which is based exclusively on control-theoretic framework (Bellman, 1966; Fleming & Rishel, 2012), we derive our DAM through a purely statistical standpoint that is more straightforward to follow, while still providing full control-theoretic analysis for interested readers. This necessitates lifting AM to a more abstract design space and, from which, specializing to other model classes such as CTMC. We demonstrate how Dynkin's 's formula (Piunovskiy & Zhang, 2020) enables such generalization by providing an estimator of function value from any stochastic pro- cess. Applying the formula to estimate the optimal CTMC rate u, (y. x)), defined formally in the next section, from CTMC processes yields discrete counterpart of adjoint (3)thhe discrete adjoint.
We validate DAM on fine-tuning diffusion-based L.Mss, particularly ---ts-S--nsttruct (Nie et al., 2025), on mathematical reasoning datasets, where fine-tuning methods have proven essential for improving performance (Shao et al., 2024; Zhao et al., 2025). To accommodate the extremely large
By construction, in (4) satisfies the rate condition: (v x) 2 for all y a x and Evex Me(y. 2) 0.
2

---

Published as a conference paper at ICLR 2026
discrete state spaces that can hinder naive implementation of discrete adjoint, exploit the masked structures embedded in all base distributions considered (Sahoo et al., 2024) and propose additional variance reduction techniques, which together enhance training stability and sampling efficiency.
In summary, we present the following contributions:
We propose Discrete Adjoint Matching (DAM), discrete counterpart of Adjoint Matching (Domingo-Enrich et al., 2025) that solves the same entropy-regularized optimization (I) but is specifically formulated for discrete state spaces with Continuous- Time Markov Chain (CTMC).
We derive DAM through a purely statistical perspective- by interpreting discrete adjoint as an estimator of the optimal colution-thhrele avoiding the convoluted control-theoretic derivation adopted original AM and providing a more general framework for adjoint-based estimators.
We address practical challenges of implementing DAM in extremely large discrete state spaces- typical of modern discrete generative modeling- by applying importance- weighting techniques for training stability and by exploring further simplification via masked diffusion models.
We validate the empirical performance of DAM in fine-tuning discrete generative models- such as diffusion-based LLMs-given prescribed reward functions on mathematical reasoning tasks.
DISCRETE ADJOINT MATCHING
We present Discrete Adjoint Matching (DAM), a computational method for learning a CTMC model p" that solves (1). It does by matching the parametrized rate ug (y, with estimator of the optimal rate u (y. z). obtained via the discrete adjoint, analogous to its continuous counterpart in (3). Such an interpretation facilitates purely statistical approach-to estimating andand allows us to bypass the seemingly convoluted SOC-based derivation, which we nevertheless present in Section 3 for interested readers. Conceptually, we seek objective of the following form:
(()) = Ex [DV (T("")),"" (+,x)()))))
(5)
where ID(:.)) is some matching function and u (- is an adjoint-based estimator of u* given the trajectory X. For instance, AM (Domingo-Enrich et al., 2025) considers (a,, b) ||a 612 and, as implied in (3), computes the adjoint- estimator by uz (x|X) := ut base (x) ( ) whereas DAM instantiates the matching objective (5) for CTMC models. All proofs are left in Appendix B.
2.1 DISCRETE ADJOINT AS ESTIMATOR OF THE OPTIMAL RATE
We begin by specifying the mathematical object of interest the optimal rate u² to the problem (1)- which the discrete adjoint is designed to estimate. Recall from (4) that the rate is uniquely deter- mined by the CTMC distribution via its first-order derivative. In the case of u*, since the optimal distribution p² has an analytic form (2), so does its optimal rate (Wang et al., 2025; Zhu al., 2025):
Vy + x, (g. u
bassse(y (y, x)e- Vo()) where Vz(x) -ol log Σ Pie (z|z)e s(*) -s(e) (6)
is the value function that quantifies the optimal cost-to-go (Bellman, 1954), u/use (y, 2) is the base rate that defines the base CTMC distribution pase and Pi Derivation of (6) is standard but somewhat convoluted; we refer to Appendix A for details.
base shorthands its conditional distribution.
With the analytic expression (6), estimating the optimal rate u² reduces to estimating the exponen- tial value difference, e Vv VV. as u hase is known a prior. Following the interpretation in (5), can design the discrete adjoint as an estimator of the exponential of the value difference, given a CTMC process X. This esproach-estimating a function value using a stochastic process-in precisely what is provided by Dynkin's 's formula (Piunovskiy & Zhang, 2020, Theorem 2.4.1):
Lemma 2.1 (Dynkin's 's formula). For a given function fi())
and a CTMC model p", it holds that
f(()) (XIXx=== ((x)) - -f' (0.-A-X-+)++ - M
dr
(7)
3

---

Published as a conference paper at ICLR 2026
Formally. Dynkin's formula applies any Feller process (Rogers & Williams, 2000), which gener- alizes both SDE and CTMC processes. On SDEs, yields the celebrated Itō lemma (Itó, 1951), a fundamental tool in the development of advanced diffusion models (Liu al., 2022; Skreta et al., 2024; Karczewski et al., 2024). In contrast, Lemma 2.1 specializes the result to CTMC, which, in its essence, enables estimation of any discrete-space function given any CTMC process.
We are now ready to present our main result. By applying Lemma 2.1 with (x y) e (V(())+V for every x * y and u := u² as the optimal rate, we obtain an estimator of the exponential value difference under the optimal distribution- the same spirit of how the continuous adjoint in (3) estimates the optimal SDE drift under p". This leads to the following result:
Theorem 2.2 (Discrete adjoint-adjoint system for CTMC). Let un (y,x) be the optimal rate the entropy-regularized problem in (1) CTMC models. Then, it holds for all y # x that
(1;xx) + J Σ ( (z
un (v, x) = -4 -. (U
āi(y:X) -xvo))(((i
(8)
where the discrete adjoint (y; X) Ry is unbiased estimator e -VV())+V(())) under p*.
Comparing (8) to (3), our proposed discrete adjoint shares a strikingly similar structure to its con- tinuous counterpart, yet differs in several profound respects. Most notably, the terminal discrete adjoint ă becomes the exponential of the terminal loss difference, as opposed to the gradient in (3). Next, better examine the differences between their dynamies, we can rewrite (8) in vector form:
ā. (X) = xx) + J. Utoae (X))d,, where ā (X) E IR [X] and Ubwe E x
respectively the discrete adjoint vector with (,; X) its y-element and the base rate matrix with base (z, y) as its (z. -entry. That is, the Jacobian matrix Vulase с Rdxdd in (3) is effectively replaced by Uhese E x Finally, the discrete adjoint perturbs the base rate in a multiplicative manner, huse E [a]. in contrast to the additive perturbation, u buse E[ā]] from continuous adjoint in (3).
Given Theorem 2.2, designing the matching function ((.,.) is straightforward application of the matching framework (Lipman al., 2024; Holderrieth al., 2024). For any Bregman divergence, Dv w) = of u) <(w) (u - w. Ju ó(w)), defined by a convex function Φ. it holds that
" (. x = arg min (x x=) [Do (ut (v, x)) u. (y. x x))]
(10)
We highlight the significance of these theoretical results, which provide a statistically rigorous foun- dation for deriving adjoint-base based estimators with any stochastic processes, including CTMC models.
2.2 IMPROVED TECHNIQUES IN PRACTICES
In this subsection, we address practical challenges of naively matching the discrete adjoint in (8) with (10), and provide improved techniques that give rises our final algorithm.
Improved discrete adjoint Despite the significance of Theorem 2.2 in characterizing the adjoint system for CTMC, the discrete adjoint in (8) and (9) is prohibitively expansive to compute, due to the need to sum over the entire state space E X' at every time step. For any reasonable discrete generative problem-with egg., vocabulary size M As 1000, context length N 100, and thus |X| N 10300_such computation becomes simply unaffordable. This necessitates further simplification for tractability, and we present one of which below:
Proposition 2.3 (Analytic discrete adjoint). The discrete adjoint in (8) admits analytic solution:
de dx X1) Pil (2 (2lg)e- (xx))
Σ base
(11)
The estimator (.)y:X1) in (11) arises from the fact that the discrete adjoint dynamics follow linear ordinary differential equation (ODE) as suggested in (n) hence admit analytic solu- tions. Although computing âz (y; X1) still involves summing over E X, it can be interpreted
4

---

Published as a conference paper at ICLR 2026
Algorithm Discrete Adjoint Matching (DAM)
Require: non-differentiblle terminal loss g(z),, parametrized rate uf 2), and base rate upase (y: x))
1: repeat
2: Sample from model {(x)) .x p", where stopgrad(u")
3: Update replay buffer B + B U {(x) )
4: for / in 1, ..2.....44 do
5: Sample time 1 ~ ((0 1J and state X; p (-1Xo, X1) where (Xo. Xi) B
6: Sample K model trajectories {x(() pd (-1x,)
7: Set (Y, Z) as the first and last jumps of one of the trajectory (
8: Compute discrete adjoint â, (Y, X2) àr xx] in (13)
9:
Take a gradient step (ODAM(() on the weighted matching objective
CDAM == P((1X DeKLL (Y, Xi), (ut ox,
(12)
10: end for
11:
until converge
probabilistically an expectation over the conditional distribution (-l: and estimated through Z pirc(1y)) which corresponds samples generated the base CTMC from X, y. Essentially, (y; X1) in (11) provides tractable solution of the original discrete adjoint ãr (y: X) in (8).
From Theorem 2.2, for (y X2) to serve as an unbiased estimator of the optimal rate u", we must sample X1 p from the optimal distribution- While one may follow original AM recipe by computing ā, (y; X1) with model samples X pd and letting the bias reduce over training, we opt for an alternative better suited for discrete domains: approximating the expectation of X1 ~ p" using model samples X1 ~ p" with importance weight. This leads the following estimator, which empirically exhibits much lower bias and variance.
its continuous variant (3) requires X ~pp
Proposition 2.4 (Importance-weiigl) discrete adjoint). Consider the following estimator
K
âz(u: 2, (x (x( = Pi.
(2(v)
-Z(
1
Pilc (x
(10 (x) (x -o(x(())
(13)
K
k=
where (x(x)) (12)) and Z ~ (-lv) are generated model samples. Then, âg is consistent
estimator of the discrete adjoint, ie., limg Elū,(y; 7, (x(x) })] "X, (ā X1))-
Comparing (13) to (11), the estimator à effectively estimates e e=((x) with single model sample generated from X, = y. and estimates (0())) with model samples from X, = using self- normalizing importance sampling (Neal, 2001). As these importance weights- the probability ratio of hase over - be easily estimated for CTMC (see Appendix A for details), we adopt the estimator (y; Z, (x(+)) in (13) the practical implementation of the discrete adjoint (e: X).
Matching with generalized KL Next, we discuss choosing the matching function (o(:)) from the family of Bregman divergence, following (10). While a naive fn-nrrm may be used, it known to neglect the necessary probabilistic structure in discrete domains, such as non-negativity (Lou al., 2023). more suitable alternative for discrete modeling the generalized KI. (gKI.):
D&KI. (u.w) w) Σ u(y, w(y, x) + w(y, 2)) log m)).x)
(14)
u(y. x)
which corresponds to setting ((a) = E, a(i) a(i)
Sampling model inputs (y,z) x) Having determined the matching function in (14) and established tractable adjoint estimator in (13), it remains to be discussed how to sample the model inputs (y, x). Intuitively, for given t, we wish to sample z = X, from high- probabbiyyyyl regions traversed by the optimal CTMC. This is commonly implemented by reciprocal projection X, p (-1Xa, X1)
5

---

Published as a conference paper at ICLR 2026
B sampled from a replay buffer (Shi al., 2023). Meanwhile, instead of sampling y uniformly as suggested in (14), we sample y from the model y ~ p (-(z) and debias the objective accordingly with P (u|z) This leads to the final matching objective presented in (12).
with X1)
Adaptation masked diffusion models While our DAM is designed for learning CTMC p" that solves the entropy-regee entropy-regularized optimization (1) given arbitrary base CTMC pruse in practice, most base processes are instantiated by masked diffusion models (Sahoo et al., 2024). These models employ simplified CTMC beginning from a fully masked state Xo
(m, ... m) with a base rate:
Vy + x, hase (v. r) == X base (x) ( (u/a), Qhine E Q
(15)
where hhuve (x) 2 denotes the jump rate with hase (z)
and X Jase (x) + oo, and Qmaak represents
the class of the probabilities where Qly(x) 2 0 if and only if unmasks exactly one token of x, e.g. (.... _,n-,,, y" x=+1 ....) with y* # xt m; otherwise, Q(y[z)= =. 0.
Simplification of the base rate u hase in (15) leads to several practical advantages. Most notably, substantially reduces the modeling complexity of Q((xx)) from O(MN) to ((MN) and en- able modern parametrization from LLMs (Vaswani al., 2017). Further, that Quae is time- independent implies a much more efficient sampling procedure via any-order unmasking (Shi al., 2024). Below, show that all these preferable features can be transferred to the optimal rate ut: Proposition 2.5 (Masked optimal rate). Given the base rate u base in (15), the optimal rate u² satisfies
Vy # x, (y x) At (x) Q C
(16)
Practically, Proposition 2.5 suggests that we can parametrize Q" (y)z) within the same model class of Qumask match Q" (r) with Quase (ylxr)(() æ)), then sample in the same any-order manner.
We highlight these discussions that lead to practical implementation of DAM in Algorithm 1.
THEORETICAL ANALYSIS
We highlight additional theoretical insights by re-deriving the discrete adjoint (r X) from a stochastic control perspective (Bellman, 1966; Fleming & Rishel, 2012), which more closely follows the original derivation in AM (Domingo-Enrich et al., 2025). All proofs are left in Appendix B.
3.1 FIXED-POINT EQUATION OF OPTIMAL RATE
We hegin by introducing new mathematical object [z;;; thethe expected accumulated cost under some CTMC distribution p" (X),, incurred from time t until the terminal, starting from the state æ:
Jr [x; u) = Expp((XXX==
buse Wt
Ue + log
basse ) (xxx)
Thy
(x
(17)
X
Since the integration in (17) corresponds to the KI. divergence between the CTMC pM and anddptn hase over the interval [t, 1] (see Appendix A for details), the cost-to-go Jr Jz(2; u) quantifies the objective in (1) that matters for optimizing u, Consequently, the optimal rate satisfies it arg min, (z u), and the value function-as the optimal cost-to-go functional-can be expressed by Vi(x)) = Je(2; u*). Below, we reveal this intertwined connection between and Je(z; u) with fixed fixed-point equation.
Lemma 3.1 (fixed-point equation of u²)) Let F be a mapping between the spaces of all possible CTMC rates that satisfy the rate condition, i.e., Footnote 2. Specifically, for all + x, we have that
F u(yzzz + u. hase (y, J((gru)+ + (m:u)
(18)
Then, the unique fixed point of the mapping is the optimal rate ui (y, x), i.e. F(u*).
Recall that M and N are respectively the vocabulary size and context length of the discrete space X.

---

Published as a conference paper at ICLR 2026
3.2 DISCRETE BASIC ADJOINT
The core of Lemma 3.1 is the observation that the optimal rate u² is the unique fixed-point of the relation ut (v,z) = buse z)) - (( As such, one can directly optimize for CTMC rate to fit this relation. This is precisely the motivation underlying the Basic Adjoint Matching (Domingo-Ean) Enniig et al., 2025), for which we provide a discrete analogue for CTMC models:
Theorem 3.2 (Discrete basic adjoint matching). For any Bregman divergence Da, the unique criti- cal point of the following matching objective is the optimal rate:
ui (y. x) = arg min (XIXX xxa)
[Do(ue(y. z), uh (y. x x)aza)a(/a (y: X, u)
(19)
where u stopgrad(u) and (y; X, u) is the discrete basic adjoint obeying backward ODE
d
(1 F(we)
ar(v: x, u) ar; X, uu)
Σ x
ur
base
+ uelog F(ue) (z, X1)
dt
Σ
uy
hase
ue + uelog F(us)
(2,
(20)
z
with the terminal condition a1 (1 X, u) = +(v)+
Intuitively, the discrete basic adjoint ax(;; X, u) provides an unbiased estimator of e (1 under the distribution p", which, upon choosing a proper matching function D3_ instantiates the mapping F(u) = Ep- [u/as (y. x)),,); X, X, = z]. In other words, the optimization in (19) can be understood as the variational form of the fixed-point equation, u" F(u"), in Lemma 3.1.
In practice, though, (y: X, u) is computationally intractable, as the dynamics in (20) involve eval- uation of the mapping F(u)-and hence Je(:; u) due to (18)-which what we wish estimate. This contrasts with the basic adjoint in continuous domains (Domingo-Enrich et al., 2025, Equation 30), which, despite computationally expansive, still remains tractable. Nevertheless, observing that all intractable terms in (20) vanish at u == u², the discrete basic adjoint on the optimal distribution (1; X, u²) becomes tractable and, as shown below, recovers the discrete adjoint ây az X).
Corollary 3.3 (Discrete adjoint). At u := u², the dynamics of discrete basic adjoint simplify to
di ur((: X; x,")) - base (2,1))a(Z; x,u") a1(y; X, u - +99)+
(21)
which coincide with the dynamics of discrete adjoint in (8). That is, by the uniqueness of ODE solutions, we have that at arly; X, u²)===== ãe (y; X).
Finally, the objective in (10) can be readily recovered from (19) by setting := u².
4 RELATED WORKS
regropy-r reward optimization. DDPO (Black et al., 2023) and DPoK (Fan et al., 2023a) consider inference in discrete-time diffusion models as a Markov Decision Process, and apply policy gradient approaches to (1). DRaFT (Clark et al., 2023) and AlignProp (Prabhudesai al., 2024) utilize reparameterization to do direct backpropagation in discrete-time. Adjoint Match- ing (AM, Domingo-Enrt Enrich et al., 2025) and ELEGANT (Uehara et al., 2024b) consider entropy- regularized optimization of continuous catiiuuuus-tmmeeee diffusion models, with path-wise KL regularization. Unlike AM, ELEGANT learns both an altered drift as well as altered initial noise distribution.
Fine-tuning discrete models. Fine-tuning for discrete diffusion models reviewed, albeit in dis- D1 (Zhao et al., 2025), Diffu-
crete time, in (Uehara et al., 2024a). Several Coder (Gong et al., 2025), DRAKES (Wang et al., 2025), and SEPO (Zekri & Boullé, 2025) policy gradient methods (Shao et al., 2024; Liu et al., 2025b) but many require approximations due to intractable likelihoods or non-differentiable rewards. In contrast, SVDD (Li et al., 2024), a value-based approach, estimates the value function via reward regression on model rollouts. while Qi (Zhou et al., 2025) learns parameterized reward distribution on rollouts. Rather than estimating the value function directly, our DAM utilizes estimator for the exponential value difference.
7

---

Published as a conference paper at ICLR 2026
(Proor) Optimal distr. (Targett)
(ours)
DAM (ablate w/
Checkerboare
Pinoshoe
Figure 1: Synthetic Examples. We compare the empirical distributions of X1 generated by the base model phose, the ground-truth optimal model pt, and four methods, including an ablation of DAM trained with the discrete adjoint in (11) instead of (13). DAM visually aligns most closely with pi.
First Jump
Second Jump
à. -=((ā
DAM (ours)
10f
DAM (ntlite w/ Fig 11)
Dclpa 1pm
1011
1014
0.2
SW33
D1
10
0.0
10
5000
10000
0
5000
10000
5000
10000
Iterations
Iterations
Iterations
Figure 2: Convergence to Optimal p' Pinwheel. Convergence of DKL(22 [p(") each jump de- scribed in (22), where DAM exhibits stable convergence compared to other methods (Ieft & middle). Our improved discrete adjoint in (13) exhibits lower bias and variance compared to (11) (right).
EXPERIMENTS
We validate the performance of DAM two benchmarks: 2D synthetic examples with analytically available optimal rates u*, and fine-tuning higher- dinher-dimensional reasoning datasets for diffusion- based LLMss. In all cases, consider masked-based CTMC models. Further details are in Appendix C.
5.1 SYNTHETIC EXAMPLES
synthetic examples discrete state space disceret {m, 1, 91 91 where each sample æ E X' represents a 2-dimensional index on a 91-by-91 grid, and hence the marginal distribution p(Xe) can be visualized 2D histogram. For masked-based CTMC models, the generative processes simplify to two distinct jumps from the mask state X, (m, m).
Setup We consider
(m, m)
first jump
(Ir;, m) or (m, 23)) second jump
(I., Le E (1, (1..... 91}
(22)
To mimic the fine-tuning setup, we design a terminal loss g(x) such that the optimal distribution pi (X1) x (X (Xi)e 9(X,) upweights or downweights certain regions, as shown in Figure 1. No- tably, the relatively low dimensionality of the problems enables exact computation of the optimal rate u² and its induced optimal distribution p", allowing us to compare convergence between our DAM and other baselines, specifically D1 (Zhao et al., 2025) and SVDD (Li al., 2024).
Results Figure presents the results on synthetic examples, where we compare the 2D histograms of samples X, ~ p(X1) generated by our DAM, its ablation, and other baselines. On both exam- ples, namely Checkerboard and Pinwheel, DAM remains visually closest to the optimal distribution, demonstrating its effectiveness in solving the reguaae-re optimization in (1). This is further confirmed quantitatively by Figure 2, where DAM converges stably to the optimal distribution for both jumps characterized in (22). In contrast, other fine-tuning methods such as D1 and SVDD tend to plateau without converging to p", likely due to the heavy approximation imposed during training. Notably, while both (11) and (13) provide theoretically sound estimators of the discrete adjoint, the latter exhibits significantly lower bias and variance. As such, adopt (13) for all remaining tasks.
8

---

Published as a conference paper at ICLR 2026
Table 1:
Performance on Mathematical Reasoning Tasks. We fine-tune the pretrained
--nstt----nstrucct model with generative sequence lengths of 128 and 256 using different meth-
ods. DAM consistently achieves higher test-time accuracy (Accu, %) and sub-rewards r. in (23) than DI on almost all setups. Best results are highlighted.
GSM8K
MATH500
Countdown
Sudoku
Method
Accu.
Teorect
Toomat
Accu.
Torreet
Tiomaa
Accu.
Pcomect
Tfomaa
Accu.
Sequenoe Length 128
LLaDA-8B-Instrur
68.6
1.37
1,46
28.8
0,58
0.72
34,8
0.31
0.10
11.5
0.11
with DI (2ho 2025)
75.6
1.51
1.74
31.2
0.62
0.73
43.8
0.39
0.10
23.8
0.24
with DAM (Ours)
75.7
1.52
1.86
32.6
0.65
0.73
60.2
0.54
0.10
89.2
0.89
Sequence Length 256
76.8
1.54
1.22
30.8
0.62
0.66
19.5
0.30
0.10
0.07
with DI (Zho al, 2025?*
79.8
1.60
1.63
37.2
0.74
0,70
31.3
0.30
0.10
12.9
0.13
with DAM (Ours)
79.9
1.60
1.82
36.4
0.73
0.92
55.5
0.499
0.10
88.1
0.88
Table 2: Test-time Generalization for Mathematical Reasoning Tasks. For each fine-tuned model in Table I, we report the test-time accuracy (%) at sequence lengths of 128, 256, and 512. DAM achieves on-par or higher accuracy than D1.
GSM8K
MATH500
Countdown
Sudoku
Method/Eval Seq Len
128
256
512
128
256
512
128
256
512
128
256
512
LLaDA-8B- 8n--nruttt
68.6
76.8
79.7
28.8
30.8
36.4
34.8
19.5
20.3
11.5
6.4
6.0
Train Seq Len
128
with DI (Zaoo et sl., 2025)
75.6
80.6
82.1
31.2
34.2
34.0
43.8
33.6
28.1
23.8
16.9
10.0
with DAM (Ours)
75.7
77.7
77.0
32.6
30.4
37.2
60.2
59.8
59.0
89.2
88.6
84.9
256
with DI (Zuo sl., 2025)
72.6
79.8
81.9
33.2
37.2
39.2
33.2
31,3
37.1
18.4
12.9
11.0
with DAM (Ours)
74.5
79.9
80.0
30.8
36.4
36.8
58.6
55.5
49.6
87.0
88.1
87.1
5.2 MATHEMATICAL REASONING TASKS
Setup We consider four standard mathematical reasoning tasks: GSM8K (Cobbe et al., 2021), 1), a dataset consisting of 8.5K grade-school math problems, MATH500 (Lightman et al., 2023), which contains 500 high-school competition math problems, Countdown, a combinatorial arithmetic game on given set of three numbers, and 4x4 Sudoku puzzles, which require filling grid with numbers to satisfy constraints. Since D1 (Zhao et al., 2025) represents the current state-of-the-an method for fine-tuning masked CTMC models these reasoning tasks, we adopt their experimental setups. (Nie et al., 2025), without supervised finetuning, as the
Specifically, we employ LLaDA-8B-Inste base CTMC model and consider the terminal loss g(-):
g(x) -r(x) where r(x) "comrea (1) Tiormat (x)
(23)
comprises two sub-reward functions, Teomeet and T'format that reward the correctness and formatting of the generated samples. Following D1, we evaluate the final checkpoint with zero-shot prompting and fine-tune separate
adopt the same decoding scheme from Nie et al. (2025). For each method, CTMC model on each task with generation lengths of 128 and 256, and evaluate each resulting model generation lengths of 128, 256, and 512. Further details are in Appendix C.2.
Results Tables and 2 summarizes the results on the reasoning tasks. Across almost all setups, except MATH500 with 256 sequence length, DAM achieves on-par higher test-time accuracy as well as higher scores on each sub-reward compared to D1, showcasing its potential for fine-tuning moderate-sized diffusion-basee LLMs. While DAM clearly outperforms DI on Countdown and Sudoku, the performance gaps narrows on GSM8K and MATH500, suggesting that the assumptions made by D1 may be reasonably effective for these tasks. Adapting such approximations for DAM- which is designed for more general poolemms-remiiss a promising future direction. Finally, Table compares samples generated by different models conditioned on the same unseen prompt, where the
model (middle column) successfully reasons and generates the correct answer.
Values taken from (7hao al., 2025).

---

Published as a conference paper at ICLR 2026
Table 3: Generated Samples on GSM8K (upper) and Countdown (bottom) from the base model and two fine-tuned models given the same unseen prompt. Correct and wrong answers are colored. More samples are in Appendix C.3.
Kenny is selling Pokemon cards to huy a ticket an amusement park, which costs $100. has a collection of cards and plans to sell them for $1.5 cach. He keeps 1/3 of them and gets to to the amusement park with $50 in spending cash. How many cards did he start with?
Prompt:
Answer: 150
Response:
LLaDA-8B SAAA--BB-Instuuu w/ DAM
w/ DI
<reasoning>>>
<reasoning>>
<rasoning
Kenny keeps 1/3 of his cards, so he sells 2/3 them. He selling 2/3 of his cards for $1.5 each, so he earns (2/3) * $1.5 = $10 from selling them. He also $50 spending cash, so he has $10
Kenny needs raise $100
Kenny needs $100 $100-$50- = $50. $1.5 per card, needs to sell $50/ / $1.5 = 33 cards. Since he keeps 1/3 of them, he sells 1/3 * 33 = 11 cards to make the necessary amount </reunnn/rassnnngg
$50 $150. He keeps 1/3 his cards, so he sells 2/3 of them.
Each card sells for $1.5, so he sells $150 S1.5 = 100 cards. Therefore, he started with 100 (2/3) = 150 cards.<<reasseeee </reasoning>>>>
+ $50 $60. Since the ticket
<answer> 33 </answer>
costs $100, he started with 60/2 30 cards. </reasoning>>
<answer> 150 </answer>
<answer>> 30 </answer>>
Prompt: Using only the numbers [67, 39, 70J. create an arithmetic expression that evaluates to exactly 98. You must use all numbers from the list, and each number must be used exactly once. You may use the operations +, -, *, and as needed.
Answer: 70 70-39) 39 67
Response: LLaDA-888 SB-Instruct
LLaDA 8B- Instruttt w/ DAM
LLaDA- PODAA-B--nnstrre w/ D1
<reasoning>>
<reasoning>>
<reasoning>>
To create an arithmetic expres- sion that evaluates to 98 using the numbers [67, 39, 70], we need to find the correct com- bination of operations and, if possible, parentheses. Given the constraints of the problem, we can start by considering the properties of the numbers and the operations operations avail- able.
To create an arithmetic expres- sion that evaluates to exactly 98 using the numbers [67, 39, 70], we need to find combination of the operations +, * and / that results in 98.
create an arithmetic expres- sion that evaluates to 98 using the numbers [67, 39, 70], we need to consider the operations of addition, subtraction. multi- plication, and division. We can start by trying to combine the numbers in a way that leads us to 98. One possible approach is to use the numbers and their oper- ations to close that sums up 98. We can also try using the and operations to get to 98. </reasoning>>
</reasoning>
<answer>> 70 39 + 67 </answer
</reasoning>
<answer>> 67*70 39 </answer>
< <answer>>
67439-770/70
</answer>>
Overall, these results highlight the effectiveness of DAM in solving entropy regulariee reward optimization yet potentially for a broader class of CTMC models, which may be beneficial for advanced CTMC models (Havasi et al., 2025).
6 CONCLUSION AND LIMITATIONS
We introduced Discrete Adjoint Matching (DAM), an adjoint-based fine-tuning method for dis- crete generative models formulated by CTMC. DAM is derived through a rigorous statistical frame- work that can be easily applied to other model classes. We demonstrate effectiveness of DAM in solving various cetropy-regularizeed optimization problems. Applying DAM to non-masked CTMC's presents an interesting future work.
10

---

Published as a conference paper at ICLR 2026
AUTHOR CONTRIBUTIONS & ACKNOWLEDGMENTS
The project was initiated and led by Guan. Guan was the primary author of the main paper, with assistance from Oswin on Section 5. Guan and Oswin co-developed the theory, co-designed the algorithm (with contributions from Ricky and Brian), and co-authored the Appendix. All experi- ments were conducted by Oswin. Chuchu provided valuable discussions during the early stages of project exploration. The authors would like to thank Jaemoo Choi for the helpful discussions and comments.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
William J Anderson.
Continuous- ttime Markov chains: An arappiicaiions--rietteddd approach. Springer Science & Business Media, 2012.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:
2021.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint WwXiv:2309.66600 2023.
Richard Bellman. The theory of dynamic programming. Technical report, Rand corp santa monica 1954.
Richard Bellman. Dynamic programming. science, 153(3731):34 37, 1966.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXi::2305. 2023.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisti- cians. Journal of the American statistical Association, 1126518):599-877
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-28277 2022.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374, 2021.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.1740 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21100 10. 1011110.44100 14168, 2021.
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky Q. Chen. Adjoint Matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In International Conference on Learning Representations (ICLR), 2025.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Systems, 08:79858-77885. 2023a.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for fine- tuning text-to-image diffusion models. In Thirth-seventh Conference on Neural Information Pro- cessing Systems (NeurIPS) 2023. Neural Information Processing Systems Foundation, 2023b.

---

Published as a conference paper at ICLR 2026
Wendell H Fleming and Raymond W Rishel. Deterministic and stochastic optimal control, vol- ume 1. Springer Science Business Media, 2012.
Shansan Gong, Ruixiang Zhang, Huangjie Zheng, Jiatao Gu, Navdeep Jaitly, Lingpeng Kong, and Yizhe Zhang. Diffucoder: Understanding and improving masked diffusion models for code gen- eration. arXiv preprint arXiv:2506. 20639, 2025.
Marton Havasi, Brian Karrer, Itai Gat, and Ricky TQ Chen. Edit flows: Flow matching with edit operations. arXiv preprint arXiv:2506.090//1 2025.
Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrinn -nrrgch, Anuroop Sriram, Brandon Wood, Daniel Levine. Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Motan-Horgg Liu, and Ricky T. Q Chen. Adjoint Sampling: Highly scalable diffusion samplers via Adjoint Matching. In International Conference on Machine Learning (ICMI., 2025.
Peter Holderrieth, Marton Havasi, Jason Yim, Neta Shaul, Itai Gat, Tommi Jaakkola, Brian Karrer, Ricky TQ Chen, and Yaron Lipman. Generator matching: Generative modeling with arbitrary markov processes. arXiv preprint arXiv:2410.20587, 2024.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 12):3, 2022.
Kiyosi Itô. On stochastic differential equations, volume 4. American Mathematical Soc., 1951.
Rafat Karczewski, Markus Heinonen, and Vikas Garg. Diffussonn models as cartoonists: The curious case of high density regions. arXiv preprint arXiv:2411.01293, 2024.
Enkelejda Kasneci, Kathrin Seßler, Stefan Küchemann, Maria Bannert, Daryna Dementieva, Frank Fischer, Urs Gasser, Georg Groh, Stephan Günnemann, Eyke Hüllermeier, et al. Chatgpt for good? opportunities and challenges of large language models for education. Learning and individual differences, 103: 102274, 2023.
Jun Hycong Kim, Seonghwan Kim, Seokhyun Moon, Hycongwoo Kim, Jeheon Woo, and Woo Youn Kim. Discrete Diffussion Schr Schr\"odinger Bridge Matching for Graph Transformation, October 2024. URL http: /ttt::://raaa /wrxivaiiivvv.////////eaa/aaaav. org/ abs / 24100. C1500. arXiv:2410.01500 [cs].
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representutions (ICLR), 2015.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: 1805 1805.0009, 2018
Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gokcen Eraslan, Surag Nair, Tommaso Bian- calani, Shuiwang Ji, Aviv Regev, Sergey Levine, and Masatoshi Uchara. Derivative- Guid- ance in Continuous and Discrete Diffusion Models with Soft Falue-Based Decoding, October 2024. URL http: http//aaxvvvv//v/////v//e../..av org/ abs / 2408 08252. arXiv:2408. .08252 [cs].
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023.
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen. David Lopez- Paz, Heli Ben- Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.066:
2024.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou. Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Guan-Homg I Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, and Ricky TQ Chen. Adjoint Schrödinger Bridge Sampler. In Advances in Neural Information Processing Systems (NeurIPS), 2025a.
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du Wee Sun Lee, and Min Lin. Understanding rl-zero-like training: critical perspective, 2025b. URL https: 20783.
12

---

Published as a conference paper at ICLR 2026
Ily Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations (ICLR), 2019.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint TrXiv:2310.66344 2023.
Radford M Neal. Annealed importance sampling. Statistics and computing, 1 11:125-139, 2001.
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992 2025.
Long Ouyang. Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models fol- low instructions with human feedback. Advances in neural information processing systems, 35:
27730-27744, 2022.
Jiay Pan, Junjie Zhang, Xingyao Wang, Lifan Yuan, Hao Peng, and Alane Suhr. https//giihubb.com/laii Panniaayriiaaainineeerr, 2025. Accessed: 2025-01-24.
Tinyzero.
Le- Tuyeet-NNii Pham, Dario Shariatian, Antonio Ocello, Giovanni Conforti, and Alain Oliviero Dur- mus. Discrete Markov probabilistic models: An improved discrete score score-based framework with sharp convergence bounds under minimal assumptions. Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu (eds.), Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pp. 199195-49255 PMLR, 13-19 Jul 2025. URL bttps::/proceedii:
mlr wwwppprssseeeee /v2 /ham22 5a. 55.mmm...................m..mm
Alexey Piunovskiy and Yi Zhang. Continuous-time markov decision processes. Probability Theory and Stochastic Modelling, 2020.
Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-to- image diffusion models with reward backpropagation, 2024. URL https://arxiv.org
abs/2310.. 03739.
Chris G Rogers and David Williams. Diffusions, Markov processes and martingales: Volume 2, Itô calculus, volume 2. Cambridge university press, 2000.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:1 30130-1110000 130184, 2024.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi. Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathemati- cal reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Kar- rer, Yaron Lipman, and Ricky T. Q. Chen. Flow Matching with General Discrete Paths: A Kinetic- Optimal Perspective, December 2024. URL http:/ arxiv. org abs///////aa /24/2 03487.
arXiv:2412.0348% [cs].
Jiaxin Shi, Kehang Han, Zhe Wang, Amaud Doucet, and Michalis Titsias. Simplified and general- ized masked diffusion for discrete data. Advances in neural information processing systems, 37: 103131-103167, 2024.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances Neural Information Processing Systems (NeurIPS), 2023.
Marta Skreta, Lazar Atanackovic, Avishek Joey Bose, Alexander Tong, and Kirill Neklyudov. The superposition of diffusion models using the it' o density estimator. arXiv preprint arXiv:2412.17762 2024.
13

---

Published as a conference paper at ICLR 2026
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew Dai, Anja Hauth, Katie Millican, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXtv:2312.11805, 2023.
Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. A generalized path integral control ap- proach to reinforcement learning. Journal of Machine Learning Research (JMLR), 11(Nov):
3137-3181, 2010.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1 13971, 2023.
Masatoshi Uehara, Yulai Zhao, Tommaso Biancalani, and Sergey Levine. Understanding rein- forcement learning-basee based fine-tuning of diffusion models: A tutorial and review. arXiv preprint arXiv:2407.13734 2024a.
Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous- time diffusion models as entropy-regularized control. arXiv preprint arXiv: 15194, 2024b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp.
2017.
Leandro Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. Trl: Transformer reinforcement learning, https: //ggtgttuuuu com http://ggggngggctt///// 2020.
Chenyu Wang, Masatoshi Uehara, Yichun He, Amy Wang, Avantika Lal, Tommi Jaakkola, Sergey Levine, Aviv Regev, Tommaso Biancalani, et al. Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design. In International Conference on Learning Representations (ICLR), 2025.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thoughh prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
Oussama Zekri and Nicolas Boullé. Fine-tuning discrete diffusion models with policy gradient 01 384.
methods, 2025. URL
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint arXiv:2504. 12216, 2025.
Jin Peng Zhou, Kaiwen Wang, Jonathan Chang, Zhaolin Gao, Nathan Kallus, Kilian Q. Weinberger, Kianté Brantley, and Wen Sun. gt: Provably optimal distributional for Ilm post-training, 2025. 20548.
URL https::/ /app:///azee /ttp//axvvvorro//e//..
Yuchen Zhu, Wei Guo, Jaemoo Choi, Guan-Horng Liu, Yongxin Chen, and Molei Tao. Mdns: Masked diffusion neural sampler via stochastic optimal control. arXiv preprint arXiv:2508. 10684, 2025.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint 1909 2019.
14

---

Published as a conference paper at ICLR 2026
A ADDITIONAL DISCUSSION & DERIVATION
A.1 ENTROPY-REGULARIZZH OPTIMIZATION
We provide additional clarification on how the memoryless condition in (24) leads to the analytic expression of the optimal solution in (2).
p hase (Xi (o)
memoryless hase (X1)
(24)
First, notice that the entropy--reuularized objective in (1) can be reformulated as:
DKL. @KI((p (x"(())((l)) basc (x)) ()) (xx)]
DKL (p(Xp) ( Jase (Xo)) + Ep(xo)
DKL (x Xx)) Ep(x [o(()))]
DKLL (p(Xo) (p base (Xo)) + Bx(xo)
(px base (xX (X(xo)e-o(x.)))
(25)
Minimizing (25) w.r.t. P yields
p X Xo) Z(Xo)
p²(oo) 1 buse (Xo) base X Xooe =(x1),
(26)
where Z(X)) is the normalization term defined by
((x) = / huse x Xn)e (xx) huse (X1
f (x (X1] -((x))xXX
(27)
That is, the optimal solution, p" (X) (X (Xo), follows the form
p² ('(X) = base (X) -XX, (Xx))
(28)
With the memoryless condition in (24), the normalization term becomes a constant, Z(Xo) = 2,
and thus (28) reduces to the analytic optimal solution presented in (2).
A.2 CONTINUOUS- TIME MARKOV CHAIN (CTMC)
Interpreting the transition rate its transition rate:
Recall from (4) that CTMC process is fully characterized by
p(X++Ae= 8|X, x) (x)
Uy (y, x) = lim
(29)
Ar+0
At
which by construction satisfies the rate conditions:
ue (y, x) 2 for y + x, and u u4 x))
Σ (1,
(30)
It is convenience to further define two mathematical objects from the transition rate ue (v. namely the jump rate X₃ (x) and the jump target probability Qe(y(z):
Qx(ra)) = ue(y,
(1, (x = Σ uely, -m,
Vy x x
(31)
Ae(x)
UIz
Intuitively, Ae(x)) = limase-o Prfjump occurs at Arl X quantifies the rate at which the jump occurs,
whereas Qu 1)) is the probability to jump from to y * z. Note that Eytr Qi ()) = 1 is by construction normalized. This leads to the following interpretation:
Ue (y)
(x
Q1(g)a)
Vy x,
u(xxx) -de(x)
(32)
how often to jump
where to jump
15

---

Published as a conference paper at ICLR 2026
The rate conditions (30) yield useful result that will become helpful in later derivation:
Lemma A.I. Given valid transition rate us (y,z) satisfying (30), it holds for all f(x) that
ue (4(y,))f Σ y
Σ
u. x) (f(w) f())
(33)
vFH
Proof.
Σ [u())
Σ (u(y,,)(()) x(())))())==
E Eu( - Em((s,))oo) V/
uf
The two most important equations that characterize the prob- ability evolution of CTMC processes are, respectively, the Kolmogorov Forward Equation (KFE), which describes the forward dynamics of the conditional distribution for any >> 5,
Kolmogorov equations for CTMC
OePt)o(yle)) = Σ uely, (2)
(34)
=E
and the Kolmogorov Backward Equation (KBE), which describes the backward dynamics:
-0.P:..
Σ
Ptja P1)s(ulz)u,zz,x)
(35)
=Ex
Note that can recover Continuity Equation- the dynamics of marginal distribution by taking expectation over (r) in KFE (34):
@ep((e) = 0 (f (nssella)ps())dd -- ve(s)))) 2)
(f (iinlcla))pa()dd = Σ us 2),z))((2)
KL divergence between CTMC Given two CTMC processes p" (X) and p (X) with the same initial distribution, their KL divergence be expressed with their transition rates (Kim al., 2024; Shaul et al., 2024):
DeL (X))) (X)) =
f( (x.)
base Ut
U1 + We log
Us ubase
(y, X,)
X.)
(36)
which, despite compactly expressed, remains difficult to interpret. The following lemma provides an alternative perspective by decomposing (36) into two components, similar to (32): one correspond- ing to the decision of how often to jump, and the other to the choice of where to jump.
Lemma A.2 (Interpretation of CTMC's KL divergence). Let pd and p buse be two CTMC distributions characterized respectively by their transition rates ue(y, x.) and whase (y. x)):
Vy # u(u,,x) Xi (") (ulx))
base ut
husse x) ( (x:)Q hase (y)a:)
where de(x)) and (x(x)) are their jump rates and jump target probabilities, c.f. (31). Then, the KL divergence in (36) can be decomposed into
DKL (p"((X (X)) S. (x) (X] (DKL(AM)|A) + DKL (Q(1Q))))) dt
(37)
where DKL (AH ([Ahase) and DKL 1QT11Qhuee are, respectively, standard KL divergences for exponen- tial and discrete distributions:
(x)
A base (x)
DEL (A) base (a:) = log
+
1,
p (3c)
Qu(()))
DKL. (Q" (-Iz) base (-1x)) = Σ Q
EQt(e (ulk)) log
16

---

Published as a conference paper at ICLR 2026
Proof. We first make the last term in (36) more explicit as
ue (v,xx
log A (z) (2))(())1z) Xamo (o) Qi (y|x)
Σ us (v, x:) log
Σ
Xi Q1)
u1 base
(y)
log 27
(x) log
+ +
Q (ul
base (x)
C (y(xz)
(x)
(x) log
+ xz) (Q9 (+xx))
(38)
base
(x)
Combing the first two terms in (38) with the leftover term in (36),
Ut
huse (v. (z,, + Xi (x) log
AN (X
(39)
buse (x)
yields the desired result.
Remark (Interpretation of cost-to-go Jr in (17)). From (36), the integration of Je(z;u) in (17) is
exactly the KL divergence between the CTMC's vd and pruse over the interval [t, 1}.
Je(x; u) = (rL( = 2) base (X = z)) + (X, (x( (x()
Estimating the probability ratio between CTMC's To estimate the probability ratio between CTMC's, we invoke their corresponding Radon-Nikodým derivative from, e.g., (Zhu et al., 2025);
X,, base (X,, Ut
log
Σ X
(u) ut
base (y. X())
+
Σ X:
log XX,
(40)
base
X
which remains tractable given a trajectory X ~ p². Note that when pd and pase are specialized to masked diffusion models with the following transition rates,
Vy + 2, xy. X (z) (ylz),
wuse (y. 2) (x)
the integration in (40) vanishes and the Radon-Nikodýnn derivative simplifies to the pog-robbbbylity ratio of between two discrete-time Markov processes:
p" (X)
(x,(xxx,- log Quase (X-)X
log base (X)
Σ
(41)
t:X XX,
A.3 ENTROPY-REGULARIZZH OPTIMIZATION FOR CTMC
Dynamics of value function Value function-as as the optimal cost-to-go functional-play a key role throughout the derivation of the discrete adjoint in Theorem 2.2. To ease later derivation, provide their dynamics in the following lemma.
Lemma A.3 (Dynamics of value function). Consider the retropp-regulariied optimization in (1). Define the value function as
Vi(x) = - -log Σ
base Pilt (2|z)e -a(e)
Vi(()) = -(x))
(42)
ZE
The dynamics of value function follow:
o.V((z) = Σ
ur
base
(2, x)e Ve())+V((x)
-v(())
Σ
u
basc bas0 (2 Vo(x)
(43)
Proof. Straightforward calculation yields
Σ. pep ( (zlzz)e-
0.V((x) = E. plase (2(zz)--gtt 9(=)
17

---

Published as a conference paper at ICLR 2026
Vi(a)
uhuse (v: 20 hase ( x(z)
by (35)
=e ve(at Σ
ut
base (y,, (1)
-ve-Vaa Velx)) - -v(x)) O.Ve()) -
Σ ui
base (y, x)) V.(u)
It is convenience to further rewrite the dynamics (z) in terms of the jump rate. Using the fact that uz (z = t (z, .e V((z))+V((z) for all + x, we can rewrite (43) by
ae (VV(()) Σ (2)
Σ aztl
ut (2, x) + un x)
That is, the dynamics of value function can be compactly presented jump rate differences:
V.V((z) - base (x)
(44)
Furthermore, the following results show that the dynamics the value function (43) can be rewritten to recover the same form the Caamilto--Jaobi-Bae Belblmn from Pham et al. (2025, Equation 85).
Corollary A.4. The dynamics of the value function (1) can be written as
a.V((z) = A (x))
Σ 2
base 2 (e)l))e-V(()++Vaall - -1)
(45)
Proof. Starting from (43):
0,V((x) = Σ
ur
base (2, (2,x)e" -VV(3))))
t (x, x)) +
Σ
buse
(2, xr V((=)+V((a))
uuaee (xr, x²) base
(ar, x.)
(2|x)e (2lx)ee
hase (2i, x)) uh
hase (cr, xr)
hasee (#|x)e- -v(o))--v(e)+vv(a)) -Vi(2)+Va(())
b (x)
Σ
Eqha 20000 (2(x)e (V())))-((())) -1) -1)
2
Optimal transition rafe We provide the derivation of the optimal transition rate u (g) x) (6).
The following lemma generalizes the initial-point conditional distribution P1 o in (26) to arbitrary time steps Prja: < t, from which the optimal transition rate can be easily obtained.
Lemma A.5 (Optimal conditional distribution). For any distinct time steps (s,t) such that < s < t < 1, the optimal conditional distribution to the entropy-regularized optimization in (1) satisfies
(y)x) = Pi|s (y)z)e -V.lx)+v()))
base
(46)
Proof. We first compute the optimal joint distribution p² y, X. z):
px y. X, .²)
Σ xX, = y, X, = xx)
18

---

Published as a conference paper at ICLR 2026
hase Xx y,, X, (0(x)) 2(Xo)
by (28)
E X
Pilt (X1 Xily)ptte
base (2/(2) p. Xo)e x(xx) -o0o 2(Xo)
by Markoviana of huse
Prls
base
(y/a
Σ X
Ps.0 (x, Xo)
(47)
Likewise, the optimal marginal distribution (x) can be obtained by
(x)
p px (X = x)
x
base b (X,X, 4 = =x)e - 2(Xo)
by (28)
x
Σ
hase
(X1 |2 ) hase (x, (o)ex-(Xxxxe loR 2(Xo)
by Markoviana of y bas
Pils
X
=e- Vz(x)
Ps.
base (z, Xo)
(48)
Xo
Combining Equations (47) and (48) yields the desired result.
Given Lemma A.5, we can derive the optimal transition rate u² by applying (29):
2 (1/, x>
O.PT e (u|z))
ae (phja buse (100))-w())+t0)))
by (46)
Or (patee (wlze) e- eViilo)+V.co) 1-=
+
Pils (ula:)e (vlaa)e-V(())+Vvtt)) (-Vi(1))
base
u have (y, x)e Vo())+V.(z) -x) - eVVa()))VV
by (34)
u. hane (y, V.(.))V.(z) (V.)))-V (v)
base Σ
u -v(0)))ve
by (43)
which recovers the same expression in (6), ie..
Vy x x, uf x) un huse (y, x)e V%((()
(49)
Note that, as expected, the optimal transition rate (z x) satisfies the rate conditions (30)
(1, 2 0 for all + x, and Σ u7 (2) =
(50)
Dynamics of value function w.r.f. optimal transition rate While it is well-known that the dy- namics of value function depend on the base transition rate base as shown in Lemma A.3, for CTMC processes, these dynamics can alternatively be reformulated in terms of the optimal rate u². This is readily hinted in (44), where V=)) = (x) - A (x) is simply the difference between the optimal and base jump rates. The following lemma formalizes this observation.
Lemma A.6 (Dynamics of value function w.r.t. u*). Let Vz(x) be the same value function in (42). The dynamics of (z(x) given the optimal transition rate in (49) follow:
0, (v)) = Σ ut (2, hile)-vaaa = a,e = Σ ut (2,
(51)
Proof.
a.V((o) = (7())) - 11
-u (I,, x)
Σ ef
u
base
(2,07
19

---

Published as a conference paper at ICLR 2026
2#2
= E Lu((e,z)ehttt-vit
which implies that Vie) = =e((=)),VV(()) -e vo(a)
Σ
-- x))
Lemma A.7. The denominator of the discrete-adjoint is an estimator for e evi(s) i.e.
e eu(()) = Eyt (Xxle))eeXXX)))
(52)
Proof. Invoking Dynkin's formula Lemma 2.1 with fe(x)= = =VVaa) and u = u², the integrand
simplifies to
(O.)((X)) + Σ (z, X,)s((z) = 0re -Vel)) + Euiiz X.)e Vi(e)
O.
by (51)
Since Vi(x) = -(() and hence fi(x) = eg(u), we have deduced that
e vi(x)) = (x((XX)))==)) eo(Xi)
(53)
A.4 BIAS AND VARIANCE COMPARISON BETWEEN (11) AND (13) ON TOY EXAMPLE
We now quantify the difference in bias and variance between the original discrete adjoint and (11)
and the improved discrete adjoint using importance weighting (13).
We consider a sequence length of N = and a vocabulary size of M = 3 with X = (m, Zeorrect; 2wrong] ). We consider an indicator function loss that penalizes wrong answers, i.e.,
g(z) 50K {- = Zeorrect
(54)
Let the base distribution Pile be
0
z = Acorrest:
Pilc (2)
(55)
0,
Zwrong
Then, the optimal distribution pije is
plase (2) (2|m)) (*#)
(z
(56)
Σ z' phoso (z' (z'mm)e- -m)))
50
50 + (1 0)
Zeorrect
(57)
1 0
0.
+ (1 0)
Zwrong
If 0 is not too small, then pip (Zcomecr mm) x 1.
We consider just the "denominator" part of (11) and (13), as the "numerator" is the same in both cases. In our setting, the true value of the denominator is
F (xx,(x//X==m/((00X1. = Pije (ccoreaa m)e -((ome)) + Pijo ( (Zwring/ mm) g(Xwe))))
(58)
1 0e + (1 (1-0) 0)*
(59)
Let the current model distribution pNje (z /m) be
(2|m) = ④2
2 Zeonrcts
(60)
0
2 Zwrong
20

---

Published as a conference paper at ICLR 2026
Bias and Variance of Original Discrete Adjoint Let X ...... x pi (Xi = m) be N i..d. samples from the model distribution pie We follow the original AM recipe and compute (11)
using these model samples. This gives us the following N-sample Monte-Carlo estimator for the denominator of (11):
E (XIx(X/ Jco(xid] ~ E P (X,]X, [c
(61)
N
eo(x!")
(62)
We first compute the mean
N N
E
ea(x{()
Epp (X, /m) (X1/m) (e]()((11))
1=1
65 -50 50 + (1 (1-
and variance
N
Var
Σ
co(x(1)
N pi. Va (x1 [e9(Xi)] (xilme/eXXi)
1
Var
N
N (xp (X, m) (e220xxx)- - -Epy, (x, (m) (ee(X.))e)
N (de -100 + +(1-0) 6) - 50 + (1 (1-
1
N Φ(1
(c 550 - 1)2.
The bias is then
N
Bias of (11) = E
Epij (X1 (x.==) (en(X.)]
N
i=1
1
( 50
(1
11-() の))
0e50 + (1 8)
1
2 (1 ②)
50
2
(63)
and the variance can be approximated as
1
co(x(0)
1
Var
N d(1 -)((c - 11)
N
N 6(1 - 6)(e 05002
(64)
where we have used the approximations that de 50 + (1 6) x 1 1 & and 0c50 +(10+(110
which hold when @ is not too big and 0 is not too small.
Bias and Variance of Importence-weiighted Discrete Adjoint The denominator part of (13)
takes the form
(X x m) c-o(xx)) 1)
base (X, Xx m)
-1
Epg, (x)
(Eng, (x, [mxxx
1
w(x("))
Sx
21

---

Published as a conference paper at ICLR 2026
where we have defined for convenience
((X1) = Pife (X = m)
Pije (Xil m) e e-((x))
(65)
Evaluating the mean and variance of SN first,
[[[SN] - Epg (x/)/XX
= E B (X1 (X|mm (xx())
= 00 + (1 0).
and
Var [SN]
Varpi, (X11m) (m) (X1))
N
1
(E (Xi,, (X) lo (xx) Epg (x (m)
N.
( (X1|X, = m)
(B (BetreXiim)
6 20((31)
R (x (=-0()))1)
Pije (X1|X, m)
1
(0_
0 100
+ (1 0)
1 1- 0
(0c + -0))))
N
Φ
1- の
1
0 e 100
- 1) + (1 0)²
1- - 1) - -20(1 - -)e 50
N
Then, using the Delta method, we can approximate the mean and variance of S as
B|S ~ E[SN]- + Var[SN]
1 02e000 - -1
()3 -1)
20 - 0)
50 (1 0) +
N
50 + (1-0))s
and
Var[SN]
Var[S
E[SN]
0 200 1) + (1 0)2( T - 1) 20 0)
N
(0e (1 0))
The bias is then
Bias of (13) = E[S 1 - -Ep (X1]Xx ( (c9(1)))
1 02:100 1 1) + (1 - () 1 -1) 20(1 - -0
1-6
N
(0e50 + - (1-0))3
1 020000
100
1)
22
N 83 (3553
1 1
e 50
(66)
N
d
while we can approximate the variance as
1 02 200 1
1) + (1 - 003(11) 1
1-4 -1) - 20(1 - 50
Var[SN]
N
(0e50 + (1-8))
1
0 10000
x -1)
N
Bd ((600))
1 1 - Φb
N 20 (e -50j3
(67)
Comparing the bias (63) and (66), we that the bias of sepportanee-wiightee discrete adjoint is much smaller than that of the orignal discrete adjoint since the bias from (66) is scaled by e -50 For the variances, both variances (64) and (67) are similar. This aligns with empirical results in Figure 2.
22

---

Published as a conference paper at ICLR 2026
B PROOFS
Lemma 2.1 (Dynkin's formula). For a given function ft(x) and a CTMC model p", it holds that
ur (v, (u, [(xx)) - S: (0. +x
f((x) - Bp-(XiX.-))
dr (7)
y
Proof. Dynkin's Formula (Piunovskiy & Zhang, 2020, Theorem 2.4.1) states that given a function f(t,x) f2(),, it holds that
Ex = + Ex [[ (AA((,, X,)dt
(68)
where A is the infinitesimal generator of the Feller process X (X1)) defined
Ex (X,, X, = x] fz(
(Af)(s, lim rie
(69)
For CTMC X ~ p" with the rate us (v. 2.), we have that (Af)(s,
[x( = -
-lm
tLe
t-8 (t,(((yyl))) (frly) - fx(x))
-lim
c/e
t-s (u)()))) (fely) fv))
lim (()))) fa(x))
-lim
tLe
t-8 Σ, Pt (Fely) fa(1))
t-s
lim
Σ P1jo((l2)ff(()
Σ (p)o(ul=))fs( (70)
t/a
dt
The first term simplifies
Σn (U)a (fe(y) - fa(u))
lim
Σ (v) 0zf((z)
ALA
t
The second term simplifies as
di Σ (el = Σ (v.(xx))..)))
The third term simplifies as
d di
Σ (() [x(()) Σ u,(y,,))
"1
Combined, (70) simplifies as
(Af(s,)) = + Σ u (z))
(71)
Rearranging (68) with the above result (71) leads to the desired result:
fo(Xo) Exxc Xx) [1(x1)) f(' (o.Ffixi))+ E uelu X.)fe()))))))) de
23

---

Published as a conference paper at ICLR 2026
Theorem 2.2 (Discrete adjoint adjoint system for CTMC). Let ui ())z) be the optimal rate the eutropy-regularized problem in for CTMC models. Then, it holds for all y + x that
S. Σ
u (y,, .²) u4 -. Ep.r) (y:
" tuve (z
ā1(y:X) = o00)+
(8)
where the discrete adjoint ä, (y; X) E Ryo is an unbiased estimator of -V(())+((()) under p".
Proof. Invoking Lemma 2.1 with ft(1; y) == e Va(o)) and u == u², the integrand simplifies to
(X.)))(Xi)+ + Σ u (()
ex(x))0 -Vi(y) + e vv(w
(aehiix)) + E (z, (i)e
evi(Xe) -V.((
by (51)
= - u, y)e (V(()+V((X))
by (43)
Since (i( g(x) and hence fi (x;,)) - we have deduced that
f.
- (V(())V((()) = (x(xx== e (-ll)to(xi))
Σ base (2, (2)5))--((
dr (72)
Given optimal trajectory X p' x) starting from at time t, estimate the expectation by solving the following backward ODE:
acly; x) dt
d
Σ
ur hase (z. v(())
(1 =e -
(73)
This implies that
Valx)++V<(w) Fxp
( X)]
by (72)
= Ex
/1 (y: X) 1. Σ
uy
hase
(2)
(1 X) J."
(x
Σ
u
base (2, (2.)) [ xi]
(72)
Ex a1 X) J.
u, //-
Σ
base
(z, ", (y; X)
(74)
which implies that the following ODE is equivalent to the one in (73) in expectation:
d iu: X) = Σ buse (z, (2; X):
ãn (y; X) = e -(()++((x))
(75)
dt
Finally, substituting (74) to (49) yield the desired result.
Proposition 2.3 (Analytic discrete adjoint). The discrete adjoint in (8) admits analytic solution:
(x X1) = Σ
Pilp (zly)e (2))))
base
(11)
24

---

Published as a conference paper at ICLR 2026
Proof. Rewrite Kolmogorov Backward Equation (35) in matrix form:
d
exp (/, Ut
1 exp ( Uboeed
Phine = U.ueee Phae dt
1/r
(76)
where U/pas, phase E | are the matrices with u (z, y) and phee (=, v) as their (z, --ntriess respectively. Likewise, rewrite the dynamics of discrete adjoint (75) in vector form:
a (x U (x)
(77)
where ā. (X) E g| is the discrete adjoint vector with ä, (y; X) as its y-element. Since (77) is linear ODE, it admits analytic solution:
( dr) ā1 (x(X) Piiee @1 (X)
(x) = exp
(78)
where the last equality is due (76). Expanding (78) for each element y yields the desired result.
Proposition 2.4 (Importance--eiighedd discrete adjoint). Consider the following estimator
K Pijc (x (x
(Zlv) arly: 2, = (2 C 9(2)
1
(x) e -=(x("))
(13)
K
P
where (x(")) ~ 111)((11)) and Z ~ (il)(-l) are generated model samples. Then, âg is a consistent
estimator of the discrete adjoint, i.e., limgo-so E[ār(y; 2, {x })] (X (v; X1)}.
Proof. Given the analytic discrete adjoint in (11), we can decompose the expectation of the discrete
adjoint into two separate expectations:
(X,1X_=)) ( X1)] (x
Σn
Σ
Pile (zlg)e-=(=) ((x))
by (11)
9(2)
E (x: [
= -E Eppen (Zws)
Both expectations can be estimated using model samples X p Specifically, applying importance sampling to the first expectation yields
e -122
Epp (2lv) ((lv)
Pijt ( (4p()
-(2)
(79)
As for the second expectation, we apply self- nelf-normalized importance sampling
w(X (k) ,x) co(x)())
E Eppss(2l)) [co(xu]
lim E K Pe (1)
(80)
w (x (i) xx)
with the unnormalized importance weight w(z,x): == plifies the estimator in (80) to
Pi (=|=))
[ -
(=
This sim-
pi (=|z)
no (*)
pr (x e ((X("
w (xx x²) E, w(x
((x() (x((())))
(5) ,z)
p
p'(X(""
Dk=l y (x(()))))
ΣR
(=(
K (x (n) - (({()))
f=1 p=(x("))1)
25

---

Published as a conference paper at ICLR 2026
(81)
R L gd (x
p=(Xp"
where we use the fact that the numerator has expectation 1, ie..
px (x)) p=((x((
Pij ((x]")a)) Ha)
Combining (79) and (81) yields the desired result.
Proposition 2.5 (Masked optimal rate). Given the base rate u hase in (15), the optimal rate u² satisfies
Yy * 2, (y, 2)) = A (x) (x) (ulz), Q Q
(16)
Proof. This proof follows from the fact that the base rate hase defined by (15) reduces its corre- sponding CTMC to tiscrete-time Markov process in which exactly N jumps occur over the time interval from 0 to 1, with each jump unmasking one of the mask tokens in Xg (m,... m). This interpretation suggests that it suffices to define N marginal distributions (gn)n where each (.) corresponds to distribution of samples which exactly tokens are unmasked, and the conditional distribution becomes time-invariant; that is,
hase Pils (y]z) = base (p(s))<( (p(o)lo(a)((la)
(82)
where we introduce the function &0 X + Zzu that counts the number of unmasked tokens r.
With cp)), we can rewrite the base jump target probability buse as
huase (3) (z) 4p(=)+ (|x()) (y)x)
(83)
With the aforementioned reasoning, the value function defined in (42) becomes time-invariant:
buse Vi(x) - -log Σ
Pile (2|z)e
log
base (N) (z|z)e o(=) 0(=)) V(x)
(84)
and negative exponential transform, e (x)) exhibits dynamic programming structure:
e -V()
base IN (xl)p(z)) (2 (2(x))=)(())
by (84)
have ΣΣ
(N(1((2)+1 (219) buse
(K(()) (v) (vlzze (1)))))))))) 9(=)
Σ
Qusse (y/xx)
(N (zly)e
base
by (83)
y
Σ
Q base (y ()lr)ee Vvv)
(85)
This implies that the optimal rate preserves the same jump rate Xhase(r) the base rate and that the the optimal jump target probability is also time-invariant, i.e.,
Vy # x, ui (y, x) = buse (2, z)e V())+V(z) A (I) (y)a)e v(x)+(())
(vle))
Since the optimal jump target probability Q' (yy .:) is reweighted of the base jump target probability Quase E Quaak (y( > 0 if and only if y unmasks exactly one token of a. Furthermore, (v(aar
normalized by construction,
Σ (r(()) = Σ Q (y|z)e V()++(())
Σy Q base (ylzr)e
V(rr)
by (85)
That is, Q* E Quask and conclude the proof.
26

---

Published as a conference paper at ICLR 2026
Lemma 3.1 (Fixed-point equation of )²). Let F be mapping between the spaces of all possible CTMC rates that satisfy the rate condition, i.e., Footnote 2. Specifically, for all & I, we have that
F u4yy,2) (y, ()lg,)))eeeeleeeee ++(ar)++++++) Jelaz))
(18)
Then, the unique fixed point of the mapping F. is the optimal rate uf (y, i.e, u² = F(u").
Proof. Recall the cost-to-go functional Jt u) defined w.r.t. a CTMC with its rate ue (y. x):
J.
ux) = XX.=x)
buse W4
uT + log ubose
Ur
(v, (,) + +x
(86)
We can rewrite Jr dynamic programming form:
J) = E
Mo 147 + ur log
base
Ur whase
(v. dd xx= x
VØX.
E [EJ++A (X++A(H) Xx = xx]
E
[17 U4 + Ue log xxx
(v, No,x,)dr
Xx
At
ElJetse(eetrriil
Xx 2) Je(2:e)
At
Taking At - 0 yields
= (A
Σ ygpler
base ut
u4 + us log
Ur
(u, x)
by (69)
-0,J(2;;))
Σ ue ue u)
Σ
ut
hase
Ue + log
hase
(v,
by (71)
-O.J(xx;u) Σ (1) (Je(V: - [)); +
Σ
base
us + ue log
buso (y.x)
ut
by (33)
-0,J(2z;u) Σ
< Uuaso ue + ue log
Ue F(u+)
(y,z)
(87)
vf
Now, consider CTMC rate ū, x) satisfying the self-consistency relation
Vy # xr ūe (y, hase (y, x (e Ja(sv))) ū, F(ū))
(88)
One can check that when u = ū, the dynamics cost-to-go Je(2; u) in (87) reduces to the ones of the value function in (43),
- (x; ū)
Σ v$
W base it + ūi log Ut
(y,
F(ūt))
>
base (M, x)) - (v)
by (88)
ut
base (Ix, x²)
Σ huse (y, z)e (y.) 51
by (30) and (88)
Σ base
ut (y. x)e - (e
That is, have shown that any policy ū satisfying the self-consistency relation in (88) has its cost- to-go Je (z; ū) corresponding to the value function Vz(x). This implies that = u² the optimal rate, and, by uniqueness, is the unique fixed point of (88).
27

---

Published as a conference paper at ICLR 2026
Theorem 3.2 (Discrete basic adjoint matching). For any Bregman divergence De, the unique criti- cal point of the following matching objective is the optimal rate:
u (y, x) = arg min (X [Dslully, x), uh M (u, ()r (y; X, ū)]
(19)
where u = stopgrad(u) and a (v: X. u) is the discrete basic adjoint obeying backward ODE
(vy X, u) arly; u2
Σ
uk
base
21 F(ue)
+ uelog F(ue)
Up
(2, X1)
X
Σ
u4
ur u,log F(ue)
(2,
(20)
zz
with the terminal condition (y;; X, u) = - +(x)))
Proof. Invoking Lemma 2.1 with fo(1;; u) == -e(leaa)+
e e-dl(ua))+hh(e))) = (x(xx (-0(()+0(x1)) S. (0re- (0.e-01(er))+000000)
uo(2, Xi)e -(oon)))))e))))tlll
dr (89)
Expanding the first term of the integrand in (89) yields
0,e (X(
=e =e (XXss)
(-O,Je(Biu) + O.J((Xeiu))
Σ
buse u4
u + uelog F(ue)
(z, v)
Σ
base
ut + u,log F(we)
(z, (i) Xi)
by (87)
X
Likewise, expanding the second term of the integrand in (89) leads to
Σ uelz, Xz)e (ibe)
Σ 11 X,)e (1),,xX))))--((::tt- dA
-(x =
Σ
ue(z, Xz) (e ((--(XKi))))tttet ++(xx:))+th(Xxii))
by (33)
x*X.
base
e (xX
Σ
Ut
F(ue)) ue) (c)
by (18)
Combined, the integrand in (89) can be obtained by
- J. +
Σ
Ue (2, Xi)e (err)) +
Σ
buse u4 + u,log F(us)) ) (2,
We
Σ
Ut
base
1
24
+ velog F(ue) (2,
Us
Flus))
X
28

---

Published as a conference paper at ICLR 2026
=c -o Gey,
where we shorthand
Ge(v, (x = Σ
base ut
ut + us uglog
F(ue) )(m)) (2,
ue
Σ x
ut
hase
U4 F(us)
Ut velog F(ue)
(z, (2,
(90)
That is, we can reformulate (89) by
+(x)+ - L.¹ (. g, (v, xz;u)dr
+
Consequently, given a model trajectory X p"(-Xx, x) starting from x time t, one can
estimate the expectation by solving the following backward ODE:
d
(y X, - (i dt
a1(y; (())+
(91)
This implies that e (a(s)))
(y: X, u)]
(92)
-xpx(XXXe=
a1 X, (.))
c (-b-buw)))th:::::ggg gr x;: u)dr
f. _x-(x(x-) [ur x x.,). u]] Gr(u. X=; u)dr
-Xp((x(x==))
a1 (1; X 11)
[( a-(y: X, u) X+: u)
=x
a1 (y: X, u)
That is, the following ODE is equivalent to the one in (91) in expectation:
d ar x, 2) (y: uy Xx; x)
a1 X, u) -(x)
(93)
dt
which readily recovers the dynamics in (20).
Finally, let ū, x) be the solution to the matching objective (19) for all y + x,
Te (v. Φ) = arg min (xx(X (oluuy x), ut huve (v. x x))x))(0 (y: X u)]
Straightforward upplication of the matching framework (Lipman al., 2024) suggests that
Üe (y, .) xxpx(X(Xe [ub (y, x)a (vi X x.u
= ut
buse (y, x) d.
by (92)
Invoking Lemma 3.1, we deduce that ū = n² is the optimal rate.
Corollary 3.3 (Discrete adjoint). At u u², the dynamics of discrete basic adjoint simplify to di 0 X, 11)) = Σ base (2, X,",, a X,u;"""')) - (xx))
(21)
which coincide with the dynamics of discrete adjoint in (8). That is, by the uniqueness of ODE solutions, we have that ar arly; X, .u')) ãe (y; X).
Proof. Notice that (e x; u) defined in (90) greatly simplifies when u == u²,
Ge u²)
Σ
(u. base
(z, u (2,y))
29

---

Published as a conference paper at ICLR 2026
Hence, the dynamics of discrete basic adjoint (93) reduce to
dt ce X, u²) ar(y; X, u (y, =a X, v
2;
hase (2,2) - -2,2)- v))
U
ayy X, yul bass (y) y)
(2, (x (2.;)a((3;XXx))
Σ
241
Σ base ut (2,))a,(2; uu))
where the third equality is due to the fact that for all 2 #
(x(xx====) [u7 y) (y: X, u u]
" (z, -=
by (92)
= uque (z, 8) e Je(zu" (tot"="))
by (49)
ut (z, y)e Jel=:u )+(((z;"*)
= (x|xx==))
have (z, y)) (2; x,))
by (92)
That is, there exists detailed balance between ut (2, (2); X, u²) and (2))))))) (2; u*)
w.r.t optimal distribution, i.e.,
[u] (M, (x((X())=x)
x. u²)] (x [u, base (2, x,")]
(94)
C EXPERIMENT DETAILS
C.I SYNTHETIC EXAMPLES
Implementation For both DAM and DI, parametrize ul (y), as the logits of tensor of size 92 x 92 x x2. The first two dimensions correspond to all possible æ E X Xdiscréie = {m, l.... 91] the third dimension corresponds to all possible y E uiiiceett conditioned on given x, and the last dimension corresponds to the two distinct jumps characterized in (22). For DAM, we use K 16 samples to estimate the discrete adjoint à in (13) for Checkerboard, and K 64 samples for Pinwheel. Since these problems are fow-dimensional, we always use on-policy samples to compute à, that we do not maintain a replay buffer B for synthetic examples. Since DI's and SVDD's codebases are specialized respectively to reasoning and molecular applications, we implement their methods ourselves for the synthetic examples. All methods are trained on CPU using Adam (Kingma & Ba, 2015) with the learning rate 1 x 10-2
Reward Functions For Checkerboard, we consider the reward function that upweights the diag- onal blocks by 4.6, superdiagonal and subdiagonal blocks by 4.0, and other off-diagona blocks by 3.4. For Pinwheel, we reweight the three circles respectively by 0.5, 1, and -20.
C.2 MATHEMATICAL REASONING TASKS
Inference We use the same decoding strategy in Zhao et al. (2025). To decode sequence of N tokens, we use N/2 denoising steps and unmask tokens each step. They decode from left to right in blocks using the semi-autores eumi-eutoreeressivve decoding strategy (Nie et al., 2025), which they found to yield slightly better performance. Specifically, the sequence is divided into blocks of length 32. In each step, the 2 tokens with the highest confidence within the current block are unmasked. Once all the tokens the current block are unmasked, we move to the next block of 32 tokens.
their official implementation, 4 which uses the TRL library (von Werra et al., 2020). We build our DAM implementation on top of the same codebases. Both fine- tuning methods use Low-Rank Adaptation (LoRA) (Hu al., 2022) with rank of 128 and
Implementation For D1,
hhttps: thgithhubb... dom/d1 h
reasoning/dl
30

---

Published as a conference paper at ICLR 2026
scaling factor 64. For all four tasks, DAM uses K 12 samples to estimate the discrete adjoint ag in (13).
Training All methods are trained on 8 Nvidia H100 GPUs with a sequence length of 128 to- kens, batch size of 12 per GPU, and 2 gradient accumulation steps. We use the Adam optimizer (Loshchilov & Hutter, 2019) with parameters 31 0.9, B 0.99, weight decay of 0.1, learning rate x 10-0, and gradient clipping at 0.2. For both methods, we fine-tune for 3000, 3500, and 4000 steps (number of gradient updates) on GSM8K and MATH500, respectively.
Reward Functions We use the same reward functions as in Zhao et al. (2025), which we describe below for completeness.
. GSM8K We conduct RL on the training split of the GSM8K dataset" and evaluate on the test split. We employ composite reward function consisting of five components following the unsloth implementation of reward functions:
XML Structure Reward: Rewards proper formatting with reasoning and answer
tags. +0.125 for each correctly placed opening and closing tag. A small penalty of -0.001 for each extraneous token after the closing tag.
Soft Format Reward: Reward of 0.5 for responses matching the pattern:
<reasoning>... (content) ... </ //rreeeniiinng.
<answer>... (content)... . ... <. </anser>>
Strict Format Reward: Reward of 0.5 for adhering to the exact prescribed format with appropriate line breaks.
- -Integer Answer Reward: Reward of 0.5 if the extracted answer a valid integer.
Correctness Reward: Reward of 2.0 if the extracted answer exactly matches the ground truth.
In Table 2, the formatting reward Fiorman (x) corresponds to the sum of the first three afore- mentioned rewards.
MATH500 We train on the train split of the MATH dataset.? ? Like GSMSK, we employ composite reward function comprising
Format Reward: We reward the presence of <answer></answee tags and
\boxed, as follows: 1.00 if answer tags are present with \boxed inside, 0.75 if answer tags are present without \boxed in them, 0.50 if answer tags are not present, but |boxed is present, and 0.25 if neither answer tags, nor \boxed is present
Correctness Reward: We reward 2.0 if the correct answer
\boxed
Countdown We train on the training split of the dataser from the TinyZero project (Pan et al., 2025), restricting to instances that use only three numbers. And we evaluate on 256 synthetically generated countdown questions with 3 numbers. We implement a reward function that checks if an arithmetic expression constructed from given numbers reaches a target value:
Formatting Reward: 0.1 points when the equation uses the right numbers but doesn't 't reach the target
Correctness Reward: 1.0 point when the equation equals the target and uses exactly the available numbers
Sudoku We use the same generated dataset of 4x4 Sudoku examples as (Zhao et al., 2025), keeping the last 500 examples for evaluation and training on the rest. We employ a reward function that rewards filling in cells correctly.
Correctness Reward: The reward computed as the average number of correctly filled in cells.
http:////ggggggeeeeeaaeeae////rp 3to4
31

---

Published as a conference paper at ICLR 2026
C.3 ADDITIONAL EXPERIMENTS
Checkerboard
Pinwheel
10a
Km4
K-16
Dalp Ip
K=12
K=64
10-3
Kw128
10
10
10-3
5000
100000
15000
20000
5000
10000
15000
20000
Figure 3: Varying K on Synthetic Examples. We investigate the effect of varying the number of samples K on the synthetic examples. Larger values of K result faster convergence from the lower bias and variance of the importance-weighted discrete adjoint in (13).
Method
D1
DAM
D1
DAM
DAY
Reward
3.5
Reward
Accuracy
60
3.0
1000
2000
3000
0
1000 2000
3000 4000
Training Steps
Training Steps
10
20
(a) GSM8K
(b) Countdown
Time (hours)
Figure 5: Accuracy vs Wall Clock Time on Sudoku. We compare D1 and DAM but plot the wall clock time in hours on the x-axis.
Figure 4: Rewards Curves GSM8K and Count- down, DAM is more effective than D1 in maximizing reward in (23).
32

---

Published as a conference paper at ICLR 2026
Table 4: Generated Samples on GSM8K from the base model and two fine-tuned models, all conditioned on the same unseen prompts. Correct and wrong answers are colored.
Prompt: Jaime is computer programmer for company that currently has employed 100 people. Senior programmers are paid S400 more than junior programmers. If the number of Junior programmers is 2/5 of the total number employees, and they are each paid S2000 per month, calculate the total amount of money the company pays to all the programmers per month.
Answer: 224000
Response: LLaDA-8B-Instin -Bsttuut
LLaDA-8B-Insinn w/ DAM
LLaDA-SB-Instrut w/ DI
<rrassoning>> >
<reasoning?>>
<reasoning>>
The number of Junior program- mers is 2/5 100 = 40. The
The number of junior program- mers is 2/5 * 100 40. The number of senior programmers is 100 40 60. Each senior pro- grammer is paid $2000 $400 $2400 per month. The total amount paid to programmers per month is 40 $2000 + 60 $2400 = S224000.
The number of Junior program-
mers is 2/5 * 100 = 40. The
number of Senior program- mers is 100-40- 60. The total salary for Junior programmers is 40 * $2000 $80000. The total salary for Senior program-
number of Senior program- mers is 100 - 40 = 60. The to- tal amount of money the com- pany pays to all the program- mers per month is 40*2000 +
mers is 60 ($2000 + $400)= =
$144000, </reasoning>>
< <answer>> 94000 </answer>
<answer>> 2240000000 </answer
</reasoning>
<answer>> 224000
</answer>>
John hires a driving service to get him to work each day. His work is 30 miles away and he has go there and back each day. He goes work 5 days week for 50 weeks year. He gets charged $2 per mile driven and he also gives his driver $150 bonus per month. How much does he pay a year for driving?
Prompt:
Answer: 31800
Response: LLaDA-883 8 Instruct
LLaDA- 8B Instruct w/ DAM
LLaDA-SB I--trrructt w/ DI
<reasoning>>
<reasoning>>
<reasoning>>
John drives 30 miles each way. he drives total of 30 miles per day. He drives 30 miles per day for 5 days a week, so he drives 150 miles per week. There are 50 weeks per year, so he drives 7500 miles per year. He gets charged $2 per mile, so he pays $15000 per year for the. He also gives his driver a $150 bonus per month. </reasoning>>
John drives 30 * #2 60 miles
John drives 60 miles each way so he drives 120 miles per day He drives 600 miles per week 50 weeks = 30,000
per day. He drives 60 * 5 300 miles per week. He drives 300 50 15000 miles per year. He pays 15000 2 $30000 for the miles. He pays 30000 150 12 $31800 per year. </reasoning>>>
miles per year He pays 30,000 miles $2/mile = $60,000 per year, </reasoning
<answer> 31800
</answer>
<answer>> 60000 </answer>
<answer>> 15000 //answer>>
Prompt: Bob is in charge of doing laundry for a large hotel. Each room has two sheets, one comforter,
twice many pillow cases as sheets and twice many towels as pillow cases. How many pieces of laun- dry are there in 80 rooms?
Answer: 1200
Response: LLaDA-8P Instructtt
LLaDA-88 8B- Instruct w/ DAM
(LaDA-8B-Instrun w/ D1
<reasoning>>
<reasoning>>
<reasoning>>
Each room has
Each room has 2 sheets + I com-
Each raom has 2 sheets, I com- forter, 2*2 = 4 pillow cases, and 2 8 towels. That, each room has 2+1++4+8 = 15 pieces laundry. So, 80 rooms have 15+80 = 1200 <//easoning>>
2 sheets
forter + 2 pillow cases + 2
1 comforter
* 2 * 2 towels = + + 8 15 pieces of laundry pieces. There- fore, rooms have 15 80 = 1200 pieces of laundry. <//eaa-
Twice as many pillow cases as
sheets: x 2 pillow cases
- Twice as many towels as pillow cases: 2 4 towels
soning>
<answer> 1200 < </answer>>
So, each room has total of 2
<answer>>
1200 </answer>>
1 + 4= 9 pieces laundry. For 80 rooms, the total number of pieces of laundry is 80 & 9- 720.
</reasoning>>
<answer>> 720
33

---

Published as a conference paper at ICLR 2026
Table 5: Generated Samples on MATH500 from the base model and two fine-tuned models, all
conditioned on the same unseen prompts. Correct and wrong answers are colored.
Prompt: In 1992, scoop gelato could be purchased in Italy for 1200 lire. The same gelato would have cost $1.50 in the U.S. At the equivalent exchange rate between the lire and the dollar, how many dollars would be equivalent to 1,000,000 lire?
Answer: 1250
Response: LLaDA- ----A------e 88--88------n Instructt
LLaDA-88 8--- -1-suutttct w/ DAM
LLaDA 8B-Instruc w/ DI
<reasooing>>>>
reesonnng
<reasoning>>>
To find how many dollars are equivalent to 1, 1,000,000 lire, we need to know the exchange rate between lire and 1 dol- lar. From the given informa- tion, we know that 1200 lire is equivalent to $1. $1.50. Therefore, 1,000,000 lire would be equiva- lent to $1,000,000/ / 1200 = 833 dollars, </reasoning>>>
To determine how many dol- lars are equivalent 1,000,000 lire, we need to know the ex- change rate between lire and dollar. Since 1200 lire is equiv- alent to $1.50, the equivalent exchange rate is lire $1.50 1200 S0.0 50.0125. Therefore, 1,000,000 lire is equivalent to: to:</reasoning>>>
To find how many dollars are equivalent to 1,000,000 lire, we need to know the exchange rate between 1 lire and dollar. Given that 1200 lire is equivalent to $1.50, can deduce that lire equivalent to $1.50 1200 = $0.00125, </reasoning>>
<answer>> 125 </answer>
<answer>> 833 </answer>
<answer>> 1250
</answer>>
Prompt: If
9. 27. + Du +1 at what value of y will there be a horizontal asymptote?
Answer: 0
Response: LLaDA ---------A-e 8B-Instruct
LLaDA- -n--rutttct w/ DAM
LLaDA- I--Itutt w/ DI
<reasoning>>
<reasoning>>
<reasoning>>
To find the horizontal asymp- tote the function, we need to compare the degrees of the nu-
To determine the horizon-
To find the horizontal asymp- tote the function, we need to compare the degrees the nu-
tal asymptote of the function
y 272 117224920 49243
+9234+233 we need to analyze the degrees of the poly- nomials in the numerator and the denominator. The degree of the numerator is 2, and the degree of the denominator is 3. Since the degree of the numerator is less than the degree of the de- nominator, there is horizontal asymptote. </reasoning>>
merator and the denominator. The degree of polynomial is the highest power of in the polynomial. In this case, the de- gree of the numerator is 2 and the degree of the denominator is 3. Since the degree of the nu- merator is less than the degree of the denominator. the horizontal asymptote is at y = 0. There- fore, the value of at which there is a horizontal asymptote. </heasoning>>
merator and the denominator.
The degree of polynomial is the highest power of z in the polynomial. Here, the degree of the numerator is 2 and the degrec of the denominator is 3. When the degrees of the numerator and the denominator are differ- ent, the horizontal asymptote is determined by the ratio of the leading coefficients. Therefore, the horizontal asymptote is 27 27.
<answer>> 0 </answer>>
</reasoning>>
<answer>
</answer>>
<answer>> </answer>
Prompt: Find the sum of the real roots x
80x 36
O.
Answer:
Response: LLaDA-8B-Insteee
-BB -B-strructt w/ DAM
MaDA--B--nsttrn w/ D1
<reasoning>
<reasoning'>
<reasoning>>
We can rewrite the given cqua- tion as x² 80r- 36
The given equation is xf - -80z-
The quartic equation x² 80x
36 0. We rewrite this
36 0 can be rewritten as
(x² 11)(z 1)_ 0. This,
as (=² 20)² O. Solving
(x 2.-r 36) (xx2++02 +200++3)
we see, gives us x² = 41 or x = -1. Since x" 41 has no real solutions, the only real root is x = 11 There- fore, the sum of the real roots is -1+ -2. </rrasoning>>>
20 0. we get x² 12.5. Therefore, the sum of the real roots is 2/5 + (-2/5) O. </reasoning>
2
O. Each factor these two factors
20, which gives
has two real roots, so the sum of
the real roots is the sum of the roots of each factor, which is 20. Therefore, the sum of the real roots of the quartic equation is 0. </reasoning>>
<answer>>
<Janswer>>
<answer>> </answer>>
<answer>> </answer>>
34

---

Published as a conference paper at ICLR 2026
Table 6: Generated Samples on Countdown from the base model and two fine-tuned models, all conditioned on the same unseen prompts. Correct and wrong answers are colored.
Prompt:
Using only the numbers [15, 78, 62). create an arithmetic expression that evaluates exactly You must use all numbers from the list, and each number must used exactly once. You may use the operations +, -, *, and / as needed.
Answer: 78- 62 15
Response: -n-ttrcctt
LLaDA-8--Instie w/ DAM
LLaDA-8B-Instran w/ D1
<reasoning>>
<reasoning>>
<reasoning>>
To create an arithmetic expres- sion that evaluates exactly I using the numbers [15. 78, 62), we need to consider the oper- ations of addition, subtraction, multiplication, and division. We start by trying to combine the numbers a way that al- lows us to reach 1. One possible approach is to use the numbers directly to form other numbers that can be further manipulated reach Another approach to use the numbers in combina- tion of operations to achieve the desired1.
To create arithmetic expres- sion that evaluates exactly using the numbers [15, 78, 62].
To create arithmetic expres-
sion that evaluates to exactly using the numbers [15, 78, 62], we need to consider the oper- ations of addition, subtraction, multiplication, and division. We can start by trying to com- bine numbers a way that leads us to the target value. One possible approach is to use the operations strategically break down the problem into smaller parts. can we find a combination of operations that will allow us to reach the target value of I? </reasoning>>
we need to find the correct com- bination of the operations +, *, and /.
</reasoning>>
<answer>> 78/ 62/22-15 15 clanswer:
</reasoning>>
<answer>> 78/15-62/15 </answer>
<answer> 15678/62
</answer>
Using only the numbers [19. 3, 5]. create an arithmetic expression that evaluates to exactly 52. You must use all numbers from the list, and each number must be used exactly once. You may the operations +, +, ", and as needed.
Prompt:
Answer: 19 -
Response: LLaDA 8B Instruct
LLaDA-8B SEEAA-BB-nnstrrn Instruct w/ DAM
LLaDA 8B Instruct w/ DI
<reasoning>>
<reasoning>>
<reasoning>>
To create an arithmetic expres- sion that evaluates to 52 using the numbers [19, 3, 51, we need consider the operations of addition, subtraction, multiplica- tion, and division. We can start by multiplying the numbers 19 and 3, which gives us 57. Then, we can subtract from 57 to get 52. This expression uses all three numhers exactly once and evalu- ates to 52.
To create an arithmetic expres- sion that evaluates to exactly 52 using the numbers [19, 3, 51, we need to consider how to combine these numbers using the opera- tions +, -, +, and to achieve 52. </reasoning>>
To create an arithmetic expres- sion that evaluates to 52 using the numbers [19, 3, 51, we need to consider the operations of addition, subtraction, multiplica- tion, and division. We can start by trying to combine the num- bers in a way that leads us to reach $2. One
<answer>> 19 + 5-5 </answer>>
way could be to use multiplica- tion and addition to We can try to
</reasoning>>
95-5=90
<answer>
19 5
/answer>>
90-3=87
87-19=68
< </reasoning>>
< <answer>>
195-3-5-9
</answer>>
35


Reviews:
[{"review_id": "ouMa5QVcfU", "reviewer": "Reviewer_Wztv", "summary": "This paper introduces Discrete Adjoint Matching (DAM), which is a novel method for solving the following problem: how to fine-tune discrete generative models (like mask-based diffusion Large Language Models) using the reward-based framework called Adjoint Matching (earlier approach works on continuous diffusion/flow models).\n\nThe work is a significant contribution, providing a solution that is not only mathematically elegant and theoretically sound but also demonstrably effective in practice, achieving state-of-the-art results on complex reasoning tasks.\n\nThe paper's context is entropy-regularized optimization, which is a formal way to fine-tune a pre-trained generative model ($p^{\\text{base}}$) to maximize a reward (or minimize a terminal loss, $g(X_1)$) without straying too far from its original training. The previous method for this, Adjoint Matching (AM), is highly effective for continuous models (like image diffusion models) because it uses gradients ($\\nabla g$) as a steering signal. However, this gradient-based approach fails for discrete models because it needs to take a gradient over a discrete set of words. The authors bridge this gap. They propose Discrete Adjoint Matching (DAM), a new formulation that adapts AM to the discrete world. The core innovation is replacing the additive gradient correction of AM with a multiplicative correction factor derived from the exponential of the terminal loss difference (i.e., $e^{-g(y)+g(X_1)}$).", "questions": "Please see the weakness above", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The authors provide a deep theoretical analysis to prove their method for discrete version of AM. They use fixed-point equations to prove that their practical algorithm is guaranteed to converge to the true, theoretically perfect optimal solution\n- The authors address a computationally impossible problem in their theoretically optimal solution. They then methodically build a practical solution: estimation via sampling and approximate the correction factor by sampling a few possible futures (K samples). They also adapts the method to masked diffusion models. \n- The experiments are clear and decisive. DAM achieves state-of-the-art accuracy on three difficult mathematical reasoning benchmarks", "weaknesses": "- The algorithm requires $K$ model-forward passes per training step to build its estimator. While this is clearly effective on an 8B model, the cost for fine-tuning much larger models (e.g., 70B+) is not discussed. A small experiment reporting training time vs. final accuracy for DAM and D1 would make the paper's practical claims much stronger.\n- A valuable addition to the empirical analysis would be an ablation study on the number of samples K used in the importance-weighted estimator.", "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": "IkYlUp62LQ", "reviewer": "Reviewer_d4QG", "summary": "This paper proposes discrete adjoint matching, which is a theoretical framework to conduct reinforcement learning on discrete diffusion models.", "questions": "See Weakness.", "limitations": "", "rating": 8, "confidence": 1, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1, The motivation is clear and significant, locating at the need of reward-guided fine-tuning of discrete diffusion-based models.\n\n2, The theoretical seems to be sound.", "weaknesses": "This seems to be a quite good paper. But I am not a theory expert. So I will be alert to any issues raised by other reviewers. Also, I want to raise a question about the performance of Llada-8b on GSM-8K. According to [A], the performance of base Llada model on GSM-8K is 80+. But in your paper, the performance is 60-70. Could you please explain this gap?\n\nReference:\n\n[A] Revolutionizing Reinforcement Learning Framework for Diffusion Large", "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": "HJsAxQgkHb", "reviewer": "Reviewer_Gkkg", "summary": "The paper introduces **Discrete Adjoint Matching (DAM)**, a framework for **fine-tuning discrete generative models**—such as diffusion-based large language models—through an **entropy-regularized optimization** approach. DAM extends the Adjoint Matching (AM) method, previously designed for continuous diffusion models, to discrete state spaces modeled as Continuous-Time Markov Chains (CTMCs).\n\nThe authors derive a **discrete analogue of the adjoint system** using Dynkin’s formula, providing a statistical interpretation that avoids the control-theoretic derivations used in AM. This leads to an estimator for the optimal transition rates between discrete states, enabling fine-tuning without requiring differentiability.\n\nTo make DAM computationally feasible, the paper introduces **variance reduction** and **importance weighting** techniques, and adapts the method to masked diffusion models, common in language modeling. Theoretically, it provides guarantees via fixed-point and adjoint equations linking DAM to optimal control formulations.\n\nExperiments on synthetic datasets show DAM closely matches the true optimal distributions, outperforming prior baselines like `D1` and `SVDD`. Applied to mathematical reasoning tasks (`GSM8K`, `MATH500`, `Countdown`), DAM yields consistent improvements in accuracy and reward metrics over `D1`.", "questions": "1. **On the discrete adjoint formulation:** Are there particular conditions (e.g., reversibility, bounded rates) under which the discrete adjoint simplifies to known forms?\n2. **On importance sampling and variance reduction:**\n- The section on importance weighting is promising but rather concise. Could the authors expand on how proposal distributions are selected or adapted during training?\n- Can the authors quantify the variance or bias behavior of the estimator as a function of time horizon or dimensionality?\n- How sensitive is DAM to poor proposal distributions, and could hybrid approaches (e.g., reweighting combined with entropy regularization) mitigate this issue?\n3. **On the link with stochastic control theory:** There are strong conceptual parallels between this work and recent formulations of diffusion models within stochastic control frameworks, such as in Pham et al. (2025) [1]. Could the authors clarify whether DAM can be interpreted as a discrete analogue of the Hamilton–Jacobi–Bellman (HJB) formulation underlying those control-based methods?\n4. **On scalability and applicability:**\n- How does the computational cost of DAM scale with the number of discrete states or sequence length?\n- Could the authors comment on potential approximations or sparsity strategies for handling large vocabularies in text diffusion models?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. **Clear conceptual motivation:**\nThe paper addresses a timely and well-motivated gap — extending adjoint-based optimization methods, previously limited to continuous diffusion models, to the discrete generative setting, which is crucial for language and symbolic models.\n2. **Principled extension of Adjoint Matching:**\nDAM is a nontrivial discrete analogue of Adjoint Matching (AM), retaining its optimization-by-simulation philosophy while adapting it to the constraints of discrete-time, discrete-state Markov processes.\n3. **Practical algorithmic contributions:**\nThe inclusion of variance reduction and importance weighting makes the method computationally tractable and numerically stable, which is key for large-scale discrete models.\n4. **Bridge between theory and practice:**\nBy grounding the derivation in both stochastic process theory and discrete optimization, DAM provides a conceptual bridge linking probabilistic control, discrete diffusion, and reinforcement-style objectives.\n5. **Empirical effectiveness:**\nExperiments show consistent and interpretable improvements over strong baselines (e.g., D1, SVDD) on synthetic reasoning and mathematical datasets, validating the theoretical claims.", "weaknesses": "1. **Clarity and depth of the theoretical exposition:**\nThe theoretical development is solid and well-motivated, but occasionally dense. Some key derivations—particularly the transition from Dynkin’s formulation to the discrete adjoint system—could be presented with more intuition and interpretive discussion, to help the reader understand the underlying mechanics beyond the formal algebra.\n2. **Limited discussion of importance sampling techniques:**\nThe paper briefly introduces importance weighting and variance reduction, which are essential for practical implementation, but this part remains somewhat underexplored.\nA more detailed analysis of variance control, proposal distribution selection, or adaptive resampling strategies would strengthen the empirical and theoretical credibility of the approach. It would also help to understand how importance sampling interacts with the entropy-regularized objective in high-dimensional discrete settings.\n3. **Positioning within the broader literature:**\nThe paper’s theoretical contribution could be better contextualized by referencing recent advances connecting stochastic control and diffusion models, such as Pham et al. (2025) [1], where the control–diffusion duality and adjoint computations bear a strong resemblance to those derived here. Acknowledging and comparing these perspectives would clarify the novelty and relevance of DAM within this rapidly developing area.\n4. **Empirical diagnostics and ablations:**\nAlthough results are consistent and promising, the paper would benefit from additional ablation studies—for example, isolating the contributions of entropy regularization, importance weighting, and the adjoint update itself. Such diagnostics would offer clearer insight into which components most directly influence performance.\n\n[1] Pham, L.T.N., et al. “_Discrete Markov Probabilistic Models: An Improved Discrete Score-Based Framework with Sharp Convergence Bounds under Minimal Assumptions._” Forty-second International Conference on Machine Learning (ICML, 2025).", "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": "CuAPewrb1W", "reviewer": "Reviewer_uj3e", "summary": "The paper under consideration proposes an algorithm to fine-tune a discrete generative model. The core idea is to develop a discrete analog of the Adjoint Matching method recently proposed by Domingo Enrich et al. for continuous models. To achieve this generalization, the authors obtain at Theorem 2.2 an equation for the adjoint discrete variable that mirrors the stochastic Maximum principle for the auxiliary stochastic control problem used in Domingo Enrich et al. The resulting algorithm (Algorithm 1) is then derived following closely the strategy used in continuous spaces, with some differences. For example, samples from the optimal law $p^{\\star}$, which are required to leverage the adjoint equation, are approximated using importance weights (Prop 2.4). The paper is completed by further theoretical results that elucidate similarities with the continuous setting (Sec. 3) and a section on numerical experiments where the method is tested both on synthetic examples and more complex mathematical reasoning tasks.", "questions": "- It seems to me that in equation (6) $f(y)$ should be replaced by $f(y)-f(X_\\tau)$\n\n\n\n- The paper puts quite some emphasis on the fact that the derivation of Theorem 2.2 avoids \"convoluted control-theoretic derivation adopted in original AM and providing a more general framework for adjoint-based estimators.\" But at the end of the day, it seems to me that equation (8) is a form of the Pontryagin maximum principle for jump process, and I don't feel that it would be so difficult to obtain following for example  the framework developed in Appendix F of [1] for the control of CTMCs. From this perspective, many aspects of (8), such as the multiplicative nature of the control, appear as rather natural than surprising. I believe that adding the statistical estimator perspective offered in this work is valuable, and I am happy to see it. However, offering also a control theoretical approach would make it even clearer why the method proposed here is the discrete counter part of AM which ultimately relies on equation (3), which is nothing but the Pontryagin maximum principle. Relegating the control-theoretic derivation of Theorem 2.2 to the appendix does not appear to be a good choice to me. \n\n*[1]  Le-Tuyet-Nhi, P. H. A. M., Shariatian, D., Ocello, A., Conforti, G., & Durmus, A. O. Discrete Markov Probabilistic Models: An Improved Discrete Score-Based Framework with sharp convergence bounds under minimal assumptions. In Forty-second International Conference on Machine Learning.*", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper extends an algorithmical framework that has drawn quite some attention for continuous models to the discrete setting in a principled way. The numerical results appear quite convincing.", "weaknesses": "- There is no theoretical guarantees of convergence for the proposed method (as there is none for its continuous counterpart). This is a sever limitation, in my opinion.\n\n* On the methodological side, AM in continuous state space requires the pre-trained model to be such that the initial and final states are independent. I find this a very strong requirement hard to satisfy in practice, unless the pre-trained model had been designed to have such property. If I am not mistaken the same limitation is present here. I encourage the authors to clarify this point.", "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)