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, domainspecific QA, and reasoning tasks, we show that the performance of the joint system exceeds that of the individual models. Through qualitative analysis of the learned latent decisions, we show models trained with our method exhibit several interesting collaboration patterns, e.g., template-filling.1

heading: 1 Introduction
text: Techniques that combine the generations of multiple large language models (LLMs) at decoding time have benefits ranging from faster decoding speed (Leviathan et al., 2023), to more controllable generations (Liu et al., 2021; Yang and Klein, 2021), to more coherent, less repetitive text (Li et al., 2023a), and even enabling a large model to be “tuned” by combining its generations with those of a smaller model from the same family (Liu et al., 2024). 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 (Mialon et al., 2023). Language models are able to produce more
1Code: https://github.com/clinicalml/co-llm
faithful and accurate generations when equipped with external APIs (Schick et al., 2023; Qin et al., 2023; i.a.), search engines or retrievers (Izacard et al., 2022; Asai et al., 2023; Nakano et al., 2021; i.a.), or code executors (Gao et al., 2023; i.a.).
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
12974
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.
where the models learn to interleave their generations token-by-token. Each token is generated by one model, so the models collaborate to generate a token sequence together. We represent the decision of which LLM generates the next token as a latent variable, assuming no direct supervision on the decision of which model to use at each decoding step. This enables an effective collaboration pattern for a given task to be learned organically from data.
Figure 1 shows example generations from our method, Co-LLM. In the top example, LLAMA7B 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, Azerbayev et al., 2023) 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 (Wei et al., 2022), 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 (Shazeer et al., 2017).

heading: 2 Latent-Variable Framework for Collaborative Generation
text: 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 (X1, . . . , XT ), we represent the choice of which model generates token Xt as a discrete latent variable Zt ∈ {0, 1, . . . ,M}, where i = 0 denotes the base model, and i ∈ {1, . . . ,M} refers to the i-th of M assistant models. We assume access to the conditional distributions Pi(Xt|X<t), i ∈ {1, . . . ,M}, of the assistant models,2 and full access to the base model. Using these distributions, we represent the
2Conditionals can be obtained via either locally deployed models or remote API calls. At training time, we only require access to the assistant model’s probability for the ground-truth tokens, not the full conditional distribution.
joint sequence-level likelihood as:
P (X,Z) =
T∏
t=1
( Pθ(Zt|X<t)PZt (Xt|X<t) ) (1)
The (learned) categorical distribution Pθ models the token-level discrete decision Zt. For efficiency, each Zt is conditionally independent of Z<t given X<t in our design. The latent variable Zt is designed in the same spirit as the defer variable in Mozannar and Sontag (2020) and classical Mixtureof-Experts models (Jordan and Jacobs, 1994) or ensemble models (Saunders et al., 2019).
Unsupervised Learning. In practice, the token level decisions Zt 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,
P (X) =
T∏
t=1
( M∑
Zt=0
Pθ(Zt|X<t)PZt (Xt|X<t) ) , (2)
which can be computed efficiently during training due to the conditional independence structure.
Collaborative Decoding. During inference time, our goal is to find the best sequence X along with the best decision Z on which assistant LM to use.
X̂, Ẑ = argmax X,Z P (X,Z) (3)
The exact argmax in Eq. (3) is intractable, so we follow the common practice of using the greedy strategy for decoding both Zt and Xt in a token-by-token, autoregressive manner (see Fig. 2 for an example). In greedy decoding, for each token position, we first choose Ẑt = argmaxZt Pθ(Zt|X<t) to determine which model to decode from, then decode greedily from that model: X̂t = argmaxXt PẐt (Xt|X<t). Compared with standard greedy decoding for a single LM, decoding in this model is performed in collaboration, coordinated by Pθ. An alternative strategy to argmax decoding is to marginalize out Zt:
X̂t = argmax Xt
∑
Zt
P Zt (Xt|X<t)Pθ(Zt|X<t)
which closely aligns with the marginal likelihood training objective. However, this requires calling all M + 1 models every token, slowing down decoding. Our empirical results show that greedily choosing Zt based on Pθ(Zt|X<t) performs well and enables interpretable collaboration, since each token is generated by a single model.
Remark. The design of the collaborative decoding is natural from the probabilistic modeling perspective, and in this work we mainly focus on important empirical questions: how should we parameterize Pθ so that it can be learned in a dataefficient way; can the base model learn how to cooperate with (larger or domain-specific) assistant models with only access to conditional probabilities (no internal hidden states or weights); and what kind of collaboration pattern can be induced if the collaboration is learned without supervision? The latent-variable framework allows us to answer these questions via the exposed interpretable variable Zt.

