Don't Throw Away Your Pretrained Model

TL;DR

This work tackles the alignment tradeoffs in language models by proposing model collaboration through Switch Generation, where a dedicated switcher LM dynamically assigns next-segment generation to pretrained, finetuned, or aligned checkpoints. By training on outcomes of candidate-segment generation and using patch-level routing, Switch Generation achieves strong improvements over individual checkpoints and baselines, with an average 12.9% gain across 18 datasets and 13 datasets where it dominates. The approach generalizes to unseen models and tasks and can distill its collaborative gains back into a single aligned model, reducing inference costs while preserving strengths. Overall, the work demonstrates the value of reusing byproducts from the model development lifecycle and provides a practical framework for compositional, adaptable AI systems.

Abstract

Alignment training has tradeoffs: it helps language models (LMs) gain in reasoning and instruction following but might lose out on skills such as creativity and calibration, where unaligned base models are better at. We aim to make the best of both worlds through model collaboration, where different models in the training pipeline collaborate and complement each other. Since LM responses feature interleaving skills that favor different models, we propose Switch Generation, where pretrained and aligned model versions take turns to ``speak'' in a response sequence. Specifically, we train a switcher LM by learning from outcomes of choosing different models to generate the next segment across diverse queries and contexts. At inference time, the switcher LM guides different model checkpoints to dynamically generate the next segment where their strengths are most needed. Extensive experiments with 8 model collaboration baselines and 18 datasets show that 1) model collaboration consistently outperforms individual models on 16 out of 18 tasks, and 2) Switch Generation further outperforms baselines by 12.9% on average. Further analysis reveals that Switch Generation discovers compositional skills to solve problems where individual models struggle and generalizes to unseen models and tasks, reusing and repurposing by-products in expensive model training pipelines that are otherwise discarded.

Figures (8)

• Figure 1: Model responses are not monolithic: they switch across diverse skills which favor different model checkpoints in the training pipeline, thus we introduce model-guided collaborative inference to optimally use models with diverse skills for different segments of response generation.
• Figure 2: Overview of Switch Generation, where multiple model checkpoints in the training pipeline (e.g., pretrained, finetuned, and aligned LM checkpoints) are dynamically selected to generate text segments in a sequence. (Up) We derive training data for the switcher LM $f$ by rolling out which model would lead to the best average outcome for a particular query and trace. (Down) At inference time, multiple models are guided by the trained switcher LM to generate text segments as part of a response when their skills and strengths are most needed.
• Figure 3: Distilling the collaboration patterns of Switch Generation back into the aligned model. Distillation recovers 58% of the collaboration gains with only one fourth of the inference cost.
• Figure 4: Correlation between the performance and helpfulness of the pretrained model. While not being the best individual model, it is consistently helpful in the model collaboration system.
• Figure 5: Frequency and treatment effect of 2-length (left) and 3-length (middle) switching sequences, as well as their correlation across three task categories for 3-length sequences (right). Switch Generation learns to identify helpful switching patterns and frequently leverages them.