Diffusion-GR2: Diffusion Generative Reasoning Re-ranker

Abstract: Generative reasoning re-rankers achieve strong recommendation accuracy by emitting a chain-of-thought before re-ordering a candidate list, but they are slow at inference: an autoregressive (AR) decoder spends one sequential forward pass per reasoning token, and the reasoning trace far exceeds the ranking it produces. To reduce this cost, block-diffusion LLMs decode many positions in parallel over a few denoising steps and are substantially faster, yet naively converting an AR re-ranker into one opens two accuracy gaps: (1) a structural gap: answer positions are denoised in parallel and scored independently, so the decoder emits invalid rankings (duplicated, dropped, or out-of-set identifiers) that AR avoids through left-to-right masking; and (2) a distributional gap: fine-tuning the converted model on fixed teacher trajectories is off-policy relative to its own decoding at inference, leaving a residual accuracy gap. To close both gaps while keeping the speedup, we propose \textbf{Diffusion-GR2}, a recipe that converts our AR reasoning re-ranker (GR2) into a block-diffusion re-ranker. First, conversion fine-tuning (CFT) adapts the AR-initialized diffusion model to denoise the answer into a valid permutation on its own, without an external constrained decoder. Next, on-policy distillation (OPD) then supervises the model on its own decoded trajectories with dense per-token targets from the AR teacher. Finally, we apply a reinforcement-learning (RL) stage against a re-ranking reward on top of OPD's on-policy policy. Experiments on Amazon Beauty demonstrate that Diffusion-GR2 recovers to near-parity with the AR re-ranker, while block-parallel decoding raises decode throughput by $2.4$--$3.5\times$ at the model's reasoning output length. Ablations show that CFT recovers most of the conversion gap, and that on-policy distillation further closes it to the AR reference.

• The paper introduces Diffusion-GR2, which converts AR generative re-ranking into a fast block-diffusion decoder through conversion fine-tuning, on-policy distillation, and reinforcement learning.
• It achieves near-parity Recall@1 (0.2951 vs 0.2960) while reducing latency by 2.4–3.5×, thereby making LLM-based re-ranking viable for high-frequency production settings.
• The method preserves chain-of-thought reasoning quality and structural accuracy, addressing challenges like candidate permutation errors inherent in naive diffusion models.

Diffusion-GR2: High-Accuracy, Low-Latency Generative Reasoning Re-Ranking via Block Diffusion

Motivation and Background

Generative reasoning re-rankers built on LLMs have emerged as a strong paradigm for high-accuracy re-ranking in industrial recommendation pipelines. These models generate explicit multi-step chain-of-thought (CoT) traces grounded in candidate semantics before emitting a permutation of the retrieved candidate set. Despite their empirical success, state-of-the-art models such as GR2 employ autoregressive (AR) decoding, incurring substantial latency and throughput costs—each reasoning token requires a sequential forward pass, and the CoT trace is typically an order of magnitude longer than the ranking produced. This severely limits deployability in high-frequency, production re-ranking scenarios.

Recent advances in block-diffusion LLMs offer significantly improved decoding speeds by resolving multiple output positions in parallel through iterative denoising. However, naively transplanting block-diffusion decoders into re-ranking models is detrimental to ranking accuracy due to two key factors: loss of autoregressive structural constraints (leading to malformed outputs) and train/inference distribution mismatch in sequence generation.

Figure 1: Conversion pipeline from the AR GR2 reasoning re-ranker into a block-diffusion model, closing accuracy gaps via staged fine-tuning, on-policy distillation, and RL to produce a fast yet accurate re-ranker.

The Diffusion-GR2 Architecture and Conversion Pipeline

The paper introduces Diffusion-GR2, a pipeline to convert an existing, high-performing AR reasoning re-ranker (GR2) into a block-diffusion decoder while minimizing accuracy sacrifice. The conversion proceeds in three primary stages:

