Orchestrating LLMs with Different Personalizations

The MoPE inference setting assumes that the user provides an instruction $x$ such as "Give me a 5 day itinerary for Switzerland" as well as a preference $\xi$ ("I want a helpful, concise, and funny response!") that consists of $n$ individual preference dimensions $\{p_{1},...,p_{n}\}$ (3 in Figure 1: helpfulness, conciseness and humour). We are also provided with $n$ black-box LLMs $\{M_{1},...,M_{n}\}$, each specialized along the corresponding preference dimension $p_{i}$ only. In other words $M_{3}$ is for example an LLM optimized to be humorous. We refer to these models as relevant expert models (REM). Given a partial response generated so far $y_{<t}$ (e.g. "Sure! Here’s a Swiss"), we want to construct a next token probability distribution in order to decode the subsequent token $y_{t}$.

Since we aim to achieve effective multi-objective personalization, we introduce a trainable neural network called the preference control model (PCM), parametrized by $\theta$. PCM takes as input the instruction $x$, the partial response $y_{<t}$ and the preference vector $\xi$ and outputs a weight vector with length $n$ (one weight for each REM) whose entries are non-negative and sum to 1. Assuming we already have a well-trained PCM, the inference of MoPE is very similar to the standard autoregressive decoding mechanism of LLMs. Specifically, we can construct the following next token probability distribution as a weighted sum of all experts:

where $\pi_{M_{i}}(y_{t}|x,y_{<t})$ is the next token probability distribution from expert $M_{i}$ and $\alpha_{\theta}(x,y_{<t},\xi)$ is the output of the PCM. As each individual $\pi_{M_{i}}()$ outputs a probability distribution and $\pi_{\theta}$ is a convex combination of them, it itself is also a well-defined distribution over all tokens in the vocabulary. As each $M_{i}$ is specialized for one specific preference dimension, the weights in $\alpha_{\theta}()$ specify how much importance should be given to each dimension $p_{i}$ at time step $t$. Individual experts $M_{i}$ are frozen and can be treated as black-box APIs since only their output probabilities are needed. Various decoding methods such as greedy or temperature sampling can be used to decode the next token $y_{t}$ from the distribution $\pi_{\theta}$. As the PCM only outputs a distribution over $n$ dimensions instead of the vocabulary size, it can be a relatively small model that effectively orchestrates the large models $M_{i}$.

An overview of the MoPE training procedure is shown in Figure 2. Our framework uses online RL algorithms such as PPO [35] and REBEL (REgression to RElative REward Based RL) [6] to train the PCM $\alpha_{\theta}$. Similar to prior work [2], we assume we have access to a black-box reward model for each individual dimension (e.g., a reward model that can quantify the level of helpfulness). Alternatively, one can also train such reward models from existing pairwise comparison data such as Cui et al. [36]. We define $y\sim\pi_{\theta}(\cdot|x,\xi)$ as the method for generating a response $y$ (a sequence of tokens) following MoPE inference procedure. Furthermore, $\pi_{\theta}(y|x,\xi)$ represents the probability of generating the response $y$. Note that since our models are autoregressive, we have $\pi_{\theta}(y|x,\xi)=\prod_{t}\pi_{\theta}(y_{t}|x,\xi,y_{<t})$, i.e., the likelihood of the whole response is the product of the likelihood of each token. Below we give our formulation for the reward modeling for RL training.

We evaluate how good baselines and MoPE can satisfy multifaceted preferences on open-ended generations. We evaluate on the same subset of 50 instances of Koala dataset [8] from [2]. Following [2], we construct eight distinct preferences from six individual dimensions: elementary, knowledgeable, concise, informative, friendly and unfriendly. The six dimensions can be categorized into three groups of two where they are opposite of each other (A and B). A preference is then formed by drawing one dimension from each group. The preference dimensions and instructions can be found in Table 1.

We perform pairwise evaluation of all methods for each preference and Koala instruction instance. We use the same win rate calculation method as [2]. Specifically, for a pair of responses $y_{A}$, $y_{B}$, we use GPT4 to simulate human judge and evaluate each dimension separately. GPT4 can assign either WIN, TIE or LOSE to the pairwise comparison, which translates to a numerical score of 1, 0 and -1. To evaluate the overall performance for the pair of generations, we sum the numerical scores across individual dimensions. $y_{A}$ is considered to be better / equally good / worse than $y_{B}$ if the overall score is greater / equal / less than 0. Finally, we calculate the overall win rate of two methods from all pairwise comparisons that do not lead to TIE. We use greedy decoding to generate responses from all methods.

