Diffusion-GR2: Block Diffusion Re-Ranking

• Diffusion-GR2 is a re-ranking framework that converts an autoregressive GR2 into a block-diffusion model, reducing latency while maintaining high ranking accuracy.
• It employs weight-preserving conversion with conversion fine-tuning, on-policy distillation, and reinforcement learning to bridge structural and distributional gaps.
• Experimental results on Amazon Beauty show near-parity with GR2 and throughput speedups up to 3.5×, underscoring its practical viability for recommendation re-ranking.

Searching arXiv for the specified paper and closely related context. Diffusion-GR2 is a conversion recipe that turns an autoregressive generative reasoning re-ranker, GR2, into a block-diffusion re-ranker while preserving nearly all of the autoregressive model’s ranking quality and reducing sequential decoding cost. In the reported formulation, the task is recommendation re-ranking with chain-of-thought reasoning: the model receives a user history and a pre-ranked candidate list, emits a free-form reasoning trace, and then outputs a ranked list that must be a permutation of the candidate identifiers. The central claim is that block-parallel denoising can substantially accelerate inference, but naive conversion from autoregressive decoding creates a structural gap in permutation validity and a distributional gap from off-policy training; Diffusion-GR2 addresses these gaps through conversion fine-tuning, on-policy distillation, and reinforcement learning (Zhang et al., 1 Jul 2026).

1. Problem formulation and latency bottleneck

Diffusion-GR2 inherits the problem setting of GR2. For each impression, the input consists of a user history $H = (s_{v_1}, \dots, s_{v_k})$ and a pre-ranked candidate list $D = (s_{y_1}, \dots, s_{y_N})$, with $N=10$ fixed in the experiments. Each item is represented by a semantic identifier derived from an RQ-VAE tokenizer. The re-ranker is a conditional model

$\pi_\theta(o \mid P(H, D)),$

where the output $o$ decomposes into a chain-of-thought reasoning trace $\tau = (r_1, \dots, r_M)$ and an answer sequence $a = (a_1, \dots, a_N)$, where the answer must be a permutation of the candidate identifiers in $D$. The prompt $P(H,D)$ is a structured chat template containing system instructions, SID-grounded history, candidate metadata, and an output schema in JSON form (Zhang et al., 1 Jul 2026).

The computational motivation is straightforward. In GR2, reasoning traces are long—approximately 130 tokens in the Beauty experiments—while the answer segment is short. Because the original model is autoregressive, it performs one sequential forward pass per generated token. In the stated production-like regime of long prompts, approximately $2.2$k tokens, plus chain-of-thought outputs and many impressions per second, this token-by-token decoding dominates latency.

Block-diffusion LLMs alter this cost profile by denoising many positions in parallel over a small number of iterative refinement steps. This reduces the dependence of sequential passes on output length. However, the paper identifies two conversion-specific failure modes. The first is a structural gap: answer positions are denoised in parallel and scored independently, so the decoder may emit duplicated items, missing items, or out-of-set identifiers. The second is a distributional gap: if the converted model is fine-tuned only on fixed teacher trajectories, training is off-policy relative to the student’s own diffusion-time states at inference, leaving a residual exposure-bias-like error (Zhang et al., 1 Jul 2026).

2. Diffusion formulation and block-causal architecture

The starting point is the autoregressive GR2 teacher, built on a Qwen3-8B decoder-only LLM trained with supervised fine-tuning on high-quality chain-of-thought traces and permutations generated by a larger LLM, followed by reinforcement learning against a re-ranking reward. Diffusion-GR2 preserves all transformer weights and changes only the decoding configuration: block size, masking behavior, and the special mask token. The conversion is therefore explicitly weight-preserving rather than architecturally redesigned (Zhang et al., 1 Jul 2026).

The diffusion model is a masked-diffusion LLM in which the assistant output span—reasoning plus answer—is treated as a discrete sequence containing [MASK] tokens. At denoising step $D = (s_{y_1}, \dots, s_{y_N})$0, the model operates on a partially filled sequence $D = (s_{y_1}, \dots, s_{y_N})$1 and predicts a token distribution at masked positions,

$D = (s_{y_1}, \dots, s_{y_N})$2

