Learning to Decode Collaboratively with Multiple Language Models (2024)

Abstract: We propose a method to teach multiple large language models (LLM) to collaborate by interleaving their generations at the token level. We model the decision of which LLM generates the next token as a latent variable. By optimizing the marginal likelihood of a training set under our latent variable model, the base LLM automatically learns when to generate itself and when to call on one of the “assistant” language models to generate, all without direct supervision. Token-level collaboration during decoding allows for a fusion of each model’s expertise in a manner tailored to the specific task at hand. Our collaborative decoding is especially useful in cross-domain settings where a generalist base LLM learns to invoke domain expert models. On instruction-following, domain-specific QA, and reasoning tasks, we show that the performance of the joint system exceeds that of the individual models. Through qualitative analysis, we show models trained with our method exhibit several interesting collaboration patterns, e.g., template-filling, by visualizing the learned latent decisions.

## introduction techniques that combine the generations of multiple large language models (llms) at decoding time have benefits ranging from faster decoding speed @xcite , to more controllable generations @xcite , to more coherent, less repetitive text @xcite , and even enabling a large model to be "tuned" by combining its generations with those of a smaller model from the same family @xcite . a parallel thread of work has aimed to equip language models with the ability to infuse external tools into their generations, with the goal of incorporating outside knowledge and capabilities @xcite . language models are able to produce more 1 code: https://github.com/clinicalml/co-llm faithful and accurate generations when equipped with external apis @xcite @xcite , search engines or retrievers @xcite @xcite , or code executors @xcite . while powerful, these methods all require a prescription on how to combine the models and when to use the tools, either via specific formulas for combining the logits of multiple models, or through (weak) supervision on where to insert tool/api calls in the training data. in this work, we explore a different type of model combination figure 2 : illustration of the decoding procedure in co-llm, where a base (llama-7b) and assistant model (meditron-70b) collaborate to generate a correct response for a medical question. for each token, the deferral control predicts the probability of switching to the assistant model to decode the next token given the context: it defers when the probability is above some threshold η (indicated by a), and uses the decoded token as the context (highlighted with orange border). when using the base model alone, it may make factual mistakes (indicated by b); co-llm learns to use the assistant model at these positions to produce correct generations. figure 1 shows example generations from our method, co-llm. in the top example, llama-7b collaborates with llama-70b on instructionfollowing by generating a list template for the answer and then calling on llama-70b to fill in each element of the list. using the larger model as an assistant allows the smaller model to make effective use of a larger knowledge base and focus its efforts on learning the correct "scaffolding" for instruction responses. in the bottom example, llama-7b collaborates with llemma-34b (a domain-specific math model, @xcite by treating the latter as an api call to fill in parts of a latex formula. in both cases, the model predicts when to call the assistant by itself, behavior it learns from training without direct supervision on which contexts suit the assistant model well. this enables the emergence of qualitatively different collaboration methods (e.g., learning to scaffold, calling the large model as an api) based on what the task demands. section 2 describes our latent-variable model for collaboration during decoding, and section 3 describes the training and decoding procedures for co-llm in this model. in, section 4, we evaluate co-llm on instruction-following, mathematical reasoning, and domain-specific question-answering tasks. our results indicate that teaching models to collaborate improves performance across all of these tasks compared to using the individual mod-els alone, and can sometimes match or exceed the performance of fine-tuning the large model. by using chain-of-thought reasoning @xcite , co-llm can also be applied to classification tasks, where our experiments show that it boosts performance by enabling improved reasoning capability. our results show that co-llm is especially useful in cross-domain settings where a generalist base llm learns to invoke domain expert models and that co-llm can be effectively combined with other ensemble models, such as mixture of experts models @xcite . ## latent-variable framework for collaborative generation given a set of lms with different expertise or sizes, we propose a latent-variable framework that enables their collaboration in a cost-efficient way. the framework centers around a finetunable base model, which itself is a relatively small lm. it decides which other assistant models (which are typically larger and/or more specialized models) to use per token. when the base model calls on an assistant to generate the next token, we say it defers generation for that token. to generate a sequence of tokens (x 1 , . . . , x t ), we represent the choice of which model generates token x t as a discrete latent variable z @xmath0, 1, . . . , m }, where @xmath1, and @xmath2, . . . , m } refers to the i-th of m assistant models. we assume access to the conditional distributions p i (x t |@xmath3, @xmath4, . . . , m }, of the assistant models, 2 and full access to the base model. using these distributions, we represent the joint sequence-level likelihood as: the (learned) categorical distribution p θ models the token-level discrete decision z t . for efficiency, each z t is conditionally independent of @xmath5. the latent variable z t is designed in the same spirit as the defer variable in @xcite and classical mixtureof-experts models (jordan and jacobs, 1994) or ensemble models @xcite . unsupervised learning. in practice, the token level decisions z t are unknown, and collecting human annotation is not scalable. our latent-variable framework offers a natural way of handling the issue with unsupervised learning. in particular, we aim to optimize the following marginal likelihood, ## co-llm: learning to decode collaboratively with lms in this work, we focus on the basic case where we only have one base model and one assistant model. in this setting, the base model needs to make a binary decision of whether to generate from itself or defer to the assistant model, i.e., z @xmath6, 1}. for clarity, we use p base and p asst to denote the base and assistant model, respectively. in the rest of this section, we explain our design for the parameterization of the model selector, and the training and inference procedure of the joint model. ## experimental setup in our experiments, we fine-tune models for specific tasks and we test the models in-domain, comparing the end-task performance between co-llm and multiple single-or multi-model baselines. we test on 4 datasets ranging from instructionfollowing to solving expert problems, trying to understand when and how model collaboration can be beneficial. we investigate the collaboration between different models (e.g., between llama models of multiple scales, and between models finetuned on different domains). overall, we find that co-llm can learn a successful collaboration between different base and reference models, leading to better results than tuning base models alone. ## related work learning to compose models. composing expert models has been a recurring strategy to improve large models, and the way of composing expert models is largely related to the underlying learning settings. mixture of experts @xcite @xcite requires that all experts are trained simultaneously using the same data. in contrast, post-hoc methods assume the presence of pretrained language models but usually pose implicit constraints on the model families to compose. proxy tuning @xcite works best when all models are all pretrained on the same data mixture; phatgoose @xcite requires that all models are lora-finetuned (hu et al., 2021) models; contrastive decoding @xcite requires an amateur model, which is not clear for tasks such as math reasoning. co-llm is more flexible, mainly because our base model learns to interact with assistant models instead of relying on prescribed strategies. our experiments also indicate that proxy tuning (concurrent work) only performs well when all models are pretrained in the same domain, whereas co-llm can effectively combine general and domain-specific models. compared to cd and proxy tuning, co-llm also makes fewer calls to the language models at inference time, as described in section 4.1. perhaps most similar to our work are speculative decoding @xcite and comblm @xcite . like our approach, speculative decoding generates some tokens with one model and some tokens with another, but our method differs in that the goal is to improve generation quality rather than to sample more quickly from a large model by approximating its generations with a smaller one. comblm also learns to combine the probability distributions of multiple models, but their combination is not time-varying, and they mainly demonstrate wins on perplexity. our approach could be seen as a special case of toolformer @xcite , where the assistant model is the tool/api called by the base model. however, our latent variable formulation allows fine-grained deferral decisions to evolve over the course of fine-tuning, whereas toolformer derives a fixed dataset prescribing which tool gets called (and where) before training. co-llm enables varying the frequency of calling the assistant model, whereas toolformer has no provision for flexibly adjusting the amount of tool use at inference time. is this just mixture of experts? in mixture of experts (moe) llms @xcite @xcite @xcite , the goal is to train a model that can be partially executed during inference time. a common choice in moes is to decompose ffn layers (feed forward network) in a transformer into modular "experts". these experts are subnetworks-they cannot be used standalone, and the experts are typically expected to have the same size (parameter count). moe also requires gradient access to all experts and assumes every expert has access to the same training data. in contrast, we aim to learn how to collaborate with existing off-the-shelf models. first, we do not assume gradient access to assistant models, and we only fine-tune the base model. this saves in training cost, but poses a more difficult learning problem, since we only assume access to logprobs from the assistant model. co-llm is able to learn effective collaboration patterns between models despite this limited access to the assistants. second, because of our less restrictive assumptions on assistant model access and architecture, co-llm can be applied in more constrained scenarios. for example, co-llm can still be applied when the assistant is trained on proprietary data (e.g., due to copyright issues @xcite ). third, our framing allows flexible collaboration between models of different sizes or even architectures. ## conclusion we propose co-llm, a latent variable approach that learns to collaboratively decode using multiple lms. co-llm can interleave token generations from different lms in patterns most suitable for the task at hand (e.g., scaffolding the generation, or using the other model like an api), and learns the collaboration pattern organically from data. we empirically show that co-llm can produce generations of better quality compared to using either of the models alone, in tasks ranging from math reasoning, medical question answering, and instruction following. the latent "defer" variable offers a flexible and interpretable way for adjusting the frequency for invoking other lms at inference time without re-training. in the future, we plan to extend co-llm to integrate more than two lms and investigate potentially more complex collaboration strategies emerging from this setting.