PaperlayerView exact revision on arXiv

Source-linked AI summary

Decoding-Time Language Model Alignment with Multiple Objectives

Ruizhe Shi, Yifang Chen, Yushi Hu, Alisa Liu, Hannaneh Hajishirzi, Noah A. Smith, Simon S. Du

arXiv:2406.18853v3cs.LG

TL;DR

Existing alignment methods focus on a single reward, limiting support for differing objective preferences. The paper introduces MOD, which combines predictions from single-objective models at decoding time. It reports improved reward over parameter merging and benefits when combining models trained for different objectives.

  • Problem

    Existing alignment approaches such as PPO and DPO optimize a single reward, despite users needing different weightings across objectives.

    Source: §1 · p. 1

  • Method

    MOD uses a Legendre-transform-based closed form to combine predictions from single-objective models for preference-weighted decoding.

    Source: §1 · p. 2

  • Results

    12.8% overall relative reward improvement over parameter merging on three equally weighted Helpful Assistant objectives; combining three TÜLU models nearly eliminates Toxigen toxicity and improves Codex@1, GSM-COT, and BBH-COT by 7.9%–33.3%.

    Source: §1 · p. 3

  • Takeaways & Limitations

    MOD allows users to vary objective weightings at inference time without extensive retraining iterations.

    Source: §1 · p. 2

    Because MOD does not yield a sampling policy, the authors cannot directly measure KL divergence from the reference policy and instead show example generations.

    Source: §4.1 · p. 5

Abstract

from arXiv · show

Aligning language models (LMs) to human preferences has emerged as a critical pursuit, enabling these models to better serve diverse user needs. Existing methods primarily focus on optimizing LMs for a single reward function, limiting their adaptability to varied objectives. Here, we propose $\textbf{multi-objective decoding (MOD)}$, a decoding-time algorithm that outputs the next token from a linear combination of predictions of all base models, for any given weightings over different objectives. We exploit a common form among a family of $f$-divergence regularized alignment approaches (such as PPO, DPO, and their variants) to identify a closed-form solution by Legendre transform, and derive an efficient decoding strategy. Theoretically, we show why existing approaches can be sub-optimal even in natural settings and obtain optimality guarantees for our method. Empirical results demonstrate the effectiveness of the algorithm. For example, compared to a parameter-merging baseline, MOD achieves 12.8% overall reward improvement when equally optimizing towards $3$ objectives. Moreover, we experiment with MOD on combining three fully-finetuned LLMs of different model sizes, each aimed at different objectives such as safety, coding, and general user preference. Unlike traditional methods that require careful curation of a mixture of datasets to achieve comprehensive improvement, we can quickly experiment with preference weightings using MOD to find the best combination of models. Our best combination reduces toxicity on Toxigen to nearly 0% and achieves 7.9--33.3% improvement across other three metrics ($\textit{i.e.}$, Codex@1, GSM-COT, BBH-COT).

1 Introduction

Existing alignment methods optimize a single reward, while users may need different objective weightings without retraining for each preference. The paper introduces training-free multi-objective decoding (MOD), with theoretical guarantees and empirical improvements over parameter merging.

1 Introduction

  • Existing alignment methods such as PPO and DPO optimize a single reward, while users may prefer different weightings across objectives.Mixed-dataset curation is challenging and resource-intensive, and prior approaches that combine reward functions still require retraining for possible weightings.Source: §1 · pp. 1–2
  • MOD combines predictions from models trained for individual objectives to produce a training-free policy for a weighted interpolated reward.The method derives a closed-form solution using a Legendre transform and efficiently decodes it; the paper reports optimality guarantees under its assumptions.Source: §1 · pp. 2–3
  • MOD supports arbitrary objective weightings at inference time and can combine models of different scales, including PPO, DPO, and supervised-finetuned models.The paper reports experiments combining 13B DPO models with a 7B SFT model for Open Instruction-Following.Source: §1 · p. 2
  • 12.8% overall relative reward improvement over parameter merging when MOD equally optimizes three objectives on the Helpful Assistant task.The paper also reports nearly zero Toxigen toxicity and 7.9%–33.3% relative improvements across Codex@1, GSM-COT, and BBH-COT when combining three TÜLU models.Source: §1 · p. 3

2 Preliminaries

The paper formulates multi-objective alignment as maximizing a preference-weighted sum of rewards using predictions from existing single-objective models. It places these models in an f-divergence regularization framework and identifies a condition enabling closed-form mappings between policies and rewards.

Multi-objective reward functions

  • The multi-objective reward is the preference-weighted sum of reward functions, with preferences represented by a normalized vector.The paper focuses on maximizing this reward through decoding that combines outputs from existing single-objective aligned language models.Source: §2 · p. 3

