Discovering Knowledge-Critical Subnetworks in Pretrained Language Models (2024)

Abstract: Pretrained language models (LMs) encode implicit representations of knowledge in their parameters. However, localizing these representations and disentangling them from each other remains an open problem. In this work, we investigate whether pretrained language models contain various *knowledge-critical* subnetworks: particular sparse computational subgraphs that can, if removed, precisely suppress specific knowledge the model has memorized. We propose a multi-objective differentiable masking scheme that can be applied to both weights and neurons to discover such subnetworks and show that we can use them to precisely remove specific knowledge from models while minimizing adverse effects on the behavior of the original model. We demonstrate our method on multiple GPT2 variants, uncovering highly sparse subnetworks (98%+ sparsity) that are critical for expressing specific collections of relational knowledge. When these subnetworks are removed, the remaining network maintains most of its initial abilities but struggles to represent the suppressed knowledge.

## introduction large-scale language models (llms) encode large amounts of relational knowledge @xcite @xcite , which they transfer to successfully adapt to downstream tasks @xcite . following this success, considerable research focuses on better understanding the extent to which llms capture this knowledge @xcite @xcite . in these works, relational triplets (e.g., (car, isa, vehicle)) are converted to natural language (e.g., "a car is a vehicle.") before being presented to a model. key tokens in these input sequences are 1 the code is made available at https://github.com/ bayazitdeniz/know-subnet controllm controlkg targetkg input examples a cafe is a type of @xcite a cafe is a type of @xcite a cafe is a type of @xcite a cafe is a type of @xcite a cafe is a type of @xcite a car is a @xcite a cafe is a type of @xcite a cafe is a type of @xcite it is closely @xcite pruned model ## background & considerations to find a knowledge-critical subnetwork in a pretrained language model, we learn a differentiable parameter mask ( §4) using a prediction task where the lm is prompted for relational knowledge. prompting lms with kgs we define a global relational knowledge graph (kg) as the set of knowledge triplets, ## methodology this section defines our methodology for discovering knowledge-critical subnetworks using differentiable weight or neuron masking. notation we define a subnetwork as in §3: f (x, m ⊙ θ), where θ is the set of parameters of the network f and m is the mask over a portion of that network's parameters. to learn a mask over neurons, we jointly mask all the weights connecting to the same neuron. we assume a target set of knowledge k t ⊂ k (targetkg) for which we want to identify the critical parameters. ## experimental results we first evaluate the degree to which discovered subnetworks are knowledge-critical. ## conclusion in this paper, we conceptualize knowledge-critical subnetworks, sparse computational subgraphs within larger language models that are responsible for expressing specific knowledge relationships. we discover these subnetworks using a multi-objective differentiable masking approach that jointly optimizes a criterion designed to suppress the expression of target knowledge when knowledge-critical subnetworks are removed from a language model, and maintenance criteria that ensure the language model retains its initial capacity to model other relational knowledge and general language. our results show that when knowledgecritical subnetworks are removed, a model loses its ability to express the knowledge encoded in the subnetwork, and to transfer it when finetuned on downstream tasks requiring the knowledge.