Paper: Discovering Knowledge-Critical Subnetworks in Neural Language Models (ICLR-2024)

Reviewer comment (Evaluation):
The effectiveness of knowledge removal is not quite clear due to limited metrics: the paper mainly uses "perplexity increase on verbalized triplet prompts" as a measure of knowledge removal but does not provide a solid interpretation of this metric. For example, one does not know how much perplexity increase corresponds to a complete (or near-complete) removal of the knowledge. Therefore, it may be hard to evaluate whether knowledge is truly removed with this metric alone. Also, there is a possibility that the perplexity drop is specific to the verbalizing template, so it may be helpful to evaluate on different templates as well.

Author response:
> The effectiveness of knowledge removal: Due to limited metrics, it is not clear whether knowledge removal is effective. Perplexity on verbalized triplet prompts alone is not enough.

As mentioned in the general message to all reviewers, the reviewers wondered whether the rank of the gold tail entity can be used as an alternative metric. We also wondered about this and we report the rank differences between the remaining model and the original pretrained model in **Appendix D**, Table 10 (originally Table 11). These results align with the suppression and maintenance criteria described in Section 4.1. We added a forward pointer to these results in the main body in **Section 5** - “Success Metrics” paragraph.

> The effectiveness of knowledge removal: Could the perplexity drop be specific to the verbalization template?

We provide performance on paraphrases in **Appendix H** (originally Appendix G) for three WordNet hypernym TargetKGs. We have found that the average perplexity differences are high on paraphrases of TargetKG, and the perplexity differences for ControlKG either are near 0 or are negative. We added a forward pointer to these results in our main body in Section 6.3.

Reviewer comment (Evaluation):
The effectiveness of knowledge preservation is also unclear due to possible overfit: the loss function is designed to preserve performance on ControlKG and ControlLM, and results show a negligible decrease in performance on them. However, it is possible that besides removing knowledge in TargetKG, the pruned model also sacrifices some other knowledge that is not in ControlKG and ControlLM. This would go against the goal of finding subnetworks specific to TargetKG. To rule out this possibility, the pruned model should be evaluated on a separate KG and a different corpus than those used in the loss function to make sure that the maintenance criterion is not overfitted to ControlKG and ControlLM.

Author response:
> The effectiveness of knowledge preservation: Besides removing knowledge in TargetKG, the pruned model could sacrifice some other knowledge that is not in ControlKG and ControlLM. To make sure that the maintenance criterion is not overfitted to ControlKG and ControlLM the authors should use a separate corpus for ControlKG and ControlLM.

As stated at the end of Section 5, paragraph “Datasets”, all results on ControlKG and ControlLM are on a held-out test set.

Reviewer comment (Experiments):
Lack of analysis on the discovered knowledge subnetworks with respect to trading off suppression and maintenance: for two competing goals like suppression and maintenance, there is usually a tradeoff rather than a single best solution. It may be helpful to show this tradeoff by varying the weights in Equation 6, which could also give justification for the choice of the weights.

Author response:
> Lack of analysis on trading-off suppression and maintenance: For two competing goals like suppression and maintenance, there is usually a tradeoff rather than a single best solution. It may be helpful to show this tradeoff by varying the weights in Equation 6, and it could also justify the choice of the weights.

To address this concern, we run a minimal experiment on giving importance to one goal at a time for two KGs and one seed. Specifically when we set one of the $\lambda$ weights in Eq. 6 to a value of 3, we set the value of the rest to 1. We find that giving more weight to the Suppresion loss indeed finds checkpoints with higher perplexity differences on TargetKG while simultaneously satisfying the maintenance criteria. Giving more weight to the sparsity regularization on the other hand, ensures a high sparsity.

This fits with the setting we use in all of our experiment, in which we give a larger weight to suppression ($\lambda_1=1.5$) and sparsity ($\lambda_4=3$) while the rest stay the same ($\lambda_{2,3}=1$).

