Source-linked AI summary
Don't Throw Away Your Pretrained Model
Shangbin Feng, Wenhao Yu, Yike Wang, Hongming Zhang, Yulia Tsvetkov, Dong Yu
TL;DR
Alignment can strengthen reasoning and instruction following while trading off skills such as creativity and calibration, motivating collaboration across model checkpoints. SWITCH GENERATION trains a switcher to select models for successive response segments, and experiments report that it beats collaboration baselines on 13 datasets by 12.9% on average.
Problem
Alignment training improves skills such as reasoning and instruction following but can trade off creativity and calibration, where unaligned base models may be stronger.
Source: §1 · p. 1Method
SWITCH GENERATION trains a switcher LM to choose among model checkpoints for successive response segments using the query, generated trace, and candidate models.
Source: §1 · pp. 1–2, §2 · p. 3Results
Takeaways & Limitations
The analysis reports generalization to unseen models and tasks and reuse of training-pipeline checkpoints, including distilling generated collaboration texts into one model.
Source: §1 · p. 2, §7 · p. 9Using multiple checkpoints incurs extra cost compared with using only the aligned model.
Source: LIMITATIONS · p. 15
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Alignment can improve reasoning and instruction following while trading off skills such as creativity and calibration. SWITCH GENERATION has different model checkpoints contribute successive response segments, and the paper reports gains over individual models and collaboration baselines.
1 INTRODUCTION
- SWITCH GENERATION trains a switcher LM to select which checkpoint generates each next segment, using outcomes from candidate continuations across queries and traces.At inference, the switcher guides different checkpoints to contribute successive segments.Source: §1 · pp. 1–2
- Model collaboration outperforms individual models on 16 of 18 datasets and is a close second on the other two.Source: §1 · p. 2
- 12.9% average improvement: SWITCH GENERATION outperforms all collaboration baselines on 13 datasets.The experiments span 8 model-collaboration baselines and 18 datasets.Source: §1 · p. 2
- Further analysis finds generalization to unseen tasks and models, solutions to some problems individual models cannot solve, and collaboration texts that can be distilled into one model.Source: §1 · p. 2
- The paper frames reuse of otherwise discarded training-pipeline byproducts as a promising direction for model development.Source: §1 · p. 2
2 METHODOLOGY
SWITCH GENERATION treats the choice of which model speaks next as a query-, trace-, and candidate-dependent decision. A trained switcher selects checkpoints to generate successive patches, using training examples derived from the outcomes of candidate continuations.
2 METHODOLOGY
- SWITCH GENERATION is a collaborative inference algorithm that dynamically selects checkpoints to generate successive response segments and combine their complementary strengths.Source: §2 · pp. 2–3
- The Query-Trace-Candidate problem asks which candidate model should generate the next segment, given the query and what has already been generated.The switcher receives the query, trace, and candidate pool; its repeated segment-level decisions differ from one-time full-response routing.Source: §2 · p. 3
- The switcher prompt marks which model generated each part of the trace, and the switcher predicts the model ID for the next segment.Its predicted logits form the distribution over candidate models.Source: §2 · p. 3
- Training examples extend a sampled partial-response trace with candidate segments, then use average scores across k sampled continuations to select the target model.Scores can use the query's evaluation metric; the resulting supervised examples train the switcher.Source: §2 · p. 3
- At inference, the switcher chooses a model for each fixed-size patch, reducing switching calls compared with token-by-token selection.Top-p sampling selects from the model distribution, balancing utility and exploration.Source: §2 · p. 4
3 EXPERIMENT SETTINGS
The experiments compare model collaboration across 18 datasets and a range of model settings. The results caption reports that collaboration beats individual models on most tasks, with Switch Generation best on 13 tasks.
3 EXPERIMENT SETTINGS
- The default setup uses pretrained, finetuned, and aligned Tulu-v3 models; the switcher is trained globally or per task and generates in 50-token patches by default.The switcher uses 10,000 SFT instances per task, k = 32, five epochs, and a 2e-4 learning rate.Source: §3 · p. 4
- The 11 baselines span API-, text-, logit-, and weight-level collaboration, alongside the pretrained, finetuned, and aligned models used individually.The listed methods include prompt-based routing, debate, logit merge, and greedy soup.Source: §3 · p. 4
- The 18 datasets cover skills where base models may help, may be worse, or have an unclear effect, including creativity, reasoning, safety, and general QA.The evaluation is organized into three categories to test collaboration across capabilities that may favor different checkpoints.Source: §3 · pp. 4–5
- Switch Generation achieves the best performance on 13 tasks, with a 12.9% relative improvement over baselines on average; collaboration beats individual models on 16 of 18 tasks.The table caption describes performance across individual models and collaboration methods.Source: Table 1 · p. 5
4 RESULTS
Model collaboration improves on individual models across most tasks, and Switch Generation leads the reported comparisons. Its flexible segment-level routing also improves on other collaboration approaches.
4 RESULTS
- Model collaboration outperforms individual pretrained, finetuned, or aligned models on 16 of 18 tasks, with a 31.0% relative improvement on average.The authors attribute the gains to checkpoints complementing one another and contributing distinct strengths.Source: §4 · p. 5
- Different tasks may favor different generation patch sizes, while fine-tuning the switcher is consistently helpful.This comparison is reported as an ablation of patch size and switching strategy.Source: Table 2 · p. 5
- 13 datasets show Switch Generation outperforming all individual models and collaboration baselines, with a 12.9% average relative improvement; it also gains 6.58 points across categories 2 and 3.The reported 6.58-point average covers tasks where the helpfulness of including the base model was uncertain.Source: §4 · pp. 5–6
- Routing-based baselines average 31.15 points, ahead of weight-based at 29.91, text-based at 26.32, and logit-based at 18.97.The authors describe Switch Generation as finer-grained and more flexible because it routes at the segment level.Source: §4 · p. 6
5 ANALYSIS
The analyses examine how Switch Generation combines model strengths, generalizes across settings, and can be distilled for more efficient inference. They report gains on problems no individual model solves and identify useful switching patterns.
5 ANALYSIS
- Distilling Switch Generation outputs into the aligned model recovers 57.5% of its gains at one-fourth the inference cost of using three models and a switcher.The procedure generates Switch Generation outputs, fine-tunes the aligned model on them, and evaluates that model individually.Source: §5 · pp. 6–7, Figure 3 · p. 6
- Although the pretrained model is not the strongest individual model on most tasks, it has positive collaboration helpfulness on the vast majority.The analysis distinguishes the pretrained model’s individual performance from its contribution when collaborating.Source: Figure 4 · p. 6, §5 · p. 7
- The trained switcher generalizes across four other model settings, with reported average relative improvements of 5.8%, 14.3%, 13.1%, and 3.1%.The settings vary model family, number of checkpoints, and use of specialized experts.Source: §5 · p. 7, Table 3 · p. 7
- On six unseen tasks, Switch Generation generalizes and outperforms baselines by 3.9% on average.The evaluation applies the trained switcher directly to the additional tasks.Source: §5 · pp. 7–8, Table 4 · p. 7
- The most helpful two- and three-model switching sequences are among the most frequent, and their frequency and treatment effect have a positive correlation (R2 = 0.017).Treatment effect is performance when a sequence is used minus performance when it is not used.Source: §5 · p. 8, Figure 5 · p. 8
- Switch Generation answers 10.7% of problems no individual model could solve, while losing out on 8.2% solvable by one model, for a net 2.5% benefit.The comparison measures multi-model correctness against the individual models’ correctness.Source: §5 · p. 8, Table 5 · p. 8
- Segment annotations show the pretrained model most often handling knowledge recall and the aligned model reasoning, with 86% human–LLM agreement.The analysis labels segments for five frequent skills and “others”; agreement was lowest for the “others” category.Source: Figure 6 · p. 8, §5 · p. 9
6 RELATED WORK
The related work situates model collaboration among approaches that exchange information at different levels. Switch Generation applies this idea to checkpoints from the training pipeline and frames those by-products as reusable.
6 RELATED WORK
- Alignment improves skills such as reasoning but has tradeoffs, while base models may struggle with instruction following and safety guardrails.The passage names creativity, uncertainty, pluralism, and knowledge among skills where base models may have advantages.Source: §6 · p. 9
- Prior collaboration methods exchange information through API-, text-, logit-, or weight-level approaches, including routing, debate, logit fusion, and model merging.The passage also describes Mixture-of-Experts as a weight-level approach.Source: §6 · p. 9
- Switch Generation trains a switcher LM to select among pretrained, finetuned, and aligned models at patch granularity, with generalization to unseen candidate models.The work also argues that existing model-development by-products can be reused for collaboration.Source: §6 · p. 9
7 CONCLUSION
SWITCH GENERATION dynamically selects training-pipeline models to contribute segments in one response. It outperforms collaboration baselines and individual models, with further analyses showing generalization and reuse of otherwise discarded checkpoints.
7 CONCLUSION
- SWITCH GENERATION outperforms individual constituent models and eight collaboration baselines on 13 datasets by 12.9% on average.Source: §7 · p. 9
- The approach generalizes to unseen models and tasks and identifies helpful collaboration patterns.Source: §7 · p. 9
- The paper highlights the potential of reusing otherwise discarded by-product models and checkpoints from LM training pipelines.Source: §7 · p. 9
LIMITATIONS
SWITCH GENERATION trades additional inference cost for compositional model strengths; the paper describes ways to mitigate that cost and notes that fixed patch sizes may not suit every task or context.
LIMITATIONS
- Collaborative inference with multiple training-pipeline checkpoints incurs extra cost compared with using only the aligned model.The exchange is for compositional model strengths.Source: LIMITATIONS · p. 15
- Parallel generation across instruction batches, patch-level switching, and distillation can mitigate the added inference cost.Distillation into the aligned model recovers part of the performance gains while returning inference to a single model.Source: LIMITATIONS · p. 15
- Fixed-size patches are used because patch switching works better than token switching, although optimal patch size may vary by task and context.Flexible, dynamic patch-size adjustment is left to future work.Source: LIMITATIONS · p. 15
ETHICS STATEMENT
The ethics discussion identifies risks from malicious model contributions and from combining aligned models with mostly non-safety-aligned base or finetuned models.
ETHICS STATEMENT
- A compromised alignment datastore and malicious aligned model could seriously impact the multi-model collaboration system.Source: ETHICS STATEMENT · p. 15
- Mostly non-safety-aligned base and finetuned models in collaboration might override aligned models’ safety guardrails.The paper identifies this possibility as relevant to red teaming and adversarial language modeling.Source: ETHICS STATEMENT · p. 15
A ANALYSIS (CONT.)
The analysis examines where different model types are used, how frequently switching occurs, and qualitative examples of successful collaboration.
A ANALYSIS (CONT.)
- Aligned models appear more often at response beginnings and endings, while pretrained and finetuned models appear more often in the middle.The authors suggest the middle favors exploration, whereas the beginning and end favor instruction following and summarization.Source: A ANALYSIS (CONT.) · p. 15, Figure 7 · p. 15
- Switching frequency is consistently high, but the analysis finds no consistent relationship between more or less switching and performance.Source: A ANALYSIS (CONT.) · pp. 15–16
- In qualitative examples, SWITCH GENERATION succeeds where pretrained, finetuned, and aligned models do not respond well individually.The authors describe more extensive deliberation and explanation, and suggest aligned-model context can make pretrained and finetuned models helpful.Source: A ANALYSIS (CONT.) · p. 16
B EXPERIMENT DETAILS
The experiments evaluate diverse datasets with zero-shot prompting and document dataset-specific evaluation protocols, baseline configurations, and supplementary analyses.
B EXPERIMENT DETAILS
- Figure 8 tracks switching frequency and P-helpfulness across tasks, while Tables 7–9 provide three examples comparing individual-model and Switch Generation responses.Source: Figure 8 · p. 16, Table 7 · p. 18, Table 8 · p. 19
- 18+6 datasets span multiple LM capability areas, use zero-shot prompting, and are accompanied by dataset statistics and significance tests in Table 6.The table caption defines one-tailed z-test thresholds of p < 0.1, p < 0.05, and p < 0.01 for significant improvements over baselines.Source: B EXPERIMENT DETAILS · p. 16, Table 6 · p. 16
- Evaluation protocols vary by task, including abstain accuracy for AbstainQA, NLI for ScienceMeter, multiple-choice TruthfulQA, and connectivity and shortest-path problems for NLGraph.Other protocols include GPT-4o evaluation for Sycophancy, regex judging for CocoNot, image-description conversion for GuessBench, and average pairwise distance for movie reviews.Source: B EXPERIMENT DETAILS · p. 16
- Switcher training searches eight epoch counts and five learning rates, selecting the configuration with the best dev-set performance for test evaluation.Epochs range from 1 to 8, and learning rates range from 1e−3 to 5e−5.Source: B EXPERIMENT DETAILS · p. 16
- Baselines include individual pretrained, finetuned, and aligned models alongside prompt- or text-based collaboration, logit methods, and model merging.The individual models use top-p = 0.7 and a 512-token generation limit; the described comparisons also include Prompt Route, RouteLLM, Text Collaborate, Text Debate, Logit Merge, Proxy Tuning, Greedy Soup, and DARE-TIES.Source: B EXPERIMENT DETAILS · pp. 16–17