TS$^2$: Training with Sparsemax+, Testing with Softmax for Accurate and Diverse LLM Fine-Tuning
Authors: XuZiyang
Ananthu Rajendran Pillai
Yinghua Yao
Yuangang Pan
Venue: ICLR 2026 Poster (2026)
Keywords: Large Language Models
Supervised Fine-Tuning
Output Diversity
Primary area: foundation or frontier models, including LLMs

Abstract: Large Language Models typically rely on Supervised Fine-Tuning (SFT) with Cross-Entropy (CE) loss to specialize in downstream tasks. However, CE forces the distribution toward one-hot targets and ignores alternative continuations, thereby limiting output diversity, a key drawback for generative applications that rely on sampling-based exploration.
In this paper, we propose ``Training with Sparsemax$+$, Testing with Softmax (TS$^2$)''. Intuitively, sparsemax and its tailored loss mask the gradients of probabilities outside the support set, leaving excessive probability mass on irrelevant tail classes when evaluating with softmax. To address this issue, we propose an improved variant, Sparsemax$+$, for training, which augments the sparsemax loss with a suppression term that penalizes the out-of-support probabilities. At testing, we decode with softmax, yielding calibrated, non-degenerate probabilities where plausible near-ties survive.
We fine-tuned Llama-3.1-8B and Qwen-2.5-7B with TS$^2$, achieving consistent improvements in accuracy and output diversity across chat, code, and open-domain benchmarks. Together, these results demonstrate that TS$^2$ provides a practical, drop-in solution for fine-tuning LLMs that are both more accurate and more creative. 
The code is available at https://github.com/xzy-bit/TS-2-ICLR-2026.