| $\lambda_1$ | $\lambda_2$ | $\lambda_3$ | $\lambda_4$ | Sparsity (%) | TargetKG $\Delta$ PPL | ControlKG $\Delta$ PPL | ControlLM $\Delta$ PPL | # of Valid Checkpoints |
|---|---|---|---|---|---|---|---|---|
| 1 | 1 | 1 | 3 | 99.5 [99.4, 99.5] | 350.2 [43.0, 657.4] | 0.1 [-0.1, 0.2] | 0.2 [0.2, 0.2] | 83.0 [7, 159] |
| 1 | 1 | 3 | 1 | 96.9 [95.6, 98.2] | 17.3 [-7.8, 42.4] | 1.0 [-2.0, 4.1] | 0.5 [0.2, 0.8] | 2.5 [0, 5] |
| 1 | 3 | 1 | 1 | 97.3 [96.4, 98.3] | 80.9 [59.2, 102.6] | -1.5 [-8.9, 5.8] | 1.0 [0.5, 1.5] | 53.5 [5, 102] |
| 3 | 1 | 1 | 1 | 98.5 [97.6, 99.3] | 12516.4 [682.2, 24350.6] | 0.0 [-1.4, 1.5] | 0.4 [0.2, 0.6] | 145.0 [133, 157] |

Reviewer comment (Theory):
Clarification is needed on how the expression criteria in Section 6.2 differ from the suppression criteria; it seems that Equation 7 is (approximately) just the reverse of Equation 3.

Author response:
> How do the expression criteria in Section 6.2 (now in Appendix G) differ from the suppression criteria? It seems that Equation 7 is (approximately) just the reverse of Equation 3.

Eq. 3 and 7 are not approximately opposite equations. Eq. 3 suppresses the expression of the tail token given the head and relation tokens from the remaining model by optimizing for  a uniform prediction distribution. Eq. 7, on the other hand, minimizes the tail token cross-entropy loss on the removed subnetwork. The point we wanted to make with this experiment is that one objective does not guarantee the other, and in particular, the Eq 7. or in other words, the Expression loss, combined with the other objectives we proposed will lead to larger knowledge-critical subnetworks, which does not as fulfill the sparsity criterion.

Reviewer comment (Experiments):
Can you provide more details about the filtering processes applied to the sampled connected KGs? How did these processes ensure the quality and balance of the sampled graphs?

Author response:
We describe the filtering processes applied to the sampled connected KGs in **Appendix A**. In summary, once a small connected KG is sampled, we apply two filtering processes. The first one enforces many-to-one relationships in the $K_T$ graph to avoid head entities with multiple tails. This ensures that the dataset is not dominated by certain head entities and balances the graph. The second filtering process reduces the tail-entity imbalance to avoid over-fitting to a small set of tokens. For this, we count the frequency of the tail tokens in the sampled graph and keep at most a quartile amount of triplets with shared tail entities.

Reviewer comment (Evaluation):
The authors are showing the sparsity level over the masked parameters. Therefore 98% sparsity does not mean 98% of parameters in LLMs are not used. It would be more clear if the authors display the real sparsity levels.

Author response:
Indeed, 98% sparsity means the sparsity of the removed subnetwork mask over the masked layers (in this case, the upper-half layers of the model). To make this more intuitive, we suggest reporting the **density of the removed subnetwork across all parameters** in the network, including those that are not masked. As it would affect the phrasing in many parts of the paper, we suggest making this revision for camera-ready.

Reviewer comment (Evaluation):
The authors only use PPL to measure the effectiveness of their method. However, the ppl values can be a bit confusing as its range can be very big. Would it be more sensible to use ranking metrics? For example, the rank of the target entity token?

Author response:
As mentioned in the general message to all reviewers, the reviewers wondered whether the rank of the gold tail entity can be used as an alternative metric. We also wondered about this, and so in the submission, we report the rank differences between the remaining model and the original pretrained model in **Appendix D**, Table 10 (originally Table 11). These results align with the suppression and maintenance criteria described in Section 4.1. We added a forward pointer to these results in the main body in **Section 5** - “Success Metrics” paragraph.

We chose to display perplexity differences in the main body instead of rank as they better capture model certainty on the gold-tail token. For example, there could be cases where the remaining model has a larger rank than the original model leading to a large difference, but the probability differences between the two can be very small, and therefore not enough information about model certainty is reflected. Particularly, in our setting, where we optimize for a small difference between the TargetKG tail-token prediction distribution and a uniform distribution over all tokens in the vocabulary, we are bound to have multitudes of tokens being considered equally probable. This makes the rank metric less interpretable than perplexity. However, we agree that both metrics provide complementary views on this question, which is why we reported rank in the appendix initially.

Reviewer comment (Experiments):
There is a lack of ablation on sparsity levels. Does varying the sparsity level from 98% to 20% or 60% change the conclusion?