A commitment rule $D = (s_{y_1}, \dots, s_{y_N})$3, controlled by a confidence threshold $D = (s_{y_1}, \dots, s_{y_N})$4, fixes high-confidence positions and remasks the rest: $D = (s_{y_1}, \dots, s_{y_N})$5 Positions whose maximum token probability exceeds $D = (s_{y_1}, \dots, s_{y_N})$6 are committed to their arg-max token, while at least one position is always forced to commit so that decoding progresses. The formulation is discrete masked diffusion rather than Gaussian diffusion; there is no continuous noise schedule of the $D = (s_{y_1}, \dots, s_{y_N})$7 type.

A central systems choice is the block-causal attention pattern. The assistant output is partitioned into contiguous blocks of fixed size $D = (s_{y_1}, \dots, s_{y_N})$8, with $D = (s_{y_1}, \dots, s_{y_N})$9 in the experiments. Attention is bidirectional within a block and causal across blocks: a block can attend to the prompt and all earlier blocks, but not later ones. Decoding proceeds block by block. The prompt is prefixed once and its KV states are cached; the active block is then denoised over a few steps until all positions in that block are committed; decoding then advances to the next block. Setting $N=10$0 recovers ordinary autoregressive decoding with the same weights, which gives a direct baseline for throughput comparisons (Zhang et al., 1 Jul 2026).

This block structure matters because it retains KV caching for the long prompt while still allowing many positions to be resolved per step. The reported argument is that fully bidirectional masked-diffusion LMs would repeatedly re-encode the approximately $N=10$1k-token prompt, which would negate the speed benefit in the target operating regime. Diffusion-GR2 instead amortizes prompt prefill in an AR-like manner while parallelizing generation inside each block.

3. Conversion pipeline: CFT, OPD, and RL

Diffusion-GR2 consists of three sequential training stages applied after the weight-preserving AR-to-diffusion conversion: Conversion Fine-Tuning (CFT), On-Policy Distillation (OPD), and Reinforcement Learning (RL). Each stage targets a distinct failure mode left by the previous one (Zhang et al., 1 Jul 2026).

Conversion Fine-Tuning is the first and most important stage for closing the structural gap. It reuses the original GR2 supervised data: full prompts with history, candidates, metadata, and expert system role, paired with target reasoning traces and rankings. During CFT, the assistant response is treated as the span to be denoised via masked diffusion, and the model is trained with a masked language-model objective over randomly masked assistant tokens, with separate weighting for reasoning and answer tokens. The stated goal is for the AR-initialized diffusion model to learn to denoise the answer span into a valid permutation of the candidate set $N=10$2 without any external constrained decoder at inference.

The mechanism proposed for validity is entirely learned rather than hard-coded. The AR teacher never emits invalid permutations because left-to-right masking suppresses already-used candidate identifiers. Copying those weights provides a strong prior over JSON structure, candidate ordering, and permutation-like behavior. CFT then adapts this prior to the new masking and attention regime, always training on valid permutations. The paper explicitly states that it does not introduce special combinatorial machinery such as Sinkhorn layers or stepwise candidate constraints. Instead, validity emerges from AR initialization, masked-diffusion training on valid sequences, and the structured prompt and output format (Zhang et al., 1 Jul 2026).

On-Policy Distillation addresses the distributional gap that remains after CFT. Rather than training on teacher-forced trajectories from the autoregressive model, OPD trains the diffusion student on its own decoded rollouts. For each sampled prompt, the diffusion model generates outputs using the same block-diffusion procedure used at inference. For each committed token position, the frozen AR teacher is then queried on the same prompt and prefix to produce a next-token distribution, and the student is optimized to match that distribution via forward KL: $N=10$3 The important distinctions are that the KL is forward rather than reverse, and that the state distribution is on-policy because trajectories are sampled from the student itself. In the reported implementation, $N=10$4 trajectories are generated per prompt for OPD, and teacher logits are cached per trajectory (Zhang et al., 1 Jul 2026).

Reinforcement Learning is applied last, starting from the OPD checkpoint rather than from CFT directly. The rationale given is that a strong on-policy imitation checkpoint stabilizes RL because its decoded trajectories are already non-degenerate and reward-bearing. The RL stage uses a GRPO/DAPO-style objective to optimize re-ranking behavior directly.

4. Reward design and optimization objective

