Detecting and Pruning Prominent but Detrimental Neurons in Large Language Models

Thank you for reviewing our work. Below are our replies to your concerns.

Limited theoretical justification provided regarding the specific selection of Integrated Gradients versus other neuron attribution methods.

The core idea of our approach is detecting a subset of neurons that contribute most positively to the model’s prediction which is general and can, in principle, be applied with various attribution methods, such as saliency maps, SHAP, or other gradient-based techniques, as long as they reliably identify influential neurons. This flexibility underscores the robustness of our DSM neuron pruning framework, which prioritizes identifying domain-specific mechanisms (DSMs) regardless of the specific attribution tool. Our choice of Integrated Gradients (IG) was motivated by both its axiomatic guarantees and practical suitability for the neuron-level attribution task in our setting: Axiomatic Justification: IG satisfies important theoretical properties such as sensitivity and implementation invariance (Sundararajan et al., 2017) 1, which are particularly relevant when comparing neuron contributions in large models where functional equivalence and saturation issues can confound other gradient-based methods (e.g., saliency or raw gradients). Applicability to Neuron-Level Attribution: IG extends naturally to internal neurons by treating weights (or activations) as the path of interpolation. Prior works (e.g., Dai et al., 2022; Ali et al., 2024) [2,3] have successfully used IG to analyze or prune neurons, and our method builds on this foundation by adapting IG for few-shot neuron scoring without needing labeled data.

The paper lacks detailed analysis or visualization of exactly how the pruned neurons were detrimental, limiting interpretability claims.

Can you provide specific examples or qualitative analysis showing exactly how pruned DSM neurons were contributing detrimentally to shortcut learning?

Following the review, we provide a qualitative analysis using an SST2 example with the LLAMA2-7B-CHAT model, showing Integrated Gradients (IG) attribution scores at the input space to highlight changes in token importance before and after pruning.

Example Sentence: "The movie was absolutely dreadful, a complete disaster." (Ground Truth: Negative)

Before Pruning:

Prediction: Positive (incorrect)

IG Attribution Scores (for the maximum logit, "positive"):

1. "absolutely": 0.50
2. "dreadful": 0.25
3. "disaster": 0.15

Observation: The model incorrectly predicts positive due to high attribution to "absolutely," an intensifier often correlated with positive sentiment in SST2 training data. This shortcut causes the model to overlook the negative sentiment conveyed by "dreadful" and "disaster," leading to poor generalization

After Pruning (DSM neurons removed from layer 10):

Prediction: Negative (correct)

IG Attribution Scores:

1. "dreadful": 0.35
2. "disaster": 0.40
3. "absolutely": 0.20

Observation: Post-pruning, the model correctly predicts negative by shifting focus to "dreadful" and "disaster," which carry the semantic weight of negative sentiment. The reduced attribution to "absolutely" indicates diminished reliance on the shortcut, enhancing generalization.

We will add an in-depth analysis in the next revision of the paper.

Experiments are predominantly on multiple-choice tasks; it remains unclear whether the approach generalizes effectively to more complex language tasks like generation or open-ended reasoning.

Have you tested or do you have results indicating if this neuron-level pruning approach generalizes to open-ended generation or dialogue tasks beyond multiple-choice questions?

Our focus on multiple-choice tasks leverages their structured outputs, enabling precise application of Integrated Gradients (IG) to identify DSM neurons by attributing contributions to the maximum logit. This clarity facilitates efficient detection of domain-specific mechanisms, as shown across diverse benchmarks. Extending IG, which is not designed for autoregressive tasks, to open-ended generation like summarization is challenging due to the complexity of attributing neuron contributions across sequential outputs. While we hypothesize that our pruning approach can generalize by adapting IG to focus on key output tokens, we prioritized multiple-choice tasks to establish a robust foundation. We are exploring adaptations for generative tasks and will include this discussion in the revised paper.

We appreciate your thoughtful feedback. Please find our responses to each of your concerns below.

Discrepancy between the paper's goal and its evaluation. The goal of the paper is to remove domain-specific spurious correlation of the model. However, the evaluation is based on test accuracy in one domain. Ideally experiments should be conducted with a source and a target domain in mind, and use accuracies in both domains to demonstrate the effectiveness of removing domain-specific spurious correlations.

We thank the reviewer for highlighting this issue. We acknowledge that the introduction’s reference to “original” and “target” domains may suggest a cross-domain evaluation, which could confuse readers given our within-domain experiments. To clarify, our method aims to mitigate domain-specific mechanisms (DSMs) that exploit spurious correlations within a task’s training data (e.g., BoolQ, SST2), thereby enhancing in-domain generalization. In this paper, we examine a setup similar to those studied by SADI, CAA, and ITI, focusing on improving generalization within a single task / domain. To avoid confusion, in the next revision we will rephrase the goals - for example: “we hypothesize that this allows us to identify mechanisms that exploit spurious correlations in the training data, improving generalization to the test data” - to avoid this confusion and better align with our experimental design. While improving cross-domain performance with a low-resource method is also an important goal, we believe that improving in-domain performance is likewise an important research direction, and the consistent performance improvements validate our approach. We hope this clarification addresses the reviewer’s concern, and we leave cross-domain evaluation for future work.

