Optimizing Diversity and Quality through Base-Aligned Model Collaboration

AI summary

Overview

• Research area: Natural Language Processing / large language model generation, specifically inference-time decoding and the diversity-quality trade-off in open-ended text generation.
• Technical level: Intermediate. The core idea is intuitive, but the evaluation uses multi-objective optimization concepts (Coverage, Dominance, Pareto frontiers) and a large set of diversity metrics.
• Scope (one sentence): The paper proposes BACo, an inference-time token-level framework that routes each decoding step between a base LLM and its aligned counterpart to jointly improve output diversity and quality across instruction following, dialogue, and creative writing.

What This Paper Is About

Alignment improves the quality of LLM outputs but sharply reduces their diversity, so aligned models tend to produce very similar responses to repeated open-ended prompts, while base models produce varied but lower-quality responses. BACo attacks this trade-off not by retraining or re-prompting a single model, but by having a base model and its aligned counterpart collaborate token by token during a single decoding pass, using a lightweight router to decide which model generates each next token. The goal is a controllable method that can be moved along the diversity-quality spectrum on demand.

Key Contributions

1. The BACo framework. An inference-time, token-level model collaboration framework that combines a base model (as a source of diversity) with its aligned counterpart (as a source of quality), together with a family of lightweight routing strategies, requiring no fine-tuning or prompting.
2. A reformulation of the trade-off as bi-dimensional evaluation spaces. The authors frame evaluation as a collection of 11 × 2 diversity-quality spaces (11 diversity metrics paired with 2 quality metrics) and evaluate across them, including long-form diversity (plot structure and emotional flow) and human evaluation.
3. Extensive experiments across three open-ended tasks. Instruction following (NoveltyBench), dialogue (WildChat), and creative writing (Narrative-Discourse), showing BACo consistently outperforms strong inference-time baselines.
4. Controllability as a first-class property. An adjustable routing threshold γ lets users tune the base/aligned contribution, analogous to tuning decoding temperature, with a family of routers offering an additional axis of control.

Main Findings

• The trade-off is stark in a pilot study. On a subset of WildChat with Llama-3, the base model (Llama-3-8B) scored 8.13 clusters and 1.28 reward, while the aligned model (Llama-3-8B-Instruct) scored 2.58 clusters and 7.62 reward. The base model was 3.15x more diverse, while the aligned model had 5.95x higher quality. Diversity was measured as the number of semantic equivalent clusters of the output group; quality as the average reward per output from Skywork-Reward-Gemma-2-27B, with 10 outputs sampled per prompt.
• BACo beats inference-time baselines overall. With the best router, BACo achieved a 21.3% joint improvement in diversity and quality, an overall Coverage of 0.403 and 32.7% Dominance, versus 0.261 Coverage and 9.6% Dominance for the strongest baseline, nudging.
• Semantic diversity gains are especially large. BACo reached 40.5% Dominance on semantic spaces versus 24.9% on lexical spaces, suggesting it produces meaning-level rather than surface-level variation. BACo improved Coverage by 0.142 over the strongest baseline, which the authors describe as expanding the achievable diversity-quality area by over 30%.
• Gains widen on NoveltyBench. Coverage improved by 0.274 and Dominance reached 39.9% on that dataset, and results were consistent for the Olmo2 model family.
• Routing choice matters, but the framework is robust. On NoveltyBench, the combined -P-Punc router performed best (Overall Coverage 0.474, Dominance 31.0%). The random router (-Rand) looked competitive in aggregate (Coverage 0.451) but collapsed on semantic metrics (0% Dominance on Semantic Entropy), and the prompt-based -Judge router underperformed simpler heuristics at much higher cost.
• Base-aligned collaboration beats aligned-aligned collaboration. Under the same -P-Punc router on NoveltyBench, base-aligned BACo reached 0.474 Overall Coverage and 50.9% Dominance, versus 0.014 Coverage and 5.7% Dominance for aligned-aligned collaboration (AACo, using Llama-3-8B and Llama-3.2-11B).
• Gains hold on verifiable tasks. On IFEval (instruction following) and GSM8K (mathematical reasoning), BACo achieved higher diversity at matched quality or accuracy than the aligned baseline, indicating results are not artifacts of open-ended metrics.
• Multi-turn and repetition results. On MT-Bench, BACo-P-Punc achieved Coverage 0.681 vs. 0.320 and Dominance 72.8% vs. 27.2% against the aligned baseline. On Artificial Hivemind with Llama-3.1-70B-Instruct, at a threshold matched to Min-p's quality (average reward 8.08; γ = 0.2), 15% of BACo generations fell in the genuinely diverse regime (similarity < 0.6) versus 8% for Min-p and 5% for Top-p.
• Human evaluation supports the automatic results. Four expert annotators rated BACo higher on quality despite near-identical automatic quality scores (aligned 5.93; BACo 5.85): 4.04 vs. 2.83 on NoveltyBench and 3.83 vs. 3.44 on WildChat (1–5 Likert). BACo won 79.0% of NoveltyBench prompts and 63.9% of WildChat prompts overall, and was judged most creative in 79.6% and 61.8% of prompts respectively.
• Long-form diversity also improves. On Narrative-Discourse, BACo achieved substantially higher turning-point and arousal diversity than the aligned baseline at comparable quality.
• Switching declines over the course of generation. Base-model contribution and switching frequency were high at the start and decreased over time across all three tasks.
• An emergent failure mode exists. Tuning the router more aggressively toward diversity increases premature end-of-sequence termination ("inherent early stopping"), which the authors note can both truncate valid continuations and guard against incoherent repetition.

