VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing
arXiv:2606.28301v1 [cs.LG] 26 Jun 2026
Kijung Jeon Georgia Tech [email protected]
Thuy-Duong Vuong∗ UCSD [email protected]
Molei Tao∗ Georgia Tech [email protected]
Abstract Inference-time scaling is a promising paradigm to improve generative models, especially when outputs must satisfy structural constraints or optimize downstream rewards. We consider Masked Diffusion Model (MDM) and introduce MDM-VGB, a discrete diffusion sampler that augments unmasking generation with theoretically principled reward-guided remasking. Inspired by the recent success of the classical Jerrum-Sinclair backtracking Markov chain [25] in reward-tilted generation [23], MDM-VGB extends the backtracking random walk from a fixed prefix tree to a masked-state graph, allowing tokens to be unmasked and remasked at arbitrary positions. The resulting sampler favors unmasking and remasking moves that lead to higher-value partial configurations, enabling both effective high-reward generation and efficient repair of low-reward samples. We prove that MDMVGB is robust to process-verifier noise and achieves quadratic complexity, while popular test-time heuristics such as best-of-N can incur exponential complexity due to error accumulation. Our theoretical findings are corroborated by strong empirical performance, particularly on popular constraint-satisfaction and scientific benchmarks such as Sudoku and QM9 [22].1
1
Introduction
Many sequence-generation tasks require outputting configurations that satisfy global structural properties. Such properties are often easy to verify for complete sequences, but difficult to reliably extrapolate from partially generated configurations. Examples include mathematical proofs, computer programs, Sudoku puzzles, protein sequences for designing drug candidates: a completed configuration can often be accurately evaluated by a theorem prover, compiler, (Sudoku) constraint checker, or drug-property oracle, while evaluating whether a partial configuration can be extended to a valid full configuration is much harder. Verifier-guided inference-time algorithms are a promising direction for improving structured generation [5, 15, 19, 23, 26–29]. In many realistic settings, one has access to a reliable, possibly perfect, outcome verifier for complete configurations, together with weaker and noisier process verifiers that estimate the potential value of partial configurations [17, 28, 32]. Empirically, non-trivial gains have been shown even with simple methods that only use an outcome verifier. A canonical example is Best-of-N , which selects the highest-reward full configuration among n sampled candidates [5]. Access to a process verifier may unlock additional gains [17, 27, 28]; however, a recent work by Rohatgi et al. [23] points out that these methods suffer from the curse of horizon as even mild errors in process verifiers can lead to outcome-level errors that grow exponentially with sequence length. They further observed that the classical Jerrum–Sinclair backtracking Markov chain [25] can avoid this issue, yielding a theoretically grounded inference-time algorithm with promising performance. * Joint mentorship on this project. 1 All code is provided in the following repository: https://github.com/KraitGit/MDM-VGB.
Preprint.
(a) Root-start reward tilting
depth
(b) Leaf-start editing / repair
depth
root
depth 0
depth 0
depth 1
depth 1
depth 2
depth 2
depth 3
depth 3
depth 4
root reveal re-mask
depth 4
y low
leaf reward
y0 high
low
(a) Generation (Root-start)
yrepair leaf reward
high
(b) Editing (Leaf-start)
Figure 1: Visualization of MDM-VGB for generation and editing. The depth of a state is the number of unmasked tokens, and same-depth states are ordered from left to right by estimated value. MDMVGB is a random walk on the masked state graph: root-start generation begins at the fully masked state, while leaf-start editing begins at a fully revealed configuration to repair. Since value estimates are noisy, greedy unmasking can be suboptimal; re-masking can correct earlier mistakes. However, this chain is fixed-order: it reveals tokens from left to right and can backtrack only by deleting the most recently generated token. This restriction limits its ability to revise earlier decisions, which is especially problematic in editing, and in general constraint-generation tasks, where early mistakes can often be fatal (see Figure 7). In this work, we introduce Masked-Diffusion-Model Value-Guided Backtracking (MDM-VGB), a verifier-guided discrete diffusion sampler with remasking. MDM-VGB generalizes the classical Jerrum-Sinclair backtracking Markov chain [25] from fixed-order generation to any-order generation. The Markov chain underlying MDM-VGB operates on partial configurations, i.e., masked states, where some coordinates are revealed, and the remaining coordinates are masked. Given an outcome reward on fully revealed configurations, we define the value of a partial configuration as the expected reward of compatible full configurations. At each step, the Markov chain either reveals masked token(s) or remasks already revealed token(s), with transition probabilities biased toward configurations with larger estimated value. This rule gives a principled randomized-greedy way to reveal and remask tokens: tokens whose removal increases the value of the partial configuration are treated as unreliable and are more likely to be revised; similarly, unmasking decisions that lead to higher-value configurations are prioritized (see Definition 4.2). This Markov chain supports both generation and editing. For generation, the chain starts at the fully masked root state (see Figure 1a); for editing, it starts from a low-reward fully revealed configuration (see Figure 1b). Unlike the fixed-order autoregressive backtracking from [23], MDM-VGB can unmask and remask tokens at arbitrary locations, allowing it to efficiently fix fatal early mistakes (see Figure 7); this flexibility yields significant performance gains in practice. While MDM-VGB also supports flexible block updates in the sense that each transition may reveal or re-mask an arbitrary block of coordinates, we will place particular emphasis on the singleton-block case as in the original masked diffusion model, which we call Any-Order-Auto-Regressive VGB (AOAR-VGB). This simplest variant already yields strong empirical gains while admitting the cleanest theoretical guarantees. This Markov chain provably targets the reward-tilted target distribution in the asymptotic time limit. In particular, in hard-constraint settings with a perfect outcome verifier, where invalid full configurations are certifiably rewardless, the sampler is guaranteed to never produce an invalid full configuration. Moreover, as in the fixed-order case, the chain mixes rapidly, even when subjected to substantial multiplicative error in the intermediate value estimates provided by the process verifiers. In particular, we show that AOAR-VGB converges to the target reward-tilted distribution within a number of Markov steps2 that is quadratic in the sequence length n, whereas best-of-N incurs an exponential cost. In addition, motivated by non-reversible liftings of Markov chains [11] and the momentum lift of the AR-VGB chain considered by Rohatgi et al. [23], we also study MDM-VGB-Momentum, the momentum lift of MDM-VGB. By augmenting each state with a direction variable, the momentum lift mitigates the diffusive and oscillatory behavior of MDM-VGB and improves first-leaf hitting time, a practical quantity of interest under finite-budget sampling. 2 In our implementation, the cost per step matches a standard forward generation step, up to a small constant factor
2
Empirically, MDM-VGB and MDM-VGB-Momentum deliver consistent gains on scientific-design benchmarks and standard combinatorial constraint generation tasks such as QM9 and Sudoku. For generation tasks, MDM-VGB-Momentum achieves higher reward at matched inference cost than best-of-N , and reward-guided MDM (MDM-VGR, see Section 3), with especially impressive gains for harder tasks where unguided MDM has lower accuracy. For editing tasks, MDM-VGB and MDM-VGB-Momentum outperform fixed-order backtracking algorithms from [23, 25]: on Dyck grammar editing, they reach near-perfect accuracy using substantially fewer moves compared to AR-VGB and AR-VGB-Momentum. Moreover, MDM-VGB and MDM-VGB-Momentum improve over non-editing baselines in both output quality and cost across scientific-design benchmarks, e.g., DNA and Protein synthesis. Contributions. • We introduce MDM-VGB, a reward-guided masked discrete diffusion sampler with a principled remasking strategy that supports both generation and editing, with especially strong gains in editing. • Theoretically, we prove that MDM-VGB efficiently targets the reward-tilted distribution over complete configurations, even in the presence of substantial process-verifier error. • Across structured-generation benchmarks, MDM-VGB outperforms standard baselines under matched inference budgets. We also observe substantial gains for editing compared to the fixed-order backtracking baselines (AR-VGB and AR-VGB-Momentum) from [23].
2
Other related works
Recent works, such as Lee et al. [15], Wang et al. [26], have considered alternative formulations of masked diffusion models (MDM) with remasking. Even without an additional reward signal, remasking can empirically improve sample quality by correcting errors introduced by MDM-style parallel decoding [26]. However, it is less clear how to remask in a principled way with quantitatively significant benefits. In fact, the Markov-chain formulation of Lee et al. [15] also targets a reward-tilted stationary distribution, but unlike our method, it does not provide mixing-time or error-tolerance guarantees. Moreover, Lee et al. [15] remasks coordinates uniformly at random rather than by confidence, which reduces its ability to correct earlier mistakes. Remasking diffusion (ReMDM) [26] can use confidence scores to guide remasking decisions, but its updates follow a factorized parallel decoding structure: all coordinates are resampled in parallel from coordinate-wise conditionals. This can be limiting for structured outputs with strong cross-coordinate dependencies. For example, changing 1 + 2 = 4 to 1 + 2 = 3 gives a valid expression, whereas simultaneously changing 4 → 3 and 2 → 1 does not. ReMDM does not target the reward-tilted distribution, and as with Lee et al. [15], does not provide mixing-time or robustness guarantees. Fixed-order backtracking (AR-VGB) [23, 25] comes with strong theoretical guarantees, as MDMVGB does, but it does not allow earlier mistakes to be edited without undoing a long suffix. Likewise, sampling-and-search methods like Sequential Monte Carlo [20] and Monte-Carlo Tree-Search [19, 20], as well as standard masked diffusion [2, 3, 14, 24], do not allow backtracking, and thus cannot repair earlier mistakes in the generative process. This is particularly problematic for structured languages, where an early error can make every subsequent extension invalid (see Figure 7).
3
Preliminaries
Let Y = V n be the set of fully revealed sequences of length n over vocabulary V. Let Z = (V ∪ {[mask]})n be the set of masked states, where ∅ = (([mask], . . . , [mask])) denotes the fully masked state. For z ∈ Z, let R(z) = {i : zi ̸= [mask]} be the revealed set and k(z) = |R(z)| be its depth. Let C(z) = {y ∈ Y : yi = zi for all i ∈ R(z)} denote the compatible completion set. Throughout, we fix a conditioning context x. Given P a reference distribution πref (·|x) over Y, we define the marginal at z ∈ Z as πref (z | x) := y∈C(z) πref (y | x). For B ⊆ [n] \ R(z) and aB ∈ V B , let z B→aB be the state obtained by setting the tokens in B to aB . We define the conditional marginal at B conditioned on z, i.e., πref (YB = aB | x, z), by: πref (YB = aB | x, z) :=
πref (z B→aB | x) = Pπref (YB = aB | Y ∈ C(z)). πref (z | x) 3
Masked diffusion model (MDM). Following [24], if a pretrained MDM is unavailable for a task, we train a discrete masked diffusion model πref (·|x) ≡ πθ over Y using a continuous masking process by minimizing the following objective function: ′ X αt LMDM (θ) = E(x,y)∼D,t∼Unif[0,1],zt ∼qt (·|y) log πθ (Yi = yi | x, zt ) . 1 − αt i: zt,i =[mask]
where D is the training data, αt ∈ [0, 1] is a decreasing masking schedule, αt′ is its time derivative, and zt is sampled from the distribution qt (·|y) over Z defined by: qt (zt | y) =
n Y
(αt 1{zt,i = yi } + (1 − αt )1{zt,i = [mask]}) .
i=1
At inference time, MDM samples from πref (·|x) by starting from ∅ and iteratively revealing a new block of tokens according to the marginal of πref (·|x) conditioned on the revealed tokens. In other words, at a current masked configuration z, the process reveals a new block B ⊂ [n] \ R(z) by sampling from πref (YB = · | x, z). Reward-tilted sampling. Given a base reference model πref (· | x) and a nonnegative terminal score, i.e., reward function, τ (x, ·) : Y → R≥0 , define the reward-tilted target distribution π ⋆ (· | x) over Y where: π ⋆ (y | x) :=
πref (y | x)τ (x, y) , Z(x)
Z(x) :=
X
πref (y | x)τ (x, y).
y∈Y
P We also define the marginal of π ⋆ at z ∈ Z by π ⋆ (z | x) := y∈C(z) π ⋆ (y | x). In particular, binary reward τ (i.e., τ (x, y) ∈ {0, 1}, ∀y ∈ Y) corresponds to the hard-constraint case. Our goal is to sample from π ∗ using a given base model πref and terminal score τ at inference time, without accessing or changing the base model’s parameters, i.e., via test-time scaling rather than RL fine-tuning. Value of partial configurations. For a given terminal score τ , the ideal value of a partial configuration or masked state z ∈ (V ∪ {[mask]})n is V ⋆ (x, z) := EY ∼πref (·|x) [τ (x, Y ) | Y ∈ C(z)] . In practice, V ∗ is not available, and therefore we train a process verifier Vb to approximate V ∗ by fitting Vθ with mean squared error, as in Rohatgi et al. [23]: h i 2 Lreg (θ) := Ex,z,y (Vθ (x, z) − τ (x, y)) , where z is a partial configuration sampled from the reference model πref and y is a completion of z. After training, we use Vb = Vθ . In some settings, we instead use a lightweight heuristic to reduce computational cost. For example, in constraint-satisfaction tasks such as Sudoku, where the terminal reward is the indicator that a full configuration satisfies a certain constraint, we take Vb (x, z) to be the indicator that the partial assignment z has not yet violated the constraint. We note that this approximation can be far from the true value V ⋆ (x, z), since there exist partial assignments that do not violate the constraint but cannot be extended to valid full configurations. Nevertheless, such heuristic verifiers can still perform strongly on some constraint-satisfaction tasks; see Section 5. Value-guided Masked Diffusion Model (MDM-VGR). Given a base MDM πref and the value function V ⋆ with respect to a terminal reward τ , MDM-VGR samples from π ⋆ (·|x) by starting from ∅ and iteratively revealing a new block of tokens according to the reward-tilted marginal conditioned ⋆ B→aB ) on the current state, i.e., π ⋆ (YB = aB | x, z) = πref (YB = aB | x, z) · V (x,z . In practice, V ⋆ (x,z) MDM-VGR uses a process verifier V̂ that approximates V ⋆ . 4
(a) Masked state graph
depth
(b) Prefix tree
depth
root
root
depth 0
depth 0
depth 1
depth 1
depth 2
depth 2
depth 3
depth 3
depth 4
depth 4
low
leaf reward
high
low
reveal re-mask
leaf reward
high
Figure 2: Masked-state graph vs. prefix tree. MDM-VGB operates on the any-order masked-state graph, which allows arbitrary-coordinate re-masking. AR-VGB operates on a fixed-order prefix tree, where backtracking only removes the latest revealed token. Masked states graph. Consider a masked state z ∈ Z. For B ⊆ [n] \ R(z) and aB ∈ V B , we call z B→aB a forward child of z; let C(z) be the set of children of z. For B ⊆ R(z), write z −B for the state obtained by re-masking B, which we call a backward parent of z; we note that a state can have multiple parents and let P (z) be the set of parents of z. The parent-child relation between states defines the masked states graph, an undirected graph whose vertices are masked states and edges are {z, w}, where w is a backward parent of z. We write z ∼ w and say z and w are neighbors when {z, w} is an edge in the graph. We call the fully masked state ∅ the root, and the set of full configurations Y the leaves of this graph. Prefix tree. The AR-VGB Markov chain from [23, 25] is a random walk on the following prefix tree graph. Fix an ordering σ of {1, 2, · · · , n}. Let Zσ denote the set of masked states in which the revealed coordinates form a prefix in this order. The prefix tree on Zσ is the graph where {x, y} is an edge whenever x is the back-parent of y obtained by remasking its last unmasked coordinate in the order σ. See Figure 2 for an illustration.
4
Methods and Theoretical Guarantees
4.1
Formulation of MDM-VGB
To address the goal of Section 3, we now describe the MDM-VGB algorithm. At a high level, MDM-VGB is a random walk on the weighted masked state graph with appropriately chosen edge weights. The edge weights depend geometrically on the estimated values of endpoint states, so that remasking or unmasking moves that lead to higher-value states are prioritized. We also rescale each edge weight by a scalar depending on the depths of its endpoint states. This depth-dependent normalization is designed to balance forward moves that unmask new coordinates and backward moves that remask coordinates at each state and is crucial for ensuring the chain’s fast convergence. The MDM-VGB Markov chain is defined by either of the following equivalent definitions. −1 Definition 4.1. Fix λ ≥ 0 and depth coefficients sk,r := n−r , which satisfy n−k sk,r = k r k r sk−r,r . The MDM-VGB Markov chain is the random walk on the weighted masked state graph, where the edge weights are as follows: for neighboring states u ∼ v, let p denote the lower-depth endpoint and c denote the higher-depth endpoint among u, v. Let the symmetric edge weight between u, v be ( sk(p),k(c)−k(p) πref (c | x)Vb (x, c)Vb (x, p)λ if c ∈ /Y f (u, v) = f (v, u) := . b sk(p),k(c)−k(p) πref (c | x)V (x, c) if c ∈ Y Equivalently, the same chain can be specified through its local forward and backward move weights. In both views, sk,r is the depth-dependent normalization factor that balances forward moves revealing r new coordinates with backward moves re-masking r revealed coordinates. 5
Definition 4.2. For a state z at depth k, a child c = z B→aB ∈ C(z) where B ⊆ [n] \ R(z) with ′ |B| = r, and a backward parent p = z −B ∈ P (z) where B ′ ⊆ R(z) and |B ′ | = r′ , set: ( sk,r Vb (x, z)λ Vb (x, c)πref (YB = aB | x, z), c ∈ / Y, wfwd (z → c) := sk,r πref (YB = aB | x, z)Vb (x, c), c ∈ Y, ( sk−r′ ,r′ Vb (x, z)Vb (x, p)λ , z ∈ / Y, wbwd (z → p) := b sk−r′ ,r′ V (x, z), z ∈ Y, P P Let W (z) := c∈C(z) wfwd (z → c) + p∈P (z) wbwd (z → p) be the normalization factor at z. The MDM-VGB Markov chain’s transition probabilities at z are given by: ∀c ∈ C(z) : P[z → c] =
wfwd (z → c) and ∀p ∈ P (z) : W (z)
P[z → p] =
wbwd (z → p) W (z)
The first definition is more helpful for proving the theoretical properties of the Markov chain, while the second is more helpful for implementation. Remark 4.1 (Role of λ). The parameter λ ≥ 0 controls how strongly re-masking decisions depend on the value of the resulting parent state: larger λ makes the chain more likely to erase a block that yields a higher-value parent state. Our ablation study found that a moderate λ yields the best performance, as the chain can effectively edit problematic tokens while still retaining sufficient flexibility to explore alternative edits (see Section 5.3). In the appendix, we further distinguish between the case λ = 0, which we call Balanced MDM-VGB (see Section C.3), and the general case λ ≥ 0, which we call Geometric Balanced MDM-VGB (see Sections C.5 and E.2). The λ = 0 chain is closer to existing backtracking walks [23, 25], but it selects a uniformly random block for re-masking and thus is less effective at editing problematic tokens. In Section C.2, we also show that the direct generalization of AR-VGB to the diffusion setting, which corresponds to setting λ = 0 and having no depth-normalization factor sk,r , has exponentially slow convergence. Remark 4.2 (Implementation details: shortlisting rule). To improve efficiency, instead of exploring the entire neighborhood of a masked state, we follow the practical shortlisting convention of Rohatgi et al. [23]: each MDM-VGB (or MDM-VGB-Momentum) step considers at most Lf choices of blocks to unmask, Lb choices of blocks to re-mask, and K token vocabulary candidates for each coordinate to unmask. Unless otherwise stated, we use the singleton block size |B| = 1. 4.2
Asymptotic behavior: properties of the stationary distribution
We now state the assumptions on the estimated masked-state values used in our theoretical analysis. Following Rohatgi et al. [23], we assume that the process verifier approximates the true expected value of every masked state within a multiplicative factor κ. We also assume that the estimated values are uniformly bounded between m and M . This boundedness condition is mild in practice, since these estimated values are typically clipped during implementation. Assumption 4.1. Assume terminal anchoring at leaves (Vb (x, y) = τ (x, y), ∀y ∈ Y) and the following κ-accurate verifier condition on every reachable non-leaf state z: κ−1 V ⋆ (x, z) ≤ Vb (x, z) ≤ κV ⋆ (x, z) Define the positive target-support component by Z+ (x) := {z ∈ Z : π ⋆ (z | x) > 0}, and assume additionally that every non-leaf z ∈ Z+ (x) satisfies 0 < m ≤ Vb (x, z) ≤ M < ∞. It is easy to see that MDM-VGB is an irreducible Markov chain. Thus, its lazified version3 is ergodic and converges to a unique stationary distribution µ. Next, we show that the stationary distribution µ, conditioned on the leaves Y, is exactly the target distribution π ⋆ .
6
Theorem 4.1. Assume terminal anchoring Vb (x, y) = τ (x, y) for all y ∈ Y. For any λ ≥ 0, the stationary law µ of the lazy MDM-VGB chain satisfies µ(· | Y) = π ⋆ (· | x). In particular, in hard-constraint settings with τ ∈ {0, 1}, µ(· | Y) assigns zero probability to invalid configurations. Our analysis focuses on the simplest case, MDM-VGB with singleton-block updates, i.e., AOARVGB. The next theorem lower bounds the probability mass of the full configurations under the stationary distribution. Theorem 4.2. Under Assumption 4.1, the stationary distribution µ of the lazy AOAR-VGB chain places substantial mass on the set of full configurations Y: for some constant C > 0, µ(Y) =
1 . CM λ κn
Mixing time guarantee. We study AOAR-VGB’s speed of convergence to the stationary distribution, commonly known as the mixing time. The following theorem lower bounds the spectral gap of the AOAR-VGB Markov chain (see [16] for the connection between spectral gap and mixing time). Theorem 4.3. Under Assumption 4.1, let PAOAR denote the transition matrix of the lazy AOARVGB Markov chain. Define the spectral gap [16] by γ(PAOAR ) := 1 − max{|λ| : λ ∈ Λ(PAOAR ), λ ̸= 1}, where Λ(PAOAR ) denotes the set of eigenvalues of PAOAR . Then m 2λ 1 γ(PAOAR ) = Ω . M κ4 n2 By standard Markov chain theory, this together with Theorem 4.2 gives an upper bound on the mixing time to the target reward-tilted distribution π ⋆ (·|x) over Y, starting from any initial state. Corollary 4.1. Assume Assumption 4.1. For any target accuracy level δ > 0, the distribution ZT obtained after T steps of the lazy AOAR-VGB chain initializedat a masked state z satisfies dTV (L(ZT | ZT ∈ Y), π ⋆ (· | x)) ≤ δ and P[ZT ∈ Y] ≥ Ω M λ1κn , when 2λ λ ! M M κn a T ≥Θ κ4 n2 log m δµ(z) 2λ 4 2 (M/m)λ κn In particular, when z is the root ∅, it suffices to take T = Θ M κ n log m δ a where µ(z) is the probability mass of z under the stationary distribution µ of the lazy chain AOAR-VGB.
Remark 4.3. The bound above controls the target law conditioned on reaching a leaf, but the the stationary distribution µ of the current chain’s mass on the leaves, i.e. µ(Y), is only of order 1/n. Thus, obtaining a leaf with constant probability would require running the chain O(n) times. By using the modification from [11, 23], i.e. increasing the holding probabilities at the leaves, we can instead make µ(Y) at least a constant, and obtain a quadratic-time sampler for π ⋆ (·|x). This modification is mainly a technical step to improve the theoretical analysis; in practice, the sampler simply stops after seeing a few leaves, so we keep the simpler construction of the Markov chain. 3 For a Markov chain with transition matrix P and a parameter γ ≥ 0, its γ-held lazy version has transition matrix γ γ 1 P lazy,γ = 1+γ I + 1+γ P. Each step of P lazy,γ stays in-place with probability 1+γ , and makes moves according to P 1 with probability 1+γ . When γ = 1, we refer to P lazy,γ as the lazy version of P . For our theoretical analysis, it is convenient to work with the lazy MDM-VGB.
7
Proof sketch. We discuss the analysis of AOAR-VGB. The proof is by relating the properties of AOAR-VGB to those of AR-VGB. The key observation is that the weighted masked-state graph defining AOAR-VGB can be decomposed as a sum of weighted prefix-tree graphs that define ARVGB. Consequently, the properties of AOAR-VGB can be compared with those of the corresponding AR-VGB walks. The stationary distribution probability at a state z is proportional to the total weight of the edges incident to z. First, the stationary distribution of AOAR-VGB is the average of the stationary distributions of the AR-VGB walks, and thus induces the same target distribution π ⋆ on the leaves as the AR-VGB walks. Second, the weight on the leaves under the stationary distribution of AOAR-VGB can be lower bounded by those of the AR-VGB chains. Third, the conductance of the masked-state graph, which corresponds to the relaxation time and mixing time of AOAR-VGB, can be lower bounded by the conductance of the prefix-tree graphs, which in turn can be lower bounded via the existing analysis in [23, 25]. For details, see Theorem C.5. 4.3
Formulation of MDM-VGB-Momentum
We introduce MDM-VGB-Momentum, a momentum lift of MDM-VGB in the spirit of [11] that targets the same distribution over the full configurations, while reducing oscillatory behaviors. The lift augments each masked state z with a direction variable, yielding two lifted states (z, ↓) and (z, ↑). In reveal mode (z, ↓), the chain can only reveal tokens, while in re-mask mode (z, ↑), it can only re-mask tokens. At each state, the chain switches between these two modes with a probability governed by a parameter χ chosen to reduce reveal/re-mask oscillations. See Figure 3 for an illustration. For AR-VGB, momentum has been shown empirically to improve performance [23]; theoretical work by [11] also shows mixing-time improvements in the low-error regime, where the approximation parameter κ is sufficiently close to 1. On the theory side, we prove that at a sufficiently low error– when κ = 1 + O(1/n)–MDM-VGB-Momentum reaches a leaf in expected O(n) time, achieving a quadratic speedup over the non-momentum version, which requires Ω(n2 ) time4 . We expect analogous mixing-time improvements in this low error regime, but leave a full mixing time analysis to future work. In our experiments, MDM-VGB-Momentum tends to perform better on generation tasks, whereas the non-momentum version performs slightly better on editing tasks. Definition 4.3 (MDM-VGB-Momentum). For a masked state z, a child c ∈ C(z), and a parent p ∈ P (z), let wfwd (z → c) and wbwd (z → p) be as in Definition 4.2. Let X X wfwd (z → c), B(z) := wbwd (z → p), F (z) := c∈C(z)
M (z) := min {F (z), B(z)} ,
p∈P (z)
W (z) := F (z) + B(z).
Fix a parameter χ ∈ [0, 1]. The MDM-VGB-Momentum with parameter χ is a Markov chain on the lifted state space Ze = Z × {↓, ↑}a . e the Markov chain can either transition to (c, ↓) for c ∈ C(z), or (z, ↑), or stays at At (z, ↓) ∈ Z, (z, ↓). The transition probabilities are given by: P[(z, ↓) → (c, ↓)] =
wfwd (z → c) χM (z) B(z) − χM (z) , P[(z, ↓) → (z, ↓)] = , P[(z, ↓) → (z, ↑)] = W (z) W (z) W (z)
e the Markov chain can either transition to (p, ↑) for p ∈ P (z), or (z, ↓) or stays at At (z, ↑) ∈ Z, (z, ↑). The transition probabilities are given by: P[(z, ↑) → (p, ↑)] =
wbwd (z → p) χM (z) F (z) − χM (z) , P[(z, ↑) → (z, ↑)] = , P[(z, ↑) → (z, ↓)] = W (z) W (z) W (z)
a Blue ↓ denotes forward/reveal momentum, and red ↑ denotes backward/re-mask momentum.
The following theorem characterizes asymptotic behavior of AOAR-VGB-Momentum.
4We suppress dependency on other parameters M, m, λ.
8
Figure 3: MDM-VGB and MDM-VGB-Momentum. MDM-VGB alternates value-guided reveal and re-mask moves, while the momentum variant augments the state with a direction variable to reduce immediate reversals. Theorem 4.4 (Stationary distribution of AOAR-VGB-Momentum). AOAR-VGB-Momentum with parameter χ ∈ [0, 1] as defined in Definition 4.3 is an ergodic Markov chain, with stationary distribution µ e defined by 1 µ e(z, ↓) = µ e(z, ↑) = µ(z) 2 where µ is the stationary distribution of the AOAR-VGB chain. Hence the projection proj# µ e of µ̃ on Z satisfies proj# µ e = µ. In particular, if µ(· | Y) = π ⋆ (· | x) then proj# µ e(· | Y) = π ⋆ (· | x). The following theorem shows that in the low error regime κ = 1 + O(1/n), AOAR-VGB-Momentum takes only O(n) steps to hit a leaf, whereas AOAR-VGB requires Θ(n2 ) steps. Theorem 4.5. Consider the AOAR-VGB-Momentum with parameter χ as defined in Definition 4.3. Let p↓ (z) =
F (z) , F (z) + B(z)
p↑ (z) =
B(z) , F (z) + B(z)
∆ :=
sup |p↓ (z) − p↑ (z)| , ◦ (x) z∈Z̃+
◦ where the supremum is over reachable interior states Z̃+ (x) := {z ∈ Z+ (x) | 0 < k(z) < n} × {↓, ↑} . Suppose ∆ = O(1/n) and 1 − χ = O(1/n). Let Tn and TnMo be the time of first hitting the set of leaves Y for AOAR-VGB and AOAR-VGB-Momentum initialized at ∅ and (∅, ↓) respectively, then E[Tn ] = Θ(n2 ), E[TnMo ] = O(n), In particular, under the multiplicative value-error bound κ−1 V ⋆ ≤ Vb ≤ κV ⋆ , the condition ∆ = O(1/n) is satisfied when κ = 1 + O(1/n).
5
Experiments
We evaluate our methods MDM-VGB and MDM-VGB-M OMENTUM on reward-guided generation and editing tasks across several domains. Task descriptions. We briefly summarize each task, its terminal score R, the base model, and the learned verifier used by our method. When applicable, the Pass@95 terminal reward is defined from the terminal score R as described in the Evaluation metrics paragraph. Molecular design aims to design drug-like molecules using the QM9 benchmark [22]. For the base model, we train a 92.4M-parameter QM9 MDLM model with the MDLM objective using 9
the UDLM-QM9 architecture. We use the quantitative estimate of drug-likeness (QED) [4] as the terminal score R. The terminal reward is the Pass@95 indicator defined from QED as described in the Evaluation metrics paragraph. To approximate the values of partial configurations, we train a process verifier using a token Transformer over SMILES tokens with 3.24M parameters. DNA enhancer design evaluates regulatory-sequence generation. We use D3LM [33], a DNA diffusion language model initialized from Nucleotide Transformer [8], as the base model, and DeepSTARR developmental enhancer activity [9] as the terminal score R. The terminal reward is the Pass@95 indicator defined from DeepSTARR developmental activity. The goal is to generate DNA sequences with high predicted developmental enhancer activity. We learn the values of partial configurations using a small Transformer over D3LM tokens. We test multiple verifier model sizes (see Section 5.3) and report results for the verifier with 0.97M parameters. Protein motif scaffolding evaluates whether generated protein sequences can scaffold a prescribed structural motif. We use EvoDiff OADM as the base model [1], and fold generated sequences with OmegaFold [30]. For terminal scoring, we measure the motif Cα root-mean-square deviation (RMSD) against the reference structure [18]. For evaluation, we report motif RMSD and Success@1Å, the percentage of generated configurations whose motif RMSD is below 1.0Å. We learn the values of partial configurations using a small Transformer over amino-acid tokens with 1.94M parameters. To avoid training the verifier on a sparse Success@1Å signal, we use a smoothed RMSD-derived target for verifier training. Sudoku is a typical constraint-satisfaction generation task. The goal is to generate a complete 9 × 9 grid filled with digits 1, . . . , 9 so that every row, column, and 3 × 3 subgrid contains each digit exactly once. The terminal reward is the binary indicator that all constraints are satisfied. For partial configurations, we use a heuristic process verifier that returns the indicator that no constraint has yet been violated. For the base model, we use DiT-MDM, i.e., a masked discrete diffusion model [24] with a Diffusion Transformer backbone [21]. Letter avoidance is a synthetic constrained natural-language generation task: the model must generate a one-sentence story without using the letter “e”, a task also considered in AR-VGB [23]. We use Qwen3-0.6B-diffusion-mdlm-v0.1 as the base model [10]. The terminal reward is the binary indicator that the sequence does not contain the forbidden character. For partial configurations, we use a heuristic process verifier that returns this same indicator. Dyck grammar is a formal language over bracket symbols {(, ), [, ]}, commonly used as a sandbox test (see e.g. [23]). The terminal score is the binary indicator that the (full) configuration is a valid well-balanced expression: every opening bracket must be closed by the matching bracket type in the correct nested order. For example, (), [], and [()] are valid, whereas (, ], [(], ((), and [) are not. We focus on the editing task, which highlights the benefit of arbitrary-coordinate repair, since due to long-range dependencies, an early mistake can invalidate an otherwise plausible subsequent generation. For the base model, we use a BERT-style masked language model with 12.9M parameters. For the process verifier, we train a small token-state Transformer with 3.24M parameters. Evaluation metrics. For a given task, all test-time inference methods use the same reference model πref and terminal reward τ . We compare MDM-VGB and MDM-VGB-M OMENTUM against the baselines BASE, B O N, MDM-VGR, AR-VGB, and AR-VGB-M OMENTUM. BASE is the unguided MDM sampler. B O N (best-of-N ) generates N full configurations using unguided MDM and returns the one with the highest terminal reward [5, 12]. To match the complexity of B O N, we allow up to n × N MDM-VGB steps, and upon reaching this budget, perform forward-only sampling using πref . For editing, all methods edit the same pool of full configurations of lower reward. See Section A.1 for further details on the baselines. For tasks with hard constraints, i.e., binary terminal score/reward, we report the success rate (percentage of configurations satisfying the constraint), and, when appropriate, a secondary quality metric: for Sudoku, we report the number of violated constraints and for Letter, we report the overall quality score (Ovl.) of the generated sentence provided by Qwen3-32B model [31]; see the Letter prompt details in Section F.1. For tasks with a real-valued terminal score (QM9, DNA, Protein), we report a calibrated metric that measures how often generated samples reach the high-quality region of the base model’s distribution. Given an independent BASE calibration set and task-specific terminal score R(y), let qαBase (R) be 10
2
4
8
32
1
90 85
MDM-VGB-Momentum
MDM-VGB
QM9
32 8 4
2
16
dotted: FLOP-adjusted NFE)
: raw NFE /
32 16
50 1
8
2 4
0
0
2
Dev
Valid (%)
1
0
27
28
29
210 211
25
NFE / sample
26
27
28
29
210
32
16
8
50
4 2
1
0
1.5
1
1.0
60 26
Protein
100
Mean RMSD (A)
100
80
: FLOP-adjusted NFE)
DNA
100
50 1
(solid: raw NFE /
MDM-VGB-Momentum (
0
2
Viol (%)
MDM-VGB
100
Pass@95 (%)
Solve (%)
MDM-VGR
Sudoku 16
100 95
MDM-VGR
Base
Success@1A (%)
Base
BoN
Pass@95 (%)
BoN
26
NFE / sample
27
28
210 211
29
26
NFE / sample
27
28
29
210
211
NFE / sample
Figure 4: Generation quality–cost frontiers for Sudoku, QM9, DNA, and Protein. Solid curves count raw reference-model forward calls, and dotted curves show FLOP-adjusted NFE including learned-verifier overhead. Numeric labels on the B O N curve denote the number of independent rollouts N . Heuristic verifier Method
Letter
Learned verifier
Sudoku
QM9
DNA
Protein
Avoid ↑ Ovl.↑ Solve ↑ Viol.↓ Pass@95 ↑ Valid ↑ Pass@95 ↑ Dev.↑ Success@1Å ↑ RMSD ↓ Base BoN MDM-VGR MDM-VGB MDM-VGB-Momentum
0.0 0.1 100.0 100.0 100.0
– 7.00 5.72 5.86 5.72
92.9 99.3 86.9 92.0 100.0
0.791 0.087 1.659 0.917 0.000
3.2 41.0 39.1 61.3 90.6
66.8 79.5 92.9 82.2 95.7
0.041 0.500 0.058 0.411 0.705
-0.50 1.37 -0.42 1.06 1.55
8.7 73.7 15.0 62.0 82.0
1.732 0.926 1.780 1.101 0.982
Table 1: Generation results at a representative budget N = 16. Letter and Sudoku use heuristic verifiers; QM9, DNA, and Protein use learned verifiers. Columns report primary and secondary metrics for each task. Bold and underline mark best and second-best. the α-quantile of the distribution over terminal scores produced by the base model. We define 1{R(y) ≥ qαBase (R)}, if higher values are better, Pass@α(y) = Base 1{R(y) ≤ q1−α (R)}, if lower values are better. Across QM9 and DNA we use α = 0.95: QM9 reports Pass@95 based on QED and DNA reports Pass@95 based on developmental activity. Protein is evaluated with Success@1Å and motif RMSD. We report additional metrics as appropriate: for QM9, the average QED of valid generated molecules and validity (percentage of valid molecules); for DNA, the average developmental activity score (Dev); and for Protein, mean motif RMSD and Success@1Å, the percentage of generated configurations whose RMSD is below 1.0Å. For editing, we report the reward of the edited configuration and the total computational cost of making the edits. For tasks with learned verifiers, we also report adjusted NFE, which accounts for the total cost of verifier and base-model calls. Further details about the experiments are given in Section F. BoN Base
Generation
Figure 4 and Table 1 evaluate MDM-VGB and MDM-VGBMomentum as generative mechanisms under matched compute. Overall, the main trend is that MDM-VGB and MDM-VGBMomentum improve the quality–cost frontier, especially on tasks with learned verifiers. On QM9 and DNA, MDM-VGBMomentum reaches high Pass@95 at substantially smaller NFE than B O N; this trend largely holds even after FLOP-adjusting for verifier calls. Protein has a more competitive B O N frontier at large rollout budgets, but MDM-VGB-Momentum still gives the best success rate (Success@1Å) under matched budget (see Table 1). 11
Unique Pass@95 Ratio (%)
5.1
MDM-VGR MDM-VGB
MDM-VGB-Momentum
60
16
40
8
20
1 0
25
2 26
4
27
28
29
210
FLOP-adjusted NFE / sample
Figure 5: QM9 unique Pass@95: fraction of distinct generated molecules satisfying Pass@95 versus per-sample compute.
LeafMDM-VGB-Momentum VGB-Momentum NoNo repair Leaf repair
QM9
RootMDM-VGB VGB Root VGB-Momentum Root Root MDM-VGB-Momentum
DNA
800
1250
600
1000
Protein 1000
500
750
400
Pass@95 (%)
100
Pass@95 (%)
BoN BoN(N=16) (N=16)
50
0
500
0
100
100
Success@1A (%)
Adjusted NFE / sample
Leaf VGB Leaf MDM-VGB
50
0 60 65 70 75 80 85 90 Seed reward percentile α
50
0
60 65 70 75 80 85 90 Seed reward percentile α
60 65 70 75 80 85 90 Seed reward percentile α
Figure 6: Editing on QM9, DNA, and Protein. Initial configurations are grouped by seed reward percentile α, i.e., base samples that fail Pass@95 but pass Pass@α. Larger α indicates higher-reward initial outputs. Top: adjusted NFE per sample (↓). Bottom: edited quality (↑). Dashed lines show the reward of the initial configuration (no-edit) and of baselines that discard the initial configurations and generate from scratch: BoN, root-start MDM-VGB, and root-start MDM-VGB-Momentum. Task
Metric
NoEdit
Generation (Root start)
Editing (Leaf start)
Gain
BoN MDM-VGB MDM-VGB-Momentum MDM-VGB MDM-VGB-Momentum QM9 Pass@95 DNA Pass@95 Protein Success@1Å
0.0 0.0 7.4
41.0 50.0 73.7
61.3 41.1 62.0
90.6 70.5 82.0
87.5 69.0 97.7
92.5 88.5 95.6
1.9 pp 18.0 pp 15.7 pp
Table 2: Comparison of generation and editing methods on QM9, DNA, and protein-design tasks. Editing methods are initialized from the same low-reward configurations, whereas generation methods discard these initial configurations and generate new samples from scratch. Leaf-start entries are averaged over seed-leaf reward percentile groups α ∈ {60, 65, . . . , 90}. BoN uses N = 16 independent rollouts, and the MDM-VGB variants use up to the corresponding matched budget. Gain denotes the absolute percentage-point improvement of the best editing method over the best non-editing baseline. For Sudoku, a popular constraint-satisfaction benchmark, MDM-VGB-Momentum reaches perfect accuracy with much lower compute than the baselines, whereas MDM-VGB performs slightly worse. For QM9, Figure 5 plots the fraction of generated samples that are both unique and satisfy Pass@95 versus compute. The MDM-VGB variants still achieve a higher fraction of high-reward molecules at similar FLOP-adjusted NFE, suggesting that their gains are not simply due to repeatedly sampling the same few high-reward molecules. 5.2
Editing
Next, we evaluate MDM-VGB and MDM-VGB-M OMENTUM as editing procedures. We initialize the Markov chain at a completed low-reward output y0 ∈ Y sampled from the reference model. All editing methods start with the same y0 ∈ Y. N O E DIT returns y0 , B O N discards y0 and samples fresh completions from πref , and root-start MDM-VGB variants restart from the fully masked configuration under the same maximum budget. Leaf-start MDM-VGB variants instead re-mask selected coordinates of y0 and return the best full-configuration observed within the allowed edit budget. Figure 6 and Table 2 show that using MDM-VGB and MDM-VGB-M OMENTUM as editing procedures improves over non-editing procedures that generate samples from scratch. Across the three scientific tasks, MDM-VGB and MDM-VGB-M OMENTUM as editing procedures match or improve output quality relative to the best non-editing baseline, with gains of 1.9, 18.0, and 15.7 percentage points on QM9, DNA, and Protein, respectively. The required adjusted NFE generally decreases as the initial reward percentile α increases, indicating that the editing methods make fewer edits when starting closer to the targeted high-reward region. 12
re-mask suffix AR
(
)
)
)
[
]
(
MDM
(
)
)
)
[
]
(
)
M
M
M
M
)
M
)
[
]
AR
re-mask coordinate
Figure 7: MDM-VGB can edit an early mistake directly, while AR-VGB must erase the suffix. The red token marks the local error, and gray tokens denote positions selected for re-masking.
MDM
Method
Acc. ↑
Cost ↓
Acc. ↑
Cost ↓
VGB +Momentum
25.64% 80.92%
127.63 99.41% 75.50 97.54%
29.98 26.24
Table 3: Dyck grammar editing. Cost is the average number of forward and backward moves. MDM-VGB variants achieve higher accuracy with fewer edit moves than AR-VGB variants.
Figure 7 and Table 3 show the advantage of MDM-VGB and MDM-VGB-M OMENTUM over the fixed-order backtracking baselines AR-VGB and AR-VGB-M OMENTUM as editing procedures. On Dyck grammar editing, MDM-VGB reaches 99.41% accuracy with an average cost of 29.98 moves, compared with 25.64% accuracy and 127.63 moves for AR-VGB. Momentum reduces edit cost for both AR and MDM variants: MDM-VGB-M OMENTUM achieves 97.54% accuracy with 26.24 moves, while AR-VGB-M OMENTUM reaches 80.92% accuracy with 75.50 moves. For MDM-VGB, momentum trades a small drop in accuracy for fewer moves, while the any-order variant remains substantially more efficient overall. Component Ablations
AOAR-VGB+Momentum
We highlight two main ablations: one on the re-masking parameter λ, which controls how strongly re-masking decisions are guided by reward values, and one on verifier model size. Additional ablations, including block size and shortlisting rule, are provided in Section F.4.
accuracy (%)
100
accuracy moves
95
50 45 40
90
35 85
30
80
25
mean repair moves
AOAR-VGB
5.3
20 0 2 4
Re-masking parameter λ. We ablate the re-masking parameter λ on Dyck grammar editing. We observe that moderate values of λ give the best accuracy–efficiency tradeoff, because the algorithm can focus on editing problematic tokens while retaining enough flexibility to explore alternative moves. Very large λ makes the algorithm more deterministic, producing shorter trajectories but lowering accuracy; see Figure 8.
8
16
32
geometric exponent λ
Figure 8: Dyck grammar editing with varying re-masking strength λ. Moderate re-masking strength improves editing accuracy while reducing the average number of moves. 100 2.0M
1.5M
Pass@95 (%)
1.0M
0.5M
Pass@95 FLOPs ratio
6
Verifier / Base FLOPs (%)
95 Verifier size and amortized cost. We ablate the size of 3 90 the learned DNA verifier in Figure 9. Increasing verifier 85 2 size from 0.5M to 2.0M parameters improves Pass@95, 80 while adding less than 3% overhead relative to the base 75 1 generator FLOPs. This suggests that a stronger verifier can 70 provide more useful guidance for deciding which tokens to 65 0 0.5M 1.0M 1.5M 2.0M reveal or re-mask. To keep this guidance inexpensive, we Verifier parameters (M) evaluate only a small shortlisted set of candidate edits: for each selected coordinate, we sample K candidate tokens Figure 9: DNA verifier size tradeoff. from the reference model, remove duplicates, and batch Larger verifiers improve Pass@95 at the cost of additional verifier inference. the verifier calls.
Conclusion and Future Work
We introduced MDM-VGB, a reward-guided backtracking sampler for masked diffusion and any-order discrete generation. By extending value-guided backtracking from fixed prefix-tree to masked-state graph, MDM-VGB can reveal and re-mask arbitrary coordinates, allowing it to effectively repair faulty tokens. Theoretically, MDM-VGB efficiently target the reward-tilted law on full configurations and tolerate process verifiers noise. Empirically, MDM-VGB and MDM-VGB-Momentum improve the quality–cost frontier over baselines such as best-of-N and action level sampling (MDM-VGR) for 13
reward-tilting generation and editing across scientific-design and constraint-satisfaction tasks, with especially large gains when the base model performs poorly. A key limitation of our work lies in the cost of training high-quality process verifiers, especially in domains with long sequence length and large vocabulary. We plan to improve verifier training through distillation, or to directly learn the MDM-VGB transition weights via fine-tuning. On the inference-time side, we plan to incorporate hierarchical update and combine our method with parallel Monte Carlo sampling Lee et al. [14]. We hope these improvements will make MDM-VGB applicable on larger reasoning and coding benchmarks such as GSM8K, MATH500, HumanEval, and SWE-bench.
References [1] Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex Lu, Nicolo Fusi, Ava Amini, and Kevin Yang. Protein generation with evolutionary diffusion: sequence is all you need. In NeurIPS 2023 Generative AI and Biology (GenBio) Workshop, 2023. 10, 70 [2] Marianne Arriola, Aaron Gokaslan, Justin T. Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregressive and diffusion language models. In The Thirteenth International Conference on Learning Representations, 2025. 3 [3] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, volume 34, pages 17981–17993, 2021. 3, 58 [4] Richard Bickerton, Gaia Paolini, Jérémy Besnard, Sorel Muresan, and Andrew Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, 4:90–8, 02 2012. doi: 10.1038/nchem.1243. 10, 69 [5] Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.21787, 2024. 1, 10, 19 [6] Noam Chomsky and Marcel P. Schützenberger. The algebraic theory of context-free languages. In Computer Programming and Formal Systems, pages 118–161. North-Holland, 1963. 68 [7] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. 19 [8] Hugo Dalla-Torre, Liam Gonzalez, Javier Mendoza-Revilla, Nicolas Lopez Carranza, Adam Henryk Grzywaczewski, Francesco Oteri, Christian Dallago, Evan Trop, Bernardo P. de Almeida, Hassan Sirelkhatim, Guillaume Richard, Marcin Skwark, Karim Beguir, Marie Lopez, and Thomas Pierrot. Nucleotide transformer: building and evaluating robust foundation models for human genomics. Nature Methods, 22:287–297, 2025. doi: 10.1038/ s41592-024-02523-z. 10, 69 [9] Bernardo P. de Almeida, Franziska Reiter, Michaela Pagani, and Alexander Stark. DeepSTARR predicts enhancer activity from DNA sequence and enables the de novo design of synthetic enhancers. Nature Genetics, 54:613–624, 2022. 10, 69 [10] dLLM Hub. Qwen3-0.6b-diffusion-mdlm-v0.1. https://huggingface.co/dllm-hub/ Qwen3-0.6B-diffusion-mdlm-v0.1, 2026. 10, 68 [11] Thomas P. Hayes and Alistair Sinclair. Liftings of tree-structured markov chains. In Approximation, Randomization, and Combinatorial Optimization. Algorithms and Techniques, pages 602–616. Springer, 2010. 2, 7, 8, 46 [12] Audrey Huang, Adam Block, Qinghua Liu, Nan Jiang, Akshay Krishnamurthy, and Dylan J. Foster. Is best-of-N the best of them? coverage, scaling, and optimality in inference-time alignment. arXiv preprint arXiv:2503.21878, 2025. 10, 19 14
[13] Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, and Sitan Chen. Fine-tuning masked diffusion for provable self-correction. arXiv preprint arXiv:2510.01384, 2025. 69 [14] Jihoon Lee, Hoyeon Moon, Kevin Zhai, Arun Kumar Chithanar, Anit Kumar Sahu, Soummya Kar, Chul Lee, Souradip Chakraborty, and Amrit Singh Bedi. Test-time scaling in diffusion llms via hidden semi-autoregressive experts. arXiv preprint arXiv:2510.05040, 2025. 3, 14 [15] Sanghyun Lee, Sunwoo Kim, Seungryong Kim, Jongho Park, and Dongmin Park. Effective testtime scaling of discrete diffusion through iterative refinement. arXiv preprint arXiv:2511.05562, 2025. 1, 3 [16] David A. Levin and Yuval Peres. Markov Chains and Mixing Times, volume 107. American Mathematical Society, 2017. 7, 28, 32, 38, 39, 44 [17] Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri 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, 2024. 1 [18] Vladimir N. Maiorov and Gordon M. Crippen. Significance of root-mean-square deviation in comparing three-dimensional structures of globular proteins. Journal of Molecular Biology, 235(2):625–634, 1994. ISSN 0022-2836. doi: https://doi.org/10.1006/jmbi.1994.1017. URL https://www.sciencedirect.com/science/article/pii/S0022283684710175. 10 [19] Kou Misaki and Takuya Akiba. Unmaskfork: Test-time scaling for masked diffusion via deterministic action branching. arXiv preprint arXiv:2602.04344, 2026. 1, 3 [20] Zijing Ou, Chinmay Pani, and Yingzhen Li. Inference-time scaling of discrete diffusion models via importance weighting and optimal proposal design. arXiv preprint arXiv:2505.22524, 2025. 3 [21] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2023. 10, 69 [22] Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1:140022, 2014. 1, 9, 69 [23] Dhruv Rohatgi, Abhishek Shetty, Donya Saless, Yuchen Li, Ankur Moitra, Andrej Risteski, and Dylan J Foster. Taming imperfect process verifiers: A sampling perspective on backtracking. arXiv preprint arXiv:2510.03149, 2025. 1, 2, 3, 4, 5, 6, 7, 8, 10, 19, 24, 32, 46 [24] Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems, volume 37, 2024. 3, 4, 10, 58, 69 [25] Alistair Sinclair and Mark Jerrum. Approximate counting, uniform generation and rapidly mixing markov chains. Information and Computation, 82(1):93–133, 1989. 1, 2, 3, 5, 6, 8 [26] Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference-time scaling. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. 1, 3 [27] Kaiwen Wang, Jin Peng Zhou, Jonathan Chang, Zhaolin Gao, Nathan Kallus, Kianté Brantley, and Wen Sun. Value-guided search for efficient chain-of-thought reasoning. In Advances in Neural Information Processing Systems, 2025. 1 [28] Peiyi Wang, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce LLMs step-by-step without human annotations. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 9426–9439, Bangkok, Thailand, 2024. Association for Computational Linguistics. 1 15
[29] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023. 1 [30] Ruidong Wu, Fan Ding, Rui Wang, Rui Shen, Xiwen Zhang, Shitong Luo, Chenpeng Su, Zuofan Wu, Qi Xie, Bonnie Berger, Jianzhu Ma, and Jian Peng. High-resolution de novo structure prediction from primary sequence. bioRxiv preprint, 2022. 10, 70 [31] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang, Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yuqiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. 10, 69 [32] Kevin Yang and Dan Klein. FUDGE: Controlled text generation with future discriminators. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 3511–3535, Online, 2021. Association for Computational Linguistics. 1 [33] Zhao Yang, Hengchang Liu, Chuan Cao, and Bing Su. D3LM: A discrete DNA diffusion language model for bidirectional DNA understanding and generation, 2026. URL https: //arxiv.org/abs/2603.01780. 10, 69
16
TABLE OF C ONTENTS 1
Introduction
1
2
Other related works
3
3
Preliminaries
3
4
Methods and Theoretical Guarantees 4.1 Formulation of MDM-VGB . . . . . . . . . . . . . . . . . . . . . . . . . 4.2 Asymptotic behavior: properties of the stationary distribution . . . . . . . 4.3 Formulation of MDM-VGB-Momentum . . . . . . . . . . . . . . . . . .
5 5 6 8
5
Experiments 5.1 Generation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5.2 Editing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5.3 Component Ablations . . . . . . . . . . . . . . . . . . . . . . . . . . . .
9 11 12 13
6
Conclusion and Future Work
13
A Table of Key Notation and Baseline Algorithms A.1 Baseline Algorithms . . . . . . . . . . . . . . . . . . . . . . . . . . . .
18 19
B Shared Setup and Verifier Training
23
C AOAR-VGB C.1 Setup and Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.2 Primitive AOAR-VGB . . . . . . . . . . . . . . . . . . . . . . . . . . . C.3 Balanced AOAR-VGB . . . . . . . . . . . . . . . . . . . . . . . . . . . C.4 Theoretical Framework and Inference Algorithms . . . . . . . . . . . . . C.5 Geometric Balanced AOAR-VGB . . . . . . . . . . . . . . . . . . . . .
24 24 26 29 34 39
D Momentum-Balanced AOAR-VGB D.1 Setup and Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.2 Momentum-Balanced AOAR-VGB . . . . . . . . . . . . . . . . . . . . . D.3 Momentum Lift for Geometric Balanced AOAR-VGB . . . . . . . . . . . D.4 Theoretical Analysis: Depth-Speed Gain with Momentum . . . . . . . .
46 47 48 51 53
E Balanced Parallel VGBs for Masked Diffusion E.1 Balanced MDM-VGB . . . . . . . . . . . . . . . . . . . . . . . . . . . . E.2 Geometric Balanced MDM-VGB . . . . . . . . . . . . . . . . . . . . . . E.3 Momentum Framework for Balanced and Geometric MDM-VGB . . . .
58 58 62 65
F Experimental Setting and Supplementary Results F.1 Task Descriptions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.2 Leaf-start repair protocol . . . . . . . . . . . . . . . . . . . . . . . . . . F.3 Verifier Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.4 Additional Ablations . . . . . . . . . . . . . . . . . . . . . . . . . . . .
67 68 70 70 71
17
A
Table of Key Notation and Baseline Algorithms Table 4: Table of Key Notations Symbol
Definition
Descriptions
x n V Y Z R (z) k (z) C (z) πref (y | x) τ (x, y) π ⋆ (y | x) Mref (x, z) V ⋆ (x, z) U ⋆ (x, z) Vb (x, z) z j←a z −i sk λ d χ
conditioning context sequence length finite vocabulary Vn n (V ∪ {[mask]}) {i ∈ [n] : zi ̸= [mask]} |R (z)| {y ∈ Y : yi = zi , i ∈ R (z)} base sequence model nonnegative reward ∝ πref (y | x) τ (x, y) P y∈C(z) πref (y | x) Eπ [τ (x, Y ) | x, Y ∈ C (z)] P ref y∈C(z) πref (y | x) τ (x, y) learned approximation of V ⋆ (x, z) single-site reveal update single-site re-mask update depth-dependent edge coefficient geometric value-gating exponent d ∈ {↑, ↓} cancellation strength
Prompt or source input Number of token positions Token alphabet Fully revealed sequences Masked state space Observed coordinates of z Depth of masked state z Compatible full completions of z Reference distribution Leaf-level verifier or reward Tilted target distribution Base completion mass Conditional-expectation value Tilted completion mass Practical plug-in verifier Fill masked coordinate j with token a Re-mask observed coordinate i Balanced AOAR-VGB coefficient Strength of geometric remasking gate Momentum direction state Flow-cancelled momentum parameter
B
subset of coordinates with r = |B| n−r −1
Reveal or re-mask block
sk,r z B←aB z −B
k
block reveal update block re-mask update
Block-update depth coefficient Fill block B with assignment aB Mask every coordinate in block B
Throughout the appendix, AOAR-VGB denotes the singleton-block any-order case, while MDM-VGB denotes the block-update method used in the main text. Table 5: Overview of baseline methods. The final column marks whether the theoretical sampler has the exact leaf-conditioned target π ⋆ (y | x) ∝ πref (y | x)τ (x, y). Here △ denotes conditional exactness: BoN matches the success-conditioned target when τ ∈ {0, 1}, while VGR is exact only when Vb = V ⋆ . Exactness here refers to the ideal asymptotic sampler with full-neighborhood transitions. Method Base BoN VGR VGB VGB-Momentum
Uses Vb ?
Backtracking?
Exact π ⋆ law?
X X O O O
X X X O O
X △ △ O O
18
AR-VGB
leaf hitting
[23]
AR-VGB Momentum [23]
AOAR extension
Primitive AOAR-VGB (Section C.2)
leaf-mass dilution (Remark C.1
Balanced AOAR-VGB
value-gated remasking
(Section C.3)
Geometric Balanced AOAR-VGB
leaf hitting
(Section C.5)
& Prop. C.3)
Momentum-Balanced Geometric AOAR-VGB (Section D.3)
MDM extension
Geometric Balanced MDM-VGB
leaf hitting
(Section E.2)
Momentum-Balanced Geometric MDM-VGB (Section E.3)
Figure 10: Roadmap of the VGB variants studied in the appendix. AR-VGB and its momentum variant follow Rohatgi et al. [23]; Primitive AOAR-VGB gives the any-order extension; Balanced AOARVGB addresses leaf-mass dilution; Geometric Balanced AOAR-VGB adds value-gated remasking; momentum-balanced geometric variants target leaf hitting under finite budgets; and Geometric Balanced MDM-VGB gives the block-update extension.
A.1
Baseline Algorithms
We compare MDM-VGB algorithms against baselines that use the same reference model πref . BoN is the terminal-verifier version of the standard best-of-N reranking baseline. VGR keeps the verifierweighted forward reveal rule from VGB, but disables all re-mask (backtracking) moves. In this subsection, Vb (x, z) denotes the learned verifier score used by value-guided baselines on non-leaf states, while the true reward τ (x, y) is used at leaves. Best-of-N with terminal verifier. BoN is the usual verifier-reranking or best-of-N selection baseline [5, 7, 12]. It samples N independent complete candidates from a forward-only reference rollout of πref (· | x), and returns the candidate with the largest terminal verifier value τ (x, y). It therefore tests what can be gained by spending compute on more complete reference samples, without modifying local transitions online. Algorithm 1 Best-of-N with terminal verifier (BoN) (
= parallelizable area)
Require: context x, reference model πref , reward τ , number of samples N 1: Initialize D ← ∅ Parallel candidate generation. For i = 1, . . . , N , independently sample a complete candidate y (i) by a forward-only reference rollout from πref (· | x), and compute si ← τ (x, y (i) ). 2: for i = 1, . . . , N do 3: Add (y (i) , si ) to D 4: end for 5: i⋆ ← argmaxi∈[N ] si ⋆ 6: return y (i )
Value-guided rollout. VGR is a forward-only value-guided baseline. It is the AOAR/MDM analogue of action-level rejection sampling in Rohatgi et al. [23]: the verifier is used online to choose forward reveal actions, but all re-mask moves are disabled. Thus VGR isolates online value-guided forward selection from the stochastic backtracking mechanism that distinguishes VGB. Since the full forward action space can be large, VGR uses the same forward shortlisting approximation as VGB. At state z, it constructs a shortlisted forward candidate set A(z). For a candidate 19
u = z B←aB , with singleton blocks giving the AOAR case, the VGR weight is w(z, u) := sk(z),|B| πref (YB = aB | x, z) Vb (x, u). The geometric factor Vb (x, z)λ is omitted because it is common to all forward candidates from the same state and cancels under normalization. We do not report a separate momentum version of VGR: once backward moves are disabled, the reveal/re-mask direction variable is degenerate. Algorithm 2 Value-guided rollout (VGR) (
= parallelizable area)
Require: context x, initial masked state z0 , reference model πref , verifier score Vb , depth coefficients sk,r , step limit T 1: z ← z0 2: for t = 1, . . . , T do 3: if z ∈ Y then 4: return z 5: end if Parallel forward construction. Construct a shortlisted forward candidate set A(z). For every u ∈ A(z), write u = z B←aB , r = |B|, compute w(u) = sk(z),r πref (YB = aB | x, z) Vb (x, u). 6:
if
P
u∈A(z) w(u) > 0 then
7: Sample u ∈ A(z) proportional to w(u) 8: else 9: Sample u = z B←aB ∈ A(z) proportional to πref (YB = aB | x, z) 10: end if 11: z←u 12: end for 13: if z ∈ / Y then 14: Complete remaining masked coordinates using πref (· | x, z) 15: end if 16: return z
MDM-VGB samplers. We next give the block-update VGB samplers. The first algorithm is the generic MDM-VGB with admissible block sizes M; the second applies the lifted momentum construction to the same local forward and backward weights. Formal definitions of these weights and the stationarity guarantees are given in the following sections. The block sets Br+ (z), Br− (z) and the full MDM neighborhoods CMDM (z), PMDM (z) are defined in Section E.1.
20
Algorithm 3 MDM-VGB sampler with block coefficients, without a self-loop ( area)
= parallelizable
Input: context x, initial state z, reference block conditionals πref (· | x, z), verifier Vb , reward τ , admissible sizes M, coefficients sk,r , geometric exponent λ ≥ 0, step limit step_limit Output: leaf sample list L 1: Initialize L ← ∅ and step ← 0 ▷ Balanced MDM-VGB uses λ = 0; Geometric Balanced MDM-VGB allows λ > 0. 2: while step < step_limit do 3: Set k ← k(z) = |R(z)| 4: Initialize an empty weighted candidate set A 5: if k < n then Parallel forward construction. For all r ∈ M with r ≤ n − k, all A ∈ Br+ (z), and all aA ∈ V A , independently set uA,aA ← z A←aA and add uA,aA to A with weight uA,aA ∈ Y, sn−r,r πref (YA = aA | x, z)τ (x, uA,aA ), sk,r πref (YA = aA | x, z)Vb (x, uA,aA ), uA,aA ∈ / Y, λ = 0, s Vb (x, z)λ π (Y = a | x, z)Vb (x, u / Y, λ > 0. k,r ref A A A,aA ), uA,aA ∈ 6: 7:
end if if k > 0 then Parallel backward construction. For all r ∈ M with r ≤ k and all A ∈ Br− (z), independently set uA ← z −A and add uA to A with weight z ∈ Y, sn−r,r τ (x, z), sk−r,r Vb (x, z), z∈ / Y, λ = 0, s λ b b / Y, λ > 0. k−r,r V (x, z)V (x, uA ) , z ∈
8: end if 9: Sample u ∈ A proportional to its weight 10: Set z ← u 11: if z ∈ Y then 12: Append z to L 13: end if 14: step ← step + 1 15: end while 16: return L
21
Algorithm 4 Momentum MDM-VGB sampler (
= parallelizable area;
= momentum area)
Input: context x, initial lifted state (z, d) with d ∈ {↓, ↑}, base local weights wfwd , wbwd for balanced or geometric MDM-VGB, cancellation strength χ ∈ [0, 1], step limit step_limit Output: leaf sample list L 1: Initialize L ← ∅ and step ← 0 2: while step < step_limit do 3: Set k ← k(z) = |R(z)| 4: if k = 0 and d =↑ then 5: Set d ←↓ ROOT D IRECTION R ESET 6: end if 7: if k = n and d =↓ then 8: Set d ←↑ L EAF D IRECTION R ESET 9: end if 10: Initialize empty weighted candidate sets Afwd , Abwd , and Ae 11: if k < n then Parallel forward construction. For all c ∈ CMDM (z), independently add c to Afwd with weight wfwd (z → c). 12: end if 13: if k > 0 then Parallel backward construction. For all p ∈ PMDM (z), independently add p to Abwd with weight wbwd (z → p). 14: end if Momentum construction. Set X F (z) = wfwd (z → c), c∈Afwd
B(z) =
X
wbwd (z → p),
p∈Abwd
M (z) = min {F (z), B(z)} . Populate Ae with the following lifted-state weights: mode
lifted state
weight
transition mode
d =↓ d =↓ d =↓
(c, ↓), c ∈ Afwd (z, ↑) (z, ↓)
wfwd (z → c) B(z) − χM (z) χM (z)
M OVE S WITCH U P S TAY D OWN ( OPTIONAL )
d =↑ d =↑ d =↑
(p, ↑), p ∈ Abwd (z, ↓) (z, ↑)
wbwd (z → p) F (z) − χM (z) χM (z)
M OVE S WITCH D OWN S TAY U P ( OPTIONAL )
15: Sample (u, d′ ) ∈ Ae proportional to its weight 16: Set (z, d) ← (u, d′ ) 17: if z ∈ Y then 18: Append z to L 19: end if 20: step ← step + 1 21: end while 22: return L
22
B
Shared Setup and Verifier Training
We first collect the notation shared by all VGB variants. Fix a conditioning context x ∼ ρ, a sequence length n, and a finite vocabulary V. The fully revealed state space and the masked state space are Y := V n ,
n
Z := (V ∪ {[mask]}) ,
where [mask] ∈ / V is the masking symbol. For z ∈ Z, let R(z) := {i ∈ [n] | zi ̸= [mask]} ,
k(z) := |R(z)| .
Here R(z) is the set of revealed coordinates and k(z) is the depth of z. The compatible completions of z are C(z) := {y ∈ Y | yi = zi for all i ∈ R(z)} . Equivalently, observing the partial state z restricts the unknown terminal sequence Y to the completion set C(z). As more coordinates are revealed, this set becomes more informative. Tilted target. Let πref (y | x) be a reference sequence model and let τ (x, y) ≥ 0 be a terminal reward. These two objects define the tilted target distribution π ∗ (y | x) :=
πref (y | x)τ (x, y) , Z(x)
Z(x) :=
X
πref (y | x)τ (x, y).
y∈Y
The goal of the VGB procedures is to sample from π ∗ (· | x) while using πref as the proposal backbone. Completion values. To describe local moves on masked states, we need to know both how much reference probability remains under a partial assignment and how promising that partial assignment is under the reward. The first quantity is the reference completion mass X πref (y | x) = PY ∼πref (Y ∈ C(z) | x). Mref (x, z) := y∈C(z)
This mass determines which masked states are reachable under the reference model: Zreach (x) := {z ∈ Z : Mref (x, z) > 0} . If z ∈ Zreach (x), we say that z is πref -reachable. When the context is fixed, we write Zreach for Zreach (x). The second quantity is the conditional-expectation value V ∗ (x, z) := Eπref [τ (x, Y ) | x, Y ∈ C(z)] , which is the object approximated by the verifier. For the analysis, it is also convenient to introduce the unnormalized tilted completion mass X U ∗ (x, z) := πref (y | x)τ (x, y). y∈C(z)
We note that these two views are related by U ∗ (x, z) = Mref (x, z)V ∗ (x, z). The quantity U ∗ is useful for stationarity proofs, but learning it directly would also require estimating the completion mass Mref . The VGB framework discussed in later sections avoids this extra estimation problem: its local transitions can be implemented using reference conditionals together with a learned approximation Vb of V ∗ . 23
Verifier training. When a learned verifier is used, we train a masked-state verifier Vb (x, z). During training, we fit a parametric model Vθ (x, z) and then set Vb to the trained model. For a sampled context x, we run several reference rollouts, each producing a terminal completion y. Along the way, the sampler reveals coordinates, producing a trajectory of intermediate masked states. We subsample several snapshots from each trajectory and assign all of them the same terminal label r = τ (x, y). The verifier is trained with the Monte Carlo (MC) rollout regression objective h i 2 Lreg (θ) := Ex,z,y (Vθ (x, z) − τ (x, y)) , where z is a sampled intermediate state from a rollout that terminates at y. The following proposition reveals the population target of this regression. Proposition B.1 (MC regression recovers V ⋆ ). Consider the population rollout law Pref over triples (x, z, y). Suppose this law is reference-consistent: whenever z ∈ Zreach (x), Pref (Y = y | x, z) = πref (y | x, Y ∈ C(z)). Then the population regression objective Lreg (V ) := EPref
h
V (x, z) − τ (x, Y )
2 i
recovers V ⋆ as its population optimum: V ⋆ = argmin Lreg (V )
Pref -a.e.
V
Proof. For fixed (x, z), the squared loss is minimized by the conditional mean EPref [τ (x, Y ) | x, z] . Therefore, reference consistency gives EPref [τ (x, Y ) | x, z] = Eπref [τ (x, Y ) | x, Y ∈ C(z)] = V ⋆ (x, z), which proves the claim. By Proposition B.1, under a reference-consistent population rollout law Pref , the population target of MC rollout regression is exactly V ⋆ (x, z). Thus, Vb plays the role of an intermediate-state value function, analogous to the prefix-state value function in Rohatgi et al. [23], but lifted here to any-order partial states. The training procedure is summarized in Algorithm 5.
C
AOAR-VGB
We now extend the VGB perspective of Rohatgi et al. [23] from left-to-right prefix trees to any-order masked states. The key difference is that the state is no longer a prefix. Instead, a state is a masked partial sequence z ∈ Z, and a forward move reveals one currently masked coordinate. A backward move re-masks one currently revealed coordinate. This gives a natural random walk on the graph of masked states. C.1
Setup and Notation
We use the shared notation from the previous section. In particular, for a masked state z ∈ Z, R(z) denotes the set of revealed coordinates, k(z) = |R(z)| denotes its depth, and C(z) denotes the set of full sequences compatible with z. We write ∅ for the fully masked state, so R(∅) = ∅. The fully revealed states are exactly the leaves Y = V n. For an unrevealed coordinate j ∈ / R(z) and token a ∈ V, define the single-site reveal update z j←a by a, i = j, z j←a i = zi , i ̸= j. 24
Algorithm 5 Training the shared masked-state verifier Require: context distribution ρ, reference model πref , reward τ , parametric verifier Vθ , snapshots per rollout K, number of rollouts N 1: Initialize snapshot dataset Dsnap ← ∅ 2: for i = 1, . . . , N do 3: Sample x ∼ ρ 4: Run a reference rollout using πref (· | x) to obtain a terminal completion y and snapshots (1) z , . . . , z (K) 5: Compute r ← τ (x, y) K 6: Add (x, z (k) , r) k=1 to Dsnap 7: end for 8: Update θ by minimizing X 1 2 θ ← argmin (Vθ′ (x, z) − r) ′ |D | snap θ (x,z,r)∈Dsnap
b (x, z) ← Vθ (x, z) 9: Define V b 10: return trained verifier V For a revealed coordinate i ∈ R(z), define the single-site re-masking update z −i by [mask], ℓ = i, z −i ℓ = zℓ , ℓ ̸= i. The AOAR neighbor sets are CAOAR (z) := z j←a | j ∈ / R(z), a ∈ V, Mref (x, z j←a ) > 0 , −i PAOAR (z) := z | i ∈ R(z) ,
(Child neighborhood) (Parent neighborhood)
and NAOAR (z) := CAOAR (z) ∪ PAOAR (z). (AOAR neighborhood) Thus, CAOAR (z) contains one-step reveal states, and PAOAR (z) contains one-step re-masked states. The AOAR oracle assumption is that for every reachable state z and every j ∈ / R(z), we can evaluate the single-coordinate reference conditional πref (Yj = a | x, z) := Pπref (Yj = a | x, Y ∈ C(z)) ,
a ∈ V.
Equivalently, πref (Yj = a | x, z) =
Mref (x, z j←a ) . Mref (x, z)
We first record the basic identities that will be used in the weighted-graph construction and the stationary analysis below. Lemma C.1 (Single-site partition identity). Fix a reachable state z and coordinate j ∈ / R(z). Then G C(z) = C(z j←a ). a∈V
Consequently, X
X
U ∗ (x, z j←a ) = U ∗ (x, z),
a∈V
Mref (x, z j←a ) = Mref (x, z).
a∈V
Proof. Every full sequence y ∈ C(z) has a unique value yj = a at coordinate j. Hence y belongs to exactly one set C(z j←a ). This proves the disjoint union. Summing the masses πref (y | x)τ (x, y) over this partition gives the identity for U ∗ . Setting τ ≡ 1 gives the identity for Mref .
25
Proposition C.1 (AOAR Bellman identity). For every reachable state z and every j ∈ / R(z), X V ∗ (x, z) = πref (Yj = a | x, z)V ∗ (x, z j←a ). a∈V
Proof. By the law of total expectation, V ∗ (x, z) = Eπref [τ (x, Y ) | x, Y ∈ C(z)] . Conditioning further on the value of Yj , we obtain X V ∗ (x, z) = Pπref (Yj = a | x, Y ∈ C(z))Eπref [τ (x, Y ) | x, Y ∈ C(z), Yj = a] . a∈V
The event {Y ∈ C(z), Yj = a} is exactly Y ∈ C(z j←a ) . Therefore the inner conditional expectation is V ∗ (x, z j←a ). Proposition C.2 (Exact target conditional). For every reachable state z with U ∗ (x, z) > 0, every j∈ / R(z), and every a ∈ V, π ∗ (Yj = a | x, z) =
U ∗ (x, z j←a ) V ∗ (x, z j←a ) = πref (Yj = a | x, z) . ∗ U (x, z) V ∗ (x, z)
Proof. By definition of conditional probability under the tilted target, π ∗ (Yj = a | x, z) =
Pπ∗ (Y ∈ C(z j←a ) | x) . Pπ∗ (Y ∈ C(z) | x)
Since Pπ∗ (Y ∈ C(z j←a ) | x) = and Pπ∗ (Y ∈ C(z) | x) =
U ∗ (x, z j←a ) Z(x) U ∗ (x, z) , Z(x)
the first equality follows. For the second equality, use U ∗ (x, z) = Mref (x, z)V ∗ (x, z) and
C.2
Mref (x, z j←a ) = Mref (x, z)πref (Yj = a | x, z).
Primitive AOAR-VGB
The most direct AOAR extension of AR-VGB uses the same local rule as the prefix-tree AR-VGB walk: forward moves are weighted by the reference conditional times the verifier value of the child, and backward moves are weighted by the verifier value of the current state. Let Vb (x, z) denote the trained verifier. We use terminal anchoring: whenever the terminal reward is available at a leaf, the verifier is fixed to the true reward, Vb (x, y) := τ (x, y),
y ∈ Y.
Definition C.1 (γ-held primitive AOAR-VGB local kernel). Fix a holding parameter γ ≥ 0. For a reachable state z, define the primitive AOAR-VGB local weights by F ORWARD : BACKWARD :
w(z → z j←a ) := πref (Yj = a | x, z) Vb (x, z j←a ), w(z → z −i ) := Vb (x, z),
26
j∈ / R(z), i ∈ R(z).
Let
X
WPM (z) :=
w(z → v).
v∈NAOAR (z)
The associated non-self move kernel is KPM (z, u) :=
w(z → u) , WPM (z)
u ∈ NAOAR (z).
The γ-held primitive AOAR-VGB kernel is (γ)
PPM (z, z) =
γ , 1+γ
and, for u ∈ NAOAR (z), (γ)
PPM (z, u) =
1 w(z → u) 1 P KPM (z, u) = . 1+γ 1 + γ v∈NAOAR (z) w(z → v)
When γ = 1, this recovers the usual 1/2-lazy kernel. When γ = 0, it is the non-lazy local random walk. Although Mref is not needed to implement the transition, it is useful to temporarily reintroduce it in the analysis. Multiplying the learned value by the reference completion mass Mref gives the plug-in tilted mass b (x, z) := Mref (x, z)Vb (x, z). U Theorem C.1 (Weighted-graph representation and exact leaf law). Construct an undirected graph on reachable masked states by connecting z and z j←a whenever j ∈ / R(z) and z j←a is reachable. Assign the symmetric edge weight b (x, z j←a ) = Mref (x, z j←a )Vb (x, z j←a ). fPM (z, z j←a ) := U For a state z, let its weighted degree be DPM (z) :=
X
fPM (z, u),
u∼z
where u ∼ z means that u and z are connected by an edge in the AOAR weighted graph. (γ)
Then, for every γ ≥ 0, the γ-held primitive AOAR-VGB kernel PPM from Definition C.1 is exactly the γ-held weighted random walk on this graph. Also, its stationary distribution is DPM (z) . ′ z ′ DPM (z )
µPM (z) = P
Moreover, if Vb (x, y) = τ (x, y) on all leaves y ∈ Y, then the stationary law conditioned on the leaves is exactly the tilted target: µPM (y | y ∈ Y) = π ∗ (y | x). Proof. Fix a reachable state z. For a forward neighbor z j←a , we have fPM (z, z j←a ) Mref (x, z j←a ) b = V (x, z j←a ) = πref (Yj = a | x, z)Vb (x, z j←a ). Mref (x, z) Mref (x, z) This is exactly the forward local weight. For a backward neighbor z −i , the edge between z −i and z has weight b (x, z) = Mref (x, z)Vb (x, z). fPM (z −i , z) = U Dividing by the same state-dependent factor Mref (x, z) gives the backward local weight Vb (x, z). 27
Thus, from every state z, the local weights in Definition C.1 are exactly the incident edge weights divided by the common positive factor Mref (x, z). This common factor cancels under normalization. Therefore the non-self move kernel is the weighted random walk on the graph with edge weights fPM : fPM (z, u) KPM (z, u) = , u ∼ z. DPM (z) This is the standard random walk on a network with conductances fPM (z, u) in the sense of Levin and Peres [16]. We now verify stationarity directly. Let Dtot :=
X
DPM (z ′ ),
µPM (z) :=
z′
DPM (z) . Dtot
For two distinct neighboring states z ∼ u, the γ-held kernel satisfies (γ)
PPM (z, u) =
1 fPM (z, u) . 1 + γ DPM (z)
Therefore DPM (z) 1 fPM (z, u) · Dtot 1 + γ DPM (z) fPM (z, u) = . (1 + γ)Dtot
(γ)
µPM (z)PPM (z, u) =
Since fPM (z, u) = fPM (u, z), the same calculation gives (γ)
µPM (u)PPM (u, z) =
fPM (z, u) . (1 + γ)Dtot
Thus detailed balance holds on every non-self edge. For the self-loop terms, (γ)
µPM (z)PPM (z, z) = µPM (z)
γ , 1+γ
(γ)
so detailed balance is immediate. Therefore PPM is reversible with respect to µPM , and hence µPM is stationary. Now consider a leaf y ∈ Y. It has exactly n parents, one for each coordinate that can be re-masked. For each parent y −i , the edge weight is fPM (y −i , y) = Mref (x, y)Vb (x, y) = πref (y | x)τ (x, y). Here we used terminal anchoring. Hence DPM (y) = n πref (y | x)τ (x, y). The factor n is constant across leaves. Therefore µPM (y | y ∈ Y) = P
πref (y | x)τ (x, y) = π ∗ (y | x). πref (y ′ | x)τ (x, y ′ )
y ′ ∈Y
When the verifier is exact, every γ-held primitive AOAR-VGB chain has a particularly simple stationary law. Corollary C.1 (Exact primitive stationary law). For any γ ≥ 0, if Vb = V ∗ , then the stationary (γ) distribution of PPM satisfies µ∗PM (z) ∝ U ∗ (x, z). 28
b = U ∗ . For a state z of depth k = k(z), every backward Proof. When Vb = V ∗ , we have U ∗ edge contributes U (x, z), and there are k such edges. For each unrevealed coordinate j ∈ / R(z), Lemma C.1 gives X U ∗ (x, z j←a ) = U ∗ (x, z). a∈V
There are n − k unrevealed coordinates, so the total forward incident weight is (n − k)U ∗ (x, z). Hence DPM (z) = kU ∗ (x, z) + (n − k)U ∗ (x, z) = nU ∗ (x, z). By Theorem C.1, for every γ ≥ 0, µ∗PM (z) ∝ DPM (z). Therefore µ∗PM (z) ∝ U ∗ (x, z).
However, we note that the same stationary law also exposes a depth-allocation problem: although the conditional law on leaves is correct, the primitive chain assigns exponentially small marginal mass to the leaf layer, as follows: Proposition C.3 (Exponential leaf-mass dilution of primitive AOAR-VGB). For any γ ≥ 0, assume (γ) Vb = V ∗ . Then the stationary mass of depth k under PPM is n µ∗PM (k(z) = k) = In particular,
k
2n
.
µ∗PM (Y) = 2−n .
Proof. For any fixed depth k, double-counting identity gives X X X U ∗ (x, z) = πref (y | x)τ (x, y) z: k(z)=k
z:k(z)=k y∈C(z)
=
X
πref (y | x)τ (x, y)
y∈Y
X z:k(z)=k
n 1 {y ∈ C(z)} = Z(x). k
Indeed, each full sequence y ∈ Y contributes the mass πref (y | x)τ (x, y) to exactly states of depth k, one for each choice of k revealed coordinates. Therefore n n X X X X n U ∗ (x, z) = U ∗ (x, z) = Z(x) = 2n Z(x). k z∈Z
k=0 z:k(z)=k
n k
masked
k=0
By Corollary C.1, the primitive stationary law is proportional to U ∗ , so n n ∗ k k Z(x) µPM (k(z) = k) = n = n. 2 Z(x) 2 Taking k = n gives µ∗PM (Y) = 2−n . Thus, primitive AOAR-VGB has the correct leaf-conditioned distribution, but the stationary chain spends exponentially little mass on leaves. We therefore introduce Balanced AOAR-VGB, a depthbalanced variant designed to resolve this leaf-mass dilution problem. C.3
Balanced AOAR-VGB
Primitive AOAR-VGB treats every single-site edge symmetrically. In the AOAR graph, however, the number of states at each depth grows as nk . As a result, the primitive stationary law puts most of its mass around the middle depths. Balanced AOAR-VGB rescales forward and backward moves as a function of depth to suppress this combinatorial depth bias.
29
Definition C.2 (γ-held depth-rescaled AOAR-VGB local kernel). Fix a holding parameter γ ≥ 0 and positive coefficients β0 , β1 , . . . , βn−1 > 0,
α1 , α2 , . . . , αn > 0.
For a reachable state z with k = k(z), define the depth-rescaled local weights by F ORWARD : BACKWARD :
w(k) (z → z j←a ) := βk πref (Yj = a | x, z)Vb (x, z j←a ), w(k) (z → z −i ) := αk Vb (x, z),
Let WBAL (z) :=
X
j∈ / R(z), i ∈ R(z).
w(k) (z → u).
u∈NAOAR (z)
The associated non-self move kernel is KBAL (z, u) :=
w(k) (z → u) , WBAL (z)
u ∈ NAOAR (z).
The γ-held balanced AOAR-VGB kernel is γ (γ) , PBAL (z, z) = 1+γ 1 1 w(k) (z → u) (γ) P PBAL (z, u) = KBAL (z, u) = , 1+γ 1 + γ v∈NAOAR (z) w(k) (z → v)
u ∈ NAOAR (z).
When γ = 1, this recovers the usual 1/2-lazy kernel. When γ = 0, it is the non-lazy local random walk. In the Proposition C.4 below, we use the word “move” to mean a non-self transition. Thus, conditioning on move means conditioning on the event that the kernel does not take the holding self-loop and instead chooses a neighbor in NAOAR (z). The events forward and backward mean that the selected neighbor lies in CAOAR (z) or PAOAR (z), respectively. Since the holding probability is common to all non-self moves, the conditional directional probabilities below do not depend on γ. Proposition C.4 (Exact directional masses and balanced condition). Assume Vb = V ∗ . Let z be a reachable state of interior depth 1 ≤ k ≤ n − 1 with V ∗ (x, z) > 0. Then the exact total forward and backward local masses are (k),∗
Wfwd (x, z) = (n − k)βk V ∗ (x, z),
(k),∗
Wbwd (x, z) = kαk V ∗ (x, z).
Therefore, conditional on making a non-self move, (n − k)βk , (n − k)βk + kαk kαk P(backward | x, z, move) = . (n − k)βk + kαk P(forward | x, z, move) =
In particular, the forward and backward directions are exactly balanced at depth k if and only if (n − k)βk = kαk . Proof. For each unrevealed coordinate j ∈ / R(z), Proposition C.1 gives X πref (Yj = a | x, z)V ∗ (x, z j←a ) = V ∗ (x, z). a∈V
There are n − k unrevealed coordinates, so the total forward mass is X X (k),∗ Wfwd (x, z) = βk πref (Yj = a | x, z)V ∗ (x, z j←a ) = (n − k)βk V ∗ (x, z). a∈V j ∈R(z) /
30
There are k revealed coordinates, and each backward move has weight αk V ∗ (x, z). Hence X (k),∗ Wbwd (x, z) = αk V ∗ (x, z) = kαk V ∗ (x, z). i∈R(z)
The conditional directional probabilities follow by normalizing the forward and backward masses among non-self moves. The common holding factor 1/(1 + γ) cancels after conditioning on move. Exact balance is therefore equivalent to (n − k)βk V ∗ (x, z) = kαk V ∗ (x, z), or simply (n − k)βk = kαk . Remark C.1 (Primitive AOAR-VGB as the unrescaled case). Primitive AOAR-VGB corresponds to the unrescaled choice βk = 1, αk = 1 at every interior depth. Under this choice, Proposition C.4 gives k n−k , P(backward | x, z, move) = . n n Thus, the primitive chain strongly favors forward moves near the root and strongly favors backward moves near the leaves. P(forward | x, z, move) =
This local depth bias is consistent with the global pathology in Proposition C.3: under the exact primitive stationary law, the depth marginal is binomial and the leaf layer receives only 2−n stationary mass. Balanced AOAR-VGB removes this local directional bias by choosing αk , βk so that (n − k)βk = kαk , making the forward and backward directions equally likely after conditioning on a non-self move. We next give the stationary law for the balanced chain. The following edge-coefficient representation is the cleanest way to see that terminal exactness is preserved even when Vb is imperfect. Let g1 > 0 be arbitrary, and define g2 , . . . , gn recursively by βk gk+1 = gk , k = 1, . . . , n − 1. αk We set β0 = αn = 1. These boundary choices do not affect the normalized transition, since only forward moves are available at the root and only backward moves are available at the leaves. Theorem C.2 (Weighted-graph representation of balanced AOAR-VGB). Construct an undirected graph on reachable masked states with single-site AOAR edges. For an edge between a depth-k parent z and its child z j←a , assign the symmetric edge weight fBAL (z, z j←a ) := gk+1 Mref (x, z j←a )Vb (x, z j←a ). Then, for every γ ≥ 0, the balanced AOAR-VGB kernel from Definition C.2 is exactly the γ-held weighted random walk on this graph. Hence its stationary distribution is X DBAL (z) µBAL (z) = P , D (z) := fBAL (z, u). BAL ′ z ′ DBAL (z ) u∼z If Vb (x, y) = τ (x, y) for all leaves y ∈ Y, then µBAL (y | y ∈ Y) = π ∗ (y | x). Proof. Let z be a state of depth k. First consider a forward edge z → z j←a . Then Mref (x, z j←a ) b fBAL (z, z j←a ) = gk+1 V (x, z j←a ) Mref (x, z) Mref (x, z) = gk+1 πref (Yj = a | x, z)Vb (x, z j←a ). 31
For a backward edge from z to z −i , the child in the edge is z, so fBAL (z −i , z) = gk Vb (x, z). Mref (x, z) For an interior state, the recursion βk gk+1 = gk αk implies gk+1 /βk = gk /αk . Thus, at every interior depth k, the incident edge weights divided by Mref (x, z) are proportional to the local depth-rescaled weights in Definition C.2. The proportionality factor is common across all outgoing neighbors of z, and hence cancels under normalization. At the root and leaves, we use the boundary convention β0 = αn = 1. Since only forward moves are available at the root and only backward moves are available at the leaves, these boundary choices do (γ) not change the normalized transition. Therefore PBAL is exactly the γ-held weighted random walk on the graph with edge weights fBAL : (γ)
PBAL (z, z) =
γ , 1+γ
(γ)
PBAL (z, u) =
1 fBAL (z, u) , 1 + γ DBAL (z)
u ∼ z.
We now verify the stationary distribution directly. This is the standard random walk on a network with conductances fBAL (z, u), for which the stationary distribution is proportional to the vertex conductance DBAL (z) [16]. For completeness, we check detailed balance. Let X DBAL (z) Dtot := DBAL (z ′ ), µBAL (z) := . Dtot ′ z
For two distinct neighboring states z ∼ u, (γ)
µBAL (z)PBAL (z, u) =
fBAL (z, u) DBAL (z) 1 fBAL (z, u) = · . Dtot 1 + γ DBAL (z) (1 + γ)Dtot
Since the edge weight is symmetric, fBAL (z, u) = fBAL (u, z), the same calculation gives (γ)
µBAL (u)PBAL (u, z) =
fBAL (z, u) . (1 + γ)Dtot
Thus, detailed balance holds on every non-self edge. For the self-loop terms, γ (γ) µBAL (z)PBAL (z, z) = µBAL (z) , 1+γ (γ)
so detailed balance is trivial. Therefore PBAL is reversible with respect to µBAL , and hence µBAL is stationary. Now let y ∈ Y be a leaf. Each of its n incident edges has weight gn Mref (x, y)Vb (x, y) = gn πref (y | x)τ (x, y). Thus DBAL (y) = ngn πref (y | x)τ (x, y). The factor ngn is constant over leaves, so conditioning on Y gives πref (y | x)τ (x, y) = π ∗ (y | x). ′ ′ y ′ ∈Y πref (y | x)τ (x, y )
µBAL (y | y ∈ Y) = P
Remark C.2 (Zero holding in the practical sampler). The holding parameter γ is useful for connecting the balanced AOAR-VGB transition to the standard lazy random-walk formalism, but explicit self-loop steps have limited algorithmic value in implementation and directly add computational cost. Following the practical convention in Rohatgi et al. [23], we therefore set γ = 0 in our AOAR-VGB implementation and use the non-lazy local move kernel. The stationarity argument in Theorem C.2
32
is stated for every γ ≥ 0, so this practical choice does not change the leaf-conditioned stationary law. When the verifier is exact, the weighted-degree expression from Theorem C.2 simplifies to the tilted mass U ∗ multiplied by a depth-dependent coefficient. Corollary C.2 (Exact balanced stationary law). Assume Vb = V ∗ . Define d0 := ng1 ,
dn := ngn ,
dk := kgk + (n − k)gk+1
(1 ≤ k ≤ n − 1).
Then, for every γ ≥ 0, the exact balanced stationary law satisfies µ∗BAL (z) ∝ dk(z) U ∗ (x, z). Proof. Let z have depth k. If 1 ≤ k ≤ n − 1, its backward incident edges contribute kgk U ∗ (x, z). For each unrevealed coordinate j ∈ / R(z), Lemma C.1 gives X gk+1 U ∗ (x, z j←a ) = gk+1 U ∗ (x, z). a∈V
There are n − k unrevealed coordinates, so the forward incident weight is (n − k)gk+1 U ∗ (x, z). Thus
∗ DBAL (z) = (kgk + (n − k)gk+1 ) U ∗ (x, z). The root and leaf cases give d0 = ng1 and dn = ngn , respectively. By Theorem C.2, the stationary ∗ distribution is proportional to DBAL for every γ ≥ 0, which proves the claim.
Corollary C.3 (Canonical balanced AOAR-VGB). Set βk = k,
αk = n − k,
1 ≤ k ≤ n − 1.
Then, the exact forward and backward directional probabilities are both 1/2 at every interior depth. Choosing g1 = 1, the corresponding edge coefficients are gk =
1
1 ≤ k ≤ n.
,
n−1 k−1
Moreover, under Vb = V ∗ , for every γ ≥ 0, µ∗BAL (Y) =
1 . 2n
Proof. The balanced condition follows immediately: (n − k)βk = (n − k)k = k(n − k) = kαk . The recursion for gk becomes gk+1 = gk
k . n−k
Starting from g1 = 1, this gives gk =
k−1 Y
r (k − 1)!(n − k)! = = n − r (n − 1)! r=1
1
.
n−1 k−1
For 1 ≤ k ≤ n − 1, dk = kgk + (n − k)gk+1 = kgk + kgk = 2kgk =
2n n . k
33
At the boundaries, d0 = dn = n. Using the double-counting identity X X U ∗ (x, z) = z: k(z)=k
X
πref (y | x)τ (x, y)
z: k(z)=k y∈C(z)
=
X
πref (y | x)τ (x, y)
y∈Y
X
1 {y ∈ C(z)} =
z: k(z)=k
n Z(x), k
we the unnormalized stationary weight layer by layer. For the root layer, P can compute U ∗ (x, z) = Z(x) and d0 = n, so its total stationary weight is nZ(x). For the leaf z: k(z)=0 P layer, z: k(z)=n U ∗ (x, z) = Z(x) and dn = n, so its total stationary weight is also nZ(x). For each interior depth 1 ≤ k ≤ n − 1, the total stationary weight is X 2n n ∗ dk U (x, z) = n Z(x) = 2nZ(x). k k z: k(z)=k
Hence the total weight over all depths is nZ(x) + nZ(x) + (n − 1)2nZ(x) = 2n2 Z(x). The leaf layer has weight nZ(x), so µ∗BAL (Y) =
nZ(x) 1 = . 2n2 Z(x) 2n
Thus, canonical balanced AOAR-VGB keeps the exact target law on leaves while increasing the stationary leaf mass from 2−n to order 1/n. C.4
Theoretical Framework and Inference Algorithms
We now state the mixing framework for canonical balanced AOAR-VGB. The argument follows the same high-level principle as VGB on a prefix tree: stochastic backtracking defines a reversible walk whose leaf-conditioned stationary law is exact, and whose conductance can be controlled under a multiplicative value-approximation assumption. Assumption C.1 (Uniform multiplicative verifier accuracy). There exists κ ≥ 1 such that for every reachable non-leaf state z, κ−1 V ∗ (x, z) ≤ Vb (x, z) ≤ κV ∗ (x, z). At leaves, we use terminal anchoring: Vb (x, y) = τ (x, y),
y ∈ Y.
The canonical balanced AOAR graph can be viewed as an average of fixed-order VGB trees. Let Permn denote the set of coordinate orders, i.e., permutations of [n]. For ω = (ω1 , . . . , ωn ) ∈ Permn , define the fixed-order state set Zω := {z ∈ Z | R(z) = {ω1 , . . . , ωk } for some k ∈ {0, . . . , n}} . This is a tree whose forward move at depth k reveals coordinate ωk+1 . On this tree, define the edge weight fω (z, z ωk+1 ←a ) := Mref (x, z ωk+1 ←a )Vb (x, z ωk+1 ←a ). Lemma C.2 (Fixed-order tree conductance). We call X Dω (z) Dω (z) := fω (z, u), µω (z) := P . ′ z ′ Dω (z ) u∼z
34
the weighted degree and the degree-normalized stationary law, respectively. For any nonempty set A ⊊ Zω , define its set conductance and the tree conductance by P fω (z, u) z∈A, u∈A P / Φω (A) := , Φω := inf Φω (A). A⊆Zω z∈A Dω (z) 0<µω (A)≤1/2
Assuming Assumption C.1, the random walk on each fixed-order tree Zω satisfies Φω ≥
1 . 4κ2 n
Furthermore, its stationary leaf mass satisfies X 1 . µω (Y) = µω (y) ≥ 4κn y∈Y
Proof. Write b (x, z) := Mref (x, z)Vb (x, z). U For non-leaf states, Assumption C.1 gives b (x, z) ≤ κU ∗ (x, z). κ−1 U ∗ (x, z) ≤ U b (x, y) = U ∗ (x, y). On leaves, terminal anchoring gives U For a fixed-order tree, the weighted degree Dω (z) is the sum of the parent edge and the child edges incident to z. If k = k(z), then X b (x, z) + 1 {k < n} b (x, z ωk+1 ←a ). Dω (z) = 1 {k > 0} U U a∈V
b ≤ κU ∗ and the partition identity for the next reveal coordinate, Using U X Dω (z) ≤ κ1 {k > 0} U ∗ (x, z) + κ1 {k < n} U ∗ (x, z ωk+1 ←a ) a∈V
= κ 1 {k > 0} + 1 {k < n} U ∗ (x, z) ≤ 2κU ∗ (x, z). Let Dω,tot := Then
P
z∈Zω Dω (z) denote the total unnormalized stationary weight of the fixed-order tree.
Dω,tot ≤ 2κ
n X
X
U ∗ (x, z).
k=0 z: k(z)=k, z∈Zω
For each fixed depth in a fixed-order tree, the compatible states partition the leaves, so X U ∗ (x, z) = Z(x). z: k(z)=k, z∈Zω
Hence Dω,tot ≤ 2κ(n + 1)Z(x). Let Dω,leaf :=
X
Dω (y)
y∈Y
denote the total degree of the leaf layer. In the fixed-order tree, each leaf y has exactly one incident b (x, y). From the terminal anchoring in non-self edge, namely its parent edge, whose weight is U Assumption C.1, X Dω,leaf = U ∗ (x, y) = Z(x). y∈Y
Since the stationary distribution of the fixed-order weighted walk is degree-normalized, µω (Y) =
Dω,leaf Z(x) 1 1 ≥ = ≥ , Dω,tot 2κ(n + 1)Z(x) 2κ(n + 1) 4κn 35
where the last inequality uses n ≥ 1. We next lower bound the conductance. Let A be a connected subset of the fixed-order tree that does not contain the root. Let v ∈ A be the unique minimum-depth node in A. Since A does not contain the root, v has a parent p(v). By the minimality of k(v), we have p(v) ∈ / A, and hence the edge {p(v), v} leaves A. By Assumption C.1, its weight is b (x, v) ≥ κ−1 U ∗ (x, v). U For ℓ ∈ {0, . . . , n − k(v)}, let Cωℓ (v) := {u ∈ Zω | u is a descendant of v, k(u) = k(v) + ℓ} be the set of descendants of v at relative depth ℓ. Define the total weighted degree of the descendant subtree of v by n−k(v) X X Dω,sub (v) := Dω (u). ℓ=0
ℓ (v) u∈Cω
Applying the pointwise degree bound above to the descendant subtree of v, we have n−k(v)
Dω,sub (v) ≤ 2κ For every descendant level ℓ, X X U ∗ (x, u) = ℓ (v) u∈Cω
X
X
X
ℓ=0
ℓ (v) u∈Cω
πref (y | x)τ (x, y) =
ℓ (v) y∈C(u) u∈Cω
because the completion sets
C(u) : u ∈ Cωℓ (v)
U ∗ (x, u).
X
πref (y | x)τ (x, y) = U ∗ (x, v),
y∈C(v)
partition C(v). Therefore
Dω,sub (v) ≤ 2κ(n − k(v) + 1)U ∗ (x, v) ≤ 2κ(n + 1)U ∗ (x, v). Since A is connected in the tree is contained in the Pand v is its unique minimum-depth node, AP descendant subtree of v. Hence z∈A Dω (z) ≤ Dω,sub (v). The cut numerator z∈A, u∈A / fω (z, u) in the tree conductance Φω (A) contains the single boundary edge from v to its parent, because v ∈ A and p(v) ∈ / A. Therefore X b (x, v) ≥ κ−1 U ∗ (x, v). fω (z, u) ≥ fω (v, p(v)) = fω (p(v), v) = U z∈A, u∈A /
Combining this boundary lower bound with the descendant-subtree volume upper bound gives P fω (z, u) 1 1 κ−1 U ∗ (x, v) z∈A, u∈A P / Φω (A) := = 2 ≥ 2 . ≥ ∗ (x, v) D (z) 2κ(n + 1)U 2κ (n + 1) 4κ n ω z∈A This proves the bound for connected sets that do not contain the root. If A is disconnected and does not contain the root, decompose it into connected components A1 , . . . , Am . Let X X Br := fω (z, u), Vr := Dω (z). z∈Ar
z∈Ar , u∈A / r
Since the fixed-order graph is a tree, there are no edges between distinct components of A, and therefore Pm m X Br V Pm r Φω (A) = Pr=1 = Φω (Ar ) ≥ min Φω (Ar ). m r∈[m] r=1 Vr s=1 Vs r=1 If set A contains the root, then Ac does not. Since µω (A) ≤ 1/2, we have P an admissible P c z∈A Dω (z) ≤ z∈Ac Dω (z), and therefore Φω (A) ≥ Φω (A ). Applying the previous arguc 2 ment to the connected components of A proves Φω ≥ 1/(4κ n). The next two statements show that conductance survives summing weighted graphs, and that canonical balanced AOAR is exactly such a sum over fixed coordinate orders.
36
Lemma C.3 (Conductance of a sum of weighted graphs). Let G1 , . . . , Gm be weighted graphs on the same vertex set with nonnegative symmetric weights fG1 , . . . , fGm . Let G be the weighted graph with edge weight m X fG (z, u) := fGr (z, u). r=1
Then for every set A, ΦG (A) ≥ min ΦGr (A). r∈[m]
Consequently, ΦG ≥ min ΦGr . r∈[m]
Proof. Write DGr (z) :=
X
fGr (z, u),
X
DG (z) :=
u∼Gr z
fG (z, u) =
u∼G z
m X
DGr (z).
r=1
For any set A, P
fG (z, u) z∈A, u∈A P / . ΦG (A) = D z∈A G (z) Pm Pm Using fG = r=1 fGr and DG = r=1 DGr , this becomes Pm P r=1 / fGr (z, u) Pm z∈A, P u∈A . ΦG (A) = r=1 z∈A DGr (z) We use the elementary fact that for ar ≥ 0 and br > 0, Pm ar ar Pr=1 ≥ min . m br r∈[m] b r=1 r Therefore
P ΦG (A) ≥ min
r∈[m]
z∈A, u∈A / fGr (z, u)
P
z∈A DGr (z)
= min ΦGr (A). r∈[m]
Taking the infimum over admissible A gives the conductance bound. Proposition C.5 (Canonical balanced AOAR as a sum of fixed-order VGB trees). Consider the canonical balanced AOAR-VGB weighted graph whose depth-k parent-child edge coefficient is gk+1 =
1
0 ≤ k ≤ n − 1.
,
n−1 k
Then, the sum of the fixed-order tree graphs over all ω ∈ Permn is exactly (n − 1)! times this canonical balanced AOAR-VGB weighted graph. Proof. Consider an AOAR edge between a depth-k state z and a child z j←a , where j ∈ / R(z). This edge appears in exactly those fixed orders ω for which the first k coordinates are precisely R(z), in any order, and the (k + 1)-st coordinate is j. The number of such permutations is k!(n − k − 1)!. Therefore, in the sum of fixed-order tree graphs, the coefficient multiplying Mref (x, z j←a )Vb (x, z j←a ) is (n − 1)! k!(n − k − 1)! = n−1 . k
This is exactly (n − 1)!gk+1 for the canonical balanced coefficient in the proposition. Combining the fixed-order conductance bound with the sum-of-graphs representation gives the main canonical balanced mixing guarantee.
37
Theorem C.3 (Mixing guarantee for canonical balanced AOAR-VGB). Assume Assumption C.1. Let PBAL be the canonical balanced AOAR-VGB lazy kernel, and let µBAL be its stationary distribution. For a set A, write X QBAL (A, Ac ) QBAL (A, Ac ) := µBAL (z)PBAL (z, Ac ), Φ(PBAL ) := min A⊆Zreach µBAL (A) z∈A
0<µBAL (A)≤1/2
for conductance. We write γ(PBAL ) := 1 − max{|λ| : λ ∈ Λ(PBAL ), λ ̸= 1} for the spectral gap, where Λ(PBAL ) denotes the set of eigenvalues of PBAL . Then 1 1 1 , µ (Y) = Ω Φ(PBAL ) = Ω , γ(P ) = Ω . BAL BAL κ2 n κ4 n2 κn Moreover, the stationary law conditioned on the leaves is exactly µBAL (· | Y) = π ∗ (· | x). Consequently, if ZT is the state after T lazy balanced AOAR-VGB Markov chain steps initialized at Z0 ∈ Z with µBAL (Z0 ) > 0, then for any δ > 0, κn T = O κ4 n2 log δµBAL (Z0 ) suffices to guarantee dTV (L(ZT | ZT ∈ Y), π ∗ (· | x)) ≤ δ For example, if Z0 = ∅, then
a L(Z
and
P[ZT ∈ Y] = Ω
1 κn
a
.
κn . T = O κ4 n2 log δ
T ) denote the law of ZT and dTV denote the total variation distance.
Proof. By Proposition C.5, the canonical balanced AOAR weighted graph is a sum of fixed-order VGB tree graphs. By Lemma C.2, each fixed-order graph has conductance Ω(1/(κ2 n)). Lemma C.3 therefore implies 1 . Φ(PBAL ) = Ω κ2 n Since the chain is lazy and reversible, Cheeger’s inequality [16] gives 1 1 2 γ(PBAL ) ≥ Φ(PBAL ) = Ω . 2 κ4 n2 The leaf-mass lower bound follows by the same ratio-of-sums argument. Indeed, let X X Dω,leaf := Dω (y), Dω,tot := Dω (z). y∈Y
z∈Zω
Then µω (Y) = Dω,leaf /Dω,tot , or equivalently Dω,leaf = Dω,tot µω (Y). The global factor (n − 1)! in Proposition C.5 cancels after normalization, so the leaf mass of the canonical balanced AOAR graph satisfies P P Dω,tot µω (Y) ω∈Permn Dω,leaf P n µBAL (Y) = P = ω∈Perm ω∈Permn Dω,tot ω∈Perm Dω,tot n 1 1 ≥ min µω (Y) ≥ =Ω . ω∈Permn 4κn κn The exact leaf law follows from Theorem C.2, since terminal anchoring gives Vb (x, y) = τ (x, y) on leaves. 38
Finally, since PBAL is lazy and reversible, the standard spectral estimate for reversible chains [16] gives the mixing time bound from Z0 ∈ Z as follows: t ∥PBAL (Z0 , ·) − µBAL ∥TV ≤
1 −tγ(PBAL ) e µBAL (Z0 )−1/2 . 2
In particular, the mixing time from Z0 is, (Z ) Tmix0 (ε) = O
1 1 log p γ(PBAL ) ε µBAL (Z0 )
! .
(1)
T Let νT:= PBAL (Z0 , ·). If dTV (νT , µBAL ) ≤ ε and ε ≤ µBAL (Y)/2, then νT (Y) ≥ µBAL (Y)/2 = 1 Ω κn ; this follows from |νT (Y) − µBAL (Y)| ≤ dTV (νT , µBAL ). Then, from the definition of total variation (TV) distance, µBAL (y) 1 X νT (y) − dTV (νT (· | Y), µBAL (· | Y)) = 2 νT (Y) µBAL (Y) y∈Y 1 X νT (y) − µBAL (y) 1 1 = + µBAL (y) − 2 νT (Y) νT (Y) µBAL (Y) y∈Y
≤
X 1 |νT (Y) − µBAL (Y)| |νT (y) − µBAL (y)| + 2νT (Y) 2νT (Y) y∈Y
ε 3ε 2ε + ≤ . ≤ µBAL (Y) µBAL (Y) µBAL (Y) Here the partial-sum bound used above is X X |νT (y) − µBAL (y)| ≤ |νT (z) − µBAL (z)| = 2dTV (νT , µBAL ) ≤ 2ε. y∈Y
z∈Z
Thus, taking ε := δ µBAL (Y)/3 ensures that the leaf-conditioned law is within δ of the stationary leaf-conditioned law. Plugging this choice into (1), and using µBAL (Y) = Ω(1/(κn)), gives ! 1 κn 4 2 4 2 (2) T = O κ n log p = O κ n log δµBAL (Z0 ) ε µBAL (Z0 ) Since µBAL (· | Y) = π ∗ (· | x), the claim follows. Next, consider the case Z0 = ∅. It remains to lower bound its stationary mass µBAL (∅). For a fixed order ω, the root degree satisfies X X X b (x, ∅ω1 ←a ) ≥ κ−1 Dω (∅) = U U ∗ (x, ∅ω1 ←a ) = κ−1 πref (y | x)τ (x, y) = κ−1 Z(x), a∈V
a∈V
y∈Y
where the middle equality uses that the first-reveal completion sets {C(∅ω1 ←a )}a∈V partition Y. The total-degree bound in Lemma C.2 gives Dω,tot ≤ 2κ(n + 1)Z(x). Therefore µω (∅) ≥ 1/(2κ2 (n + 1)). Averaging over ω as above, with the global factor (n − 1)! canceling after normalization, yields 1 µBAL (∅) = Ω . κ2 n Plugging into (2) finishes the proof.
C.5
Geometric Balanced AOAR-VGB
Geometric balanced AOAR-VGB formulation. We take one further step beyond the depthbalanced AOAR-VGB local kernel in Definition C.2. In that kernel, a backward move from a non-leaf state z has weight proportional to Vb (x, z) for every revealed coordinate i ∈ R(z). Hence, once 39
Geometric AOAR-VGB transition on 𝑵𝑨𝑶𝑨𝑹 𝒛
𝒙, 𝒛 𝑽 𝒙, 𝒛−𝒊 Backward / re-mask: 𝒘𝐛𝐰𝐝 𝒛 = 𝒔𝒌−𝟏 𝑽
parent states (𝒌 − 𝟏)
𝒛−𝒊𝟏
𝝀
𝒛−𝒊𝟐
𝒛
current state (𝒌)
Self-loop for laziness (optional)
masked state child states (𝒌 + 𝟏)
𝒛𝒋𝟏 ←𝒂
𝒛𝒋𝟐 ←𝒃
𝒙, 𝒛 𝝀 𝑽 𝒙, 𝒛𝒋←𝒂 𝝅𝒓𝒆𝒇 𝒀𝒋 = 𝒂 𝒙, 𝒛 Forward / reveal: 𝒘𝐟𝐰𝐝 𝒛 = 𝒔𝒌 𝑽
Figure 11: Illustration of one geometric balanced AOAR-VGB transition from a masked state z. Forward reveal moves use the reference conditional times the child verifier value, with the common gate Vb (x, z)λ . Backward re-masking moves use the current verifier value and the parent gate Vb (x, z −i )λ . Setting λ = 0 removes the geometric gates and recovers balanced AOAR-VGB.
the chain decides to move backward, the index to re-mask is selected uniformly among revealed coordinates. This is a natural depth-balancing rule, but it ignores an extra degree of freedom that is absent in AR prefix trees: in AOAR, the sampler can choose which revealed coordinate should be erased. This motivates a value-gated remasking rule. When considering a backward move z → z −i , we would like to favor coordinates whose remasked parent z −i remains promising under the verifier. We implement this through a clean geometric edge-weight construction: the child side carries the tilted b (x, z j←a ), while the parent side contributes through Vb (x, z −i )λ . This preserves completion mass U the usual VGB forward ranking and makes the backward index choice depend on the verifier signal of the remasked parent. Figure 11 illustrates the resulting local transition rule. The geometric gate leaves the forward candidate ranking unchanged up to the state-wise factor Vb (x, z)λ , while the backward weights distinguish revealed coordinates through the verifier values of their remasked parents. Definition C.3 (Geometric balanced AOAR edge weights). Fix λ ≥ 0 and positive depth coefficients s0 , s1 , . . . , sn−1 > 0. For neighboring states u ∼ v, let p(u, v) denote the lower-depth endpoint and c(u, v) denote the higher-depth endpoint. Thus c(u, v) = p(u, v)j←a for a unique revealed coordinate and token. If c(u, v) ∈ / Y, define the non-leaf edge weight by b (x, c(u, v))Vb (x, p(u, v))λ . fBG (u, v) = fBG (v, u) := sk(p(u,v)) U If c(u, v) = y ∈ Y, define the leaf-corrected edge weight by fBG (u, v) = fBG (v, u) := sn−1 U ∗ (x, y) = sn−1 πref (y | x)τ (x, y).
b and U ∗ , the local Local implementation weights. Although Definition C.3 is written in terms of U weights below do not require evaluating these tilted masses directly. Let z have depth k = k(z). 40
For a non-leaf child c = z j←a ∈ / Y, we have b (x, c) = Mref (x, z)πref (Yj = a | x, z)Vb (x, c). U Dividing all incident non-leaf edge weights at current state z by the common factor Mref (x, z) gives wBG,fwd (z → c) = sk Vb (x, z)λ πref (Yj = a | x, z)Vb (x, c). The factor Vb (x, z)λ is common across all forward candidates from the same current state z. Hence the forward candidate ranking is the same as in ordinary VGB: wBG,fwd (z → c) ∝ πref (Yj = a | x, z)Vb (x, z j←a ). For a backward move from a non-leaf state z to u = z −i , the relevant edge has child endpoint z and parent endpoint u. Thus b (x, z)Vb (x, u)λ = sk−1 Mref (x, z)Vb (x, z)Vb (x, u)λ . fBG (u, z) = sk−1 U Dividing by the current-state common factor Mref (x, z) gives wBG,bwd (z → z −i ) = sk−1 Vb (x, z)Vb (x, z −i )λ . Thus backward moves prefer revealed coordinates whose remasked parent has high verifier value. At the leaf boundary, the leaf-corrected edge gives the local forward weight directly. If y = z j←a ∈ Y, then fBG (z, y) = sn−1 U ∗ (x, y) = sn−1 Mref (x, z)πref (Yj = a | x, z)τ (x, y), and dividing by Mref (x, z) gives wBG,fwd (z → y) = sn−1 πref (Yj = a | x, z)τ (x, y). Similarly, from a leaf y to a parent y −i , wBG,bwd (y → y −i ) = sn−1 τ (x, y). Remark C.3 (λ = 0 recovers balanced AOAR-VGB). When λ = 0, the geometric gate disappears and the local weights reduce to the balanced AOAR-VGB weights with depth coefficients sk . Theorem C.4 (Stationarity and exact leaf law of geometric balanced AOAR). Let PBG be the random walk on the reachable AOAR graph with edge weights fBG from Definition C.3. That is, X fBG (z, u) PBG (z, u) = fBG (z, v). , DBG (z) := DBG (z) v∼z Then PBG is reversible with stationary distribution DBG (z) . ′ z ′ DBG (z )
µBG (z) = P
Moreover, under terminal anchoring, its leaf-conditioned stationary law is exact: µBG (y | y ∈ Y) = π ∗ (y | x). Proof. Since fBG is a nonnegative symmetric edge weight, the random walk is a standard weightedgraph random walk. For neighboring distinct states z ∼ u, fBG (z, u) µBG (z)PBG (z, u) = P = µBG (u)PBG (u, z), ′ z ′ DBG (z ) so detailed balance holds. Now let y ∈ Y be a leaf. Its AOAR parents are y −i for i ∈ [n]. By the leaf correction, fBG (y −i , y) = sn−1 U ∗ (x, y) 41
for every i. Hence DBG (y) =
n X
fBG (y −i , y) = nsn−1 U ∗ (x, y).
i=1
The factor nsn−1 is constant over leaves, and therefore U ∗ (x, y) = π ∗ (y | x). ∗ (x, y ′ ) U ′ y ∈Y
µBG (y | y ∈ Y) = P
Mixing analysis of geometric balanced VGB. We next transfer the canonical balanced mixing guarantee to the geometric kernel. The only care point is that zero target-mass states should not create spurious components. We therefore work on the positive target-support component, which is stable under all positive-weight moves. Lemma C.4 (Positive target support is closed under AOAR moves). Assume Z(x) > 0 and Assumption C.1. Define the positive target support Z+ (x) := z ∈ Zreach : U ⋆ (x, z) > 0 . Then ∅ ∈ Z+ (x). Moreover, Z+ (x) is closed under all positive-weight balanced and geometric AOAR moves. More precisely: 1. If z ∈ Z+ (x) and i ∈ R(z), then z −i ∈ Z+ (x). 2. If z ∈ Z+ (x), j ∈ / R(z), and c = z j←a ∈ / Z+ (x), then the AOAR edge {z, c} has zero balanced edge weight and zero geometric edge weight. Proof. Since Z(x) > 0,
U ⋆ (x, ∅) = Z(x) > 0,
so ∅ ∈ Z+ (x). For the first claim, let p = z −i . Since p is obtained by re-masking one revealed coordinate of z, we have C(z) ⊆ C(p). Thus X X U ⋆ (x, p) = πref (y | x)τ (x, y) ≥ πref (y | x)τ (x, y) = U ⋆ (x, z) > 0. y∈C(p)
y∈C(z)
Hence p ∈ Z+ (x). For the second claim, let c = z j←a ∈ / Z+ (x). If c is not πref -reachable, then it is not an admissible child edge. Otherwise, Mref (x, c) > 0 and U ⋆ (x, c) = 0, hence V ⋆ (x, c) = 0. By Assumption C.1, Vb (x, c) = 0,
b (x, c) := Mref (x, c)Vb (x, c) = 0. U
b (x, c) for non-leaf children and to Therefore the balanced edge weight, which is proportional to U U ⋆ (x, c) for leaf children, is zero. The same is true for the geometric edge weight: for non-leaf children it is proportional to b (x, c)Vb (x, z)λ , U and for leaf children it is proportional to U ⋆ (x, c). Hence no positive-weight edge leaves Z+ (x). On this component, the geometric graph differs from the canonical balanced graph only through the parent value gate Vb (x, p)λ on non-leaf edges. The next assumption prevents this gate from creating an additional bottleneck.
42
Assumption C.2 (Bounded geometric verifiers). Consider the canonical geometric balanced AOARVGB graph with 1 sk = n−1 , 0 ≤ k ≤ n − 1. k
Fix λ > 0. Assume there exist constants 0<m≤1≤M <∞ such that, for every non-leaf state z ∈ Z+ (x), m ≤ Vb (x, z) ≤ M. Theorem C.5 (Mixing stability for geometric balanced AOAR-VGB). Assume Assumption C.1 and Assumption C.2. Let PBG be the lazy canonical geometric balanced AOAR-VGB kernel on Z+ (x) with Z(x) > 0, and let µBG be its stationary distribution. Then µBG (· | Y) = π ∗ (· | x). Moreover, its conductance satisfies Φ(PBG ) = Ω
m λ 1 , M κ2 n
and hence, by Cheeger’s inequality for lazy reversible chains, its spectral gap satisfies m 2λ 1 γ(PBG ) = Ω . M κ4 n2 The stationary leaf and root masses satisfy 1 , µBG (Y) = Ω M λ κn
µBG (∅) = Ω
m λ 1 M κ2 n
.
Consequently, if ZT is the state after T lazy geometric AOAR-VGB transitions initialized at Z0 ∈ Z+ (x) with µBG (Z0 ) > 0, then for any δ > 0, ! 2λ M M λ κn 4 2 T =O κ n log m δµBG (Z0 ) suffices to guarantee dTV (L(ZT | ZT ∈ Y), π ∗ (· | x)) ≤ δ
and
P[ZT ∈ Y] = Ω
1 λ M κn
.
For example, if Z0 = ∅, then T =O
M m
2λ
(M/m)λ κn κ n log δ 4 2
! .
In particular, if m, M are constants independent of n and λ is fixed, the geometric chain’s mixing time is of the same order as canonical balanced AOAR-VGB. Proof. Let fBAL be the canonical balanced AOAR edge weight and let fBG be the canonical geometric balanced AOAR edge weight. For an edge e = {p, c}, let p be the lower-depth endpoint and c the higher-depth endpoint. If c ∈ / Y, then by Definition C.3, fBG (e) = Vb (x, p)λ fBAL (e). If c ∈ Y, the leaf correction gives fBG (e) = fBAL (e). Therefore, using Assumption C.2 and m ≤ 1 ≤ M , every positive-weight edge in the component satisfies mλ fBAL (e) ≤ fBG (e) ≤ M λ fBAL (e). 43
For any set A, write QBG (A, Ac ) :=
X
fBG (e),
volBG (A) :=
XX
fBG (u, v),
u∈A v∼u
e={u,v} u∈A, v ∈A /
and define QBAL and volBAL analogously. The edge comparison gives QBG (A, Ac ) ≥ mλ QBAL (A, Ac ),
volBG (A) ≤ M λ volBAL (A).
Thus, for any admissible A, ΦBG (A) =
QBG (A, Ac ) m λ QBAL (A, Ac ) m λ ≥ = ΦBAL (A). volBG (A) M volBAL (A) M
Taking the infimum over admissible A and applying Theorem C.3 gives m λ 1 Φ(PBG ) = Ω . M κ2 n Since the lazy geometric chain is reversible, Cheeger’s inequality gives m 2λ 1 γ(PBG ) = Ω . M κ4 n2 The exact leaf-conditioned law follows from Theorem C.4, since laziness does not change the stationary distribution. For the leaf mass, leaf-corrected edges have the same weight as in the canonical balanced graph. Writing DBG and DBAL for weighted degrees, P P 1 1 y∈Y DBAL (y) y∈Y DBAL (y) . µBG (Y) = P ≥ λP = λ µBAL (Y) = Ω M M M λ κn z DBG (z) z DBAL (z) Similarly, the root degree is at least mλ times the balanced root degree, while the total degree is at most M λ times the balanced total degree, so m λ m λ 1 µBG (∅) ≥ . µBAL (∅) = Ω M M κ2 n Finally, the standard start-state spectral bound for lazy reversible chains [16] gives 1 T PBG (Z0 , ·) − µBG TV ≤ e−T γ(PBG ) µBG (Z0 )−1/2 . 2 As in the proof of Theorem C.3, taking the full-chain TV error to be at most ε :=
δ µBG (Y) 3
implies dTV (L(ZT | ZT ∈ Y), µBG (· | Y)) ≤ δ
and
P[ZT ∈ Y] ≥ µBG (Y)/2 = Ω
1 M λ κn
.
Substituting the bounds on γ(PBG ) and µBG (Y), while retaining the start-state factor µBG (Z0 ), yields the general mixing time. The lower bound on µBG (∅) gives the root-start example. Since µBG (· | Y) = π ∗ (· | x), the claim follows. Forward/backward mass analysis. We now return from global mixing to the local behavior of the geometric transition. The next proposition measures how the forward-versus-backward move probabilities change under the exact verifier setup. Proposition C.6 (Directional masses under the exact verifier). Assume Vb = V ∗ on all non-leaf states. Let z be a non-leaf interior state of depth 1 ≤ k ≤ n − 2 with V ∗ (x, z) > 0, so that its one-step forward children are also non-leaf. Define the parent-value ratio ri∗ (z) :=
V ∗ (x, z −i ) , V ∗ (x, z)
44
i ∈ R(z).
Then the exact forward and backward masses are Fλ∗ (z) = sk (n − k)V ∗ (x, z)1+λ , and
X
Bλ∗ (z) = sk−1 V ∗ (x, z)1+λ
ri∗ (z)λ .
i∈R(z)
Consequently, P(forward | x, z, move) =
sk (n − k) P , sk (n − k) + sk−1 i∈R(z) ri∗ (z)λ
and P(backward | x, z, move) =
sk−1
∗ λ i∈R(z) ri (z)
P
sk (n − k) + sk−1
∗ λ i∈R(z) ri (z)
P
.
Proof. For the forward mass, using the AOAR Bellman identity, X X Fλ∗ (z) = sk V ∗ (x, z)λ πref (Yj = a | x, z)V ∗ (x, z j←a ) a∈V j ∈R(z) /
X
= sk V ∗ (x, z)λ
V ∗ (x, z) = sk (n − k)V ∗ (x, z)1+λ .
j ∈R(z) /
For backward moves, Bλ∗ (z) =
X
sk−1 V ∗ (x, z)V ∗ (x, z −i )λ
i∈R(z)
= sk−1 V ∗ (x, z)1+λ
X V ∗ (x, z −i ) λ V ∗ (x, z)
i∈R(z)
.
Normalizing the forward and backward masses gives the probabilities in the proposition. Remark C.4 (Reference-target mismatch interpretation). By the exact target conditional, π ∗ (Yi = zi | x, z −i ) = πref (Yi = zi | x, z −i )
V ∗ (x, z) , V ∗ (x, z −i )
ri∗ (z) =
πref (Yi = zi | x, z −i ) . π ∗ (Yi = zi | x, z −i )
Thus ri∗ (z) > 1 means πref favors zi more than π ∗ and increases the remasking weight of coordinate i, while ri∗ (z) < 1 means π ∗ gives stronger support and the remasking weight is reduced. In other words, the chain preferentially erases tokens that look base-plausible but target-weak, and keeps tokens that are target-supported. Remark C.5 (Penultimate-depth directional masses). When z is a penultimate-depth state with k(z) = n − 1, any forward move reaches a leaf. The leaf correction resets the forward edge to sn−1 U ∗ (x, y) rather than the non-leaf product edge. Hence, under the exact verifier, Fλ∗ (z) = sn−1 V ∗ (x, z), whereas
Bλ∗ (z) = sn−2 V ∗ (x, z)
X
V ∗ (x, z −i )λ .
i∈R(z)
This boundary rule is intentional: it keeps the leaf degree proportional to U ∗ (x, y), and therefore preserves the exact leaf-conditioned stationary law.
45
Remark C.6 (Canonical depth anchor versus exact 1/2 balance). The linear balanced choice sk (n − k) = ksk−1 does not give exact 1/2–1/2 balance for λ > 0 in the clean geometric chain. Under this choice, away from the leaf boundary, P ∗ λ i∈R(z) ri (z) P , k = |R(z)| . P(backward | x, z, move) = ∗ k + i∈R(z) ri (z)λ Thus sk (n − k) = ksk−1 is best interpreted as the canonical depth-neutral anchor: it recovers exact balance when λ = 0, while the λ > 0 correction changes the backward probability according to the verifier value of the remasked parent states. More quantitatively, for small λ, ri∗ (z)λ = exp(λ log ri∗ (z)) = 1 + λ log ri∗ (z) + O λ2 log2 ri∗ (z) . Define the first empirical log-moment of the parent-value ratios 1 X V ∗ (x, z −i ) log . k V ∗ (x, z)
m̄1 (z) :=
i∈R(z)
Then
X
ri∗ (z)λ = k + λk m̄1 (z) + O λ2 k m̄2 (z) ,
i∈R(z)
where
1 X V ∗ (x, z −i ) log2 . k V ∗ (x, z)
m̄2 (z) :=
i∈R(z)
Substituting this into the backward probability gives P(backward | x, z, move) =
1 λ + m̄1 (z) + O λ2 m̄2 (z) + λ2 m̄1 (z)2 . 2 4
Equivalently, P(forward | x, z, move) =
1 λ − m̄1 (z) + O λ2 m̄2 (z) + λ2 m̄1 (z)2 . 2 4
Therefore, the leading-order deviation from 1/2 is not a combinatorial depth bias. It is the value-gated remasking signal induced by the geometric parent factor.
D
Momentum-Balanced AOAR-VGB
Balanced AOAR-VGB fixes the depth-allocation pathology of primitive AOAR-VGB by balancing forward and backward directions at each interior depth in the ideal exact-value case. However, the resulting walk can still be locally diffusive: after revealing a coordinate, the chain may immediately remask a coordinate, and after re-masking it may immediately reveal again. Such re-masking/revealing oscillations waste Markov steps and can be computationally inefficient, especially for long sequences where each transition may require evaluating many candidate child states. This issue is closely related to the diffusive behavior of reversible random walks studied in the lifting literature. A lifting replaces each state by several copies and defines a larger Markov chain whose projection preserves the desired stationary law, while using non-reversible momentum to reduce unnecessary direction changes [11]. In the present AOAR setting, the two natural directions are downward reveal moves and upward re-masking moves. We therefore introduce a two-copy momentum lift of balanced AOAR-VGB. The construction is inspired by the momentum version of VGB in Rohatgi et al. [23] and by the flow-cancellation principle of Hayes and Sinclair [11], but is adapted here to the any-order masked-state graph. 46
D.1
Setup and Notation
Throughout this section, we use the balanced AOAR-VGB notation from Section C.3. In particular, KBAL denotes the self-loop-free balanced AOAR-VGB move kernel induced by the local weights in Definition C.2 with γ = 0, and µBAL denotes its stationary distribution from Theorem C.2. Directional neighborhoods and masses. Recall the AOAR child and parent neighborhoods CAOAR (z) := z j←a : j ∈ / R(z), a ∈ V, Mref (x, z j←a ) > 0 , −i PAOAR (z) := z : i ∈ R(z) . A move to CAOAR (z) is called a downward (↓) or forward move, because it reveals one additional coordinate. A move to PAOAR (z) is called an upward (↑) or backward move, because it re-masks one coordinate. For the self-loop-free balanced kernel, define the directional masses X X p↓ (z) := KBAL (z, c), p↑ (z) := c∈CAOAR (z)
KBAL (z, p).
p∈PAOAR (z)
Thus p↓ (z) + p↑ (z) = 1. At the boundaries, these definitions give p↓ (∅) = 1,
p↑ (∅) = 0,
p↓ (y) = 0,
p↑ (y) = 1
(y ∈ Y).
In the exact canonical balanced case, where βk = k, αk = n − k for 1 ≤ k ≤ n − 1, and Vb = V ∗ , Proposition C.4 gives 1 p↓ (z) = p↑ (z) = 2 for every interior reachable state. In practical plug-in samplers, however, Vb may break this exact equality. Lifted tuple state space. Momentum is represented by augmenting the masked state with a direction variable. We use tuple states (z, ↓), (z, ↑). As introduced above, the label ↓ means that the chain is currently in a downward, forward-revealing momentum mode, while ↑ means that it is currently in an upward, backtracking momentum mode. Let ZBAL,+ be a reachable component of Zreach on which WBAL (z) > 0 for every state z. We keep both direction labels for every state in this component: ZeBAL := ZBAL,+ × {↓, ↑} . Thus, states in this component, including its root and leaf boundary states, have two momentum copies. Some boundary copies are only used as switch states. For example, from (∅, ↑) the chain switches to (∅, ↓) before moving downward, while from (y, ↓) a leaf switches to (y, ↑) before moving upward. We write proj(z, d) := z, d ∈ {↓, ↑} , for the projection from a lifted state to its current masked state. For a probability distribution νe on ZeBAL , its current-state marginal is proj# νe(z) = νe(z, ↓) + νe(z, ↑). Stationary split. The goal of the momentum lift is to change path geometry without changing the stationary law after projection. Thus the lifted chain should have a stationary distribution on the two-copy space whose projection is exactly the balanced AOAR-VGB stationary law µBAL . We take the simplest such target: an equal split of the base stationary mass across the two momentum modes, 1 1 µ eBAL (z, ↓) := µBAL (z), µ eBAL (z, ↑) := µBAL (z). 2 2 This choice does not bias the chain toward either momentum direction at stationarity, and projecting out the momentum label immediately recovers the desired base stationary law: proj# µ eBAL (z) = µBAL (z). 47
D.2
Momentum-Balanced AOAR-VGB
We now define the flow-cancelled switch lift. When the chain is in downward mode, it keeps the child moves of the balanced AOAR kernel and replaces the total parent-move mass by a switch to upward mode. When the chain is in upward mode, it keeps the parent moves and replaces the total child-move mass by a switch to downward mode. For a fixed masked state z, the two lifted copies are (z, ↓) and (z, ↑). If both momentum modes assign positive probability to switching between these two copies, the smaller of the two switch probabilities is merely a symmetric back-and-forth exchange. We remove this common exchange and turn it into same-copy holding probability, leaving only the residual imbalance as an actual momentum switch. This is the direct analogue of crossing-flow cancellation in lifted Markov chains: when there is crossing flow in both directions between two lifted copies, the common part can be cancelled without changing the projected stationary flow.
χm
(z, ↓) p↓ − χm
p↑ − χm
(z, ↑) χm m = min{p↓ , p↑ }
Figure 12: Flow cancellation at state z.
Definition D.1 (Flow-cancelled momentum-balanced AOAR-VGB kernel). Fix a cancellation strength χ ∈ [0, 1]. For each reachable state z, define m(z) := min{p↓ (z), p↑ (z)}. e fc,χ on ZeBAL is defined as follows. The flow-cancelled lifted kernel K BAL Downward copy. with probability
From (z, ↓), the chain moves to a child while keeping downward momentum
e fc,χ (z, ↓), (c, ↓) := KBAL (z, c), K BAL
c ∈ CAOAR (z).
(C HILD M OVE)
It switches to upward momentum with residual opposite-direction probability e fc,χ (z, ↓), (z, ↑) := p↑ (z) − χm(z), K BAL
(S WITCH U P)
and it stays in the same lifted state with the cancelled crossing mass e fc,χ (z, ↓), (z, ↓) := χm(z). K BAL
(S TAY D OWN)
Upward copy. probability
From (z, ↑), the chain moves to a parent while keeping upward momentum with
e fc,χ (z, ↑), (p, ↑) := KBAL (z, p), K BAL
p ∈ PAOAR (z).
(PARENT M OVE)
It switches to downward momentum with residual opposite-direction probability e fc,χ (z, ↑), (z, ↓) := p↓ (z) − χm(z), K (S WITCH D OWN) BAL and it stays in the same lifted state with the cancelled crossing mass e fc,χ (z, ↑), (z, ↑) := χm(z). K BAL
(S TAY U P)
All other transition probabilities are zero. The parameter χ acts as a cancellation knob, analogous to a friction coefficient in underdamped Langevin dynamics: it controls how strongly the lifted chain suppresses direction changes and preserves momentum persistence. Larger χ removes more symmetric back-and-forth switching and keeps more probability in the current momentum copy. The case χ = 0 performs no cancellation. 48
The case χ = 1 maximally cancels opposing switch flows. In particular, if p↓ (z) = p↑ (z) = 1/2 as in the exact canonical balanced case with Vb = V ∗ , then the switch probabilities at z become zero and the cancelled mass becomes a same-copy holding probability 1/2. Remark D.1 (Implementation as a categorical distribution). Recall that w(k) denotes the depthrescaled balanced AOAR-VGB local weight from Definition C.2. Let X X F (z) := w(k(z)) (z → c), B(z) := w(k(z)) (z → p), c∈CAOAR (z)
p∈PAOAR (z)
and M (z) := min {F (z), B(z)} . Then p↓ (z) =
F (z) , F (z) + B(z)
p↑ (z) =
B(z) , F (z) + B(z)
m(z) =
M (z) . F (z) + B(z)
Equivalently, the kernel can be implemented by unnormalized categorical weights. In downward mode, sample from n o A↓ (z) := c ∈ CAOAR (z) : weight w(k(z)) (z → c) ∪ {switch_up : weight B(z) − χM (z)} ∪ {stay_down : weight χM (z)} . If a child is sampled, the chain moves to (c, ↓). If switch_up is sampled, it moves to (z, ↑). If stay_down is sampled, it remains at (z, ↓). In upward mode, sample from n o A↑ (z) := p ∈ PAOAR (z) : weight w(k(z)) (z → p) ∪ {switch_down : weight F (z) − χM (z)} ∪ {stay_up : weight χM (z)} . If a parent is sampled, the chain moves to (p, ↑). If switch_down is sampled, it moves to (z, ↓). If stay_up is sampled, it remains at (z, ↑). The sampler implementation is the singleton-block restriction of the momentum MDM-VGB sampler in Algorithm 4: set M = {1} and use the AOAR local forward/backward weights. Thus each MDM child/parent block is a single revealed or re-masked coordinate. e fc,χ is a Markov kernel on Proposition D.1 (Markov property). For every χ ∈ [0, 1], the kernel K BAL ZeBAL . Proof. Since m(z) ≤ p↓ (z) and m(z) ≤ p↑ (z), all transition probabilities are nonnegative. For the downward copy, X e fc,χ (z, ↓), (c, ↓) + K e fc,χ (z, ↓), (z, ↑) + K e fc,χ (z, ↓), (z, ↓) K BAL BAL BAL c∈CAOAR (z)
= p↓ (z) + p↑ (z) − χm(z) + χm(z) = p↓ (z) + p↑ (z) = 1. The upward row is identical: X e fc,χ (z, ↑), (z, ↓) + K e fc,χ (z, ↑), (z, ↑) e fc,χ (z, ↑), (p, ↑) + K K BAL BAL BAL p∈PAOAR (z)
= p↑ (z) + p↓ (z) − χm(z) + χm(z) = 1. Thus every row sums to one.
49
Theorem D.1 (Stationarity of flow-cancelled momentum-balanced AOAR-VGB). Assume terminal anchoring, Vb (x, y) = τ (x, y), y ∈ Y. Let KBAL be the self-loop-free (γ = 0) balanced AOAR-VGB move kernel from Definition C.2, restricted to a reachable component where WBAL (z) > 0 for every state z. Then, from Theorem C.2 KBAL is reversible with respect to µBAL , and µBAL (· | Y) = π ∗ (· | x). For every χ ∈ [0, 1], the lifted measure µ eBAL (z, ↓) = µ eBAL (z, ↑) =
1 µBAL (z) 2
e fc,χ . Moreover, is stationary for K BAL proj# µ eBAL = µBAL . Consequently, the flow-cancelled lifted chain has the exact tilted target as its leaf-conditioned current-state law: proj# µ eBAL (· | Y) = π ∗ (· | x). Proof. We verify stationarity by checking incoming flow into each lifted copy. First consider (z, ↓). Incoming flow comes from three sources: parents of z in downward mode, the switch from (z, ↑), and the self-loop at (z, ↓). The incoming move-flow from parent states u ∈ PAOAR (z) is X 1 µBAL (u)KBAL (u, z). 2 u∈PAOAR (z)
Using reversibility of KBAL , X 1 1 µBAL (u)KBAL (u, z) = µBAL (z) 2 2 u∈PAOAR (z)
X
KBAL (z, u) =
u∈PAOAR (z)
1 µBAL (z)p↑ (z). 2
The incoming switch-flow from (z, ↑) is 1 µBAL (z) p↓ (z) − χm(z) , 2 and the self-loop flow at (z, ↓) is 1 µBAL (z)χm(z). 2 Therefore the total incoming flow into (z, ↓) is 1 1 1 µBAL (z)p↑ (z) + µBAL (z) p↓ (z) − χm(z) + µBAL (z)χm(z) 2 2 2 1 1 eBAL (z, ↓). = µBAL (z) p↑ (z) + p↓ (z) = µBAL (z) = µ 2 2 The proof for (z, ↑) is symmetric. Incoming move-flow comes from children v ∈ CAOAR (z): X X 1 1 1 µBAL (v)KBAL (v, z) = µBAL (z) KBAL (z, v) = µBAL (z)p↓ (z), 2 2 2 v∈CAOAR (z)
v∈CAOAR (z)
where the first equality uses reversibility of KBAL . The incoming switch-flow from (z, ↓) is 1 µBAL (z) p↑ (z) − χm(z) , 2 and the self-loop flow at (z, ↑) is 1 µBAL (z)χm(z). 2 50
Geometric AOAR-VGB-Momentum transition on 𝑵𝑨𝑶𝑨𝑹 𝒛
𝒙, 𝒛 𝑽 𝒙, 𝒛−𝒊 Backward / re-mask: 𝒘𝐛𝐰𝐝 𝒛 = 𝒔𝒌−𝟏 𝑽
𝒛−𝒊𝟏
parent states (𝒌 − 𝟏)
𝒛−𝒊𝟐
current state (𝒌)
𝝀
Momentum change: 𝒘𝒇𝒍𝒊𝒑 𝒛, ↑ = 𝑭 𝒛 − 𝝌𝑴(𝒛)
(𝒛, ↑)
(𝒛, ↓)
Self-loop for laziness (optional)
𝒘𝒇𝒍𝒊𝒑 𝒛, ↓ = 𝑩 𝒛 − 𝝌𝑴(𝒛)
𝒛𝒋𝟏 ←𝒂
child states (𝒌 + 𝟏)
𝒛𝒋𝟐 ←𝒃
𝒙, 𝒛 𝝀 𝑽 𝒙, 𝒛𝒋←𝒂 𝝅𝒓𝒆𝒇 𝒀𝒋 = 𝒂 𝒙, 𝒛 Forward / reveal: 𝒘𝐟𝐰𝐝 𝒛 = 𝒔𝒌 𝑽
Figure 13: Illustration of the momentum lift for geometric balanced AOAR-VGB. Forward moves preserve downward momentum, backward moves preserve upward momentum, and residual oppositedirection mass switches between lifted copies; the cancelled common mass may be kept as an optional stay probability. Thus the total incoming flow into (z, ↑) is 1 1 1 1 µBAL (z)p↓ (z) + µBAL (z) p↑ (z) − χm(z) + µBAL (z)χm(z) = µBAL (z). 2 2 2 2 Therefore µ eBAL is stationary. Finally, proj# µ eBAL (z) = µ eBAL (z, ↓) + µ eBAL (z, ↑) = µBAL (z). By Theorem C.2, terminal anchoring implies µBAL (· | Y) = π ∗ (· | x). Since the lifted current-state marginal is µBAL , the same leaf-conditioned law holds after projection. D.3
Momentum Lift for Geometric Balanced AOAR-VGB
Let KBG denote the non-lazy move kernel induced by Definition C.3, and let µBG be its stationary law from Theorem C.4. Define directional probabilities X X p↓ (z) := KBG (z, c), p↑ (z) := KBG (z, p). c∈CAOAR (z)
p∈PAOAR (z)
Then p↓ (z) + p↑ (z) = 1 at non-isolated states. The same flow-cancelled lift used above applies to this geometric base chain. Figure 13 illustrates the lifted transition. The downward copy keeps forward reveal moves, the upward copy keeps backward re-masking moves, and the residual opposite-direction mass switches the momentum label. Definition D.2 (Flow-cancelled momentum geometric AOAR kernel). Fix χ ∈ [0, 1] and set m(z) := min {p↓ (z), p↑ (z)} . On the lifted state space ZeBG := Z+ (x) × {↓, ↑} , e fc,χ as follows. define the lifted kernel K BG
51
Downward copy. with probability
From (z, ↓), the chain moves to a child while keeping downward momentum
e fc,χ (z, ↓), (c, ↓) := KBG (z, c), K BG
c ∈ CAOAR (z).
(C HILD M OVE)
It switches to upward momentum with residual opposite-direction probability e fc,χ (z, ↓), (z, ↑) := p↑ (z) − χm(z). K BG
(S WITCH U P)
and it stays in the same lifted state with the cancelled crossing mass e fc,χ (z, ↓), (z, ↓) := χm(z). K BG
(S TAY D OWN)
Upward copy. probability
From (z, ↑), the chain moves to a parent while keeping upward momentum with
e fc,χ (z, ↑), (p, ↑) := KBG (z, p), K BG
p ∈ PAOAR (z).
(PARENT M OVE)
It switches to downward momentum with residual opposite-direction probability e fc,χ (z, ↑), (z, ↓) := p↓ (z) − χm(z). K (S WITCH D OWN) BG and it stays in the same lifted state with the cancelled crossing mass e fc,χ (z, ↑), (z, ↑) := χm(z). K BG
(S TAY U P)
All other transition probabilities are zero. Theorem D.2 (Momentum preserves the geometric AOAR leaf law). Define µ eBG (z, ↓) = µ eBG (z, ↑) :=
1 µBG (z). 2
e fc,χ for every χ ∈ [0, 1]. Moreover, its current-state projection is Then µ eBG is stationary for K BG µBG , and therefore proj# µ eBG (· | Y) = π ∗ (· | x). Proof. We prove stationarity for the downward copy; the upward proof is symmetric. Incoming flow into (z, ↓) has three parts. First, parent states u ∈ PAOAR (z) in downward mode contribute X 1 µBG (u)KBG (u, z). 2 u∈PAOAR (z)
Since KBG is reversible with respect to µBG , this equals X 1 1 µBG (z) KBG (z, u) = µBG (z)p↑ (z). 2 2 u∈PAOAR (z)
Second, the switch from (z, ↑) contributes 1 µBG (z)(p↓ (z) − χm(z)). 2 Third, the self-loop at (z, ↓) contributes 1 µBG (z)χm(z). 2 The total incoming flow is therefore 1 1 µBG (z) {p↑ (z) + p↓ (z)} = µBG (z) = µ eBG (z, ↓). 2 2 The upward-copy calculation is identical with parents and children exchanged. Projection gives µ eBG (z, ↓) + µ eBG (z, ↑) = µBG (z). The leaf law follows from Theorem C.4. 52
D.4
Theoretical Analysis: Depth-Speed Gain with Momentum
We now focus on the depth-wise behavior of Balanced AOAR-VGB (the λ = 0, non-geometric chain) and its flow-cancelled momentum lift. We note that this theoretical analysis is restricted to this balanced chain, rather than the geometric variants above. In AOAR-VGB without momentum, even when forward and backward moves are locally balanced, the reversible chain still behaves like a diffusive random walk along the depth axis, leading to a quadratic first-leaf hitting time. We show that, under a near-balance condition on the realized forward and backward masses, the no-momentum chain remains diffusive with first-leaf hitting time of order n2 , whereas the maximally flow-cancelled momentum lift can accelerate this depth traversal to order n steps. This is the quantity most aligned with practical VGB-style implementations, which often stop the Markov chain once it first reaches a leaf. Definition D.3 (Depth process and hitting time). For a masked state z, let k(z) = |R(z)|. For the balanced chain with current masked state ZtBAL , define τtBAL := k(ZtBAL ),
TnBAL := inf{t ≥ 0 : τtBAL = n}.
For a lifted momentum chain Zet = (ZtBAL,Mo , Dt ), define τtBAL,Mo := k(ZtBAL,Mo ),
TnBAL,Mo := inf{t ≥ 0 : τtBAL,Mo = n}.
Assumption D.1 (Uniform directional near-balance). There exists εn ∈ [0, 1) such that for every reachable interior state z, |p↓ (z) − p↑ (z)| ≤ εn . Equivalently, since p↓ (z) + p↑ (z) = 1, 1 + εn 1 − εn ≤ p↓ (z), p↑ (z) ≤ . 2 2 The exact canonical balanced case corresponds to εn = 0. The near-balanced speedup regime is εn = O(1/n). Intuitively, this says that the plug-in verifier may break exact 1/2, 1/2 balance, but only by the same order as the inverse sequence length 1/n. Lemma D.1 (Biased reflecting walk hitting time). Consider a birth-death chain on {0, 1, . . . , n} with absorbing state n, deterministic transition 0 → 1, and for 1 ≤ k ≤ n − 1, k → k + 1 with probability p,
k → k − 1 with probability q := 1 − p.
Let (Kt )t≥0 denote this birth-death chain, and define the hitting time of depth n by Tn(p) := inf {t ≥ 0 : Kt = n} . Set (p)
H0
h i := E Tn(p) | K0 = 0 .
If p ̸= q, then (p) H0 =
n 2pq − p − q (p − q)2
n q 1− . p
If p = q = 1/2, then (1/2)
H0
= n2 .
p−
1 c ≤ , 2 2n
Moreover, for every fixed c > 0, if
53
then
(p)
H0
= Θc (n2 ).
Proof. Let h i hk := E Tn(p) | K0 = k . Then hn = 0, and the reflecting root gives h0 = 1 + h1 . For 1 ≤ k ≤ n − 1, hk = 1 + phk+1 + qhk−1 . Define finite differences ∆k := hk − hk−1 for 1 ≤ k ≤ n. The boundary equation gives ∆1 = −1, and the interior recurrence is equivalent to p∆k+1 − q∆k = −1. If p ̸= q, set r := Then the solution is
q , p
Sn (r) :=
rℓ .
ℓ=0
∆k =
n−1 X
−1 +
1 p−q
rk−1 −
1 . p−q
Since hn = 0, −h0 = hn − h0 =
n X
∆k =
−1 +
k=1
1 p−q
Sn (r) −
n . p−q
Rearranging and using Sn (r) = (1 − rn ) / (1 − r) gives h0 = Sn (r) +
n − Sn (r) n 2pq = − p−q p − q (p − q)2
1−
n q . p
If p = q = 1/2, the recurrence becomes hk = 1 + 12 hk+1 + 12 hk−1 , and the solution satisfying hn = 0 and h0 = 1 + h1 is hk = n2 − k 2 . Thus h0 = n2 . Finally suppose |p − 1/2| ≤ c/(2n). Then |p − q| ≤ c/n. If p = q, the claim follows from h0 = n2 . It remains to consider p ̸= q. Recall that r = q/p; equivalently, r=
1 − (p − q) q = p 1 + (p − q)
because p + q = 1. For n ≥ 2c, we have |p − q| ≤ 1/2. The mean-value theorem then gives |log r| = |log (1 − (p − q)) − log (1 + (p − q))| ≤ 4 |p − q| . Therefore, for 0 ≤ ℓ ≤ n − 1, e−4c ≤ e−4|p−q|n ≤ rℓ ≤ e4|p−q|n ≤ e4c . Summing these bounds gives ne−4c ≤ Sn (r) =
n−1 X
rℓ ≤ ne4c ,
ℓ=0
so Sn (r) = Θc (n). The finitely many cases n < 2c are absorbed into the constants hidden by Θc . Moreover, n − Sn (r) =
n−1 X
n−1 ℓ−1 XX
ℓ=0
ℓ=1 j=0
(1 − rℓ ) = (1 − r)
rj ,
|1 − r| = Θc (|p − q|).
The double sum is Θc (n2 ), again because all powers rj stay within constant factors. Thus |n − Sn (r)| = Θc (|p − q| n2 ), and (n − Sn (r))/(p − q) = Θc (n2 ). Substituting into h0 = Sn (r) + (n − Sn (r))/(p − q) gives h0 = Θc (n2 ).
54
Proposition D.2 (No-momentum remains diffusive under near-balance). Assume Assumption D.1 with εn ≤ c/n for a fixed constant c > 0. Consider the self-loop-free no-momentum balanced AOAR-VGB depth process, started from depth 0. Then E TnBAL | τ0BAL = 0 = Θc (n2 ). Proof. At every interior state, the probability of increasing the depth is between p− :=
1 − εn 2
p+ :=
and
1 + εn . 2
Because the conditional probability of increasing depth at each interior step is always in [p− , p+ ], the no-momentum depth process can be coupled between two reflecting birth-death chains using common uniform random variables at each step. The favorable chain increases depth with probability p+ and reaches depth n no later than the AOAR depth process, while the unfavorable chain increases depth with probability p− and reaches depth n no earlier. Therefore, (p ) (p ) H0 + ≤ E TnBAL | τ0BAL = 0 ≤ H0 − . (p )
(p )
Since |p± − 1/2| ≤ c/(2n), Lemma D.1 gives H0 + = Θc (n2 ) and H0 − = Θc (n2 ). Hence E[TnBAL | τ0BAL = 0] = Θc (n2 ). Theorem D.3 (First-leaf hitting for near-maximal flow cancellation). Assume Assumption D.1. Run e fc,χ from the initial lifted state (∅, ↓), where χ ∈ [0, 1]. the flow-cancelled momentum chain K BAL Then n h i 4n 2 − χ(1 − εn ) BAL,Mo e E Tn | Z0 = (∅, ↓) ≤ +1 . 1 − εn 1 − εn In particular, if εn ≤ min {c/n, 1/2} and 1 − χ ≤ cχ /n for fixed constants c, cχ > 0, then h i E TnBAL,Mo | Ze0 = (∅, ↓) = Oc,cχ (n). For χ = 1, this recovers the maximally flow-cancelled bound. Proof. Let qε := (1 − εn )/2. Under Assumption D.1, at every interior state z, p↓ (z), p↑ (z) ≥ qε . From (z, ↓), the flow-cancelled kernel has downward move probability p↓ (z), switch-up probability p↑ (z) − χm(z), and same-copy holding probability χm(z), where m(z) := min{p↓ (z), p↑ (z)}. Thus, while in downward mode, the probability of a downward depth increment is at least qε . We decompose the trajectory into downward attempts. A downward attempt starts in downward mode. It succeeds if the chain reaches depth n while remaining in downward mode, and it fails if the chain switches from downward mode to upward mode before reaching depth n. Same-copy holding steps do not change either the depth or the momentum label. For the successprobability calculation, condition them out and look only at decisive steps, namely steps on which either a downward move or an upward switch occurs. Given the current lifted state (z, ↓), p↓ (z) p↓ (z) + p↑ (z) − χm(z) p↓ (z) qε 1 − εn = ≥ := αχ,ε = . 1 − χm(z) 1 − χqε 2 − χ(1 − εn )
P(downward move | decisive, z, ↓) =
To succeed from depth k, an attempt only needs the next n − k decisive outcomes to be downward n moves. Therefore, with sχ,ε,n := αχ,ε , n−k n P(attempt succeeds from depth k) ≥ αχ,ε ≥ αχ,ε .
Thus every downward attempt succeeds with probability at least sχ,ε,n , uniformly over the past. Next, we bound the expected duration of one attempt cycle, counting Markov transitions including same-copy holding steps. During the downward part, every transition in downward mode increases 55
the depth with probability at least qε . To upper-bound the time cost, ignore early switch failures and wait until n downward depth increments have occurred. If Gr is the waiting time for the r-th such increment after the (r − 1)-st one, then P(Gr > t) ≤ (1 − qε )t and E[Gr ] ≤ 1/qε . Hence the expected duration of the downward part is at most n/qε . If the attempt fails, the chain has switched to upward mode before reaching depth n. We upper-bound the cost of preparing the next attempt by a conservative reset: move upward until the root and then take the boundary switch back to downward mode. Returning to the root requires at most n upward depth decrements, each occurring with probability at least qε . Thus the reset duration is at most n/qε + 1 in expectation, and one complete attempt cycle has expected duration at most Lε := 2n/qε + 1. Let Natt be the number of attempts until the first successful one. Since each attempt succeeds with conditional probability at least sχ,ε,n , we have P(Natt > m) ≤ (1 − sχ,ε,n )m and therefore E[Natt ] ≤ 1/sχ,ε,n . Let Cj denote the duration of the j-th attempt cycle. The cycle-length bound aboveP is uniform over the past, so E[Cj | Hj−1 ] ≤ Lε whenever the j-th cycle starts. Since TnBAL,Mo ≤ j≥1 1{Natt ≥ j}Cj , we obtain X 2n BAL,Mo −n E[Tn ] ≤ Lε P(Natt ≥ j) = Lε E[Natt ] ≤ + 1 αχ,ε . qε j≥1
Substituting qε = (1 − εn )/2 and αχ,ε = (1 − εn )/(2 − χ(1 − εn )) gives n 2 − χ(1 − εn ) 4n BAL,Mo +1 . E[Tn ]≤ 1 − εn 1 − εn Finally, suppose εn ≤ min {c/n, 1/2} and 1 − χ ≤ cχ /n. Then 2 − χ(1 − εn ) 1 − χ + εn (1 + χ) 1 =1+ = 1 + Oc,cχ . 1 − εn 1 − εn n Therefore its n-th power is Oc,cχ (1), while (4n)/(1 − εn ) + 1 = Oc (n). Hence h i E TnBAL,Mo | Ze0 = (∅, ↓) = Oc,cχ (n).
Corollary D.1 (Depth-axis speedup with near-maximal momentum). Assume Assumption D.1 with c 1 εn ≤ min , n 2 for a fixed constant c > 0. Let the flow-cancelled momentum parameter satisfy cχ 1−χ≤ n for a fixed constant cχ > 0. Then the no-momentum balanced AOAR-VGB depth process satisfies E TnBAL | τ0BAL = 0 = Θc (n2 ), while the flow-cancelled momentum lift satisfies h i E TnBAL,Mo | Ze0 = (∅, ↓) = Oc,cχ (n). Thus, under the same near-balance condition and near-maximal cancellation, E[TnBAL | τ0BAL = 0] = Ωc,cχ (n). e0 = (∅, ↓)] E[TnBAL,Mo | Z The maximally flow-cancelled case χ = 1 is included by taking cχ = 0.
56
Proof. This follows immediately from Proposition D.2 and Theorem D.3. Corollary D.2 (Ratio form of the near-balance condition). Suppose there exists ρn ≥ 1 such that for every reachable interior state z, 1 p↓ (z) ≤ ≤ ρn . ρn p↑ (z) Then
ρn − 1 . ρn + 1 Consequently, if there are fixed constants c, cχ > 0 such that c 1 cχ ρn − 1 ≤ min , , 1−χ≤ , ρn + 1 n 2 n |p↓ (z) − p↑ (z)| ≤
then the depth-axis speedup in Corollary D.1 applies. In particular, the ratio condition holds asymptotically when ρn = 1 + O(1/n), provided the cancellation is near-maximal. Proof. Let p = p↓ (z) and q = p↑ (z). Since p + q = 1 and p/q ≤ ρn , we have p ≤ ρn q = ρn (1 − p), so p ≤ ρn /(1 + ρn ). Similarly, q/p ≤ ρn gives p ≥ 1/(1 + ρn ). Therefore, p−
1 ρn − 1 ≤ . 2 2(ρn + 1)
Since |p − q| = |2p − 1|, we obtain |p − q| ≤ (ρn − 1)/(ρn + 1). The displayed condition above gives Assumption D.1 with εn ≤ min {c/n, 1/2} and the near-maximal cancellation condition 1 − χ ≤ cχ /n, so Corollary D.1 applies. If ρn = 1 + O(1/n) and 1 − χ = O(1/n), then these conditions hold for all sufficiently large n. Corollary D.3 (Implication of multiplicative verifier accuracy). Assume the canonical balanced coefficients βk = k and αk = n − k for 1 ≤ k ≤ n − 1. Suppose the plug-in verifier satisfies the multiplicative accuracy condition κ−1 V ∗ (x, z) ≤ Vb (x, z) ≤ κV ∗ (x, z) on all relevant non-leaf states. Then, for every reachable interior state z, p↓ (z) 1 ≤ ≤ κ2 . 2 κ p↑ (z) Therefore, by Corollary D.2, the linear first-leaf hitting guarantee is certified whenever, for fixed constants c, cχ > 0, κ2 − 1 c 1 cχ ≤ min , , 1−χ≤ . κ2 + 1 n 2 n Asymptotically, certifying this condition from κ alone requires κ2 = 1 + O(1/n), equivalently κ = 1 + O(1/n), together with near-maximal cancellation. Proof. Let F (z) and B(z) denote the plug-in forward and backward total masses at z. Under the exact value V ∗ , the canonical balanced choice gives F ∗ (z) = B ∗ (z) by Proposition C.4. Under the multiplicative verifier condition, the plug-in forward total satisfies κ−1 F ∗ (z) ≤ F (z) ≤ κF ∗ (z), because each child value Vb (x, c) is within a factor κ of V ∗ (x, c). Similarly, κ−1 B ∗ (z) ≤ B(z) ≤ κB ∗ (z). 57
Since F ∗ (z) = B ∗ (z), it follows that F (z) 1 ≤ ≤ κ2 . κ2 B(z) But
p↓ (z) F (z)/(F (z) + B(z)) F (z) = = . p↑ (z) B(z)/(F (z) + B(z)) B(z) Applying Corollary D.2 with ρn = κ2 gives the stated sufficient condition. In particular, to certify the linear first-leaf hitting guarantee from this multiplicative bound alone, it is enough asymptotically that κ2 = 1 + O(1/n), equivalently κ = 1 + O(1/n), together with 1 − χ = O(1/n).
E
Balanced Parallel VGBs for Masked Diffusion
The previous AOAR sections study single-site reveal and re-mask moves. In this section, we extend the same VGB principle to parallel masked-diffusion style updates, motivated by discrete and masked diffusion language models [3, 24]. The resulting MDM-VGB family allows a move to reveal or re-mask multiple coordinates at once. The construction proceeds in three steps. First, we define a leaf-corrected balanced MDM-VGB, which is the direct block-update analogue of Balanced AOAR-VGB. Second, we introduce a clean geometric balanced MDM-VGB whose backward remasking rule is value-gated while avoiding any explicit Mref correction in the local implementation. Third, we add a unified flow-cancelled momentum framework that applies to either base kernel once its local forward and backward weights have been specified. In particular, AOAR-VGB is recovered by the singleton-block restriction M = {1}. Under this restriction, the MDM algorithms below also serve as the implementation templates for the AOAR samplers. E.1
Balanced MDM-VGB
Setup and notation. We reuse the shared notation from Sections B and C. In particular, x is the n conditioning context, Y = V n is the set of fully revealed sequences, Z = (V ∪ {[mask]}) is the masked state space, R(z) is the set of revealed coordinates, k(z) = |R(z)| is the depth of z, and C(z) is the compatible completion set. Recall that U ∗ (x, z) = Mref (x, z)V ∗ (x, z) is the tilted completion mass, and recall the reachable masked-state set Zreach := {z ∈ Z : Mref (x, z) > 0} . Let M ⊆ {1, . . . , n} be a nonempty set of admissible update sizes. We assume 1 ∈ M, so every non-root state has at least one possible backward update and every non-leaf state has at least one possible forward update. For a reachable state z and a block size r ∈ M, define Br+ (z) := {A ⊆ [n] \ R(z) : |A| = r} ,
Br− (z) := {A ⊆ R(z) : |A| = r} .
For a forward block A ∈ Br+ (z) and assignment aA ∈ V A , write z A←aA for the state obtained by filling all coordinates in A with aA . For a backward block A ∈ Br− (z), write z −A for the state obtained by re-masking all coordinates in A. The MDM oracle assumption is that, for every reachable z, every A ⊆ [n]\R(z), and every aA ∈ V A , we can evaluate the joint reference conditional πref (YA = aA | x, z) := Pπref [YA = aA | x, Y ∈ C(z)] . Equivalently, πref (YA = aA | x, z) =
58
Mref (x, z A←aA ) . Mref (x, z)
Remark E.1 (Practical joint conditionals). In a practical masked diffusion model, the exact joint conditional πref (YA = aA | x, z) may not be directly available. One may replace it by a factorized approximation, for example Y πg πref (Yi = ai | x, z). ref (YA = aA | x, z) := i∈A
This factorized plug-in is a practical approximation: the exact stationarity statements below use the exact joint conditional. Block identities. identities.
We first collect the block analogues of the single-coordinate partition and Bellman
Lemma E.1 (Block partition identity). Fix a reachable state z and a block A ⊆ [n] \ R(z). Then G C(z A←aA ). C(z) = aA ∈V A
Consequently, X
X
U ∗ (x, z A←aA ) = U ∗ (x, z),
aA ∈V A
Mref (x, z A←aA ) = Mref (x, z).
aA ∈V A
Proof. Every full sequence y ∈ C(z) has a unique restriction yA ∈ V A . Hence y belongs to exactly one completion set C(z A←aA ), namely the one with aA = yA . Summing πref (y | x)τ (x, y) over this disjoint partition gives the U ∗ identity. Setting τ ≡ 1 gives the Mref identity. Proposition E.1 (Block Bellman identity). For every reachable state z and every block A ⊆ [n] \ R(z), X V ∗ (x, z) = πref (YA = aA | x, z)V ∗ (x, z A←aA ). aA ∈V A
Proof. By the law of total expectation, V ∗ (x, z) = Eπref [τ (x, Y ) | x, Y ∈ C(z)] X = Pπref [YA = aA | x, z] Eπref [τ (x, Y ) | x, Y ∈ C(z), YA = aA ] . aA ∈V A
The event {Y ∈ C(z), YA = aA } is exactly Y ∈ C(z A←aA ) . Therefore the inner conditional expectation equals V ∗ (x, z A←aA ). Proposition E.2 (Exact block target conditional). For every reachable state z with U ∗ (x, z) > 0, every block A ⊆ [n] \ R(z), and every aA ∈ V A , π ∗ (YA = aA | x, z) =
V ∗ (x, z A←aA ) U ∗ (x, z A←aA ) = π (Y = a | x, z) . ref A A U ∗ (x, z) V ∗ (x, z)
Proof. The first equality follows from conditional probability under the tilted target: π ∗ (YA = aA | x, z) =
Pπ∗ (Y ∈ C(z A←aA ) | x) U ∗ (x, z A←aA ) = . Pπ∗ (Y ∈ C(z) | x) U ∗ (x, z)
For the second equality, use U ∗ (x, z) = Mref (x, z)V ∗ (x, z) and Mref (x, z A←aA ) = Mref (x, z)πref (YA = aA | x, z).
59
Leaf-corrected balanced MDM edge weights. For each update size r ∈ M, define the depth-size edge coefficient 1 sk,r := n−r , 0 ≤ k ≤ n − r. k
Then
n−k k sk,r = sk−r,r r r whenever r ≤ k ≤ n − r. This is the multi-token analogue of the single-site balanced coefficient sk = 1/ n−1 k . For a reachable state z, define the MDM child, parent, and full neighborhoods by CMDM (z) := z A←aA : r ∈ M, A ∈ Br+ (z), aA ∈ V A , z A←aA ∈ Zreach , PMDM (z) := z −A : r ∈ M, A ∈ Br− (z) , NMDM (z) := CMDM (z) ∪ PMDM (z). Definition E.1 (Leaf-corrected balanced MDM edge weights). Consider a parent state z of depth k = k(z) and a child u = z A←aA , where A ∈ Br+ (z) and r = |A|. Thus u has depth k + r. Define the symmetric edge weight ( sk,r Mref (x, u)Vb (x, u), u∈ / Y, fMDM (z, u) := sn−r,r Mref (x, u)τ (x, u), u ∈ Y. b (x, u) := Mref (x, u)Vb (x, u) and U ∗ (x, y) = Mref (x, y)τ (x, y) for leaves, Equivalently, using U ( b (x, u), sk,r U u∈ / Y, fMDM (z, u) = sn−r,r U ∗ (x, u), u ∈ Y. The edge weight is understood symmetrically on the undirected block-update graph. Local implementation weights. The edge weights contain Mref , but the local implementation does not require estimating it. From a fixed current state z, all incident edge weights share the common factor Mref (x, z), which cancels after normalization. Definition E.2 (γ-held balanced MDM-VGB local kernel). Let z be a reachable state of depth k = k(z). We use terminal anchoring on leaves, Vb (x, y) = τ (x, y),
y ∈ Y.
Define the local implementation weights by F ORWARD : BACKWARD : where
wMDM,fwd (z → u) := sk,r πref (YA = aA | x, z)Vb (x, u), wMDM,bwd (z → u) := sk−r,r Vb (x, z),
u = z A←aA , u=z
−A
,
A ∈ Br+ (z)
A ∈ Br− (z)
for forward moves, for backward moves.
Let wMDM (z → u) denote the corresponding forward or backward weight, and set X X WMDM (z) := wMDM,fwd (z → u) + wMDM,bwd (z → u). u∈CMDM (z)
u∈PMDM (z)
The non-lazy move kernel is KMDM (z, u) :=
wMDM (z → u) , WMDM (z)
60
u ∈ NMDM (z).
This formula is used only at states with WMDM (z) > 0. The γ-held kernel is obtained by setting (γ)
PMDM (z, z) =
γ , 1+γ
(γ)
PMDM (z, u) =
1 KMDM (z, u). 1+γ
Theorem E.1 (Weighted-graph representation and exact leaf law). The local kernel in Definition E.2 is the weighted random walk on the graph with symmetric edge weights from Definition E.1. Therefore, its stationary distribution is X DMDM (z) µMDM (z) = P , DMDM (z) := fMDM (z, u). ′ z ′ DMDM (z ) u∼z Moreover, the leaf-conditioned stationary law is exact: µMDM (y | y ∈ Y) = π ∗ (y | x). Proof. Fix a state z and a forward child u = z A←aA of size r. If u ∈ / Y, then Mref (x, u) b fMDM (z, u) = sk(z),r V (x, u) Mref (x, z) Mref (x, z) = sk(z),r πref (YA = aA | x, z)Vb (x, u), which is the forward local weight. If u ∈ Y, the same calculation gives fMDM (z, u) = sn−r,r πref (YA = aA | x, z)τ (x, u). Mref (x, z) For a backward move from z to u = z −A , the edge is the same edge with child endpoint z. Dividing its edge weight by Mref (x, z) gives sk(z)−r,r Vb (x, z) if z ∈ / Y, and sn−r,r τ (x, z) if z ∈ Y. Thus the local weights are proportional to incident edge weights by a state-dependent common factor, which cancels under normalization. The stationary law of a random walk on a symmetric weighted graph is proportional to weighted degree. It remains to check the leaf law. For any leaf y ∈ Y, every admissible backward block A ⊆ [n] of size r ∈ M gives a parent y −A , and the corresponding leaf edge has weight fMDM (y −A , y) = sn−r,r U ∗ (x, y). Therefore DMDM (y) = U ∗ (x, y)
X n r∈M r≤n
r
sn−r,r .
The multiplicative factor is constant over leaves. Since U ∗ (x, y) = πref (y | x)τ (x, y), conditioning on Y gives πref (y | x)τ (x, y) µMDM (y | y ∈ Y) = P = π ∗ (y | x). ′ ′ y ′ ∈Y πref (y | x)τ (x, y )
Proposition E.3 (Exact directional balance under exact values). Assume Vb = V ∗ on all non-leaf states and terminal anchoring on leaves. Let z be an interior state of depth k. For a fixed block size r ∈ M with r ≤ min {k, n − k}, the exact forward and backward masses of size r satisfy ∗ ∗ Wfwd,r (z) = Wbwd,r (z).
Consequently, if at an interior state only paired sizes Mpair (z) := {r ∈ M : r ≤ min {k(z), n − k(z)}} are used, then the exact move-conditioned forward and backward probabilities are both 1/2. 61
Geometric MDM-VGB transition on 𝑵𝑴𝑫𝑴 𝒛 parent states (𝒌 − |𝑩|)
𝒛−𝑩𝟏
𝒙, 𝒛 𝑽 𝒙, 𝒛−𝑩 Backward / re-mask: 𝒘𝐛𝐰𝐝 𝒛 = 𝒔𝒌− 𝑩 , 𝑩 𝑽
child states (𝒌 + |𝑩|)
𝒛−𝑩𝟐
masked state
𝒛
current state (𝒌) 𝒛𝑩𝟏 ←𝒂𝑩𝟏
𝝀
Self-loop for laziness (optional)
𝒛𝑩𝟐 ←𝒂𝑩𝟐
𝒙, 𝒛𝑩←𝒂𝑩 𝝅𝒓𝒆𝒇 𝒀𝑩 = 𝒂𝑩 𝒙, 𝒛 𝒙, 𝒛 𝝀 𝑽 Forward / reveal: 𝒘𝐟𝐰𝐝 𝒛 = 𝒔𝒌, 𝑩 𝑽
Figure 14: Illustration of one geometric balanced MDM-VGB transition. Forward moves reveal an admissible block using the joint reference conditional and the child verifier value, while backward moves re-mask an admissible block using the verifier value of the remasked parent. Setting λ = 0 recovers balanced MDM-VGB. Proof. Fix r ≤ min {k, n − k}. Using Proposition E.1, for each A ∈ Br+ (z), X πref (YA = aA | x, z)V ∗ (x, z A←aA ) = V ∗ (x, z). aA ∈V A
There are
n−k r
forward blocks of size r, so ∗ Wfwd,r (z) = k r
n−k sk,r V ∗ (x, z). r
choices, and each backward weight is sk−r,r V ∗ (x, z). k ∗ Wbwd,r (z) = sk−r,r V ∗ (x, z). r The coefficient identity in Section E.1 gives equality. Summing over paired sizes preserves equality. For backward blocks of size r, there are Hence
E.2
Geometric Balanced MDM-VGB
Balanced MDM-VGB gives the block-update analogue of the balanced AOAR walk. We now add the same value-gated remasking idea used in the geometric AOAR variant. The goal is to preserve the forward VGB ranking while making the backward block choice depend on the verifier value of the remasked parent. Figure 14 shows the corresponding block-update transition. The only change from the AOAR case is that each reveal or re-mask candidate is now an admissible block A ∈ M, with joint reference conditionals used on forward block reveals. Definition E.3 (Clean geometric balanced MDM edge weights). Fix λ ≥ 0. For every admissible size r ∈ M, let −1 n−r , 0 ≤ k ≤ n − r, sk,r := k be the canonical balanced depth-size coefficient. Consider neighboring states z, u with |k(u) − k(z)| = r. Let p(z, u) be the lower-depth endpoint and c(z, u) be the higher-depth
62
endpoint, so that c(z, u) = p(z, u)A←aA for a unique block A of size r and assignment aA . If c(z, u) ∈ / Y, define b (x, c(z, u))Vb (x, p(z, u))λ . fBGM (z, u) = fBGM (u, z) := sk(p(z,u)),r U If c(z, u) = y ∈ Y, define the leaf-corrected edge by fBGM (z, u) = fBGM (u, z) := sn−r,r U ∗ (x, y). Local implementation weights. Let z have depth k. For a non-leaf forward child u = z A←aA ∈ /Y with |A| = r, wBGM,fwd (z → u) = sk,r Vb (x, z)λ πref (YA = aA | x, z)Vb (x, u). The factor Vb (x, z)λ is common across all forward candidates from the same current state z, so the forward ranking remains πref (YA = aA | x, z)Vb (x, z A←aA ). For a backward move from a non-leaf state z to u = z −A , wBGM,bwd (z → z −A ) = sk−r,r Vb (x, z)Vb (x, z −A )λ . Thus backward moves prefer blocks whose remasked parent has high verifier value. At the leaf boundary, if y = z A←aA ∈ Y, then leaf wBGM,fwd (z → y) = sn−r,r πref (YA = aA | x, z)τ (x, y),
and from a leaf y to a parent y −A , wBGM,bwd (y → y −A ) = sn−r,r τ (x, y). Remark E.2 (λ = 0 recovers balanced MDM-VGB). When λ = 0, the geometric gate disappears, and the local weights in Section E.2 reduce to the balanced MDM-VGB weights in Definition E.2. Theorem E.2 (Stationarity and exact leaf law for clean geometric MDM). Let PBGM be the weighted random walk on the block-update graph with edge weights from Definition E.3. Then X DBGM (z) fBGM (z, u), µBGM (z) = P , DBGM (z) := ′ z ′ DBGM (z ) u∼z is stationary. Moreover,
µBGM (y | y ∈ Y) = π ∗ (y | x).
Proof. The stationarity proof is the standard detailed-balance proof for a symmetric weighted graph. For neighboring states z ∼ u, DBGM (z) fBGM (z, u) fBGM (z, u) =P . ′) D ′ D (z (z) ′ BGM BGM z z ′ DBGM (z )
µBGM (z)PBGM (z, u) = P
This expression is symmetric in z and u, hence detailed balance holds. Now take a leaf y ∈ Y. For every admissible block size r ∈ M and every block A ⊆ [n] with |A| = r, the parent is y −A and the leaf edge has weight fBGM (y −A , y) = sn−r,r U ∗ (x, y). Therefore ∗
DBGM (y) = U (x, y)
X n r∈M r≤n
63
r
sn−r,r .
The multiplier is independent of y, so conditioning on leaves yields U ∗ (x, y) = π ∗ (y | x). ∗ ′ y ′ ∈Y U (x, y )
µBGM (y | y ∈ Y) = P
Proposition E.4 (Exact block directional masses). Assume Vb = V ∗ on non-leaf states. Let z have depth k with V ∗ (x, z) > 0, and fix an admissible size r with r ≤ k and k + r < n, so that the size-r forward children are non-leaf. For a backward block A ⊆ R(z) with |A| = r, define V ∗ (x, z −A ) . V ∗ (x, z)
∗ rA (z) :=
Then the exact size-r forward and backward masses are n−k ∗ Fλ,r (z) = sk,r V ∗ (x, z)1+λ , r and
∗ Bλ,r (z) = sk−r,r V ∗ (x, z)1+λ
X
∗ rA (z)λ .
A⊆R(z) |A|=r
Proof. For each forward block A ∈ Br+ (z), the block Bellman identity gives X πref (YA = aA | x, z)V ∗ (x, z A←aA ) = V ∗ (x, z). aA ∈V A
There are
n−k
forward blocks, hence n−k ∗ ∗ λ n−k ∗ Fλ,r (z) = sk,r V (x, z) V (x, z) = sk,r V ∗ (x, z)1+λ . r r r
For backward blocks, ∗ Bλ,r (z) =
X
sk−r,r V ∗ (x, z)V ∗ (x, z −A )λ
A⊆R(z) |A|=r ∗
X V ∗ (x, z −A ) λ . V ∗ (x, z)
1+λ
= sk−r,r V (x, z)
A⊆R(z) |A|=r
Remark E.3 (Canonical MDM depth anchor). For the balanced MDM chain, the canonical coefficient 1 sbal k,r = n−r k
satisfies
n − k bal k bal sk,r = s . r r k−r,r
For the clean geometric MDM chain with λ > 0, this coefficient no longer gives exact 1/2–1/2 balance in general. Instead, it is the canonical depth-neutral anchor: it recovers exact balance at λ = 0, while the λ > 0 correction changes backward probabilities according to the verifier values of remasked parent states. The sampler in Algorithm 3 specializes to the AOAR sampler when M = {1}, since each admissible block is a singleton coordinate. 64
Geometric MDM-VGB-Momentum transition on 𝑵𝑴𝑫𝑴 𝒛
𝒙, 𝒛 𝑽 𝒙, 𝒛−𝑩 Backward / re-mask: 𝒘𝐛𝐰𝐝 𝒛 = 𝒔𝒌− 𝑩 , 𝑩 𝑽 parent states (𝒌 − 𝑩 )
𝝀
𝒛−𝑩𝟐
𝒛−𝑩𝟏
Momentum change: 𝒘𝒇𝒍𝒊𝒑 𝒛, ↑ = 𝑭 𝒛 − 𝝌𝑴(𝒛)
(𝒛, ↑)
current state (𝒌)
Self-loop for laziness (optional)
(𝒛, ↓)
𝒘𝒇𝒍𝒊𝒑 𝒛, ↓ = 𝑩 𝒛 − 𝝌𝑴(𝒛)
child states (𝒌 + 𝑩 )
𝒛𝑩𝟐 ←𝒂𝑩𝟐
𝒛𝑩𝟏 ←𝒂𝑩𝟏
𝒙, 𝒛 𝝀 𝑽 𝒙, 𝒛𝑩←𝒂𝑩 𝝅𝒓𝒆𝒇 𝒀𝑩 = 𝒂𝑩 𝒙, 𝒛 Forward / reveal: 𝒘𝐟𝐰𝐝 𝒛 = 𝒔𝒌, 𝑩 𝑽
Figure 15: Illustration of the flow-cancelled momentum lift for geometric MDM-VGB. The downward copy keeps block reveal moves, the upward copy keeps block re-mask moves, and residual oppositedirection mass switches momentum; the cancelled common mass may be retained as an optional stay probability. E.3
Momentum Framework for Balanced and Geometric MDM-VGB
We now add momentum after both base MDM kernels have been defined. This separation is useful because the same lifted transition only needs the local forward and backward weights of the chosen base kernel. Figure 15 illustrates the momentum construction for the geometric MDM base kernel. The same diagram also applies to balanced MDM-VGB after setting λ = 0 in the local block weights. Let B ∈ {MDM, BGM} denote either the balanced MDM-VGB base kernel or the clean geometric balanced MDM-VGB base kernel. Let wB,fwd (z → c), wB,bwd (z → p) be its local forward and backward weights, as defined in Definition E.2 for B = MDM and in Section E.2 for B = BGM. Define the forward and backward total masses X FB (z) := wB,fwd (z → c),
BB (z) :=
c∈CMDM (z)
X
wB,bwd (z → p).
p∈PMDM (z)
The non-lazy base move kernel is wB (z → u) , u ∈ NMDM (z), FB (z) + BB (z) where wB denotes the appropriate forward or backward local weight. Let µB be its stationary law. By Theorem E.1 for B = MDM and Theorem E.2 for B = BGM, the projected leaf-conditioned law is exact: µB (· | Y) = π ∗ (· | x). Let ZB,+ be the positive-weight component of KB , i.e., the component on which FB (z) + BB (z) > 0. KB (z, u) :=
For z ∈ ZB,+ , define directional probabilities X pB↓ (z) := KB (z, c) = c∈CMDM (z)
65
FB (z) , FB (z) + BB (z)
and pB↑ (z) :=
X
KB (z, p) =
p∈PMDM (z)
BB (z) . FB (z) + BB (z)
Thus pB↓ (z) + pB↑ (z) = 1 whenever FB (z) + BB (z) > 0. Definition E.4 (Flow-cancelled momentum MDM kernel). Fix a base kernel B ∈ {MDM, BGM} and a cancellation strength χ ∈ [0, 1]. Set mB (z) := min pB↓ (z), pB↑ (z) . On the lifted state space ZeB := ZB,+ × {↓, ↑} , e fc,χ as follows. define K B Downward copy. with probability
From (z, ↓), the chain moves to a child while keeping downward momentum
e fc,χ (z, ↓), (c, ↓) := KB (z, c), K B
c ∈ CMDM (z).
(C HILD M OVE)
It switches to upward momentum with residual opposite-direction probability e fc,χ (z, ↓), (z, ↑) := pB↑ (z) − χmB (z), K B and stays in the same lifted state with cancelled crossing mass e fc,χ (z, ↓), (z, ↓) := χmB (z). K B Upward copy. probability
(S WITCH U P)
(S TAY D OWN , OPTIONAL)
From (z, ↑), the chain moves to a parent while keeping upward momentum with
e fc,χ (z, ↑), (p, ↑) := KB (z, p), K B
p ∈ PMDM (z).
(PARENT M OVE)
It switches to downward momentum with residual opposite-direction probability e fc,χ (z, ↑), (z, ↓) := pB↓ (z) − χmB (z), K (S WITCH D OWN) B and stays in the same lifted state with cancelled crossing mass e fc,χ (z, ↑), (z, ↑) := χmB (z). K B
(S TAY U P, OPTIONAL)
All other transition probabilities are zero. e fc,χ is Proposition E.5 (Markov property). For every χ ∈ [0, 1] and every B ∈ {MDM, BGM}, K B a Markov kernel. Proof. All transition probabilities are nonnegative because mB (z) ≤ pB↓ (z) and mB (z) ≤ pB↑ (z). For the downward copy, X e fc,χ (z, ↓), (c, ↓) + K e fc,χ (z, ↓), (z, ↑) + K e fc,χ (z, ↓), (z, ↓) K B B B c∈CMDM (z)
= pB↓ (z) + pB↑ (z) − χmB (z) + χmB (z) = 1. The upward row is identical. Theorem E.3 (Momentum preserves the MDM leaf law). Fix B ∈ {MDM, BGM} and χ ∈ [0, 1]. Define 1 µ eB (z, ↓) = µ eB (z, ↑) := µB (z). 2
66
e fc,χ . Moreover, Then µ eB is stationary for K B proj# µ eB = µB , and hence
proj# µ eB (· | Y) = π ∗ (· | x).
Proof. We verify stationarity for the downward copy; the upward copy is symmetric. Incoming flow into (z, ↓) comes from three sources: parents of z in downward mode, the switch from (z, ↑), and the self-loop at (z, ↓). The incoming move-flow from parents is 1 2
X
µB (u)KB (u, z) =
u∈PMDM (z)
=
1 µB (z) 2
X
KB (z, u)
u∈PMDM (z)
1 µB (z)pB↑ (z), 2
where reversibility of KB was used. The switch-flow from (z, ↑) is 1 µB (z) pB↓ (z) − χmB (z) , 2 and the self-loop flow is 1 µB (z)χmB (z). 2 The total incoming flow is therefore 1 1 µB (z) pB↑ (z) + pB↓ (z) = µB (z) = µ eB (z, ↓). 2 2 The upward-copy argument is identical after exchanging parents and children. Projection gives µ eB (z, ↓) + µ eB (z, ↑) = µB (z). The leaf law follows from the corresponding base-chain leaf law. The sampler in Algorithm 4 also specializes to the momentum AOAR sampler when M = {1}.
F
Experimental Setting and Supplementary Results
This section gives the full task definitions, data sources, backbone models, reward functions, rootstart and leaf-start protocols, and hyperparameter grids used in our experiments. Across tasks, the reference model is denoted by πref and the masked-state verifier by Vb (x, z), following the notation in Section A. For hard-constraint tasks, the terminal reward τ (x, y) is binary. For QM9 and DNA, we first define a real-valued terminal score R(y) and then use the corresponding Pass@95 indicator as the terminal reward when a binary reward is needed. Protein reports Success@1Å and motif RMSD, while its learned verifier uses the smoothed training target defined below. Unless otherwise stated, VGB uses terminal anchoring: at a complete leaf y, Vb (x, y) is replaced by the corresponding exact terminal reward or score. Experiments were run in an Ubuntu PyTorch environment on NVIDIA RTX PRO 6000 Blackwell 96GB and NVIDIA H200 141GB GPUs. Common evaluation protocol. For root-start sampling, every method starts from the fully masked state ∅ and produces a completed sample under a fixed step-limit budget. For leaf-start repair, every method starts from the same set of low-reward or invalid completed leaves y0 . All comparisons use the same reference model πref and the same task-specific terminal score/reward convention. In matched-budget root-start comparisons, B O N uses N independent full rollouts, while MDM-VGB variants use up to an n × N step budget before forward-only completion. 67
Adjusted NFE. Raw NFE counts realized reference-model forward calls. For tasks with learned verifiers, we additionally report FLOP-adjusted NFE to account for verifier scoring cost. Let Nbase be the number of reference-model forward calls, Nverif the number of verifier evaluations, and Cbase and Cverif the FLOPs of one base-model and verifier forward pass. We define NFEadj = Nbase + Nverif
Cverif . Cbase
Thus each verifier query is charged as its equivalent fraction of a base-model forward call. For task-provided exact checkers, this adjustment is not applied unless otherwise stated. Pass@α normalization. For tasks with continuous rewards, let R(y) denote the target task-specific score. To make success rates comparable across tasks, we convert the continuous score into a rareevent indicator using a fixed threshold from an independent BASE calibration set. Specifically, for α ∈ (0, 1), let qαBase (R) be the α-quantile of R(y) on this calibration set, where y ∼ πref . We define 1{R(y) ≥ qαBase (R)}, if higher is better, Pass@α(y) := Base 1{R(y) ≤ q1−α (R)}, if lower is better. The reported Pass@95 metric corresponds to α = 0.95. F.1
Task Descriptions
Dyck suffix-edit repair. The Dyck language is the canonical language of well-balanced brackets studied in formal-language theory [6]. It is simple enough to serve as a controlled toy problem, but it still requires hierarchical dependency tracking: for example, [()] is valid, whereas ([)] has the right local symbols in the wrong nesting order. We use this structure to illustrate the difference between prefix-tree backtracking and any-order re-masking. The terminal reward is binary: reward 1 means that all opening brackets are properly paired with closing brackets of the same type and in the correct nested order; reward 0 means that this wellbalanced-bracket condition fails. Each repair example is initialized from the fixed corrupted leaf template B((([(((([(((XXXXXXXX)))]))))])))E. Here B and E are special begin/end tokens. The prefix before the marked block is fixed and used as the prompt, while the X’s denote the initially corrupted span of length eight. We generate 10,000 corrupted leaves by sampling random bracket strings for this span and retaining only examples with reward 0. Thus every example starts from an invalid fully revealed leaf. In the main suffix-edit protocol, all bracket positions from the first corrupted token through the token immediately before E are editable. The suffix is not manually restored: if a sampler re-masks or changes a suffix token, it must also repair that token before the exact Dyck checker accepts the sequence. Equivalently, the reward is the exact bracket-validity checker implemented as membership in the accepted terminal set. The base model is a 12.9M-parameter BERT-style masked language model, and the process verifier is a 3.24M-parameter token-state Transformer. Letter avoidance.
Letter is a simple lexical positive-control task. The fixed task prompt is
Task prompt. Generate a one-sentence story without using the letter ‘e’:
This is a lipogram-style constraint, where a text is deliberately written while avoiding a specified letter. The terminal reward is exact: 1.0, y is non-empty and contains no ‘e’, τ (x, y) = 0.0, otherwise. All Letter experiments use the Qwen3-0.6B-diffusion-mdlm-v0.1 reference model [10]. For partial configurations, we use a heuristic process verifier that checks whether the revealed tokens currently 68
contain the forbidden character. The primary metric is the avoidance rate. We additionally report a Qwen3-32B [31] judge score on outputs satisfying the lexical constraint. The judge prompt below is used with the task prompt above and each candidate answer; the reported Ovl. score is the returned overall field. Qwen3-32B overall-score prompt: You are evaluating a candidate answer to a writing task. Score the candidate on four dimensions from 1 to 10, where 10 is best. Dimensions: 1. naturalness: Does the text read like fluent natural English? 2. coherence: Is it internally coherent as a one-sentence story? 3. semantic_plausibility: Does it make semantic sense rather than feeling like token salad or broken text? 4. overall: Overall quality as an answer to the task prompt. Use the task prompt as context when judging overall quality. Return ONLY a JSON object with integer fields naturalness, coherence, semantic_plausibility, overall.
Sudoku. Sudoku is a structured constraint-satisfaction benchmark. Each puzzle is a flattened 9 × 9 grid with fixed clues and editable empty cells. The output is a completed board. The terminal checker is exact: a board solves the puzzle if every row, column, and 3 × 3 box is a permutation of {1, . . . , 9}. The primary metric is exact solve rate. We also report the normalized mean violation count, shown as a percentage: 1 X Viol(y) = (9 − |unique(yU )|) , 27 U ∈U
where U is the set of the 27 Sudoku units: 9 rows, 9 columns, and 9 boxes. Thus, for a fully filled board with fixed clues preserved, Solve(y) = 1
⇐⇒
Viol(y) = 0.
We use the Sudoku data provided by Kim et al. [13], with 48,000 training puzzles and 2,000 test puzzles. Our base model is a 28.6M-parameter DiT-MDM, i.e., a masked discrete diffusion model [24] with a Diffusion Transformer backbone [21], trained on this data. For partial configurations, we use a heuristic process verifier that returns the indicator that no Sudoku constraint has yet been violated. QM9 molecule generation. QM9 is a small-molecule generation benchmark over organic molecules with up to nine heavy atoms [22]. We use a split with 127,190 training molecules and 6,695 test molecules, and train a QM9 MDLM 92.4M base model with the MDLM objective using the UDLM-QM9 architecture. The terminal score RQM9 uses QED [4], a quantitative estimate of drug-likeness: QED(y), y is a valid molecule, RQM9 (y) = 0, y is invalid. Since higher QED is better, the terminal reward used for rare-success guidance and evaluation is Pass@95 from the common definition above with R(y) = RQM9 (y). We also report mean QED over valid molecules, validity, and unique Pass@95 where appropriate. DNA enhancer design. DNA is a regulatory-sequence design task. The reference model is D3LM [33], a DNA diffusion language model initialized from Nucleotide Transformer [8]. The terminal scorer is DeepSTARR [9], which predicts developmental and housekeeping enhancer activity from a 249 bp DNA sequence. The terminal score uses the developmental enhancer activity: RDNA (y) = sDev (y). For verifier training, we use a smoothed target rather than the sparse Pass@95 indicator. Let sDev (y) − µDev zDev (y) = clip , −3, 3 , r̃train (y) = exp(ηzDev (y)). σDev 69
Here µDev and σDev are computed from base-model calibration samples, and we use η = 1. This target is used only for process-verifier regression. Since higher developmental enhancer activity is Base better, the evaluation terminal reward is the Pass@95 indicator 1{sDev (y) ≥ q0.95 (sDev )}. Protein motif scaffolding. Protein is a motif-scaffolding sequence-design task. The reference model is EvoDiff OADM 38M [1]. We use official EvoDiff motif-scaffolding task metadata whenever available: PDB code, motif ranges, fixed motif residues, scaffold length range, and generated motif positions. In our task, all reported Protein experiments use the 1YCR motif-scaffolding target. The generated sequence is folded with OmegaFold [30]. We compute motif Cα root-mean-square deviation (RMSD) against the reference motif: rmotif (y) = RMSDCα (Fold(y)M , PDBM ) . Here M denotes the motif residue positions, Fold(y) is the OmegaFold-predicted structure, PDBM contains the reference motif Cα coordinates, and RMSD is computed after optimal rigid alignment. The hard success metric is Success@1Å(y) = 1{rmotif (y) ≤ 1.0Å}. For verifier training, we use a smoothed RMSD-derived target " 2 # rmotif (y) , τtrain (y) = exp − 4.0Å where rmotif (y) is the motif Cα RMSD. In implementation, the verifier regresses log τtrain (y) for numerical stability, and we exponentiate its output at inference time to obtain the positive value used by MDM-VGB. This verifier-training target is distinct from the reported Protein metrics: Success@1Å and mean motif RMSD. F.2
Leaf-start repair protocol
For leaf-start repair, we choose bad but repairable leaves. This avoids turning the repair experiment into full regeneration. For the α-sweep in Figure 6, Pass@α is defined by the same base-calibration rule as in the common evaluation protocol: higher-is-better scores use qαBase , while lower-is-better Base scores use q1−α . Each starting group contains completed samples with Pass@95(y0 ) = 0 but Pass@α(y0 ) = 1. Thus larger α selects easier near-miss leaves that are closer to the rare-success boundary. The main-text leaf-start table entries average over α ∈ {60, 65, . . . , 90}. Unless otherwise stated, the bad-leaf pools are summarized in Table 6. Table 6: Leaf-start bad-leaf pools for the α-sweep.
F.3
Task
Bad-leaf pool
QM9 DNA Protein
Base Base y0 valid, qα (QED) ≤ QED(y0 ) < q0.95 (QED) Base Base qα (RDNA ) ≤ RDNA (y0 ) < q0.95 (RDNA ) Base Base q0.05 (rmotif ) < rmotif (y0 ) ≤ q1−α (rmotif )
Verifier Architecture
We use task-specific process verifiers. Letter and Sudoku use heuristic process verifiers with no learned parameters, while Dyck, QM9, DNA, and Protein use learned lightweight Transformer verifiers. The learned verifiers are trained by regression to the task-specific targets described in the task descriptions above. • Dyck: a learned token-state Transformer process verifier (3.24M parameters). • Letter: a heuristic process verifier that checks whether the revealed tokens currently contain the forbidden letter; it has no learned parameters. • Sudoku: a heuristic process verifier that checks whether the revealed entries already violate row, column, box, or fixed-clue constraints; it has no learned parameters. 70
• QM9: a learned token Transformer process verifier over SMILES tokens (3.24M parameters). • DNA: a learned Transformer process verifier over D3LM tokens (0.97M parameters). • Protein: a learned Transformer process verifier over amino-acid tokens with motif/editable/mask features (1.94M parameters). We summarize the main inference hyperparameters in Table 7. Task
|V|
Length
Dyck Letter Sudoku QM9 DNA Protein (1YCR)
34 32 89 32 48 73
Lf , Lb K
8 (8, 8) 8 151,670 (8, 8) 8 11 (8, 8) 8 35 (8, 32) 8 4107 (8, 8) 32 20 (8, 8) 20
Table 7: Main inference hyperparameters. Lf and Lb are the numbers of shortlisted forward and backward coordinates or blocks. K is the number of token candidates per forward coordinate.
F.4
Additional Ablations
The main text focuses on geometric re-masking and verifier amortization. Here we collect additional ablations on shortlisting and the block size. Shortlisting budget. We ablate the local shortlisting budget for MDM-VGB-M OMENTUM by increasing Lf = Lb while holding the task-specific token budget K fixed. On QM9, increasing the shortlist from 1 to 4 substantially improves Pass@95 while also reducing realized search cost, and larger budgets give smaller additional gains. On DNA, Pass@95 peaks around Lf = Lb = 4, whereas Lf = Lb = 8 substantially increases adjusted NFE without improving quality. Thus the gains are not explained by simply evaluating an arbitrarily large local neighborhood; a moderate shortlist is sufficient for the verifier to expose useful moves. MDM-VGB-Momentum
DNA
Adjusted NFE ↓
Adjusted NFE ↓
QM9
350 300 250
2600
1800
1000 Lf = Lb
Lf = Lb
Pass@95 (%) ↑
Pass@95 (%) ↑
80
60
40 1
2
4
8
96
92
88 2
Lf = Lb
4 Lf = Lb
8
Figure 16: Shortlisting-budget ablation for MDM-VGB-M OMENTUM on QM9 and DNA. Increasing Lf = Lb helps up to a moderate budget, after which quality saturates while adjusted NFE can continue to grow.
71
Block size / MDM precision. We next study how MDM block size affects the finite-budget reward– compute tradeoff for MDM-VGB-M OMENTUM. Singleton AOAR updates give the finest local control, while larger MDM blocks can reduce the number of sequential decisions by revealing or re-masking multiple coordinates at once. This creates a precision–parallelism tradeoff: moderate block sizes may amortize verifier and proposal costs, but overly coarse blocks can make the local edit decisions too blunt.
90
500 Pass@95 (%) ↑
Adjusted NFE ↓
MDM-VGB-Momentum
400 300 200 1
2
4 MDM block size |B|
6
80 70 60 50 1
2
4 MDM block size |B|
6
Figure 17: QM9 root-start MDM-VGB-M OMENTUM block-size ablation under a maximum budget corresponding to N = 16, with Lf = Lb = 4, K = 4, and λ = 4. The left panel reports adjusted NFE, including verifier FLOPs, and the right panel reports Pass@95. Moderate block sizes improve the compute–quality tradeoff, whereas overly large blocks increase cost and degrade reward satisfaction.
72