The comparison with baselines is confusing.

The compared methods are using more datapoints than the proposed method (Line 236). The number of datapoints are not discussed in experiment section but instead in the related work section. Ideally all methods should be compared with the same number of datapoints.

Our DSM neuron pruning method is designed to be data-efficient, achieving strong performance with only 10 samples per task, a significant advantage over baselines like SADI, which, as noted in the SADI paper [1], requires substantially larger datasets for contrastive pair construction: 2500 for NLI, 500 for MMLU (testing set, transductive), 2000 for SST2, 2000 for BoolQ, 2000 for SST5, and 1500 for COPA. This data efficiency is a core strength of our approach, as baselines like SADI, ITI, and CAA often rely on thousands of samples to achieve comparable results and may not perform effectively with as few as 10 samples. To clarify this, we will revise the experiment section (Section 4) to explicitly state the number of data points used by each method, highlighting our method’s minimal data requirements. To further demonstrate that our approach does not benefit from larger adaptation sets, we conducted an experiment on BoolQ using LLaMA3.1-8B-Instruct with 5 different seeds, testing 10, 50, and 100 samples, as shown below. The results show stable performance with 10 samples (82.93 ± 0.65) and no significant improvement with more samples (82.57 ± 0.05 for 50, 82.43 ± 0.02 for 100), with a slight decline possibly due to overfitting to larger adaptation sets. This underscores our method’s efficiency with minimal data.

[1] : https://arxiv.org/pdf/2410.12299

One of the advantages of the paper is the data efficiency of the method. The proposed method uses at most 10 datapoints. A straightforward baseline is to fine-tune the model on these 10 (or less) datapoints.

To address this, we conducted a preliminary experiment fine-tuning LLaMA3.1-8B-Instruct on SST2 and SST5 using the same 10 datapoints as our DSM neuron pruning method, employing zero-shot prompting and updating all model parameters. The results show minimal improvements: SST2 accuracy increased from 88.87% to 89.0% (+0.13%), and SST5 from 48.83% to 48.90% (+0.07%), compared to our method’s gains of +3.27% on SST2 (92.14%) and +2.1% on SST5 (50.93%, Table 2). This suggests that supervised fine-tuning with such a small dataset is less effective than our targeted pruning approach. We will include a comprehensive comparison with this fine-tuning baseline across all tasks in the next revision version

Thank you for sharing your valuable feedback. We have addressed each of the points you raised below.

The paper would benefit from a discussion regarding how the choice of the adaptation set affects the accuracy and how robust the method is to this random choice. Specifically it would be interesting to see what is the variance of the accuracies given different random samples of the adaptation set.

To evaluate the robustness of our DSM neuron pruning approach to the choice of the 10-sample adaptation set, we conducted an experiment with LLaMA3.1-8B-Instruct, performing five runs with different random seeds to assess the variance in performance across distinct randomly selected adaptation sets:

The results on BoolQ show a mean accuracy of 82.93% with a standard deviation of ±0.65% for 10 samples, 82.57% ±0.05% for 50 samples, and 82.43% ±0.02% for 100 samples, indicating low variance and high stability even with small, randomly chosen sets. Additionally, we tested the impact of sample correctness by comparing 10 randomly selected samples (82.35% on BoolQ, 92.14% on SST2) against sets of 10 correctly predicted (85.25% on BoolQ, 92.14% on SST2) or wrongly predicted samples (79.81% on BoolQ, 87.5% on SST2).

These results demonstrate that our method is robust to the random selection of adaptation samples, with minimal performance variability, and performs comparably to curated sets. We will incorporate this analysis in the next revision.

Additional details and clarifications on the experiments and adaptation set would also be very helpful, specifically regarding the following questions: 1. Are all the adaptation sets chosen at random? 2. Is the adaptation set specific for each task or is it taken from multiple tasks and the pruning is done once simultaneously for all the tasks in the figure/table? e.g. in Figure 1, Table 1 & 2.

To address your questions regarding the adaptation sets used in our experiments:

1. Random Selection of Adaptation Sets: Yes, all adaptation sets are chosen randomly from the task-specific training data without access to labels, ensuring an unbiased selection process. For each task (e.g., BoolQ, SST2, MMLU), we sample 10 unlabeled examples to compute Integrated Gradients (IG) attributions for DSM neuron detection, as described in Section 3.3. To further validate robustness to this random selection, we conducted an ablation study with five runs using different random seeds for the 10/50/100-sample adaptation sets on BoolQ (see response for the first concern) , this low variance confirms that our method is stable across different random selections of the adaptation set.
2. Task-Specific Adaptation Sets: The adaptation sets are specific to each task, and pruning is performed independently for each task reported in Figure 1 and Tables 1 and 2. For each task (e.g., NLI, MMLU, SST2, BoolQ, SST5, Balanced COPA), we select a unique 10-sample adaptation set from the corresponding training dataset (Table 4) to identify DSM neurons and determine the optimal layer and pruning percentage via grid search. The pruning process is not performed simultaneously across tasks; instead, a separate pruned model is created for each task to ensure task-specific optimization. This approach allows our method to target dataset-specific mechanisms (DSMs) tailored to each task’s characteristics, contributing to the consistent performance improvements observed across all tasks

