Mode-conditioning unlocks superior test-time compute scaling
Authors: Chen Henry Wu
Sachin Goyal
Aditi Raghunathan
Venue: ICLR 2026 Poster (2026)
Keywords: test-time compute
reasoning
diversity
Primary area: foundation or frontier models, including LLMs

Abstract: Parallel sampling is essential to test-time scaling and reinforcement learning (RL), but its effectiveness is sharply limited by diversity collapse, where models concentrate on a few modes and repeated samples produce the same mistakes. We propose the mode-conditioning (ModC) framework, which explicitly allocates sampling compute across reasoning modes using either specialist models or mode-specific prefixes. With predefined mode labels, ModC consistently improves test-time scaling (Pass@k) across controlled graph-search tasks and math reasoning benchmarks, spanning model families and sizes from 0.5B to 7B. On OpenThoughts, fine-tuning Qwen2.5-7B with ModC achieves an 4× efficiency gain over standard training while also improving the maximum attainable Pass@k. We further show that gradient clustering enables ModC without predefined mode labels, yielding up to 10% gains on datasets such as NuminaMath. Finally, we show that ModC improves Pass@k after RL training and can further boost the Pass@k gains of diversity-inducing RL methods. These results demonstrate that standard training underutilizes the diversity in data, and that ModC provides a simple, effective remedy for unlocking the full benefits of diversity in parallel sampling.

Published as a conference paper at ICLR 2026
MODE-CONDITIONING UNLOCKS SUPERIOR TEST-TIME COMPUTE SCALING
Chen Henry Wu, Sachin Goyal, Aditi Raghunathan Cameeiie Mellon University {chenwu2, sachingo, adi ti rag @css cmu
ABSTRACT
Parallel sampling is essential to test-time scaling and reinforcement learning (RL), but its effectiveness is sharply limited by diversity collapse, where models concen- trate on a few modes and repeated samples produce the same mistakes. We propose the mode-condee condiiiooin (ModC) framework, which explicitly allocates sampling compute across reasoning modes using either specialist models or mode-specife prefixes. With predefined mode labels, ModC. consistently improves test-time scaling (Pass@ k) across controlled graph-search tasks and math reasoning bench- marks, spanning model families and sizes from 0.5B to 7B. On OpenThoughts, fine-tuning Qwen2.5-7BB with ModC achieves an 4x efficiency gain over standard training while also improving the maximum attainable Pass @ k. We further show that gradient clustering enables ModC without predefined mode labels, yielding up 10% gains on datasets such NuminaMath. Finally, we show that ModC improves Pass @ k after RL training and can further boost the Pass@k k gains of diversity-nduccigg RL methods. These results demonstrate that standard training underutilizes the diversity in data, and that ModC provides a simple, effective remedy for unlocking the full benefits of diversity in parallel sampling.
INTRODUCTION
Parallel sampling is essential to test-time scaling and reinforcement learning (RL), driving major
advances in capability (Wang et al., 2023; Snell al., 2025; DeepSeek-Al et al., 2025). parallel scaling, the model is given multiple independent attempts, and it has proven especially effective in domains like mathematics, coding, and scientific discovery, where candidate solutions can verified automatically, making it a backbone of systems such AlphaEvolve (2025).
Despite its promise, parallel scaling relies on crucial assumption: the model must generate diverse and creative solutions. In practice, however, finetuning (Dang al., 2025) and reinforcement learning (Yue et al., 2025) are well-documented to induce diversity collapse, where models tend to produce the same mistakes across trials. As result, additional samples often reproduce the same errors or converge on indistinguishable strategies, leading to diminishing returns as compute is scaled. While recent works propose ways to mitigate diversity collapse, such as Pass@k k training (Chen et al., 2025b) and weight-space regularization (Dang et al., 2025), we argue that key challenge lies in using joken-by-token temperature sampling for diversity, which leads to lack of diversity in high-level reasoning modes.
In this work, we put forward mode-conditioning (ModC), a new paradigm that explicitly structures test-time scaling around multiple reasoning modes. Consider a simple well-defined task of Countdown from Gandhi al. (2024) which involves a search problem that admits two clear modes: breadth breadth-first search (BFS) and depth-first search (DFS). Depending on the graph structure, both modes have their own strengths. However, as we will show, temperature sampling can fail in covering both modes for given problem. Rather than drawing repeatedly from collapsed distribution, we enforce coverage across strategies by conditioning on modes and allocating samples to cover diverse modes. This simple yet powerful and, to the best of our knowledge, previously unexplored idea provides principled way to boost test-time scaling. The advantage becomes especially pronounced on inputs where the dominant mode fails but lower-probability succeeds.

---

Published as a conference paper at ICLR 2026
Standard Training on Search Problems
Standard Inference
Mixed teachers
QwQ-328 (best teacher)
IM
ply [x
Mode- cond-cioniiionng
BFS samples
samples
Task by DFS
Qwen2.5-78 (Instruct)
0.8
Mode rode-condittioned Training
Made-conditioned Inference
Pass@k
0.6
BFS samples
0.4
ply x
0.2
LM
DPS samples
ply |x
solvable by DFS
k/2
2
2
2
2
Mode-conditioning for test-time scaling. Modern LLMs often collapse to a single strategy, making Pass @ k scaling suboptimal: if the chosen strategy is wrong, every attempt fails. (Left) In controlled graph task solvable by DFS or BFS, models trained on both still often commit to just one. To address this, we introduce mode-conditioning (ModC) that explicitly allocates test-time compute across modes. We study two training methods that enable this: separate models or a single model with mode-specific prefixes. (Right) 4x efficiency gains with ModC training. We apply ModC to long chain-of of- of-thhught reasoning distillation on the TpenThoughts dataset. With ModC, the model achieves the same Pass @ 1024 as standard training using only k = 256 samples, yielding an ~4x improvement in inference efficiency. Moreover, ModC also improves the maximum attainable @ass@k, k. pushing the frontier of test-time scaling.
Figure I:
While the principle of ModC' is appealing, the key question is how to implement it in practice. One option is to prompt the model with explicit instructions for which mode use, but this requires extensive manual effort to characterize modes and, more critically, models often fail to follow instructions for low-probability strategies. To move beyond such ad-hoc prompting, we train models to provide explicit control over modes. We explore two natural instantiations: (i) training separate specialist models, each dedicated to a distinct mode, and (ii) training a single model with mode- specific prefixes, where modes can be sampled reliably by using the corresponding prefix. Although both approaches are conceptually simple and have appeared in other contexts, we find that even these straightforward implementations yield substantial gains for test-time scaling.
On the task of Coun tdown, we see that standard training does in fact struggle to sample from both modes (BSF and DFS) for several inputs. With predefined mode labels, ModC with both separate models and prefixes shows consistent superior test-time scaling to standard training, with gains especially pronounced Countdown task, we observe that ModC with separate models outperforms ModC with prefixes, suggesting that knowledge transfer across modes is less erucial.
those instances that only be solved by either BFS or DFS. In this
We apply ModC to real-world LLM reasoning datasets with both short- and long-form CoT. Our experiments fine-tune two different base LLMs, Qwen2.5-Base and OLMo2-Base across model scales from 0.5B to 7B, using distillation from two distinct teachers, each representing different mode. Across all settings, we observe a clear and consistent trend: ModC significantly improves test- time scaling compared to both standard mixed-teacher training and the best single-teacher baseline. In particular, on Qwen2.5-7B, ModC training yields up to 4x efficiency gains on both long-CoT CoT (Figure 1) and short-CoT (Figure 4). Notably, standard training often fails to exploit teacher diversity, counterintuitive outcome,
with single-teacher models outperforming those trained on mixtures since more diverse data should in principle help. ModC, by contrast, effectively harnesses this diversity, translating it into stronger test-time scaling. other words, diverse training data is most useful when paired with mechanisms that preserve and control its modes.
Next, we turn to ask: can ModC better hameess the implicit diversity of existing post-training datasets (such NuminaMath; Li et al., 2024) that standard training might be missing? However, applying mode conditioning requires knowing the modes apriori as well as mode annotations on the training data. We find that gradient clustering (inspired by Xia al. (2024); Jung et al. (2025)) enables ModC.
2

---

Published as a conference paper at ICLR 2026
without predefined mode labels. Once again, 10% gains on the NuminaMath dataset with no additional side information.
consistent gains with ModC achieving up to
We show that ModC improves Pass @ k after RL training and can further boost the Pass @ gains of inversity-inducing RL methods. While RL brings both models to the same Pass @ 1, ModC excels immediately at A = 2 with higher Pass & scores. This result shows that unlike standard SFT, ModC successfully enriches the solution space without degrading the accuracy of its top output. We see similar benefits of ModC on top of Pass@ RL (Chen et al., 2025b), which explicitly prevents diversity collapse during RL. This suggests that interventions designed to prevent diversity collapse during RL can be further boosted by SFT interventions such as our ModC.
These results demonstrate that standard training underutilizes the diversity in data, and that ModC provides a simple, effective remedy for unlocking the full benefits of diversity in parallel sampling. In summary, we make the following contributions.
1. Introducing the mode-conditioning (ModC) framework. We propose simple but powerful paradigm to address diversity collapse in LLM reasoning and improve test-time scaling (Section 2). ModC explicitly allocates test-time compute across reasoning modes. We propose two training methods to allow for such test-time allocation: (i) specialist models and (ii) mode-specific prefixes.
ModC demonstrates consistent gains across tasks. Through controlled graph- search experi- ments (Section 3) and large-scale reasoning benchmarks (short- and long-form CoT, distillation from multiple teachers) (Section 4), we show that ModC achieves substantial and consistent improvements in test-time scaling, including up 4x efficiency gains. We also carefully analyze tradeoffs between different ModC training methods, effect of model size, data composition etc.
3. ModC on training data without predefined mode labels. We find that gradient clustering enables ModC without predefined mode labels, yielding up to 10% gains on datasets such as NuminaMath. These results suggest that standard training underutilizes the diversity in data an inefficiency that the ModC framework directly addresses.
4. The benefit of ModC holds on top of standard RL and diiverity-induccio RL. Unlike standard SFT, ModC successfully enriches the solution space and mitigates diversity collapse during RL. Results on top of Pass & RL suggest that interventions designed to prevent diversity collapse during RL can be further boosted by SFT interventions such as our ModC.
2 THE MODE-CONDITIONING FRAMEWORK
2.1
PRELIMINARIES
Large language models (LLMs) generate outputs by sampling from a probability distribution over continuations. more complex tasks, instead of producing single output, we can allocate additional test-time compute by drawing k independent samples for the same input and selecting the best candidate. This strategy, known as parallel scaling, especially effective in tasks where solutions can be automatically verified (e.g.. mathematics or programming). The performance of this standard approach is captured by the Pass ksta metric on an input
Pass@ kista = 1 pz))
where is the Pass@1 or probability of sampling a trajectory that is successful on input x.
In practice the gains with parallel scaling depend strongly the underlying success probability Ps- Modern training pipelines such as supervised fine-tuning and reinforcement learning often induce mode collapse, where the model commits to a small set of strategies (Dang al., 2025; Yue et al., 2025; Sessa et al., 2024; Chow et al., 2024). On some prompts, this collapse drives the probability pa of sampling a successful strategy to be very small, so that an impractically large number of samples is required to obtain good performance.
2.2 MODE-CONDITIONED TEST-TIME SCALING
One approach is to modify the finetuning objective to prevent collapse and maintain higher Pa- We take complementary route: rather than sampling from single collapsed distribution, we explicitly allocate test-time compute across diverse modes, enforcing coverage so samples include not only the dominant strategy but also alternatives that may succeed where it fails.
3