heading: 3 Co-LLM: Learning to Decode Collaboratively with LMs
text: 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., Zt ∈ {0, 1}. For clarity, we use Pbase and Passt 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.

heading: 3.1 Modeling Pθ(Zt|X<t)
text: Since we have full access to the base model, we build Pθ on top of the base model for parameter efficiency. Specifically, we represent θ as a linear binary classification head in the last layer. Formally, if ht(X<t) ∈ Rd is the base model’s last hidden state at time step t for inputs X<t, and θ ∈ Rd is the weight vector, we set:
Pθ(Zt|X<t) = σ(⟨θ, ht(X<t⟩),
where σ(·) is the sigmoid function. This introduces only d new parameters to the base model, where d is the base model’s hidden dimension size.

heading: 3.2 Training
text: Our training objective minimizes the negative logmarginal likelihood −∑Tt=1 logP (Xt|X<t), where
P (Xt|X<t) =Pbase(Xt|X<t)Pθ(Zt = 0|X<t)+ Passt(Xt|X<t)Pθ(Zt = 1|X<t) (4)
is the likelihood of the next token after marginalizing out the latent variable Zt. The training procedure updates both θ and the base model parameters; it requires the forward pass of both the base
and assistant model to obtain next-token probabilities, and the backward pass of the base model for gradient computation. This implies that any commercial LLMs or locally deployed LMs exposing next-probability output can be used as the assistant LMs. The marginal likelihood aligns well with the typical pretrained objective of maximizing the nexttoken probs. Moreover, offloading “difficult tokens” can potentially alleviate hallucination issues of the base model, thus leading to better generalization even without much help from the assistant model, as we will show in the experiments.
Initialization of θ. In our experiments, we found that appropriately initializing θ helps encourage the base LM to quickly switch from generating by itself to generating collaboratively. Instead of collecting direct supervision for Zt values to initialize θ, we use weak supervision to collect initial Zt values, use these pseudolabels to initialize the parameters θ, then allow the Zt’s to change during the training procedure.
Intuitively, the prime location for Zt = 1 is when the assistant model correctly predicts the target token Xt but the base model does not. To that end, we set the pseudolabels for Zt as:
Ẑt := 1[Xt = argmaxv∈V Passt(v|X<t)∧ Xt ̸= argmaxv∈V Pbase(v|X<t)], (5)
and initialize the d parameters θ by maximizing the likelihood of logPθ(Ẑt|X<t) while holding the rest of the base model fixed. In our experiments, we compare to a baseline that combines this initialization with the usual language-model fine-tuning loss; our results show that the marginal likelihood objective leads to better performance by enabling the base model to learn better Zt values from data (see ablation results in Section 5.2).

heading: 3.3 Decoding
text: In initial experiments, we found that performance of the joint model is sensitive to the choice of Zt; due to the exposure of Zt in our latent-variable framework, we can impose extra priors over Zt for better performance. Specifically, we follow the general greedy decoding strategy, and set a threshold η for decoding Zt:
Ẑt = 1[Pθ(Zt = 1|X<t) > η], (6)
which means that when Pθ(Zt = 1) > η, we execute the assistant model to predict the next token.
The hyperparameter η is picked via grid search on a small validation set per dataset. The choice of threshold for the decoding probability also allows for inference-time control over the amount of collaboration, in contrast with other approaches such as DExperts (Liu et al., 2021), and our performance degrades gracefully as the threshold increases.

heading: 4 Experimental Setup
text: In our experiments, we fine-tune models for specific tasks and we test the models in-domain, comparing the end-task performance between CoLLM 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.
Models Used. Our primary experiments are concerned with whether smaller models can collaborate with expert models that have been specialized to different domains. In Section 5.1, we experiment with collaboration between the finetuned LLAMA-7B and the LLEMMA family (Azerbayev et al., 2023, finetuned for math and reasoning), as well as the MEDITRON family (Chen et al., 2023, finetuned for biomedicine). It is also possible to use differently-sized models from the same family with Co-LLM: in Section 5.2 we experiment with the 7B and the 70B model from the same LLAMA2 family (Touvron et al., 2023) as the base and assistant model, respectively.
Datasets. We train on the full Tülu v2 mix data (Wang et al., 2023) for instruction following, GSM8k (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021), each with 7.5k samples, for reasoning and math problem solving, and the BioASQ (Tsatsaronis et al., 2015) (4.7k samples) for medical question answering. We train and test the model on the corresponding data and evaluation suites separately.
Evaluation. We only compare greedy decoding results for all models, as they are commonly used in real-world applications. We evaluate the instruction following model using the AlpacaEval (Li
et al., 2023b) benchmark; we use the GPT-4 annotator and compute the win rate of the testing model as judged by GPT-4-0613 when compared to the outputs produced by Davinci-003. For GSM8k, following Wang et al. (2023),we extract the last numerical answer from the model’s output, and calculate the exact match of the model prediction on a 200-example subset of the test set. The MATH dataset provides 5 levels of math problems from 7 categories: similar to Wu et al. (2023), we skip the geometry category since it largely relies on Asymptote code, and sample 25 questions per level per category for evaluation, resulting in a 750-example subset. We adopt the prompting and evaluation code from Azerbayev et al. (2023) and Lewkowycz et al. (2022), extracting the last generated math after “The answer is”. 3 The BioASQ comes with 330 test examples of 4 categories: factoid, list, summary, and yes/no questions, evaluated using strict accuracy (SAcc.), F1, accuracy (Acc.), and Rouge2 (R2). We test on 310 examples (saving 20 for validation) and reimplement the evaluation code.4

heading: 4.1 Baselines
text: Single models. The performance of the base and assistant models can inform whether the learned collaboration is beneficial. We report 0-shot performance of the original untuned models and their finetuned counterparts. The same data and hyperparameters are used for model finetuning; for 70B models, we fine-tune using QLoRA (Dettmers et al., 2023) with the hyperparameters in Appendix A.
Other collaborative models. We use two collaborative strategies from the literature. Contrastive Decoding (Li et al., 2022b; O’Brien and Lewis, 2023, CD) combines the output of the untuned “expert” (e.g., a 70B model) and “amateur” (e.g., a 7B model) models by subtracting their logits and sampling from the resulting distribution. We follow the setup in O’Brien and Lewis (2023), setting the same α = 0.1 and β = 0.5, and use unmodified LLAMA-70B as the expert model and LLAMA-7B as the amateur model.5 Proxy Tuning (Liu et al., 2024, PT) proposes to approximate finetuning a (large) base model M by composing the outputs
3We use greedy decoding for GSM8k and MATH; some literature refers to this as “maj@1”.
4See details in Appendix A. 5In their paper, O’Brien and Lewis (2023) use a 1.5B parameter amateur model; as this model is not released, we use the 7B model instead. Different from their paper, we use 0- or 1-shot prompting in congruence with our other results.
of smaller expert M+ and anti-expert models M− with M. We include CD and PT results as ways to enhance untuned models and to simulate finetuning 70B models, respectively. Both CD/PT require calling the smaller and larger models at each time step. In contrast, Co-LLM may generate multiple tokens from the base model before calling the large model, which can be seen as a form of speculative decoding (Leviathan et al., 2023). For example, for a sequence of length L, Proxy Tuning makes L calls6 to the large model and 2L calls to the small model, whereas Co-LLM makes fL calls to the large (assistant) model and L calls to the small (base) model, where 0 ≤ f ≤ 1 is the empirical frequency of deferral, the percent of Zt = 1 tokens.
Ablated Co-LLM. We consider different variants of Co-LLM to verify the necessity of a learned model selector Pθ. First, we consider two simple heuristics as model selectors: Co-LLM-Random randomly chooses the base or the assistant model to produce a token with probability p = 0.5; Co-LLMGreedy runs both models in parallel for each token and selects the token with the higher probability from either model. This is a strong baseline since it requires observing next-token probabilities from both models at every decoding step. Similar to our default setting, the base model is finetuned on target datasets while the assistant model is frozen.
Weakly-supervised Co-LLM. Finally, we consider a different weakly supervised training procedure for Co-LLM. This baseline is inspired by the process used to derive tool-use labels in Toolformer (Schick et al., 2023): the weak supervision for when to call the tool is chosen and fixed before updating the language model parameters. Specifically, the training procedure is two-stage: first, we collect pseudo-labels Ẑt using Eq. (5). Second, we jointly train Pθ(Zt|X<t) and the base model by optimizing the weighted sum of logP (Ẑt|X<t) and the usual language modeling loss. This trains the base model to defer to the assistant when Ẑt = 1 while also fine-tuning it on the full training set. In contrast, our marginal likelihood training Section 3.2 only uses the Ẑt values to initialize θ and allows the Zt values to evolve during training.
6Here each “call” corresponds to a token-level decoding step, which is a sensible unit with which to measure inference latency as the context encoding portion can be easily parallelized and thus grows slowly with respect to context length in small batch regimes.

heading: 5 Results

heading: 5.1 Collaboration across domains
text: Table 1 shows that Co-LLM enables collaboration between LLAMA and domain-specific models and that this collaboration improves performance compared to the individual models themselves. For example, being able to utilize LLEMMA as an assistant leads to improved performances on math and reasoning tasks. On the MATH dataset, even invoking a small, 7B-scale LLEMMA assistant (17.2) outperforms fine-tuned LLAMA-7B (7.6), fine-tuned LLAMA-70B (11.7), and LLEMMA-34B (6.3). Similarly, cooperation with MEDITRON models leads to performance gains on some BioASQ subtasks (e.g., List, Summary), and outperforms fine-tuned LLAMA-7B, fine-tuned LLAMA-70B, and base MEDITRON-70B on average.
In addition, Co-LLM with LLAMA-7B and LLEMMA-34B can achieve similar performance as fine-tuned LLEMMA-7B, which scores 43.5 on GSM8k and 23.5 on MATH. Co-LLM allows the base 7b model access to collaborate with a domain expert (LLEMMA-34B), which surprisingly leads to similar performance as performing a large amount of domain-specific fine-tuning plus further taskspecific fine-tuning on the base model (finetuned LLEMMA-7b).
These results suggest that Co-LLM enables a modular approach to continued pretraining and task-specific finetuning: one can pretrain a large model on a domain-specific corpus, then fine-tune smaller models with Co-LLM to leverage the knowledge from the larger models and attain improved performance on the downstream tasks.
Comparison against Proxy Tuning. While PT and our work are differently motivated and con-
structed, they both leverage multiple models during generation. Table 1 also provides an in-depth comparison between the two methods in the context of combining models from different domains. PT only performs well when all three models (M,M+, M−) are pretrained on the same domain mix (compare, e.g. “LLEMMA + LLAMA” to “LLEMMA + LLEMMA”). This is due to the implicit assumption that the difference between the base model M and a hypothetical, tuned version of the base model is the same as the difference between the smaller expert M+ and the anti-expert M−. Our results show that Co-LLM is more effective at enabling collaboration between models from different domains. PT also requires more calls to the larger model, thus resulting in slower inference. Co-LLM makes fewer calls to both large and small models.
In the following section, we show that in addition to enabling collaboration across domains, Co-LLM also allows collaboration across model scales.

heading: 5.2 Collaboration across scales
text: Table 2 shows that using Co-LLM leads to a successful collaboration between the base and assistant models of different sizes within the LLAMA family. In particular, compared to using the (finetuned) base model LLAMA-7B alone, Co-LLM7B + LLAMA-70B, which occasionally calls the unmodified assistant model during decoding, consistently achieves significantly better performance across all tasks and datasets (2.6, 10.5, 7.5, and 3.3 absolute improvements for the 4 datasets, respectively). Co-LLM is sometimes better than the QLoRA finetuned assistant model (on MATH and BioASQ), suggesting that our method can effectively combine the best of the models and achieve better performance than the “sum” of their parts. Training with Co-LLM does not hurt the perfor-
mance of the base model: when we prohibit using the assistant model during inference, performance is comparable to the base model finetuned with the usual language modeling objective (LLAMA-7B (Finetuned)): for example, getting 33.0 and 6.4 for GSM8k and MATH, respectively. Co-LLM thus degrades gracefully as the amount of deferral is changed, which we explore further in §5.4.
Comparing with the model collaboration baselines, we show that interleaving two model generations is not a trivial task: randomly switching between the two models lead to worse than singlemodel performances in Co-LLM-Random. Even if running two models in parallel, Co-LLM-Greedy does not consistently yield better performance than using either model alone, and in some cases, they are worse (e.g., in the case of GSM8k). As discussed in Section 5.1, Proxy Tuning performs very well when used with models of different sizes in the same family, but Co-LLM also performs well despite using far fewer calls to the language models.
Finally, the Toolformer-style baseline with fixed, weakly supervised values of Zt leads to overall worse performance compared to Co-LLM, indicating the benefits of our latent variable formulation and marginal likelihood loss, which allow the best deferral patterns to be learned from data.

heading: 5.3 Collaboration across architectures
text: Since Co-LLM only assumes access to token logprobs, it can easily be used for collaboration between models of different architectures. For example, Table 3 shows that Co-LLM can be adopted to collaborate between a dense model (MISTRAL-7B) and a sparse MoE model (MIXTRAL-8×7B), and the joint model achieves strong accuracy gains compared to either the finetuned MISTRAL-7B model or the MIXTRAL-8×7B model. These results indicate that Co-LLM is still useful when used together with other model combination methods like MoE.

heading: 5.4 Qualitative Analysis
text: The exposure of the interpretable variable Zt in our latent-variable framework makes it easy to visualize the model composition patterns. As shown in Figs. 1 and 2, our unsupervised training can induce an interesting template-filling collaboration strategy: the base model learns to generate a template with slots that requires deeper knowledge (in question answering) or reasoning (in math problems) for assistant models to fill in. Table 4 further illustrates how the collaboration between the assistant
7See Section 5.4 for a detailed analysis of the frequency of invoking the assistant model and the end performance.
(LLEMMA-34B) and base (fine-tuned LLAMA-7B) model leads to correct responses while either model alone fails to do so. When the base model is tasked to solve the math question alone (i.e, deferral frequency f = 0), it fails to produce a valid answer, generating 125 rather than 3125. As we increase the frequency of deferral by lowering the threshold η, we observe Co-LLM starts to invoke the assistant model to generate latex code and compute the results, and when f = 0.4, the joint model produces the correct answer. More deferral in this case does not yield better generations. When we fully rely on the assistant model (f = 1), since it is not tuned or aligned, it produces no helpful solutions.
We also evaluate the joint model at different deferral frequencies on small validation sets for GSM8k, MATH, and BioASQ, and plot the results in Fig. 3 and Fig. 4 in the Appendix. Across different domains and scales, the models exhibit a similar concave performance curve: at some optimal deferral frequencies, the joint model can achieve
better performance than using either of the models alone, with decreased performance at both extremes. The optima vary across different datasets, corresponding to the different patterns of invoking assistant models (e.g., API-call or “leading” style). In practice, one can pick the proper η to balance the accuracy and the efficiency/cost.

heading: 6 Related Work
text: 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 (Shazeer et al., 2017; Jiang et al., 2024; Dai et al., 2024; Xue et al., 2024, MoE) 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 (Liu et al., 2024, 2021) works best when all models are all pretrained on the same data mixture; PHATGOOSE (Muqeeth et al., 2024) requires that all models are LoRA-finetuned (Hu et al., 2021) models; Contrastive Decoding (Li et al., 2022b, CD) 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, CoLLM 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 (Leviathan et al., 2023) and CombLM (Ormazabal et al., 2023). 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 (Schick et al., 2023), 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 (Zhou et al., 2022; Li et al., 2022a; Jiang et al., 2024; Xue et al., 2024; Dai et al., 2024), 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 (Duetting et al., 2023)). Third, our framing
allows flexible collaboration between models of different sizes or even architectures. Table 3 shows performance gains from combining a dense model with an MoE model using Co-LLM, indicating that the gains are orthogonal to MoE. Finally, in our approach, each “expert” is a full-fledged LLM rather than a sub-network. This allows for interpretable generation patterns (as illustrated in Figs. 1 and 2), as well as flexible inference-time tradeoffs between how much each model is used (e.g., in Table 4, we show Co-LLM can work well at different frequencies of deferral).
Learning to defer. A large body of literature focuses on the related problems of prediction with rejection (Cortes et al., 2016), where the goal is to train a model that can predict on some inputs and decline to predict on others, and learning with deferral (Mozannar and Sontag, 2020; Mozannar et al., 2023), where the goal is to train a model to make predictions on some inputs and defer to a human expert on others. Mohri et al. (2023) combine prediction with example-level rejection and LLMs on the text decontextualization problem. We use a latent variable formulation inspired by Mozannar et al. (2023), initially developed for learning to defer to a human expert on classification problems. We replace the human expert with a fixed LLM assistant model and extend the loss to allow for token-level rather than sequence-level deferral.

