{"id":"f4c47b18-231a-44c2-bad9-4b40d458198f","arxiv_id":"2411.14055","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"DRPruning uses distributionally robust optimization with scaling-law-predicted reference losses and adaptive data ratios to improve domain-balanced performance recovery in LLM pruning and continued pretraining.","lead":"The paper proposes DRPruning, a training method that rebalances data across domains during large language model pruning and continued pretraining, using distributionally robust optimization with automatically adjusted reference losses and data ratios. The authors report lower perplexity and better downstream task scores than a reproduction of Sheared Llama in most settings, though the 2.7B model still trails the official Sheared Llama on downstream tasks.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The scaling-law reference-loss extrapolation (Eq. 4) is untested under the feedback loop of dynamic data scheduling; if biased, DRO weights are miscalibrated and claimed gains vanish.","rationale":"The reader's weakest assumption correctly identifies the scaling-law extrapolation as the critical fragile premise, and I agree that miscalibrated reference losses would make the DRO weights unreliable. The paper's own admission that official Sheared Llama 2.7B outperforms DRPruning on the domain benchmark (Table 5) provides empirical support for this concern, though that comparison is confounded by data differences. The controlled comparison against ReSheared is the strongest evidence for the method, but it does not isolate the contribution of the scaling-law prediction; an ablation with an oracle reference loss would directly test whether the prediction is the source of the improvement or a source of error. My read does not change the reader's verdict of CONDITIONAL: the method is well-engineered and promising, but the central mechanism—the predicted reference loss—should be validated against actual final losses before accepting the broader claims in the abstract, which currently state that DRPruning surpasses similarly sized models without noting the 2.7B domain-benchmark shortfall.","tokens_in":28001,"tokens_out":11015,"duration_ms":88414,"concrete_test":"For the 1.3B and 2.7B runs, compute the per-domain reference loss predicted at 20% of training (from Eq. 4) and compare it to the actual per-domain evaluation loss at the end of training; report signed errors and rank correlation. Then rerun the DRO scheduling with the reference loss replaced by the actual final loss (oracle) while keeping all other hyperparameters fixed. If the oracle version yields materially better PPL or downstream scores, the scaling-law prediction is a bottleneck and the method's claims should be qualified; if it yields similar results, the prediction is not the limiting factor.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central mechanism of DRPruning depends on Eq. 4: a four-parameter scaling-law fit to each domain's evaluation losses, started after 20% of training, extrapolated to the end of training, and then collapsed to the minimum of the fitted curve as the reference loss for DRO. This extrapolation is load-bearing for two reasons. First, the scaling-law form assumes a fixed training distribution and model size, but here the data distribution is itself reweighted every m steps on the basis of these predictions, creating a feedback loop that can bias the fitted curve. Second, the paper provides no direct validation that the predicted reference losses match actual final per-domain losses; the only evidence offered is that the reference losses decrease gradually (Figure 6), which demonstrates conservatism in sign, not accuracy in magnitude. If the extrapolation systematically misranks domains, the DRO weights will over- or under-emphasize the wrong domains, eroding the method's advantage over fixed scheduling. The empirical failure in Table 5 is consistent with this concern: on the authors' own domain-level benchmark, official Sheared Llama 2.7B scores 89.32 versus DRPruning's 86.00, so the claimed robust domain recovery does not hold in the larger configuration; a biased reference loss for certain domains could explain this gap.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes DRPruning, a data-scheduling method for structured pruning and continued pretraining. It maintains per-domain validation losses, fits a four-parameter scaling law (Eq. 4) to predict the end-of-training loss as the DRO reference loss, and uses an EMA update (Eq. 6) to move the reference data ratio toward high-loss domains. The method is evaluated on Llama2-7B pruned to 1.3B/2.7B with SlimPajama data, on Qwen2-based multilingual models, and through instruction tuning and a GPT-constructed sentence-continuation benchmark. The paper claims lower perplexity, higher downstream scores, and better domain robustness than Sheared Llama-style scheduling and a controlled reproduction called ReSheared.","tokens_in":28269,"tokens_out":6359,"duration_ms":54624,"significance":"If the gains hold, the method is attractive: it adds no GPU overhead, releases code, and automates the two DRO hyperparameters that are normally hard to set. The controlled comparison against ReSheared and the pruning-phase improvements over the official Sheared Llama are concrete strengths, and the five-prompt significance analysis is a useful attempt at robustness checking. However, the headline claim that DRPruning 'surpasses similarly sized models' is not supported for continued pretraining at the 2.7B scale, and the scaling-law reference-loss mechanism is not validated against actual final losses. The paper's value currently rests on a fair-comparison baseline that may be weaker than the official Sheared Llama, so the scope of the claims needs to be narrowed or the missing evidence supplied.","major_comments":[{"comment":"The abstract's claim that DRPruning 'surpasses similarly sized models' in continued pretraining is contradicted by the paper's own data. For the 2.7B target, official Sheared Llama averages 52.53 on the 15-task suite, while DRPruning averages 51.63; on the domain-level benchmark in Section 5.2, official Sheared Llama scores 89.32 versus DRPruning's 86.00. The paper discloses the latter in Section 5.2, but the abstract and conclusion do not carry the necessary qualification. Please either restrict the headline claim to comparisons against ReSheared and other open models under identical data, or provide a 2.7B comparison under the official Sheared Llama data regime.","section":"Section 4.3, Table 2"},{"comment":"The reference loss, which is load-bearing for the DRO weights, is obtained by extrapolating a four-parameter scaling-law fit to partial loss curves, but the paper never validates the predicted end-of-training loss against the actual final per-domain loss. Figure 6 only shows that reference losses decrease gradually, which is a statement about sign, not about accuracy or calibration. This matters because Eq. (4) assumes a fixed training distribution, whereas DRPruning reweights the distribution every m steps based on these predictions. A concrete test would be to report predicted versus observed final per-domain losses under a fixed schedule, or to compare DRPruning with oracle reference losses computed from the actual final losses. Without such evidence, the mechanism's advantage over naive DRO remains unexplained.","section":"Section 3.2, Eq. (4)"},{"comment":"The claimed statistical significance (t=2.99, p=0.0032) is computed from five prompt variants per task on the same training run. These are not independent samples: task scores within a model are correlated, and prompt perturbation does not re-estimate model variance from independent training runs or seeds. The t-test should be run over independent training seeds or, failing that, over tasks with a mixed-effects model. In addition, the significance test compares DRPruning only with ReSheared, not with the official Sheared Llama model that the abstract claims to surpass.","section":"Appendix B.3"},{"comment":"In the multilingual setting, the pruned-and-recovered model (Qwen2-7B to 1.8B with DRPruning, average 51.87) underperforms continued pretraining from the smaller unpruned Qwen2-1.5B with DRPruning (52.35). The paper acknowledges this in Section 5.3, but it limits the scope of the central 'pruning' claim: in this setting the benefit of DRPruning accrues to continued pretraining, not to pruning per se. Please state this qualification in the abstract and conclusion.","section":"Section 5.3, Table 6"}],"minor_comments":[{"comment":"The column header '7B 2.7B 1.3B' does not match the actual column groups; please label the model-size groups explicitly above the Sheared/ReSheared/DRPruning columns.","section":"Table 2"},{"comment":"The phrase 'finer domain segmentatio' is missing the final 'n'.","section":"Appendix B.5"},{"comment":"Equation (6) uses delta for the EMA update factor, while Appendix A.1 says 'lambda set to 0.1, which is also used for updating the reference data ratio.' Please align the notation.","section":"Eq. (6) and Appendix A.1"},{"comment":"The caption contains a missing space in 'y = 0indicating no change'.","section":"Figure 6"},{"comment":"The 1.3B Constant baseline (55.29) is higher than both Sheared Llama (42.79) and ReSheared (47.64) on the generated continuation benchmark; please discuss whether the GPT-4-based filtering introduces a bias that favors the Constant schedule.","section":"Table 5"}],"recommendation":"major_revision","confidential_remarks":"The paper has a strong controlled comparison against ReSheared and a clean engineering contribution, but the abstract's global claims are not supported by the 2.7B results. I would ask the editor to request a revision that (i) validates Eq. (4) with predicted versus actual final losses or an oracle-reference ablation, (ii) recomputes the significance test over independent seeds or with an appropriate multilevel model, and (iii) either reproduces the official Sheared Llama 2.7B data setting or narrows the headline claim to the ReSheared comparison and the pruning phase. The inability to reproduce the 2.7B Sheared Llama setting is a practical constraint, but the paper should not therefore claim to surpass Sheared Llama in continued pretraining."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The method behind DRPruning is genuinely useful: applying group DRO to pruning recovery, with the reference loss set by a scaling-law extrapolation and the reference data ratio updated gradually, is a sensible new combination. The paper does the controlled comparison the right way — reproducing Sheared Llama on the same data (ReSheared) as the primary baseline, adding a multilingual setting, instruction tuning, a domain-level benchmark, and reporting that the overhead is under 1.5% of training time. Code is promised, and the authors are explicit that they could not reproduce the official 2.7B Sheared Llama.\n\nThe soft spots are in the claims, not the engineering. The abstract says DRPruning surpasses similarly sized models, but at 2.7B the official Sheared Llama still beats it on the 15-task average (52.53 vs 51.63) and on the authors' own domain benchmark (89.32 vs 86.00). The controlled win over ReSheared is real and consistent, but the headline needs to be narrowed to what the evidence supports.\n\nThe deeper methodological worry is the reference loss. Eq. 4 is a four-parameter fit to the same loss curve the DRO loop is shaping. Calling it a prediction of the end-of-training loss is generous; it is an extrapolation that has not been validated against actual final losses. Figure 6 shows the references drift down, but not that they are accurate. If the extrapolation misranks domains, the DRO weights will be miscalibrated, and the 2.7B domain-level deficit is consistent with that. This should be the referee's main push.\n\nOne smaller issue: the significance test uses five prompt variants from a single model, not multiple seeds. The p-values (0.046, 0.031) are marginal, and the paired t-test across tasks treats tasks as independent observations, which is questionable. The consistency across PPL, task averages, and win rates is more persuasive than the t-test.\n\nThis paper deserves a serious referee. It is a solid engineering contribution with a new combination and an honest controlled evaluation; it needs a revised abstract, a direct validation of the reference-loss extrapolation, and seed-level replication before publication.","headline":"A solid, well-controlled data-scheduling method for pruning recovery, but the headline overclaims against official Sheared Llama at 2.7B and the reference-loss prediction needs direct validation.","tokens_in":28847,"tokens_out":3437,"would_cite":true,"duration_ms":31750,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"DRPruning claims that dynamically reweighting training data by a scaling-law-predicted reference loss restores balanced performance across domains after structured pruning, outperforming static scheduling in perplexity and downstream tasks.","keywords":["structured pruning","distributionally robust optimization","data scheduling","scaling laws","continued pretraining","multilingual robustness","perplexity","LLM efficiency"],"falsifier":"Compare the predicted reference loss from the scaling-law fit at the 20% checkpoint with the actual end-of-training loss per domain; if the predicted minimum for slow-converging domains is systematically too low or too high, the DRO weight updates are miscalibrated and the perplexity gain over fixed scheduling should disappear when the extrapolation is replaced by the true final loss.","tokens_in":27744,"feed_emoji":"✂️","tokens_out":7573,"duration_ms":58875,"temperature":0.7,"pith_summary":"DRPruning is a data-scheduling method for structured pruning and continued pretraining of large language models that targets the uneven performance drop pruning causes across domains. The paper argues that the standard Sheared Llama-style schedule, which forces the pruned model to track the loss profile of a larger model, does not constrain the space of possible data distributions and leaves some domains behind. DRPruning instead runs distributionally robust optimization with two automatic adjustments: a scaling-law fit predicts each domain's loss at the end of training and uses that as the reference loss, and the reference data ratio itself drifts toward domains whose losses deviate most. In monolingual and multilingual experiments, the method reports lower perplexity, higher downstream task scores, and better instruction-tuning win rates than reproduced baselines, with almost no added training cost. The point is that a pruned model's recovery can be steered by its own predicted loss landscape rather than by hand-set DRO hyperparameters.","feed_headline":"DRPruning beats static schedules by reweighting data on predicted loss","feed_subtitle":"Scaling-law loss forecasts guide dynamic data ratios that cut perplexity and lift downstream scores.","key_machinery":"The carrying object is the DRO update loop that couples three quantities per domain: the current validation loss, a predicted reference loss, and a reference data ratio. The reference loss comes from a four-parameter scaling-law fit, refit after each evaluation once 20% of training is done, and the reference ratio is updated after 40% of training by an exponential moving average of the current DRO weights, constrained by a chi-square divergence ball around the initial ratio. That loop is what lets the method reallocate data toward underperforming domains without collapsing into training only on the highest-loss domain.","core_discovery":"On its own terms, the paper's central claim is that dynamic DRO data scheduling, with a scaling-law-predicted reference loss and a gradually updated reference data ratio, restores balanced performance across domains after structured pruning. During training, evaluation losses per domain refit the curve $\\hat{\\ell}(P,T)=A\\cdot P^{-\\alpha}T^{-\\beta}+E$, and the fitted curve's value at the final training step is taken as the minimum acceptable performance for that domain; the DRO update then raises weight for domains whose current loss exceeds this reference. At the same time, the reference data ratio is moved toward the current DRO weights via $p^{t+1}_R=\\delta q^t+(1-\\delta)p^t_R$, inside a chi-square ball constraint, with each domain's reference ratio bounded between $1/n$ and $n$ times its initial value. The paper reports average gains of $-5.59\\%$ in perplexity, $+1.52\\%$ in downstream tasks, and a $55.4\\%$ instruction-tuning win rate, with the largest downstream gain ($+2.95\\%$) under multilingual distribution shift, plus a $+17.9\\%$ improvement on a domain-level sentence continuation benchmark it constructs.","pith_inferences":["The paper's head-to-head evidence is against ReSheared, its own reproduction of Sheared Llama, not against the officially released Sheared Llama 2.7B model; the appendix states the 2.7B reproduction was not possible under the available data, so the superiority claim over the original 2.7B model rests on the reproduction's fidelity.","If the scaling-law reference loss is as reliable as reported, the same recipe could replace hand-set DRO hyperparameters in other imbalanced training settings, including pretraining from scratch and cross-domain post-training, which the paper lists as untested future directions.","The instruction-tuning machine translation experiment suggests the dynamic data schedule transfers beyond language modeling; a direct test would be applying it to RLHF-style preference data, where domain balance is known to matter.","The mask-similarity analysis shows pruning decisions vary substantially across seeds with no significant perplexity difference, so part of the measured benefit may come from the continued-pretraining data schedule rather than from which substructures get pruned; the paper does not fully separate those effects."],"forward_implications":["Pruned models trained with DRPruning converge faster during the pruning phase and reach lower perplexity than constant or Sheared Llama-style schedules, with the gap widening in later training.","The same dynamic schedule transfers from monolingual pruning to multilingual continued pretraining, where it yields a +2.95% average downstream gain over the reproduced Sheared Llama baseline.","DRPruning introduces no extra GPU compute; data-ratio updates run on CPU and add less than 1.5% to total training time, and parallelizing them does not change performance.","Domain-level sentence continuation built from unlabeled data shows the method improves underperforming domains while keeping other domains stable, so perplexity gains are not achieved by sacrificing balance.","Fine-grained domain segmentation, tested by splitting CC and C4 into smaller clusters, further accelerates early convergence and improves final cross-entropy, suggesting the method benefits from finer-grained supervision."],"supporting_citations":[{"why":"Supplies the structured pruning framework, the target architectures, and the Sheared Llama scheduling baseline that DRPruning builds on and compares against.","marker":"Xia et al., 2024"},{"why":"Defines group DRO and the worst-case loss objective that the paper's dynamic weighting scheme adapts.","marker":"Sagawa et al., 2019"},{"why":"Introduces the reference loss idea for DRO language modeling that DRPruning makes automatic.","marker":"Oren et al., 2019"},{"why":"Provides the f-divergence ball constraint around a reference data ratio that keeps DRO from collapsing to a single high-loss domain.","marker":"Zhou et al., 2021"},{"why":"Establishes the scaling-law form used to predict end-of-training loss as the reference loss.","marker":"Kaplan et al., 2020"},{"why":"Supplies the Huber-loss and L-BFGS fitting procedure for the scaling-law curve.","marker":"Hoffmann et al., 2022a"},{"why":"Supports the prune-then-continue-pretraining pipeline that DRPruning operates within.","marker":"Ma et al., 2023"}],"fun_headline_variants":["DRPruning uses DRO to keep pruned LLMs balanced across domains","Scaling-law loss targets drive robust data reweighting for LLM pruning","DRPruning: dynamic data mix fixes pruning's uneven domain drop","Balanced pruning: DRO reweights data to protect rare domains","Robust LLM pruning via distributionally robust data scheduling"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that the four-parameter scaling-law fit extrapolates each domain's final loss reliably from partial curves taken after 20% of training, and that the reproduced Sheared Llama baseline is a faithful stand-in for the original; if either fails, the DRO weights and the head-to-head comparison are miscalibrated.","fun_headline_variants_meta":{"raw":{"variants":["DRPruning uses DRO to keep pruned LLMs balanced across domains","Scaling-law loss targets drive robust data reweighting for LLM pruning","DRPruning: dynamic data mix fixes pruning's uneven domain drop","Balanced pruning: DRO reweights data to protect rare domains","Robust LLM pruning via distributionally robust data scheduling"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000638,"raw_usage":{"total_tokens":2934,"prompt_tokens":937,"completion_tokens":1997,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":553,"completion_tokens_details":{"reasoning_tokens":1904}},"tokens_in":553,"tokens_out":1997,"duration_ms":13392,"temperature":1.0,"reasoning_tokens":1904,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-12T15:35:21.064524+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Compare the predicted reference loss from the scaling-law fit at the 20% checkpoint with the actual end-of-training loss per domain; if the predicted minimum for slow-converging domains is systematically too low or too high, the DRO weight updates are miscalibrated and the perplexity gain over fixed scheduling should disappear when the extrapolation is replaced by the true final loss.","supporting_citations":[],"review_version":1}