Technical Summary: Escaping the Mode Lottery: Multi-Response Training Improves Language Model Generalization

Problem Statement

Modern language model fine-tuning typically adheres to a convention of pairing a single prompt with a single response. While suitable for classification tasks, this approach is ill-suited for open-ended generation where prompts often admit multiple valid completions (e.g., "write a function that sorts a list"). By training on only one response per prompt, the model is exposed to a narrow slice of a multi-modal conditional distribution. The authors term this phenomenon the "mode lottery," where training data may repeatedly emphasize a subset of plausible modes (often the most common, easiest-to-generate, or highest-rewarded) while leaving other valid and valuable modes statistically invisible.

Existing pipelines often generate multiple candidates but select a single "best" response for training (e.g., via reward maximization). This paper argues that the field lacks a principled account of when retaining multiple responses is beneficial, how much gain can be expected, and which responses should be retained.

Methodology and Theoretical Framework

The paper formalizes Multi-Response Training (MRT) as training on $K > 1$ responses per prompt, contrasting it with Standard Single-Response Training (SRT, where $K=1$). The core theoretical insight is that prompts and responses are distinct statistical resources:

• Additional prompts reduce uncertainty about the marginal input distribution ($P(x)$).
• Additional responses to the same prompt reduce uncertainty about the conditional output distribution ($P(y|x)$).

1. Variance–Budget Law

The authors derive a structural decomposition of the empirical loss variance. For a fixed model $\theta$, the variance of the MRT estimator decomposes as:
$$\text{Var}[\hat{L}_{MRT}(\theta)] = \frac{V_x(\theta)}{N_p} + \frac{V_{y|x}(\theta)}{N_p K}$$
Where:

• $V_x$ is the across-prompt variance (uncertainty about the input distribution).
• $V_{y|x}$ is the expected within-prompt variance (uncertainty about the conditional output distribution).
• $N_p$ is the number of distinct prompts.
• $K$ is the number of responses per prompt.

This decomposition reveals a "prompt floor": increasing $K$ reduces the second term but cannot reduce the first term ($V_x/N_p$). Consequently, MRT yields diminishing returns once the prompt-level uncertainty dominates.

By introducing acquisition costs ($C_p$ for prompts, $C_r$ for responses) and a fixed budget $B = N_p(C_p + K C_r)$, the authors derive the cost-aware optimal multiplicity $K^\star$:
$$K^\star = \sqrt{\frac{V_{y|x}}{V_x} \cdot \frac{C_p}{C_r}}$$
This rule suggests MRT is most valuable when responses are cheap, prompts are expensive, and the conditional response distribution is highly diverse (high $V_{y|x}/V_x$).

2. Generalization Bounds

The paper provides two complementary generalization bounds to quantify the potential gains:

• Bernstein Bound: Shows that for a fixed model, increasing $K$ reduces the generalization error bound by a factor dependent on $V_{y|x}/K$, while the prompt-level term remains unchanged.
• Hierarchical Rademacher Bound: Addresses the non-i.i.d. structure of MRT data (repeated prompts). It proves that despite the hierarchical structure, the response-level complexity scales with the total number of prompt-response pairs ($N_p K$) rather than just $K$, provided the supremum over the model class is handled correctly. This certifies that the gains from MRT survive model selection.

3. Target-Aware Response Selection

The paper analyzes how different selection strategies alter the training target:

• Random-K-of-N (RKoN): Samples $K$ responses uniformly at random. This is the unbiased default for distributional fine-tuning, preserving the candidate distribution and reducing gradient variance at a $1/K$ rate.
• Best-K-of-N (BKoN): Selects responses based on reward scores. The authors show this induces a persistent upper-tail bias. As the candidate pool grows, BKoN concentrates training on the reward-dominant mode, potentially collapsing valid modes and misaligning gradients with the true distributional objective.
• GRADES (Greedy Reward-Aware DiversE Subsetting): A submodular optimization approach that balances quality and diversity. It provides a $(1 - 1/e)$ approximation guarantee for retaining a balance of high-quality and diverse responses, acting as a practical proxy for stratified sampling.

Experimental Results

The authors validate their theory through controlled simulations and large-scale LLM fine-tuning across four domains: instruction following (Gold), dialogue (Nectar), preference alignment (UltraFeedback), and code generation (CodeContests).

1. Variance Decomposition: Controlled simulations confirmed that increasing $K$ reduces only the within-prompt variance component, approaching the predicted prompt floor asymptotically.
2. Optimal Multiplicity: Under a fixed budget, test performance followed a U-shaped curve, with the minimum occurring near the predicted $K^\star$. Datasets with high response-level variability relative to prompt-level variability (e.g., CodeContests) showed the largest gains from MRT, while prompt-heavy datasets (e.g., Nectar) showed smaller improvements.
3. Selector Impact:
   • RKoN consistently improved reference loss (fit to the conditional distribution) and coverage as $K$ increased.
   • BKoN achieved higher reward scores but suffered from lower diversity and worse reference likelihood, confirming the target shift.
   • GRADES/DBKoN (Diverse-Best-K-of-N) successfully balanced quality and coverage, outperforming pure diversity or pure reward selection in reference coverage metrics.
4. Implicit MRT: Experiments with prompt redundancy (paraphrases) demonstrated that highly redundant corpora can mimic MRT effects if the paraphrases expose distinct conditional modes. However, if paraphrases are paired with near-duplicate responses, the benefit is negligible.

Significance and Claims

The paper reframes response multiplicity not as a heuristic but as a data-allocation problem with clear statistical guidance. Its primary claims are:

1. Context-Dependent Value: MRT is not universally superior; its value is contingent on the ratio of within-prompt to across-prompt variance and the relative costs of data acquisition.
2. Statistical Grounding: When responses are cheap and diverse, and prompts are expensive, retaining multiple responses is a statistically grounded choice that escapes the "mode lottery" by training on more of the conditional structure.
3. Selection Matters: The choice of selector is critical. While reward-based selection optimizes for a specific metric, it can collapse valid modes. For distributional generalization, unbiased selection (RKoN) or diversity-aware selection (GRADES) is required to preserve the conditional structure.
4. Implicit vs. Explicit: Large, redundant corpora can provide an "implicit MRT" effect, but explicit multi-response training is most beneficial in low-redundancy regimes where each prompt covers a distinct region of the task distribution.

The authors conclude that MRT offers a principled way to decide when response multiplicity is worth the cost, moving beyond the "one prompt, one response" convention to better model the multi-modal nature of open-ended generation.