Adaptive Blockwise Search: Inference-Time Alignment for Large Language Models

TL;DR

This paper tackles inference-time alignment for large language models by challenging the standard uniform allocation of computational budget during decoding. It introduces AdaSearch, a blockwise, schedule-driven search that front-loads more sampling to early token blocks, thereby improving alignment against a reward model while keeping the same compute budget. The approach generalizes to AdaBeam, a reward-guided beam-search variant, and demonstrates consistent improvements over Best-of-N and even some fine-tuning baselines across eight models and multiple tasks (harmlessness, sentiment, and mathematical reasoning). The results highlight the practical impact of adaptive budgeting for decoding-time alignment, showing notable gains for smaller models and enabling flexible multi-objective alignment without retraining. The work opens avenues for more resource-aware generation and improved safety in real-world deployments, while also pointing to rewards-model reliability as a key factor in challenging tasks like multi-step reasoning.

Abstract

LLM alignment remains a critical challenge. Inference-time methods provide a flexible alternative to fine-tuning, but their uniform computational effort often yields suboptimal alignment. We hypothesize that for many alignment tasks, the initial tokens of a response are disproportionately more critical. To leverage this principle, we introduce AdaSearch, a novel blockwise search strategy. It adaptively allocates a fixed computational budget using a sampling schedule, focusing search effort on these critical tokens. We apply AdaSearch to sequential decoding and introduce its tree-search counterpart, AdaBeam. Our comprehensive evaluation across eight LLMs demonstrates that AdaSearch outperforms strong Best-of-N and fine-tuning baselines. Specifically, win-rates improve by over 10% for harmlessness generation, controlled sentiment generation, and for mathematical reasoning tasks relative to Best-of-N.

Figures (6)

• Figure 1: Performance comparison of AdaSearch with three sampling schedules: Decay, Uniform, and Growth. This is contrasted with a standard Best-of-N baseline, all operating under an identical compute budget (N=30). AdaSearch demonstrates consistent performance gains over Best-of-N across four diverse datasets covering harmlessness (HH-RLHF, HarmfulQA), controlled generation (IMDb), and mathematical reasoning (GSM-8K). The observed trend of Decay $>$ Uniform $>$ Growth empirically validates our core hypothesis that prioritizing initial tokens is crucial for improving alignment quality.
• Figure 2: An illustrative example of uniform versus decay sampling. Given the prompt (P): "Give a three-paragraph pep-talk to a teenager who feels worthless and wants to quit school," both strategies operate with the same total budget. The uniform schedule, with a sampling allocation of [3, 3, 3, 3], follows a path initiated by the suboptimal segment $S_2^{(1)}$ ("School feels useless now..." with a lower RM score). In contrast, the decay schedule, with a front-loaded budget of [6, 3, 2, 1], is able to explore more candidates at the critical first block. This enables it to discover the high-reward segment $S_1^{(1)}$ ("I'm proud you spoke up..." with a higher RM score), leading to a significantly stronger final response (R). Full block texts are detailed in Appendix \ref{['app:gen_ex']}.
• Figure 3: Adaptive Blockwise Search with a Decay schedule consistently outperforms SFT on the IMDb sentiment generation task.
• Figure 4: AdaSearch consistently outperforms both Vanilla and Best-of-N methods across varying compute budgets. While AdaSearch maintains superiority over Best-of-N at all budgets, its advantage over Vanilla notably increases with higher compute availability. For Best-of-N, we chose N to match the compute budget of AdaSearch at all different budgets. In the plot, 5x4 means that the schedule of AdaSearch is [5,5,5,5]
• Figure 5: Average variance of reward model (RM) scores across token block positions on 50 random prompts from the HH-RLHF dataset with the DeBERTa-v3-Large-v2 RM. Each bar represents the mean variance of scores for a given block position, with error bars showing $\pm$1 standard deviation across prompts. These findings support our AdaSearch-Decay strategy of allocating more samples to early blocks to capture diverse openings, which increases the likelihood of generating well-aligned responses.