Speculate, then Collaborate: Fusing Knowledge of Language Models during Decoding

We greatly appreciate the reviewer's recognition of the advantage of CoSD over the existing frameworks. Regarding the weaknesses raised, we address all the concerns in detail below.

W1: I'm not sure if the goal of this algorithm is to A) achieve performance comparable to the assistant model but in a more efficient way, or if it's aimed at B) outperforming both the draft model and the assistant model individually (1+1>2). How do these objectives apply to the four scenarios of knowledge fusion discussed in section 4.1? If the goal is A, since the draft models in complementary knowledge fusion and catastrophic forgetting recovery scenarios are about the same size as the assistant model, and the algorithm involves the autoregressive generation of the draft model, I doubt the algorithm improves efficiency. If the goal is B, I can't see improvement based on Table 2.

We thank the reviewer for clearly pointing out goals A and B. In fact, both A and B are our goals. We expect our algorithm to achieve different objectives when applied to different tasks and models. Specifically:

(1) When the assistant model has a much larger parameter size than the draft model, we expect CoSD to achieve goal A, which achieves performance comparable to the assistant model but in a more efficient way. This is also the goal of speculative decoding, and we expect our CoSD algorithm to retain this functionality while achieving results closer to the assistant model. In our experiments, pair 4 in Table 2 shows the effectiveness of CoSD in this scenario:

We display the average score across all 3 benchmarks in this table. It shows that using a 1b draft model and a 7b assistant model in CoSD inference can achieve similar performance to the single 7b model (24.39 -- 25.16) and be 2 times faster.

(2) When the two models have similar sizes and complementary knowledge, we expect CoSD to have higher averaging performance across all the tasks. Pair 2 and pair 3 in Table 2 show the effectiveness of CoSD in this scenario:

CoSD-Rule outperforms both the draft and the assistant model in these two pairs and has significant improvement when the areas of expertise of the two models are entirely distinct (pair 3).

(3) When the two models are of similar size but one is significantly stronger overall, CoSD can achieve the performance level of the stronger model but cannot save computational costs. However, considering that in real-world applications, we cannot predict the performance of different models across all tasks, we believe it is still worthwhile to attempt fusing the knowledge of two similarly sized models to ensure that the combined performance is close to the better model across various domains. For instance, pair 1 in Table 2 shows this scenario:

CoSD still outperforms all the baselines.

Here, we want to emphasize that the superiority of our approach lies in its ability to consistently achieve excellent cross-domain performance, regardless of the type or performance of the given models. By using the CoSD algorithm, we eliminate the need for users to evaluate and select models or consider goals like A and B. Instead, the system can automatically adapt to either efficient inference tasks or knowledge fusion tasks. This represents a significant advantage and contribution compared to previous works focused solely on single objectives like A: efficient inference (e.g., speculative decoding) or B: knowledge integration (e.g., CoLLM).

Thanks for clarifying these points. I'd suggest directly highlighting them in the next version. I will keep my score.

We greatly appreciate the reviewer's recognition of the advantage of CoSD over the existing frameworks. Regarding the weaknesses raised, we address all the concerns in detail below.

W1: The paper only does the experiment in each pair of the LLMs. It would be interesting to see more LLMs collaboratively fuse knowledge.

Thank you for your valuable comment! Our proposed algorithm indeed supports collaboration among multiple LLMs, including scenarios involving three or more models. As an example, in a three-model CoSD, one draft model generates the draft and 2 assistant models verify the draft. The rule-based verification process will be:

$x_{t+i} \neq x_{t+i}^{0}$ and $x_{t+i} \neq x_{t+i}^{1},$

$M_p(x_{t+i}) < \alpha,$

max($M_q(x_{t+i}^{0}$), $M_q(x_{t+i}^{1})) > \beta \cdot M_p(x_{t+i}),$

The assistant token with a higher probability will replace the draft token if all the 3 conditions are met.

For the tree-based CoSD with $x$ LLMs, we extend the decision tree to $x$-class classification and select the model that predicts the next token with the highest probability as the ground truth label, with all other settings remaining the same as the two-LLM setting. We have conducted experiments with this three-model CoSD to validate the algorithm’s capability in fusing knowledge from multiple LLMs. The results are shown in the table below:

where Draft model = TinyLlama, Assist. 1 = Llama 2 Chat 7b, Assist. 2 = Llama-7b.

The current experimental results of this model group demonstrate that when three models collaborate if one significantly outperforms the other two, the final system will achieve performance close to that of the best model. This indicates that our algorithm is effective when applied to more than two models. We are conducting collaborative experiments with other model groups. We will keep you updated in real time if we obtain new results and conclusions. Once the experiments for all model groups are completed, we will incorporate this section into the main body of the paper.

W2: It would be better to show more details about the limitations of the proposed method and show some error analysis.

A potential limitation of our approach is that in some cases when the two collaborating models have similar parameter sizes but one is significantly more powerful than the other, it is not necessary to use CoSD since directly using the more powerful model is enough. For example, in Table 2 pair 1, the assistant 8B model has much better overall performance than the draft 8B model, and CoSD can only achieve similar performance to the assistant model, and cannot save computation cost since the models have the same parameter size.

About the error analysis, CoSD cannot guarantee the assistant token is better than the draft token it replaced. Sometimes it may drop some good draft tokens with a worse assistant token when the draft model is not confident enough. Here is an example question in MMLU:

Rowena can paint a room in $14$ hours, while Ruby can paint it in $6$ hours. If Rowena paints for $x$ hours and Ruby paints for $y$ hours, they will finish half of the painting, while if Rowena paints for $y$ hours and Ruby paints for $x$ hours they will paint the whole room. Find the ordered pair $(x,y)$. A. $(\frac{11}{10}, \frac{11}{10})$ B. $(\frac{231}{20}, \frac{21}{20})$ C. $(\frac{231}{40}, \frac{21}{40})$ D. (1,1)

The draft model gave the correct answer C with 0.32 probability but was rejected by the answer D from the assistant model with 0.51 probability. This example illustrates that low probability does not necessarily indicate an incorrect token. Therefore, replacing tokens with higher-confidence alternatives is merely a heuristic algorithm and is subject to error. After all, during the inference stage and model deployment, the ground truth is unknown.

Thank you for your feedback so far. We have provided a detailed rebuttal and an updated draft addressing your comments with additional results. Please feel free to share any further questions or thoughts before the discussion period ends.

We greatly appreciate the reviewer's recognition of the advantage of CoSD over the existing frameworks. Regarding the weaknesses raised, we address all the concerns in detail below.

W1: In Table 2, I notice that in most cases, the fused model underperforms the draft model and the assistant model. For instance, ... Then I wonder what is the point of fusing knowledge in these cases if we can simply adopt one model instead of the other.

The reviewer mentioned that in most cases, the fused model underperforms the draft model and the assistant model, we disagree with this conclusion. Here we list the averaging accuracy across all 3 benchmarks of Table 2:

For pair 2 and pair 3, CoSD outperforms both the draft model and the assistant model, and all the baselines. It means that applying CoSD can achieve better and more balanced performance across multi-domain tasks.

For pair 4, CoSD achieves comparable performance to the assistant model and outperforms all baselines. Considering that the assistant model is much larger than the draft model, CoSD can significantly improve the inference speed by 2 to 3 times.

For pair 1, CoSD has comparable performance to the assistant model. This is because the assistant model has much better performance across domains than the draft model. In this case, we agree with the reviewer that we can simply adopt the assistant model. However, in practice, it is hard to know in advance which model will perform better. CoSD alleviates the need to choose a model, providing performance superior or comparable to the best model.

We further clarify that model fusion cannot guarantee consistently achieving better performance than the best individual model in the fusion across all benchmarks. Instead, we aim to:

(1)Improve the average accuracy across all the benchmarks if the merged models have complementary knowledge.

(2)If one model performs significantly better than the other across different benchmarks, the fusion aims to approximate the performance of the better-performing model as closely as possible.

W2: It seems that for Pair 3, CoSD-Rule performs exceptionally well on GSM8K, ... Could you give more insights into such a result? Could you present some examples that CoSD-Rule excel at under this situation that cannot be addressed by either the draft nor the assistant model?

Here is an example of the generation of the draft model, the assistant model, and the CoSD-Rule for the GSM8K dataset in pair 3:

Draft:

The cost of the subs is 2 x $7.50 = <<2*7.5=15>>15. The cost of the chips is 2 x $1.50 = <<2*1.5=3>>3. The cost of the cookies is 2 x $1.00 = $<<2*1=2>>2. The total cost of the food items is $15 + $3 + $2 = $<<15+3+2=20>>20. The cost of the delivery fee is 20% of the total cost of the order, so 20% of 20 = <<20% of 20=10>>10. The total cost of the food items and delivery fee is $20 + $10 = $<<20+10=30>>30. The cost of the tip is $5.00. The total cost of the order is 15 + 3 + 2 + 10 + 5 = <<15+3+2+10+5=45>>45.

Assistant:

The cost of the subs is 2 x $7.50 = $<<2*7.5=15>>15. The cost of the chips is 2 x $1.50 = $<<2*1.5=3>>3. The cost of the cookies is 2 x $1.00 = $<<2*1=2>>2. The total cost of her order without any fees is $15 + $3 + $2 = $<<15+3+2=20>>20. The delivery fee is 20% of the total order, so $20 * 0.20 = $<<20*0.20=4>>4. The tip is an additional $5. Therefore, her delivery order will cost $15 + $3 + $2 + $4 + $5 = $<<15+3+2+4+5=35>>35.

CoSD-Rule:

The cost of the subs is 2 x $7.50 = $<<2*7.5=15>>15. The cost of the chips is 2 x $1.50 = $<<2*1.5=3>>3. The cost of the cookies is 2 x $1.00 = $<<2*1=2>>2. The total cost of the food items is $15 + $3 + $2 = $<<15+3+2=20>>20. The delivery fee is 20\% of $20, so the delivery fee is $20 x 0.20 = $<<20*0.2=4>>4. The total cost of the food items and delivery fee is $20 + $4 = $<<20+4=24>>24. The tip is $5.00, added to the total cost of the food items and delivery fee, so the total cost of the order is $24 + $5 = $<<24+5=29>>29.

Where the correct answer is 29, only CoSD-Rule gets the correct answer. We found that in this case, the draft model helps establish a better chain of thought (e.g., clearly calculating the total cost first, followed by delivery fee and tips). However, the draft model's mathematical computation ability is weak, making errors in multiple steps within the chain of thought. At this point, the assistant model ia involved to correct the computational mistakes, resulting in an excellent final outcome. We believe this situation occurs when the two models have highly complementary capabilities in different domains (e.g., pair 3 in the table of W1)