1. Conversion Fine-Tuning (CFT): After directly initializing a block-diffusion model from the AR checkpoint (preserving weights), CFT is performed on the re-ranking task. This adapts the model to denoise the answer block into a valid permutation of candidates, overwhelmingly addressing the structural errors (e.g., duplicate, dropped, or out-of-set candidate identifiers) that are inherent to naive diffusion decoding.
2. On-Policy Distillation (OPD): Standard SFT or distillation would be off-policy—the student is trained on teacher rollouts, not samples generated under its own (block-diffusion) decoding. OPD closes this distributional gap by having the block-diffusion student generate its own outputs and then directly supervising these trajectories with dense per-token targets from the frozen AR teacher. This alignment over the student's own decode distribution is critical to matching on-distribution accuracy.
3. Reinforcement Learning (RL): An RL stage, starting from the OPD-initialized policy, optimizes a ranking-specific reward (lifting the ground-truth candidate's rank) directly, as in RLVR for reward-verifiable tasks. RL here is most effective applied atop OPD, as OPD delivers a non-degenerate, healthy student policy.

The outcome is that Diffusion-GR2 maintains the original re-ranking interface, preserves CoT reasoning, but amortizes and parallelizes decoding computation, yielding substantial inference speedups.

Inference Acceleration via Block-Diffusion

A primary advantage of Diffusion-GR2 lies in its block-diffusion decoder, which attends bidirectionally within a block and causally across blocks. This enables reuse of the prompt and committed context's key-value cache—after a one-time prefill, denoising is restricted to the current active block, substantially decreasing redundancy.

Figure 2: Diffusion-GR2’s inference traces the speed-accuracy path: AR emits one token per pass; Diffusion-GR2 denoises blocks in parallel with dynamic commitment based on confidence, amortizing prompt/KV cache and reducing sequential computation.

Crucially, the commitment rule (thresholded by maximum per-position logit confidence) determines at each denoising step which positions to commit, balancing accuracy and speed. The total number of forward passes is reduced from per-token to per-denoising-step times per block, with the performance controlled by the confidence threshold and block size. Unlike fully bidirectional diffusion models (e.g., LLaDA, Dream), which must re-encode the prompt at each diffusion step, the block-causal attention structure preserves AR-style prefill amortization—a decisive advantage for long-prompt, short-answer re-ranking use cases.

Empirical Results and Strong Claims

On Amazon Beauty under the TIGER protocol, Diffusion-GR2 demonstrates that structural and distributional gaps post-conversion can be closed to near-parity with AR GR2. Without CFT, virtually no valid rankings are produced, and re-ranking gains vanish. CFT restores nearly the full accuracy penalty, with OPD closing the residual off-policy gap and RL making up the final margin. The major claims of the paper are as follows:

• Accuracy: Diffusion-GR2 (Qwen3-8B) reaches Recall@1 within $0.001$ of the AR reference (0.2951 vs. 0.2960), indistinguishable on broader top-3 metrics (NDCG@3, Recall@3), and in some settings, slightly surpasses the teacher due to ranking reward optimization.
• Latency: Diffusion-GR2 achieves a $2.4$--$3.5\times$ speedup in tokens-per-second throughput at the reasoning output length over the AR decoder at comparable accuracy operating points. The decode speedup scales with output length, further favoring long reasoning traces.
• Reasoning Quality: Diffusion-GR2's CoT traces match AR GR2 on human/LMM-based evaluation axes: history grounding, logical flow, internal consistency, and identifier correctness, with no systematic degradation.

  Figure 3: Diffusion-GR2's decode speedup relative to AR increases with output length, demonstrating scaling advantages in scenarios with long reasoning spans.

Theoretical and Practical Implications

Diffusion-GR2 directly addresses a core barrier in deploying generative reasoning models for real-time ranking: the prohibitive compute and latency costs of generating long CoT traces per search/impression in commercial systems. By disambiguating and explicitly quantifying the conversion accuracy loss, and then closing the gap via CFT and OPD, Diffusion-GR2 establishes a practical template for transitioning high-signal generative re-rankers onto hardware-optimized, parallel decoding regimes.

Practically, this enables LLM-based, reasoning-aware re-ranking to run at production scale with tight latency budgets. The methodology sidesteps the need for cumbersome, task-specific constrained decoders in diffusion LMs, relying entirely on conversion from a valid AR base and stagewise distillation. The flexibility of OPD as a general distillation interface also suggests that future draft-verify speculative decoding approaches can be constructed, where the fast diffusion model generates drafts at aggressive speed, and an AR re-ranker serves as the final verifier and quality gate.

On the theoretical side, the work elucidates how structured generative tasks, such as permutation output, can be effectively realized via diffusion-based LMs despite their local independence assumptions, provided that ongoing on-policy fine-tuning and RL are incorporated.

Future Directions and Extensions

Key avenues for further development include integration with speculative decoding frameworks, wherein the diffusion draft's speed is maximized, and the AR verifier ensures final correctness. Dynamic block-sizing and pivot-aligned verification could yield even more efficient parallelization. Additionally, further research is needed to assess the impact of Diffusion-GR2 in broader recommendation domains, multi-modal tasks, and with larger candidate sets or multi-objective ranking constraints.

Conclusion

Diffusion-GR2 provides a robust, empirically-validated framework for converting slow but accurate AR generative re-rankers into high-throughput, block-diffusion decoders with negligible loss in ranking quality. By interleaving conversion fine-tuning, on-policy distillation, and RL, the model closes structural and distributional accuracy gaps, yielding models suitable for production deployment at previously unattainable latencies. The proposed pipeline is generalizable and provides a foundation for further innovation in scalable, reasoning-aware re-ranking with LLMs.