d2-StepMerge in DLMs and Discrete Morse Theory

• d2-StepMerge is a parallel merging method that composes partial likelihoods in masked diffusion language models and assembles discrete gradient fields in Morse theory.
• The algorithm partitions global problems into manageable blocks to reduce computation while controlling bias and ensuring robust reinforcement learning and topological accuracy.
• Empirical validations demonstrate state-of-the-art performance with significant accuracy improvements on reasoning benchmarks and reliable assembly of global gradient fields.

d2-StepMerge refers to parallel merging procedures arising independently in the contexts of policy-estimation for masked diffusion LLMs (DLMs), as well as the assembly of discrete Morse vector fields from local computations on subcomplexes. The unifying theme is the efficient composition of partial, locally computed information—either likelihoods or gradient fields—into a consistent global object, under constraints of compute or locality. The d2-StepMerge methods combine compute-efficient, blockwise processing with explicit bias or error control, supporting large-scale reinforcement learning for DLMs and highly parallel discrete Morse theory computations.

1. StepMerge in Diffusion LLMs and Discrete Morse Theory

In masked DLMs, the d2-StepMerge estimator is an efficient, approximately unbiased method for estimating trajectory likelihoods needed for policy-gradient reinforcement learning (Wang et al., 25 Sep 2025). For discrete Morse theory, d²–StepMerge is a parallel merging protocol for assembling discrete gradient vector fields, as described in Lenseth–Goldfarb (Lenseth et al., 2021). Both contexts rely on decomposing the global problem—whether the log-likelihood of a long denoising trajectory or the construction of a gradient field on a large complex—into smaller, manageable blocks ("trajectory segments" or "patches"), with systematic recombination.

2. d2-StepMerge for Diffusion LLMs: Definition and Derivation

Given a DLM backward process forming a length-$T$ Markov chain over $L$ token positions, the exact trajectory likelihood is

$$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$

where $1_{t,l}$ indicates unmasking of token $l$ at step $t$. Computing this exactly incurs $T$ full forward model evaluations per trajectory, which is computationally prohibitive for large $T$. d2-StepMerge divides the $T$ steps into $N$ contiguous blocks ($L$0), computes the likelihood only at block boundaries, and reuses it within each block:

$L$1

with the log-likelihood

$L$2

This reduces computation from $L$3 to $L$4 model passes per trajectory with a bias that admits analytic control (Wang et al., 25 Sep 2025).

3. Algorithmic Implementation and Complexity

The d2-StepMerge algorithm for policy learning in DLMs proceeds as follows:

• Sample complete denoising trajectories under the current policy.
• Partition each trajectory into $L$5 blocks; evaluate $L$6 at block boundaries only.
• Compute approximate log-likelihoods and GRPO-style importance weights in blockwise fashion.
• Use these to form policy-gradient and KL-regularized updates.
• Repeat for multiple updates per batch until convergence.

Pseudocode from (Wang et al., 25 Sep 2025) concretizes these steps. The overall complexity per on-policy step is $L$7, where $L$8 is model parameter count, compared to $L$9 for exact likelihood and $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$0 for any-order decoding. The chief hyperparameter is the block number $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$1, controlling the bias–compute trade-off. As $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$2, bias vanishes; as $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$3, this recovers the diffu-GRPO one-block heuristic (highest bias, lowest cost).

Estimator	# Model Passes	Applicability
Exact	$$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$4	Any masked DLM
d2-AnyOrder	$$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$5	Only models with any-order decoding support
d2-StepMerge	$$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$6	Any masked DLM; bias controlled by $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$7

Empirically, d2-StepMerge achieves state-of-the-art performance on logical and mathematical reasoning tasks (e.g., Sudoku, GSM8K, MATH500), with accuracy increases from ∼10% (base LLaDA) to above 90% (Sudoku, $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$8) under fixed computational budget (Wang et al., 25 Sep 2025).