Methodology in Plain English

The key motivation is the "superficial alignment" observation: a base model and its aligned version usually agree on what the next token should be, so the two models can be mixed without producing incoherent text. BACo formalizes this as a per-token weighted combination of the base and aligned model distributions, where the weight is 0 or 1 — the router picks exactly one model per token.

Routers use two kinds of signals. Logit-based routers look at the base model's next-token distribution: BACo-P switches to the base model when its maximum token probability falls below a threshold γ, and BACo-H switches when entropy exceeds γ. Content-based routers look at what kind of word is being generated: BACo-Punc sends the decision to the aligned model when the top-ranked token is punctuation or a formatting token, and BACo-FC routes function words to the aligned model to keep style coherent. These can be combined, with one strategy taking priority (for example, -P-Punc checks punctuation first, then falls back to probability). Content-based strategies work even for black-box models because they do not need logits. Switching is restricted to word boundaries to avoid problems when the two models tokenize differently.

Evaluation treats each method not as a single point but as a curve traced out as its control configuration varies. Two multi-objective indicators summarize performance: Coverage (the area under a method's trade-off curve) and Dominance (the share of the global Pareto frontier the method contributes, measured with the C-metric). These are computed over 11 × 2 diversity-quality spaces and averaged. The 11 diversity metrics include Distinct-n, EAD-n, Self-BLEU, embedding cosine dissimilarity, Vendi Score, NLI diversity, and Semantic Entropy; the 2 quality metrics are perplexity under the aligned model and reward scores from Skywork-Reward-Gemma-2-27B. Baselines span single-model sampling at varying temperatures, prompting-based methods (in-context resampling, paraphrase prompting), decoding-based Diverse Beam Search, ensemble methods (response ensemble, logit ensemble), and the collaboration-based nudging method.

Why This Matters

• Research impact: The paper reframes the diversity-quality trade-off as a controllable multi-objective problem rather than a fixed property of a single model, and shows that a training checkpoint and its aligned version can be treated as complementary collaborators. It also provides an evaluation protocol (Coverage and Dominance over many metric pairs) that could be reused for other trade-off problems.
• Real-world applications:
  • Creative writing assistants that need varied plots, characters, and emotional arcs rather than formulaic prose.
  • Chatbots and dialogue systems where repeated user queries currently return near-identical responses.
  • Ideation and brainstorming tools in human-AI interaction, where diversity suppression has been shown to hurt idea generation.
  • Systems needing verifiable output, since BACo also preserved accuracy on GSM8K and IFEval.
• Industry relevance: BACo requires no training, no prompt engineering, and runs in a single decoding pass, so it can be retrofitted onto existing deployed model pairs. An adjustable threshold gives product teams a single knob to tune how adventurous versus reliable the model should be per use case.

Future Directions

• Position-aware routing. Because switching and base contribution decline as generation proceeds, the authors suggest dynamic, position-aware thresholds, noting the current pattern may be less suitable for list-structured tasks that need uniform exploration over time.
• Learned routers. The authors explicitly leave learned routers for future work, citing conflicting diversity objectives, unstable joint training, and the high cost of group-level sampling during training.
• A "breadth thinking" mode for LLMs. The paper positions BACo as a preliminary step toward such a mode, though the discussion is cut off in the available content.
• Deeper investigation of the creativity link. The creativity analysis is described as preliminary, and the authors leave a larger-scale and deeper investigation to future work.
• Robustness and mitigation questions. The inherent early-stopping failure mode raises open questions about restart strategies and mitigating truncation, as well as safety and helpfulness implications of mixing in base-model tokens.

Target Audience

Researchers and practitioners working on LLM decoding, alignment, and open-ended generation will get the most from this paper, particularly those studying diversity, creativity, or human-AI co-writing. It is also relevant to evaluation researchers interested in multi-objective metrics for generation, and to engineers who want a training-free, deployable way to tune the diversity-quality balance of an existing base/aligned model pair.

Authors’ abstract

Alignment has greatly improved large language models (LLMs)' output quality at the cost of diversity, yielding highly similar outputs across generations, especially in open-ended generation tasks. We propose Base-Aligned Model Collaboration (BACo), an inference-time token-level model collaboration framework that dynamically combines a base LLM with its aligned counterpart to optimize diversity and quality. Using uncertainty and content-based signals, BACo employs routing strategies to determine, at each token, which model to decode from. Prior diversity-promoting methods often improve diversity at the expense of quality or require expensive decoding or post-training. In contrast, BACo achieves both high diversity and quality post hoc within a single pass, while offering strong controllability. We introduce a family of effective routing strategies and evaluate them across three open-ended generation tasks with 13 diversity and quality metrics. BACo consistently surpasses state-of-the-art inference-time baselines. With our best router, BACo achieves a 21.3% joint improvement in diversity and quality, which is further supported by human evaluations. Overall, our results demonstrate that collaboration between base and aligned models provides an effective and controllable mechanism for optimizing the diversity-quality trade-off.