We use Tulu-7B [7], an instruction-tuned language model, as the base model for reward models and expert models. Specifically, we use the reward model training data from [2] to first train a reward model for each of the six preference dimension in Table 1. We note that although we trained reward model ourselves, the reward model can technically be off-the-shelf classifiers or even black-box API models. Then, to obtain the experts that specialize in each preference dimension, we perform RLHF by using PPO to fine-tune a separate base model for each preference. After training, the six experts are always frozen without any update. Note that we selected these design choices to maintain consistency with Jang et al. [2] to permit a fair comparison. For the preference control model, to illustrate that we can control large expert models, we use a much smaller LLaMA based model¹¹1https://huggingface.co/JackFram/llama-160m that has 160M parameters in total. The final linear layer of this model that originally outputs the next token probability distribution is replaced by another randomly initialized linear layer with a size equal to the number of preference dimensions in a preference, i.e. 3 in our setting which is significantly smaller than the size of the vocabulary. We use LoRA [5] for all model training.

We compare MoPE with several baselines ranging from directly prompting the model to full-blown multi-objective RL training. In Table 2, we categorize methods with two characteristics: 1. whether weight access to the individual experts is needed. 2. whether additional training (of the experts or other models) is required. Below we give more details and explanations for each baseline.

In Table 3, we summarize the performance of baselines and MoPE. The first thing to note is that the average win rate for vanilla prompting is significantly lower than other methods, indicating providing preferences in addition to instructions is crucial for preference personalization. Because Tulu-7B is an instruction-tuned model, preference prompting is a very competitive baseline and achieves an average win rate of more than 54%. Personalized Soup is slightly better than the prompting baseline, suggesting merging on the expert weights could slightly improve personalization. Multi-Objective RL is the best baseline since it can train the large 7B model directly with RLHF. MoPE outperforms the baselines, achieving an average win rate of 62.00%. This shows that although MoPE does not directly fine-tune the large experts, learning how to control the specialized experts via output merging using a much smaller model is already capable of achieving strong performance.

In Table 5, we ablate MoPE through several axes of configuration. We first study how MoPE performs without any training in the preference control model. In this case, we simply uniformly merge the outputs from the preference experts. We study two spaces to merge the outputs: logit and probability space. As seen in the first two rows from Table 5, merging on the probability space outperforms its counterpart in logit space. This is possibly due to logit space is not normalized and directly adding the logits from different experts could lead to drastic change of distribution. Merging on the probability space outperforms the prompting baselines and weight merging in Table 3, confirming the empirical benefit of merging in the output probability spaces over merging in the parameter spaces. Next, instead of modeling the reward with the Bradley-Terry calculation, directly averaging the reward from individual dimensions achieves much lower performance. This empirically confirms our intuition of using the reward signal from the BT model: these reward signals are always normalized at the same scale and are more interpretable (i.e., probability of winning over a reference response). Finally, we also show that other online RL algorithms, such as PPO, can also be directly applied to MoPE and achieve competitive performance against baselines, indicating the flexibility of our framework in terms of integrating different RL black-box algorithms.

In Table LABEL:tab:qualitative_comparison_1, we provide a few examples of response generated by Preference Prompting, Personalized Soup, and MoPE to qualitatively demonstrate the difference between the methods. In the first example, the generations from Preference Prompting and Personalized Soup both overuse repetitive metaphors, which makes the comparisons more confusing and harder to understand, while the generation from MoPE keeps the use of metaphors relatively simple and straightforward, keeping the message clear and easier to understand while also being more concise, while the more casual tone of the conclusion makes it feel more friendly than the other generations. In the second example, the generations from Preference Prompting and Personalized Soup both use slightly more advanced terminology that could be difficult to understand, while the generation from MoPE uses simpler language in a more straight-forward explanation that is both easier to understand and more concise, while its blunt denial gives it a more unfriendly tone than the other two, which give more mixed responses with more neutral tones. In the third example, the generation from MoPE offers a slightly more in-depth explanation and uses more technical language that an expert is likely to be familiar with without significantly sacrificing conciseness, being of similar length to the other two generations.

Compared to prompting and weight merging, output merging requires more compute and memory to generate responses linear to the number of models contributed to the merged composition. In addition, since output merging has not been used widely, a custom efficient implementation is usually required for faster decoding. We have made an efficient implementation for Huggingface library via explicit key value cache and plan to open source it for future research.

Despite the compute and memory overhead, output merging has two important benefits that motivate our work. First, compared to weight merging and other fine-tuning methods, output merging does not require access to the expert parameters. Since many state-of-the-art language models [40, 41] are closed-sourced, output merging can be potentially applied to merging of such models whereas other approaches are not applicable. Besides, even with the compute and memory overhead, output merging is more friendly to a large scale deployment setting than weight merging. Given a batch of requests with many distinct preferences, because one set of merged weights can only process requests for a specific preference, weight merging either requires sequential serving the requests or spawning resource that can host as (possibly exponentially) many distinct preferences as requested for parallel serving. Output merging, however, can perform parallel serving with just the resource needed to host linear number of individual experts.