{"id":"0b3201bd-12d8-43a0-9eb1-288f55040256","arxiv_id":"2501.00961","paper_version":3,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"A small subset of neurons memorizes minority-group examples during training, and a pruning-based contrastive fine-tuning framework that targets them improves worst-group accuracy.","lead":"This paper studies why machine learning models perform poorly on minority groups when training data contains spurious correlations, such as image backgrounds. It reports that a small set of neurons memorizes minority examples during training, and proposes a fine-tuning method that removes this memorization to improve worst-group accuracy.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Stage II fine-tuning gains lack a random-mask or no-pruning control, so the WGA improvement cannot yet be attributed to removing the identified critical neurons.","rationale":"The paper's argument has two parts: a descriptive claim that a small number of neurons disproportionately affect minority-group training accuracy, and a causal claim that removing these neurons during fine-tuning improves worst-group test accuracy. The descriptive part is genuinely supported by the group-agnostic magnitude criterion, the random-initialization and random-noise controls in Figures 4-5, and Tables 13-14, which show training accuracy being affected more than test accuracy. The causal part is the load-bearing step, because the proposed fine-tuning changes several ingredients at once relative to ERM, and no control separates the adaptive pruning mask from the contrastive training procedure itself. The reader's weakest assumption identified exactly this missing-control problem, and I agree with that assessment. I would not reject the paper: the phenomenon is plausible, the descriptive evidence is non-trivial, and the missing experiments are straightforward to specify and run. A conditional accept requiring these controls remains the right verdict, so I recommend no change to the reader's judgment.","tokens_in":17886,"tokens_out":7880,"duration_ms":75995,"concrete_test":"On Waterbirds/ResNet-50, repeat the exact Section 2.2.2 protocol (same augmentation, optimizer, scheduler, lambda=0.2, 40-epoch ERM kick-in, 20 fine-tuning epochs, 10 seeds) but replace the gradient/magnitude mask with a mask selecting 0.01% of neurons uniformly at random, redrawn at the start of every fine-tuning epoch. Report mean +/- std WGA. If the random-mask WGA is within error bars of 80.9 (magnitude) or 76.6 (gradient), the fine-tuning improvement does not depend on identifying critical neurons and the central causal claim is unsupported. If random-mask WGA is substantially lower, the targeted selection is necessary, but a further control with an unpruned auxiliary model would still be needed to separate the pruning perturbation from the contrastive objective and to confirm the memorization interpretation.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central causal claim — that spurious memorization concentrated in a few neurons drives imbalanced group performance — rests on the Section 2.2.2 fine-tuning experiments, which compare only against ERM. The proposed loss (Eq. 2.3) changes several things at once: continued fine-tuning, an MSE term, an NT-Xent alignment to an auxiliary model, hard-example-gradient selection, and an adaptively recomputed 0.01% gradient/magnitude mask. The final model is never pruned; pruning only defines the auxiliary branch that provides the contrastive target. No ablation isolates the mask as the cause of the WGA gains. In particular, there is no control with a same-density mask over randomly selected neurons (redrawn each epoch) and no control with an unpruned auxiliary. If a random-mask control reproduces the reported 80.9% (ResNet-50/Waterbirds) or 76.5% (ViT/Waterbirds) WGA, the improvement is a regularization feature of the dual-branch contrastive objective rather than evidence for critical neurons. Stage I's descriptive result (minority training accuracy is unusually sensitive to a few neurons) is reasonably supported — the magnitude criterion is group-agnostic, and the random-initialization/noise controls plus Tables 13–14 support a memorization rather than generalization effect — but it does not establish that removing those neurons improves test imbalance. The Stage I gradient criterion is also partly circular, since it selects top-gradient neurons for a group and then measures that same group's accuracy change, although the magnitude criterion mitigates this concern.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"This paper studies the role of memorization in spurious-correlation group imbalance. Using Waterbirds and CelebA with ResNet-50 and ViT-Small, the authors define 'spurious memorization' as the concentration of minority-group training accuracy in a small set of critical neurons. In Stage I (Section 2.1), gradient- and magnitude-based tracing identify and prune or perturb top-k neurons, showing that minority group training accuracy drops substantially (up to 11.96%) while majority groups drop at most 2.15%; random-initialization and noise ablations plus Tables 13-14 support a memorization interpretation. In Stage II (Section 2.2), a fine-tuning framework contrasts a target model with an adaptively pruned auxiliary model under NT-Xent and MSE losses, reporting worst-group accuracy gains from 64.0 to 80.9 on Waterbirds/ResNet-50 and from 52.7 to 76.5 on Waterbirds/ViT-Small. The paper concludes that removing spurious memorization of critical neurons causes the robustness improvement.","tokens_in":18175,"tokens_out":4101,"duration_ms":36141,"significance":"The descriptive finding is valuable and, in part, well supported: the magnitude-based criterion is group-invariant, the random-initialization and noise controls show that minority training accuracy is unusually sensitive to a few neurons, and the train/test discrepancy in Tables 13-14 supports a memorization rather than a generalization role for those neurons. If the causal claim were established, the paper would be a strong contribution to mechanistic accounts of group robustness. However, the causal claim is not yet established, because the Stage II framework varies several components jointly and lacks controls for random masks or an unpruned auxiliary. The code link is included, which is useful for reproducibility, but the missing intervention controls are the main barrier to the paper's central conclusion.","major_comments":[{"comment":"The worst-group accuracy improvements in Figure 6 cannot be attributed to removal of the identified critical neurons, because the comparison is only against ERM. The proposed fine-tuning simultaneously changes the training schedule, replaces cross-entropy with MSE, adds an NT-Xent contrastive term against an auxiliary branch, selects hard examples for gradient computation, and adaptively recomputes a 0.01% mask each epoch. There is no control with an unpruned auxiliary branch, no control with a same-density random mask redrawn each epoch, and no control with a fixed non-adaptive mask. If a random-mask condition reproduces the 80.9% or 76.5% WGA, the result would be explained by the dual-branch contrastive regularizer rather than by critical-neuron memorization. Since the central causal claim in Section 3 rests on this experiment, these controls are required before the claim can be accepted.","section":"Section 2.2.1-2.2.2, Eq. (2.3)"},{"comment":"The gradient-based tracing criterion is partially circular: neurons are selected as the top-gradient neurons for group j and then evaluated on the same group j. The magnitude criterion does not share this issue, and Figure 3 shows overlap between the two sets, so the descriptive claim survives; nevertheless, the textual claim that 'both gradient and magnitude-based criteria are effective' should be qualified, and the authors should present the magnitude-based results as the primary non-circular evidence for critical-neuron existence.","section":"Section 2.1.1, Figure 2"},{"comment":"The ablation in Table 6 shows that 10 fine-tuning epochs yields WGA 54.98%, which is below ERM, while 20 and 30 epochs yield 80.89%. This large and non-monotonic dependence on fine-tuning length suggests that the reported gains are not solely a simple function of the mask; without disclosing the checkpoint-selection details, this sensitivity undermines the attribution of the final gain to spurious memorization removal and should be discussed explicitly.","section":"Section 2.2.1, Table 6"}],"minor_comments":[{"comment":"The Methods text states that 'all accuracy metrics reported in this paper are derived from the test set', which contradicts the explicit statement in Section 2.1 that all group accuracy evaluated before and after pruning is evaluated on the training set; please clarify which numbers are train and which are test.","section":"Section 4"},{"comment":"In the definition of cosine similarity, the denominator reads 'u· v/(||u||·|| b||)' where b is undefined; it should presumably be v.","section":"Eq. (2.1)"},{"comment":"The caption contains a typo: 'L NE-Xent' should be 'NT-Xent'.","section":"Figure 9 caption"},{"comment":"The paper says 0.01% pruning percentage serves as a hyperparameter, but Table 9 shows extreme sensitivity (10% gradient-based pruning drops WGA to 0.04%); a brief discussion of this sensitivity and its implications for the 0.01% default would strengthen the presentation.","section":"Section 2.2.1, Table 9"}],"recommendation":"major_revision","confidential_remarks":"The manuscript's central weakness is the missing intervention control for the Stage II framework; however, the descriptive tracing evidence is credible and the framework is clearly presented. I recommend major revision rather than rejection because the missing controls are feasible within the paper's scope and the descriptive contribution is solid enough to warrant a revision."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: the descriptive result is genuinely interesting and reasonably supported. A small set of neurons — channels in ResNet, units in ViT — can be identified by gradient or magnitude such that zeroing them drops minority-group training accuracy by up to ~10-12% while majority groups move a couple of points. The magnitude criterion is group-agnostic, and the random-initialization/noise ablations plus the train/test contrast in Tables 13–14 strengthen the memorization reading: modifying these neurons moves training accuracy much more than test accuracy. That is a real empirical phenomenon, and it is not something I had seen demonstrated cleanly in the spurious-correlation setting before.\n\nThe new twist is the two-stage framing: trace critical neurons, then use them to define an auxiliary pruned branch in a contrastive fine-tuning loss. The WGA gains are large (80.9 vs 64.0 on Waterbirds/ResNet-50; 76.5 vs 52.7 on ViT). But this is where the paper goes soft. The fine-tuning changes several things at once: continued training, an MSE term, an NT-Xent alignment to an auxiliary model, hard-example selection, and an adaptively recomputed 0.01% mask. No control with a random mask of the same density, and no control with an unpruned auxiliary. Without those, you cannot attribute the WGA improvement to removing the identified critical neurons. It could be a generic regularization effect of the dual-branch contrastive objective. The ablations in Tables 4–10 tune hyperparameters but do not isolate the mask as the causal ingredient. That is a load-bearing gap for the mechanistic claim, not a cosmetic one.\n\nMinor issues: the main figures lack error bars (Table 12 gives means±std, which helps), and the gradient-based tracing criterion is partly circular because it selects neurons by a group's gradient and then measures that same group's accuracy. The magnitude criterion mitigates this, so I would not call it fatal. Also, the 'first-ever evidence' language in the abstract is over-claimed; the paper extends existing localization results (Maini et al. 2023, Stephenson et al. 2021) rather than inventing the concept.\n\nWho this is for: people working on spurious correlations, memorization, and interpretability. The descriptive tracing result is worth engaging with seriously. The fine-tuning method needs a revision with the missing controls before I would believe the causal story. But that is fixable, and the paper deserves a serious referee. I would send it out rather than desk-reject.","headline":"The descriptive finding on critical neurons is solid and worth taking seriously; the fine-tuning story needs controls before the mechanistic claim lands.","tokens_in":18753,"tokens_out":1875,"would_cite":true,"duration_ms":16363,"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":"The paper claims that spurious memorization, concentrated in a tiny subset of neurons, drives worst-group accuracy gaps, and that pruning 0.01% of those neurons during fine-tuning substantially closes those gaps.","keywords":["spurious correlations","memorization","critical neurons","spurious memorization","worst-group accuracy","group robustness","neural network pruning","contrastive learning"],"falsifier":"Run the same fine-tuning framework on Waterbirds/ResNet-50 with no pruning, or with 0.01% of neurons chosen at random for the auxiliary model; if worst-group accuracy still rises by roughly 17 points over ERM, then the identified critical neurons are not the cause of the improvement.","tokens_in":17679,"feed_emoji":"🧠","tokens_out":6902,"duration_ms":50180,"temperature":0.7,"pith_summary":"The paper argues that when models trained on spurious correlations perform badly on minority groups, the root cause is memorization: a small set of critical neurons learns to predict atypical (minority) training examples, so minority accuracy is high in training but collapses at test time because it depends on a few fragile neurons. It establishes this by showing that zeroing out one to three of these neurons drops minority training accuracy by up to 11.96% while majority groups drop at most 2.15%. It then proposes a fine-tuning framework that prunes 0.01% of neurons and contrasts the pruned model with the target model, raising worst-group test accuracy from 64.0% to 80.9% on Waterbirds/ResNet-50 and from 52.7% to 76.5% on Waterbirds/ViT-Small. If correct, this means spurious-correlation failure is localized in a tiny fraction of parameters and can be mitigated with a very small intervention.","feed_headline":"Pruning 0.01% of neurons lifts worst-group accuracy by 17 points","feed_subtitle":"Contrastive fine-tuning against a pruned twin raises Waterbirds worst-group accuracy from 64.0% to 80.9%.","key_machinery":"The central object is the 'critical neuron': a channel (in ResNet-50) or neuron (in ViT-Small) whose removal disproportionately affects minority group accuracy. They are identified by two criteria: gradient norm (group-specific loss gradient) and weight magnitude (group-invariant). The mechanism that carries the argument is a dual-branch fine-tuning framework: an auxiliary model is created by masking the top 0.01% of neurons by gradient or magnitude, and the target model is trained to align its features with the auxiliary model using NT-Xent contrastive loss plus MSE classification loss, so the model learns not to rely on the spurious memorization encapsulated in those neurons.","core_discovery":"On the paper's own terms: neural networks trained by empirical risk minimization with spurious correlations store minority-group information in a very small set of 'critical neurons'. Evidence comes from three sources: unstructured tracing (zeroing top-k neurons by gradient or magnitude across the whole network), structured tracing (same within each layer), and perturbation experiments (random initialization and random noise). All three show minority-group training accuracy is disproportionately sensitive to these neurons, while majority groups are barely affected. The paper further claims that contrasting the target model with a pruned auxiliary model during fine-tuning removes this spurious memorization and improves worst-group accuracy.","pith_inferences":["If the mechanism generalizes, measuring the concentration of critical neurons (e.g., how few neurons carry minority accuracy) could predict a model's vulnerability to spurious correlations before test-time evaluation.","The paper's gradient-based criterion evaluates neurons on the same group whose accuracy they affect; a stricter test would check whether neurons identified on one minority group also matter for another, which would separate memorization from generic feature importance.","No control runs the fine-tuning framework with zero pruning or random pruning, so part of the reported WGA gains may be attributable to contrastive fine-tuning itself rather than specifically to removing critical neurons.","The 0.01% pruning ratio is surprisingly small; testing whether the optimal ratio scales with dataset size or number of minority examples would clarify whether the memorization set is truly fixed or grows with data."],"forward_implications":["Worst-group robustness can be improved by intervening on 0.01% of parameters, suggesting that spurious memorization is highly localized rather than diffuse.","The critical-neuron criterion (gradient norm or weight magnitude) provides a cheap diagnostic for which groups are memorized rather than generalized.","Contrasting a model with its pruned twin can serve as a training signal that discourages reliance on a few neurons, without needing group labels for every sample.","The reported gains on Waterbirds and CelebA across ResNet-50 and ViT-Small indicate the mechanism may transfer across architectures and spurious-correlation benchmarks."],"supporting_citations":[{"why":"Motivates the auxiliary-model idea by showing pruning can improve predictions on rare and atypical examples.","marker":"Hooker et al. [2019]"},{"why":"Provides evidence that memorization can be localized in specific neurons, grounding the search for critical neurons.","marker":"Maini et al. [2023]"},{"why":"Introduces the Waterbirds benchmark and the group-robustness formulation of spurious correlations.","marker":"Sagawa et al. [2020]"},{"why":"Supplies the simCLR contrastive learning framework that the fine-tuning loss is built on.","marker":"Chen et al. [2020]"},{"why":"Defines the NT-Xent loss used to align target and pruned auxiliary model features.","marker":"Sohn [2016]"},{"why":"Provides the CUB dataset that Waterbirds is constructed from.","marker":"Wah et al. [2011]"}],"fun_headline_variants":["Tiny neuron subset drives spurious correlation failures","Pruning 0.01% of neurons fixes worst-group accuracy","Memorization hides in 0.01% of neurons","Spurious correlations traced to a few critical neurons","New method prunes spurious memorization for fairer models"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing assumption is that the accuracy gains come specifically from removing the identified critical neurons, rather than from the fine-tuning procedure or from pruning any random neurons; the paper does not run a control that applies the same contrastive fine-tuning without pruning or with randomly selected neurons.","fun_headline_variants_meta":{"raw":{"variants":["Tiny neuron subset drives spurious correlation failures","Pruning 0.01% of neurons fixes worst-group accuracy","Memorization hides in 0.01% of neurons","Spurious correlations traced to a few critical neurons","New method prunes spurious memorization for fairer models"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00018,"raw_usage":{"total_tokens":1294,"prompt_tokens":923,"completion_tokens":371,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":539,"completion_tokens_details":{"reasoning_tokens":290}},"tokens_in":539,"tokens_out":371,"duration_ms":4294,"temperature":1.0,"reasoning_tokens":290,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T22:38:00.153392+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the same fine-tuning framework on Waterbirds/ResNet-50 with no pruning, or with 0.01% of neurons chosen at random for the auxiliary model; if worst-group accuracy still rises by roughly 17 points over ERM, then the identified critical neurons are not the cause of the improvement.","supporting_citations":[{"cited_title":"Can neural network memorization be localized? In Proceedings of the 40th International Conference on Machine Learning, volume 202, pages 23536--23557","cited_arxiv_id":null,"evidence_quote":"Provides evidence that memorization can be localized in specific neurons, grounding the search for critical neurons."}],"review_version":1}