PaperlayerView exact revision on arXiv

Source-linked AI summary

ARGS: Alignment as Reward-Guided Search

Maxim Khanov, Jirayu Burapacheep, Yixuan Li

arXiv:2402.01694v1cs.CLcs.AIcs.LG

TL;DR

LLMs can generate misinformation and harmful outputs, while PPO-based RLHF can require unstable, costly training and retraining when reward models change. ARGS incorporates reward signals into decoding, and the paper reports improved rewards across tasks and models, including a 19.56% average-reward gain over standard decoding.

  • Problem

    LLMs trained on varied data can generate misinformation and harmful outputs, motivating alignment with human objectives and safety.

    Source: §1 · p. 1

  • Method

    ARGS guides decoding by adjusting token predictions with reward signals, allowing alignment changes without PPO retraining.

    Source: §1 · p. 1

  • Results

    ↑19.56% average reward over standard decoding, with 64.33% preference or tie in GPT-4 evaluation; improvements also span model architectures, sizes, and alignment tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Decoding-time alignment offers a way to customize model outputs to changing reward signals and user requirements without extensive retraining.

    Source: §1 · p. 2

    Current evaluations use HH-RLHF and SHP benchmarks; the authors identify more intricate tasks, such as multi-step reasoning, as future work.

    Source: E LIMITATIONS AND FUTURE WORK · p. 17

Abstract

from arXiv · show

Aligning large language models with human objectives is paramount, yet common approaches including RLHF suffer from unstable and resource-intensive training. In response to this challenge, we introduce ARGS, Alignment as Reward-Guided Search, a novel framework that integrates alignment into the decoding process, eliminating the need for expensive RL training. By adjusting the model's probabilistic predictions using a reward signal, ARGS generates texts with semantic diversity while being aligned with human preferences, offering a promising and flexible solution for aligning language models. Notably, ARGS demonstrates consistent enhancements in average reward compared to baselines across diverse alignment tasks and various model dimensions. For example, under the same greedy-based decoding strategy, our method improves the average reward by 19.56% relative to the baseline and secures a preference or tie score of 64.33% in GPT-4 evaluation. We believe that our framework, emphasizing decoding-time alignment, paves the way for more responsive language models in the future. Code is publicly available at: \url{https://github.com/deeplearning-wisc/args}.

1 INTRODUCTION

Large language models can generate misinformation and harmful outputs, while PPO-based RLHF entails unstable, costly training and repeated retraining when reward models change. ARGS instead guides decoding with reward signals, and the paper reports higher rewards and preference alignment across tasks and model dimensions.

1 INTRODUCTION

  • LLMs’ varied training data can lead to misinformation and harmful outputs, making alignment with human objectives and safety a central challenge.Source: §1 · p. 1
  • PPO-based RLHF is widely adopted, but its training can be unstable and costly, and changing the reward model requires repeating training.Source: §1 · p. 1
  • ARGS adjusts token predictions using reward signals during decoding, avoiding the need for PPO retraining to realign outputs with changing preferences.Its reward weight trades off contextual relevance and preference alignment, with zero weight recovering standard maximum-likelihood decoding.Source: §1 · pp. 1–2
  • ↑19.56% average reward over standard decoding, with 64.33% preference or tie in GPT-4 evaluation.The paper also reports consistent improvements across model architectures, sizes, and alignment tasks.Source: §1 · p. 2
  • The paper frames ARGS as decoding-time alignment that supports flexible customization without expensive RL training and evaluates it qualitatively and quantitatively.It presents post-training adjustments as a complement to traditional training-phase alignment.Source: §1 · p. 2

2 ARGS: ALIGNMENT AS REWARD-GUIDED SEARCH

ARGS aligns generation at decoding time by combining language-model predictions with reward scores, then selecting or sampling among likely next tokens. Its search considers the top-k candidates, updates context iteratively, and adds a factor of k to decoding complexity.

2.1 REWARD-GUIDED SCORING

  • ARGS adds a weighted reward-model score to each candidate token’s language-model score, trading off contextual coherence and relevance against alignment criteria.The reward model assigns a scalar to a prompt-response pair; the coefficient w controls the reward’s weight.Source: §2.1 · p. 3

2.2 TOKEN SELECTION

  • ARGS-greedy selects the highest-scoring token from the k most likely predictions, while ARGS-stochastic samples from a renormalized distribution over those candidates.As temperature τ approaches zero, stochastic selection approaches greedy selection; higher τ makes sampling more uniform.Source: §2.2 · p. 3
  • After selection, ARGS appends the token to the prior context and repeats decoding until the desired sequence length is reached.The greedy procedure scores each top-k candidate, selects the maximum, and updates the context at each step.Source: §2.2 · pp. 3–4