The reinforcement-learning stage treats the diffusion re-ranker as the policy and a full decoded assistant response as the action, although reward depends only on the answer permutation together with a format check. The primary scalar is the rank-promotion reward $N=10$5. Let $N=10$6 denote the rank of the true next item in the original pre-ranked list $N=10$7, let $N=10$8 denote its rank in the model’s re-ranked output, and let $N=10$9 be the number of candidates. Then

$\pi_\theta(o \mid P(H, D)),$0

This reward is positive when the model moves the target upward, negative when it pushes it downward, and is described in the source as a simple normalized improvement in rank closely related to Recall@K and NDCG when there is a single relevant item (Zhang et al., 1 Jul 2026).

A second term, the conditional format reward $\pi_\theta(o \mid P(H, D)),$1, is granted only when re-ranking strictly improves the target’s rank or when the target was already rank-1 and remains rank-1. The total reward is

$\pi_\theta(o \mid P(H, D)),$2

with $\pi_\theta(o \mid P(H, D)),$3. Invalid or overflowing permutations receive reward $\pi_\theta(o \mid P(H, D)),$4 but are retained in the training group so that the policy can learn to avoid them.

Optimization uses group-based PPO-style learning. For each prompt, the old policy samples a group of $\pi_\theta(o \mid P(H, D)),$5 trajectories, each receives a scalar reward, and the reward is converted into group-normalized advantages that are constant across all tokens in the trajectory. Importance ratios are computed by replaying the denoising trace, reconstructing the masked canvas at the moment each token was committed, and rescoring it under the current policy. The objective is

$\pi_\theta(o \mid P(H, D)),$6

with $\pi_\theta(o \mid P(H, D)),$7. The trace-replay mechanism is diffusion-specific and ensures that ratios are computed under the actual masking and attention patterns used at inference rather than under a teacher-forced proxy (Zhang et al., 1 Jul 2026).

5. Empirical evaluation

The reported experiments use Amazon Review Beauty under the TIGER protocol with 5-core filtering, chronological ordering, and leave-one-out train/validation/test splitting. The dataset statistics are 22,363 users, 12,101 items, and average sequence length 8.87. For each test user, a retriever provides a top-10 candidate set, which defines the pre-rank floor. The test set contains 1,615 users. Each item is represented by a 4-token semantic ID, and the re-ranker receives history, candidate SIDs, and metadata before reordering the 10 candidates. The evaluation metrics are Recall@1, Recall@3, and NDCG@3 (Zhang et al., 1 Jul 2026).

The systems compared are the retriever pre-rank floor, AR GR2, naive diffusion conversion without CFT, Diffusion-GR2 with CFT, Diffusion-GR2 with CFT plus OPD, and Diffusion-GR2 with CFT plus OPD followed by RL.

Method	Recall@1	Recall@3
Pre-rank (retriever)	0.2811	0.5591
AR GR2 (ref)	0.2960	0.5651
Diffusion naive	0.2811	0.5591
+ CFT	0.2930	0.5651
+ CFT + OPD	0.2944	0.5658
+ CFT + OPD→RL	0.2951	0.5671

The corresponding NDCG@3 values are 0.4401 for the pre-rank floor, 0.4497 for AR GR2, 0.4401 for naive diffusion, 0.4497 after CFT, 0.4497 after CFT plus OPD, and 0.4517 after OPD followed by RL. These numbers support three conclusions. First, naive diffusion conversion collapses to the retriever floor because malformed outputs cannot be parsed. Second, CFT recovers most of the gap to the AR reference. Third, OPD narrows the remaining difference and RL closes it further, reaching near-parity on Recall@1 and slightly exceeding the AR teacher on Recall@3 and NDCG@3 (Zhang et al., 1 Jul 2026).

The paper also reports explicit evidence about structural validity. Before CFT, naive diffusion decoding yields a valid JSON rate of 0.001, or approximately $\pi_\theta(o \mid P(H, D)),$8, and Recall@1 remains at the retriever floor of 0.2811. After CFT, valid permutations become the norm, and Recall@1 rises to 0.2930 versus 0.2960 for AR GR2. Although a post-CFT invalid-rate table is not provided, the paper describes outputs after CFT as essentially always valid permutations.

Throughput is measured at the model’s reasoning output length of approximately 130 tokens on a single H100-80G using torch.compile. The reported tradeoff is governed by the confidence threshold $\pi_\theta(o \mid P(H, D)),$9.