heading: 7 Conclusion
text: 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.
Limitations
As shown in Fig. 4 and Table 6 of the Appendix, the optimal deferral threshold may be different across different datasets and models. Co-LLM thus requires picking the deferral threshold per task, which can be inconvenient in practice. However, the threshold also enables inference-time control over the amount of collaboration. Second, not every deferral matters: for some positions, the assistant model may generate an identical token as the base model does. This suggests the development of more fine-grained control of deferral strategies, potentially via more sophisticated modeling of deferral model parameters θ.
Another limitation in our method comes from fully relying on an assistant model at some point in decoding. For example, if the assistant model is not well-tuned or aligned, it may unintentionally break the generation due to occasional mistakes. As shown in the example below, one erroneous token might lead to a cascade of errors, causing repetition patterns or generating irrelevant content. One future work is to develop a more robust deferral strategy that allows backtracking when the assistant model fails to generate a proper response. Here’s a recipe for Kubdari, a traditional Georgian dish: Ingredients: * 1 lb ground beef * 1 onion, finely cho pped * 2 cloves garlic, minced * 1 cup chopped parsley * 1 cup chopped cilantro * 1 cup chopped dill * 1 cup chopped ...
[...repeating the same pattern...]

heading: Acknowledgements
text: SS and DS were supported by the National Science Foundation (NSF award no. IIS-2205320 Conceptualizing ML for Dynamic Information Retrieval of EHR Notes). HL was funded by an NDSEG fellowship. BW and YK were partially supported by MIT-IBM Watson AI and an Amazon award. We thank CloudBank (Norman et al., 2021) for supplying computational resources, which is supported by the National Science Foundation under award #1925001. And thanks to Jacob Andreas, Lucas Torroba Hennigen, Hussein Mozannar, Ilker Demirel, Andreas Haupt, Tiwalayo Eisape, Alex Gu, Ruochen Zhang, and Doug Downey for feedback on the draft of this paper.