2.3 IMPLEMENTATION AND COMPLEXITY

  • ARGS has aggregate time complexity O(k · m2), adding a factor of k over classical O(m2) decoding because it evaluates top-k candidates.Its initial step costs O(k · n2), while subsequent steps cost O(k · t) when previously calculated attention is reused.Source: §2.3 · p. 4

3 EXPERIMENTS

The experiments evaluate ARGS for helpful and harmless response generation using open-source models and datasets. They report higher reward and diversity than standard decoding, GPT-4 comparisons, and further analyses of parameter and decoding choices.

3.1 SETUP

  • The setup uses HH-RLHF crowd preferences, a LLaMA-7B model fine-tuned on preferred responses, and a reward model with 74.58% validation accuracy.Evaluation uses test prompts; decoding settings include a maximum 2,048-token prompt and 128-token continuation.Source: §3.1 · p. 5

3.2 RESULTS

  • ARGS improves average reward by 19.56% over greedy decoding, generally improves diversity, and maintains comparable contextual consistency at mild weights.Contextual consistency degrades when the weight becomes too large.Source: §3.2 · p. 5, Figure 2 · p. 6

3.3 GPT-4 EVALUATION

  • GPT-4 evaluation compares ARGS with alternative decoding methods using win-tie percentages, with higher percentages indicating more favorable responses.The evaluation scores helpfulness, harmlessness, relevance, accuracy, and insightfulness; Table 2 reports the comparison.Source: §3.3 · p. 6, Table 2 · p. 7

3.4 FURTHER ANALYSIS

  • ARGS-greedy generally improves average reward and human preference more effectively, while ARGS-stochastic can produce more diverse text through probabilistic token selection.Table 3 compares the two variants using the same default k and w; stochastic decoding uses temperature 0.7.Source: §3.4 · p. 7, Table 3 · p. 7
  • On another evaluation, ARGS consistently outperforms a greedy baseline, and its GPT-4 comparison with PPO achieves a 72.33% win-tie rate.The supplied passage also describes ARGS as model- and task-agnostic.Source: §3.4 · pp. 7–8

4 DISCUSSION

ARGS moves alignment to decoding time, supporting rapid realignment without retraining. On SHP, it achieves reward comparable to PPO while outperforming PPO and DPO on diversity and coherence, with an inference-speed tradeoff.

4 DISCUSSION

  • Figure 3 compares ARGS with a baseline on OPT models trained on SHP, with reward models used to guide ARGS on the x-axis.The passage describing the figure says a smaller reward model such as OPT-125m does not significantly change performance.Source: Figure 3 · p. 8, §4 · p. 9
  • Decoding-time alignment supports quick realignment when reward models change, avoiding exhaustive RL-model retraining and enabling customization to evolving needs.The framework is described as compatible with a wide range of models.Source: §4 · p. 8
  • ARGS achieves a comparable average reward as PPO while significantly outperforming PPO and DPO in diversity and coherence on SHP.The comparison uses the same reward and language models for ARGS as for PPO, but ARGS requires no further training.Source: §4 · p. 8
  • A 1.9× generation-time increase with OPT-2.7b and k = 10 accompanies a 6.8% reward improvement over conventional greedy decoding.The added complexity is a constant factor of k; the inference gap can be reduced with a smaller reward model or parallel reward computation.Source: §4 · pp. 8–9

5 RELATED WORKS

Related work spans training-time preference alignment and decoding strategies. ARGS is distinguished by applying reward guidance at the token level during decoding to address alignment without expensive RL training.

5 RELATED WORKS

  • RLHF uses external reward models to refine language models through reinforcement learning, including PPO; its instability and resource demands have also motivated supervised fine-tuning alternatives.The discussion cites Christiano et al. (2017), Ziegler et al. (2019), Stiennon et al. (2020), Lee et al. (2021), Nakano et al. (2022), and Snell et al. (2023) for RLHF, and Askell et al. (2021), Ouyang et al. (2022), Bai et al. (2022), and Glaese et al. (2022) for PPO.Source: §5 · p. 9
  • Deterministic decoding selects continuations by probability or scoring criteria, while top-k and nucleus sampling introduce stochasticity that can diverge from or contradict the human-written prefix.The cited examples include greedy beam search, contrastive search, top-k sampling, and nucleus sampling.Source: §5 · p. 9
  • ARGS uses token-level reward guidance for alignment, unlike step-level verifiers; its pairwise-ranking reward model contrasts with Deng & Raffel (2023)'s unidirectional model trained with cumulative squared error.Earlier token-level guided decoding work had not connected decoding directly to the alignment problem addressed here, according to the passage.Source: §5 · p. 9

6 CONCLUSION

The conclusion presents ARGS as a decoding-time alignment framework that reduces the need for resource-intensive RL training. The paper also describes potential access and applicability benefits, while directing readers to Appendix E for limitations and future work.