Single objective alignment with f-divergence regularization

  • PPO, DPO, and related methods regularize single-objective policies against a reference policy using f-divergences, including KL-divergence.The paper also situates forward KL, JSD, and α-divergence among proposed alternatives to reverse KL.Source: §2 · p. 3

Strong-barrier function benefits multi-objective decoding

  • When the barrier function is continuously differentiable and strongly convex on R+, a closed-form bijection links each aligned policy to its reward.The paper calls an f with the relevant properties a strong-barrier function and uses the resulting closed forms to motivate combining model outputs.Source: §2 · p. 4

Formal problem formulation

  • The formal problem asks for retraining-free decoding that optimizes a weighted multi-objective reward given a reference policy and single-objective base policies.The setup assumes the individual reward functions are inaccessible.Source: §2 · p. 4

3 Proposed Method: Multi-Objective Decoding

The proposed method first derives a direct combination of base-policy outputs, then removes an intractable normalization factor to enable efficient decoding. Its main algorithm uses greedy token-level decoding; the approach trades increased memory use for avoiding retraining across preferences.

3.1 Warm-up: an inefficient decoding version

  • The warm-up decoding form combines base-model outputs, but computing its exact normalization factor is nearly impossible because it requires looping over all possible outputs..Source: §3.1 · p. 4

3.2 Towards an efficient algorithm: reformulation and approximation

  • The Legendre-transform reformulation removes the normalization factor, and greedy search converts response-level decoding into efficient token-level decoding.The revised optimization has an optimal solution under the theorem's stated threshold condition.Source: §3.2 · pp. 4–5

3.3 Main algorithm: efficient decoding with optimality for strong-barrier function

  • Greedy token-level decoding turns the closed-form combination into an efficient decoding rule, with the main experiments using reverse KL, equivalent to linearly combining model logits by preference weights.The paper states that greedy search improves efficiency without much performance loss and that MOD requires multiple models loaded simultaneously, increasing memory cost; low-rank adapters or distributed deployment can mitigate it.Source: §3.3 · p. 5, §3.2 · p. 5

4 Experiments

Experiments evaluate MOD across Reddit Summary, Helpful Assistant, Safety Alignment, and Open Instruction-Following. Results generally favor MOD over parameter merging, with some comparisons and evaluation constraints qualifying the findings.

4.1 Experiment setup

  • Experiments compare MOD with parameter-merging and retraining baselines across four tasks, evaluating Pareto-frontier area with off-the-shelf reward models.Because MOD does not yield a sampling policy, reference-policy KL is not measured directly; example generations are used to indicate whether outputs deviate much from the reference policy.Source: §4 · p. 5, §4.1 · p. 5

4.2 Results

  • On Reddit Summary, MOD’s Pareto frontier generally lies over RS and MORLHF across preference weightings.The experiment uses two PPO-tuned models for summary quality and faithfulness.Source: Figure 2 · p. 6, §4.2 · p. 6
  • MOD prominently beats RS for every Helpful Assistant reward pair, but trails MORLHF when balancing harmlessness and humor.MORLHF is more costly.Source: §4.2 · p. 6, Figure 3 · p. 6
  • MOD generally outperforms RS across Reverse KL, JSD, 0.3-divergence, and 0.5-divergence safety-alignment experiments.MODPO is applicable only to KL-divergence; MOD is more convenient but has a slight performance gap against this retraining baseline.Source: §4.2 · p. 6, Figure 4 · p. 6
  • For safety alignment, decreasing w2 makes MOD increasingly harmful; at w = (5, −4), RS fails to produce meaningful responses.MOD can use non-all-positive preference weights that sum to one, allowing optimization for a reward function −R.Source: §4.2 · p. 6, Table 2 · p. 7, Table 3 · p. 7
  • With weights (0.75, 0.1, 0.15), MOD reduces Toxigen to nearly 0 and improves the other three metrics by 7.9–33.3% versus CODETÜLU-2-7B.The table covers combinations of CODETÜLU-2-7B, TÜLU-2-HH-13B, and TÜLU-2-ULTRA-13B; combinations show better overall performance than single models.Source: Table 4 · p. 7, §4.2 · p. 7, Table 5 · p. 7

5 Theoretical Analysis

The analysis identifies settings where parameter merging lacks guarantees, establishes barrier-function regularization as necessary for optimal multi-objective alignment, and bounds error propagation when base policies are sub-optimal. It also extends the framework to explain proxy-tuning and related decoding approaches under assumptions.

5.1 Failures of parameter-merging paradigm

  • Parameter merging lacks a theoretical guarantee in practical scenarios, and choosing merging weights equal to preference weights can be sub-optimal.The reduced reward mis-specification hypothesis fails for almost all f-divergence-regularized policies under the stated theorem conditions.Source: §5.1 · pp. 7–8

