Source-linked AI summary
Controlled Decoding from Language Models
Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, Jilin Chen, Alex Beutel, Ahmad Beirami
TL;DR
The paper addresses how to align language-model responses with rewards while retaining a frozen base model. It proposes controlled decoding, which trains a separate prefix scorer to steer tokenwise or blockwise generation under a KL-regularized RL objective. CD is reported to improve several rewards, support multi-objective and unseen-model control, and provide blockwise control with a practical connection to best-of-K.
Problem
The paper asks how to align generated content with rewards when the pretrained representations of a generative language model cannot be controlled.
Method
Controlled decoding trains a separate prefix scorer to learn a reward value function and uses it at inference time to control generation from a frozen base model.
Results
CD improves dialog length, helpfulness, harmlessness, and summarization quality while maintaining a small deviation from the base policy, and supports multi-objective and unseen-base-model control.
Takeaways & Limitations
CD provides configurable inference-time alignment through tokenwise or blockwise control without changing the base language model.
Takeaways & Limitations
The tokenwise RL formulation is more restrictive than the sequence-level RL used for RLHF and DPO, and evaluation uses a different sequence-level KL measure.
Abstract
from arXiv · showhide
KL-regularized reinforcement learning (RL) is a popular alignment framework to control the language model responses towards high reward outcomes. We pose a tokenwise RL objective and propose a modular solver for it, called controlled decoding (CD). CD exerts control through a separate prefix scorer module, which is trained to learn a value function for the reward. The prefix scorer is used at inference time to control the generation from a frozen base model, provably sampling from a solution to the RL objective. We empirically demonstrate that CD is effective as a control mechanism on popular benchmarks. We also show that prefix scorers for multiple rewards may be combined at inference time, effectively solving a multi-objective RL problem with no additional training. We show that the benefits of applying CD transfer to an unseen base model with no further tuning as well. Finally, we show that CD can be applied in a blockwise decoding fashion at inference-time, essentially bridging the gap between the popular best-of-K strategy and tokenwise control through reinforcement learning. This makes CD a promising approach for alignment of language models.
1. Introduction
The paper asks how to align generated content with rewards while keeping a pretrained language model's representations unchanged. It proposes controlled decoding as a modular inference-time approach that preserves a frozen base model while enabling configurable reward control.
- Inference-time add-on methods preserve a frozen base model and offer reward configurability, unlike generator-improvement methods that update model weights.
- Controlled decoding (CD) learns a reward prefix scorer that steers generation from a partially decoded path while solving a KL-regularized RL objective.
- CD-FUDGE and CD-Q are presented as variants that provably sample from a solution to the tokenwise RL objective.
- CD supports blockwise control, combining prefix-scored path selection with best-of-K to bridge sequence-level sampling and tokenwise reinforcement learning.
- The framework combines multiple rewards at inference time and transfers to an unseen base model without further training.
2. KL-Regularized Reinforcement Learning
The paper formulates alignment as tokenwise KL-regularized reinforcement learning, balancing reward improvement against divergence from a frozen base language model. Its theorem gives a unique optimal policy that can be implemented by combining base-model logits with a learned value function.
- Objective: The value function estimates the expected cumulative reward of a fully decoded response when generation continues from a partial sequence under the base model.Rewards are assigned only after decoding completes.
- Objective: The KL-regularized objective trades off expected reward advantage against tokenwise divergence from the frozen base model.The KL penalty discourages excessive drift because such drift is expected to degrade other top-line language-model metrics.
- Optimal policy: Theorem 2.1 states that the RL objective has a unique optimal policy, characterized by the paper’s tokenwise solution.At λ = 0, the optimum reduces to the reference policy, while larger λ values trace reward–KL tradeoffs.
- Optimal policy: The optimal policy enables inference-time alignment by combining logits from a frozen language model with those of a value function.This preserves the frozen base model while using the learned controller to alter token selection.
- Scope: The tokenwise formulation is more restrictive than the sequence-level objectives used in RLHF and DPO, although the paper compares their sequence-level reward–KL tradeoffs.The distinction limits what the tokenwise objective directly represents relative to sequence-level RL.
3. Controlled Decoding
Controlled decoding learns a prefix scorer for the reward value function and uses it with a frozen base model through tokenwise or blockwise inference procedures. The framework supports theoretically grounded control, and blockwise decoding can match best-of-K reward with substantially smaller K while reducing latency.
- Controlled decoding: CD trains a prefix scorer to approximate the optimal value function and then uses it at inference time to control a frozen base language model.The paper presents two training approaches and two inference-time uses.
- Prefix-scorer training: Under regularity assumptions, SGD for CD-FUDGE converges to a stationary point of the target objective when its training data comes from base-model rollouts.This provides a formal connection between FUDGE prefix scorers and the paper’s RL problem.
- Prefix-scorer training: CD-Q trains the prefix scorer with a Bellman-identity-based loss using expected next-prefix values under the reference policy, with a stop gradient on the target.The method is presented as an off-policy solver that does not require rolling out the base model during training.
- Inference-time control: Tokenwise CD combines base-model and prefix-scorer logits at every token, while blockwise CD samples K length-M continuations and accepts the highest-scoring block.Blockwise decoding repeats this process until a candidate containing EOS is accepted.
- Scope: Neither blockwise CD nor best-of-K is designed to optimally solve the sequence-level KL-regularized objective targeted by RLHF methods such as PPO and DPO.The paper motivates blockwise control partly through empirical reward–KL behavior of best-of-K.
- Inference-time control: Blockwise CD has best-of-K-like throughput but only M-token decoding latency, enabling lower-latency and potentially streaming applications for long sequences.Best-of-K must fully decode all K sequences before selecting one to serve.
- Controlled decoding: Blockwise CD achieves the same reward level as best-of-K with up to 10x smaller K.This combines reinforcement-learning-based scoring with blockwise candidate selection.
4. Experimental Setup
The experiments evaluate controlled decoding and baselines across response length, helpfulness and harmlessness, and summarization tasks. They compare reward outcomes against KL divergence using matched experimental procedures and several inference-time metrics.
- Experiments use PaLM 2-XXS as the base model and prefix-scorer model unless otherwise specified.
- Response length uses rlength([x, yT ]) = log(T/Tmax) with Tmax = 1024, while HH and summary rewards are trained from pairwise preferences using the Bradley-Terry model.
- Compared methods include CD-Q, CD-FUDGE, blockwise CD-FUDGE, KL-regularized PPO, online DPO, online IPO, and best-of-K.Best-of-K samples K responses from the base model, ranks them by reward, and selects the highest-ranking response.
- The evaluation focuses on KL values below 10 because larger values show significant signs of overfitting, and sequence-level evaluation favors PPO, DPO, and IPO.
- Performance is assessed with normalized expected reward, win-rate against the base policy, and reward-versus-KL tradeoff curves.KL divergence is swept through scorer strength, implicit regularizers, or K to compare aligned policies at matched divergence levels.
- The study evaluates response length on Reddit conversations, helpfulness and harmlessness on Anthropic HH, and summarization quality on TL;DR.Reddit conversations optimize response length, while Anthropic HH and TL;DR provide preference data for reward models.
5. Experimental Results
Across dialog length, HH, and summarization, blockwise CD-Q generally achieves the strongest reward–KL tradeoffs, often matching best-of-K with smaller K and supporting transfer and multi-objective control.
- Experiment 1: Increasing dialog response length: K=6 blockwise CD-Q achieves similar length and KL divergence to best-of-K with K=50, while matching best-of-K’s best length–KL tradeoff.Best-of-K also outperforms KL-regularized PPO on the reward–KL tradeoff in this experiment.
- Experiment 1: Increasing dialog response length: Tokenwise CD-FUDGE and CD-Q yield more favorable reward–KL tradeoffs than all baselines, including DPO and IPO.
- Experiment 1: Increasing dialog response length: Blockwise CD-Q is on par with best-of-K and outperforms blockwise CD-FUDGE because CD-FUDGE’s prefix-score predictions are much noisier.
- Experiment 2: Improving dialog helpfulness and harmlessness (HH): Blockwise CD-Q and CD-FUDGE substantially improve HH over baselines, but neither matches best-of-K; CD-Q and CD-FUDGE classification accuracy is approximately 0.6 versus Reward-XXS at approximately 0.7.
- Experiment 3: Improving summarization quality: For summarization, blockwise CD-Q outperforms IPO but does not match best-of-K, while combining HH and length scorers enables live adjustment of the objective trade-off.
- Transfer and blockwise-control studies: CD-Q transfers to PaLM 2-S and PaLM 2-XS without prefix-scorer retraining, matching the strongest baseline, while larger block sizes improve HH win-rate–KL tradeoffs.
6. Related Work
The paper situates controlled decoding among inference-time control, tree search, KL-regularized reinforcement learning, and supervised generator-improvement methods.
- Controlled decoding/generation: Inference-time prefix scorers control a frozen language model, preserving configurability compared with training-time generator interventions.Related methods include FUDGE and COLD, while DPO, PPO, SliC, and IPO update model weights and offer less reward configurability.
- Tree search: Controlled decoding is conceptually related to tree search, but uses a fixed search depth of one during decoding.Prior work applies Monte Carlo tree search or heuristic tree search to guide language-model generation.
- Reinforcement learning: KL-regularized reinforcement learning aligns language models while penalizing divergence from the language model, and related work applies PPO with language-model regularization.The paper also connects reward configurability to combining reward functions through reward soups.
- Supervised learning from negative examples: Supervised interventions such as unlikelihood training, contrastive losses, DPO, and IPO improve generator behavior during training rather than controlling a frozen model at inference time.These approaches are related alternatives to controlled decoding.
7. Concluding Remarks
The paper concludes that controlled decoding solves a KL-regularized alignment objective through inference-time prefix scoring and supports tokenwise, blockwise, multi-objective, and transferred control.
- 7. Concluding Remarks: Controlled decoding improves dialog length, helpfulness and harmlessness, and summarization quality while remaining close to the base language-model policy.The framework supports tokenwise and blockwise control and extends to multiple rewards and an unseen base model without retraining.
- 7. Concluding Remarks: Blockwise CD and best-of-K consistently achieve better practical reward-KL tradeoffs than tokenwise CD and KL-regularized RL.This empirical pattern is also reported by prior studies and supported by theoretical analysis of best-of-K.
- 7. Concluding Remarks: A comprehensive understanding of throughput, latency, and performance tradeoffs remains future work, potentially involving speculative decoding.The paper explored these tradeoffs only in a narrow set of experiments.
Impact Statement
The paper presents inference-time alignment as more configurable than training-time intervention, while warning that alignment methods can behave inconsistently on safety and other socially consequential issues.
- Impact Statement: Inference-time control offers finer-grained and more flexible alignment, potentially supporting configurable and personalizable behavior.The paper contrasts this with commonly used training-time interventions for KL-regularized reinforcement learning.
- Impact Statement: Alignment methods showed inconsistent behavior on safety and other socially consequential issues, requiring extreme caution in nuanced applications.The statement limits the broader implications of the proposed alignment techniques.
A. Additional details on experimental setup
The appendix provides additional reward-model and evaluation-prompt details for helpfulness, harmlessness, and summarization experiments.
- A. Additional details on experimental setup: The appendix includes additional experimental-setup material for the reported evaluations.The section heading identifies these passages as supplementary setup details.
- Reward Model training setup: The appendix describes reward-model training for helpfulness, harmlessness, and summarization using pairwise human-preference objectives.Helpfulness and harmlessness use a combined dataset, while summarization uses the TL;DR preference dataset and one training epoch at learning rate 1e-5.
- Evaluation prompts: Evaluation prompts ask PaLM 2-L to rank two assistant outputs according to helpfulness, harmlessness, or summarization quality.The prompts present the dialogue or post together with two candidate responses and request a majority-human-style ranking.
- Evaluation prompts: The appendix specifies the structured ranking format, requiring a response of 1 or 2 to indicate the better output.The same output convention is used for the shown assistant-comparison prompts.
B. Additional experimental results
Additional experiments examine how well CD-Q and CD-FUDGE prefix scorers predict response length and how blockwise CD-Q interacts with DPO.
- CD-Q predictions align much better with actual Reddit response lengths than CD-FUDGE predictions, which are noisy, especially for pairwise comparisons.
- Blockwise CD-Q is evaluated against DPO across checkpoints with different KL divergence values, including when applied to DPO without retraining.
- Combining blockwise CD-Q with DPO reaches a similar win rate at a lower apparent KL divergence, with win-rate points around 0.7 shown in the comparison.
C. Proofs
The proofs establish uniqueness of the optimal policy and analyze convergence of the prefix-scorer training procedure under stated regularity assumptions.
- The objective is strongly convex in π, yielding a unique maximizer for the optimal-policy characterization.
- The gradient estimator used for training is unbiased for the gradient of the optimal objective, supporting the convergence analysis of CD-Q and CD-FUDGE.
- Under Lipschitzness, a non-empty solution set, the PL inequality, and bounded expected gradient norms, SGD on ℓF converges to θ⋆.
- The proof of the convergence theorem follows from the unbiased-gradient lemma and an existing SGD convergence result.