6 CONCLUSION

  • ARGS frames alignment as a decoding-stage problem and leverages reward signals to reduce the need for resource-intensive RL training.The conclusion characterizes its performance gains as consistent and describes the approach as potentially supporting more flexibly aligned language models.Source: §6 · p. 9
  • The conclusion places discussion of limitations and future work in Appendix E.Source: §6 · p. 9
  • Post-training adaptation may let smaller institutions or businesses tailor pretrained models without large-scale training, and compatibility with different reward models may extend use across domains.The ethics statement presents these as potential benefits, particularly where resources are limited or data change rapidly.Source: §6 · p. 10

A EXPERIMENTATION DETAILS

The experimentation details describe the computing environment and training configurations for LLaMA-7B, OPT-family models, PPO, and DPO.

A EXPERIMENTATION DETAILS

  • Experiments used RTX A6000 and A100 GPUs, Ubuntu 22.04.2, CUDA 11.8, cuDNN 8.9, Python 3.11.4, and PyTorch 1.12.1.Source: A EXPERIMENTATION DETAILS · p. 15
  • LLaMA-7B was trained on the HH-RLHF training split using LMFlow, AdamW, and DeepSpeed ZeRO stage 3, with hyperparameters summarized in Table 5.The training scheme follows Dong et al. (2023).Source: A EXPERIMENTATION DETAILS · p. 15, Table 5 · p. 15
  • OPT-family models used DeepSpeed-Chat and the Ouyang et al. (2022) scheme, with supervised fine-tuning on 20% and reward modeling on the subsequent 40% of SHP training data.The configurations were applied across OPT-125m, OPT-350m, OPT-1.3b, and OPT-2.7b.Source: A EXPERIMENTATION DETAILS · p. 15
  • PPO training followed DeepSpeed-Chat default configurations, while DPO experiments used Huggingface TRL's DPOTrainer module.The configurations are detailed in Tables 7 and 8, respectively.Source: A EXPERIMENTATION DETAILS · p. 15

B GPT-4 EVALUATION DETAILS

This section describes the prompt structure used for GPT-4 evaluation and points to supplementary training-hyperparameter tables.

B GPT-4 EVALUATION DETAILS

  • Each GPT-4 request combines a system prompt describing the proxy and task with a user prompt containing a question and two responses.The responses are from the two methods being compared.Source: B GPT-4 EVALUATION DETAILS · p. 15
  • Table 6 summarizes supervised fine-tuning and reward-modeling hyperparameters for OPT-family models.Source: Table 6 · p. 16
  • Table 7 summarizes PPO hyperparameters.Source: Table 7 · p. 16

C COMPARISON WITH ALL BASELINES

ARGS outperforms the best baseline in both non-finetuned and fine-tuned settings, with the comparison spanning decoding methods and model tuning conditions.

C COMPARISON WITH ALL BASELINES

  • ↑27% in average reward for non-finetuned ARGS and ↑18% for fine-tuned ARGS over the best baseline.Table 10 compares greedy and stochastic ARGS with vanilla greedy decoding, top-k sampling, nucleus sampling, and contrastive search, using base and fine-tuned models.Source: C COMPARISON WITH ALL BASELINES · p. 16
  • ARGS-greedy and ARGS-stochastic outperform established baselines in both non-finetuned and fine-tuned scenarios.Source: C COMPARISON WITH ALL BASELINES · p. 16
  • Table 8 summarizes Direct Policy Optimization (DPO) hyperparameters.Source: Table 8 · p. 17
  • Table 9 shows a sample GPT-4 evaluation prompt, distinguishing the prompt from the responses being compared by color.Source: Table 9 · p. 17

D ADDITIONAL QUALITATIVE EXAMPLES

The section presents additional qualitative examples comparing standard greedy decoding with ARGS-greedy.

D ADDITIONAL QUALITATIVE EXAMPLES

  • Table 11 presents additional qualitative examples illustrating how ARGS can steer decoded outputs toward greater alignment with human preference; the referenced setup for LLaMA-7B is described in Section 3.1.Source: D ADDITIONAL QUALITATIVE EXAMPLES · p. 17

E LIMITATIONS AND FUTURE WORK.

The paper evaluates ARGS on established public alignment benchmarks, while identifying more complex tasks and reward-model approaches as directions for future work.

E LIMITATIONS AND FUTURE WORK.

  • Current evaluations use standard alignment benchmarks, including HH-RLHF and Stanford Human Preferences, both described as large and publicly available.The authors identify multi-step reasoning as a more intricate task for future evaluation.Source: E LIMITATIONS AND FUTURE WORK · p. 17
  • Table 10 compares decoding methods across models with and without fine-tuning.Source: Table 10 · p. 18

LLaMA-7B-SFT

  • The authors note that a higher-quality reward model results in enhanced generation quality.Source: LLaMA-7B-SFT · p. 18
  • Table 11 gives comparative examples of greedy and ARGS-greedy decoding, with ARGS settings w = 1.5 and k = 10.Source: Table 11 · p. 19

Related papers

Loading 2402.01694v1…