Published as a conference paper at ICLR 2026
TS TRAINING WITH SPARSEMAX+, TESTING WITH SOFTMAX FOR ACCURATE AND DIVERSE LLM FINE- TUNING
Ziyang Xu'; Ananthu Rajendran Pillai?; Yinghua Yao* Yuangang Pan
National University of Singapore, Singapore
2 "Nanyang Technological University, Singapore
3 Centre for Frontier AI Research, Agency for Science, Technology and Research, Singapore
Institute of High Performance Computing, Agency for Science, Technology and Research, Singapore el 4499203uu nus. edu, anant hurpilaas 7 com,
{yao. man.yyinhua, pan yuangang) (a aattarrst..r... edu. sg
ABSTRACT
Large Language Models typically rely Supervised Fine-Tuning (SFT) with Cross-Entropy (CE) loss to specialize in downstream tasks. However, CE forces the distribution toward one-hot targets and ignores alternative continuations, thereby limiting output diversity, a key drawback for generative applications that rely sampling-bon based exploration. In this paper, we propose "Training with Sparsemax+, Testing with Softmax (TS²)". Intuitively, sparsemax and its tailored loss mask the gradients probabilities outside the support set, leaving excessive probability mass on irrelevant tail classes when evaluating with softmax. To ad- dress this issue, we propose an improved variant, Sparsemax+,, for training, which augments the sparsemax loss with a suppression term that penalizes the out-of- support probabilities.
testing, we decode with softmax, yielding calibrated, non-degenerate probabilities where plausible near-ties survive. We fine-tuned Llama-3.1-8B and Qwen-2.5-7B with TS², achieving consistent improvements in accuracy and output diversity across chat, code, and open-domain benchmarks. Together, these results demonstrate that TS² provides a practical, drop-in solution for fine-tuning LLMs that are both more accurate and more creative. The code is available at pestps:/gqithinn wwwp:/yoooxccxxeeeee bit 2 XCLR-2028
INTRODUCTION
Supervised fine-tuning (SFT) is one of the major steps in the Large Language Models (LLMs) post- training stage: with a small amount of high-quality annotated data, it teaches models to organize language better and produce instruction-followwe followiigg responses. The default loss function is cross- entropy loss, mainly because it coincides with maximum likelihood and a strictly proper scoring rule, minimizing it recovers the data generating conditional under well-specification (Gneiting & Raftery, 2007). However, the same geometry drives the probability mass toward the one-hot target and away from plausible alternatives, yielding overconfident posteriors and reduced useful diversity. A large body of work seeks to counteract this overconfidence and recover useful diversity. One branch changes only the decoding, e.g., nucleus sampling and best-of-N. leaving training dynamics and calibration untouched (Holtzman et al., 2020). Another branch alters the training signal itself. The recent GEM framework reframes SFT as reverse- KL minimization with an entropy regular- izer, improving variety and mitigating overfitting (Li et al. (025). These approaches highlight fundamental issue: promoting diversity can conflict with keeping probabilities calibrated and tails disciplined.
We argue that the field lacks a precise operational notion of useful "diversity" for instruction follow- ing. In many tasks, we do not want to "spread probability" indiscriminately over the entire vocab-
"Equal contribution. Work was done during their internship at Centre for Frontier Al Research, A SSTAR. "orresponning author.

---

Published as a conference paper at ICLR 2026
1.0
Pretmined
GEM
TSPoours
Meanlingless ess Taken
C.5
Pooabiliity
0.2
0.
Figure 1: Token Distribution for single digit generation (detailed in Appendix C.4
ulary. Instead, want probability mass concentrated among a handful of semantically plausible next tokens, those with real chance of leading to a high quality continuation, while aggressively deflating the long tail of obviously incorrect tokens toward (near) zero. The right diversity within the plausible set, not across the whole simplex. The forward KL KL (ipp|1 q) is mean- mean-seeking, incen- tivizing probability wherever the data has support; the reverse KI. (T.(g I p) is mode-seeking, con- centrating mass on promising regions (Minka 2005). This lens helps explain why CE with entropy maximization (a -L-aall-e--l objective under softmax) can inflate low-probability tokens, while reverse-KL flavored objectives like GEM avoid gratuitous tail mass. Yet even reverse- KL does not guarantee that clearly implausible tokens go to zero.
Our approach takes a geometric route by decoupling the mapping from logits probabilities between training and testing. Specifically, we train with spansemax and optimize a modified Fenchel-Young loss tailored this mapping (Martins & Astudillo, 2016) Blondel al. 2019), while at inference we revert to softmax, which restores calibrated and smooth probabilities on the same logits. Our tailored loss contains tail penalty that drives non-support tokens
while ensuring the gold token is never penalized, even if it lies outside the instantaneous sparsemax sup- port. Notably, CE with softmax collapses diversity: all non-gold logits, even plausible ones, are pushed toward zero. In contrast, sparsemax maintains sparse support set by zeroing gradients of non-support tokens, preserving plausible candidates. However, if sparsemax were also used at inference, a converged model would still produce one-hot outputs, similar to CE with softmax de- coding, thus limiting diversity (Martins Astudillo 2016) Blondel et al.) 2019). Figure [sshows that the pretrained model naturally exhibits diversity, but such diversity is lost during SFT: CE training drives the distribution into a one-hot solution, while GEM manages to retain only a few candidates, with most remaining probability mass assigned to irrelevant symbols. In contrast, method de- livers stronger and more stable diversity: the candidate set is both meaningful and varied, striking a balance between plausibility and coverage.
Figure[sshwss
This decoupled recipe. Train with Sparsemaxt, Test with Softmax (TS?), has two key effects. Dur- ing training, sparse gradients dates on already separated tail candidates. At inference, reverting to softmax restores smooth, cal- ibrated probabilities so that plausible near-ties survive and sampling can explore them without ag- gressive temperature tuning. By construction, our method achieves local diversity among plausible tokens while assigning near-zero confidence to implausible ones. We position TS² among comple- mentary strategies. Inference-only methods (e.g., nucleus, top-k, best-of-N) improve sample variety but leave training untouched; our approach reshapes training dynamics while remaining fully com- patible with such decoders (Holtzman et al. 2020). Entropy targeting methods (e.g., GEM) promote spread but do not enforce exact zeros on implausible tokens; our penalty term supplies this "hard" suppression, while sparsemax ensures spread occurs where it matters (Li et al. 2025). Finally, be- cause T decouples mappings rather than altering model architecture, it integrates seamlessly into existing SFT pipelines. The contribution of this paper summarized in the following:
principled early-stopping mechanism by avoiding wasted up-
We frame the problem as achieving Suppeespeee Plausible Diversity (TSPD) and propose TS? which decouples training and inference by using Sparsemax+ loss with tail penalty for training and standard softmax for decoding.
We provide a theoretical analysis showing how TS² avoids the distributional collapse common to CE training via gradient-masking mechanism, thereby preserving diversity inference.
We demonstrate in practice that our T significantly improves winrates, sample efficiency in code generation, and output diversity across multiple benchmarks compared to existing methods.
2

---

Published as a conference paper at ICLR 2026
2 DISTRIBUTION COLLAPSE AND OUR INSIGHT
Recent studies have observed an "alignment tax" in large language models (LLMs): while super- vised tine-tuning (SFT) improves faithfulness and task adherence (Brown et al. 2020 often comes
at the cost of reduced output diversity and partial forgetting of pre-trained knowledge (0 O'Mahony et eta.. 2024, Kim et al. 2025). Pre-trained LLMs naturally exhibit broad generative repertoire, pro- ducing multiple semantically valid outputs for the same prompt (Wang et al. 2025). However, after SFT, models tend to respond with highly deterministic and homogeneous output (Li et al. (025), weakening their utility in downstream applications such planning (Song et al. (023), writing (Lee) ct al. 2022p, or code generation (Liu et al., 2023), all of which fundamentally rely on the ability to explore diverse candidate responses.
A central obstacle in supervised fine-tuning is that cross-entropy (CE), driving the predictive distri- bution towards one-hot distribution, causing all probability mass to collapse onto the gold token. This distribution collapse ensures convergence, but comes at a severe cost: the model suppresses all alternatives nearly zero, leading to deterministic outputs, both in the choice of tokens and in the semantic content of the whole responses. The mechanism destroys diversity, erasing helpful variations preserved in the pre-trained distribution and thereby yielding monotonous generations.
2.1 OUR GUIDING INSIGHT: FAIL-SUPPRESSED PLAUSIBLE DIVERSITY
In generative modeling, output diversity is essential. A target distribution should retain a compact set of plausible candidates with non-negligible probability while suppressing irrelevant long-tail tokens toward zero. We formalize this as Tail-Suppressed Plausible Diversity (TSPD), which remedies the distribution collapse commonly observed in existing SFT.
Notation. We consider prompt-response pairs (ar, v) E DD from a supervised dataset Let so denote a pre-trained LLM parameterized by 0. For prompt v, let fo(z)
(x) R. denote the
corresponding logit vectod ! We define the probability simplex as A AK- {p E R Ps 2
0, E Ps = 1}. where = g(z) denotes probability distribution obtained from the logits via a probability mapping function g(-).
K
Definition SSppr-Sup Plausible Diversity (m, head, Etail)). Given prompt-response pair (x,y), let p = g(fo (z)) A be distribution over a vocabulary Fix an integer m 2 2 and thresholds 0 Ehead S and 0 S Stail < 1 Sheadd Let Top, (p) denote the indices of the m largest coordinates of p. If E Top, (P), let Topm (p); otherwise, S = 1
( (y))
We say that p satisfies TSPD of order m if
(Head Preservation)
min Pj 2
Ehead
(1a)
JES
(Tail Suppression)
Pj
S
Etail-
(Ib)
JES
which ensures that candidates in S retain non-negligible probability, whereas tokens outside S re- ceive essentially zero probability, thereby preserving uncertainty and transferable knowledge at in- ference. one chooses Ehead 1/mm exactly, then the strict requirement Etail 2 forces mEhead and p; = Vj & S: therefore, in practice one can take Ehead < 1/m and relax Etail > 0.
Corollary 1. Definition holds and Etail < Chead, then max Pj Etaill < Eheud S minniss Pi. so each plausible sample has strictly higher probability than any tail sample.
Corollary 2. If all probability mass collapses onto the ground-trute truth token, i.e., Py and Py 0 Vy' # then p fails to qualify the TSPD (m(2 2), Ehead, Etail..
In the next section, we motivate our method that operationalizes this principle, directly countering the diversity-reducin bias of CE loss while retaining the benefits of supervised fine-tuning.
y and he sequential, where an auto-regressive formulation used.
3

---

Published as a conference paper at ICLR 2026
3 ACHIEVING TAIL-SUPPRESSED PLAUSIBLE DIVERSITY
A natural way to realize TSPD in Equation 1 is to exploit the sparsity of the sparsemax mapping sparsemax(2) (Martins & Astudillo, 2016). which projects logits z I2 onto the probability simplex A KK1 and can assign exact zeros to non non-support tokens.
Let E R and let 2(1) 2 2(2) 2 2 2(K) denote the sorted coordinates of z. Define k(z) :=
max k: [l,...,.......... K} + 2())) Σラ= #63 }. and the threshold r(2)
>((a) RG)
k(z)
The sparsemax probabilities are then given elementwise by
p. (=) = (2)a) = max ( - (=2), 0}, i= 1,...,,, K.
The (data-dependent) support set is (=) = {i E (I....,,K) (2) > 0}. which follows from the definition above and circularity arises. Equivalently, in vector form, (=) [= - (2))1]
where [o]+ == ma 0} is applied elementwise.
In effect, sparsemax automatically identifies a compact support set of plausible candidates (29(2)
and prunes away the long tail. Compared to the softmax probability mapping p" (z) exp(x) == softmax(z), its Jacobian is sparse; in particular, gradients vanish outside (*9(2)
Ee exp(x.)
when the target lies the support (Lemma 3).
Lemma (Gradients vanish outside the sparsemax support (Martins & Astudillo 2016)). Ler
sparsemax(2) and (z) ) be its support. Define L(p, v) as a supervised loss between the sparsemax probability p and the target y. If (*)(2), then Vi & S*P(2),
0C(z.v) 0.
While sparsemax provides margin-induced sparsity, it nonetheless tends to collapse into nearly one-hot distribution once the leading logit surpasses the margin threshold. Such collapse inevitably reduces sampling diversity, making sparsemax undesirable for inference.
This motivates vanishing property established in LemmaBremain advantageous during training: by nullifying gra- dients outside the active support whenever the target is included, it mitigates the cross-entropy style erosion of plausible near-optimal alternatives, thereby inducing an implicit early-stopping effect.
instead carry out decoding with softmax. Under this choice, the gradient-
Theorem (Sparsemax expands pairwise gaps faster than softmax). Let z E TR paf =
softmax(2), 2)), and pdd sparscmax(z). For any indices i + j. let u 24 2j and have
0 sp
(p.
Po 50 ) =1 -
Vi,j E 5TP
sparsemax
Ou
ou (v 13) <
softmax
Given the same logits, Theorem shows that sparsemax linearly preserves pairwise probability gaps within its active support and collapses one-hot prediction once finite margin attained, whereas softmax strictly contracts such gaps. Consequently, sparsemax induces sharp discrimination and faster label collapse during training, while applying softmax to the same logits at inference preserves non-degener mass on plausible candidates, maintaining output diversity that is desirable for generative tasks.
Corollary 5 (Softmax remains TSPD valid when sparsemax is one-hot). Let z E R- with y
arg max; 2j. and 8y := Zy -z. Assume sparsemax is one-hot at i.e., 8min == minj by 2 > 0 (e.g., Y = 1), and the top-m head is bounded: 6(k) := By 2(k) I B Vk 2...., Set
Am = m + (K m)e Then for pd = softmax(z) we have
B
(K - m)e
py
Am
P(A)
Am
(Vk- = 2, 2....,m,,
Σ k>m
Am
Consequently, pd satisfies TSPD of order m with any thresholds < Eboad
S Am
B
(K (K-m)e- 2
Am
Etail < 1 m Shead-
4

---

Published as a conference paper at ICLR 2026
Remark 1. Without the head bound 6(z) S B (vk S m), Plm) can be made arbitrarity small even when min 2 1, so only vanishingly small head floor Ehead can be guaranteed for general
Remark 2 (Existence of a tight upper bound for the cumulated tail mass Eksm P(k) in Corollary 5 Under the assumption of Corollary
1 + (m -1)e-B + (K m))7 then we have With the new upper bound, yd still satisfies TSPD of order mn with any
5 let Pmin
Exsm P())) 1 (K (K-m)eT
Skaii
thresholds 0- < Ehead derivation.).
(K (K-m)e"
S Etail
1 1 Ehead- (see Appendix D for detailed
Am
According to Remark 2 the upper bound on the cumulated tail mass Σk>m P(a) is strictly increas- ing in K and approaches 1 as K - Importantly, this bound is tight (see Proposition c in Appendix DD: there exists a worst-case configuration where the tail mass grows monotonically with K. Thus, for large vocabularies, the admissible tail under softmax at inference becomes nearly 1, indicating that sparsemax training has not theoretically guaranteed the suppression of irrelevant tail mass, contradicting the goal of suppressing irrelevant tail mass.
To address these issues, we propose a fine-tumin finettuing strategy of Training with Sparsemax Testing with Softmax. Sparsemax+ builds Sparsemax, inheriting margin- -inguced sparsity to introduce gradient masking during training, thereby implicitly enforcing an early-stopping effect once the top-1 candidate is clearly separated. It further incorporates lightweight Tail-suppressing Loss to explicitly penalize residual probability on tail tokens, ensuring that tail mass is sharply suppressed. At inference, revert to softmax over the same logits, which restores smooth, calibrated proba- bilities across the plausible candidates within the support set, while keeping the irrelevant tail mass negligible due to the additional suppressing effect. In this way, the model learns to separate and prune the logits during training, yet preserve and diversify the output distribution during inference, achieving the desired support-awan diversity.
4 TS²: TRAINING WITH SPARSEMAX+, TESTING WITH SOFTMAX
In the following, we present supervised fine-tuning based on the Fenchel-Young loss, which encom- passes both the softmax and sparsemax mappings. It then motivates our Sparsemax+ loss.
4.1 DIFFERENT PREDICTION MAPPINGS WITH THE UNIFIED FENCHEL- Loss
For any strictly convex regularization function K-1 RR the corresponding regularized prediction function is (z) = arg maxpeaK-1 (p, 2) [[()). The associated Fenchel- Young loss can be represented as
Ln(2zy) I(e,)) II(p.) ) + (z.p--or ey)-
(3)
where y is the gold label and e, the corresponding one-hot vector. Different choices of II yield different prediction mappings and losses.
Softmax Softmax corresponds to using the negative Shannon entropy as regularizer TT(p)
ΣK Ps log Pt. which gives (x) F exp(2s) == softmax(z). The Fenchel- Young loss re- duces to the standard CE loss Leottmax
exp(z)
(2: (z: = log Σ exp(z;) - = - log > E exp(2:)
exp(2.)
Sparsemax Sparsemax corresponds to using the negative Gini entropy as regularizer II(p) "K Pr(1 - P.). which gives (x) = (=-r(z)1]__ = sparsemax(z). The cor- responding Monchel-Yuung loss, called the sparsemax loss, is Luparsemma (z:1) = -2y + ; E)es((<)) (o) - (2))) (=))) + 1-
In conclusion, when training with sparsemax but performing inference with softmax, although softmax(z) does not yield one-hot output like sparscmax(z). it still assigns the highest probabil- ity to the correct class. Importantly, it naturally enables early stopping and preserves distributional diversity across all classes, which consistent with the goal of diversifying plausible candidates.
Given pronpt-responss response pairs (z, y) from supervised dataset, let 2 E R be a logit vector and p be probability mapping either via softmax sparsemax. If the gradient of the sparsemax loss
5

---

Published as a conference paper at ICLR 2026
vanishes, i.., V Spaasprseaax(22 (2; 0, then it follows that sparsemax(z), 1. For any index Vj * y, sparsemax = 0, holds that softmax(=), > O. That is, softmax assigns non-zero probability to all entries, including those which sparsemax maps zero. According to Corollary [
the cumulated tail mass of softmax outside the top-m satisfies Σx>m P((k) S [K=m)e- With
A...
large vocabularies, the admissible tail under softmax at inference becomes nearly 1. This behavior
undesirable, as assigning non-negligible probabilities to clearly incorrect classes may lead the model to produce semantically meaningless outputs.
Sparsemax+ + To address this issue, we introduce a lightweight tapp-suppressing loss that explicitly suppresses probabilities assigned to the non-plausible candidates. Given logits E IR let pd softmax(2) E A-1 The tail-suppressing loss is defined as
Caup (p: y) log(1 (1- > EPdr.
$4S
where S is defined in Definition] 1 thereby avoiding residual mass on clearly implausible candidates.
This penalty drives the probabilities of tail tokens toward zero,
Remark 3. The tail suppressing loss can be interpreted as a direct generalization of the standard sofimax CE to the group-label setting. Specifically, given logits and softmax distribution pd softmax(2)) z), the suppressing term can be written as
[z)p(2) log(11
ΣE - log Σ P. $
iES
which is exactly the softmax cross-entropy with the larget label being the merged "super-class" S. In the special case where $ = {y] is a singleton, this reduces to the usual CE loss log p.j Thus, the suppressing loss can be viewed as encouraging the softmax probability mass to concentrate on a set of plausible candidates while retaining the probabilistic interpretation of cross-entropy.
Combining sparsemax with the suilsuppressing loss yields our proposed Sparsemax+ loss:
Lspm (z:y) Zy +
Σ ( + (-log log (1
Σ
pi
p:))
(4)
(ES*(=)
igS*(=)
where r(z) is the sparsemax threshold and O controls the strength of the suppression. For sim- plicity, we find that directly implementing the candidate set from Definition I support (z) achieves superior performance.
Jusing the sparsemax
We summarize our fine-tuning strategy of Training with Sparsemax+,, Testing with Softmax in that it prevents CE-style erosion of plausible near-ties by amplifying relative ratios among top logits while nulling the rest, thereby achieving two goals: sparsemax selects stable support set with early stopping of gradient flow, and the suppress- ing term explicitly drives unreasonable tokens toward zero to prevent spurious mass at inference.
Algorithm [ From Lupm+ (z; y) in equation 4
5 EXPERIMENTS
To situate our work within the current state-of-the-arr, we build upon the experimental foundation of GEM (Li al. 2025). adopting similar training setup. Our primary methodological difference is the substitution of the GEM objective with our proposed TS² loss. Furthermore, while GEM eval- uates OpenLLM Leaderboard tasks using standard one shot setting, we employ a multi-response, best-of-N protocol. We argue this is a more faithful and informative evaluation for diversity aware models, as it measures model's 's latent ability to find the correct answer rather than penalizing it for plausible "hesitation" in a single attempt.
Setup. We conduct experiments on two powerful, open source base models: Llama-3, 1-18B and Qwen-2-7B For supervised finetuning, we use the high quality UltraFeedback dataset (Cui) et al. 2024). large-scale corpus of preference aligned responses generated by a diverse set of models. All models are finetuned for 3 epochs using the AdamW optimizer with an effective batch
size of 128. We employ cosine learning rate schedule with an initial rate of 2 x 10 1 and a
warm-up ratio of 0.03, standard practice for fine-tuning modern LLMs Yu et al. 2024 Liu al.

---

Published as a conference paper at ICLR 2026
Algorithm TS²:
Training with Sp Testing with Softmax
Input: pre-trained model fø; training dataset Dir = {(x, y)}; test dataset Die {z}.
Hyperaarameters: epochs T'; batch size B; learning rate 7j 0; suppression weight 0.
1: for to 7 do
b. Training Phase
2:
for mini-batch {(rr, 96) 36)) Dr do
3:
Compute logits 2B < fo( 1.2.......
4: Compute loss Lb < Lisput (21; 9b).vb 1,2, 1,2,....
Sparsemax+ loss
5: end for
6: Update - VV Σ
2b=1
7:
end for
8: for test input Die do
D. Testing Phase
9: Compute logits < fo(
Predict probability p <- softmax(z)
11:
Evaluation on
Use for decoding
12:
end for
The maximum sequence length capped at 2,048 tokens. For our proposed TS? method, the suppression weight (see Equation 4 is empirically determined for each model architecture, with optimal values reported alongside results. Further implementation details are provided in the Appendix B
2024].
We compare TS² against a suite of strong and relevant baselines to provide a comprehensive evalu- ation: Cross-entropy (CE): The standard SFT objective, which serves as our primary baseline. CE with Weight Decay (CE+WD): A common regularization technique shown to help preserve diver- sity in instruction tuning (Ouyang et al. 2022, Bai et al. 2022). We use weight decay coefficient of 0.1. NEFTune (NEFT): A regularization method that adds noise to word embeddings during training to mitigate overfitting and improve generalization (Jain et al. 2023). GEM: The current state-of-the-art method for diversity preserving SFT, which we use as main point of comparison Li et al. 2025).
5.1 IMPROVING ACCURACY AND DIVERSITY IN OPEN-ENDED GENERATION
We first evaluate TS² in open ended domains to assess its ability to navigate the critical trade- off between response quality and diversity. While standard fine-tuning often improves quality at the cost of collapsing output distribution, we hypothesize that TS² can break this trade-off by simultaneously enhancing generation quality and fostering rich, useful diversity beneficial for sampling-base based decoding. To test this, we evaluate on two distinct benchmarks. For conversational chat, use the AlpacaEval dataset (Dubois et al. 2024) with a best-of-33 3f-32 (BoN @ 32) protocol; a state-of-the reward model, FsfairX- Mittaiix-LLLMM LLaMA3 RM-v0.1 (Lambert et al. 2024), selects the best response, which is then compared against GPT-4 determine win rate. For code generation, we measure the pass@k metric on the EumanEval benchmark (Chen et al. 2021). which assesses the model's ability generate functionally correct Python code via execution.
Model
Method
Win Rate (%) f
N-gram
100
Self-BLEU
Sent-BERT
CE
29.77
17.78
47.04
9.97
CE+WD
29.72
17.78
47.14
10.03
LLaMA-3.1-8B
NEFT
29.77
17.74
47.41
10.07
GEM
31,53
20.32
49.82
11.16
TS² (Ours)
33.12
23.78
53.87
12.80
CE
31.41
17.23
16.77
7.95
CE+WD
31.05
17.43
17.08
8.06
Qwen-2-7B
NEFT
30.36
16.59
24.59
8.06
GEM T (Ours)
33.89
24.35
31.19
9.25
37.48
30.15
39.04
9.81
Table I: Win rate (Best N@32) and diversity metrics for Llama-3. 1-B and Qwen-2-7B on Al- pacaEval. T achieves the best results across both quality and diversity on both architectures.

---

Published as a conference paper at ICLR 2026
(a) LLaMA- Chat Task
(b) LLaMA Code Task
90
32
NEFT
80
30
GEM
(S<()urs)
Pass Rate (%)
Win Hate
26
24
20
@2
@4
28
816
@92
@10
@20
@50
@100
Figure 2: Performance of Llama-3.1-88 open-ended tasks. Left: Win rate AlpacaEval sampling budget (N). Right: Pass rate on HumanEval vs. sampling budget (k). TS² consistently outperforms baselines.
Performance Chat and Code Generation. As shown in Figure 2 TS² demonstrates clear performance advantage on Llama-3.1-8P In chat generation, its win rate at budget of N=32 re- sponses. reaches 33.12%, which is improvement of 11.2% relative over the bascline cross entropy loss and a 5.0% relative improvement over the strong GEM baseline. This advantage extends to structured problem solving, on HumanEval, TS² achieves a pass 100 of 87. 87.00%, which is 4.3% increase relative to GEM and 19.8% to that of CE. Notably, the diversity fostered by method translates to superior sample efficiency: the pass @ rate for TS (82.70%) nearly matches GEM's pass @ 100 performance (83.40%), indicating that correct solutions can be found with fewer sam- ples. Similar results are also observed for Qwen-2-7B model. Detailed breakdown of results for both Llama-3.1-8B and Qwen-2-7B are detailed in the Table
Crucially, these performance gains do not come at the cost of diversity. As detailed in Table T TS² not only achieves the highest win rate but also scores best across all three diversity metries. improves N-gram diversity by 17.0% ,BLEU diversity by 8.1% and sentence-bert diversity by 10.7% over GEM for LLama-3.1-8D Similarly for Qwen-2-7B, the same metrics are improved by 23.8%, 25.1 % and 6% respectievely over GEM. This result confirms that TS² successfully breaks the quality-diversin trade-off, producing responses that are simultaneously judged as higher quality by reward model while being measurably more varied.
5.1.1 DIVERSITY ON CREATIVE WRITING TASKS
To further probe the nature of the diversity generated by TS² we evaluate it on purely creative tasks: generating, poems from 573 titles in the poetry8 dataset and stories from 500 prompts from ROCStories (Mostafazadeh et al., 2016). As shown in Table 2 TS² once again achieves the highest scores across all three diversity metrics on both tasks, confirming its ability to produce wider range of high-quality, creative outputs compared to all baselines.
Poem
Story
Method
N-gram
100 Sel--BIfBLL Sent-BERT
N-gram
100 - Self-BLE-UB Sent-BERT
CE
38.87
$5.38
14.83
67.20
22.15
CE+WD
38,92
55.69
14.17
44.43
67.26
22.22
NEFT
38.80
55.68
14.13
44.31
67.21
22.04
GEM
46.59
57.50
14.70
50.05
69.15
24.02
TSa (Ours)
49.70
59.41
16.52
52.10
70.36
24.98
Table 2: Diversity evaluation on creative writing tasks for Llama-3.1-8B Higher is better.
5.2 PRESERVING PRE-TRAINED CAPABILITIES ON STANDARD BENCHMARKS
To assess generalization and knowledge retention, evaluate models six tasks from the Open- LLM Leaderboard: ARC, GSM8K, HellaSwag, MMLU, TruthfulQA, and WinoGrande. Instead of the standard greedy one-shot decoding that penalizes models preserving multiple reasoning paths,
8

---

Published as a conference paper at ICLR 2026
propose a best-of-n (BoN) strategy on the OpenLLM leaderboard, which is better aligned with evaluating the capabilities of Dnversity-preservve models.
(a) LLaMA-3.1-8B
(b) Qwen2-7B
95
CE+WD
90
NEFT
CE
85
Accuracy (%)
GEM
80
(S*(Ours)
75
70
60
@
@4
@8
@16
@32
@
@4
@8
@16
@32
Sampling Budget
Sampling Budget
Figure 3: Average Best-of-N accuracy across six OpenLLM Leaderboard tasks. While competitive few-shot settings @ 2), TS performance scales far more effectively with the sampling budget, revealing its superior knowledge retention.
We argue that a more faithful metric is Best-of-N (BoN) accuracy. This protocol measures the model's 's latent ability identify the correct answer within small sampling budget, which better reflects the true underlying capabilities a well-calibrated, diverse model. For fair comparision, all methods are evaluated under the same BoN protocol and report the average accuracy across all tasks.
Figure validates this hypothesis. While all methods competitive small sampling budget, TS²'S performance scales significantly better as 'N (responsees) increases. On Llama-3. 1-8B, the massive 13.2-point absolute (+17.4% relative)
average accuracy of 'T N=32 reaches 88.88% improvement over GEM (75.69) (7).69)). The trend is consistent on Qwen-2-7B, where TS² again achieves the highest accuracy, demonstrating the robustness of our TS across different model architectures.
This shows that TS² effectively preserves the model's pre-trained knowledge. Unlike CE, which collapses the distribution and discards valid alternatives, TS² maintains clean, calibrated set of high-quality reasoning paths. With sampled responses, the model consistently finds the correct solution. detailed breakdown of performance on each of the six tasks is provided in the Table 10
5.3 ABLATION STUDY
To assess the contribution of each component, we run ablation study AlpacaEval, comparing win rate against GPT-4 and BLEU diversity. TS² integrates three elements: (1) sparsemax-based training, (2) softmax decoding, and (3) tail-suppression penalty. We evaluate three variants: De- coupling Only (sparsemax training, softmax inference. no penalty), Unified Sparsemax (sparsemax for both training and inference), and Suppression Only (CE loss with suppression term).
Figure demonstrates that all components of TS² are essential. First, using the Decoupling Only strategy results massive increase in diversity, high BLEU diversity score, but a catastrophic drop in win rate. This shows that while decoupling unlocks variety, the suppression penalty is crucial for ensuring that this diversity high-quality and not just uncalibrated noise.
Conversely, the Unified Sparsemax approach achieves competitive win rate but offers lesser diver- sity than our full method. This confirms that the switch to softmax at inference is key to translating the leamned logit geometry into rich, sample-able probability distribution. Finally, applying the Suppression Only penalty to a standard CE baseline fails on both metrics, proving it not stan- dalone improvement but works in synergy with the maassemax-defined support set.
Meanwhile, the TS? method successfully integrates these components, achieving the best balance of high win rate and high diversity across both model architectures. This analysis confirms that the sparsemax objective, the decoupled inference, and the suppression penalty are all necessary and synergistic elements of our approach.

---

Published as a conference paper at ICLR 2026
Liama-3.1-88
Qwen2-7B (Ours)
Aupeosent Infer
35
TD (Ours)
Cuest
30
Winrate
25
20
15-
10
20
40
60
80
0
10
20
30
40
50
60
100
Self-BLEU
100
Self-BBEU
Figure 4: Ablation study on Llama-3.1-8B and Qwen-2-7B.
CONCLUSION
In this work, we make the first step toward decoupling training and inference by adopting different prediction mappings in supervised finetuning. By combining Sparsemax+ loss; tailored design that leverages margin induced sparsity with an additional suppression term non plausible tokens; with softmax decoding at inference, our approach achieves significant improvements over existing SFT paradigms. preserves support-aware diversity while maintaining high accuracy, thereby al- leviating the alignment tax. Despite its simplicity, our method consistently outperforms CE and GEM across both chat and code tasks, achieving the highest win rates and more diverse generations. Unlike prior methods that inevitably trade off diversity against accuracy, our paradigm improves both, providing a natural remedy to distribution collapse and open up new directions for advancing alignment with broad and long term impact.
ETHICS STATEMENT
This work investigates new algorithms for supervised fine-tuning of large language models. Our objective is to improve training stability and output diversity, thereby broadening the range of down- stream applications. The methods introduced in this paper are purely algorithmic and evaluated on public datasets.
REPRODUCIBILITY STATEMENT
Experiment details for reproducing our numerical results can be found in Appendix B and Ap-
pendix C
LLM USAGE STATEMENT
We used large language model to correct grammar errors, polish the writing, and adjust the format- ting of the paper.
ACKNOWLEDGMENT
This work was supported by the National Research Foundation, Singapore under its National Large Language Models Funding Initiative (AISG Award No: AISG-NMLP-2024-000 findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore.
Any opinions,
REFERENCES
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislay Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson
10

---

Published as a conference paper at ICLR 2026
Kernion, Tom Conerly, Sheer El-Showk, Nelson Elhage, Zac Hatfiel--Dodte Danny Hernan- dez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Ka- plan. Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022. Utt.tttss:://// http//aavvvv//.o///vv//..///.avl wrg. abs 2204. 05862
Mathieu Blondel, André F. T. Martins, and Vlad Niculae. Learning classifiers with fenchel- young losses: Generalized entropies, margins, and algorithms. In Proceedings of the Eventy-Secons International Conference on Artificial Intelligence and Statisties, volume 89 of Proceedings of Machine Learning Research, pp. 606-615. PMLR, 2019. URL https: http:///oooediieee ngs.
mlr. wwtpp/eeeeeeeoeeeeeeeeeeeeeeeeee
html
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger. Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings. Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss,, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin. Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Mc- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code, 2021. URL htttt::/
org/ / abs / 21 0.007.. 03374
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Bingxiang He, Wei Zhu, Yuan Ni, Guotong Xie, Ruobing Xie. Yankai Lin, Zhiyuan Liu, and Maosong Sun. Ultrafeedback: Boosting language models with scaled ai feedback, 2024. URL
Yann Dubois, Xuechen Li, Rohan Taori, Tianyi Zhang, Ishaan Guirajani. Jimmy Ba, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpacafarm: A simulation framework for methods that learn from human feedback, 2024. URL httphhttt////itiii wwtpwawwwwo.////333.3...
14387
Tilmann Gneiting and Adrian E. Raftery. Strictly proper scoring rules, prediction, and estima- tion. Journal of the American Statistical Association, 102(477):359-378, 2007. doi: 10.1 1198 01621450600001437.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1321-1330. PMLR, 2017.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL https: //openreview. httt//orrmmid=rroooyyo.
Neel Jain, Ping yeh Chiang, Yuxin Wen, John Kirchenbauer, Hong-Min Chu, Gowthami Somepalli, Brian R. Bartoldson, Bhavya Kailkhura, Avi Schwarzschild, Aniruddha Saha, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Neftune: Noisy embeddings improve instruction finetuning, 2023.
Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds. Nova DasSarma, Eli Tran-Johnson, Scott Johnston, Sheer El-Showk, Andy Jones, Nelson Elhage, Tristan Hume, Anna Chen, Yuntao Bai, Sam Bow- man, Stanislav Fort, Deep Ganguli, Danny Hermandez, Josh Jacobson, Jackson Kernion, Shauna Kravec, Liane Lovitt, Kamal Ndousse, Catherine Olsson, Sam Ringer, Dario Amodei, Tom Brown, Jack Clark, Nicholas Joseph, Ben Mann. Sam McCandlish, Chris Olah, and Jared Ka- plan. Language models (mostly) know what they know, 2022. URL https: : /arxiv. org/ abs/ /20//227... 05221
11

---

Published as a conference paper at ICLR 2026
Jiyeon Kim, Hyunji Lee, Hyowon Cho, Joel Jang, Hyeonbin Hwang. Seungpil Won, Youbin Ahn, Dohaeng Lee, and Minjoon Seo. Knowledge entropy decay during language model pretraining hinders new knowledge acquisition. In The Thirteenth International Conference on Learning Representations, 2025.
http://///peeerveeeeeeeeeee net.
Robert Kirk, Ishita Mediratta, Christoforos Nalmpantis, Jelena Luketina, Eric Hambro, Edward Grefenstette, and Roberta Raileanu. Understanding the effects of rlhf on Ilm generalisation and diversity, 2024. URL httptt:ss////xx...ooom
abs/ 2310. 06452.
Nathan Lambert, Valentina Pyatkin, Jacob Morrison, LJ Miranda, Bill Yuchen Lin, Khyathi Chandu, Nouha Dziri, Sachin Kumar, Tom Zick, Yejin Choi, Noah A. Smith, and Hannanch Hajishirzi. Rewardbench: Evaluating reward models for language modeling, 2024. URL
btttss///arxiis org. attp://///////..../.3
Mina Lee, Percy Liang, and Qian Yang. Coauthor: Designing a human-ai collaborative writing dataset for exploring language model capabilities. In Proceedings of the 2022 CHI conference on human factors in computing systems, pp. 1-19, 2022.
Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Ruoyu Sun, and Zhi-Quan Luo. Preserving diversity in supervised fine-tuning of large language models. International Confer- ence Learning Representations, 2025. URL
net wwweffoumm" forum?id
94CYQccSMM
Stephanie Lin, Jacob Hilton, and Owain Evans, Teaching models to express their uncertainty in 14334
words, 2022. URL
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chat- gpt really correct? rigorous evaluation of large language models for code generation. Advances 2023.
in Neural Information Processing Systems, 36:21558-2141
Wei Liu, Weihao Zeng, Keqing He, Yong Jiang, and Junxian Hc. What makes good data for align- ment? a comprehensive study of automatic data selection in instruction tuning, 2024. URL
aittps://axiiv otgg abs/ 22112 15685
André F. T. Martins and Ramón Fernandez Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. In Proceedings of the 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1614-1623. PMLR, 2016. URL/tttps: http:////ooeeeeeeeeeee ml prwss v48/v/48/////ee /tt./maaiinnsnle htmml
Sabrina J. Mielke, Arthur Szlam, Emily Dinan, and Y-Lan Boureau. Reducing conversational agents' overconfidence through linguistic calibration, 2022. URL https://arxiv..rm
abs/2012.14983
Thomas Minka. Divergence measures and message passing. Technical Report MSR-TR-2005-173, Microsoft Research, 2005.
Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vander- wende, Pushmeet Kohli, and James Allen. A corpus and evaluation framework for deeper under- standing of commonsense stories, 2016. /RLLtttps://aatt //axiii wwwp:xiwwworr.oooo...... org/ hwtp./11......1100
Laura O' O'Mahony, Leo Grinsztajn, Hailey Schoelkopf, and Stella Biderman. Attributing mode col- lapse in the fine-tuning of large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024. URL https: /ttp::///eeeeeeeeeeeeeeee .ne.//
Eorum? 3pDMY jpOxk.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models follow instructions with human feedback, 2022.
Chan Hee Song, Jiaman Wu, Clayton Washington, Brian M Sadler, Wei-Lun Chao, and Yu Su. Llm-planner: Few-shot grounded planning for embodied agents with large language models. In Proceedings ofthe the IEEE//VF international conference on computer vision, pp. 2998-3009, 2023.
12

---

Published as a conference paper at ICLR 2026
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818-2826 2016.
Katherine Tian, Eric Mitchell, Allan Zhou, Archit Sharma, Rafael Rafailov, Huaxiu Yao, Chelsea Finn, and Christopher D Manning. Just ask for calibration: Strategies for eliciting calibrated con- fidence scores from language models fine-tuned with human feedback. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023, URL net//ozzue htt//oorummici ii- geldddggaaacccc 3faCErwm?
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.092. 2023.
Evan Z Wang. Federico Cassano, Catherine Wu, Yunfeng Bai, William Song, Vaskar Nath, Ziwen Han, Sean M. Hendryx, Summer Yue, and Hugh Zhang. Planning in natural language improves LLM search for code generation. In The Thirteenth International Conference on Learning Rep- resentations, 2025. URL (tRp:tt/: /owpnreeiveieeeeeeee
net htt.///rrmmmiiiiimeemmmoo (tw.rrruunn...44////////HwH.
Yuxiang Wei, Zhe Wang, Jiawei Liu, Yifeng Ding, and Lingming Zhang. Magicoder: Empowering code generation with oss-instruct, 2024. URLhctps. //tp///aavvvv/vvv///a//eee
02120
Sean Welleck, Ilia Kulikov, Stephen Roller, Emily Dinan, Kyunghyun Cho, and Jason Weston Neural text generation with unlikelihood training. In International Conference on Learning Rep- resentations, 2019.
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T. Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models, 2024. htthstts::/ www.xii.oor/gaaoo abtpg//ba///// 12284
A RELATED WORK
Our work, T intersects with three primary areas of research: supervised finetuning (SFT) and its inherent limitations, methods for enhancing generative diversity in large language models (LLMs), and the use of sparse activation functions in neural networks.
A.1 SUPERVISED FINETUNING AND THE ALIGNMENT TAX
Supervised finetuning is a major landmark in adapting pre-trained LLMs to downstream applica- tions, enabling them to follow instructions and adhere to specific conversational styles (Ouyang ct al. 2022 Touvron et al., 2023). The standard practice involves minimizing a cross-entropy (CE) loss dataset of high quality datasets. While effective, this approach is known to induce an "aligmment tax" " (O Mahony et al.) 2024). where models become overly specialized to the finetuning distribution. This often leads "knowledge forgetting" or a collapse in output diversity (Kim et al., 2025) Li et al.) 2025). The CE objective, by driving the model's posterior towards one-hot representation aggressively penalizes all alternative continuations, including those that are semantically plausible. This results in overconfident and deterministic models. Our work directly addresses this limita- tion by replacing the CE objective with a loss that preserves a set of plausible next-tokens, thereby mitigating the distributional collapse and retaining more of the pre-trained model's capabilities.
reduction in creative capacity, phenomenon sometimes termed
the target token,
Overconfidence and miscalibration have been extensively studied in deep neural networks more broadly, where modern architectures are known to produce poorly calibrated probability estimates even when their accuracy is high (Guo et al. (017)). Recent work has begun to document and address similar phenomena in large language models and conversational agents. Mielke et al. (2022) show that dialogue agents can be systematically overconfident and propose linguistic calibration strategies to temper their stated certainty, Kadavath et al.] (2022) analyze when LLMs "know what they know" and how their confidence estimates relate to factual correctness, and Tian al.) (2023) and Lin ct al. (2022) develop post-hoe calibration methods for RLHF and SFT-trained LLMs by verbalized confidence calibration probes. These approaches primarily operate top of fixed model, either
13

---

Published as a conference paper at ICLR 2026
by rescaling scalar confidences modifying prompts at inference time. By contrast, TS² targets the source of overconfidence in SFT itself: we reshape the token-level logit geometry during training so that plausible alternatives preserved within compact support set while the softmax tail is explicitly suppressed.
A.2 ENHANCING GENERATIVE DIVERSITY
Efforts to counteract the loss of diversity finetuned LLMs can be broadly categorized into decoding-time and training-time strategies.
Decoding- Strategies: popular line of work focuses on modifying the sampling process at inference. Techniques such as temperature scaling, top-k sampling, and nucleus (top-p) sam- pling (Holtzman et al., (2200) manipulate the output probability distribution to encourage variety. Similarly, best-of-N sampling. where multiple candidate responses are generated and ranked by a reward model (Bai et al. 2022), can improve output quality by exploring a wider search space. While widely used and effective, these methods are applied post-hoc and do not address the under- lying overconfidence of the model's learned distribution. TS² is complementary to these techniques but fundamentally different, as reshapes the logit geometry during training to produce a more inherently diverse and well-calibrated posterior.
Training- Trame Strategies: Another branch of research modifies the training objective itself. Label smoothing (Szegedy et al. 2016) is a regularization technique that discourages overconfidence by training on soft targets. More recently, unlikelihood training was proposed to explicitly penalize undesirable tokens or repetitive patterns (Welleck et al. 2019), Closest to work is the recent GEM framework (Li et al. 2025), which recasts SFT as a reverse-KL minimization problem with an entropy regularizer. GEM successfully improves diversity by preventing the model's posterior from collapsing. However, it does not enforce a hard separation between plausible and implausible tokens, potentially leaving residual mass the long tail of the distribution. TS² offers more direct approach: the sparsemax function provides a principled mechanism for identifying compact sup- port set of plausible tokens, while our proposed suppression penalty explicitly drives the probability of out-of-support tokens to zero, achieving a cleaner and more decisive separation.
A.3 SPARSE ACTIVATIONS IN NEURAL NETWORKS
The sparsemax function, which we leverage for our training objective, is a projection onto the proba- bility simplex that produce exact zeros (Martins & Astudillo 2016). It was originally introduced as a sparse alternative to softmax for attention mechanisms and structured prediction tasks, valued for its ability select a small subset of relevant inputs. The sparsemax loss is specific instance of a Monchel-Young loss, broader class of losses that provides a unified framework for various structured prediction mappings (Blondel al. 2019). While sparsemax has been explored for clas- sification and attention. its application to generative LLM fine-tuning for diversity preservation is novel. Critically, our work is the first to propose decoupled paradigm: we use the desirable prop- erties of sparsemax (e.g.. gradient masking for non-support tokens) during training but revert to the smooth, fully-supported softmax for inference. This decoupling is the key to unlocking calibrated diversity, concept not explored in prior work that typically uses the same mapping for both training and testing.
B EXPERIMENT DETAILS
We conduct all training on 1200-141GB GPUs, employing the DeepSpeed framework with ZeRO-2 optimization and gradient checkpointing enabled. Offloading disabled. For efficient and repro- ducible training, we adopt flash-attention-2 with deterministic backward passes. Our base models
are Llama- 3 5.1 1-8B and Qwen-2 2-7B, optimized using Adam W with a total batch size of 128. The learning rate is initialized decay schedule, as suggested by prior work (Yu et al., 2024 Liu et al. 2024 Cui et al., 2024)... Training is run for 3 epochs. All supervised datasets reformatted into the chat style with the Llama-3, 1 1-8B-- Instruct and Owen2. 5-77------- Instruct tokenizer. For inference, we em- ploy vI.LM to accelerate response generation.
2 x 10-5 with a warm-up ratio of 0.03 and follows a cosine
14

---

Published as a conference paper at ICLR 2026
The supervised finetuning is done the binarized UltraF eetraFeedbbck dataset curated by the Hug- gingfaceH4 team which contains 61,135 training examples and 1,000 held-out test prompts. Inputs longer than 2,048 tokens are truncated, while shorter ones padded. To achieve global batch size of 128, we use 4 GPUs, each with per-device batch size of and gradient accumulation of 4. A single training run requires roughly 12 GPU hours. For CE+WD baselines, the weight decay coefficients is 0.1. For NEFT, we set the noise scale to 5, consistent with Jain et al. (2023).
Evaluation Protocol For chatting,
805 prompts from Al AapaacEval and score outputs with the F Fsfairt aMA3 -MM--MM-------- 1 reward model. The maximum decoding length is 2,048, and each prompt yields 32 samples using temperature=0.6. top-k=50, and top-p=0.9. Win rate is com- puted against GPT-4 rreponssss via the Brudley-Terry model:
exp(r(x..))
Ply > y/(=)= exp(r(æ, v)) + exp(r(x,
For code generation, we adopt the HumanEval benchmark (164 Python problems). Prompts follow the template of (Wei al., (024)).
You are an exceptionally intelligent coding assistant that consistently delivers accurate and reliable respenses to user instructions.
ea Instruction (instruction)
For each task, we sample 200 outputs with the same decoding configuration. The evaluation metric pass rates, which are computed using execution-based evaluation scripts from Magicode
C ADDITIONAL RESULTS
C.I SENSITIVITY OF PAIPER--PRAMETERS
Our T framework provides two hyper-parameter to control the final behavior:
Training-time control. The hyperparameter a in Sparsemax+ equation determines the strength of suppression during training.
Training time control via logit scale c. There is also an internal logit scale in I (c2; y), which plays role analogous temperature. In all reported experiments we simply set = which already yields strong performance.
Inference-time control. The inference of our TS² relies on softmax. Therefore, the tem- perature parameter can also be adjusted influence the final output. We fixed it to I which works well.
In our formulation, the hyperparameter a directly controls how aggressively the Sparsemax+ trans- formation suppresses negative (low-logit) components. When loss is the pure sparsemax loss, lots of low logits keep their mass, yielding a extreme diverse out- put(see Fig As increases, the suppression let distribution becomes more concentrated on the top few logits, which will let the diversity The results are showing in Tab In practice, we typically choose 0 < a 1, which allows the first part dominates when updating.
there's no supersession, the
C.2 EXTRA BASELINES
For completeness, additionally evaluate two overeonfidence--ittee mittigttinn baselines that com- monly used as alternatives to standard cross-entropy: CE + label smoothing(O.2) and o-Entmax (a 1.5). Both methods modify the training distribution to reduce overconfident predictions. Tab reports the AlpacaEval (Chat) win-rate(BoN@ 32) results across all methods on Llama-3.1-8b.
zed
preview
py
15

---

Published as a conference paper at ICLR 2026
Table
Effect on Chat (c=1.0, Llama-3.1-8b)
Setting
Winrate (%)
n-gram
100-self-BLEU
Sent-BERT
o=0.25
33.12
23.78
53.87
12.80
a=0.50
31.44
20.54
51.34
11.61
o=0.75
30.84
19.07
49.29
10,96
a=1.00
30.39
18.13
48.05
10.30
a=5.00
25.34
13.44
41.42
8.12
Table Comparison of TS", GEM, CE, sparsemax, Entmax, and CE with label smoothing on chat task MlkAAlaaaa-vall
TS (ours)
GEM
CE
CE label smoothing
sparsemax
1.5-Entmax
Winrate (%)
33.12
31.53
29.77
28.25
12.87
13.51
While both CE + label smoothing and 1.5- ----Ettmxx introduce additional diversity compared to vanilla CE, this comes at a substantial loss in win-rate, and both remain far below GEM and TS² Sparsemax and Entmax also do not explicitly regulate cumulative tail probability: label smoothing over-flattens the distribution, whereas sparsemax/entmax produce sparse supports without suppressing softmax- style tails at inference. These behaviors reduce helpfulness and alignment quality.
In contrast, TS combines sparse training objective with an explicit tail-suppression mechanism while retaining softmax decoding. This yields the highest win-rate among all compared methods, while still improving diversity, suggesting that TS achieves more favorable balance between confidence calibration, output variability, and instruction-followii performance.
C.3 MULTI-TURN TASK PERFORMANCE
We additionally evaluate our method on MT-Bench-100. fine-grained multi-turn dialogue bench- mark covering detailed conversational capabilitie After generating the response of each question, we using GPT-4.2 4.1 nano to score them from 13 metrics, see Table 5 Across all capability groups, method achieves consistent improvements, for example, +0.95 in CM, +1.67 in CC, and +1.18 in deCoonstttraiing gains in reasoning coherence, contextual consistency, and conversational sta- bility. Complete results shown in Table These findings indicate that approach generalizes effectively to multi-turn dialogue scenarios.
C.4 MACRO- AND MICRO-LEVE ANALYSIS OF TOKEN DISTRIBUTIONS
To understand why our method simultaneously improves accuracy and diversity, we analyze token probability distributions from two complementary perspectives: (i) a macro-level analysis of model outputs real benchmark, and (ii) micro-level controlled probing task.
Macro-level distribution. We evaluate the models(L) ama) on the AlpacaEval dataset. For each generated response, record the probability of every selected token and compute the average probability of that response. We then plot these values across all responses to obtain global view of the distribution. As shown in Figure CE exhibits the highest mean probability (~ 0.90) with the smallest variance, indicating collapsed and overly uniform predictions. GEM lowers the mean probability to about 0.86 with larger variance, consistent with its entropy-regularized updates that discourage overconfidence. Moving along the sequence CE + GEM + Sparsemax (sparse inference) + TS² we observe a systematic trend: mean probability decreases (remaining above 0.8), while variance increases, revealing more balanced allocation of probability mass to plausible alternatives.
16

---

Published as a conference paper at ICLR 2026
Abbr.
Task
Deseription
CM
Context Memory
Recall early dialogue details address the user's current query.
SI
Separate Input
First turn provides task requirements; later turns supply the actual input. Identify pronoun referents across dialogue turns.
AR
Anaphora Resolution
TS
Topic Shift
Detect and follow the new topic when users change subjects unexpectedly. Avoid interference from similar-looking but semantically distinct queries.
CC
Content Confusion
CR
Content Rephrasing
Rephrase the content of the previous response based new requirements.
FR
Format Rephrasing
Reformat previous response into structure requested by the user.
SC
Self-Correction
Correct the previous response according to user feedback.
SA
celf-Affimmaton
Retain the previous response when the user provides inaccurate feedback. Solve multi-turn mathematical reasoning tasks collaboratively. Solve multi-turn general reasoning tasks collaboratively. Ask clarifying questions when the user's query is ander-specified. Proactively ask questions to sustain or deepen the conversation.
MR
Mathematical Reasoning
GR
General Reasoning
IC
Instruction Clarification
Pl
Proactive Interaction
Table 5: Descriptions of the 13 MTT-eenc- 101 dialogue capability dimensions.
Model
Avg.
Perceptivity Understanding
Adaptability Reflection
Interactivity Questioning IC
Memory CM
Interference TS
Rephrasing CR
Reasoning MR
SI
AR
CC
FR
SC
SA
GR
CE
5,99
5,01
4.59
6.03
5.10
5.03
7.33
7.02
6.34
7.38 6.37
4.67
7.49
5.46
GEM
6.24
4.63
4.43
6.88
5.95
5.36
7.19
7.76
7.2
8.05
6.09
5.05
6.54
5.98
TS² (Ours)
6.65
5.96
4.76
6.58
5.94
6.70
8.27
6.60
7.42
8.53
6.03
6.14
6.80
6.76
Table 6: Comparison on MTTBeench-10
sparso sparsomax Infor
Probabiity.
Hoen 1.0
0.8
Avarage
6432
12879
10010
20759
6432
12379
18012
20759
6423
22279
19329
20750
6430
32879
19019
20759
Response idx
Figure 5: Macro-level analysis: average selected token probability distribution on AlpacaEval.
Micro-level probing. To complement the macro-level view, we design controlled probing task to test whether models can distribute probability mass across relevant candidates. We prompt the model with the few-shot instruction to generate a single-digit number. Each model is queried 100 times. Whenever a digit is generated, we record the probability distribution of the top-300 tokens. Finally, compute the average probability of each token across the 100 trials, resulting in a fine-grained view of how probability mass is allocated.
You' re an Al assistant, will give you an example of following quest on. Example:
User: Give me word of fruit.
Assistant: Apple.
Now you follow the format the example, Give me a single-digit number,
Answer:
The results, shown in Figure D reveal stark differences. CE collapses to one-hot distribution: the chosen digit monopolizes probability, while the tail is filled with irrelevant tokens. GEM re- tains few candidate digits but remains nearly one-hot, yielding limited diversity. Sparsemax (Sparsemax-infer distributes mass across more digits, but still assigns non-negligible probability to spurious tokens. In contrast, T combines sparsemax, which preserves probability relevant
17

---

Published as a conference paper at ICLR 2026
digits, with the suppressing loss, which eliminates unrelated characters. This synergy results in distributions that are both diverse and accurate.
0.1
1.0
0.8
Probablitty
0.6
0.2
0.4
0.1
Std: 0 5005
5000
10000
15000
20000
25000
3
Figure 6: Macro- and Micro-level probling: Sparsemax Training and Softmax Inference
As a special case, we also examine the strategy of sparsemax training with softmax infer- ence(shown in F6gure(6). In the macro-level probing, this setting produces distribution that is close to uniform, suggesting that model does not exhibit clear preferences over candidate tokens. In the micro-level probing lask, we observe that although some valid numerical answers (such as "42" or "125") appear, a large number of irrelevant tokens also receive comparable probability mass. As result, the model's outputs become difficult interpret, and its effective generation ability is di- minished. This illustrates why conventional diversity metries may report artificially high scores in this case: while probability is spread across many tokens, much of it corresponds to spurious rather than meaningful outputs.
C.5
HARD THRESHOLD
0.25-top3
0.5-top3
1.0-to33
0.5-top5
0.5-top10
1.0-top10
0.25-TS²
0.5-TS2
50.00
50.16
52.58
48.80
49.86
47.58
46.61
53.73
51.03
Table 7: Results on Ll ala 3. 2-1B with different a and sparsification strategies. All tokens except target are thrown out the support set.
We also evaluate the Ll ama- 3.2 2- under different values and supersession strategies. Here, "top-k" means only the largest logits preserved in the defined support set, while all other tokens (except the target) are thrown out the support set, and "TS2" denotes our proposed two-stage suppression method. We take the vanilla cross-entropy (CE) training the baseline, which yields score of 50.
From the results Tab. 7 observe two main trends: (i) Increasing from 0.25 to 1.0 generally decreases performance, indicating that larger a values reduce the model's output diversity. (ii) Within the top-k setting, smaller (e.g., top-3 scores, while larger values dilute the distribution and hurt performance. Overall, both reducing a and carefully selecting smaller encourage the model maintain useful diversity, while our T method further boosts results beyond simple top-k truncation.
top-10) leads higher diversity and better
C.6 CHAT AND CODE GENERATION
Table g details the performance of both models on the open-ended generation tasks. For chat gen- eration, it presents the win rate against GPT-4 across various best-of-N sampling budgets (N = 2, 4, 8. 16. 32). For code generation, it shows the corresponding pass@k rates for = 1, 10. 20, 50, and 100.
C.7 CREATIVE WRITING
We further investigate output diversity on two creative writing tasks: poetry and short stories. For poetry, 573 titles drawn from the Huggingface poet ry8 dataset, which covers themes such
18

---

Published as a conference paper at ICLR 2026
LLaMA-3,1-8B
Qwen-2-7B
CF+WD
NEFT
CR
GEM
( 0.25)
CF+WD
NEFT
CR
GEM
TS 0.5)
@-2
20.14
19.35
19.88
20.26
19.43
18.35
18.13
18.4
18.06
18.72
@4
23.02
22.33
22.78
23.34
23.37
21.59
21.49
21.78
21.93
22.54
25.44
25.19
25.74
26.67
26.61
24.39
27.9
27.66
@6
27.82
27.51
27.97
29.26
29.85
27.76
27.02
27.9
30.02
32.77
@32
29.77
29.72
29.77
31.53
33.12
31.05
30.36
31.41
33.89
37.48
LLaMA-3.188B
Qwen-2-7B
Code
CE+WD
NEFT
CE
GEM
(a 0.25)
CE+WD
NEFT
GEM
TS 0.5)
@1
30.30
28.50
24.60
31,90
32.80
45.10
45.30
44,90
41.80
42.20
@10
60.40
57.00
55.30
64.80
67.00
76.80
76.50
76.00
78.50
78.20
@0
65.90
63.60
62.20
71.80
74.70
81.30
81.00
81.10
84.50
84.60
50
71.30
70.00
69.10
79.00
82.70
84.10
83.20
83.50
87.20
$7.80
74.50
73.30
72.60
83.40
87.00
87.20
85.40
90.20
91.50
Table Performance comparison different methods on LLaMA-3.1 1-8B and Qwen-2-7BB models for the chat and code geneartion tasks
as love, nature, and mythology. For stories, we construct 500 prompts from the ROCStories dataset [Mostafazadeh et al., 2016). In both settings, the instruction to write a piece titled "[X]" in under 200 words, where [X] sampled from the corresponding dataset.
Diversity is measured along three dimensions following Kirk et al. (202)): N-gram, the fraction of distinct n-grams within single response (intra-diversity): Self-BLEU, computed by treating each sample as the reference for the others (inter-diversit)): Sentence-BERT dissimilarity, the mean cosine distance between generated responses in the embedding space. All scores are scaled the range [O, 100). with higher values indicating greater diversity.
For evaluation, each model generates 16 completions per prompt using temperature=00.6 top-k=50, and top-p=0.9. Results are summarized in Table 9 It is evident that methods such as CE+WD and NEFT bring only marginal improvements diversity. GEM consistently improves intra- and inter-diversity, while TS² achieves the highest scores.
Method (Llama-3.1-8B)
Poem
Story
N N-gram
100
Self-BLEU Sent-BERT
N-gram
100
Self- BLEU
Sent-BERT
CE+WD NEFT
38.92
55.69
14.17
44.43
67.26
22.22
38.80
55.68
14.13
44.31
67.21
22.04
CE
38.87
55.38
14.83
44.47
67.20
22.15
GEM TS (a 0.25)
46.59
57.50
14.70
50.05
69.15
24.02
49.70
59.41
16.52
52.10
70.36
24.98
Poem
Story
Method (Qwen-2-7B)
N-gram
100 Selt-BLLLL
Sent-BERT
N-gram
100 Self-BLEU
Sent-BERT
CE+WD
44.29
44.9
8.66
56.62
50.06
19.01
NEUT CE GEM
44.37
45.09
8.55
59.66
52.2
18.94
43.94
44.92
8.56
56.44
49.83
18.86
50.29
48.62
60.91
56.05
20.98
TSa (a
53.46
51.10
10.26
62.13
57.17
20.95
Table 9: Diversity evaluation on creative writing tasks (poem and story). Higher values indicate greater diversity (N-gram, 100 Self-BLEU,, and Cintence-BERT
C.8 OPENLLM LEADERBOARD TASKS
Table 10 reports results six representative OpenLLM leaderboard tasks under varying sampling budgets. These benchmarks collectively reflect a broad spectrum of model capabilities: ARC focuses grade-school science questions, reflecting commonsense reasoning; GSM8K requires multi-step solutions, capturing mathematical reasoning; HellaSwag emphasizes physical commonsense and narrative continuation, probing contextual understanding; MMLU spans 57 subjects, testing broad
19

---

Published as a conference paper at ICLR 2026
factual knowledge; TruthfulOA challenges models with common mnscconceptions, measuring bustness; and WinoGrand ande coreference benchmark, assessing pronoun disambiguation and line-grained language understanding.
Building on this setup, we observe that other methods exhibit only limited gains as the sampling budget increases. In contrast, TS² consistently improves performance across tasks, achieving the largest boosts under larger budgets, often surpassing all baselines by a substantial margin. The improvements are especially pronounced for LLaMA-33 8B, where diversity-oriented training translates into 10-15 point gains under BoN sampling. For Qwer: 2 2-7B, whose baseline win rates already exceed 90%, the relative gains appear smaller but still confirm the benefits of preserving diversity during training.
(a) Llama-3.1-88
Method
ARC882
ARC#4
ARCØ8
ARC@ 16
ARC@32
Hellaw0
Hella@
Hella8
Hellaädi 16
Hella
CE+WD
75.59
80.27
80.27
83.28
83.61
66.95
70.55
72.95
74.37
75.05
NEFT
81.61
65.45
CE
76.59
79.50
71.27
82.60
83:00
67.03
61.96
63.06
63.85
64.33
GEM
78.60
82.27
83.94
85.28
85.61
66.51
71.71
74,84
76.96
77.95
TS 0. 0.25)
78.93
85.95
88.96
90.30
91.63
63.47
76.76
84.43
88.55
90.81
Method
Wino6e2
Wino604
Wino&8
Winood 16
Winooe 32
MMI 1/84
MMLIE0
MMLUG 16
MML 11
CE+WD
59.65
61.48
63.30
64.01
64,96
60.75
63.12
65.42
66.93
67.98
63.85
69.19
CE
59.75
80.27
80.27
83.28
83.61
60.36
66.23
67.90
68.90
GEM
61.80
64.64
66.69
68.43
69.85
62.01
66.12
69.32
71.85
73.54
[a 0.25)
66.14
75.77
80.51
83.98
87.21
64.19
73.64
81.24
85.82
88.42
Method
Truth622
Truth@24
Truth
Trith
Truth@:32
GSM@2
GSM@4
GSM@8
GSM@ 16
GSM@: 32
CE&WD
43.02
45.29
46.88
48.59
49.20
53.84
64.39
74.37
82.03
86.66
NEPT
46.74 74
50.06
51.41
52.39
53.12
54.21
66.87
75.06
82.03
85.67
CE
43.21
45.04
47.86
49.82
50.67
54.66
66.87
76.72
82.49
86.96
GEM
47.86
51.29
55.32
57.04
39.73
54.44
66.72
75.59
82,64
87.49
(a 0.25)
51.16
62.42
71.85
80.05
87.39
50.27
64.67
75.06
82.49
87.87
(b) Qwen2-7B
Method
ARC@2
ARC@4
ARC
ARC
ARC
Hella@22
Hella@84
Hella
Hella@ 16
Hella
CF&WD
84.61
87.62
88.96
89.96
89.96
80.20
85.50
88.76
90.67
91.44
NEFT
84.94
87.29
88.62
89.96
89.96
80.20
85.39
88.67
90.59
91.40
CF
84.61
87.95
89.62
90.30
90.30
80.45
85.60
88.79
90.62
91.42
GEM
84.94
88.96
92,30
92.60
93.64
80.01
87.54
91.94
94.52
95
(a
92.97
93.97
96.97
98.21
Method
Winot22
Wino
Wino Wino
Winooe32
MMLU8
MMLU8
CE+WD
70.79
77.34
81.84
83.89
69.42
77.86
$4.53
90.37
93.56
NEFT
71.19
77.26
$2.00
83.58
83.97
69.28
71.49
79.83
85.97
91.28
CE
70.63
77.26
82.16
83.66
84.13
69.37
76.65
83.89
89.46
93.40
GEM
73.63
81.76
89.50 39.50
90.37
69.91
93.
82
[o 0.5)
71.11
84.37
91.31
95.26
95.97
69.01
82.11
90.67
94,84
97.65
Method
Truth02
Truth
Truth/68
Truth 16
Truthe032
GSMG2
GSM@4 4
GSMG88
GSM 16
GSM@
55.07
55.93
66.72
NEFT
45.41
49.32
$2.50
$4.46
$5.69
67.32
71.49
79.83
85.97
91.28
GEM
46.87
52.14
$7.89
62.66
65.97
68.08
80.36
88.86
93.86
96.82
(a 0.5]
82.11
Table 10: Pass Rate (%) of Different Methods 6 OpenI.LM leaderboard tasks under Various Sampling Budgets.
C.9 DIVERSE OUTPUT EXAMPLES
Chatting Examples. To qualitatively illustrate the behavioral differences between CE-SFT and SFT, randomly sampled 32 generations from each model AlpacaEval for the prompt "I want to get better at networking at work". From these sets, we selected representative examples to showcase typical response patterns for both training methods. The results are shown in Tab. 11
20

---

Published as a conference paper at ICLR 2026
Prompt
I want get better at networking at work.
CE- SFT
Resp. I: Networking is an essential skill for personal and professional growth. Here are some tips to help you improve your networking skills at work: (1) identify your goals.... re- search...(3) be authentic....4) listen more than you talk, (5) follow up....66) be sclective....(7) offer value....(8) stay positive....11))kee track of contacts....(10) be patient....
crucial aspect of building a successful career. Here are some tips to help you improve your networking skills at work: (1) identify your goals...(2) research....33) dress appropriately....4) show up early.... (5) introduce yourself....(6) listen actively....7)) exchange business cardss...(8) follow (9)...9) be (lnjine....10))scck out oppor- tunities....
Resp. 2: Networking
Resp. 3: Networking a valuable skill that can help you advance in your career. Here are some tips to help you improve your networking skills at work: (1) identify your goals... (2) do your research....3)) be authentic....44) be preparcd....55) follow up....6) be patient.... All three CE-SFT samples reuse the same "Networking a a... skill" template followed by near-identical numbered-tip list, with only minor wording changes.
TSa.
Resp. I'd be happy to help you improve your networking skills at work! Networking a valuable tool for career growth and development, and there are many ways to do it ef- Hectively.Here are some strategies that may help: (I)Identify your g2))s....(2)Do your re- search...(3)Networrf within your arganizaiion...4]eevrraa events and conferencest... Volunteer for industry-related sonal brand....8)) good networking etiquette...
SFT
social (S)ten industry
(T)Esaab your per-
Resp. 2:Networking can greatly contribute to personal and professional growth. To become better at networking at work, consider following these tips: (1)Re- ceach...(9Comfidnnee.....
M4dia...(4)Bc present...(S)Iniiatte (fnggee actively...(7))e persistent....8)TTak your progress....
conversa-
Resp. 3: To improve your networking skills at work. consider the following tips: (1) Practice, Practice. Practice....2)Knovv Your Audience... Goal... (4) the 90/10 Rule....5))Akk Open-Ended Questions....((IPaaiiee Good Follow-Up....))fffee Val 't Be Afraid to Reach Out:...
TS² -SFT produces distinct yer coherent angles on the same prompt (internal networking, ex- ternal visibility, conversational tactics), illustrating higher response diversity while remaining on-topic and helpful.
Table 11: Qualitative comparison on an open-ended carcer-advice prompt. CE-SFT collapses to a single response template, while TS2-SFT yields diverse but still relevant completions.
For the prompt "I want to get better at networking at work", CE-SFT produces several highly similar completions: all begin with variations of "Networking is an essential/crucial skill..."" fol- lowed by list of roughly 10 generic recommendations such identify your goals, do your search, be authentic, follow up, be patient, appearing in nearly identical order and wording across samples. This reflects CE's tendency toward distributional collapse, yielding templated and homo- geneous responses.
In contrast, TS²-SFT generates much broader range of coherent yet distinct answers that emphasize complementary aspects of workplace networking. One response focuses on internal networking and personal brand building (leveraging colleagues, volunteering for cross-team work, building personal brand); another emphasizes confidence, presence, and communication skills; and third highlights concrete interaction heuristies such as the >90/10 rule" for active listening, using open-ended questions, and maintaining systematic follow-up routines. All TS² samples remain on-topic, helpful, and safety-aligned, yet differ substantially in structure, emphasis, and phrasing. This diversity aligns with our quantitative findings, showing that TS² preserves response variation without sacrificing coherence.
Math Reasoing Examples. TS² is applied at the token level and does not truncate or shorten reasoning chains; it only reshapes the logit geometry within each step by suppressing implausible tail tokens. To check that this does not harm multi-step reasoning, we inspected GSM8K generations with full thain-of-thought.
Table 12 shows three TS² samples for a GSM8K question. All completions follow the same two- step reasoning pattern (5 = 40 total pieces, then 40
14 26 pieces taken) and arrive the
21

