Optimizing Diversity and Quality through Base-Aligned Model Collaboration

TL;DR

This work tackles the diversity decline that accompanies LLM alignment by introducing BACo, a token-level inference-time framework that collaborates a base model (diversity) with its aligned counterpart (quality). A lightweight router at each decoding step selects which model generates the next token, formalized as $P_{BACo}(y_t|c_t) = w_{base} P_{base}(y_t|c_t;\theta_{base}) + (1-w_{base}) P_{aligned}(y_t|c_t;\theta_{aligned})$, with controllable thresholds to navigate the diversity-quality frontier. Across three open-ended tasks, BACo consistently surpasses inference-time baselines on 11 diversity metrics and 2 quality metrics, achieving a best reported joint improvement of $21.3\%$ in diversity and quality, and is validated by human evaluations showing enhanced creativity and perceived quality. The design space includes logit-based and content-based routers, and results demonstrate that combining strategies yields the strongest performance, establishing base-aligned collaboration as a practical, controllable approach to multi-objective open-ended generation. The work further demonstrates robustness across long-form narrative tasks and highlights potential future directions in efficiency, partially aligned checkpoints, and user-driven interfaces for breadth-oriented AI systems.

Abstract

Alignment has greatly improved large language models (LLMs)' output quality at the cost of diversity, yielding highly similar outputs across generations. 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. Inspired by prior work (Fei et al., 2025), BACo employs routing strategies that determine, at each token, from which model to decode based on next-token prediction uncertainty and predicted contents' semantic role. Prior diversity-promoting methods, such as retraining, prompt engineering, and multi-sampling methods, improve diversity but often degrade quality or require costly decoding or post-training. In contrast, BACo achieves both high diversity and quality post hoc within a single pass, while offering strong controllability. We explore a family of routing strategies, across three open-ended generation tasks and 13 metrics covering diversity and quality, 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. Human evaluations also mirror these improvements. The results suggest that collaboration between base and aligned models can optimize and control diversity and quality.

Figures (13)

• Figure 1: BACo is an inference-time token-level model collaboration framework that combines a base model's diversity with its aligned counterpart's quality. (A) A comparison of generated outputs. The aligned model produces high-quality but low-diversity outputs, while the base model produces high-diversity but low-quality outputs. BACo optimizes both diversity and quality by dynamically routing between them. The probabilities of token(s) are in grey next to text boxes. (B) Illustration of the diversity-quality trade-off space. Single models face a steep trade-off, where improving diversity by adjusting configuration (e.g., by increasing temperature) degrades quality. BACo achieves a better Pareto curve and allows for easy traversal across this frontier by adjusting the router's threshold. The examples in this figure are modified for simplicity.
• Figure 2: Illustration of the two indicators on diversity-quality space: Coverage is the area under a method’s trade-off curve (blue shading for the blue method); Dominance is the proportion of the global Pareto frontier (highlighted curves) contributed by the method. In practice, high Coverage is preferable for general-purpose design, as it ensures a method offers a good trade-off across a wide range of tasks or user preferences. In contrast, high Dominance is desirable when selecting a specialized method to achieve the optimal trade-off within a specific, known target region of the space.
• Figure 2: Averaged performance of all methods across all datasets and diversity–quality spaces. BACo consistently outperforms baselines across all semantic and most lexical spaces, demonstrating stronger controllability and substantially improving the semantic diversity–quality trade-off. The overall gains, as driven primarily by improvements in semantic, suggest that BACo produces more meaningful and content-level diversity, rather than superficial word-level changes, compared to other methods. See full results at Appendix \ref{['app:detail_results']}.
• Figure 3: Averaged performance of routers within BACo on NoveltyBench across all diversity–quality spaces. The -P-Punc router achieves the best overall performance. While the random router (-Rand) attains moderately strong results, mainly from increased surface-level lexical diversity, its performance drops sharply on semantic metrics, confirming that unguided switching fails to produce meaningful diversity. In contrast, -P-Punc delivers the most balanced and consistent results across both lexical and semantic evaluations, showing combination of designed routing strategies leads to more meaningful diversity.
• Figure 4: Comparison of BACo's and baselines' discourse-level diversity-quality trade-off curve on Narrative Discourse. BACo obtains a larger Coverage, achievable in the high-diversity, high-quality region (top-left). The results demonstrate it has richer discourse-level diversity without sacrificing quality largely, compared with baselines. The x-axis is quality (perplexity; lower is better), and the y-axis is discourse-level diversity, either turning-point diversity (left figure) or arousal diversity (right figure) (higher is better).