heading: A Additional Experimental Details
text: A.1 Model Training and Computation
Training details. In our experiments, we finetune the base model and learn the latent variable parameters θ jointly: i.e., we optimize both the “base” and θ parameters by optimizing the marginal likelihood from Eq. (4) with AdamW (Loshchilov and Hutter, 2018). We train our models primarily using 4 A100 80G GPUs, and we use FlashAttention (Dao, 2023) and DeepSpeed ZeRO Stage 2 (Rasley et al., 2020) during training to reduce the GPU memory usage.
We follow Ivison et al. (2023) and Liu et al. (2024), using similar hyperparameters and settings detailed in Table 5. In all training experiments, we compute the marginal likelihood loss on only target tokens (i.e., we mask out the input tokens). It takes around 2 hours to finish the finetuning experiments for GSM8k, and we estimate a total of 3,000 GPU hours used for all experiments.
Datasets and prompts. Table 8 details the sizes of the training and evaluation datasets. We format the data using the same simple prompts during training and evaluation, with examples shown in Table 11.
Model and data licenses. We use three different LLMs in our experiments, i.e., LLAMA (Touvron et al., 2023), LLEMMA (Azerbayev et al., 2023), and MEDITRON (Chen et al., 2023), and they all share the same LLaMA 2 community license.8 The licenses for the dataset are listed in Table 9.
Threshold search. For all methods with a learnable model selector (including Co-LLM), we choose the best η for Eq. (6) using a small validation set before testing, described in Algorithm 1. For GSM8k, MATH, and BioASQ, we conduct the search in-domain, using a small hold-out subset for that dataset; for AlpacaEval, we use a subset of a separate instruction following benchmark MTBench (Zheng et al., 2023) to search and pick the η∗. We report the used values under different settings in Table 6 (for Co-LLM) and Table 6 (for Weakly-supervised Co-LLM).
A.2 BioASQ Evaluation
There are four different substasks in the BioASQ dataset, and they are evaluated using different met-
8https://ai.meta.com/llama/license/.
Algorithm 1: Find Optimal Deferral Th. η Input :Base Model, Asst. Model, Model
Selector ϕ, Validation Dataset D Let P = {} for i = 0 to |D| do
Given the input prompt in Di, generate response X(i) using the base model Predict per-token deferral probability Pϕ(Z (i) t = 1 | X(i)<t), and append it to P
end Sort P in ascending order Set current best threshold η = 0 and
evaluation score s = 0 for j = 0 to 100 by 10 do
Get the j-th quantile pj in P and use it as the deferral threshold ηj Generate responses X(i) for i in from 0 to |D| using the base and asst. model controlled by Pϕ(Z(i)t = 1 | X(i)<t) > ηj Evaluate the responses, and if the evaluation score is better than s, set η = ηj and s to the new score
end Return η
rics according to the guideline:9
• Factoid: It require a particular entity name (e.g., of a disease, drug, or gene, or a number) to be generated per the question, which might not appera in the original question. The model may generate a list of candidate, and we pick the first generation (as the model often only generates one) and search for matching among the allowed candidate answers. This is the Strict Accuracy (SAcc.) metric in Krithara et al. (2023). • List: Similar to Factoid questions but the model is expected to generate a list of entities. The model is required to produce the answer in a bullet list format, and we use F1 score to evaluate the performance. • Summary: The answer is expected to be a longform text like the description of a treatment. Following Krithara et al. (2023), we report ROUGE2 (Lin, 2004) (using the implementation from the Huggingface Evaluate library10) to measure the
9http://participants-area.bioasq.org/general_ information/Task9b/
10https://huggingface.co/spaces/ evaluate-metric/rouge
textual overlapping between the generation and the ground truth. • Yes/No: The model needs to provide a binary
answer yes or no to the given question, and the classification accuracy is reported.