---

Published as a conference paper at ICLR 2026
conrect answer, while differing in stylistic framing and level of explanation. In our broader manual inspection, we consistently observe this behavior: TS² maintains coherent, stepwise reasoning and correct final answers, with diversity manifesting in wording and structure rather than in contradictory or incoherent chains. A systematic, benchmark-level study of long-trace reasoning under TS² is an interesting direction; here we provide qualitative evidence that TS² preserves the integrity of reasoning traces while improving sample diversity.
GSM8K Ques-
Grandma Jones baked apple pies for the fireman's luncheon. She cut each pie into 8 pieces, At the the evening there were 14 pieces remaining. How many pieces were taken by the guests? Gold: 26
tion
-S=-SFT
Thank you for asking me to assist you! I'm happy to help you with your question. To answer your question, we need to determine the total number of pieces of pie that were taken by the guests. Let's analyze the information given: Grandma Jones baked apple pies, and she cut each pie into pieces. So, the total number pieces of pie she made is 5 = 40 pieces. The question states that at the end of the evening, there were 14 pieces of pie remaining. This means that 40 14 26 pieces of pie were taken by the guests. Therefore, the answer your question is: The guests took 26 pieces of pie. hope this answer is helpful and makes sense to you. Please let mc know if you have any other questions or if there's anything else I can assist you with!
Resp. 1:
Resp. 2: Okay, let's dive into the problem. First, let's establish the information we have: Grandma Jones baked 5 apple pies. She cut each pie into 8 pieces. At the end of the evening, there were 14 pieces of pie remaining. Now, let's determine the information we need to find out: How many pieces of pie were taken by the guests. To find this out, we can use simple subtraction problem: The total number of pieces of pie before the guests started eating = 5 pies x 8 pieces per pie 40 pieces. The number of pieces of pie remaining after the guests finished eating 14 pieces. Now, let's solve the problem: 40 pieces 14 pieces 26 pieces. The answer is: The guests took 26 pieces of pie. Here's 's the final answer: The answer is: 26 pieces.
Both the responses that TS² samples express complete and coherent two-step reasoning chain (compute total pieces, then subtract remaining) and agree the correct answer 26. Diversity appears in phrasing and level of detail, not in the underlying logic or facts.
Table 12: Qualitative T generations GSM8K problem with multi-step reasoning trace.
C.10 EMPIRICAL ANALYSIS OF SOFTMAX TAIL SUPPRESSION
To complement the theoretical analysis presented in Corollary 5 and Proposition 6 we measure the realized softmax tail mass for models trained with Sparsemax versus Sparsemax
We conduct an evaluation on 100 randomly selected AlpacaEval test examples. For each example, we record the average softmax probability assigned to the top-1, top-2, and top-3 tokens, as well as the cumulative probability mass assigned to the "Tail" (defined here as tokens falling outside the sparsemax support set).
Figure| visualizes these results. The comparison reveals sharp contrast in distribution behavior:
Sparsemax-only (Blue): Exhibits heavy tail, with substantial probability mass (RS 0.553) lcaking into non-support tokens during softmax inference.
Sparsemax+ (Red): Effectively concentrates mass on the lead candidate (1st Prob Ro 0.853) and suppresses the tail leakage to negligible level (* 0.069).
These empirical measurements validate that, in practice, Sparsemax+ substantially suppresses the softmax tail mass that is theoretically admissible in the worst-case construction, thereby enforcing the intended tail-suppressed plausible diversity.
D DETAILED PROOFS
Corollary If Definition I holds and Etail < Eheads then maxjes Pj S Etail < Ehead S minies Pis so each plausible sample has strictly higher probability than any tail sample.
22