Author response:
The reviewer brings an interesting perspective from the point of view of compression. Differentiable-weight masking is popular as a compression technique (Sanh et al 2020, Zhao et al 2020). Varying the sparsity level would require implementing a different type of regularization which could be possible. However, we instead designed an aggressive sparsity regularization as our goal in finding a knowledge-critical subnetwork is to find the minimum compression ratio to achieve semantic suppression of the knowledge. As a result, while an interesting research direction on its own, we think this analysis may not align with the core inquiry of our paper.

Reviewer comment (Experiments):
Pruning methods based on hard concrete are usually sensitive to hyper-parameters. Do you have any hyper-parameter sensitivity analysis?

Author response:
The reviewer is concerned about the sensitivity of hard-concrete hyperparameters. Past studies have shown the mask initialization hyperparameter tends to be sensitive (Zhao et al. 2020). We initially had run a small study to decide on the initialization. We re-ran this study with 2 KGs (location, communication) for the same seed (735) for 4 initializations (0.25, 0.45, 0.5, 0.75).  

We find that when the masking probability, or in other words, the probability of being included in the critical subnetwork is low – such as 0.25 – the hyperparameter often leads to worse maintenance on ControlKG and lower suppression of TargetKG. 
On the other hand, a high initialization such as 0.75 can be a hit or miss for TargetKG suppression (as seen by the minimum and maximum values shown in the Table below), with overfitting on the maintenance datasets (i.e. decrease of perplexity in predicting ControlKG triplet tails compared to original model). 
Between 0.5 and 0.45, where the latter is the initialization we used in our experiments, we don’t find a significant difference other than a slightly higher suppression effect for 0.45.

Therefore our choice for mask initialization seems reasonable.

| Initial Mask Prob | Sparsity (%) | TargetKG $\Delta$ PPL | ControlKG $\Delta$ PPL | ControlLM $\Delta$ PPL |
|---|---|---|---|---|
| 0.25 | 99.5 [99.5, 99.6] | 332.1 [119.5, 544.6] | 3.7 [2.6, 4.7] | 0.1 [0.1, 0.1] |
| 0.45 | 99.5 [99.5, 99.5] | 1287.8 [80.0, 2495.6] | 0.9 [-1.1, 2.9] | 0.2 [0.1, 0.2] |
| 0.5 | 99.4 [99.4, 99.5] | 939.1 [59.9, 1818.2] | -0.2 [-0.2, -0.1] | 0.2 [0.2, 0.2] |
| 0.75 | 98.9 [98.8, 99.1] | 13043.1 [4.6, 26081.5] | -1.9 [-1.9, -1.8] | 0.3 [0.3, 0.4] |

Reviewer comment (Reproducibility):
What are the computational infrastructure for your experiments? If the readers want to reproduce your results, how many GPUs do they need?

Author response:
All details on the computational infrastructure are in **Appendix B**, in particular in the “Hyperparameters” paragraph, where we define how many steps are required for how many A100 40GB GPUs. For all experiments, one can use a single A100 40GB GPU. However, we chose to speed up the process for the larger models (GPT2-medium and GPT2-large) by using data parallelism across 3 GPUs.

Reviewer comment (Experiments):
Does combining two knowledge-critical subnetworks lead to suppression of both pieces of knowledge?

Author response:
As mentioned in the general message to all reviewers, to address this inquiry, we created a new section in **Appendix J** called “Knowledge-Based Analysis”. We calculated the Intersection-over-Union of the subnetworks across different layers and layer types for the same seed but different KGs (Fig. 7) and composed subnetworks for the same seed and different KGs (Table 18). 

Using the union of the subnetworks increased the suppression effect across all the TargetKGs on average, but it also reduced the maintenance effect. Therefore, it may be possible to naively combine subnetworks across seeds or KGs, however, they may not guarantee the maintenance criteria to the same extent.

Reviewer comment (Evaluation):
Does predicting a missing entity fully represent this knowledge triplet? Even if it can correctly predict the missing entity, the prediction might be only based on the (subject, object) pair instead of based on the specific relation. In general, a knowledge triplet can be rephrased in multiple ways, e.g., swapping the order of subject and object, or missing relation prediction [2]. Can the proposed method deal with the various rephrasings of a certain knowledge?

Author response:
We provide performance on paraphrases in **Appendix H** (originally Appendix G) for WordNet three hypernym TargetKGs. We have found that the average perplexity differences are high on paraphrases of TargetKG, and the perplexity differences for ControlKG either are near 0 or are negative. To clarify the takeaway, we have rephrased it and the changes can be seen in purple. We also added a forward pointer to these results in our main body in Section 6.3, as it is also a robustness analysis.