Appendix D states that it explores the transferability from the adaptation set to the test set. Aren't the results throughout the paper presented for a held out test set? Please clarify the difference between the experiments in the main paper and the results in Appendix D.

To clarify, the main paper presents results on held-out test sets across various benchmarks (e.g., MMLU, SST2, BoolQ) to evaluate the overall effectiveness of our DSM neuron pruning method, where the pruning configuration (layer and percentage) is determined using a small adaptation set and then applied to these test sets. Since the adaptation set is (1) very small, and (2) used both for IG and for the grid search, we verify in Appendix D that the parameters (layer and pruning amount) selected using the small 10-sample are similar to what would have been selected with the entire test set.

Thank you for your insightful feedback. Below, we address each of the concerns you raised.

Incomplete Mechanistic Justification: The focus on the gate_proj layer lacks empirical or theoretical evidence comparing it to other layers, for instance, up_proj, weakening the rationale for layer selection.

How does the choice of the gate_proj layer compare to pruning other layers in terms of performance and computational cost. To address the concern regarding the justification for selecting the gate_proj layer, we conducted an ablation study comparing the performance of pruning neurons in the gate_proj, up_proj, and down_proj layers of the LLaMA3.1-8B-Instruct model. The results, summarized below, demonstrate the effectiveness of our choice:

Llama 3.1 8B Instruct:

The results show that pruning the gate_proj layer achieves the highest performance improvements, with 82.35% on BoolQ (vs. 81.84% for up_proj and 79.53% for down_proj) and 92.14% on SST2 (vs. 91.02% for up_proj and 90.36% for down_proj). This supports our hypothesis that gate_proj plays a critical role in encoding dataset-specific mechanisms (DSM) that can be pruned to enhance generalization.

The computational overhead of IG scales with the number of neurons, making gate_proj and up_proj (11,008 neurons each) slightly more expensive than down_proj (4,096 neurons). However, the superior performance of gate_proj pruning justifies this modest increase in cost, as it maximizes generalization without requiring broader structural changes. Prior work (Ali et al., 2024) [1] successfully used IG to analyze and prune bias neurons in gate_proj layer, and our method builds on this foundation.

[1]: Ali et al. Mitigating Copy Bias in In-Context Learning through Neuron Pruning. 2024 https://arxiv.org/abs/2410.01288

Sample Sensitivity: Reliance on 10 unlabeled samples for DSM detection risks bias if samples are unrepresentative. No analysis of variability across different sample sets is provided. Could the method’s reliance on 10 samples lead to instability if those samples are noisy or biased? Please provide sensitivity analysis across varying sample sets.

To address the potential instability due to reliance on 10 unlabeled samples, we conducted five runs with different random seeds for 10, 50, and 100 samples, averaging performance on BoolQ: 82.93 ± 0.65% (10 samples), 82.57 ± 0.05% (50 samples), and 82.43 ± 0.02% (100 samples). The low variance across runs, particularly with larger sample sizes, demonstrates that our method is stable and minimally affected by sample variability.

Scalability Concerns: The grid-search approach for layer/pruning percentage may not scale efficiently to larger models (e.g., 70B+ parameters).

To address this, we extended our evaluation to the LLaMA3.1-70B-Instruct model, which has significantly more parameters (70B) than the 8B models. Our results below demonstrate that our method remains effective even for larger models, achieving improvements from 93.69% to 94.51% on SST2 and from 89.14% to 91.11% on BoolQ. While the grid search for optimal layer and pruning percentage (5% to 40% in 5% increments) incurs higher computational costs for larger models due to the increased number of neurons, the process is mitigated by our use of only 10 unlabeled samples for DSM neuron detection, keeping the overall computational overhead manageable. We will include additional results on more datasets and tasks in the next revision.

Llama 3.1 70B Instruct:

What modifications would enable scaling the grid-search strategy to models with 100B+ parameters?

Our method can be scaled to 100B+ models through the following lightweight modifications:

1. Parallelism: Grid evaluations are parallel and can be distributed across GPUs.
2. Coarse-to-Fine Search: A two-stage grid search-starting with coarse pruning percentages and refining only promising configurations-greatly reduces compute.
3. Layer Subsampling: Instead of searching all layers, we can pre-select a small subset based on IG statistics or transfer insights from smaller models (e.g., LLaMA 7B → 70B).