Tree Reward-Aligned Search for TReASURe in Masked Diffusion Language Models

• URL: http://arxiv.org/abs/2509.23146v1
• Date: Sat, 27 Sep 2025 06:22:45 GMT
• Title: Tree Reward-Aligned Search for TReASURe in Masked Diffusion Language Models
• Authors: Zichao Yu, Ming Li, Wenyi Zhang, Weiguo Gao,
• Abstract summary: Tree search has emerged as a powerful framework for aligning generative models with task-specific rewards at test time.<n>We propose TReASURe, a tree-search test-time alignment method that addresses these issues.<n>TReASURe achieves state-of-the-art results on perplexity, linguistic acceptability, and control of sentiment and toxicity.
• Score: 13.433506313486701
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Tree search has recently emerged as a powerful framework for aligning generative models with task-specific rewards at test time. Applying tree search to Masked Diffusion Language Models, however, introduces two key challenges: (i) parallel unmasking yields highly correlated branches, limiting exploration, and (ii) reward evaluation via sampled completions produces high-variance estimates, making pruning unstable. We propose TReASURe, a tree-search test-time alignment method that addresses these issues. It introduces (i) UnmaskBranch, a branching strategy based on first-hitting unmasking that diversifies both token content and reveal order with a single model call per parent node, and (ii) ResubstituteScore, a pruning rule that uses deterministic resubstitution to score partially masked sequences with low-variance proxy completions. Theoretically, we quantify branching efficiency gains in NFEs (number of function evaluations), show that the scoring rule approximates the true reward with error bounded by predictive uncertainty, and prove improvements with larger tree widths. Empirically, TReASURe achieves state-of-the-art results on perplexity, linguistic acceptability, and control of sentiment and toxicity, outperforming prior methods under matched compute budgets, with especially strong gains in low-NFE regimes.

Related papers

• Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models [96.0074341403456]
  Inference-time compute has re-emerged as a practical way to improve LLM reasoning.<n>Most test-time scaling (TTS) algorithms rely on autoregressive decoding.<n>We propose Prism, an efficient TTS framework for dLLMs.
  arXiv  Detail & Related papers  (2026-02-02T09:14:51Z)
• Diffusion Language Model Inference with Monte Carlo Tree Search [22.7649405246503]
  Diffusion language models (DLMs) have emerged as a compelling alternative to autoregressive generation.<n>We introduce MEDAL, a principled search mechanism for DLMs inference.<n>Across multiple benchmarks, MEDAL achieves up to 22.0% improvement over existing inference strategies.
  arXiv  Detail & Related papers  (2025-12-13T04:30:02Z)
• Ensembling LLM-Induced Decision Trees for Explainable and Robust Error Detection [24.742137117129502]
  Error detection is important for ensuring data quality.<n>Recent state-of-the-art ED methods leverage the pre-trained knowledge and semantic capability embedded in large language models (LLMs) to directly label whether a cell is erroneous.<n>We propose an LLM-as-an-inducer framework that adopts LLM to induce the decision tree for ED (termed TreeED) and further ensembles multiple such trees for consensus detection (termed ForestED)<n>Our methods are accurate, explainable and robust, achieving an average F1-score improvement of 16.1% over the best baseline.
  arXiv  Detail & Related papers  (2025-12-08T07:40:48Z)
• Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models [51.12873073612084]
  Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance depends on the inference time order of unmasking.<n>We propose Lookahead Unmasking (LookUM), which addresses these concerns by reformulating sampling as path selection over all possible unmasking orders.<n>LookUM requires only two to three paths to achieve peak performance, demonstrating remarkably efficient path selection.
  arXiv  Detail & Related papers  (2025-11-04T02:37:37Z)
• Parallel Sampling from Masked Diffusion Models via Conditional Independence Testing [4.707859580472452]
  Masked diffusion models (MDMs) offer a compelling alternative to autoregressive models (ARMs) for discrete text generation.<n>They enable parallel token sampling, rather than sequential, left-to-right generation.<n>We present PUNT, a model-agnostic sampler that reconciles this trade-off.
  arXiv  Detail & Related papers  (2025-10-24T18:41:26Z)
• TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling [65.46347858249295]
  TreePO is a self-guided rollout algorithm that views sequence generation as a tree-structured searching process.<n>TreePO essentially reduces the per-update compute burden while preserving or enhancing exploration diversity.
  arXiv  Detail & Related papers  (2025-08-24T16:52:37Z)
• TreeRPO: Tree Relative Policy Optimization [65.51935468270916]
  name is a novel method that estimates the mathematical expectations of rewards at various reasoning steps using tree sampling.<n>Building on the group-relative reward training mechanism of GRPO, name innovatively computes rewards based on step-level groups generated during tree sampling.
  arXiv  Detail & Related papers  (2025-06-05T15:56:38Z)
• RASD: Retrieval-Augmented Speculative Decoding [5.3926068062773895]
  Speculative decoding accelerates inference in large language models (LLMs)<n>This paper proposes RASD (Retrieval-Augmented Speculative Decoding), which adopts retrieval methods to enhance model-based speculative decoding.
  arXiv  Detail & Related papers  (2025-03-05T12:10:14Z)
• Don't Get Lost in the Trees: Streamlining LLM Reasoning by Overcoming Tree Search Exploration Pitfalls [83.89771461061903]
  Recent advancements in tree search algorithms guided by verifiers have significantly enhanced the reasoning capabilities of large language models (LLMs)<n>Recent advancements in tree search algorithms guided by verifiers have significantly enhanced the reasoning capabilities of large language models (LLMs)<n>We identify two key challenges contributing to this inefficiency: $textitover-exploration$ due to redundant states with semantically equivalent content, and $textitunder-exploration$ caused by high variance in verifier scoring.<n>We propose FETCH, a flexible, plug-and-play system compatible with various tree search algorithms.
  arXiv  Detail & Related papers  (2025-02-16T16:12:01Z)
• Token-Supervised Value Models for Enhancing Mathematical Problem-Solving Capabilities of Large Language Models [56.32800938317095]
  Existing verifiers are sub-optimal for tree search techniques at test time.<n>We propose token-supervised value models (TVMs)<n>TVMs assign each token a probability that reflects the likelihood of reaching the correct final answer.
  arXiv  Detail & Related papers  (2024-07-12T13:16:50Z)
• Feature Attenuation of Defective Representation Can Resolve Incomplete Masking on Anomaly Detection [1.0358639819750703]
  In unsupervised anomaly detection (UAD) research, it is necessary to develop a computationally efficient and scalable solution. We revisit the reconstruction-by-inpainting approach and rethink to improve it by analyzing strengths and weaknesses. We propose Feature Attenuation of Defective Representation (FADeR) that only employs two layers which attenuates feature information of anomaly reconstruction.
  arXiv  Detail & Related papers  (2024-07-05T15:44:53Z)
• Recursive Speculative Decoding: Accelerating LLM Inference via Sampling Without Replacement [11.91629418177851]
  Speculative decoding is an inference-accel method for large language models. Recent works have advanced this method by establishing a draft-token tree. We present Recursive Speculative Decoding (RSD), a novel tree-based method that samples draft tokens without replacement.
  arXiv  Detail & Related papers  (2024-02-21T22:57:49Z)
• Lookback for Learning to Branch [77.32867454769936]
  Bipartite Graph Neural Networks (GNNs) have been shown to be an important component of deep learning based Mixed-Integer Linear Program (MILP) solvers. Recent works have demonstrated the effectiveness of such GNNs in replacing the branching (variable selection) in branch-and-bound (B&B) solvers.
  arXiv  Detail & Related papers  (2022-06-30T02:33:32Z)
• Complex Event Forecasting with Prediction Suffix Trees: Extended Technical Report [70.7321040534471]
  Complex Event Recognition (CER) systems have become popular in the past two decades due to their ability to "instantly" detect patterns on real-time streams of events. There is a lack of methods for forecasting when a pattern might occur before such an occurrence is actually detected by a CER engine. We present a formal framework that attempts to address the issue of Complex Event Forecasting.
  arXiv  Detail & Related papers  (2021-09-01T09:52:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.