5.2 Necessity of barrier function

  • A barrier function is necessary to find an optimal policy for multi-objective alignment using base policies.The theorem gives a counterexample for non-barrier f-divergences; the authors therefore identify barrier regularization as crucial for bridging policies.Source: §5.2 · p. 8

5.3 Sub-optimality error propagation

  • When base policies are sub-optimal, the performance bound depends on their log-probability distance from optimal policies and their divergence from the reference policy.The authors interpret these distances as small when policies are near-optimal and remain close to the reference, as can occur with divergence constraints and early stopping.Source: §5.3 · pp. 8–9

5.4 Beyond f-divergence regularized alignment and multi-objective decoding

  • Under certain assumptions, the framework theoretically justifies proxy-tuning, which shifts a larger model’s predictions using the logit difference between tuned and untuned smaller models.The authors also derive connections to other decoding approaches in Appendix C.3.Source: §5.4 · p. 9

6 Related Work

Related work spans preference-alignment objectives, controllable decoding, and multi-objective alignment strategies. MOD is positioned as a decoding-time alternative to retraining, parameter merging, and preference-conditioned prompting.

Algorithms for aligning LMs to human preferences

  • PPO and DPO optimize reward-based alignment with divergence penalties, while related methods alter the reward mapping, divergence, training signal, or use of human feedback.The paper focuses on PPO, DPO, and their extensions.Source: §6 · p. 9

Decoding-time algorithms for controllable generation

  • Token-level decoding methods guide individual token choices, whereas response-level methods optimize or sample whole outputs; proxy-tuning steers a larger model through logits from a smaller tuned model.The paper describes DeRa as exploring a special case of MOD with a per-token distribution approximation.Source: §6 · p. 9

Multi-objective LMs alignment

  • Multi-objective alignment approaches include retraining on combined rewards, parameter merging, and preference-conditioned prompting, with the latter two relying on assumptions or out-of-distribution generalization.The cited approaches include MORLHF and MODPO for retraining, rewarded soups for parameter merging, and prompting methods conditioned on preference weights.Source: §6 · p. 9

7 Conclusion

The paper concludes that MOD offers a training-free decoding-time approach to multi-objective alignment, supported by theoretical analysis and empirical evidence. It also acknowledges that its analysis and decoding setup have limitations.

7 Conclusion

  • MOD is presented as a practical, training-free decoding-time method for multi-objective alignment under the f-divergence framework.The conclusion frames it as applicable across diverse tasks and use cases.Source: §7 · p. 10
  • The analysis is primarily tabular and omits function-approximation error, while token-level decoding from a response-level distribution may degrade performance.The authors note energy-based approaches as a possible way to alleviate the latter issue.Source: §7 · p. 10

Appendices

The appendices include an impact statement, notation, a decoding pipeline, and divergence-specific policy forms. They also warn that misuse could lead to harmful or privacy-infringing outputs.

A Impact Statement

  • Misuse in downstream tasks could induce harmful, offensive, or privacy-infringing generation, though the authors distinguish this risk from their focus on a general algorithm with theoretical guarantees.Source: A Impact Statement · p. 16

B Notation

  • The notation defines the policy set as mappings from prompts to probability distributions over responses.It also defines the prompt set, alphabet, response set, probability simplex, and Kronecker delta.Source: B Notation · p. 16

C.1 Pipeline

  • The decoding pipeline takes a prompt, beam count, maximum length, divergence, preference weights, and reference and base policies, then returns the highest-scoring completed sequence.It iteratively expands candidate tokens and retains the top K sequences by score.Source: §C.1 · p. 16

C.2 Divergence measures and closed-form policies

  • The appendix presents divergence measures alongside optimal sampling policies and approximated decoding policies for multi-objective alignment.The listed measures include reverse KL, forward KL, Jensen–Shannon, total variation, and chi-squared divergences.Source: §C.2 · pp. 16–17

C.3 Extended variants

The appendix extends MOD’s decoding framework to SFT and proxy-tuning variants, and defines the divergence conditions used to characterize its alignment setting.

C.3 Extended variants

  • MOD’s Eq. (7) also applies to SFT models under the stated assumption that SFT implicitly optimizes a reward with respect to Reverse KL-divergence.The assumption treats supervised fine-tuning from a pretrained model as reward optimization under Reverse KL regularization.Source: §C.3 · p. 17
  • Proxy-tuning, jail-breaking, and δ-unlearning are identified as instances of related model-tuning operations derived from the stated policy relationship.The passage identifies proxy-tuning with tuning a model from a tuned version, and relates reversing the model positions to jail-breaking.Source: §C.3 · p. 17
  • Multi-objective proxy-tuning uses smaller-model variants tuned for separate objectives to decode outputs from a larger pretrained model.The variants may be tuned using PPO, DPO, or SFT before decoding.Source: §C.3 · p. 17