4. Theoretical Bias–Compute Trade-Off and Empirical Validation

The theoretical bias of the StepMerge estimator, measured in KL divergence $$\pi(x_{0:T}^{1:L}) = \prod_{t=0}^{T-1} \prod_{l=1}^L 1_{t,l} \cdot \pi_\theta(x_t^{l} \mid x_{t+1}^{1:L}),$$9 between the true and blockwise factorization, is given by (Theorem 4.2 in (Wang et al., 25 Sep 2025)):

$1_{t,l}$0

where $1_{t,l}$1 captures softmax logit change across skipped steps. Empirical plots (Figure 1, (Wang et al., 25 Sep 2025)) confirm that $1_{t,l}$2 decreases monotonically with $1_{t,l}$3, following the log-bound. Ablations over $1_{t,l}$4 identify $1_{t,l}$5 as a robust regime balancing computation and estimator fidelity on reasoning benchmarks, with diminishing returns for larger $1_{t,l}$6.

5. StepMerge for Discrete Morse Vector Fields: Principles and Formalism

In discrete Morse theory on a finite regular CW-complex $1_{t,l}$7, d²–StepMerge provides a method for assembling a global gradient field from local patchwise computations (Lenseth et al., 2021). For a cover by patches $1_{t,l}$8 and a $1_{t,l}$9-local gradient algorithm $l$0, one seeks a global field $l$1. The two-stage merge, for $l$2 (ProcessLowerStars), proceeds via:

• Patchwise computation of $l$3 and all 1-overlap fields.
• Successive horizontal and then vertical merges along rows and columns, removing pairs from overlap regions that were erroneously created due to incomplete local neighborhoods.
• The general two-set merge theorem (Theorem 3.1, (Lenseth et al., 2021)) is

$l$4

This ensures no spurious pairs introduced by locality violations survive in the global field.

The parallel complexity is $l$5 wall-time for $l$6 cells across $l$7 patches, with strong-scaling up to $l$8.

6. Comparative Analysis: d2-StepMerge versus AnyOrder and Naive Approaches

In DLMs, d2-AnyOrder estimator is exact, requiring only one model pass, but is restricted to models supporting any-order decoding (i.e., satisfying independent-mask and order-causality). Most practical masked DLMs, such as LLaDA-8B-Instruct, do not meet these constraints; d2-AnyOrder yields highly erroneous importance weights in these settings (log-likelihoods diverge by an order of magnitude as per Table 1 in (Wang et al., 25 Sep 2025)). Conversely, d2-StepMerge is model-agnostic, providing the first high-fidelity RL signal for general masked DLMs. On any-order compatible models (e.g., Esoteric-LM), d2-AnyOrder remains superior under matched FLOPs (Table 4, (Wang et al., 25 Sep 2025)), but for standard models, only d2-StepMerge is applicable with high accuracy.

In discrete Morse theory, StepMerge formalizes the merging process, guaranteeing global correctness—something that naive union of local fields fails to achieve, as border cells within incomplete stars can incur spurious pairings (Lenseth et al., 2021). The d²–StepMerge protocol systematically subtracts these inconsistent contributions.

7. Applications and Significance

d2-StepMerge in DLMs is central to reinforcement learning for models not supporting any-order likelihood estimation, enabling high-accuracy reasoning on logical/mathematical tasks under manageable compute budgets. It provides a tunable bias-compute framework, with analytic control and demonstrated state-of-the-art empirical results (Wang et al., 25 Sep 2025).

In discrete Morse theory, d²–StepMerge enables scalable, parallelizable computation of discrete gradients, suitable for high-resolution images or large-scale complexes, preserving correctness in the assembled global field (Lenseth et al., 2021). The method is foundational for parallel persistent homology calculations via local-to-global assembly.

In both contexts, d2-StepMerge exemplifies principled merging of local information, balancing global accuracy with algorithmic efficiency.