Decoding	Recall@1	Throughput (tok/s)
AR GR2	0.2960	71
Diffusion ($o$0)	0.2950	172
Diffusion ($o$1)	0.2942	234
Diffusion ($o$2)	0.2936	246

These correspond to speedups of $o$3, $o$4, and $o$5 over the AR baseline. At $o$6, the valid-JSON rate is 1.0 and Recall@1 is approximately equal to GR2 while throughput improves by $o$7. Lowering $o$8 increases parallelism and speed, with only modest degradation in Recall@1 until parsing starts to degrade below $o$9. The paper further states that speedup grows with output length and is almost flat in input length, which is consistent with the claim that gains derive primarily from block-parallel decoding rather than from prompt prefill (Zhang et al., 1 Jul 2026).

Reasoning quality is evaluated in a blind LLM-as-judge comparison over 50 AR versus diffusion pairs. Mean scores are 4.50 versus 4.34 for history grounding, 4.94 versus 4.94 for internal consistency, 4.54 versus 4.40 for logical flow, and 1.00 versus 1.00 for identifier correctness. Pairwise preference counts are AR 17, Diffusion-GR2 9, tie 24. The reported interpretation is that no systematic degradation of reasoning coherence or correctness is observed.

6. Interpretation, limitations, and prospective extensions

The empirical pattern suggests a division of labor across the three training stages. CFT is the dominant mechanism for structural alignment because it repairs permutation validity under the new decoding regime. OPD addresses the student’s exposure to its own partially denoised states. RL then operates in a regime where structural output quality is already reliable and can therefore optimize ranking metrics directly. This suggests that the three-stage procedure is not redundant but sequentially targeted to different pathologies of AR-to-diffusion conversion (Zhang et al., 1 Jul 2026).

The paper emphasizes several conditions that make block diffusion effective in this setting. The reasoning span is long and latency-dominating, whereas the answer span is short and structurally constrained. The block-causal design preserves KV caching for long prompts while enabling parallel denoising inside each block. AR initialization transfers JSON formatting, SID grounding, and ranking logic into the diffusion model. Finally, the same diffusion model must generate both free-form chain-of-thought and a tightly constrained permutation over a small candidate set; the results indicate that the answer-span constraint can be learned without explicit combinatorial decoding machinery.

Several limitations are explicitly acknowledged. Diffusion-GR2 depends heavily on a strong AR teacher that has already been trained with SFT and RL. The conversion procedure is operationally complex, involving CFT, OPD, and RL, each with its own hyperparameters and GPU cost. The experiments are restricted to a single dataset, Amazon Beauty, with $\tau = (r_1, \dots, r_M)$0 candidates, so the behavior for larger candidate sets or more heterogeneous domains remains open. Even with $\tau = (r_1, \dots, r_M)$1 speedup, reasoning-heavy re-ranking remains expensive in absolute terms. The method also provides no explicit combinatorial guarantee: permutation validity is learned rather than enforced (Zhang et al., 1 Jul 2026).

The reported future-facing ideas remain tentative. The paper explicitly points to speculative decoding with a diffusion draft and an AR verifier as a natural next step. It also proposes adaptive block sizing and adaptive confidence thresholds $\tau = (r_1, \dots, r_M)$2, where easier regions could be decoded with larger blocks or lower thresholds and harder regions with smaller blocks or stricter thresholds. A plausible implication is that Diffusion-GR2 functions as a general recipe for AR-to-diffusion conversion in structured reasoning tasks, but the source text presents this as a suggestion rather than as an experimentally established claim.

In summary, Diffusion-GR2 is a block-diffusion re-ranking framework that preserves the interface and much of the behavior of an autoregressive reasoning re-ranker while addressing two conversion-specific problems: invalid parallel decoding of structured answers and off-policy mismatch between training and inference. Its defining contribution is methodological rather than architectural: weight-preserving AR conversion, CFT for structural validity, OPD for on-policy alignment, and RL for direct metric optimization, together yielding near-parity with GR2 and $\tau = (r_1, \dots, r_M)$3–$\tau = (r_1, \dots, r_M)$4 decode-throughput gains at realistic reasoning lengths (Zhang et al., 1 Jul 2026).