{"id":"cd6c8382-df43-402f-aa51-3de76c078200","arxiv_id":"2506.03978","paper_version":2,"verdict":"REJECT","confidence":"HIGH","novelty_score":6.0,"correctness_risk":"high","formal_verification":"none","parameter_count":4,"one_line_summary":"SPRINT learns to select which attention heads to prune per question, improving Pass@N over random head selection and multinomial sampling on MATH500 and GSM8K.","lead":"The paper shows that pruning certain attention heads in a math language model can improve accuracy on some questions, and proposes SPRINT, a contrastive-learning method that picks which heads to prune for each new question. It suggests a new, cheap way to create diverse candidate answers for reasoning tasks without training multiple models, though the experiments are small-scale and the comparison to best-of-N is misleading.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The paper never validates that the contrastive selector transfers to unseen questions; Pass@N gains could be pruning diversity rather than the learned per-question alignment.","rationale":"The central claim assumes that the contrastive embeddings trained on 1,500 questions generalize to new questions, positioning each test question near the embeddings of heads whose pruning yields a correct answer. That assumption is the load-bearing part of the proposed mechanism, and it is never directly checked. Tables 1–4 show that pruning can help, but they select the best head per subject, so they do not establish that the best head for a question is predictable from a sentence embedding. Figure 4 provides only aggregate Pass@N curves, with no error bars and an ambiguous random-head baseline, so the reported gains could come from the diversity of any set of pruned heads rather than from the learned alignment. The proposed test directly measures selector precision and coverage on held-out questions, which would separate the method's actual contribution from the trivial 'pruning adds diversity' effect. This concern supports the reader's rejection; the verdict is unchanged.","tokens_in":21120,"tokens_out":10137,"duration_ms":104411,"concrete_test":"On a held-out set of 300 MATH training questions not used to train θ and V, compute the true correctness z_ij for every pruned head (greedy decode, oracle/exact-match check). Then compare: (a) SPRINT top-1 hit rate, the fraction of questions where the nearest head embedding is in M_i^+; (b) SPRINT top-N coverage, the fraction where at least one of the top-N heads is in M_i^+; (c) the same coverage for N randomly chosen heads and for the N fixed heads with the highest training-set solve rates. Repeat over 5 random training subsamples and report means with 95% CIs. If SPRINT's top-1 hit rate is not significantly above the random-head solve rate, and its top-N coverage is not above the fixed-best-N baseline, the Pass@N gains in Figure 4 cannot be attributed to dynamic selection.","verdict_should_be":"UNCHANGED","load_bearing_attack":"SPRINT's central novelty is the per-question selection of heads (Section 4, Eq. 1), but the paper never measures whether the learned embeddings actually rank, for an unseen question, the heads whose pruning leads to a correct answer. The Pass@N curves in Figure 4 conflate three effects: (1) pruning alone can improve accuracy (Section 3.2), (2) using any set of N different pruned heads adds diversity relative to multinomial samples from one model, and (3) the learned question-to-head selector. The only pruning-aware baseline, 'random head selection', is described ambiguously ('we use a greedy approach to identify the heads that solve the most samples in the training set. During testing, we randomly select N heads') and is reported without error bars, so it does not isolate effect (3). Without a held-out test of the selector's top-1/top-N hit rate against random and fixed-best-head baselines, the observed gains cannot be attributed to the proposed contrastive alignment. If the selector degenerates to random selection on new questions, the central claim of dynamic pruning optimization fails.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"This paper investigates whether selectively pruning individual attention heads in a transformer LM can improve mathematical reasoning. It reports that for several Qwen and Llama models, pruning certain heads raises per-category accuracy on MATH500. It then proposes SPRINT, a contrastive framework that learns embeddings for head-layer pairs and, at test time, selects the top-N nearest heads to a question embedding to generate N pruned-model candidates. Experiments on MATH500 and GSM8K with Qwen2.5-Math-1.5B-Instruct report Pass@N curves comparing SPRINT with multinomial sampling at three temperatures and a random head selection baseline. The authors claim SPRINT significantly outperforms these baselines at low N.","tokens_in":21357,"tokens_out":8329,"duration_ms":76931,"significance":"The core idea—using pruning to create diverse reasoning candidates and selecting the head per question via contrastive embeddings—is interesting and, if validated, could offer a lightweight way to improve best-of-N reasoning without model retraining. The paper also documents a potentially surprising phenomenon that head pruning can improve accuracy. However, the current evidence is not strong enough: the main experimental claim lacks statistical support, the baselines do not isolate the contribution of the learned selector, and the motivating phenomenon is vulnerable to post hoc selection artifacts. The method itself is described at a high level but lacks implementation details needed for reproduction.","major_comments":[{"comment":"The claim that SPRINT 'significantly outperforms' multinomial sampling and random head selection is not supported by any statistical analysis. The Pass@N curves in Figure 4 are single-point estimates with no error bars, confidence intervals, or significance tests. With 500 and 1,319 test questions, differences of a few percentage points are within the range of sampling noise; the authors should report bootstrap confidence intervals or repeated-seed results and apply a paired test (e.g., McNemar's test on question-level correctness) to support the headline claim.","section":"Section 5.1, Figure 4"},{"comment":"The experimental design does not isolate the contribution of the learned per-question head selection. The only pruning-aware baseline is described as 'we use a greedy approach to identify the heads that solve the most samples in the training set. During testing, we randomly select N heads.' This is ambiguous: if the greedy step selects a small set of training-accurate heads and random selection is within that set, the baseline is a fixed-best-heads-plus-random-diversity method; if selection is over all heads, it may include catastrophic heads. Neither variant controls for the diversity effect of using N different pruned heads. The paper should add a fixed-best-heads baseline (e.g., the N heads with the highest training-set accuracy, used for all test questions) and an oracle per-question upper bound, and should directly report the top-1/top-N head retrieval accuracy of the learned embeddings on held-out questions. Without these, the Pass@N gains cannot be attributed to the contrastive alignment.","section":"Section 5 (Baselines) and Section 4"},{"comment":"The motivating observation that 'pruning certain heads leads to improvements' is based on the best pruned head per category (violin plots of additive gains in Figure 2 and Tables 1–4). Because the maximum is taken over many head-layer pairs (e.g., 48 for Qwen2.5-Math-1.5B in Table 2) and category sample sizes in MATH500 are small (often 30–100 questions per category), positive maximum gains are expected under pure noise even if no head truly improves accuracy. A multiple-comparison correction or a permutation test against the no-prune baseline is needed before the phenomenon is treated as established.","section":"Section 3.2"},{"comment":"The description of SPRINT is not reproducible. The sentence embedding model φ is never identified (Figure 3 hints at Sentence-BERT, but no model name or dimension p is given), the dimension and initialization of head embeddings V are not specified, and the training procedure (optimizer, learning rate, number of epochs, batch size, and the value of λ) is entirely absent. Without these details, the joint optimization in Eq. (1) cannot be reimplemented or compared with future work.","section":"Section 4, Eq. (1) and Appendix A"}],"minor_comments":[{"comment":"Section 3.1 states that L=4 layers are selected (first layer, layers 5 and 15, and the last layer), while Section 5 uses five layers {5, 10, 15, 20, 25}; the paper should clarify which configuration applies to the main experiments and why the motivating analysis uses different layers.","section":"Section 3.1 vs. Section 5"},{"comment":"The reference list contains a suspicious placeholder-style entry (John Doe, Jane Smith, and Alan Turing, 2023) and a duplicate entry for Li et al. 2024a/2023; the authors should verify and clean the bibliography.","section":"References"},{"comment":"The abstract uses 'best-of-$N$' with math mode while the rest of the paper uses 'best-of-N'; please make the notation consistent.","section":"Abstract"},{"comment":"The per-category tables do not report category sample sizes, which makes it hard to judge the reliability of the reported differences; adding the number of questions per category would improve interpretability.","section":"Tables 1–4"},{"comment":"The sentence 'In this experiment 1, we study...' contains a stray footnote marker, and the y-axis label of Figure 2 repeats the subtitle; both should be cleaned up.","section":"Section 4, Figure 2"},{"comment":"The statement that SPRINT works 'without introducing inference-time overhead' should be made precise: generating N candidates from N pruned models still requires N forward passes, just like multinomial best-of-N, so the intended meaning is presumably that no additional reward model or retraining is needed.","section":"Section 5"}],"recommendation":"major_revision","confidential_remarks":"The draft is not fully polished: the reference list contains an apparent placeholder, and the implementation details are too thin for a methods paper. The central idea is interesting, but the evidence is currently insufficient because the key baselines are ambiguous and no statistical support is provided. I would not recommend rejection if the authors can add the missing ablation, direct retrieval evaluation, and error bars; however, without those additions, the paper's central claim remains unsubstantiated."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Short version: the paper has a genuinely new selection mechanism and a solidly documented pruning phenomenon, but the main claim is not backed by enough evidence. I would let a referee take a look, but only with a clear request for a transfer test and proper baselines.\n\nThe new thing is SPRINT's contrastive question-to-head mapping. I do not see that in the cited literature, and the idea is concrete: learn head embeddings, push question embeddings toward heads whose pruning leads to a correct answer, and add a diversity term that separates heads with similar outcome profiles. The motivating section is also worth something. The per-category tables across four models are a lot of work, and they do show that pruning a single head can beat the unpruned model on specific categories. The held-out evaluation is separated from the training split, so the core result is not circular.\n\nThe soft spots are real, though. The central experiment compares Pass@N curves against multinomial sampling and a vaguely described 'random head selection' baseline, with no error bars or significance tests. The gains are a few points on one 1.5B model at small N. More importantly, the Pass@N comparison conflates three effects: pruning alone can help, any set of N different pruned heads adds diversity, and the learned per-question selector. The random-head baseline does not isolate the selector, and the paper never measures whether the contrastive embeddings actually rank beneficial heads correctly on held-out questions. That is exactly the stress-test concern, and it lands. The sentence encoder, embedding dimension, lambda, and training details are underspecified; the reference list also has a placeholder citation ('Doe et al. 2023'), which should have been caught.\n\nNone of this kills the idea. The method is testable and the held-out Pass@N results are suggestive. But as written, the central claim that SPRINT 'significantly outperforms' best-of-N is not verifiable. A revision measuring top-1/top-N hit rate of the selector against random and fixed-best baselines, adding error bars, and comparing to a fixed-best-head baseline would make the evidence match the claim.\n\nWho is this for? People working on inference-time diversity or structured pruning for reasoning. It is not going to change how reasoning models are built, but it is a reasonable short-paper contribution. I would send it to reviewers rather than desk-reject: the mechanism is new, the data is there, and the missing analyses are fixable. I would not cite it in its current form.","headline":"A new contrastive head-selection mechanism with a solid motivating study, but the main Pass@N claim is not yet supported because the evaluation never isolates the learned selector.","tokens_in":21836,"tokens_out":2935,"would_cite":false,"duration_ms":29837,"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":"Selectively pruning attention heads in a transformer can improve, not just shrink, a language model's mathematical reasoning, and the paper proposes a contrastive method to choose which head to prune for each question.","keywords":["attention head pruning","best-of-N sampling","contrastive learning","mathematical reasoning","large language models","Pass@N","structured pruning","inference-time diversity"],"falsifier":"Give a held-out set of math questions and record, for each question, the accuracy of the single best pruned head. Then compare the Pass@1 of SPRINT's selected head against the distribution of randomly chosen heads on those questions; if the selected head is not reliably better than a random head, or if a fixed head chosen once on the training set matches its performance, the alignment is doing no work.","tokens_in":20903,"feed_emoji":"🧠","tokens_out":6729,"duration_ms":55332,"temperature":0.7,"pith_summary":"Pruning attention heads in a transformer is usually seen as damage, but the paper reports that it can improve mathematical reasoning: some heads add noise, and which head to remove depends on the question. To exploit this, the paper proposes SPRINT, which learns one embedding per head-layer pair and prunes the heads nearest to the question embedding when generating candidates. On MATH500 and GSM8K, this dynamic selection beats temperature-based best-of-N sampling and random head selection, with the largest gains when only a few candidates are generated. The practical payoff is a more diverse set of candidates from one base model, with no extra inference-time cost.","feed_headline":"Selectively pruning attention heads improves math reasoning at low N","feed_subtitle":"SPRINT picks which head to prune per question, outperforming best-of-N and random selection on math benchmarks.","key_machinery":"The carrying object is the set of learnable head embeddings $v_j$ for every head-layer pair, aligned with question embeddings $q_i = \\theta(\\varphi(x_i))$ through a contrastive loss with a diversity regularizer. The first loss term is a softmax over negative squared Euclidean distances $\\lVert q_i - v_j\\rVert_2^2$, which pulls each question toward heads whose pruning solves it; the second term $\\lambda \\sum_{j<k} s_{jk} \\lVert v_j - v_k\\rVert_2^2$, weighted by the training-set agreement score $s_{jk}$, pushes similar heads apart. At inference, the selector simply takes the head (or top-N heads) closest to the question embedding and prunes them before greedy decoding. This makes pruning choices cheap and input-dependent.","core_discovery":"The paper's central claim is that a single-head pruning choice is a per-question decision, not a global one, and that a contrastive embedding can learn to make it. For each training question, the method records which heads, when pruned, yield a correct answer; it then optimizes head embeddings and a linear projection of a sentence embedding so that the question embedding sits close to those beneficial heads, while a diversity term spreads apart heads with similar correctness patterns. At test time, pruning the top-N nearest heads produces several answers from structurally different model variants. The paper asserts that this approach consistently outperforms multinomial sampling at temperatures 0.1, 0.6, and 0.9 and beats random head selection across both datasets.","pith_inferences":["Editorial inference: an extension the paper does not pursue is combining pruned candidates with a verifier or reward model, since the pruned variants are cheaply generated and a second-stage scorer could select among them without extra forward passes.","Editorial inference: a testable extension is pruning sets of heads instead of one head at a time; single-head gains are small per question, and multiple simultaneous removals might create larger, more useful diversity.","Editorial inference: the unspecified sentence encoder is a variable worth isolating, because the method's transfer on out-of-domain questions would reveal whether the learned alignment is about question content or only about dataset-specific patterns."],"forward_implications":["A single base model yields a diverse pool of candidate generators by pruning different heads, so best-of-N diversity does not require sampling temperature or multiple models.","The gains concentrate at $N \\in \\{1,2,3,4\\}$, where ordinary multinomial sampling supplies the least diversity, indicating pruning adds a kind of diversity temperature cannot.","Because the selector is a lookup over learned embeddings, applying it at inference adds negligible overhead relative to generating the candidates.","If the alignment generalizes, the same recipe applies to any reasoning dataset where head-level correctness can be measured on training questions."],"supporting_citations":[{"why":"Defines best-of-N sampling with a reward model, the baseline that SPRINT is designed to improve.","marker":"Stiennon et al., 2020"},{"why":"Establishes self-consistency and majority voting as a multi-sample reasoning baseline that motivates generating diverse candidates.","marker":"Wang et al., 2022"},{"why":"Shows structural pruning of large language models can preserve performance, supporting the claim that heads can be removed without loss.","marker":"Ma et al., 2023"},{"why":"Introduces structured pruning of attention heads, the operation SPRINT applies at inference.","marker":"Wang et al., 2019"},{"why":"Identifies 'superweights' whose removal is catastrophic, which the paper cites to explain why some early-layer heads are irreplaceable.","marker":"Yu et al., 2024"},{"why":"Provides the MATH benchmark whose 500-question split is used for the pruning analysis.","marker":"Hendrycks et al., 2021"},{"why":"Provides GSM8K, the second evaluation dataset for the method.","marker":"Cobbe et al., 2021"},{"why":"Shows accuracy on MATH-500 scales with test-time compute, the setting that best-of-N belongs to and that SPRINT targets.","marker":"Beeching et al., 2024"}],"fun_headline_variants":["Per-question head pruning outperforms best-of-N on math","Contrastive learning picks the right head to prune per question","Pruning the right head per question improves math reasoning","SPRINT: per-question head selection beats random pruning"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that the contrastive embeddings learned from 1,500 training questions generalize, so a new question's embedding lands near heads whose pruning will actually help that question; if that alignment fails, the selector is no better than random head choice.","fun_headline_variants_meta":{"raw":{"variants":["Per-question head pruning outperforms best-of-N on math","Contrastive learning picks the right head to prune per question","Pruning the right head per question improves math reasoning","SPRINT: per-question head selection beats random pruning"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000594,"raw_usage":{"total_tokens":2709,"prompt_tokens":798,"completion_tokens":1911,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":414,"completion_tokens_details":{"reasoning_tokens":1842}},"tokens_in":414,"tokens_out":1911,"duration_ms":13391,"temperature":1.0,"reasoning_tokens":1842,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T10:49:48.989015+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Give a held-out set of math questions and record, for each question, the accuracy of the single best pruned head. Then compare the Pass@1 of SPRINT's selected head against the distribution of randomly chosen heads on those questions; if the selected head is not reliably better than a random head, or if a fixed head chosen once on the training set matches its performance, the alignment is doing no work.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Shows accuracy on MATH-500 scales with test-time compute, the setting that best-of-N belongs to and that SPRINT targets."}],"review_version":1}