Weight ensembling improves reasoning in language models
Authors: Xingyu Dang
Christina Baek
Kaiyue Wen
J Zico Kolter
Aditi Raghunathan
Venue: COLM 2025 (2025)
Keywords: test-time scaling
RL
reasoning
diversity
decoding
TL;DR: Weight ensembling improves pass@k of reasoning models.

Abstract: We investigate a pitfall during the training of reasoning models where the diversity of generations begins to collapse, leading to suboptimal test-time scaling. Notably, Pass@1 reliably improves during supervised finetuning (SFT), but Pass@k rapidly deteriorates. Surprisingly, a simple intervention of interpolating the weights of the latest SFT checkpoint with an early checkpoint, otherwise known as WiSE-FT, almost completely recovers Pass@k while also improving Pass@1. The WiSE-FT variant achieves better test-time scaling (Best@k, majority vote) and achieves superior results with less data when tuned further by reinforcement learning. Finally, we note that WiSE-FT provides complementary gains across performance metrics that is not achievable by diversity-inducing decoding strategies alone, like temperature scaling. We formalize a \emph{bias-variance tradeoff} of Pass@k with respect to the expectation and variance of Pass@1 over the test distribution. We find that WiSE-FT can reduce bias and variance simultaneously, while temperature scaling and possibly other decoding strategies face an inherent tradeoff between decreasing variance with increasing bias.

Published as a conference paper at COLM 2025
Weight Ensembling Improves Reasoning in Language Models
Xingyu Dang*1
Christina Back*2
Kaiyue Wen
Zico Kolter*
Aditi Raghunathan"
'Tsinghua University CCanneiie Mellon University
Stanford University
dangxy20@mails, wtapaailss.ssssiisiieeeee tsinghua. edu. cn
kbaek@andrew. cmu. edu
Abstract
We investigate failure mode that arises during the training of reasoning models, where the diversity of generations begins to collapse, leading to suboptimal test-time scaling. Notably, the Pass@1 rate reliably improves during supervised finetuning (SFT), but Pass@k rapidly deteriorates. Sur- prisingly, a simple intervention of interpolating the weights of the latest SFT checkpoint with an early checkpoint, otherwise known WiSE-FT, almost completely recovers Pass@k while also improving Pass@1. The WiSE-FT variant achieves better test-time scaling (Best@k, majority vote) and achieves superior results with less data when tuned further by rein- forcement learning. Finally, we find that WiSE-FT provides complementary performance gains that cannot be achieved only through diversity-inducinn decoding strategies, like temperature scaling. We formalize a bias-variance tradeoff of Pass@k with respect to the expectation and variance of Pass@1 over the test distribution. We find that WiSE-FT can reduce bias and vari- ance simultaneously, while temperature scaling inherently trades off be- tween bias and variance.
Introduction
Recent advances in large language models (LLMs) have showcased their remarkable ability to perform complex reasoning, yet these successes often hinge on test-time scaling strategies (Lightman et al., 2023; Snell et al., 2024; Wu et al., 2024). In many applications, such as math problems, puzzles, and logical reasoning, LLMs employ verification framework where is significantly easier for the model to verify candidate solution than to generate one from scratch. This distinction has given rise to strategies that sample multiple "reasoning traces" or sequences of reasoning steps during inference, selecting the best final guess through an outcome reward model (ORM) or majority vote. In this setting, an upper bound on the performance a model could achieve is measured by Pass@K, or the probability that at least one out of K independently sampled reasoning traces is correct.
Unfortunately, while the standard training pipeline of supervised finetuning (SFT) followed by reinforcement learning (RL) dependably improves Pass@1 for reasoning, Pass@K tends to drop early into finetuning (Cobbe et al., 2021; Chow al., 2024a; Chen et al., 2025). This mismatch arises from symptom of finetuning called diversity collapse, where overtuned models yield less diverse generations. This is detrimental Pass@K since the model wastes K attempts on only a handful of guesses. In fact, by analyzing the model's error rate i.e., 1 Pass@1, 1, across the test distribution, we derive Pass@K bias-variance trade-off. To improve expected lest Pass@K, one can either reduce the bias which is the expected error rate or how much the model's 's error rate varies across problems. The latter term connected to diversity more diversity allows models to hedge and do uniformly well across all test questions. In particular, during SFT, Pass@1 improves (bias 4) at the cost of diversity collapse (variance 1).
Surprisingly, common ways of alleviating diversity collapse, such as early stopping at peak Pass@K decoding with high temperature, suffer from the reverse trade-off: diversity improves (variance L) at the cost of overall Pass@1 degrading (bias t). Consequently, in this paper are concerned with central question:

---

Published as a conference paper at COLM 2025
Pass@1
Pass@4
Pass@1 16
0.80
0.90
0.6
0.88
Pass@1
Pass@4
P
0.5
0.86
0.70
0.84
0.4
0.65
0.82
0.80
500
1000
1500
500
1000
1500
500
1000
1500
Step
Step
Step
SFT T=0.7
SFT
IWSE-FTT=1.0
SFT T=1.3
SFT T=1.6
Figure 1: Pass@k of WiSE-FT versus SFT on GSM8k Gemma-2-2B supervised finetuned and evaluated on GSM8k. At each SFT timestep t, we evaluate Pass@k of checkpoint w, (in dashed) with its WiSE-FT variant 1/2. w: 1/2 wo (in solid), where traces are independently sampled with temperature [0.7, 1.0, 1.3, 1.6).
Is it possible to simultaneously improve both Pass@1 and Pass@K, thereby overcoming the bias-variance tradeoff inherent in current approaches?
In work, we introduce simple, scalable and effective intervention that allows models to achieve both high Pass@K and Pass@1 across mathematical reasoning tasks GSM8K, MATH,, and AIME. The specific technique we use variant of WISE-FT (Wortsman et al., 2022) where we interpolate the weights of the latest SFT checkpoint w, with an early checkpoint wp as @WiSE(() 1/2 200 + 1/2 wp Our key finding that WiSE-FT successfully merges the diverse sampling capabilities of earlier checkpoints while retaining or surpassing the Pass@1 of later checkpoints. In Figure 1, we observe that the WiSE-FT model achieves both higher Pass@K and Pass@1 with more SFT steps f, unlike naive SFT which suffers from an early decay in Pass@K. Moreover, the gains with WiSE-FT is unachievable by early-stopping or diversity-aware decoding alone.
Thus, we propose
paradigm of training reasoning models: 1.) Train extensively using SFT as long as Pass@1 improves, 2.) Perform WiSE-FT with an earlier SFT checkpoint, 3.) Continue tuning the WISE-FT variant using RL. Overall, the WISE-FT model has the following immediate practical benefits:
Better Test-Time Scaling Across all datasets and base models, the WiSE-FT variant achieves the highest performance with test-time scaling (Majority Vote, ORM) compared to an overtrained SFT model paired with diversity-aware decoding.
Better Reinforcement Learning Since RL uses self-generated data to tune models, to generalize reliably, it is important for generations to provide sufficient learning signal while also having high coverage over the data space. We find that continued RL train- ing starting from WiSE-FT weights achieves superior results with less synthetic data compared to initializing RL from the last SFT checkpoint and even early-stopped SFT.
In summary, we provide comprehensive analysis of how reasoning models suffer from diversity collapse during SFT and its negative downstream impact during RL and test-time scaling. We first discuss our WiSE-FT findings in $4. Motivated by this discovery, we investigate two fundamental questions. First, we investigate diversity collapse during SFT and RL of reasoning models in $5. Diversity collapse not only impacts the model's ''' ability to attempt different guesses. In fact, we make an even stronger observation - the generations of reasoning models converge towards single reasoning trace for each test question. We theoretically prove that standard RL algorithms (i.e., REINFORCE and GRPO) fail to recover lost diversity in simplified discrete bandit setting.
Second, we formalize the competing goals of Pass@1 and Pass@K as a bias-variance trade- off in $6. We empirically measure and compare the bias and variance of WiSE-FT versus
2

---