---

Published as a conference paper at ICLR 2026
5 hobocblilee, fop3,
141
06
0.4
0.3
41
00
Nox
Figure 7: Comparison of top-3 softmax probabilities and cumulative tail mass between Sparsemax (blue) and Sparsemax+ (red). The "Tail" category represents the cumulative probability mass signed to tokens that would have been zeroed out by the sparsemax transformation (non-support tokens). Sparsemax+ significantly reduces this tail mass compared to standard Sparsemax.
Proof. From tail suppression, Ejgs Pj S. Stail, hence maxjes S Etail. From head preservation, minjes Pi 2 Chead- Combine with the condition étail < Ehead, we complete the proof.
Corollary 2 If all probability mass collapses onto the ground-truth token, i.e., Py 1 and Py 0 Vy # y. then p fails to qualify the TSPD (m(2 2), Chead, Etail).
Proof. For 2, S Top (p) contains y and some y with Py 0, violating minjes Pj Z Ehoad 0.
Lemma (Gradients vanish outside the sparsemax support (Martins & Astudillo 2016)). Let p sparsemax(2) and (2) (z) be its support. Define L(p. y) as supervised loss between the sparsemax probability p and the target y. If y (*)(2), then Vi & (2), 0C(x.y) = 0.
Proof. The gradient satisfies vzL(z,v) p Cy- For i $ STP(2) have p. and under the assumption E (*P(z) we have i + y. hence 0L(z, /0z 0.
Theorem 4 (Sparsemax expands pairwise gaps faster than softmax). Let E RRKK par softmax(2) =), and pd = sparsemax(z). For any indices i # j. let u = 2; 2j and we have
0
du (1; ) =1
Vi,j S
sparsemax
du (p:
pj <1
softmax
Proof. Inside the sparsemax support, we have P3 2j T(2) and p. -リ (z: T())) (zj 0 T(2)) = 2 Zj. thus du (p:" - 2))) = For softmax, using the Jacobian Vp = diag(p5)
of (p=t) and differentiating only the direction 2, 1, 2₁ (other logits fixed) yields (vi 189)
pt + p - -ppt ) which is strictly < for finite z.
Corollary 5 (Softmax remains TSPD-valid when sparsemax is one-hot). Let x E R with
arg max, 2y. and 6, := - 2z. Assume sparsemax is one-hot y, i.c., 8min == min / d3 2 2
(e.g., ? = 1), and the top-m head is bounded: Ō(k) == by - 2(K) I B Wk = 2,... m. Set
Am m + (K m)e Then for pa softmax(z) we have
B
(K mee
(vk 2....,
Σ rily k>m
Am
Am
Consequently, pd satisfies TSPD of order m with any thresholds < Ehead
S -B
_
(K (K-m)e" 2
Am
Etail < m7 Shead-
23

---

Published as a conference paper at ICLR 2026
Proof. For any j.
e²) pj Σx E
(2y
e s2
where Q == 1+
Σ
e 6n.
1+ Eiety e- (zy-
Then, v2 5 k & m, we have e (( E [e-B.11 according to the head bound 6(τ) S B; 4k > m, we
have e (A() S E according to the sparsemax one-hot margin 6(k) 2
To lower-bound pzf, we upper-bound C. by taking the largest possible contributions in each group:
m
C 1+
Σ
6(4)
Σ
e e-
600 I 1 (m 1 (K -mm e Am
Therefore, we have
Py C
of 2
Am P(k)
e
c
(k= 2,....... m).
C
Am
For h >> m. 8(n) 2 gives Σn>m
P(x)
(K-m))
Am
We complete the proof.
Remark 2 (Existence of a tight upper bound for the cumulated tail mass ΣxSm P(ke in Corollary
Under the assumption of Corollary 5 let Qmin
1 + (m 1)e B + (K -mmeee then we have
Ensm P(k)) < [K-m)e-
With the new upper bound, pd still satisfies TSPD of order m with any
Dmin
thresholds < Ehead S A,..
(K-m)e
S Etail S m Ehead-
Puin
Proof. For any the softmax probability is given by:
e=
pj
where Ω 1+
Σ
Ω
kyg
Based on the assumptions, the logit gaps satisfy:
Head (2 < k < m): d()) S B = e =BS 5e 6000 $ 1.
Tail (I: >> m): 5(m) 27
e B(2) Se
1. Lower bounds for the head candidates. To lower-bound D and pkky we must upper-bound the partition function Ω. By taking the largest possible contributions from every token (setting 5(,) for the head and (k) = 7 for the tail), we obtain Ami
Ω < 1 I+
e m + (K - mm) Am
k=
k>m
Therefore,
6(0
c B
of Pyy
of ((k)
(vk = 2, ......m.. m))
Am
Ω
Am
2. Upper bound for the tail mass. The cumulated tail mass is given by:
Σ
kksm 660
T(2) Σ
((k)
Ω
kmm
To find the maximum tail mass, we must maximize the numerator and minimize the denominator.
The numerator is maximized when tail gaps are minimal (61z) = 7). The denominator Ω 1+
Σm e 66 + Ekom e -600 minimized when the head contributions are minimal (6(z) = B).
Let Puin = =1 + (m 1)e-B B + (K - m)e-" Then:
P(M) S (K m)e
Σ k>m
of
2min
This completes the proof of Remark
24

---

Published as a conference paper at ICLR 2026
Proposition (Tightness of the Tail Bound). Under the assumptions of Corollary the upper bound on the softmax tail mass is right.
Proof. Consider the extremal configuration:
By O,
((2
(m) - B,
(mm+1)
2(K)
This configuration satisfies all assumptions. The softmax tail mass becomes:
(K - m))
max (K) = 1 (m 1)e- (K m
Let = e- (m B)- and K m. We write T Tmax (K) f()==
b-at" The
ob derivative f'(1) (b+ar)? > 0 shows that the tail mass is strictly increasing in K. Moreover,
at at
lim T max (K)
lim t-
1.
K000
Thus, the maximum softmax tail mass can grow monotonically with and approach 1, confirming the tightness of the bound.
E ADDITIONAL TECHNICAL ANALYSIS
E. A NEW TRAINING PARADIGM
Training with CE loss leads to distribution collapse: under gradient descent, the predictive distribu- tion p converges to the target y. This causes over-confidence and degenerate predictions at inference.
To address this issue, we discuss a new paradigm consisting of three steps:
1. Inflation during training. Given pa mox) we define an inflated distribution
p. f(p i 1,2, K, E/=1 f(1)))
where /0 1] >> R2 is strictly increasing and satisfies a ratio amplification property:
f(p)) f(p,
P
2. Loss applied on the inflated distribution. We train by minimizing tailored loss ((p. v). Ratio-amplifying inflation accelerates the collapse of p to one-hot.
3. Softmax inference. At test time, predictions made with the original p, which remains smooth and calibrated.
This paradigm improves optimization dynamics while preserving smooth probabilistic predictions.
Proposition 7. Given the predictive distribution and an infated distribution P = =b(p) where p: f(p)/ E, f(p,), let f satisfy the ratio amplification property. For any finite number of training steps T:
1. (Training: Faster Convergence) The inflated distribution p converges the one-hot label strictly faster than the original distribution p.
2. (Inference: Distribution Collapse Prevention) training stops at step T such that the inflated residual is Djtyř
O, then the original residual Ej = e satisfies e >> 6.
In particular, for any finite T, p remains strictly inside the simplex, preventing distribution collapse at inference.
25

---

Published as a conference paper at ICLR 2026
Proof. Consider target class and an incorrect class + y. Let Py 1 and By = 1 6. By the definition of the inflated distribution, the ratio of probabilities is:
Pu
f(Pw), Pj f(ps)
Pu Ni
Pj
Swap the numerator and the denominator:
Pj Pi Pivj
Pj
Pi
6 1 1- 1-6
5
Σههt Py
< E Eizッy
Py Py
Py
Let = Since g' = (I-z}²
Vr (0, (xx) strictly inereasing on (0, 1). The inequality 9( < (()) implies 8 < This proves that at any iteration t, p always leads the convergence toward the one-hot vertex.
Therefore, for any finite number of training steps, while p may appear to have collapsed to a one- hot vector, the underlying distribution p remains strictly inside the simplex, preserving a smooth probabilistic prediction for inference.
L
Theorem (Sparsemax as piecewise ratio-amplifying d-mapping of softmax). Let 2 E RK be logit vector, pdf softmax(z) E A K-1 with
i= 1,2..... K;
===1" ez
and pd sparsemax(2)) E AK with
00 max - [[2),}} Σ p =1
Then define
50 = = f(pi)
where f [0,1] + Ryo is the piecewise function f(x) max,log 0,0), 40 R.
Proof. According to p = Σn F this into the definition of sparsemax,
e we have z; log + with log E, es. Substituting
p. = mogg + +z)
Letting -(2)=- - C, we obtain
= max 0,0).
Since E.r,p 50 1, normalizing yields
maxflog p; 0,0}
p. Σ, max{logp, 0,03
We now analyze the following two cases.
Case I (support set (5) {i log pff
0}.. For i E S, {p:f)
= log pff 0 > 0. Recall that
0 = T(z) C, hence log pff 0 - log vif
+ C - T(2) = - T(z) = Thus for i,j E S,
50
(p)):
log off 0 p.
T (p))
log 0 p,
On the other hand, using the softmax identity
exp(2: - 24), and noting that for i,j E S,
zi - 2; = (z; - T) - (zz - 7) = D. - sp we obtain 2 = exp((?? - -,)))
26

---

Published as a conference paper at ICLR 2026
Since Pi p." Pj E (0,1] and p; Pj whenever p. > Pj, the function g(x) log. is strictly increasing on (0, 1], which implies f explp, ))
Py
Therefore,
p Tpp),
Thus & strietly amplifies the relative ratios within the support.
Case 2 (outside the support (z)) For j & (9P(2), we have log p5¹ S 0 and hence f(25)) 0. Therefore,
0
(())) = Ex f(p5t) = 0.
By contrast, p. since p softmax(z) has full support. Thus sparsemax(z) coincides with (p)), where Φ is generated by the piecewise ratio-amplifying function f.
Overall, sparsemax(z) is piecewise ratio-amplifying inflation of softmax(z). Training on $(p) drives the inflated distribution to collapse one-hot its support, while inference with the original softmax preserves strictly positive mass on all coordinates. This prevents the predictive distribu- tion from degenerating into an exact one-hot vector inference.
Having established sparsemax as concrete instance of ratio-amplifying inflation, it is natural to ask whether other mappings f might be equally effective, or perhaps even more suitable in specific contexts. To answer this, we next examine the general collapse condition in the binary case.
E.2 GENERAL COLLAPSE CONDITION IN THE BINARY CASE
Consider binary classification with p = (p, - p) and label y = 1. The inflated distribution is
f())
P1 f(p) f(1 pp' P2 P1.
Define the ratio
R(p) f(1 (11p)) p)
f(p)
Then
P1 + R(p)
For precision parameter e 0, we say collapse occurs if
$1 2 1 - - R(p) S 6
1. Power inflation. For f(x) xa a > 1.
R(p) 11ppp
Collapse condition:
p
1/a
1 +
[-c
c
2. Exponential inflation.
For f(x) E² > 0,
R(p) (1-
Collapse condition:
1 1-e
2 +
log
27
27

---

Published as a conference paper at ICLR 2026
3. Logarithmic inflation.
For f(x) log(1 + +) - log 6 with > 0, so that f(z) 0 [0. 1].
R(p) log(1 + 8) - log
log(p 6) log
Collapse condition:
log(1 p 6) log 6
v
log(p + 6) - log
with log(p + 6) - log 0 (i.e., > 0).
E.3 GRADIENT DYNAMICS UNDER RATIO AMPLIFICATION
The ratio-amplifying property of ø-mappings only accelerates the collapse of p, but also shapes the gradient dynamics during training. For a strictly proper loss e, the gradient w.r.t logits assumed be
[zl(z;)) y) = Cy, g(z),
where g(-) denotes a probability distribution obtained from the logits x and ey is a one-hot vector with the y-th entry equals 1.
When training on the inflated distribution p = <b()), the chain rule gives
Op
(p v),
Öp
where 0e is the Jacobian of the inflation operator. Up
Effect of ratio amplification. so that
Suppose f [0, - R4 is strictly increasing and ratio-amplifyingg
Py
Py
Vj A y-
Pj
This guarantees that the relative gap between the correct and incorrect probabilities grows under 4. Hence, even if the exact magnitude of each gradient entry depends on the Jacobian structure, the ratio
V
is enlarged compared to the original probability space. In other words, the margin 2, receives stronger effective gradient pressure grow. Intuitively, because Py Py and Pj Pj for j y. the gradient signal on the correct logit zy is reinforced, while the signals on the incorrect logits 2j are diminished. This rescaling accelerates the suppression of false classes and boosts the dominance of the true class. Although the absolute gradient values are determined by both ø and the Jacobian Op the effective separation between correct and incorrect classes is consistently larger under ratio- amplifying mappings.
Op
Summary. Any é-mapping with ratio amplification reshapes the optimization dynamics by pre- conditioning the gradient flow:
The relative strength of gradients is tilted further in favor of the true class.
Incorrect classes are suppressed earlier, as their probabilities are diminished more aggres- sively.
Consequently, the system reaches effective one-hot collapse earlier than when training directly p. Crucially, since inference is carried out with the original distribution p. the final predictions remain smooth and non-degenerate, preserving diversity while benefiting from sharper supervision during training.
28


Reviews:
[{"review_id": "DHATR7ckQX", "reviewer": "Reviewer_GaaU", "summary": "The authors propose a new normalization method from logits to the probability distribution. By analyzing the behavior of the new normalization method, the authors point out that the method is only good for training but not for inference. Thus, the authors give the overall framework that the model is trained with the new method (i.e., Sparsemax+) and is inferred with softmax. The experimental results show that the framework increases the pass@N a lot when N is large.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. To preserve diversity, the authors introduce a new normalization layer called Sparsemax+.\n\n2. After analyzing the properties of Sparsemax and softmax, the authors choose to use different normalization layers for training and inference.\n\n3. The performance increases a lot when the number of rollout samples is large.", "weaknesses": "1. It is unclear what the mismatch of using Sparsemax+ and softmax will cause.  The model performance is not as good as GEM for pass@N when N is small (e.g., N = 2). Is it because of the mismatch?\n\n2. Sparsemax+ does not seem to be tunable. In softmax, there is a temperature parameter that can control the final results. Also, in GEM, there is a hyperparameter that can balance the KL term and the objective.\n\n3. For $i,j \\notin S^{sp}$, $\\frac{\\partial }{\\partial u}(p_i - p_j) < 1$, what is different between this result and softmax?", "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": "gUDNimkole", "reviewer": "Reviewer_2nRT", "summary": "This paper proposes TS2 (Train with Sparsemax+, Test with Softmax), a novel fine-tuning framework for large language models that aims to mitigate distributional collapse and preserve useful diversity. The key idea is to decouple the probability mappings used in training and inference. During training, the model uses Sparsemax+, an enhanced version of sparsemax that includes a tail-suppressing penalty to eliminate residual probability mass on implausible tokens, thereby enforcing “Tail-Suppressed Plausible Diversity (TSPD).” At inference, the model reverts to the softmax mapping to restore smooth, calibrated probabilities, allowing exploration among plausible alternatives. The authors provide theoretical analysis on gradient masking and sparsity behavior, proving that sparsemax expands pairwise logit gaps faster than softmax and yields a more decisive support set. Empirically, TS2 is benchmarked against CE, NEFT, and GEM on both chat and code-generation tasks, showing consistent improvements in win rate, pass@k, and diversity metrics. Ablation studies further confirm that the combination of sparsemax training, softmax inference, and suppression penalty is crucial to the method’s success.", "questions": "see weakness", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Empirically, the results are comprehensive and convincing, spanning multiple domains (chat, code, creative writing, and reasoning benchmarks) and multiple base models (Llama-3.1-8B, Qwen-2-7B). \n\nThe method outperforms strong baselines such as GEM not only in accuracy but also in diversity metrics, demonstrating improvement both in quality and variety. The clarity of presentation and reproducibility documentation is also commendable.", "weaknesses": "1. The main limitation is the upper bound of the cumulated tail mass of softmax outside the top-$m$,\n$$\n\\sum_{k > m} p^{\\mathrm{sf}}_{(k)} \\le \\frac{(K - m)e^{-\\gamma}}{A_m}.\n$$\nIt is said that this upper bound is strictly increasing in $K$ (for fixed $m, \\gamma$) and approaches $1$ as $K \\to \\infty$. Finally, the authors state that ''with large vocabularies, the admissible tail under softmax at inference becomes non-negligible.''\n\nHowever, the statement is not support by the theorem. The upper bounds goes to 1 does not tell anything about the probability. To support the claim, the authors should prove that the upper bound is tight, or provide a lower bound on the probability.\n\n2. The proposed loss is simple, compressing the tailed probability to 1 single probability. To support that the proposed one is better, the authors should provide a result like Corollary 5 with a tighter bound (if possible). Otherwise, it is hard to tell whether the proposed loss resolved the raised theoretical issue.\n\nWriting and presentation issues:\n\n1. The definition of sparsemax (lines 154–157) appears to be ill-defined. Specifically, the threshold function is defined in terms of the support set, yet the support set itself is defined based on the threshold function, resulting in an apparent circular definition. I suspect that sparsemax has a formal, non-circular definition in its original paper, and that the circular definition presents only two derived properties rather than the definition itself.\n\n2. Lemma 3, Lsp(z,y) is not defined before.Leaving readers to guess whether it is refered to “we train with sparsemax and optimize a modified Fenchel–Young loss tailored to this mapping” mentioned in the introduction.", "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": "PMjEWPD5F5", "reviewer": "Reviewer_uXcp", "summary": "This paper addresses the issue of output diversity loss in supervised fine-tuning (SFT) of large language models. Standard SFT with cross-entropy (CE) tends to over-fit the model's output distribution to one-hot targets, suppressing alternative correct continuations and thus reducing generative diversity. The authors propose TS$^2$ (Training with Sparsemax+, Testing with Softmax) as a drop-in replacement for the usual fine-tuning procedure. During training, they replace the softmax+CE loss with a modified Sparsemax+ loss, which encourages a sparse probability distribution over next-token outputs while adding a suppression term to push down probabilities of tokens outside the support of the target. At inference, the model still uses a normal softmax for decoding. This strategy preserves plausible near-tie candidates in the output distribution (instead of forcing a single high-probability token as CE does) while still penalizing irrelevant tokens. The paper's contributions include: (1) a new Sparsemax+ transformation and loss function that fixes a known issue of vanilla sparsemax (which left too much residual probability on tail classes), and (2) demonstrating that TS$^2$ yields consistent improvements in both accuracy and output diversity over the baseline CE fine-tuning. The results suggest that TS$^2$ is an effective, easy-to-implement approach to obtain LLMs that are more accurate while also more creative (producing varied outputs).", "questions": "The questions are the same as those listed in the Weaknesses section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Novelty:** The paper tackles the problem of maintaining output diversity in supervised fine-tuning. It introduces a creative solution by leveraging sparsemax (a non-standard probability transform) in a new way. The proposed Sparsemax+ improves upon sparsemax by addressing its known gradient masking issue for out of support set classes. This combination (train with sparsemax+, test with softmax) is a novel idea that, to my knowledge, has not been explored in prior LLM fine-tuning literature.\n\n* **Quality:** The authors fine-tune two open models (Llama 3.1 8B and Qwen 2.5 7B) on multiple benchmarks covering different domains (chat/dialog, code generation, and open-domain tasks). The results consistently show improvements over the baseline cross-entropy fine-tuning. Importantly, both accuracy (task performance) and diversity metrics improve, which demonstrates the method's effectiveness without trade-offs. The paper also provides intuition (and some theoretical discussion) for why TS$^2$ works, grounding the empirical findings.\n\n* **Clarity:** The paper is clearly written with a logical flow. The authors motivate the problem (lack of alternative continuations under CE) convincingly, and then explain their solution step-by-step. Key terms like sparsemax and the new loss are defined, and the intuition is easy to grasp.\n\n* **Significance:** The findings have practical significance for the field. Many applications of LLMs (creative writing, coding assistant, dialogue systems) benefit from diverse outputs, and TS$^2$ provides a simple drop-in technique to achieve that without sacrificing accuracy. This is significant because typically one might expect a trade-off between diversity and correctness. The method is also easy to implement for practitioners (just changing the loss function during training), which means it could be adopted widely if these results hold.", "weaknesses": "* **Missing Baseline Comparisons:** The experimental section might be missing some baselines for a fair assessment. In particular, it's important to know how TS$^2$ compares against simpler fixes to the one-hot problem, such as label smoothing during CE training or using entmax loss [1] (which is a family that includes sparsemax and softmax). If those were not evaluated, the authors should clarify why. Including such baselines would strengthen the paper by demonstrating that Sparsemax+ indeed outperforms these alternatives in balancing accuracy and diversity. Without them, it's harder to quantify the gain from the new method versus existing tricks.\n\n* **Impact on Downstream Use Cases:** The paper shows benchmark improvements, but have the authors examined qualitative examples or specific scenarios illustrating the benefit? For instance in a chat setting, does TS$^2$ produce more varied and interesting responses to open-ended questions (compared to CE which might produce more templated answers)? A few concrete examples in the rebuttal version would be great to demonstrate the practical difference in model behavior.\n\n* **Hyperparameter Sensitivity:** How sensitive is the Sparsemax+ training to its hyperparameter - the strength of the out-of-support suppression term $\\alpha$? The paper mentions that $\\alpha$ is empirically determined for each model architecture, but additional analysis on its effect (example varying $\\alpha$ to show robustness or convergence behavior) would help readers better understand how critical this choice is and how generalizable the method is.\n\n[1] \"Sparse Sequence-to-Sequence Models\", Ben Peters and Vlad Niculae and André F. T. Martins, https://arxiv.org/abs/1905.05702.", "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": "uQt6s8MHN8", "reviewer": "Reviewer_jH36", "summary": "The paper proposes an alternative to the cross-entropy loss for language model SFT. The main motivation for this is that CE loss pushes all the probability mass under the golden token, killing the diversity in the pre-trained model. The proposed method uses sparse activation (sparsemax+) to train the model on only plausible tokens. This makes the model output diverse when used with softmax while keeping it accurate. The paper shows better diversity and accuracy compared to existing methods like GEM, CE+weighting (CE regularization), and CE.", "questions": "- For baselines, did you consider label smoothing and α-Entmax since they directly address the overconfidence problem? \n- I am curious to know the effects of the TS2 loss on factual accuracy when the diversity is high? As in, when you sample multiple times, do the facts stay consistent? \n- Can authors comment on the effectiveness of this method when training on datasets with long reasoning traces? Does the model still give coherent reasoning steps when diversity is maintained?\n- Similar to above, how does it perform on a multi-turn dataset (for example, on tau-bench)?\n- Could you share the decoding hyperparameters? Since this could influence the measured diversity", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The results show strong performance and diversity on a variety of benchmarks and are measured on a variety of metrics for diversity.\n- Results show that better diversity also leads to sample efficiency when doing BON\n- Ablation studies show holding out one of the component leads to drop in performance, indicating the need for all components (sparsemax+, softmax decoding) for the best performance\n- The method is compared against existing SOTA (GEM)\n- The motivation is clearly laid out, and the research questions are made clear", "weaknesses": "- The paper discusses the overconfidence due to the SFT issue, but doesn't discuss existing work on it\n- Other relevant methods to prevent diversity collapse are mentioned, but don't have a comparison against (Label smoothing, unlikelihood)", "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)