D.1 Definitions

  • A strong-barrier function is a continuously differentiable, strongly convex barrier function, with the barrier condition requiring 0 not to lie in the domain of ∇f.The appendix defines f-divergence for probability measures using a convex f satisfying f(1)=0.Source: §D.1 · p. 18

D.2 Proofs of subsection 5.2

The proofs establish conditions for optimal decoding from single-objective policies and show where related optimality assumptions can fail.

D.2 Proofs of subsection 5.2

  • For non-barrier f, the proof constructs policies meeting the single-objective conditions while producing a strict inequality, illustrating failure of the claimed multi-objective relationship.The construction’s inequality is r+s>0.Source: §D.2 · pp. 18–20
  • Under strong-barrier regularization, Proposition 1 gives the weighted-reward optimal policy as a normalized combination of the single-objective policies.Its normalization factor is numerically computable when the response set is finite.Source: §D.2 · p. 21

D.3 Proof of key theorem

  • Theorem 5 establishes that the policy characterized by the theorem is an optimal solution for the weighted objective under strong-barrier regularization.The result applies to preference weightings in the simplex.Source: §D.3 · p. 22

D.4 Proofs of subsection 5.3

  • The paper evaluates a policy for the weighted reward using its negative KL divergence from the optimal policy.The cited identity rewrites reward minus reference-policy KL as negative KL from the optimum plus a constant.Source: §D.4 · p. 23

D.5 Proofs of subsection 5.1

  • Theorem 2 shows the reduced reward mis-specification hypothesis can fail for non-barrier divergences, Reverse KL, and strong-barrier functions with finite roots.The proof includes a non-barrier case and a Reverse-KL construction where the hypothesis does not hold.Source: §D.5 · pp. 25–26

E Implementation Details

The implementation appendix documents evaluation and training choices; supplementary experiments compare MOD with rewarded soups and show MOD generally performs better on the reported settings.

E Implementation Details

  • The implementation uses specified SFT datasets and models for Reddit Summary, Helpful Assistant, Safety Alignment, and HelpSteer.The appendix also reports PPO and DPO training and inference settings, and approximates the JSD solution using Reverse KL in the inference code.Source: E Implementation Details · p. 27
  • Evaluation uses down-sampled prompt sets and off-the-shelf reward models, with identical generation configurations across algorithms.The sampled sets contain 2,000 test prompts for Reddit Summary and Helpful Assistant, and 200 validation prompts for Safety Alignment and HelpSteer.Source: E Implementation Details · p. 28

F.1 Motivating example

  • The motivating experiment tunes two T5-LARGE models for relevance and factuality, then reverses selected Q and K matrix signs to form one policy.After that parameter reversal, the reported RS performance noticeably lags, while performance was comparable for normally trained models.Source: §F.1 · p. 28

F.2 Additional results for Helpful Assistant

  • MOD generally beats rewarded soups across the reported three-objective Helpful Assistant preference configurations.The supplementary results cover many preference weightings and include a 3D Pareto-frontier visualization.Source: §F.2 · p. 28

F.3 Additional results for BeaverTails

The supplementary experiments report MOD’s comparisons with rewarded soups across safety, HelpSteer, and instruction-following settings, and provide example generations.

F.3 Additional results for BeaverTails

  • MOD can use a negative harmfulness weight, while rewarded soups fail to generate meaningful responses when that weight is at most −2.The comparison is illustrated with Safety Alignment example generations.Source: §F.3 · p. 28, Table 7 · p. 30
  • 12.8% overall improvement over the parameter-merging baseline is reported for MOD when equally optimizing three Helpful Assistant objectives.The weighted score combines helpfulness, harmlessness, and humor.Source: Table 6 · p. 29

F.4 Additional results for HelpSteer

  • With early-stopped f-DPO base models and equal weights, MOD generally combines their advantages and achieves stronger performance than RS on HelpSteer.The experiments cover Reverse KL, JSD, 0.3-divergence, and 0.5-divergence models tuned for helpfulness, complexity, and verbosity.Source: §F.4 · p. 31, Table 8 · p. 31, Table 9 · p. 31

F.5 Additional results for Open Instruction-Following

  • Table 12 reports MOD results combining TÜLU-2-HH-13B and CODETÜLU-2-7B for safety, coding, and reasoning capabilities.The table caption identifies Toxigen, Codex, and COT metrics as the evaluation areas.Source: Table 12 · p. 32, §F.5 · p. 32

F.6 Example generations

  • The appendix provides representative generated responses for each dataset, comparing models and algorithms at equal MOD and RS weights of (0.5, 0.5).Examples are organized in Tables 13–23.Source: §F.6 · p. 33, Table 13 · p. 33, Table 14 · p. 33

Related papers

Loading 2406.18853v3…