Published as a conference paper at COLM 2025
early-stopping versus high temperature decoding. Notably, only WiSE-FT reduces both bias and variance. We conclude with a remark on the limitations of decoding strategies such as top-k (Shao et al., 2017), nucleus (Holtzman et al., 2020), and min-p (Nguyen et al., 2024), at eliciting the maximum capabilities with test-time scaling from current reasoning models.
2 Related Works
The standard pipeline for enhancing reasoning in LLMs involves an initial phase of supervised fine-tuning (SFT) followed by reinforcement learning (RL) (Guo et al., 2025; Setlur et al., 2024). SFT is critical for instilling interpretable and readable reasoning chains and ensuring that the model adheres to a consistent rollout templates (Guo et al., 2025). However, number of recent works have identified critical pitfalls of SFT that hinders the model's ability explore and ultimately it's overall problem solving ability. Notably, Cobbe et al. (2021) observe diversity collapse when finetuning on GSM8k training dataset, during which the Pass@1 continuously improves whereas Pass@k starts to fall shortly into the training. Similar diversity collapse phenomenon also exists in the self-improvement setting with SFT (Song et al., 2024), and is theoretically investigated as the sharpening effect (Huang et al., 2024). This is not desirable diverse sampling at inference is important for test-time scaling using majority voting (Wang et al., 2023) or reward model guided search (Setlur et al., 2024; Beeching et al., 2024). Yeo et al. (2025); Chu et al. (2025) attribute this behavior to overfitting, memorization of samples and overfixation to a template style leading to reduced generalization. In our work, we corroborate similar findings and propose ensembling over the course of SFT as a mitigation strategy.
Diversity collapse with SFT:
Mitigating diversity collapse: Given the importance of diversity for effectively scaling inference-time compute, several recent works have proposed auxiliary finetuning objec- tives and decoding strategies to mitigate diversity collapse. Li et al. (2025) regularize the SFT process using game-theoretic framework that encourages sparse updates, thereby preserving output diversity. Zhang et al. (2024b) directly optimizes for diversity during finetuning. Other approaches modify the finetuning procedure to directly optimize for Best-of-N sampling at inference time (Chow et al., 2024b; Sessa et al., 2024; Chen et al., 2025). Another line of work focuses on inference- solutions through modified beam search strategies (Vijayakumar et al., 2018; Olausson et al., 2024; Chen et al., 2024; Beeching et al., 2024). Li et al. (2023) improve diversity during parallel decoding by appending curated prompts to the input. In formal reasoning settings e.g,, Lean, methods such as Monte Carlo tree search have been used to diversify interme- diate reasoning steps, demonstrated in AlphaProof (AlphaProof and AlphaGeometry teams,, 2024). In this work, we identify a simple and complementary intervention during the finetuning process to maintain the diversity of generations. We especially care about enforcing diversity while preserving the overall accuracy of generations.
decoding, explicitly encouraging diverse
3 Preliminaries and Experimental Setup
3.1 Pass@k, Best@k, and Majority Vote
Given reasoning model f(-), decoding strategy D, and problem x, the model's solution is
obtained by sampling reasoning trace [x, s)), s(2), s(.,s))) s(n), , consisting of sequence of intermediate steps s(i) and a final guess ý. Given k independently sampled traces, Pass@K measures the probability that at least one guess matches the true answer y:
Pass@K(x) TE (((() [1 E [k] s.t. 91 = y = 1 p1
(1)
where p= = P(g = y x, is the Pass@1l or marginal probability of sampling the ground truth answer. Then (1-p)) xx) is the probability that all K guesses are incorrect. We will refer
to Pass@1 as Px interchangeably in our paper.
3

---

Published as a conference paper at COLM 2025
WISE-FT GSM8K
WiSE-FT on MATH500
WISE-FT on AIME24
0.75
0.47
0.75
0.84
0.46
0.5
0.4
0.70
Pass@8
0.82
Pass@
0.45
Pass@16
0.6
0.3
0.2
0:9
0.25
0.65
0.80
0.95
0.43
0.1
0.60
0.78
Gemma-2-28
0.42
Qwen-2.5-0.55
Qwen
0.575
0.600
0.625
0.18
0.19
0.28
0.30
0.32
Pass@1
Pass@1
Pass@1
Figure 2: Pass@1 vs. Pass@K across Interpolation Coefficients We perform WiSEFT with 5 E 0.91 between the first and last checkpoints of model (in legend) finetuned GSM8K, MATH, and OpenIhhuught--1144 114K, then evaluate on GSM8K, MATH500, and AIME24, re- spectively. Early SFT model observe higher Pass@K (y-axis) while later SFT model observes higher Pass@1 (x-axis). The interpolated model observe best of both metrics.
In practice, test-time compute is scaled by selecting one of K guesses either by output reward model (ORM) or Majority Vote. Then we can measure Best@K as
K
Best@K(x)
IB D(((())))(()))) 10:- = y] where 1" arg max KE|K]
Σ {9:
-;} or ORM(r;)
Notably, Pass@K is equivalent to Best@K using a perfect ORM verifier. As we will observe, WiSE-FT achieves both higher Pass@1 and Pass@K and this directly translates to achieving better Best@K with an ORM verifier and by Majority Vote.
3.2 Weight-Space Ensembling (WiSE-FT)
WiSE-FT is a weight-space ensembling technique proposed by Wortsman et al. (2022) to improve the out-of-distribution accuracy of finetuned models at no extra computational cost. In particular, while models tend to achieve better in-distribution performance after finetun- ing, they tend to be less robust to distribution shift. Surprisingly, by simply interpolating the weights of the finetuned model w, with the pretrained weights wi
(WISE(() 5 ww -- + (1 - 5) wi
(2)
WiSE-FT can achieve best of both words: the out-of-distribution accuracy of models im- proves without incurring drop in in-distribution accuracy. Similar to this philosophy, we apply weight ensembling to achieve both the diverse generation ability of early SFT checkpoints while maintaining the high Pass@1 accuracy of later SFT checkpoints.
3.3 Training and Evaluation Pipeline
The majority of our experiments are conducted on Gemma-2-2B and Qwen-2.5-0.5B We perform SFT on a 30K subset of rephrased augmentations of GSM8k (Cobbe et al., 2021) and MATH (Hendrycks et al., 2021) in MetaMath40k (Yu et al., 2023) for 1710 steps or 10 epochs. We then continue finetuning on another 30K subset of rephrased training questions from MetaMath using Group Relative Policy Optimization (GRPO) with a binary reward of the correctness of the model's 'ss final answer. Finally, we evaluate models on GSM8K and MATH500, respectively. To estimate the true Pass@K and Pass@1 marginalized over the distribution of sampled traces, we sample 100 reasoning traces per test example and average over them to estimate Pass@1, i.e. Pr. Then to calculate Pass@K, we use the theoretical formula - (1- - -x in Equation 1. Unless noted otherwise, we employ a naive decoding strategy with top-p threshold 0.9, temperature T 0.8, and top-k with K = 50.

---

Published as a conference paper at COLM 2025
MATH Majority Vote
MATH Best w/ ORM
RL Scaling on GSM8k
0.35
0.275
Best@k
0.250
Best@k
0.30
0.86
Pass
0.64
0.200
0.25
0.62
0.150
0.20
8
16
32
2
16
32
1x10*
2x10*
4x10*
SFT
WISE-FT
T=0.8
T=1.0
Synthetic Data
(a)
-
SFT Early
+
WISEFT
SFT Last
Competition Math (AIME24)
RL Scaling on MATH500
Pass@k
Pass-11
0.22
0.5
0.18
0.16
2x10"
4x10*
8x10'
16x10*
K
Synthetic Data
WisE-OpenThinker
OpenThinker
SFT Early
WISEFT
SFT Last
(b)
Figure 3: Downstream Advantages of WiSE-FT: (a) Best@K on MATH500 of the final SFT Gemma-2-28 checkpoint and its WiSE-FT counterpart. (b) Pass@K on AIME24 WiSE-FT after SFT on general purpose reasoning dataset OpenThoughts-1144 achieves higher Pass@K on AIME24. (c) RL Scaling Gemma and Qwen SFT checkpoints further tuned by GRPO on GSMSK and MATH, respectively. RL from the final WiSE-FT model achieves higher Pass@1 with less data compared to GRPO starting from both early and late SFT checkpoints.
Improving Diverse Reasoning Capabilities by WiSE-FT
We first carefully track Pass@K for E 3,33 across the SFT trajectory of Qwen-2.5-0.5B and Gemma-2-2B. Similar to findings from Cobbe al. (2021); Chen et al. (2025), we observe that Pass@1 continues to improve with longer SFT, whereas for larger K 4,32, Pass@K tends to peak much earlier on in training (in Figure 1, 17, and 19). In other words, while later SFT checkpoints achieve higher Pass@11 earlier SFT checkpoint achieve higher Pass@K. This tradeoff in model selection is not ideal downstream for test-time scaling.
Building upon this intuition, we propose weight ensembling between earlier and later SFT checkpoints. We apply a variant of WiSE-FT where instead of the pretrained model, we interpolate between the earliest SFT checkpoint (in case, after epoch of training) and the weights of later checkpoint. As shown in Figure 2, we observe a "sweet spot" of interpolation coefficients 5 E (0, 1) where the WiSE-FT model achieves both higher Pass@1 than the last SFT model and higher Pass@K than the early SFT model. We will fix 1/2, which generally performs decently for all of the datasets we've ve tested. In fact, after WiSE-FT (WiSE(1)) both Pass@1 and Pass@k grow monotonically with SFT steps f (see Figure 1).
Better Test-Time Scaling This boost in both Pass@1 and Pass@K directly translates to better performance with test-time scaling. We measure Best@K by Majority Vote and by selecting the reasoning trace with highest reward using off-the-shelf ORM ALHFlow/LLamal' Dee--RRR-Deepsekc
(Xiong et al., 2024). We evaluate the performance of the last SFT checkpoint with highest Pass@1 versus the corresponding WISE-FT variant

---

Published as a conference paper at COLM 2025
with 0 1/2. In Figure 3, we see that the performance gap MATH500 between the final Gemma-2-2B SFT checkpoint and Wise-FT model widens with larger K. The WiSE-FT model achieves 5 -7% better performance with test-time scaling.
Better RL Scaling WiSE-FT's ability to achieve both high Pass@1 and Pass@K is particu- larly advantageous for continued RL training where models are further trained by policy gradient methods using self-generated data. In particular, WiSE-FT is able to generate data rich in learning signal (high Pass@1) while still having high coverage over the data space (high Pass@K). We continue training on rephrased training questions of GSM8K and MATH using GRPO paired with a binary reward of the correctness of the final guess. Across runs, observe that continued RL training starting from the final WiSE-FT model improves performance more stably than finetuning starting from the final SFT checkpoint, Notably the final SFT checkpoint suffers low coverage over the data space, causing l'as@11 to improve slowly. We also try continued RL training from an earlier SFT checkpoint with peak Pass@4 performance. While RL scales better over the early SFT checkpoint in comparison the final checkpoint, the performance still remains subpar compared to WiSE-FT.
4.1 General Purpose Reasoning Models
we have studied the effect of WiSE-FT on models tuned on reasoning data for the same specific reasoning task (e.g,, train on GSM8k and evaluate on GSM8k). We've additionally tested how well our findings generalize to models trained on general purpose reasoning datasets and tested on a ouf-of--istribution reasoning task. We take Qwen2.5-7B-Instruu. and high-quality synthetic dataset of math, science,
So we
SFT for 5 epochs and coding questions paired with DeepSeek-RI completions, then evaluate its performance AIME24 competition problems (with ASY code for figures from Muennighoff et al. (2025)). In this setting, the Pass@K trends during SFT on more subtle. We still observe diversity collapse in Figure 12, but the affect is not strong enough for Pass@K to drop back down. However, we observe that the rate at which Pass@K improves for K E {16,32} slows down early while Pass@1 grows constant rate (Figure 10). We then perform WiSE-FT between the final and earlier checkpoint with higher diversity. We choose early checkpoint at epoch 3 where improvements in Pass@K begin to slow. Similarly, we observe that WiSE-FT improves both Pass@1 1 and Pass@K in Figure 2.
OpennToought--111
Diversity Collapse during Finetuning
In previous sections we alluded to the phenomenon where Pass@K decreases because SFT and RL induces diversity collapse in reasoning traces. To verify this hypothesis, we sample 100 traces per test GSM8k problem and measure diversity using three metrics:
1. Answer Diversity: The fraction of unique guesses & among reasoning traces.
2. Operation Diversity: The fraction of unique sequence of arithmetic operations performed among reasoning traces (In GSM8k, each intermediate step consists of basic arithmetic operation, e.g. + 8.).
3. Semantic Diversity: The average cosine similarity between the text embeddings of the reasoning traces, computed using Stella-400M-v55 (Zhang et al., 2024a)
As shown in Figure 4, we observe a stark trend where longer SFT on Gemma-2-2B incre- mentally suffers from clear diversity collapse across all diversity metrics. Specifically, the model places most of its probability mass not only on one particular guess, but on a single reasoning trace, evidenced by the reduced semantic and operation diversity.
5.1 Theoretical Discussion of Diversity Collapse During SFT and RL
We assess theoretically why diversity collapse tends to arise during SFT and RL training. Our analysis reveals that while SFT and RL operate on different principles, they share common pathways that lead to reduced generation diversity when optimizing for accuracy.

---

Published as a conference paper at COLM 2025
Diversity SFT [T=0.8]
Step 171
0.5
0.5
0.5
Step 342
0.4
0.4
0.4
Step
Density
Density
Step 684
0.3-
0.3
0.3
Step 855
0.2
Step
Step
0.1
0.1
0.1
Step 1368
Step 1539
0.0
0.0
0.0
0.0
0.5
0.00
0.02
0.04
0.0
1.0
Step
Operation Diversity
Semantic Diversity
Answer Diversity
Figure 4: Diversity Collapse The answer, semantic, and operation diversity of Gemma-2-2B reasoning traces across GSM8K test examples. Colors map to different SFT checkpoints.
K=1
K=5
K=10
K=50
0.80
0.72
0.48
0.78
Pass81
Pass@'S
0.70
Pars@10
Pass8
0.46
0.76
0.66
0.42
1000
2000
1000
2000
1000
2000
1000
2000
Steps
Steps
Steps
Steps
SFT
PPO from Step 513
PPO from Step 1026
PPO from Step 1710
Figure 5: Pass@k for SFT and RL of Qwen-2.--0.5B on GSM8K. The purple solid line measures Pass@K across SFT steps, while the dashed lines correspond to further training different checkpoints by Proximal Policy Optimization (PPO), While Pass@1 continues to improve, Pass@k for larger K can decrease even with RI.
Diversity Collapse during SFT Overparameterized models are well-known to exhibit overconfidence their predictions, an effect that has been studied extensively in classifica- tion (Guo et al., 2017). In particular, the model's confidence towards the most likely class PΦς = kmax x)) often much higher than the model's accuracy. In binary classification with linear models fx) = of (( x)) and linearly separable training data, gradient descent provably drives the norm of the weights infinity, causing probabilities to collapse or 1 (Soudry et al., 2018). We demonstrate this in linear models in Appendix A. A similar phenomenon likely arises in large reasoning models, which may also be prone overfitting during SFT, ultimately leading to overly confident solutions in spite of limited coverage over the space of traces (Cobbe et al., 2021).
Diversity Collapse during RL We further prove why applying reinforcement learning to low-diversity policy yields suboptimal results and sometimes even exacerbates diversity collapse-in in discrete bandit setting (see Figure 5). In this scenario, we assume there exist K equally good arms, corresponding to set of successful strategies, and one bad arm that the policy should learn avoid. We show two key results in this setting:
Implicit Collapse of Policy Diversity without KL Regularization. Our analysis demonstrates that when standard reinforcement learning applied without KL regularization, the training dynamics inevitably lead to a collapse in output diversity. Although multiple arms (actions) are equally optimal, the updates become self-enforcing as training progresses. Once one of the good arms is randomly reinforced, its probability increases at the expense of the others, ultimately driving the policy to converge on single-arm strategy (Theorem C.1).
and GRPO- are
Diversity Does Not Increase with KL. Regularization. When KL regularization incorporated to constrain the divergence from the initial policy in REINFORCE, the final policy longer collapses into a single-arm strategy. However, the diversity of the converged policy cannot exceed the initial diversity. Concretely, show that the probability

---

Published as a conference paper at COLM 2025
High diversity
Low diversity
Prompt: "1+1="
Samples: [2, 2, 2, 2]
Density
0.4
Prompt: "2+3="
0.2
Prompt: "2+3=" Samples: [6, 6, 6, 8]
Samples: [5, 5, 6, 8]
0.0
0
Pass@1 (px)
Pass@1 (px)
(a)
SFT induces
Bias Var
Temp Scaling: Blas t Var $
WiSE-FT Blas
Var
Step 171
Step 1710
0.4
T
1.0
1.5
WiSE-FT
0.4
0.4
0.3
0.3
Density
Density
Density
Density
0.33
0.2
0.41
0.2
0.2
0.35
0.1
0.1
0.1
0.0
0.0
0.0
0.0
0.5
1.0
0.0
1.0
1.0
Px
1 Px
1 - Px
1
(b)
Figure 6: Histogram of error 1- Px of Gemma-2-2B SFT checkpoints across GSM8K test. SFT progressively decreases bias but increases variance of error iee, 1 Pass@1, across the test distribution, causing Pass@K to fall. Applying Wise-FT reduces both bias and variance, but temperature scaling trades off decreasing variance with increased bias.
distribution over the good arms remains proportional to the initial distribution when the RL algorithm converges (Theorem C.8). This explains why initializing with diverse policy is critical for the generalization of reinforcement learning.
Bias- Variance Tradeoff of Pass@K
So far, we saw a mismatch in growth Pass@1 and Pass@K during SFT and alluded to the impact of diversity collapse to Pass@K. We now formalize the relationship between Pass@1, Pass@K, and diversity collapse. Notably, we show that the upper bound of expected Pass@K over the test distribution can be decomposed into bias and variance quantities.
6.1 Diversity Collapse leads to Bimodal Pass@1 Distribution
Consider the expected Pass@K over the entire test distribution x,y D. By Jensen's inequal- ity, we can derive a straightforward upper bound expected Pass@K that decomposes into the bias and variance of 1 Px (See proof in Appendix B). Note that the upper bound falls monotonically with larger bias and variance:
Proposition 6.1. Exy-D [Pass@K(x)) < ((Ery-D(1-p)) -aapp/11pllll px])2 + Var(px))
Bias
Variance
In Figure 6b, we plot the distribution of error 1 -xx, estimated using 100 sampled traces, over GSMSK test examples. We notice two trends with longer SFT. First, bias decreases, i.e., the expected error shifts towards 0. However, the distribution becomes increasingly bimodal with the densities converging towards the two extremes 0 and As a result, the variance increases with longer SFT. This increase in variance directly explains the drop in Pass@k.
8