---

Published as a conference paper at ICLR 2026
Consider two modes with success probabilities P1.æ and P2,x- If we split the budget evenly, sampling k/2 trajectories from each mode, the resulting probability of solving imput x is
Pass @ kinoac (x) = (1 Pı,.) //2 p2.x) kk22
(2)
In contrast, suppose the model does not know a priori which mode is better for Let W2 E [0, 1] denote the (random) probability with which the model uses the first mode, and assume that the model is unbiased between the two modes, in the sense that R[mz]] 1/2 (e.g., Wa is drawn from any distribution centered at 0.5). A single sample from the model succeeds with probability 2,,1, + (1 we P2,e, so we have
Pass kisaa (2; wx) = 1 11 WxP1, (1 w2.
The function w 1 (1 p2,2 (pl,2 w2,x)w)+ is concave [0, 1]. so by Jensen's inequality
Bw, Pass kita (z; wy)] Pass kista (2; B[w=]) Pass @ Sassskstalla (2; 1/2).
(4)
When 1/2, the single-sample success rate is (p1,. + P2.=)/2. Whenever pı., F P2,x, we have
(1 -1-18 P1,2) p2.x) 122 < (1-Pa)*,
(5)
which implies
Pass @ (xoodc(2)
Pass @ (hst 1/2) 2 Ew, Pass (2;
(6)
In other words, even when the model's mode-usage follows any distribution centered at 0.5, explicitly allocating compute evenly across modes is strictly better than sampling from the model's 's own uncertain mixture.
But how do we explicitly sample from different modes in practice? A naive baseline is to prompt the model with instructions use different strategies. However, this approach is unreliable: is unclear how to phrase the prompts, and the model may not consistently follow them.
2.3 MODE-CONDITIONED TRAINING
Instead of relying on ad-hoe prompting to elicit different behaviors, we consider scalable training objectives that explicitly enforee control over modes. This provides reliable lever at test time for allocating compute across diverse strategies. We explore two natural instantiations: training with separate specialist models and training with prefixes within single model.
In this section, we assume that the relevant modes are known a priori and that training data can be annotated with the mode used to generate each trajectory. While this assumption is convenient for exposition and testing the benefits of our paradigm, it is not strictly necessary. In practice, one could imagine automated approaches for mode discovery and annotation, for example by clustering similarity measures or other unsupervised techniques. We return to this point $5 where we discuss how mode-conditioned training can be extended to settings where the modes are not explicitly labeled in the data.
trajectories using gradient-basee
Mode-conditioned training with separate models. The most direct approach is to train distinct models, each specialized to particular mode of reasoning. Concretely, the training data is partitioned into subsets corresponding to different strategies, and a separate model is trained on each subset while keeping total training data and compute constant. At test time, the sampling budget is divided across the specialists (c.g., k/2 samples from each in the two-mode case). This design ensures strong specialization and reduces correlated errors, which translates into more effective parallel scaling.
codeicondiiioned training with prefixes. While separate models improve diversity, they prevent knowledge sharing across modes. This is significant drawback in realistic reasoning tasks, where different strategies often rely on common linguistic or mathematical foundations.
To overcome this, draw inspiration from the literature steering model behavior via explicit conditioning tokens, a technique used widely in controlled text generation (Keskar et al., 2019) and instruction tuning. We prepend discrete condition tokens (e.g., [Mode 1]. [Mode 2]) to the input,
4

---

Published as a conference paper at ICLR 2026
training the model to associate each prefix with distinct reasoning strategy. At inference, balanced compute allocation is enforced by sampling evenly across the conditioning prefixes. This allows the model to specialize into distinct modes while still sharing knowledge across them, making it more scalable than training separate specialist models.
INVESTIGATING MODE COVERAGE IN PARALLEL SAMPLING
3.1 THE Countdown TASK
Countdown is generalization of Game 24, where model must find sequence of arithmetic operations to transform a set of starting numbers into a target value (Gandhi et al., 2024). Given several starting numbers, the model can apply operations {+. - X, : to reach a target. For example. given {10, 10, 4, 6} with target 16, one solution is (10 x 10 - 4) + 6 = 16.
This task naturally admits two distinct problem-solving modes: depth-first search (DFS) and breadth- first search (BFS). This allows us to precisely control and examine which mode is used. Solutions are easily verifiable by checking if operations reach the target, which makes it an ideal testbed for studying test-time scaling with parallel sampling.
3.2
WHY MODE COVERAGE IMPORTANT?
In principle, both BFS and DFS are complete search algorithms capable of finding any solution, so why would mode coverage matter for this task? For real-worl worldw problems, however, computational constraints require using heuristics to make the search tractable. Following Gandhi et al. (2024), we use heuristics to guide and prune the search space, which introduces important asymmetry: with heuristic pruning and search budget constraints, each algorithm excels on different problem instances some problems become solvable only by the oracle DFS while others only by the oracle BFS. Since we cannot know a priori which algorithm will succeed for a given problem, maintaining coverage of both modes during test-time sampling becomes crucial for achieving high success rates.
To evaluate test-time scaling, we report Pass@ k metries on two held-out test sets. We first create natural test set of 500 problems by randomly selecting unseen target numbers and valid starting numbers that can reach the target. Second, we subsample an adversarial test set designed to require mode diversity: we filter for problems where either oracle BFS or oracle DFS (but not both) achieves less than 5% success rate across multiple runs. This adversarial set directly tests mode coverage each problem is effectively solvable by only one algorithm, so achieving high accuracy requires sampling from both DFS and BFS modes when we do not know priori which one works better. We expect that the benefit from explicit test-time balancing is much larger on the adversarial test set.
3.3 MODE COVERAGE
We rejection sampling to create training set, keeping only instances where least one search algorithm (oracle DFS or BFS) successfully finds a solution. Specifically, we uniformly generate target number from to 200 and four starting numbers that can reach the target, and uniformly choose DFS and BFS and one of the search heuristics for guidance and pruning. We note that DFS has a higher overall success rate, so our final training set ends up consisting of 163K problems, with 97K DFS solutions and 65K BFS solutions. Each training example includes the input, the search trajectory, and final operations. We train Qwen2.5-Baa models (0.5B-7B) for 4 epochs.
We first examine whether models trained on mixed DFS and BFS data can learn to balance the two algorithms under repeated sampling. Figure 2 shows the fraction of BFS used by the model for each test problem. We observe that standard training on the mixture of both algorithms (shown in gray) tends to bias toward one algorithm for many test problems, i.e., either predominantly using DFS (low BFS fraction) or BFS (high BFS fraction), rather than balancing both.
Effect of diversity of training data. Recall that we use rejection sampling to create our training set, which naturally biases the training data towards the on-average more promising algorithm (i.c., DFS in this case). What if we had 50-50 data for DFS and BFS? To answer this, in Figure 2 we also plot the distribution for this standard training with balanced data. We see that the distribution is less skewed, but still many problems have extremely imbalanced allocation of test-time compute.
5

---

Published as a conference paper at ICLR 2026
Standard training
ModC w/ prefix
Standard training (balanced data)
ModC rodels
0wen2.5-0.58 (Base)
(wen2:5-0.58 (Base)
0w
(Base)
(wen2:5-0:58 (Base)
0.34
0.8-
0.8
Density
Density
0-2
Density
Density
0.6
9.4
0.1
0.2
0.2-
BFS per problex
Frac
BFS per problen
Frac. of per probles
Frac.
BPS per problen
Figure Standard training fails to balance diverse modes per problem under repeated sampling.
This issue does not go away with balanced training data. Instead, ModC explicitly targets and successfully achieves balanced test-time compute allocation.
ModC balances test-time compute allocation. In contrast, ModC with explicit balanced allocation achieves the desired behavior. We see that for both training separate models (shown in blue) and training with prefixes (shown in red), the fraction BFS per problem is concentrated around 0.5, which demonstrate nearly perfect balance.
Qwen2.5-68 (Base)
Swen2.5-1.58 (Base)
Qwen2. (Base)
8.8
8.8
Passik
8.7
0.7
0.7
Standard training
ModC w/ prefix
0.6
0.6
0.6
ModC w/ separate models
0.5
0.5
0.4
(a) Pass@k l performance on Countdown natural test set
Qwen2.5-0.1 (Base)
Qwen2.5-1.58 (Base)
Qwen2.5-23 5-3B (Base)
0.8
8.8
Passigk
0.6
0.6
0.6
Standard training
Madt w/ prefix
MadC w/ separa te models
0.4
0.4
0.4
8.2
29
2
2
2
2
22
2
2
29
2
20
27
21
26
2
210
(b) Pass & performance on Countdown adversarial test set
Figure 3: Balanced test-time allocation improves scaling. (a) the natural test set of Countdown, balanced test-time allocation with ModC. shows consistent improvements as k: increases. (b) On the adversarial test set where each problem is challenging for one one algorithm (oracle DES or BFS), the gains from enforced mode diversity are more pronounced.
3.4 FROM MODE COVERAGE TO @
In this part, we show that ModC dramatically improves test-time scaling, particularly on problems that require diverse algorithmic modes. We start with ModC with separate models. Figure 3 shows the results on the natural and adversarial test set of Countdown across model scales. While Pass@1 is comparable or slightly lower for separate models, the scaling behavior dramatically improves by up to 8% for Pass 1024. This gap is consistent across all model sizes, with separate models maintaining superior scaling curves. The advantages are even more pronounced we see that the gap boosts up to 20% for Pass 1024. This shows that mode balance is indeed crucial for test-time scaling problems that DFS finds hard to solve potentially solvable for BFS. which enable better coverage of the problem space. On the other hand, Figure shows that ModC with prefix with balanced allocation outperforms standard training for most scales. As control, we try
the adversarial test set, where

---

Published as a conference paper at ICLR 2026
0wen2.5-0..m (Base)
Owen2.5-1.4 (Base)
0wen2. (Base)
Owen2 5- (Base)
0.8
0.6
0.8
0.8
Passok
0.6
0.4
0.6
0.6
0.4
6.4
24
2
2
k
OLMo2-10 1B
(Base)
OLMo2 11lo2-BB 7B (Base)
0.8
0.6
Pass@k
0.6
Mixed teachers
0.4
Best teacher
0.4
ModC w/ separate nodels (teacher)
0.2
ModC w/ prefix (teacher)
0.2
29
22
2
2
2
2
22
24
2
2
k
k
Figure 4: ModC improves short CoT reasoning, @ on MATH500. Naively mixing teacher data underperforms the single-teacher baseline, while ModC shows consistent gains. ModC with prefixes generally works better than ModC with separate models underscoring the importance of sharing knowledge across modes (teacher strategies) in math reasoning.
random partitioning the training data into two groups, which sometimes shows gains but does outperform ModC (see details in SB).
4 MODE-CONDITIONINS IMPROVES MATH POST-TRAINING
We saw how ModC improves Pass@ performance on Countdown and is superior for parallel scaling. In this section, we evaluate ModC when post-training for math reasoning. We investigate both long and short CoT reasoning across multiple model families and scales.
4.1 DISTILLATION FROM MULTIPLE TEACHERS
distillation from multiple stronger teacher models. This setting is particularly relevant given the existence of strong models with distinct reasoning and response styles. Recent work typically selects the single teacher that provides the best distillation performance (Muennighoff al., 2025; Guha et al., 2025). We test whether explicitly balancing compute across multiple teacher strategies improves performance. In each setting, we collect CoT reasoning traces from two teacher models for mathematical problems, where each teacher demonstrates distinct problem-solving styles. Following the methods established Section we compare four training strategies: (1) mixed teachers: mixing all teacher data together, (2) best teacher: training on the teacher data that produce the best performance (3) ModC with separate models: training separate models each teacher's data independently, and (3) ModC with prefixes: prepending teacher identity tokens (c.g., using [teacher]-style reasoning) to each CoT and using balanced test-time allocation where we sample equally from each teacher condition.
We start with a natural source of diverse modes:
4.2 SHORT CHAIN-OF-THOUGHTT
Experimental setup. We first experiment with post-training NuminaMath (Li et al., 2024) dataset where the chain-of-thought completions are relatively short and do no involve long thinking. We use the SFT traces distilled from two teacher models: DeepSeek-RI (DeepSeck-Al et al., 2025) and GPT-OSS- 120B (OpenAI, 2025), with the problems from NuminaMath. For evaluation, we use MATH500 (Hendrycks et al., 2021) and measure Pass kt We post-train Qwen2.5-B.r (0.5B-7B) and OLM Base (IB-7B) models for epochs. We tune the learning rate E {1e-4, 1e-5) and use AdamW optimizer with global batch size of 256.

---

Published as a conference paper at ICLR 2026
Naive data mixture underperforms. Figure compares the Pass@k curves MATH500 for the four training strategies mentioned above. We observe that naively mixing data from both teachers either underperforms or is at best comparable to the stronger single-teacher baseline. This is consistent with prior intuition that training on the best teacher can be more effective than mixing teachers.
ModC training unlocks superior test-time scaling. On the other hand, training both teachers' data but with mode conditioned (ModC) training and inference unlocks better test-time scaling (Figure 4). The gains are consistent across model families (Qwen and OLMo2) and scales (0.5B to 7B), offering up to 10% gain on Qwen2-0.5B and 15% on OLMo-2-7B. Comparing the two variants of ModC, we see that ModC with prefixes generally outperforms ModC with separate models suggesting that knowledge sharing across modes is crucial for math tasks.
4.3 LONG CHAIN-OF-THOUGHTS
We now examine long CoT reasoning, where models usually spend tens of thousands of tokens on extended reasoning before producing the answer more challenging problems.
Experimental setup. We use the subset prob- lems from OpenThoughts, (Guha et al., 2025) that they did ablation studies for the teacher models with. Specifically, solutions are from two teach- ers: QwQ-32B (Team, 2025) and DeepSeek-Re (DeepSeek-Al et al., 2025). For evaluation, we use AIME 2025 and measure Pass@ k. Following Guha et al. (2025), we initialize model weights with Qwen2.5-7B-Innsrrue
Mixed teachers
QwQ 32B teacher
DeepSeek R1 teacher
ModC separate models
ModC prefix
Owen2. 5B (Instruct)
0.8
Possøk
0.6
4x efficiency gains with ModC. Figure shows similar patterns to short CoT: ModC achieves con- sistently higher Pass@1 k than both single-teacher and mixed-teacher baselines. Even in the long CoT setting with extremely large token budgets per sam- ple, standard training fails to adequately cover mul- training again does not outper- form the best single teacher. In contrast, ModC not only provides an effective way to learn from mul- tiple teachers and surpass each individual teacher, but also delivers substantial efficiency gains: it matches the Pass @ 1024 of standard training with only k = 256 samples, yielding nearly 4x faster inference, while simultaneously pushing the maximum achievable @ass@k. k. We
0.4
0.2
20
2
24
2
2
200
k
tiple
Figure 5: ModC improves long CoT reason- ing. @ on AIME 2025. Standard train- ing with multiple teachers fails to outperform the best single teacher, but ModC with multiple teachers surpasses single teacher.
AUTOMATED MODE DISCOVERY
We have that ModC improves test-time scaling across wide variety of real-world settings. However, previous sections assume access to natural modes in the data: search algorithms in Countdown or teacher identities in multi-teacher distillation. In practice, most real-world training data lacks such clear segregation. Can we extend ModC to work on training data that contains mixed modes but lacks explicit labels? Furthermore, can we discover meaningful modes in training data without apriori knowledge of these modes? We explore both questions in this section, applying gradient clustering to discover and annotate modes in training data.
5.1 GRADIENT CLUSTERING
Gradient similarity has been shown effective in understanding training dynamics (Jacot et al., 201 2018), identifying influential training data (Koh & Liang, 2017), and diversifying data selection (Jung et al., 2025). That inspires us to test whether gradient cluster can discover meaningful modes in training data that we should condition on.
8

---

Published as a conference paper at ICLR 2026
For each training example (z, y), compute gradients with respect to model parameters
go (z. y) V log (yl
To reduce dimensionality,, follow prior works (Xia et al., 2024; Jung et al., 2025) to apply Rademacher random projection (Park et al., 2023) to each gradient vector. Once we get the projected gradient vectors for all training samples, cluster the vectors into C clusters, and all samples in the same cluster belong to one "mode" based on which we apply ModC.
5.2 GRADIENT CLUSTERING IMPROVES POST-TRAINING ON GENERAL DATA
We apply gradient clustering NuminaMath, real-world dataset that is probably quite diverse but we lack clear sense of what modes exist. Surprisingly, we see that training with ModC on automatedly discovered modes (via gradient clustering) yields significant improvements. Figure shows that ModC consistenty improves Pass @ compared to standard training across model scales.
Standard training
ModC w/ prefix (gradient)
ModC w/ separate nodels (gradient)
Twen2.5-0..0 (Base)
(wen2.5-1.55 (Base)
0- (Base)
Qwen2.5-78 (Base)
8.8
0.8
0.6
8.8
Pass@k
0.6
0.4
8.6
0.4
0.4
8.4
2
22
2
2
2"
2
2
2
2"
22
2
2
2
2
2
2
2
2
k
k
k
Figure 6: ModC on automatedly discovered modes via gradient clustering improves short CoT.
Pass@k k on MATH500 for Qwen2.5-Base models finetuned with standard finetuning and ModC.
MODE-CONDITIONIN IMPROVES REINFORCEMENT LEARNING
Finally, turn to verify if the benefit of ModC holds after RL, where the distribution becomes sharper and diversity decreases. For RL methods: we consider (1) standard RL, where the distribution becomes sharper and diversity decreases, and (2) Pass @k RL (Chen al., 2025b), which can improve both Pass and Pass @ k. For each RL method, we initialize the policy from either the standard SFT models or ModC models. For the prefix variant, we believe need prefix-following rewards to make sure that the learned prefix-mode binding is maintained. Therefore, we focus on the separate models variant given its simplicity and leave RL on the prefix variant for future work.
Figure shows the results of ModC vs standard SFT followed by different RL methods. Crucially, while RL brings both models the same Pass@ ModC excels immediately at k 2 with higher Pass & scores. This result shows that unlike standard SFT, ModC successfully enriches the solution space without degrading the accuracy of its top output. Moreover, we see that this observation also
RL
Pass@k RL
0.8
Standard SFT
Pass@k
ModC SFT
0.6
Standard SFT
ModC SFT RL
0.4
Standard SFT + Pass@k RL
ModC SFT Passøk RL.
0.2
20
22
24
2
28
20
22
24
26
28
k
k
Figure 7: ModC improves long COT reasoning. Pass@ on MATH500, While RL brings both models to the same Pass@1 @1, ModC excels immediately at with higher Pass scores.

---

Published as a conference paper at ICLR 2026
holds for @ RL, which explicitly prevents diversity collopse during RL training. This suggests that interventions designed to prevent diversity collapse during RL can be further boosted by SFT interventions such as our ModC.
7 RELATED WORK
Improving parallel test-time scaling. Repeated sampling significantly improves the performance of LLMs especially in domains like reasoning and coding ( (Wang et al., 2023; Brown et al., 2024; Rozière et al., 2024). To decide the final answer from the k attempts, one use either majority voting (Wang et al., 2023) or a verifier (especially in the code and scientific discovery domains) (Wang al., 2024; AlphaEvolve, 2025). However, series of works (Cobbe et al., 2021; Dang al., 2025) have identified issues during post post-training of LLMs that impair the diversity in model generation, consequently affecting the efficacy of test-time scaling. Huang et al. (2024) attribute this to the sharpening effect whereas Chu et al. (2025) highlight memorization as root cause.
Lot of recent works have in turn have proposed fixes for improved parallel test-time scaling. Beeching et al.; Snell et al. (2024) propose modifications to beam search to explcitly optimize for diversity amongst the candidates. Wang et al. (2025); Hughes et al. (2024) propose diverse prompting to improve test-time scaling. Taking step back, Sessa et al. (2024); Chow et al. (2024); Chen et al. (2025a) explicitly optimize for best-of-k performance during the finetuning process. Dang et al. (2025) propose a simple fix of ensembling the finetuned weights with the base model to mitigate diversity collapse. Goyal et al. (2025) further take step back and propose pretraining with logit distillation to improve parallel test-time scaling behaviors. In contrast, in this work propose more data centric conditioning of the finetuning process to explcitly encode specialist modes in the model.
Creativity and diversity of language models. Recent work has evaluated creativity and diversity in language models with mixed findings. Models underperform human experts in creative writing (Chakrabarty al., 2024) and humor generation (Zhang et al., 2024). While Anderson et al. (2024) find that language models increase individual idea divergence but causes group-level homogenization, Si al. (2024) report that LLMs can generate novel research ideas despite feasibility limitations. Nagarajan et al. (2025) demonstrate that global planning and seed coddiiiniiioning are crucial for creative generation. Temperature sampling weakly correlates with creativity and often introduces incoherence (Chen & Ding, 2023). On the evaluation front, various benchmarks have been proposed to measure output diversity and creativity (McLaughlin et al., 2024; Zhang et al., 2025; Jansen et al., 2024).
Specialization in model training. In this work, we proposed ModC where single model explicitly conditioned to learn separate modes of finetuning data. We do this by prepending conditioning tokens (e.g.. DFS or BFS) to the reasoning traces. Closest to our work is Mixture-of- Experts (Shazeer et al., 2017: Jiang et al., 2024; Jelassi et al., 2025) where different datapoints are routed to specific subpart of the model which is expert for the domain of the particular datapoint. However, in contrast in ModC all the datapoints are processed by the whole model and not a subpart. Mixture-of-experts are aimed at reducing the active parameter footprint of the model.
8 CONCLUSIONS AND FUTURE WORK
In this work, demonstrated that deliberate mode-conditioning (ModC) during training and infer- ence is crucial for unlocking the full benefits of test-time compute scaling. Across both controlled search problems and large-scat scale reasoning benchmarks, we showed that standard training tends to collapse onto a single strategy, while specialization through mode-condiiioning consistently yields superior scaling and more reliable gains. Beyond explicit labels such teacher identity, we further showed that gradient clustering can automatedly uncover meaningful specializations, making ModC broadly applicable. We finally show that ModC mitigates diversity collapse during both standard RL and diversity-inducing inducing RL. It suggests that interventions designed to prevent diversity collapse during RL can be further boosted by SET interventions such as our ModC. Looking ahead, exciting direc- tions include extending ModC to larger number of modes and to richer behavioral dimensions such as reasoning depth or planning style, as well as integrating ModC with adaptive allocation policies that learn how to optimally divide compute across modes at test time. Another promising avenue is further improving ModC with reinforcement learning, where balanced modes could encourage diverse exploration early in training and be gradually relaxed. Together, these directions position mode-cooniiioonss as a general and effective principle for building more powerful reasoning systems.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We have taken several steps to ensure the reproducibility of our results. Complete experimental details including model architectures, optimization hyperparameters, datasets, and evaluation protocols are provided in the relevant sections of the main paper for each experiment. In addition, we submit our code and scripts as part of the supplementary material, and will be open-soureing it along with other meta-data upon acceptance of the work.
REFERENCES
AlphaEvolve. Alphaevolve: A gemini-powered coding agent for designing advanced algorithms,
2025.
Barrett R. Anderson, Jash Hemant Shah, and Max Kreminski. Homogenization effects large language models human creative ideation. In Proceedings the 16th Conference Creativity & Cognition, Chicago, II, USA, June 23-26, 2024, pp. 413-425. ACM, 2024. URL. https: //doi.org/// 0. 1145/3635666666 3656204.
Edward Beeching, Lewis Tunstall, and Sasha Rush. Scaling test-time compute with open models.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoe V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024.
Tuhin Chakrabarty, Philippe Laban, Divyansh Agarwal, Smaranda Muresan, and Chien-Sheng Wu. Art artifice? large language models and the false promise of creativity. In Proceedings of the CHI Conference on Human Factors in Computing Systems, CHI 2024, Honolulu, HI, USA. May 11-16, 2024, pp. 30:1-30:34. ACM, 2024.
Feng Chen, Allan Raventos, Nan Cheng, Surya Ganguli, and Shaul Druckmann. Rethinking fine- tuning when scaling test-time compute: Limiting confidence improves mathematical reasoning, 2025a.
Honghua Chen and Nai Ding. Probing the "creativity" of large language models: Can models produce divergent semantic association? In Findings of the Association for Computational Linguistics: EMNLP 2023, Singapore, December 6-10, 2023, pp. 12881-12888 Association for Computational Linguistics, 2023.
Zhipeng Chen, Xiaobo Qin, Youbin Wu, Yue Ling, Qinghao Ye, Wayne Xin Zhao, and Guang Shi. @ass@k training for adaptively balancing exploration and exploitation of large reasoning models. ArXiv, abs/2508. 10751, 2025b. URL https: wtw.api..mmmaaischollsas. org 280649795.
Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Aleksandra Faust. Inference-aware fine-tunen one-tuning for best-of-n sampling in large language models, 2024.
Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shenghang Tong, Saining Xie, Dale Schuurmans, Quoe V. Lc, Sergey Levine, and Yi Ma. Sft memorizes, rl generalizes: A comparative study of foundation model post-training, 2025.
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Hcewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
Xingyu Dang, Christina Baek, Kaiyue Wen, Zico Kolter, and Aditi Raghunathan. Weight ensembling improves reasoning in language models, 2025.
DeepSeek-AA, Daya Guo, et al. Deepseck-rl: Incentivizing reasoning capability in 1lms via reinforce- ment learning. ArXiv, abs/2501.11948.. 2025.
Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, and Noah D. Goodman. Stream of search (sos): Learning search in language. ArXiv, abs/2404.03683, 2024.

---

Published as a conference paper at ICLR 2026
Sachin Goyal, David Lopez-Paz, and Kartik Ahuja. Distilled pretraining: A modern lens of data, in-context learning and test-time scaling, 2025.
Etash Kumar Guha, Ryan Marten, Sedrick Scott Keh, Negin Raoof, Georgios Smyrnis, Hritik Bansal, Marianna Nezhurina, Jean-Pierre Mercat, Trung Vu, Zayne Sprague, Ashima Suvarna, Ben Feuer, Liangyu Chen, Zaid Khan, Eric Frankel, Sachin Grover, Caroline Choi, Niklas Muennighoff, Shiye Su, Wanjia Zhao. John Yang. Shreyas Pimpalgaonkar, Kartik Sharma, Charlic Cheng- Jie Ji, Yichuan Deng, Sarah Pratt, Vivek Ramanujan, Jon Saad-Falcon, Jeffrey Li, Achal Dave, Alon Albalak, Kushal Arora, Blake Wulfe, Chinmay Hegde, Greg Durrett, Sewoong Oh, Mohit Bansal, Saadia Gabriel, Aditya Grover, Kai-Wei Chang. Vaishaal Shankar, Aaron Gokaslan, Mike A. Merrill, Tatsunori Hashimoto, Yejin Choi, Jenia Jitsev, Reinhard Heckel, Maheswaran Sathiamoorthy, Alexandros G. Dimakis, and Ludwig Schmidt. Openthoughts: Data recipes for reasoning models. ArXiv, abs/2506.0411 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. arXiv preprint arXiv:2 (0t.::/:00303884... 2021.
Audrey Huang. Adam Block, Dylan J. Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz. Jordan T. Ash, and Akshay Krishnamurthy. Self-improvement in language models: The sharpening mechanism, 2024.
John Hughes, Sara Price, Aengus Lynch, Rylan Schaeffer, Fazl Barez, Oluwasanmi Koyejo, Henry Sleight, Erik Jones, Ethan Perez, and Mrinank Sharma. Best-of-n jailbreaking. ArXiv, abs/2412.03556, 2024.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. NeurIPS, 2018.
Peter Alexander Jansen, Mare-Alexandre Cot'e, Tushar Khot, Erin Bransom, Bhavana Dalvi, Bod- hisattwa Prasad Majumder, Oyvind Tafjord, and Peter Clark. Discoveryworld: A virtual environ- ment for developing and evaluating automated scientific discovery agents. NeurlPS, 2024.
Samy Jelassi, Clara Mohri, David Brandfonbrener, Alex Gu, Nikhil Vyas, Nikhil Anand, David Alvarez-Melis, Yuanzhi Li, Sham M. Kakade, and Eran Malach. Mixture of parrots: Experts improve memorization more than reasoning, 2025.
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts, 2024.
Jaehun Jung, Seungju Han, Ximing Lu, Skyler Hallinan, David Acuna, Shrimai Prabhumoye, Mostafa Patwary, Mohammad Shoeybi, Bryan Catanzaro, and Yejin Choi. Prismatic synthesis: Gradient- based data diversification boosts generalization in Ilm reasoning. ArXiv, 2025.
Nitish Shirish Keskar, Bryan McCann, Lav R. Varshney, Caiming Xiong, and Richard Socher. Ctrl: A conditional transformer language model for controllable generation, 2019.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International Conference on Machine Learning, 2017.
Jia Li, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. Numinamath, 2024.
Aidan McLaughlin, James Campbell, Anuja Uppuluri, and Yiming Yang. Aidanbench: Stress-testing language model creativity on open-ended questions. In NeurIPS 2024 Workshop on Language Gamification, 2024.
12

---

Published as a conference paper at ICLR 2026
Niklas Muennighoff, Zitong Yang. Weijia Shi, Xiang Lisa Li, Fei-Fei Li, Hanna Hajishirzi, Luke S. Zettlemoyer, Percy Liang, Emmanuel J. Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling, ArXiv, 0tbp///51..1993..3. 2025.
Vaishnavh Nagarajan, Chen Henry Wu, Charles Ding, and Aditi Raghunathan. Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction. ICML, 2025.
OpenAl. gpt-oss- 120b & gpt-oss-20b model card, 2025.
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclere, and Aleksander Madry. Trak: Attributing model behavior at scale. In International Conference on Machine Learning, 2023.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code llama: Open foundation models for code, 2024.
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 Stanezyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amélic Héliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev, and Olivier Bachem. Bond: Aligning Ilms with best-of-n distillation, 2024.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-cxperts layer, 2017.
Chenglei Si, Diyi Yang, and Tatsunori Hashimoto. Can iims generate novel research ideas? A large-scale human study with 100+ NI.P researchers. 2024.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling I1m test-time compute optimally can be more effective than scaling model parameters, 2024.
Charlie Victor Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters. ICLR, 2025.
Qwen Team. Qwq-32b: Embracing the power of reinforcement learning, March 2025.
Peiyi Wang, Lei Li, Zhihong Shao, R. X. Xu, Damai Dai, Yifei Li, Deli Chen, Y. Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations, 2024.
Tianchun Wang, Zichuan Liu, Yuanzhou Chen, Jonathan Light, Haifeng Chen, Xiang Zhang, and Wei Cheng. Diversified sampling improves scaling Ilm inference. ArXiv, abs/2502. 1102,, 2025.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdh- ery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models, 2023.
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Selecting influential data for targeted instruction tuning. ArXiv, abs/2402.00333, 2024.
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Yang Yue, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model?, 2025.
Jifan Zhang, Lalit Jain, Yang Guo, Jiayi Chen, Kuan Lok Zhou, Siddharth Suresh, Andrew Wagen- maker, Scott Sievert, Timothy T. Rogers, Kevin Jamieson, Robert Mankoff, and Robert Nowak. Humor in AI: massive scale crowd-sourced preferences and benchmarks for cartoon captioning. CoRR, abs/2406. 10522, 2024.
Yiming Zhang, Harshita Diddee, Susan Holm, Hanchen Liu, Xinyue Liu, Vinay Samuel, Barry Wang. and Daphne Ippolito. Noveltybench: Evaluating language models for humanlike diversity. 2025. URL https:/ api. swmanticscholar... orgg/orrpusD::2776111111
13

---

Published as a conference paper at ICLR 2026
Qwen2. 1122.5-0.55 .58 (Base)
Owen2.5-1.55 (Base)
0.8
Pass@k
Standard training
0.7
0.7
ModC w/ separate models
0.6
Standard training balanced
MadC w/ randon split
0.5
0.5
0.4
2
k
(a) Pass @k performance Countdown natural test
Qwen2.5-0.35 (Base)
Qwen2.5 112..-1.58 (Base)
0.8
Pass@k
Standard training
0.6
0.6
MadC w/ separate models
Standard training balanced
0.4
randon split
0.2
2
2
k
(b) Pass@ & performance on Countdown adversarial set
Figure 8: Ablation studies on Countdown. ModC with random paritioning sometimes shows gains does not outperform ModC with DFS/BFS partition. Balanced training data DFS/BFS distribution does not show gains compared to standard training.
A LLM USAGE
We used large language model (LLL) solely to refine the writing style. The LLM played no role in research idcation, experimental design, analysis. All technical content, results, and conclusions entirely our own, and we take full responsibility for the final manuscript.
B ADDITIONAL RESULTS FOR Countdown
We see that ModC with DFS/BFS partition improves test-time scaling on Countdown. As control, we also random partitioning the training data into the same number of groups. From Figure 8, see that ModC with random paritioning sometimes shows gains but does not outperform ModC with DFS/BFS partition. Another baseline try is enforce 50-50 distribution of DFS and BFS in the training data, which we do not see any gains compared to standard training.
C ADDITIONAL RESULTS FOR GRADIENT CLUSTERING
Gradient clustering recovers teacher identity. We first validate gradient clustering on multi- teacher data where ground-trute truth labels exist. Using the short CoT dataset from Section 4.2, compute gradients using Qwen2.5- 5-Bae 1.5B and apply clustering. We first observe that gradient clustering achieves 98.7% F1 score in recovering teacher assignments. More importantly, Figure shows that ModC with these automatedly discovered gradient- clusters yields nearly identical test-time scaling benefits using teacher labels, This confirms that gradient patterns effectively capture the underlying modes.
14

---

Published as a conference paper at ICLR 2026
Mixed teachers
-
ModC w/ prefix (teacher)
ModC w/ prefix (gradlent)
Qwen2.5-0.0m
(Base)
0wen2.5-1.58 (Base)
Qwen2.5-38 (Base)
Qwen2.5 (Base)
0.8
0.8
0.6
8.8
0.6
0.6
0.6
0.2
29
22
24
25
20
22
2
2
28
20
22
2
2
29
20
2
26
29
k
k
Validating gradient clustering on multi-teacher data. ModC on training data that is distilled from multiple teachers outperforms standard training even without access to teacher identity annotation on training data. ModC with gradient clustering almost completely matches ModC with access to teacher annotations.
Figure 9:
15


Reviews:
[{"review_id": "vWCet7FgC6", "reviewer": "Reviewer_MR1M", "summary": "The paper proposes mode-conditioning, a test-time inference strategy that allocates a certain number of samples to each mode in order to improve diversity of samples, mitigating the issue of mode collapse. The authors show that ModC leads to consistent gains across tasks both when modes are fixed, as well as being able to discover modes automatically via gradient clustering.", "questions": "- Most of the ablations are comparing variants of ModC. Could you provide a comparison of ModC against other diversity-inducing techniques (see comment in weaknesses section)?\n- Do you see a clear diversity increase after ModC? For example, for MATH500, if you consider how many distinct answers are produced for each problem, how much does it increase with ModC?\n- Does the idea also apply to other domains, such as code generation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The ModC method is novel, creative, and effective, addressing the critical issue of lack of diversity.\n- The authors demonstrate that ModC training works well on a variety of tasks. They explore the idea throughout a variety of settings and domains, and it performs above standard baselines in all cases. The technique seems to be quite general and could have potential downstream applications beyond those listed in the paper.\n- The authors also compare different ways of implementing mode conditioning, and do a thorough analysis on other factors like model size, CoT length, etc.", "weaknesses": "- The work does not have any comparisons with other diversity-inducing techniques, for example pass@k training (https://arxiv.org/abs/2508.10751) or optimal sample allocation (https://arxiv.org/abs/2410.22480). While ModC is evidently effective against simple baselines, it is difficult to understand the advantages and disadvantages of this method against some of these other methods.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rf6ibebTX4", "reviewer": "Reviewer_h3Wb", "summary": "The paper presents a simple, yet effective approach to diversify the model’s outputs by providing explicit control over the modes. It explores two methods: 1) training separate specialist models and splitting test-time compute between them, 2) training a single model with mode-specific prefixes, and sample equally with the corresponding prefixes. The paper shows that these approaches surpass the mixed model trained with both modes. Morover, in the case of unspecified modes, the paper proposes an automatic mode discovery method based on gradient clustering. It shows that the method captures the labels reliably and further mode-conditioning on them recovers the improvements.", "questions": "1. Could you please explain how the heuristic prunings limit the solution to one of BFS and DFS?\n2. How is the ‘fraction of BFS used’ computed?\n3. For the distilling experiments, what kind of prefix do you use for different teachers? How does knowledge sharing happen in those 4. experiments if the model learns to follow one strategy given a prefix?\n4. Could you explain how the gradient is computed in the gradient clustering method?\n5. Did you run the gradient clustering method for the long-CoT datasets too?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper pinpoints a simple but important suboptimality in training language models with diverse data. It verifies the intuition with experiments with both toy settings such as different strategies for Countdown, and with real-world tasks and traces distilled from teachers. It is comprehensive in experimenting with different forms of chain-of-thought (short and long) generated with different models. Moreover, the work pushes its practical relavance further by providing a method for discovering unobserved modes in the data based on gradient clustering, which makes the idea more generalizable to different settings.", "weaknesses": "The paper could improve its presentation by defining its metrics more clearly. For example, it’s not clear how the “Fraction of BFS per problem” metric is computed for Figure 2. In section 5.1, p_\\theta is not defined, so it’s not obvious how the gradient is computed. \n\nI also did not understasnd how heuristic prunings and search budget constraints make the problems solvable with only one of BFS and DFS, making it unclear why this setting captures the diverse setting desired. \n\nThe novelty of the idea to learn separate models and aggregating them instead of learning from a mixed dataset is questionable given the literature around mixture of experts and other works such as “Mix Data or Merge Models? Optimizing for Diverse Multi-Task Learning”.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vjVOqbAig0", "reviewer": "Reviewer_ghYh", "summary": "This paper studies mode conditioning as a way to address the lack of diversity in LLM generations for reasoning tasks. The problem is that when we scale test time compute with parallel sampling, current models tend to collapse to one or two dominant strategies, so additional samples mostly repeat the same errors. The paper’s proposal is to make the modes explicit and to allocate test time samples across modes rather than drawing all samples from a single collapsed distribution. They instantiate this in a controlled search setting (Countdown, a generalization of Game of 24) where the target value can be found either by a DFS style search or a BFS style search, and where the search trace itself reveals which mode was used. They then extend the idea to math post training with multiple teacher models and finally to a setting where modes are discovered automatically via gradient clustering.  They show that mode conditioned training and mode conditioned inference improves Pass@k relative to standard mixed training.", "questions": "1. Unclear algorithm in Section 4.2 - The post training described for math reasoning in Section 4.2 seems to be plain SFT on two teacher traces with either mode specific prefixes or separate models. It would be good to state clearly that no RL was used here, if that is in fact the case.\n2. Figure 4 interpretation - The caption says that the dark gray line is “best teacher.” Does this mean this curve corresponds to distillation from only the best single teacher, not distillation from the union of best teacher traces across problems?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I think the problem is interesting and well chosen. How to obtain diversity in reasoning style without simply increasing sampling temperature (which has its own issues) is still not well understood, and most current post training pipelines / RL algorithms do in fact make diversity worse.\n\nI like the synthetic setup with countdown game since it cleanly isolates the question they are trying to answer, with a way to verify which mode of problem solving is used.  The experiments are pretty thorough and they also show some nice results in the math CoT setting, as well as in the automatic mode finding setting using gradient clustering.", "weaknesses": "1. Several parts of the paper are somewhat hard to follow on first read. One example is Figure 2. It is not completely clear how the per problem histograms are computed. My reading is that for each test problem the authors sample the model repeatedly, detect for each sample whether the model used DFS or BFS, compute the fraction of BFS samples for that problem, and then plot the distribution of that fraction over all problems. If that is correct, the number of samples per problem needs to be stated. If that is not correct, the figure needs a more explicit description. Right now it is difficult to tell what exactly is being compared. \n2. In the separate model setting each mode gets its own model trained on the subset of data for that mode. In Countdown the paper notes there are about 97k DFS trajectories and 65k BFS trajectories. If each of those is used to train a full model of size $N$ then the total training compute for the separate model setting is roughly $6 \\times (2 N) \\times (97+65)/165 \\approx 11.7 N D$  flops, whereas the standard training or prefix based modC is roughly 6ND .  This would mean that the separate model setting is using roughly twice the training compute.  The paper should clarify whether training budget was controlled, whether epochs were scaled down for the separate models, or whether the comparison is intentionally not compute matched. As written, it is not a fair comparison.\n3. A natural application of this work is to post train with RL, where we know that the distribution becomes sharper and diversity decreases. As far as I can tell, the paper only considers SFT / distillation like settings. A natural question is: after RL, does mode conditioning still preserve the benefits shown here, or does RL erase them. It would be useful to see in the same synthetic Countdown setting a comparison of RL that samples from the usual policy versus RL that is constrained to use mode conditioned sampling during rollouts. If the authors can show that they can keep the sampling gains after RL training (evaluating 0-shot after RL) that would be a nice finding, even if on a synthetic task.\n4.  There are a few grammatical and formatting issues. For example Section 5.1 appears to have an incomplete closing sentence. (Did not penalize for this.)\n\nI am willing to improve my score if the authors can meaningfully address these comments.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i6y0KOSS6Z", "reviewer": "Reviewer_myXD", "summary": "This work suggests a method called ModC for improving diversity in generation during parallel scaling.\n\nThe idea is that different ways to approach a problem may fall into different \"modes\", corresponding to different broad strategies. A diverse generator of possible proofs should try to sample as diversely as possible from different modes.\n\nThis paper suggests two possible ways to do this: (1) train a separate model on each mode, or (2) train a model with prefix-tuning to make it use one mode.\n\nModes can be either (1) known a priori, or (2) found automatically with a gradient clustering method. The paper tests the idea on several benchmarks such as NuminaMath, AIME, and Countdown, and finds benefits over vanilla parallel scaling -- especially when there is a large amount of parallel scalin.g", "questions": "Typos: \"this achieves up to xx% improvement\"", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* ModC has a conceptually clear motivation.\n\n* The proposed method is simple to implement, and seems to yield increased performance. This might be of interest to much of the ICLR  community, since methods to improve to model reasoning are quite popular.\n\n* The paper presents a way to find modes automatically, using gradient cluster. This makes it more broadly applicable than it would be if the modes had to be created manually.\n\n* The experimental methodology seems mostly sound (although I have a question -- see weaknesses below).", "weaknesses": "* On the methodology: I'm not sure how good of a metric pass@k is for AIME, when k = 1000, because there are only 1000 possible solutions for any problem as far as I know. \n    - Having a model that outputs a random number from 0 to 999 would give a 63% pass@k accuracy, which is roughly the accuracy  reported in Figure 5. \n    - On the other hand, having 1000 models (each of which outputs a constant number) would give a 100% pass@k accuracy.\n\n* There's a quite relevant prior work called \"Metadata Conditioning Accelerates Language Model Pre-training\" by Gao et al., 2025, that this made me think about. There they show that adding metadata of which website a text came from can improve model performance. It could be good to discuss the connection with this work.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)