heading: B Additional Generation Examples
text: Fig. 2 shows an simplified version of generation for clarity, and we show the original generation in Table 10. An additional example generation on MATH is shown in Table 12.
Dataset Prompt
AlpacaEval 
 What are some species of bears that are now extinct? 

Please solve the following math problem with detailed steps.

heading: Future Work
text: 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. Limitations
As shown in Fig. 4 and Table 6 of the Appendix, the optimal deferral threshold may be different across different datasets and models. Co-LLM thus requires picking the deferral threshold per task, which can be inconvenient in practice. However, the threshold also enables inference-time control over the amount of collaboration. Second, not every deferral matters: for some positions, the assistant model may generate an identical token as the base model does. This suggests the development of more fine-grained control of deferral strategies, potentially via more sophisticated modeling of deferral model parameters θ. Another limitation in our method comes from fully relying on an assistant model at some point in decoding. For example, if the assistant model is not well-tuned or aligned, it may unintentionally break the generation due to occasional mistakes. As shown in the example below, one erroneous token might lead to a cascade of errors, causing repetition patterns or generating irrelevant content. One future work is to develop a more robust deferral strategy that allows backtracking when the assistant model fails to generate a proper response. Here’s a recipe for Kubdari, a traditional Georgian dish: Ingredients: * 1 lb ground beef * 1 onion, finely cho pped * 2 cloves garlic, minced * 1 cup chopped parsley * 1 cup chopped cilantro * 1 cup chopped dill * 1 cup chopped. [. repeating the same pattern. ]