{"id":"80a6d34f-da48-4d96-81cf-529bdae5bcbe","arxiv_id":"2506.15021","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":4.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"SFT-GO retrains LLMs by focusing on the worst-performing group of important or unimportant tokens, yielding modest average benchmark improvements over standard supervised fine-tuning.","lead":"This paper proposes SFT-GO, a fine-tuning method that gives more training weight to groups of tokens judged important, using a worst-group loss borrowed from robust optimization. It reports small average gains over standard fine-tuning on several LLM benchmarks, though the gains are partly confounded by per-method hyperparameter tuning.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Alpaca benchmark gains may be an artifact of per-method test-set tuning: baseline and SFT-GO use different learning rates (6e-6 vs 1e-6/5e-7/6e-7 for 3B), and eta/lambda are selected on the same evaluation benchmarks; matched-configuration reruns are needed before the central claim is established.","rationale":"The reader's weakest assumption identifies the most load-bearing concern: the empirical comparison is confounded by per-method hyperparameter tuning on the same test benchmarks, including learning-rate differences of up to an order of magnitude. I agree with this diagnosis and with the conditional verdict, so I recommend no verdict change. I considered whether the proof of Proposition 1 is a more serious internal flaw: the appendix proof is terse and phrase 'we need L_CE(hat) < L_CE(avg)' is imprecise, but under exact global minimization the contradiction argument does go through (if L_worst(hat) > L_worst(avg) and L_CE(hat) >= L_CE(avg), then L_GO(hat) > L_GO(avg), contradicting the definition of hat). Proposition 2 is a standard convex SGD bound, and the paper explicitly acknowledges the convexity gap in Appendix C. Thus both theoretical issues are real limitations but secondary to the experimental protocol. The paper has compensating strengths: five-seed averages, three grouping strategies, honest compute and limitation statements, and concrete ablations. However, because the central empirical claim depends on isolating the effect of the group objective, the decisive check is a matched learning-rate comparison with held-out validation-based hyperparameter selection.","tokens_in":22228,"tokens_out":6604,"duration_ms":62222,"concrete_test":"Retrain Llama-3.2-3B on Alpaca with a factorial design: baseline-SFT and SFT-GO (TF-IDF and LLMLingua-2) each run at learning rates {6e-6, 1e-6} with all other settings identical. Select eta and lambda on a held-out 10% validation split of Alpaca, and evaluate on the seven benchmarks only after this selection. If baseline-SFT at LR=1e-6 matches or exceeds the reported SFT-GO averages (44.52 and 44.89 vs. 43.96 for baseline), the Table 2 gains are learning-rate artifacts; if SFT-GO still leads under matched learning rates and validation-based hyperparameter selection, the central claim survives.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central claim is that the SFT-GO objective consistently improves supervised fine-tuning. For that claim to hold, the observed benchmark gains must be attributable to group optimization rather than to the evaluation protocol. In Appendix F.2, on Alpaca with Llama-3.2-3B, baseline-SFT uses a learning rate of 6e-6 while TF-IDF, Rho-1, and LLMLingua-2 use 1e-6, 5e-7, and 6e-7 respectively; for Llama-3.1-8B, baseline uses 1e-7 while all SFT-GO methods use 6e-7. Method and learning rate are therefore confounded. The text states that 'we search the optimal learning rate, eta, and lambda' and reports results on the same seven benchmarks used for that search; no held-out validation split is described. SFT-GO also has additional hyperparameters (eta and lambda) beyond the baseline's learning rate, so selecting on the test benchmarks gives it more degrees of freedom that can inflate average performance even if the objective itself is neutral. For LIMA, the baseline uses the original paper's recommended hyperparameters while SFT-GO's eta and lambda are tuned, again on the same benchmarks. The ablations in Section 5.3 vary eta and lambda and report test-benchmark scores, providing further evidence of selection on the evaluation set. The theoretical results do not rescue the comparison: Proposition 1 assumes exact global optima and Proposition 2 assumes convexity, which the paper itself concedes does not hold for LLMs (Appendix C). Consequently, the empirical comparison does not currently isolate the effect of the group objective.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes SFT-GO, a supervised fine-tuning objective that combines the standard cross-entropy loss with a worst-group loss over token-importance groups: L_GO = (1-λ)L_CE + λ L_worst, where L_worst = max(L_CE(G1), L_CE(G0)). Three grouping strategies are considered: TF-IDF, LLMLingua-2, and Rho-1, with Rho-1 reformulated as a special case. The authors provide two theoretical results: Proposition 1 claims that the worst-group loss of the group-optimized solution is no worse than that of standard SFT, and Proposition 2 claims an O(1/sqrt(T)) SGD convergence rate. Empirically, the paper reports average improvements over baseline SFT across seven benchmarks using LIMA and Alpaca with Llama-3.2-3B and Llama-3.1-8B, over five seeds.","tokens_in":22593,"tokens_out":5349,"duration_ms":47957,"significance":"If the reported gains survive a controlled comparison, SFT-GO would be a simple and useful addition to the instruction-tuning toolbox: it is flexible in how token importance is defined, it subsumes Rho-1, and it is evaluated on two base models, two datasets, and seven benchmarks with multiple seeds and transparent compute reporting. The paper also usefully separates token groups rather than reweighting every token independently. However, the central empirical claim is not currently established because the evaluation protocol confounds the objective with learning-rate choices and performs hyperparameter selection on the same benchmarks used for the final numbers. The theoretical results are standard and are explicitly limited by the convexity assumption conceded in Appendix C.","major_comments":[{"comment":"The Alpaca experiments confound method with learning rate. For Llama-3.2-3B, Baseline-SFT uses a learning rate of 6e-6 while TF-IDF, Rho-1, and LLMLingua-2 use 1e-6, 5e-7, and 6e-7, respectively; for Llama-3.1-8B, Baseline-SFT uses 1e-7 while all SFT-GO variants use 6e-7. Since the learning rate differs by an order of magnitude between the baseline and each SFT-GO method, the average gains in Table 2 cannot be attributed to the SFT-GO objective rather than to a better-tuned learning rate. The authors should rerun Baseline-SFT at the learning rates used by the SFT-GO variants, and rerun SFT-GO at the baseline learning rate, to show that the objective itself, not the learning-rate schedule, drives the improvements.","section":"Appendix F.2, Table 2"},{"comment":"The paper states that the optimal learning rate, eta, and lambda are searched, but it does not describe a held-out validation split; the final numbers are reported on the same seven benchmarks used for that search. For LIMA, the baseline uses the hyperparameters recommended by the original LIMA paper while SFT-GO's eta and lambda are tuned on the evaluation benchmarks. This gives SFT-GO additional degrees of freedom and risks selection on the test set, which can inflate average performance even if the objective is neutral. The authors should either select hyperparameters on a separate validation set (e.g., a subset of the training data or a small held-out benchmark) and then report on the seven benchmarks, or show that the qualitative conclusion is unchanged across a wide grid of eta and lambda values for both baseline and SFT-GO.","section":"Section 4, Section 5, Appendix F"},{"comment":"The claim that SFT-GO models 'consistently outperform' the baseline is not supported at the individual benchmark level on Alpaca. For Llama-3.2-3B, all three SFT-GO methods score below Baseline-SFT on ARC-C (45.05-45.61 vs. 46.60) and on HellaSwag (55.54-56.10 vs. 56.54), with the average gains driven mainly by TruthfulQA and IFEval. The authors should qualify the claim to refer to average performance, or provide per-benchmark significance tests (e.g., paired tests across seeds) to show which individual differences are reliable.","section":"Table 2, Section 5.1"},{"comment":"The theoretical convergence result assumes L_GO is convex with Lipschitz subgradients and that the parameter space is convex and bounded, and Proposition 1 assumes the solutions are exact global optima of their respective objectives. Appendix C explicitly acknowledges that the convexity assumption does not hold for LLMs. As a result, the O(1/sqrt(T)) rate and the worst-group-loss guarantee do not directly apply to the actual nonconvex LLM fine-tuning setting, and the theory should not be presented as demonstrating efficiency for LLM training without a clear statement of this gap. The empirical comparison, not the theory, must carry the central claim.","section":"Appendix B, Appendix C, Proposition 2"}],"minor_comments":[{"comment":"There is a typo in the datasets paragraph: 'STF-GO' should be 'SFT-GO'.","section":"Section 4"},{"comment":"In the proof of Proposition 2, the first norm term on the right-hand side of Eq. (17) appears to be missing the square: it should read ||θ(t) - θ*||^2 for consistency with the surrounding derivation and with the telescoping sum that follows.","section":"Appendix B, Eq. (17)"},{"comment":"The HellaSwag description says it tests 'broad knowledge and reasoning in 57 diverse subjects,' which is the description of MMLU; the HellaSwag entry should be corrected to describe commonsense physical reasoning about everyday situations.","section":"Appendix E"},{"comment":"The x-axis is labeled 'Percentile' while the text refers to a 'compression rate' and the caption says that at the 90th percentile the bottom 90% of tokens are unimportant; the axis label and caption should be reconciled to avoid confusion about whether larger values mean more or fewer important tokens.","section":"Figure 3"},{"comment":"The appendix uses 'LLM Lingua' inconsistently and describes LLMLingua-2 as assigning importance based on low predictive confidence, which is not clearly the same as the keep/drop probability described in Section 3.3; the authors should make the description consistent with the method they actually use.","section":"Appendix H"},{"comment":"References [10] and [23] appear to refer to the same Rho-1 work in different venues; the authors should cite the published version once and avoid duplicate entries.","section":"References"}],"recommendation":"major_revision","confidential_remarks":"The core idea is plausible and the paper is clearly written, but the empirical protocol currently prevents clean attribution of the gains to SFT-GO. The fix is feasible within the manuscript's scope: matched learning-rate baselines and a held-out validation split for hyperparameter selection, plus a per-benchmark significance analysis. The theory is standard and should be repositioned as motivational rather than as a guarantee for the LLM setting. I recommend major revision rather than rejection because the central claim is defensible if the experimental comparison is made fair."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The core idea is worth a look: apply Group DRO's worst-group loss to token-level importance groups within each training example. The paper is honest and readable, and the three grouping strategies (TF-IDF, LLMLingua-2, Rho-1) make the framework concrete. What's genuinely new is token-level grouping with importance scores from TF-IDF and LLMLingua-2; Rho-1 is related but calling it a special case is loose, since with λ=1 your objective is the max of the two group losses, not CE on the selected tokens alone.\n\nThe empirical claim is not yet clean. You search learning rate, η, and λ on the same seven benchmarks you report, with no held-out validation split. On Alpaca, the baseline and SFT-GO run at different learning rates (6e-6 vs 1e-6/5e-7/6e-7 for the 3B; 1e-7 vs 6e-7 for the 8B), so method and training configuration are confounded. The average gains are about 1-2 points, and several individual benchmarks regress (e.g., ARC-C on Alpaca). That doesn't kill the idea, but it means the central claim—that the group objective itself improves SFT—is not yet established.\n\nThe theory is modest but not wrong. Proposition 1 is a correct and actually rather trivial consequence of minimizing a sum: the worst-group loss of the joint minimizer can't exceed that of the CE minimizer. Proposition 2 is the standard O(1/sqrt(T)) SGD bound under convexity, which the authors concede doesn't hold for LLMs (Appendix C). So the theory is a sanity check, not a selling point.\n\nThis is mainly for people working on data weighting and post-training. The paper would benefit from a matched-configuration rerun: same LR range for all methods, η/λ chosen on a validation split, and results reported with the selected hyperparameters. If the gains survive that, it's a useful practical result. As is, I'd send it to peer review—the idea is plausible and the flaws are fixable—but the reviewers should push for the rerun before acceptance.","headline":"Token-level Group DRO is a neat idea, but the empirical comparison is confounded by tuning on the test benchmarks and mismatched learning rates, so the central claim needs a matched-configuration rerun.","tokens_in":23132,"tokens_out":3901,"would_cite":false,"duration_ms":34392,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Splitting tokens by importance and adding the worse group's loss to the fine-tuning objective improves supervised fine-tuning of large language models.","keywords":["supervised fine-tuning","token importance","group optimization","worst-group loss","distributionally robust optimization","instruction tuning","large language models"],"falsifier":"Run baseline SFT and SFT-GO on the same dataset with identical learning rate, optimizer, schedule, and epoch count, and choose $\\eta$ and $\\lambda$ on a validation split that is disjoint from the final evaluation benchmarks; if the average benchmark scores no longer differ, the claim that the worst-group objective drives the improvement is refuted. Alternatively, track $L_{\\mathrm{worst}}$ on held-out tokens during training, since Proposition 1 predicts it should stay at or below the standard-SFT level and a sustained violation would contradict the theory.","tokens_in":22036,"feed_emoji":"🎯","tokens_out":10439,"duration_ms":89280,"temperature":0.7,"pith_summary":"Supervised fine-tuning of large language models normally averages cross-entropy loss over every token, so easy function words dominate the gradient while rare, semantically rich tokens remain under-optimized. The paper proposes SFT-GO, which partitions each sequence's tokens into important and unimportant groups and trains on $(1-\\lambda)L_{\\mathrm{CE}} + \\lambda L_{\\mathrm{worst}}$, where $L_{\\mathrm{worst}}$ is the larger of the two group losses. This forces the model to keep improving the token group it currently handles worst. The paper proves two supporting results: that the worst-group loss of the group-optimized solution is no larger than that of standard training, and that mini-batch SGD still converges at $O(1/\\sqrt{T})$ under convexity. It also reports that SFT-GO beats standard SFT on average across seven benchmarks with two base models and two instruction datasets, with the key insight that token importance can itself be used as a training signal rather than only for data filtering or prompt compression.","feed_headline":"Token grouping lifts LLM fine-tuning across seven benchmarks","feed_subtitle":"A worst-group loss term makes instruction-tuned models focus on important tokens and beats standard SFT on both datasets.","key_machinery":"The load-bearing object is the worst-group loss, $L_{\\mathrm{worst}}(\\theta) = \\max(L_{\\mathrm{CE}}(G_1;\\theta), L_{\\mathrm{CE}}(G_0;\\theta))$, added to the standard cross-entropy loss with weight $\\lambda$. This term makes the gradient depend on whichever token group currently has the higher loss, so training effort is repeatedly redirected to the group the model is learning most slowly. The other half of the machinery is the grouping function $g$, which fixes the partition; the paper tests a statistics-based version (TF-IDF), a semantics-based version (LLMLingua-2), and a loss-based version (Rho-1) to show that the objective works across different definitions of token importance. The annealing schedule for $\\lambda$, which starts high and decays, lets the method emphasize the worst group early and then let standard language modeling take over.","core_discovery":"The central claim is that replacing the uniform token average in supervised fine-tuning with an objective that explicitly optimizes the worse of two token groups improves the resulting model. Given a grouping function $g$ that labels each token as important or unimportant, SFT-GO minimizes $L_{\\mathrm{GO}} = (1-\\lambda)L_{\\mathrm{CE}} + \\lambda L_{\\mathrm{worst}}$ with $L_{\\mathrm{worst}} = \\max(L_{\\mathrm{CE}}(G_1;\\theta), L_{\\mathrm{CE}}(G_0;\\theta))$, so gradients are redirected toward whichever group is currently harder. The paper demonstrates the framework with three grouping functions: TF-IDF statistics, LLMLingua-2's semantics-based keep/drop probabilities, and Rho-1's excess-loss selection. It shows that Rho-1 is a special case when $\\lambda=1$ and grouping is loss-based. Empirically, all three variants outperform baseline SFT on average across eight reported columns spanning seven benchmarks for Llama-3.2-3B and Llama-3.1-8B on both LIMA and Alpaca, with the largest improvements on general-reasoning benchmarks. The paper additionally claims a bound of the form $L_{\\mathrm{worst}}(\\hat{\\theta}) \\le L_{\\mathrm{worst}}(\\theta_{\\mathrm{avg}})$ and an $O(1/\\sqrt{T})$ convergence rate, which tell a fair reader that the worst-group term is intended to balance token groups without sacrificing optimization efficiency.","pith_inferences":["An implication the authors leave implicit is that the same worst-group trick could be applied to groups defined by attributes other than semantic importance, such as token difficulty, domain, or safety-relevance, turning the objective into a general way to optimize the weakest slice of the training distribution.","Because the ablation shows performance holds over a wide range of important-token ratios, an untested prediction is that the gain should be largest precisely when the easy group dominates the token count, since that is when the uniform average most under-weights hard tokens.","A sharper test of the mechanism than final benchmark averages would be to monitor $L_{\\mathrm{worst}}$ on held-out tokens during training; Proposition 1 predicts it should stay at or below the standard-SFT trajectory, which is a direct, checkable consequence not reported in the paper."],"forward_implications":["Any method that can score token importance, from lexical statistics to a pretrained compressor to a reference-model loss, can be plugged into the same objective, so SFT-GO turns token importance into a general training signal.","Existing token-selection and token-reweighting fine-tuning methods are subsumed as special cases, giving one unified objective with which to compare and combine them.","Because the worst-group term tracks the harder group, it should reduce the under-optimization of rare, semantically rich tokens that standard average cross-entropy tends to sacrifice to frequent function words.","Empirically, the improvement appears on both LIMA and Alpaca and for two model sizes, so the benefit is not tied to a single data scale or model family.","The $O(1/\\sqrt{T})$ bound indicates that, in the convex setting, the group term does not change the optimization rate compared with standard SGD, so the added balancing behavior costs no theoretical convergence speed."],"supporting_citations":[{"why":"Supplies the worst-group loss and the distributionally robust optimization analysis that SFT-GO adapts from whole inputs to token groups.","marker":"[8]"},{"why":"Provides the LLMLingua-2 token-importance model used as the semantics-based grouping function.","marker":"[9]"},{"why":"Provides the Rho-1 loss-based token selection method that the paper reformulates as a special case of SFT-GO.","marker":"[10]"},{"why":"Supplies the LIMA instruction dataset and the training hyperparameters the paper follows for that setup.","marker":"[5]"},{"why":"Supplies the Alpaca instruction dataset used for the second set of fine-tuning experiments.","marker":"[11]"},{"why":"Provides the robust stochastic approximation bound on which the $O(1/\\sqrt{T})$ convergence proof is built.","marker":"[40]"}],"fun_headline_variants":["Worst-group loss boosts LLM fine-tuning across 7 benchmarks","Token grouping by importance improves LLM fine-tuning","Focus on hard token groups lifts LLM fine-tuning","Optimize worst token group to fine-tune LLMs better","Token-aware fine-tuning beats uniform SFT on 7 benchmarks"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The empirical comparison assumes that the reported gains come from the group-optimization objective rather than from per-method hyperparameter tuning: in the Alpaca setup the learning rates differ by an order of magnitude between baseline and SFT-GO, and $\\eta$ and $\\lambda$ were selected on the same benchmarks used for the final evaluation; the convergence theory additionally assumes convexity and exact optimality, which do not hold for large language models.","fun_headline_variants_meta":{"raw":{"variants":["Worst-group loss boosts LLM fine-tuning across 7 benchmarks","Token grouping by importance improves LLM fine-tuning","Focus on hard token groups lifts LLM fine-tuning","Optimize worst token group to fine-tune LLMs better","Token-aware fine-tuning beats uniform SFT on 7 benchmarks"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000705,"raw_usage":{"total_tokens":3247,"prompt_tokens":1084,"completion_tokens":2163,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":700,"completion_tokens_details":{"reasoning_tokens":2091}},"tokens_in":700,"tokens_out":2163,"duration_ms":16984,"temperature":1.0,"reasoning_tokens":2091,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T00:09:32.537436+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run baseline SFT and SFT-GO on the same dataset with identical learning rate, optimizer, schedule, and epoch count, and choose $\\eta$ and $\\lambda$ on a validation split that is disjoint from the final evaluation benchmarks; if the average benchmark scores no longer differ, the claim that the worst-group objective drives the improvement is refuted. Alternatively, track $L_{\\mathrm{worst}}$ on held-out tokens during training, since Proposition 1 predicts it should stay at or below the standard-SFT level and a sustained violation would contradict the theory.","supporting_citations":[{"cited_title":"Vicky Zhao, Lili Qiu, and Dongmei Zhang","cited_arxiv_id":null,"evidence_quote":"Provides the LLMLingua-2 token-importance model used as the semantics-based grouping function."},{"cited_title":"Nemirovski, A","cited_arxiv_id":null,"evidence_quote":"Provides the robust stochastic approximation bound on which the $O(1/\\sqrt{T})$ convergence proof is built."}],"review_version":1}