---

Published as a conference paper at COLM 2025
The bimodality of the 1 - fx distribution means that the Pass@1 of any test problem is either very high or very low. Interestingly, one explanation for the increased bimodality of the distribution of 1 Px is in fact when models suffer from diversity collapse. In other words, particular guess to be oversampled for each test problem. If the model places high probability on incorrect guess, Pass@1 is very low. On the other hand, if the model places high probability on the correct guess, Pass@1 is very high. We illustrate this relationship in Figure 6a. Áll in all, Pass@K can be improved in two ways either reduce bias by improving Pass@1 or reduce variance by increasing diversity.
6.2 WiSE-FT vs. Diverse Decoding
While we've proposed inducing diversity by WiSE-FT, another common alternative for inducing diversity is temperature scaling the logits. High temperature smoothens the logits allowing the model to more likely sample low probability tokens. In Figure 1, we see that while high temperatures indeed improve Pass@K, the Pass@K at any SFT timestep notably never reaches the Pass@K of our final WiSE-FT model. If temperature scaling also increases diversity, why does WISE-FT strictly outperform sampling with high temperature? In Figure 6b, we plot the distribution of 1 Px if we sample from the last SFT checkpoint with high temperature
1.5. As expected, we see that the model reasons more diversely. This smoothens the bimodal peaks and reduces the variance. However, the average accuracy of the model generations also degrades, causing the bias goes back up. We suspect bias- variance tradeoff is inherent in diversity-nnducing decoding approaches. For example, min-p (Nguyen et al., 2024) combines temperature scaling with adaptive thresholding to not sample outlier tokens. However, this additional control is unable to reduce bias (Figure 16). Surprisingly, WiSE-FT uniquely manages to reduce both bias and variance.
Discussion
In this work, we investigated the phenomenon of diversity collapse during the training of reasoning models. Our analysis reveals that standard SFT and RL pipelines can deteriorate in Pass@@K due to the convergence of model generations toward a single reasoning trace. We demonstrated that WISE-FT, which interpolates between early and late SFT checkpoints, significantly improves both Pass@1 1 and Pass@K K across multiple math datasets and model scales. This unlike alternative approaches such as temperature scaling or early stopping, which face an inherent tradeoff. Furthermore, improving on these metrics corresponded with better adaptation to test-time scaling and RL. But other limitations of WiSE-FT may exist at larger scale, which we leave for future work.
Overall, our work reveals the importance of maintaining diversity in reasoning models. Current decoding strate- gies (e.g., min-p, nucleus, and top-k) are still unable to fully extract model's 's capabilities. We estimate that significant gap, of tens of percent, remains compared to the optimal decoding strategy for Pass@K, i.e., top-K sampling over the model's marginal answer distribution P( x) (see Table 1 and Appendix G). We encourage future works to address downstream limitations more carefully in earlier stages of the training pipeline.
Method
Pass@2
Pass@4
Nucleus
0.57
0.67
Min-p
0.57
0.67
Top-k
0.56
0.67
Optimal 0.76
0.83
Table 1: Best Pass@k of Gemma on GSMSk across SFT checkpoints
Acknowledgements
We'd like to thank Aviral Kumar, Sean Welleck, Amrith Setlur and Yiding Jiang for insightful discussions about test-time scaling and reinforcement learning. We'd also like to thank Alex Li, Sachin Goyal, and Jacob Springer for their meaningful contribution to our figures and literature review. We gratefully acknowledge support from Apple, Google, Cisco, OpenAI, NSF, Okawa foundation, the A12050 program at Schmidt Sciences (Grant #G2264481), and Bosch Center for AI.

---

Published as a conference paper at COLM 2025
References
AlphaProof and AlphaGiometry teams. Ai achieves silver-medal standard solving inter-
national mathematical olympiad problems, jul 2024. URL https://deepminnd google//
discover/blog/ait sollle/sa-soooss-eesoose imo
sonoovee-----poollmm------- problems-at-siiee st--siieer--mddallleeee
Edward Beeching, Lewis Tunstall, and Sasha Rush. Scaling test-time compute
with open models, 2024. URL
blogpost-scaling secaalingeeett-eeeeeeeeeeeeeeeee wew-ttimeecopuuue
Jeff Bilmes. Submodularity in machine learning and artificial intelligence. arvi preprint axXit:2202.00132, 2022.
Feng Chen, Allan Raventos, Nan Cheng, Surya Ganguli, and Shaul Druckmann. Rethinking fine-tuning when scaling test-time compute: Limiting confidence improves mathematical reasoning. arXiv preprint ar arXio:2502.07154. 2025.
Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Alphamath almost zero: Process supervision without process, 2024. URL https://arxiv.
org/abs/2405. http:////aarvvooorrb//////////
03553.
Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagara- jan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Aleksandra Faust. Inference- aware fine-tuning for best-of-n sampling in large language models. arXiv preprint ar arXiv:2412.15287, 2024a.
Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Aleksandra Fauss. Inference-aware fine-tuning for best-of-n sampling in large language models, 2024b. URL httpps:/arxiit wrg/abs/2412.1 15287.
Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shengbang Tong, Saining Xie, Dale Schuurmans, Quoe V. Le, Sergey Levine, and Yi Ma. Sft memorizes, rl generalizes: comparative study of foundation model post-training, 2025. URL
wttpa//www....///////..... 1111/
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv. wtwp:iivvorg//////////// abpp//0bs//11..... 14168.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pp. 1321-1330. PMLR, 2017.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning, arviv preprint ar ari::r2501. 2025.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset, 2021. URL https: //axiiv/////vvva///////
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration, 2020. URL http:////aaxiv.oggb//////////// wrgg/bb/////...../// 03874.
Audrey Huang, Adam Block, Dylan J Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz, Jordan T Ash, and Akshay Krishnamurthy, Self-mmrrvemmen in language models: The sharpening mechanism. aiv preprint ar arXi::2412.01951, 2024.
Yifei Li, Zeqi Lin, Shizhuo Zhang, Qiang Fu, Bei Chen, Jian-Guang Lou, and Weizhu Chen. Making large language models better reasoners with step-aware verifier, 2023. URL
http://xxiv.oog//s////////// 02336.
Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Zhi-Quan Luo, and Ruoyu Sun. Preserving diversity in supervised fine-tuning of large language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https: http://///////ww////eemee net/
10

---

Published as a conference paper at COLM 2025
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, llya Sutskever, and Karl Cobbe. Let's verify step by step. arviv preprint ar Saii:2335.200555 2023.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. s1: Simple hwrgg/bb/25000 19393.
test-time scaling, 2025. URL
Minh Nguyen, Andrew Baker, Clement Neo, Allen Roush, Andreas Kirsch, and Ravid Shwartz-Ziv. Turning up the heat: Min-p sampling for creative and coherent 11m outputs, 2024. URL https::/aaxxi.orrg/aa////// wrgp///////0..00000/// 01082.
Theo X. Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar-L.ezama. Is self-repair silver bullet for code generation?, 2024. URL https:
warxxi..org/ab//28/e080 09896.
Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amélie Héliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev, and Olivier Bachem. Bond: Aligning lims with best-of-n distillation, 2024. URL https: org/
Amrith Setlur, Chirag Nagpal, Adam Fisch, Xinyang Geng, Jacob Fisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. Rewarding progress: Scaling automated process verifiers for Ilm reasoning, 2024. URL https://arxiv.org abs/2410.08141
Louis Shao, Stephan Gouws, Denny Britz, Anna Goldie, Brian Strope, and Ray Kurzweil. Generating high-quality and informative conversation responses with sequence-to- sequence models. Xiv preprint ar XXio:1701.03185, 2017.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters, 2024. URL https:// wwwiv.org//aee wtw../////////...../ 033144
Yuda Song, Hanlin Zhang, Carson Fisenach, Sham Kakade, Dean Foster, and Udaya Ghai. Mind the gap: Examining the self-improvement capabilities of large language models. arXiv preprint XxXio:2412.02674 2024.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018.
Ashwin K Vijayakumar, Michael Cogswell, Ramprasath R. Selvaraju, Qing Sun, Stefan Lee, David Crandall, and Dhruv Batra. Diverse beam search: Decoding diverse solutions from neural sequence models, 2018. URL h/tps://axx:i http:///aaxvvvoooros////ss// 116aaa////11... 02424.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models, 2023. URL https::/arxxv.orr//bb /200. 11171.
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo-Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models, 2022. URL https://arxiv. org/abs/2109. .10099.....0 01903.
Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. Inference scaling laws: An empirical analysis of compute-optimal inference for problem-solving with language models. Xiv preprint ar Kiv:2408.00724, 2024.
Wei Xiong, Hanning Zhang, Nan Jiang, and Tong Zhang. An implementation of generative prm. https: t/gitbb...
thub. www/RLHFOow//RL.... /wwp//LHH---. Modd-Rewad-ddllliiee 2024.
11

---

Published as a conference paper at COLM 2025
Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long chain-of-thought reasoning in lims, 2025. URL httpp:://aaiiii org.
org/ abs/ 2502.03373.
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James
Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. ar Xiv preprint ar trXi::2309.1228, 2023.
Dun Zhang, Jiacheng Li, Ziyang Zeng, and Fulong Wang. Jasper and stella: distillation of sota embedding models. aiv preprint arXio:2412.19048 2024a.
Yiming Zhang, Avi Schwarzschild, Nicholas Carlini, Zico Kolter, and Daphne Ippolito. Forcing diffuse distributions out of language models, 2024b. URL https://arxiv..or abs/2404. 10859.
12

---

Published as a conference paper at COLM 2025
A SFT in Binary Classification
Data and Model Setup We train a linear classifier fx) = (w,, x) from random initialization over binary Gaussian mixture distribution:
x y N Nyp,, pdxd)
(3)
y {1, -13 uniformly
(4)
Given a model, we sample predictions, namely 9 = 1 with probability x(wwxx)) = (1 +
exp(-(w, (w,x))) or 9 0. Then, per-example Pass@1 is equal to fx oly (w, (w,x)) ) Similarly, the expected Pass@k is equal to 1 (1 -p)x
Weight Norm Over Steps
In our experiment, we train an overparametrized linear classifier over binary Gaussian data mixture y ~ My 11 where (-1,11 and d = 1000. We then evaluate Px of 400 test samples. As training progresses, the distribution of Px over the test data become bimodal due to the norm of w monotonically increasing once it separates the training examples. Similarly, we observe that this leads to a drop in Pass@k while Pass@1 continues to improve.
1.2
1.0
250
Gradient Steps
Figure 7: Weight Norm
Pass@11
Pass@2
Pass@
Pass@8
Pass@1 16
0.55
Pass.8k
0.70
Pass@gk
Pacs@ik
0.925
Pass@Bk
0.96
0.82
0.68
0.80
0.875
0.66
250 Gradient Steps
250
250
250
250
Gradient Steps
Gradient Steps
Gradient Steps
Gradient Steps
Figure 8: Pass@k across Training in Binary Classification
Steps
Steps
Steps
Steps
Steps
30
30
75-
Count
20
Count
20
Count
Count
Count
50
10
10-
25
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Figure 9: Histogram of px across training steps
B Expected Pass@k
Proposition B.1.
E Ex [Pass@K()) 1 -(Exx-DD px])2 + Var(p.))
Proof.
> (1 -xxx22 2 k/2
E (1- px)
(5)
(1 21E [px] + [o2] k/2
(6)
((1 21E [x*] E 1px (0x12) + EE 02
E spxx1) k/2
(7)
= ((1 - E [px])2 + Var \//2
(8)
13

---

Published as a conference paper at COLM 2025
C RL Theory
C.1 Overview
We will prove that in discrete bandit setting with K equally good that is the best arm, both REINFORCE and GRPO without KL regularization will eventually collapse into single-arm strategy.
We will further prove that, with KL regularization with respect to the initial policy, the converged policy of REINFORCE have the same action distribution as the initial policy when constrained on the set of best arms. Therefore, diversity within good actions will not increase through REINFORCE training.
C.2 Notations and Setup
Formally we consider the following setting, We consider K 1-armed bandit, with arms (1,2,..., K 1}. Arms I,..., are "good," each yielding reward 1, and the other arm is "bad," yielding reward O. We use softmax parameterization:
04
Pr
i= 1= K+
K+1 2j=1
to denote the action distribution. We will use o!!) to denote the parameter at step t.
It is standard to consider using the KL divergence between the current policy with reference policy (which we set as Po here) as regularization term.
(I(p0()p00))
K-1
log
Pi
(F)
(0)
P
For REINFORCE, we will consider the following training setup. At step t:
1. We sample an arm I, according to p(-) (p] ,..., PK:1) and receive reward
2. We update using policy gradient.
el (1++1)
8!') + 1 Ve, (log ) - nBV (p 1
where 17 > is the step size and & is the hyperparameter controlling the strength of KL regulaiization
For GRPO, we will consider the following simplified training setup. This equivalent to the empirical version of GRPO with online sampling.
1. Sample G arms {1," .....,,, (...,,,()) i..dd. from the current policy p(-) and receive rewards 7
2. Compute
Nr
1 G G
10
Σ
of
(8) m1)) 1 8=1
1
(1)
8-1
and define the normalized advantage
(g)
#4
(s)
0
+ 0,
0. Tt
0 0.
We will skip the update if 0.
3. Update each parameter 0; via
G
+ 0
+
1
Σ
p
(lg ((8))
naV 0. KL (L(101)p(0))
i= -11,.,,+11.
14

---

Published as a conference paper at COLM 2025
C.3 Implicit Diversity Collapse without KL regularization
Theorem C.1 (Collapse to Deterministic Policy). Under REINFORCE or GRPO updates without KL regularization (Bo- 0), given a sufficient small y,, with probability 1:
lim sup max p 1.
(E[K]
Thus, the policy collapses to a single-arm strategy during training.
Proof. The proof two-fold.
Using Lemma C.3 and C.4, can show that bad arm probability diminishes,
(f) f-x00 K+1 = 0
lim
We will then define property named Self-enforcing Stochastic
Definition C.2 (Self-enforcing Stochastic Policy Update Rule). We define three properties of policy update rule that will lead to diversity collapse
1. The policy update takes the form of ER Ax log (i((.)) where in is the k-th sampled arm in the batch and Ax a function determined by (i) the sum of reward Σ 1 Vi. with in the batch; (ii) the reward rig and (iii) the batch size B.
2. A policy update rule is said to be self-enforcing, if [B(0 - -0f"1 is monotonous with e for all E [K] and t. Further E(o((+1) - =(() is non-positive if K and is non-negative if K.
3. A policy update rule is said to be self-enforcing stochastic if it is self-enforcing and there exists constants C₁ C2 > 0 such that for any E whenever the current pol-
icy satisfies [ P; (1) [1/2K,1 c] (i.e., single good arm dominates), for
i* = arg max;e[K] p the conditional second moment of the parameter updates for every arm i E [K + 1] and i + satisffes::
IE ((0/+t)) o!!) (0((+1) o(())) ((() 2 C.
and
l0!(T)) e!01 C2
Lemma C.5 shows that for any self-enforcing stochastic policy update rule, the final policy collapses into single-arm policy.
Using Lemma C.6 and C.7, we can show that REINFORCE and GRPO are self-enforcing stochastic policy update rules when bad arm probability is lower than 1/2. The proof is then complete.
Lemma C.3 (Bad Arm Probability Diminishes Using REINFORCE). Under the REINFORCE (t)
algorithm without KL. regularization (B 0), lim;- you Pk+1 0 almost surely.
Proof. We can first simplify the REINFORCE update rule to
oft+ o!0) + p1 = -)),, i= 1, 1.
Noted that Σ, 0 will not change with f, WLOG, assume
Σo 0.
15

---

Published as a conference paper at COLM 2025
Because TK+1 O, we can then assume without loss of generality, for all t, 4 S K. This then suggests that
of 1) K+1 =00 (f)
4PK+1
monotonically decrease.
(t)
For any E, if Pk+1 holds for infinite I, then there exists fo, where k+1 log for any t to. For any to, there exists [K], such that 0!! > O. This then suggests that
( Pk+1 S exp(0k4 K+1
of*) S c.
This leads to a contradiction. The proof is then complete.
Lemma C.4 (Bad Arm Probability Diminishes Using GRPO). Under the GRPO algorithm without KL regularization (B = 0) Pk+1 = almost surely.
Proof. For GRPO, can show that x(8) is negative iff 1,8) K + 1. Therefore, show @!))
that K+1 monotonically decreases, similar to the case in REINFORCE.
(1)
Px PK11 holds for some one can prove that KK1 will decrease by a constant depend-
ing on in expectation. Therefore, following the same line as in C.3, we can prove that (1)
lim P K+1
0 almost surely.
Lemma C.5 (Collapse Happens for All Self-enforcing Stochastic Policy Update Rule). Consider a policy update process that is self-enforcing stochastic (Definition C.2), then lim you max p! = 1 almost surely.
Proof. We will inductively prove that for different K the following induction hypotheses, for any e, 5 > there exists Te,8,K > 0,
(1) Pr( max max p; <1 - <6. KT.Ak iE[K]
We first consider the case where K Consider the stopping time,
Te = arg min max p!d > ie(K]
For any I = ( define A/- maxje(K] 0f minjez 0%.
Assume 0f = maxjelK 04 because |Z] 2 2, there exists i # i", minjez 0f > 0. We will show three properties of A
First 4 is a submartingale defined on the filtration of the distribution of (()) because
E[A
0.] A (1-1) E[( - 0%.) - (0) 07) 0.] > 0.
as the policy is self-enforcing.
16

---

Published as a conference paper at COLM 2025
Further AT has bounded growth of 2C₂ as
max 0 je[k] /
max 04) C2.
jElk
min 01+1 JEZ
max f9 <2. jCI
Furthermore, the second-momentum of 47 needs to increase with f by constant for any
t < Te.
E[(AA [((AA-1))2 0:] 2 (A + E[(AZ ([(A)t)) (1+1) -1())
2 )2 +GE2.
When < Te, holds that A log E otherwise we can prove that
max pi/pj (xp(4A 2- [JE(122
max Pi > -1c c.
1E(1,)
This a contradiction. Further, by Martingale inequality, we have that
E( (mmm(((()))) )9
Tel]
Further, as A has bounded growth, we have that
2 + 2C2) (((A A (in((txt)) 1) < (log E
This implies E[min(t, tel] < (log +202) G
for all this implies
(log 2 22
E E[Te]
C₁E2
Further, by Markov inequality, if we choose
(log 2 + 2C2)
Te
C₁e28
then,
E(Te] <8.
Pr(te > Tes2) <
Te
This concludes the proof for K 2.
Now assuming the result holds for - and consider the case for K, First, we choose
small enough constant CJÆ,K,N 0, such that when PK (0) random processes are close:
Co the following two
Running the algorithm for N steps the K arms bandit yields i [K]
Running the algorithm for N steps K 1 arms bandit yields 81"), E (K-1) 1]
with क!0) 0 ; K - and k 0x(
17

---

Published as a conference paper at COLM 2025
and there exists joint measure on I and â such that
Vi E K - 2), < N,Pr(pp) - pf1 e/2) < 8/66
Pr(1pk - PK > e/2) < 8/6.
Pr(lpk pk1 > e/2) < 8/66
This joint measure is constructed by choosing the corresponding arm for two process at each sampling step as long as the sampled arm is not K and uses the uniform convergence on V logo Pi- Now following the same argument at K 2, we can show that there exists Te,,,K such that
Pr(3t < TesK, min < <oo.KTepSSPX-)) > 1-/2.
IElK
Then can invoke the induction hypothesis and uses the coupling shown above show that if we choose Te,á,K Te,a, + T2/25/2.K1,, then there exists time step that one arm has probability higher than 1 - with probability at least 1 - 8.
Lemma C.6. The REINFORCE algorithm without KL regularization (B 0) self-enforcing stochastic (Definition C.2) once Pk+1 1/2.
Proof. The REINFORCE algorithm is self-enforcing because Bio! 1) 0f' =
Σ p7t))
je
Further,
10 11] (
oft) I
(oft+1) 0ft) (ol(+t) d!l)
and if we consider the distribution of Ai,
it holds that
7
A - (1(i I1) - 1(i" I1] -p; p++)
Pr -1 Pr (I+ =r) = Pi
Therefore
B[A 2 Pi (-1 - pi + p?)²
62
2 pi (1 pr )2 2K
This then concludes the proof with C = 1/ / and C2 =
Lemma C.7. The GRPO algorithm without KL regularization (6 0) is self-enforcing stochastic (Definition C.2) once Pk+1 1/2.
Proof. The GRPO algorithm is self-enforcing because
E[ - o! = =
11111 (1(1(s) = -- 1 (1! = 1)] = = (E7 ((1()) -]]
Noted that 11 e)]p is monotonous with p;, hence monotonous with 0,-
18

---

Published as a conference paper at COLM 2025
Further
10((11) - 0 - 1 max F/8) (1(1 = - -p)-p!") |
g
| S nIVG.
S max g
Now we only need to lower bound the second momentum of
(o/(t1) 0(-)
(of 0(0)
Noted that
1 G G g=1
0
(++1)
0
18)
1
(8) =1).
It holds that
o
1
G
Σ
(rf - p1) =
V
Σ
2
+ 122 = v 122.
8
g
Therefore when
> 0,
(()
F(8)
14
1 - 14
1 1
pt
1
2 G-1 1
0f
0f
Mt
Because all p(8) are the same when r((8
> 0, it holds that when i E
[K],
2
4 G Ht
11 141
Σ
G
1 = -1)10)))))
8-1
G
Σ 1(1(5) = - -1(((5)
G 1)
8=1
This then implies
Σ 111((0 = -i)- -1(480 = -3) )
2
E E
E
Ms + 1,0
G(G
One can without loss of generality assume (( = K + and show that
Σ ((T(5) = 1) - -1(1((0= 27
G-1
2 2 2
E C(G-1)
8=1
2 1 (1(1)) =) - 1(1,)) = :r))
Pi + P >
& E
2KG
When * K, noted that (of(tT) a similar bound can show that IE EE[4? ,1, > ZKC /2KKG and C2 = VG.
0.t) ( ( o(-1
00 (11
0.t) > Therefore,
1 This then concludes the proof with C =
19

---

Published as a conference paper at COLM 2025
C.4 Diversity Never Improves with KL regularization
Theorem C.8 (Diversity Preservation under KL Regularization). With Po as the initial policy and KL-regularization hyperparameter & 0, if the RÉINFORCE process converges to policy p Then, p" satisfies:
p²(1)
Po(i)
Vi (I,...,K}.
P-
E Po-)
Consequently, the distribution over the optimal arms under p* matches the initial distribution po restricted these arms and renormalized.
Proof. Using policy gradient theorem, know that the converged policy p" and corre- sponding parameter 0 satisfy that,
K+1
TiP: + SKL (pIPo)
0
0
8=8
This then suggests that for any
K-1
K+1
TRP
Pk
B
e, lp: p: p:l
1=1
This is equivalent to
K+1 Σ
K-1
rEPE pk
pp +
Σ
(1 = k) pt) (log pf + 1-llog
log = 0
Simplifying
K+1
K+1 1
tk+ + B(log pt + log Po)
Σ
+ B
Σ
pf (log +1 1o-oogg))
For all k E [K], we know that rk is equivalent, therefore, concluding our proof.
P (i) [6(
constant for k E [K],
C.5 Technical Lemma
Lemma C.9. For | C, it holds that
exp(x) > x Acx²
here Ac = exp(-C)+C-1 CA
exp(x)-1-x x²
Proof. Define g(x)
this function monotonically increase when
20

---

Published as a conference paper at COLM 2025
D Open-Thoughts Evaluation
We full finetune Dwen2.5-7B-Instruue over Open for 5 epochs using BF16 and A with hyperparameters -earning-rate=le55 batch-size=128, warmup= steps. We sample 40 reasoning traces with temperature set to 0.7 for each of the 30 problems in AIME24. Then we evaluate the following quantities.
Competition Math (AIME24)
@ass@1
Pass@32
Pass@64
0.4
0.8
0.8
Pass@k
0.3
Pass@K
Pass@k
0.7
0.6
0.6
0.2
0.5
0.5
Pass@64
0.4
2
3
5
1
3
4
5
1
4
Epoch
Epoch
Epoch
Figure 10: Pass@K Evaluated on AIME24 over OpenThhoghhs-114K SFT checkpoints. We plot the expected Pass@K + SD. Note that the rate at which Pass@K slows down while Pass@1 improves constant rate. Furthermore, note that the confidence interval of Pass@1 widens, meaning the variance increases during SFT.
Checkpoint 500
Checkpoint 1500
Checkpoint 2500
Checkpoint 3500
Checkpoint 4500
20
25
0.18
0.23 = 0.31
0.27
25
0.30
0.33
25
0 0.34
20
Count
Count
20
20
10
0.0
0.5
Figure 11: Histogram Pass@1 over AIME24. Variance of Pass@1 increases over finetuning on ThonThouughss- 114K. We note that since AIME24 only has 30 questions, the density plot may not be completely reliable.
Competition Math (AIME24)
Ratio of unique answers
0.500
0.475
0.450
0.425
0.400
1
2
Epoch
Figure 12: We plot the average number of unique answers sampled over the total number
samples i.e. [
Model samples less diverse number of answers as SFT progresses.
21

---

Published as a conference paper at COLM 2025
E Interpolation Coefficients
WiSE-Step672 on MATH500
WISE-Step672 on MATH500
WISE-Step672 on MATH500
WiSE-Step672 on MATH500
0.63
Pass@e
0.36
0.75
0.25
Passage
0.45
0:25
0.55
0.54
0.25
Pass
0.62
0.35 0.9
0.44
0.53
0.61
0,2
0.43
0,4
0.52
0,1
0,4
0.19
0.20
0.19
0.20
0.19
0.20
0.19
0.20
Pass@1
Passip1 1
Pass@1
Pass@]
WISE Step--5 an MATH500
WISE Stepp96 on MATH500
WWSS--Step99 on MATH500
WiSE- Step896 on MATH500
12
9.
02
0,5
0
0.55
Pa3SE4
Pass@8
0.45
0.54
Passe
0.62
0.35
0.25
bgg
0g55
0.34
0.43
0.52
0.60
0.18
0.19
0.19
0.20
0.18
0.19
0.20
0.18
0.19
Pass@1
Pass@1
WISE WISE-Step1120 MATH500
WiSE-Step on MATH500 WiSe-Stepp120 MATH500 WISE WISE-Stepl120 MATH500
0.37
0.36
P85si4
0.36
0.35
0,9
Pass@e
0.45
0.75
0.54
0.55
0,55
Passe
0,9
0,25
0.44
0.53
0.61
0.34
0.25
0.25
0.60
0.25
0.33 6.7
0.51
0.20
0.19
0.20
Pasc@1
Pass@1
Pass@1
Pass@1
Figure 13: Pass@1 versus Pass@K of WISEFT of Qwen-2.5-0.5B trained and evaluated MATH500. We interpolate between model wy at Step 112 with w, for E [672, 896, 1120] as 5wo + (1 (1-5)w; where 5 E [0.1,
22

---

Published as a conference paper at COLM 2025
StSE-Step342 GSMBK
WISE StSE-Step342 GSMBK
WisE-Step342 GSMBK
WISE-Step342 GSMBK
0.925
0.75
0,25
0.845
0.890
0,25
0.920
0.28
Passe
0.775
0,75
Pass@e
0.885
0.28
Pass832
0.915
0.25
0.770
0.035
0.
0.765
0.830
0.875
0.910
24
0.65
0.825
0.55
0.05
0.905
0.05
0.56
0.58
0.56
0.50
0.56
0.50
0.56
0.58
Pass@ 1
Pask@1
Pass@1
P
WWSE-Step684 on GSMBK
on GSMBK
WISE- WISE-Step6al Step684 on GSMBk
MASE--tep684. 84 on GSMBK
0.79
9.3
0.85
0.75
0.89 R.9
0.92
0.9
0,75
0.75
Pass@g4
0.78
0.25
Passigle
0.25
Pass91
0.88
Pass
0.91
0.25
0.83
0.77
0.87
0.90
0.82
0.76
0.95
0.05
0.55
0.p5
0.56
0.58
0.60
0.62
0.50
0.60
0.62
0.56
0.58
0.60
0.62
0.56
0.58
0.62
Pass@1 1
Pass@1
Pass@11
Pass
WiSE-Step1026 on GSM8k
WiSE-Step1026 GSM8K
WiSE-Step1026 on GSM8K
WiSE-Step1024 on GSM8K
0,5
0.85
0.79
0.9
Pass84
0,25
Pass
0.25
Pass@16
0.25
Pessi932
0.91
0,25
0.83
0.87
0.77
2.9
0.82
0.86
0.89
0.05
0.56
0.58
0.60
0.62
0.56
0.60
0.62
0.56
0.60
0.56
0.58
0.60
0.62
Pass@]
Pass@11
Pass@9
Pass@1
WISE-SKA Step1368 GSMBk
WISE WISE-Step1366 GSMBK
WISE-Step1
GSMBK
WiSE WISE-Step1368 GSMBK
0,75
0.92
0.75
0,5
0.79
0,75
0.80
Pass
0.78
0,25
Pass@aa
0.25
Pass@16
0,25
Pass@32
0.90
0.9
0.86
0.76
0.05
0.05
0.84
0.65
0.05
0.58
0.60
0.62
0.58
0.6.0
0.62
0.58
0.60
0.62
0.58
0.60
0.62
Pass@]
Pass(F1
Pass@9
Pass
WISE WISE-Step1710 on GSMBk
WISE-Step1705 110 on GSMBK 0.25
WISE: WWSS--teep17 710 on GSMBk
WiSE- WiSE-Step1710 on GSMBk
0.75
0.75
0.84
0,3
0,5
Pass@g44
0.5s
Pass-gi8
Pass@16
Poss:332
0.90
0.82
0.86
0,25
0,5
0.76
0.80
0.84
0.05
0.05
0.05
0.05
0.82
0.58 0.60 0.62
0.58 0.60
0.62
0.58
0.60
0.62
0.58
0.50 0.62
P
Pass@1
Pass@1
Pass
Figure 14: Pass@1 versus Pass@K of WISEFT of Gemma-2-2B trained and evalu- ated on GSM8K. We interpolate between model Wo at Step 171 with for t E [342, 684, 1026, 1368, 1710] as dwn- + (1-8)) 8)w, where å E [0.05, 0.9].
23

---

Published as a conference paper at COLM 2025
F Measuring Diversity of Traces
We measure the diversity of the 100 sampled traces of Gemma-2-2B across GSM8K test. We measure diversity in terms of different measures.
Output Diversity The cardinality or number of unique answers in the of all model outputs [91, 92..... .n]] over the total number of traces.
Operation Diversity In GSM8K, each intermediate step consists of basic arithmetic op- erations, e.g. 3 8. We may simply map each of the traces to the sequence of arithmetic operations the model steps through, iee + (2,,02.... .r]. This mapping extracted by code. Then, given this set, we measure unique sequence of operations over the number of total traces.
Semantic Diversity We measure the similarity of trace using cosine similarities between the text-embeddings (Bilmes, 2022; Yu al., 2023).
E1 Does temperature increase diversity?
Temperature does increase diversity, but it also increases the chances of sampling outlier
answers.
Diversity Across SFT [T=0.8]
Step
0.5
0.5
Step
0.4
0.4
0.4
Step
Density
Density
Density
Step
0.3
0.3
0.3
Step
Step
Step 1197
0.1
0.1
Step 1368
Step 1539
0.0
0.0
0.0
0.0
0.5
0.00
0.02
0.04
0.06
0.0
0.5
1.0
Step 1710
Operation Diversity
Semantic Diversity
Answer Diversity
Diversity Across SFT [T=1.0]
Step 171
0.5
0.5
0.5
Step 342
0.4
0.4
Step 513
Density
Density
Density
Step 684
Step 855
Step 1026
0.2
0.2
Step 1197
0.1
0.1
0.1
Step 1368
Step 1539
0.0
0.0
0.02
1710
Operation Diversity
Semantic Diversity
Answer Diversity
Diversity Across SFT [T=1.5)
Step 171
0.5
0.5
Step 342
Step 513
Density
Density
Density
Step 684
0.3
0.3
0.3
Step 855
Step 1026
0.2
0.2
0.2
Step 1197
0.1
0.1
0.1
Step 1368
Step 1539
0.0
0.0
0.5
1.0
0.00
0.02
0.04
0.06
0.0
0.5
1.0
Step 1710
Operation Diversity
Semantic Diversity
Answer Diversity
Figure 15: Diversity of traces sampled with Temperature SFT checkpoints on GSM8k
{0.8, 1.0, 1.5) for Gemma-2-2B
24

---

Published as a conference paper at COLM 2025
F2 How well do token-level diverse decoding strategies compare with optimal strategy with oracle?
Hyperparameter Tuning Details We grid search for optimal temperature for all baselines over T = [0.8, 1.2, 1.5, 1.8]. For nucleus, we choose the best cutoff threshold between [088, 0.9, 0.95). For min-p, we choose the best probability threshold between [0.01, 0.05, 0.1]. For tokenwise top-k, we choose best k between [12, 22,,2,,55 50).
Decoding Strategy
Pass@2
Pass@4
Pass@8
Naive
0.565
0.666
0.760
Nucleus
0.566
0.668
0.757
Min-p
0.566
0.668
0.760
Top-k
0.563
0.666
0.756
Top-k w/Oracle
0.760
0.832
0.901
Table 2: Best Pass@k of Sampling Strategies for Qwen-2.5-0.5B over SFT checkpoints
Decoding Strategy
Pass@2
Pass@4
Pass@8
Naive
0.547
0.648
0.737
Nucleus
0,528
0.617
0.694
Min-p
0,550
0.655
0.744
Top-k
0.538
0.646
0.738
Top-k w /Orralee
0.730
0.814
0.878
Table Pass@k of Sampling Strategies for Qwen-2.5-0.5B Last SFT Checkpoint
Pass@1
Pass@8
Pass 16
0.900
mm-p(f-1.0)
0.850
minpT-1133)
@k
@
Pass @K
0.875
mmnpT=1..0
0.825
0.5
0.850
0.800
0.825
0.775
0.800
500 1000 1500 Checkpoint
500 1000 1500 Checkpoint
500 1000 1500
Checkpoint
Figure 16: Pass@K over different Min-P thresholds y E [0,0.3] and temperatures T [1, 1.6) for Gemma-2-2B finetuned on GSM8K. Generally, no min-p threshold paired with high temperature T 1.6 (in light green) is able surpass the Pass@1 of T 1 with best min-p threshold (in orange). In other words, unlike WiSE-FT which increases both Pass@1 and Pass@K, Pass@1 tends to still decrease for the diverse decoding strategy of applying min-p with high temperature.
25

---

Published as a conference paper at COLM 2025
Pass@1
Pass@2
Pass
Pass
0.02
0.625
0.725
Passe1
0.600
Pars
0,700
Pa3S84
Pass88
0.575
0.82
0.630
0.80
1000
Checkpolnt
Checkpoint
Checspoint
Checkpoint
T-0.8
T-1.0
T-1.2
T-1.5
T-188
Figure 17: Pass@k of Gemma-2-2B GSM8k Naive Sampling with Replacement
Pass@1
Pass@2
Pass@4
Pass@8
0.80
0.94
0.775
0.98
0.86
0.92
Passe
0.96
0.84
0.90
0.725
0.82
0.94
0.88
0.700
0.92
0.86
0.90
0.675
1500
500
1000
Checkpoint
Checkpoint
Checkpoint
Checkpoint
0=0.8
T=10
T=12
T=1.5
T=18
Figure 18: Pass@k of Gemma-2-2B GSM8k Oracle Top K Sampling
Pass
Pass@2
Pass/24
Pass
0.5
0.6
0.70
Pessi8
0.4
Pasor
0.5
0.60
Pawag
0.4
0.50
0.2
0.55
1000
Checkpoint
Checkpoint
Checkpaint
Checkpoint
T=0.8
T=10
T=1.2
T=1.5
T=18
Figure 19: Pass@k of Qwen-2.5-0.55 GSM8k Naive Sampling with Replacement
Pass@1
Pass@e
Pass@@4
0.95
0.65
0.85
0.60
0.85
Pess
Poss@2
0.75
0.55
0.70
0.50
0.75
0.65
0.70
0.60
0.5
0.65
0.55
1000
1500
1000
1500
1500
Checkpoint
Checkpoint
Checkpoint
Checkpoint
Tw0.8
Figure 20: Pass@k of Qwen-2.5-0.5B GSM8k Oracle Top K Sampling
E3 Diversity Comparison Between SFT and WiSE-FT
26

---

Published as a conference paper at COLM 2025
Diversity Across SFT [T=1.0]
Step 171
Step 513
Density
Density
Densigy
684
Step 855
0.2
0.2
Step 1026
0.1
0.1
0.1
Stop 1368
0.0
0.0
0.5
1.0
0.00
0.02
0.04
0,06
0.0
0.5
1.0
Step 1710
Diversity
Semantic Diversity
Diversity
Wiseft Diversity SFT [T-1.0]
Step 171
0.5
0.4-
0.4-
Step 513
Density
Density
Density
684
Step 855
Step 1197
0.1
Stop 1539
0.0
0.5
1.0
0.00
0.02
0.04
0.06
0.0
0.5
1.0
Step 1710
Operation Diversity
Semantic Diversity
Answer Diversity
Figure 21: Operation, Semantic, and Answer Diversity of Gemma-2-25 checkpoints of SFT over GSM8K versus the corresponding WiSE-FT variants (with the earliest checkpoint). We decode with temperature set to 1.0.
Diversity Across [T=1.6]
Step 171
Step 513
Density
Density
Densig
Stop 855
0.2
0.2
Step
Step 1197
0.1
0.1
0.1
Step 1368
0.0
0.5
1.0
0.0
0.00
0.07
0.04
0.06
0.0
0.5
1.0
Step 1710
Operation Diversity
Semantic Diversity
Diversity
Wiseft Diversity Across SFT [T-1.6]
Step 171
0.5
0.5
0.5
342
0.4
0.4
Stop 513
Densiyy
0.3+
Density
Density
0.3
Step 855
684
0.2
Stop 1197
Step
Slep 1539
0.0
0.0
0.5
1.0
0.00
0.02
0.04
0.06
0.0
0.5
1.0
Step 1710
Operation Diversity
Semantic Diversity
Answer Diversity
Figure 22: Operation, Semantic, and Answer Diversity of Gemma-2-21 checkpoints of SFT over GSMSK versus the corresponding WiSE-FT variants (with the earliest checkpoint). We decode with temperature set to 1.6.
27

---

Published as a conference paper at COLM 2025
G Best of K Evaluation
Be5t232
0.5
0.55
Cas
1900
810
039 V2 (30
C12
133
Mish Tx3.
wbbakTT122
Figure 23: Best@K performance on MATH500 with ORM verifier, comparing different SFT and WiSE-FT checkpoints of Qwen-2.5-0.5B for K 3,
Checkpoint 224
Checkpoint 672
Checkpoint 1120
Bethkk Acturacy
0.30
0.30
0.30
0.25
0.25
0.25
0.20
0.20
0.20
wiselt T-0.8
T=seftT=1.22
st T=1.0
T=1.0
sft T,0
T=12
Checkpoint
Checkpoint 672
Checkooint
Majenty Accuracy
0.3
0.3
0.3
0.2
0.2-
0.2
0.1
10
20
30
0
10
20
30
10
20
30
k
wiseft T=1.0
sft T=0.8
T=1.2
Figure 24: Best@K performance on MATH500 with ORM (Top) and Majority Vote (Bottom) for early, middle, and late SFT checkpoints and WiSE-FT counterparts, showing Qwen-2.5- 0.5B 's scaling across K values.
Majorhyya32
0.22
0.20
0.20
0.30
245
1200
1000
200
244
1000
whot TutLi
wseft Tw122
Figure 25: Best@K performance on MATH500 with majority voting, comparing different SET and WiSE-FT checkpoints Qwen-2.5-0.5B for 2,,4,8,33
28

---

Published as a conference paper at COLM 2025
H Diversity Collapse and WiSE-FT Results for the Coding Task
To test whether coding tasks exhibit the same diversity collapse observed in reasoning benchmarks, we fine-tuned the Owen2.5-coder-0.51 model for 10 epochs on the Magicoder- Evol-Instruct-110 dataset, following the Stage 2 SFT recipe from OpenCoder LLM. We then applied WiSE-FT by interpolating the weights of the second SFT checkpoint with the initial model using interpolation ratio 0.5. Both the original SFT checkpoints and their WiSE-FT counterparts were evaluated on HumanEval for pass@k.
We found that, much like in mathematical reasoning tasks, SFT on coding data indeed suffers from diversity collapse: although pass@1 steadily improves over epochs, pass@k begins to deteriorate. And WiSE-FT still improves performance and mitigates the diversity collapse.
HumanEval- - Pass@k Acress SFT Checkpoints
Passe1
Pass@16
Passg
8.4
6.30
0.70
0.55
636
1000
Chadgolk
Che
Figure 26: Pass@K performance of SFT checkpoints on HumanEval (temperature 1.0).
HumanEval- Pass@ Across Checkpoints (SFT WiSE-FT)
Paxs@]
PaIe
Paseg
8.45
6.10
0.3
0.00
2.00
022
0.5
270
046
Chickectt
Creckpets
Chopgar
Figure 27: Comparison of pass@K for SFT checkpoints and their WiSE-FT counterparts at 1, 16, 64.
29

---

Published as a conference paper at COLM 2025
HumanEval- Last Checkpoint (1700) Comparison: SFT vs WiSE-FT
0.9
0.8
Pass@k
0.6
0.5
Original WiseFT
0.4
20
60
80
100
120
Figure 28: Pass@K performance of the final SFT checkpoint versus its WISE-FT variant.
30


Reviews:
[{"review_id": "9kQ9oGfg84", "reviewer": "Reviewer_jb4c", "summary": "This paper focuses on balancing Pass@1 and Pass@k for reasoning models. The authors present a finding that supervised fine-tuning (SFT) improves Pass@1 but reduces output diversity. Consequently, the fine-tuned model is not well-suited for subsequent reinforcement learning alignment. Based on this finding, the authors propose a weight ensemble method, which averages the weights from different stages of SFT. Surprisingly, this simple ensemble performs well on both Pass@1 and Pass@k, even outperforming early stopping in SFT. The authors present extensive results to support their conclusions.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper presents solid evidence showing how SFT affects the performance of Pass@1 and Pass@k, as well as output diversity.\n- The proposed ensemble method is simple but highly effective.\n- The authors provide theoretical explanations for their findings.\n- Extensive experimental results are presented.", "reasons_to_reject": "I just have one minor question: what optimizer is used in the experiments? Given that the ensemble method performs better than early-stopped SFT, I assume part of the reason is that the gradient direction changes significantly over time. If the gradient direction is always similar during the optimization, there should not be a significant different between ensemble and early stopping. I'm wondering whether the choice of optimizer could significantly affect this conclusion. It's worth trying other optimizer for a better insight.\n\nOther than that, I don't see any major weaknesses.", "questions_to_authors": "Please see above.", "extra_scores": {}}, {"review_id": "XjVVmLkIGl", "reviewer": "Reviewer_yBuj", "summary": "Prior work has found that during supervised finetuning of reasoning models, the single best generation (pass@1) can continue to improve, while the best of K generations (pass@K) begins to deteriorate. This is due to diversity collapse, as SFT-ed models tend to converge to a single reasoning trace.\n\nThis paper shows that by simply interpolating the weights of the latest SFT checkpoint with an early checkpoint (a method called WiSE-FT from Wortsman et al., 2022), the pass@K performance can be completely recovered while also improving pass@1. Due to preserving output diversity, there are two immediate benefits of the WiSE-FT model: it scales better with test-time compute (§4.1), and is a better starting point for RLHF (§4.2). These gains cannot be achieved using an early-stopped SFT model or high-temperature decoding from the final SFT checkpoint. The experiments mainly use math datasets like GSM8K and MATH, though they show their results also generalize to training on OpenThoughts-114k.\n\nThe authors show empirically that diversity collapse happens with more steps in SFT (§5.1), and then connect this to pass@K in some theoretical discussion (§6). They prove that the expected pass@K decreases with increased variance in pass@1. SFT causes models to become very confident, so that the pass@1 distribution becomes bimodal — it is either very low (when the model is wrong) or very high (when the model is right). This increase in variance explains the drop in pass@K.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper is very well-written and has an insightful discussion of the tradeoff between pass@1 and pass@K during SFT.\n2. The method, WiSE-FT, is extremely effective. It successfully improves *both* pass@1 and pass@K with longer SFT. The resulting models scale better with increased test-time compute, and are substantially better starting points for RLHF.\n3. There is meaningful analysis & theoretical discussion of how SFT leads to diversity collapse and compromises pass@K.", "reasons_to_reject": "1. This is not a reason to reject in my opinion, but it is worth noting that the method is not new. In the original paper, WiSE-FT interpolates the weights of a finetuned model and pretrained model to improve OOD generalization. The authors give credit to the original work and do not attempt to downplay the connection.\n2. In the related work, it would be good to discuss other successful applications of model souping for LLMs. For instance, I know that [Olmo 2](https://arxiv.org/abs/2501.00656) souped model checkpoints trained on different data permutations, which they found to match or outperform any individual checkpoint.", "questions_to_authors": "Notes on presentation:\n- Figure 2 caption is hard to parse due to formatting.\n- In Figure 2c, the color-coding of \"SFT Early\" and \"SFT Last\" switches between the top and bottom subfigure. \n- L. 252: \"diveristy\" $\\rightarrow$ \"diversity\"", "extra_scores": {}}, {"review_id": "45BMYkx0Cq", "reviewer": "Reviewer_Mseo", "summary": "This paper proposes the usage of a simple weight ensembling technique, WiSE-FT, to address the diversity collapse problem that happens after model finetuning, either SFT or RL. WiSE-FT simply proposes to mix the finetuned model weights with vanilla model weights evenly, and improves the pass@k for the fintuned model. The paper further discusses some theoretical foundations of the diversity collapse issue, and how it impacts pass@k, and further test-time scaling. The proposed method and theoretical conjectures are evaluated empirically by finetuning Gemma and Qwen models with SFT and RL.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The proposed method is simple and yet effective. It's more likely to be adapted in various finetuning recipes. \n- I like the theoretical analyses. They give more insights into how the diversity collapse problem impacts the widely used recipe of SFT followed by RL. These analyses are well supported by numbers and experiments.", "reasons_to_reject": "- The impact of finetuning RL is only backed by theory, but no empirical evidence.\n- It's not clear if the pass@1 improvement is valid. There seems to be no clear theoretical support.", "questions_to_authors": "- Line 243: It would be interesting to see if the KL loss can help SFT.", "extra_scores": {}}, {"review_id": "RXv8cut0kh", "reviewer": "Reviewer_CtBR", "summary": "The paper argues that the diversity of generations collapses during supervised finetuning (SFT) for mathematical problem solving. This leads to suboptimal test-time scaling and low pass@k. The authors make use of WiSE-FT, a straightforward intervention that interpolates the weights of the latest SFT checkpoint with an early one which recovers pass@k and boosts pass@1.\n\nThe authors provide evidence for their claims across various models and datasets, complementing it with a theoretical framework that elucidates the bias-variance tradeoff in pass@k. The paper proposes a new training paradigm (SFT → WiSE-FT → RL) that offers a promising direction for developing more capable reasoning systems.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper tackles the problem of diversity collapse in LLMs during SFT, which negatively impacts reasoning and subsequent RL training. \n\nIt proposes to use WiSE-FT, a compelling solution that is simple, practical to implement, and adds no computational overhead at inference. The method's value is demonstrated through extensive empirical results across a few models, showing consistent, and complementary gains in pass@k, pass@1, and RL scaling that are not achievable by decoding strategies alone.\n\nThe authors provide a theoretical framework that formalizes the bias-variance tradeoff, demonstrating WiSE-FT's advantages, and proposing a promising new SFT → WiSE-FT → RL training paradigm", "reasons_to_reject": "The paper primarily focuses on a  $ 1/2 \\cdot uv_0 + 1/2 \\cdot uv_t$ interpolation ratio and and the use of an early SFT checkpoint. A more comprehensive exploration of different interpolation ratios and the selection criteria for the \"early checkpoint\" could reveal further insights into the method's generalizability and optimal application. \n\nThe analysis of the sensitivity of WiSE-FT's performance to its own hyperparameters, such as the exact choice of the early checkpoint and the interpolation weights, would strengthen the empirical results. The paper mentions using \"the earliest SFT checkpoint (in our case, after 1 epoch of training)\" for interpolation in WiSE-FT. However, the justification for selecting this specific checkpoint is not explicitly provided.\n\n\n\nThe theoretical analysis of diversity collapse in RL is conducted in a simplified discrete bandit setting. While insightful, the direct applicability and implications of these theoretical results to the complex, high-dimensional space of large language models could be further discussed.\n\nThe paper states that SFT is performed on a 30K subset of rephrased augmentations of GSM8k and MATH from MetaMath40k. However, the paper does not sufficiently detail the distribution of these paraphrased augmentations or elaborate on how they might contribute to the observed decay in pass@k. It would strengthen the paper to discuss whether this degradation is an inherent characteristic of finetuning on such rephrased data, or if a different source or distribution of SFT data might mitigate this issue and prevent the diversity collapse.\n\nIn the \"General Reasoning Tasks\" section (Section 4.3), the model is trained on \"Open Thoughts-114k,\" a dataset of \"math, science, and coding questions\". However, the evaluation is only conducted on the \"AIME24 competition math dataset\". To demonstrate the robustness and generalizability of the results for models trained on a diverse dataset, the evaluation for this section should be extended to include a broader range of general reasoning tasks, encompassing science and coding problems, in addition to competition math.\n\nFigure 3 provides an analysis of diversity only for SFT checkpoints. To provide a more comprehensive understanding of how WiSE-FT addresses this issue, it would be beneficial to include a similar analysis showing these diversity metrics for the WiSE-FT interpolated checkpoints.\n\nThe paper acknowledges other related work such as methods involving exploration through provers (e.g., \"Rewarding Progress\" by Setlur et al., 2024) and inference-aware finetuning (e.g., by Chow et al., 2024). However, a direct empirical comparison of WiSE-FT against these or similar methods that also aim to improve reasoning capabilities and diversity is missing.\n\nThe experimental evaluation primarily focuses on mathematical reasoning tasks. To strengthen the paper's claims about improving reasoning in language models more generally, it would be beneficial to extend the evaluation to other complex reasoning domains, such as code generation and scientific reasoning.", "questions_to_authors": "Could you elaborate on the criteria used to select the \"early SFT checkpoint\" for WiSE-FT interpolation? Is there a principled method for determining this, or is it generally based on empirical observation?\n\nThe paper focuses on a $ 1/2 \\cdot uv_0 + 1/2 \\cdot uv_t$​ interpolation ratio. Have you explored other interpolation ratios for WiSE-FT, and if so, how did they affect the balance between pass@1 and pass@k, and overall performance?\n\nWhy was the Open Thoughts-114k dataset not used for the main experiments on diversity collapse, pass@k, and RL scaling, instead of the GSM8K and MATH rephrased datasets? Was there a specific reason for this choice, especially given that Open Thoughts-114k is described as a \"high-quality synthetic dataset\"?\n\n\nGiven that standard decoding strategies like nucleus and top-k fall short of oracle performance, have you investigated if WiSE-FT can be combined with these strategies to achieve synergistic effects that bring performance closer to the theoretical optimum?\n\n\nIn Figure 2(a), the y-axis for the \"MATH Majority Vote\" plot is labeled \"Best@K\". While \"Best@K\" is defined to encompass both ORM and Majority Vote, for clarity and specificity in that particular subplot, would it be more appropriate to label the y-axis as \"Maj@K\"?\n\n\nRegarding Table 1, why are the pass@k values only shown up to k=8, given that models often demonstrate continued increases in pass@k for much larger values (e.g., K ≥ 64)? Additionally, why is the proposed WiSE-FT method not present in this table comparing various decoding strategies?\n\nCould you elaborate on the \"Top-k w/Oracle\" methodology? How is the \"oracle\" integrated into the top-k sampling process?\n\nBeyond the current diversity metrics (answer, operation, and semantic diversity), are there other approaches or metrics that could capture a more nuanced understanding of \"meaningful\" diversity in reasoning traces, especially for open-ended reasoning tasks?", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a technique to improve the diversity of generations of a model which is post-trained using either RL or SFT, by a simple weight ensembling technique. Reviewers generally agree that the approach is simple but highly effective (jb4c, yBuj, CtBR, Mseo) and can be applicable to improve post-trained model's diversity in a general setting. Reviewers are also satisfied with the theoretical justifications provided in the paper. CtBR had several questions and comments regarding the empirical methods in the paper, which the authors respond adequately. Given the positive reviews, I'm also happy to recommend an accept.
Decision: Accept