Conditional Language Policy: A General Framework For Steerable Multi-Objective Finetuning (2024)

Abstract: Reward-based finetuning is crucial for aligning language policies with intended behaviors (*e.g.*, creativity and safety). A key challenge is to develop steerable language models that trade-off multiple (conflicting) objectives in a flexible and efficient manner. This paper presents Conditional Language Policy (CLP), a general framework for finetuning language models on multiple objectives. Building on techniques from multi-task training and parameter-efficient finetuning, CLP learn steerable models that effectively trade-off conflicting objectives at *inference time*. Notably, this does not require training or maintaining multiple models to achieve different trade-offs between the objectives. Through extensive experiments and ablations on two summarization datasets, we show that CLP learns steerable language models that outperform and Pareto-dominate the existing approaches for multi-objective

## introduction reinforcement learning (rl) is a crucial step for finetuning language models (lms) to produce desirable, human-aligned behaviors @xcite in various domains such as summarization @xcite , conversational agents @xcite and encoding social norms @xcite . in practice, defining the scalar reward function for rl is challenging due to the plurality of human preferences (e.g., succinctness vs. completeness, factuality vs. creativity) and of applications (e.g., summarization, coding, dialog) @xcite . the common practice is to linearly combine multiple, often conflicting rewards with weightings figure 1 : (left) for a fixed prompt x, a multiobjective lm can output y 1 , y 2 , y 3 for different weightings w 1 , w 2 , w 3 of two rewards r 1 and r 2 , such that the response y i for weighting w i is preferred under the weighted reward w i [1]r 1 + w i [2]r 2 . (right) pareto-fronts when using the rewards nli and rouge (sec. 4.1.2). rewarded soups (rs) @xcite is pareto-dominated by both full-clp (this paper) and prompting (say, @xcite ), but full-clp is more appealing for its steerability, evidenced by its wider pareto-front. in sum, pareto-dominance (pushing out the front) and steerability (stretching out the front) are both key desiderata for moft. a solution we explore in this paper is to frame pluralistic alignment as a multi-objective rl problem @xcite , where we train a multi-objective lm (see fig. 1 left) that can be steered to generate desirable outputs over the continuum of possible reward weightings. specifically, the multi-objective lm takes as input the prompt x and reward weighting w and responds with outputs y ∼ π mo (• | x, w) that maximize the specified reward combination j w j r j in expectation (up to kl-regularization). in contrast, the classic single-objective lm only takes input the prompt x and outputs y ∼ π so (• | x) that maximize a fixed scalar reward, such as j w fix j r j for some fixed weightings vector w fix . an important application of multi-objective language models is decision support: a system which can provide multiple diverse generations that support or cover a wide space of interests and lets the user pick their favorite @xcite , which can in turn be used to personalize the model to the user. in the context of lms, multi-objective finetuning (moft) has been explored via prompt-based approaches @xcite @xcite and parameter-based approaches @xcite . promptbased approaches finetune a lm that is steered by simply including the reward weightings into the prompt. however, prompt-based moft is suboptimal in steerability as we show in our experiments and can be sensitive to how weightings are encoded in the prompt. an alternative parameterbased approach, rewarded soups (rs) @xcite , (a) independently trains one lm per reward function, and (b) interpolates parameters from separate lms with reward weightings to perform conditional generation at inference-time. surprisingly, this zero-shot approach can effectively trade-off multiple rewards by relying on the linear mode connectivity @xcite . however, we find that zero-shot moft is sub-optimal in performance on intermediate weightings that it hasn't seen during training. this paper presents conditional language policy (clp), a general framework for that employs parameter-space conditioning with multi-task training @xcite and achieves state-of-the-art results in our extensive experiments on summarization. using parameter-conditioning from rs, clp is consistently more steerable than purely prompt-based approaches. moreover, by finetuning on a diverse set of reward weightings, clp produces higher quality responses than zeroshot approaches like rs while having comparable or better steerability. we conduct a systematic set of experiments, observing that clp both pareto-dominates rs and is more steerable than prompt-based moft (see fig. 1 right). clp robustly maintains these benefits across experimental conditions and across various reward functions and model sizes. we also conduct an automated evaluation with gemini 1.0 ultra @xcite that further supports that clp is more steerable and generates higher quality responses than existing baselines. finally, we provide novel theory proving that zero-shot methods can display near-pareto-optimal behavior when optimal policies for individual rewards are aligned, but they provably fail otherwise. in those failure cases, multi-task training (as used by clp) is necessary to learn to pareto-optimal policy. in summary, our contributions are: ## problem setup let π ref (y | x) be a base policy with parameters θ ref , where x is the input prompt and y is the output generation. in single-objective finetuning (soft), there is a fixed reward function r(x, y) and the goal is to maximize the expected reward without drifting too far from π ref . formally, soft learns a policy π θ (•) to maximize the following value: notice that the kl-term can be thought of as a special reward function to mitigate reward-hacking @xcite . the main issue with soft is that both the reward function r and the kl weightings are fixed and unchangeable after training; thus, it is not possible to offer near-optimal behavior on other reward and kl weightings without retraining soft with multiple reward or α-values. multi-objective finetuning (moft) fixes the above issues by learning a conditional policy that generates y based not only on the prompt x, but also algorithm 1 clp: conditional language policy input: weightings sampler q, number of training steps t , learning rates {η t } t @xmath0. 1: init clp: φ 0 = (θ ref [s c ], {θ ref [s]} @xmath1. 2: for t in 0 . . . t -1 do 3: sample prompt x t ∼ d and kl & reward weightings (α t , w t ) ∼ q. 4: get conditioned policy π t (•; α t , w t ) from running algorithm 2 with α t , w t , φ t . 5: conditioned generation y t ∼ π t (x t ; α t , w t ) (same computation as single lm). 6: objective: r the α and r specified at test-time without retraining. we formulate moft using multi-objective rl (roijers, 2016) and consider a vectorial reward r(s, a) ∈ r m whose linear scalarizations {w r(•) : @xmath2, where ∆ m is the (m -1)-simplex, i.e., r is a basis for all desired rewards at test-time. the goal is to learn parameters φ such that, for all weightings α ∈ [α min , 1] and @xmath3, the conditioned policy π φ (•; α, w) maximizes v α,w r , the objective with kl-regularizer α and reward function w r(•). we frame moft as multi-task training over the weighting distribution q, and aim to maximize: in theory, the optimal solution to multi-objective rl is a large (potentially exponential in m) set of policies called the convex coverage set @xcite . however, the rich representational power of a single llm may already be able to approximate such a policy cover and thus we aim to solve moft with a memory-efficient parameterization φ. ## conditional language policy (clp) this section presents the conditional language policy (clp) framework for moft, which enables a family of algorithms with varying trade-offs between quality (in terms of pareto-dominance and steerability) and cost (in terms of parameter-count). in brief, clp learns a set of parameters φ that can be processed into a conditioned lm for any given weighting across rewards and kl, via a parameteraveraging mechanism described in sec. 3.1. the learning algorithm samples a diverse set of weightings to push out its pareto-front over all weightings simultaneously. notably, this is multi-task learning across the continuum of weightings, which directly maximizes the moft objective defined in eq. ( 2 ), unlike the existing zero-shot approaches @xcite . we describe the clp algorithm in algorithm 1 (illustrated in fig. 10 ), where each training round @xmath4, 2, . . . , t consists of three steps. first, we sample a prompt x t and reward & kl weightings w t , α t ∼ q for this round (line 3). second, we condition clp on weightings (α t , w t ) and sample generations y t ∼ π t (x t ; α t , w t ) (line 5). third, we compute the conditioned objective (line 6) and update the clp parameters with policy gradient (line 7). the policy optimization step uses gradient ascent to maximize the objective in eq. ( 2 ) and can be implemented by any standard rlhf algorithms such as reinforce @xcite , @xcite , and also dpo @xcite when given pluralistic preference data. ## experiments we consider the following questions: data and models. most of our experiments and ablations are performed on summarization with the widely-used xsum dataset @xcite . we initialize the reference policy π ref and reward models from the instruction finetuned (flan) checkpoints for t5 @xcite . we use the large size (770m parameters) for reward models and we mostly use the base size (220m parameters) for policies, except in our model size ablation where we also use large size for policies. in addition to these main results, we also ran experiments on the tldr dataset @xcite with larger policy (t5-xl size) and reward models (t5-xxl size) to show that our general observations are scalable and robust; these experiments are presented in sec. 4.4. reward functions. we consider three reward functions: (1) rouge, a formulaic reward that measures similarity of generation to the ground truth summary @xcite ; (2) natural language inference (nli), a learnt reward model for textual entailment and factuality @xcite ; and (3) a reward model for summary quality learnt from the "too long; didn't read" (tldr) dataset @xcite . rouge and the quality model (referred to as tldr) tend to favor verbose and descriptive summaries while nli favors concise summaries; this gives our setup a distinct tension between various reward pairs. methods. we benchmark the three instances of clp in sec. 3.2: full-clp, attn-clp, logit-clp. as a gentle reminder, full-clp maintains one full lm per reward, akin to standard model soups; attn-clp maintains replicas of all attention layers thus being more parameter efficient; and logit-clp maintains replicas only the logit layer which is theoretically grounded but less expressive. for policy optimization, we use reinforce with control variate, which is a lighter implementation than ppo @xcite and has been successfully used for summarization @xcite . baselines. we consider (a) rewarded soups (rs) @xcite all methods and baselines are run for same number of training iterations. see app. b for related details and hyper-parameters. ## related works approaches for multi-reward alignment (or moft) can be broadly classified into two categories: prompt-based and parameter-based conditioning. prompt-based conditioning approaches include personalized soups @xcite which use hand-crafted prompts for personalizing lms based on binary weights on different rewards; cpo @xcite which employs prompting within a dpo framework; ric @xcite , dpa @xcite use prompting within a supervised finetuning setup that differs from this paper which focuses on rl finetuning. on the parameter-conditioning front, rewarded soups (rs) @xcite presents a zeroshot approach (i.e. without multi-task training) to multi-reward alignment by inference time averaging of parameters for lms that are independently trained to optimize each of the rewards. a more recent manuscript @xcite presents an approach where the reward weightings are embedded as singular values within the adalora framework @xcite ; this can be framed as an instance of the proposed clp framework. with regards to kl realignment, decoding time realignment (dera) @xcite linearly mixes logits between π ref and another lm learned via soft with the minimum kl weight α min . @xcite showed that this idea is also effective for trading off multiple rewards. finally, model souping @xcite , learning policy sub-spaces @xcite , and objective weight conditioning (dosovitskiy and djolonga, 2020) have been applied in domains beyond lms. we leverage these advances along with multi-task training to develop steerable lms at inference time. ## theory for logit mixing and clp in this section, we perform a sensitivity analysis for logit mixing and derive regret bounds for its paretofront. while clp uses parameter mixing instead of logit mixing, this analysis is still be instructive due to the similarity between parameter and ensemble mixing @xcite . ## conclusion we introduced clp, a flexible framework for moft that leverages techniques from multi-task training and parameter efficient finetuning to develop steerable lms that adapt their generations to produce near pareto optimal behavior across different weightings of individual rewards. we provide extensive benchmarking and ablations to better understand factors that enable the development of steerable lms within the clp framework. we supplement this with theoretical results that present conditions under which zero shot approaches work and when multi-task training is necessary to obtain near-optimal behavior. promising future directions include (a) understanding other conditioning mechanisms such as soft tokens @xcite , (b) automated tuning of the weight sampling distributions @xcite , (c) non-linear reward scalarizations @xcite and risksensitive rl alignment @xcite , and (d) designing adaptive algorithms with instancedependent guarantees @xcite .