Technical Summary: Optimizing Diversity and Quality through Base-Aligned Model Collaboration (BACO)

1. Problem Statement

Large Language Models (LLMs) face a fundamental diversity-quality trade-off. While alignment techniques (e.g., RLHF) significantly improve output quality, instruction following, and safety, they systematically reduce output diversity. Aligned models tend to converge on a single dominant response (mode collapse) across repeated generations, whereas their base (pre-aligned) counterparts produce diverse but often lower-quality or incoherent outputs.

Existing methods to address this trade-off suffer from significant limitations:

• Training-time approaches: Incorporating diversity objectives into preference optimization often compromise alignment properties like safety and helpfulness.
• Inference-time approaches: Techniques such as temperature scaling, diverse beam search, or prompt-based resampling often require multiple decoding passes, incur high computational costs, or degrade generation quality disproportionately.

The core challenge is to achieve both high diversity and high quality in a single decoding pass without requiring additional fine-tuning or expensive post-training.

2. Methodology: Base-Aligned Model Collaboration (BACO)

The authors propose BACO, an inference-time, token-level model collaboration framework. BACO dynamically combines a base model (source of diversity) and its aligned counterpart (source of quality) during the decoding process.

Core Mechanism

BACO operates via a lightweight router that decides, at each token generation step, which model should generate the next token. The framework relies on the "superficial alignment" hypothesis, where base and aligned models largely agree on next-token predictions, allowing for sparse interventions.

The generation probability is formulated as a gated mixture:
$$P_{BACO}(y_t | c_t) = w_{base} \cdot P_{base}(y_t | c_t) + (1 - w_{base}) \cdot P_{aligned}(y_t | c_t)$$
Where $w_{base} \in \{0, 1\}$ is determined by the routing strategy $R$ based on the context $c_t = [x, y_{<t}]$.

Routing Strategies

The router utilizes two complementary signal categories to determine the switch:

1. Logit-Based Strategies: Leverage the model's prediction distribution to infer uncertainty.
   • Example: BACO-P routes to the base model when the base model's maximum token probability falls below a threshold $\gamma$.
   • Example: BACO-H routes to the base model when the entropy of the base model's prediction exceeds $\gamma$.
2. Content-Based Strategies: Leverage the semantic roles of predicted tokens.
   • Example: BACO-PUNC routes to the aligned model when the top-ranked token is punctuation or a formatting token (preserving structural coherence).
   • Example: BACO-FC routes to the aligned model for function words to maintain stylistic consistency.

Control: A threshold parameter $\gamma$ allows users to continuously tune the balance between diversity and quality. Lower $\gamma$ biases toward the aligned model (higher quality), while higher $\gamma$ biases toward the base model (higher diversity).

Implementation Details:

• The first token is always generated by the aligned model to ensure trajectory quality.
• Switching is restricted to word boundaries to prevent tokenization mismatches.
• End-of-sequence (EOS) tokens are controlled by the aligned model to prevent premature termination.

3. Key Contributions

1. Framework Proposal: Introduction of BACO, a token-level collaboration framework that requires no fine-tuning and operates in a single decoding pass, effectively balancing the diversity-quality trade-off.
2. Evaluation Formulation: The authors formulate the diversity-quality trade-off not as a single objective but as a collection of bi-dimensional evaluation spaces (11 diversity metrics $\times$ 2 quality metrics). They introduce Coverage (area under the trade-off curve) and Dominance (proportion of the global Pareto frontier) as metrics to aggregate performance across these spaces.
3. Comprehensive Validation: Extensive experiments across three open-ended generation tasks (Instruction Following, Dialogue, Creative Writing) and verifiable benchmarks (Math, Instruction Following with constraints), demonstrating consistent improvements over state-of-the-art baselines.

4. Results

The authors evaluated BACO using 13 metrics across 11 diversity-quality spaces on datasets including NoveltyBench, WildChat, and Narrative-Discourse.

• Performance Gains: With the best router (-P-PUNC, combining probability and punctuation rules), BACO achieved a 21.3% joint improvement in diversity and quality over the strongest inference-time baselines.
• Pareto Optimality: BACO contributed to 32.7% of the global Pareto frontier across all spaces, with a particularly strong 40.5% dominance in semantic diversity spaces.
• Semantic vs. Lexical: BACO showed superior performance in semantic diversity (meaning and intent) compared to lexical diversity (surface-level phrasing), indicating it produces meaningful variations rather than just superficial noise.
• Verifiable Benchmarks: On GSM8K (math) and IFEval (instruction following), BACO maintained comparable accuracy to aligned baselines while achieving significantly higher diversity, proving the gains are not artifacts of open-ended metrics.
• Human Evaluation: In human studies on NoveltyBench and WildChat, BACO outputs were rated significantly higher in quality (4.04 vs. 2.83 on NoveltyBench) and creativity (79.6% win rate) compared to the aligned baseline, despite having similar automatic quality scores.
• Long-Form Generation: On the Narrative-Discourse dataset, BACO demonstrated superior diversity in plot structure (turning points) and emotional flow (arousal) without sacrificing coherence.

5. Significance and Claims

The paper claims that BACO provides a simple, controllable, and effective mechanism for optimizing the diversity-quality trade-off in LLMs.

• Complementarity: The results suggest that the collaboration between a less-aligned (high diversity) and a more-aligned (high quality) model is essential for jointly optimizing both objectives, outperforming even collaborations between two aligned models.
• Controllability: Unlike temperature scaling, which degrades quality to gain diversity, BACO offers a principled way to traverse the diversity-quality spectrum via the routing threshold $\gamma$.
• Broad Applicability: The framework is applicable to various open-ended tasks, including creative writing and dialogue, and extends to long-form generation.
• Future Direction: The authors position BACO as a step toward a "breadth thinking" mode for LLMs, where models explore a wide space of perspectives before converging, potentially helping humans break out of "information cocoons."

The authors remain modest regarding limitations, noting that BACO incurs approximately $2\times$ FLOPs overhead per token (though independent of group size) and that the current prototype is not fully optimized for speed. They also acknowledge that enhanced diversity could amplify harmful generations if safety safeguards are not maintained, emphasizing that BACO is complementary to, not a replacement for, safety alignment.