SlowGuess/NAACL_2022_OCR huggingface dataset:
Pretrained masked language models (PLMs) were shown to be inheriting a considerable amount of relational knowledge from the source corpora. In this paper, we present an in-depth and comprehensive study concerning specializing PLMs into relational models from the perspective of network pruning. We show that it is possible to find subnetworks capable of representing grounded commonsense relations at non-trivial sparsity while being more generalizable than original PLMs in scenarios requiring knowledge of single or multiple commonsense relations.
The past few years have witnessed the revolution of NLP methods with the advent of pretrained language models (PLMs) such as BERT (Devlin et al., 2019) and ROBERTA (Liu et al., 2019a). They are first pretrained on vast amount of unlabeled text corpora using masked language modeling (MLM) objective and then fine-tuned on task-specific data, offering a surge of improvements on a wealth of NLP tasks. However, we know very little about what and how much knowledge embedded in PLMs actually contributes to the success. Notable endeavors (Peters et al., 2018; Goldberg, 2019; Tenney et al., 2019) toward this understanding focus on probing linguistic knowledge therein. They demonstrated that pretraining did impart useful linguistic abstraction about syntax and semantics into PLMs.
More recently, several works presented intriguing results examining relational knowledge within PLMs. Relational knowledge (Speer and Havasi, 2012; Vrandecic and Krötzsch, 2014) is typically defined as describing the abstractive relationship between a pair of concepts or entities, which is crucial for facilitating language understanding. Petroni et al. (2020) first posed the LAMA probe,
an English benchmark comprising multiple sets of prompts. Each prompt is a cloze-like sentence transformed from a relational knowledge triple:
Knowledge Triple: <bus, HasA, ?>
Object Label: seats.
Prompt: you are likely to find _ in a bus.
By substituting _ with a special [MASK] token and reusing the MLM head, prompt-based relational knowledge probing provides an estimated lower bound of what PLMs know without training an additional layer which was used in the previous linguistic probes. They showed that, even without grounded supervision, PLMs capture such relational knowledge at a level competitive to supervised alternatives. Subsequent works further showed that some specific prompts, acquired either through heuristical mining (Jiang et al., 2020) or gradient-guided search (Shin et al., 2020), can better trigger the models to correctly predict the missing object.
Despite the mounting evidence for the existence of relational knowledge in PLMs, it remains unclear how such knowledge is represented internally. In light of this, we raise the core question in this paper: Given the general language representation space modeled by a PLM, can we extract its latent representation subspaces for different relations and specialize the PLM into relation-specific knowledge models? These subspaces exclusively represent knowledge inherited from different subsets of MLM data, expressing different relations between masked word and remaining context, thus can potentially benefit applications where knowledge of certain relations are explicitly required.
We study this question by first drawing inspiration from recent findings (Saunshi et al., 2020; Lee et al., 2020; Zhang and Hashimoto, 2021): the more MLM pretraining simulates downstream task, the more successful the transfer will be. For example, filling like or hate into a cloze like "I [MASK] this film, it's great." provide a clear way in which
Findings of the Association for Computational Linguistics: NAACL 2022, pages 2195 - 2207 July 10-15, 2022 ©2022 Association for Computational Linguistics
Figure 1: Querying original/pruned BERT-BASE with prompts of relation HasA. The color spectrum indicates the 12 attention heads in the last layer (Vig, 2019).
the model can implicitly learn to perform sentiment classification. Similarly, we hypothesize that training on MLM data expressing certain relation  $r$  between masked word and remaining context would lead to effective transfer on knowledge probing that targets relation  $r$ .
By exploiting this correlation conversely, we offer a new way for extracting representation subspaces responsible for different relational knowledge based on their transfer performance on knowledge probing task. Instead of introducing additional parametric transformation upon the original space, we restrict these subspaces to have correspondence with subnetworks of PLMs and propose an end-to-end differentiable weights pruning framework. Our experiments show that it is possible to find subnetworks capable of representing grounded commonsense relations at non-trivial sparsity while also exhibiting evident disentanglement. Figure 1 exemplifies a cloze prompt where the identified subnetwork produces the valid answer seats by attending to relevant context, i.e., bus, while the original BERT fails. We then examine the knowledge transfer ability of these subnetworks in scenarios requiring knowledge of single or multiple commonsense relations for reasoning. Experiment on commonsense knowledge base completion show that the identified subnetworks even outperform strong supervised knowledge base completion methods. These subnetworks also outperform the original PLMs in both many-shot and zero-shot commonsense question answering tasks, when combined properly.
Code and all pruned subnetworks are open-sourced at https://github.com/DRSY/
We first provide background on pretrained masked language models and the formulation of cloze prompt for querying these models, then we proceed to elaborate our proposed pruning procedure.
2.1 Pretrained Masked Language Models
Given a sequence of tokens  $\pmb{w} = [w_{1}, w_{2}, \dots, w_{n}]$  with length  $n$ , the model outputs a sequence of hidden states  $\pmb{h} = [h_{1}, h_{2}, \dots, h_{n}]$  corresponding to each token. In standard MLM pretraining,  $h_{i}$  is fed into a MLM head for computing the reconstruction probability  $P(w_{i} | \pmb{w}_{-i})$  of the masked  $i$ -th token  $w_{i}$ , where  $w_{-i}$  are all other unmasked tokens. We denote the pretrained masked language model  $\mathcal{LM}$  with parameter  $\theta$  as  $\mathcal{LM}_{\theta}$  in the following sections.
2.2 Knowledge Probing with Cloze Prompts
The natural language cloze prompts, such as "you are likely to find a basement in below your [MASK]", offer a straightforward means of querying pretrained masked language models that conform to their interfaces.
We follow the formulation of Petroni et al. (2020), where relational knowledge is in the form of triplets  $\langle \text{subj}, r, \text{obj} \rangle$ . Here  $\text{subj}$  refers to the subject,  $\text{obj}$  refers to the object, and  $r$  indicates their corresponding relation. To query a model  $\mathcal{LM}_{\theta}$ , each relation  $r$  is associated with a set of cloze template prompts  $T_r$ , each of which consists of a sequence of tokens, and two of which are place-holders for  $\text{subj}$  and  $\text{obj}$  (e.g., "you are likely to find [subj] in [obj]"). We can check the existence of the knowledge in  $\mathcal{LM}_{\theta}$  by substituting the [subj] place-holder with the real subject and asking the model to predict the missing object:
$$
\hat{o} j = \arg \max_{w\in \mathcal{V}}P_{\mathcal{L}\mathcal{M}_{\theta}}([obj] = w|subj,T_{r})
$$
where  $\mathcal{V}$  is the vocabulary of  $\mathcal{LM}_{\theta}$ . We say that  $\mathcal{LM}_{\theta}$  carries the knowledge if  $o\hat{b} j = o b j$
2.3 Extracting Representation Subspaces by Weights Pruning
Extracting representation subspaces for different roles/functionalities has been explored by prior works, such as attaining disentangled subspaces of style and semantic content in text style transfer task (John et al., 2019). The typical approach is to apply parametric transformation function upon
the original space and optimize through end-to-end downstream fine-tuning. However, it induces additional parameters and cannot faithfully reveal the relational knowledge originally present in PLMs since such knowledge can be stored in the newly introduced parameters outside of PLMs. We circumvent this issue by focusing on the representation spaces modeled by subnetworks of  $\mathcal{LM}_{\theta}$ . A subnetwork  $\mathcal{LM}_{\theta_r}$  of relation  $r$  is obtained by setting certain dimensions of  $\theta$  to zero.
The next step is to identify  $\mathcal{LM}_{\theta_r}$  such that the representation space it corresponds to inherits its knowledge from the MLM data expressing relation  $r$ . Based on the previous evidence (Zhang and Hashimoto, 2021) that shows positive correlation between downstream performance and task similarity with MLM data, we propose to estimate representation space for relation  $r$  by searching for the  $\mathcal{LM}_{\theta_r}$  that is the most predictive of the prompts expressing relation  $r$ . Specifically, for each weight matrix  $W^l$  from the set of all weight matrices  $W^l$  in the  $l$ -th transformer layer, we assign a learnable pruning mask generator  $G_r^l$  that is element-wise initialized from a prior distribution  $\phi(\cdot)$ . Each entry  $g_{i,j}^l \in G_r^l$  is a real-valued scalar that determines whether its corresponding weight  $w_{i,j}^l \in W^l$  should be pruned. We explore two different schemes of converting  $G_r^l$  into a masking matrix  $M_r^l$  from a probabilistic view.
2.3.1 Stochastic Pruning
The first variant is to establish a probabilistic formulation for determining the importance of individual weights. Formally,  $g_{i,j}^{l}$  is taken as the input to a sigmoid function for parametrizing a Bernoulli distribution  $B(\sigma(g_{i,j}^{l}))$ , from which a binary masking random variable  $m_{i,j}^{l}$  is sampled:
$$
m _ {i, j} ^ {l} \sim B \left(\sigma \left(g _ {i, j} ^ {l}\right)\right) \tag {1}
$$
where  $m_{i,j}^{l}\in M_{r}^{l}$ . The resulting masking matrix  $M_r^l$  can then be used to select weights within original linear layer  $W^{l}$  by a Hadamard product:
$$
W _ {r} ^ {l} = W ^ {l} \odot M _ {r} ^ {l} \tag {2}
$$
Due to the non-differentiability introduced by sampling, the gradient w.r.t. loss function (described in Section 2.3.3) cannot be back-propagated to  $g_{i,j}^{l}$ . As a remedy, we use the re-parametrization technique by Li et al. (2018) to approximate  $m_{i,j}^{l}$  with
another differentiable variable  $\tilde{m}_{i,j}^{l}$ :
$$
\tilde {m} _ {i, j} ^ {l} = \sigma \left(\frac {g _ {i , j} ^ {l} + \log U - \log (1 - U)}{\tau}\right) \tag {3}
$$
where  $U\sim \mathit{Uniform}(0,1)$  and  $\tau$  is a small positive temperature parameter. As  $\tau$  approaches zero,  $\tilde{m}_{i,j}^{l}$  will match sampled  $m_{i,j}^{l}$  more accurately (detailed proof can be found in Appendix A).
Consequently, Eq. (2) becomes:
$$
W _ {r} ^ {l} = W ^ {l} \odot \tilde {M} _ {r} ^ {l} \tag {4}
$$
2.3.2 Deterministic Pruning
While our first probabilistic pruning formulation considers flexible weights combination, the second proposed variant utilizes a hard thresholding function to directly generate the masking matrix.
Let  $t$  denote the predefined thresholding hyperparameter ranging from 0 to 1, then we have:
$$
\hat {m} _ {i, j} ^ {l} = \left\{ \begin{array}{l l} 1, & \sigma \left(g _ {i, j} ^ {l}\right) \geq t \\ 0, & \text {o t h e r w i s e ,} \end{array} \right. \tag {5}
$$
where  $\sigma$  is the sigmoid function. Similar to Section 2.3.1, the resulting binary masking matrix  $\hat{M}_r^l$  is then used to select weights relevant to relation  $r$  by a Hadamard product:
$$
W _ {r} ^ {l} = W ^ {l} \odot \hat {M} _ {r} ^ {l} \tag {6}
$$
Note that the hard thresholding operation in Eq. (5) also blocks the gradient propagation to  $g_{i,j}^{l}$ . Here we employ the Straight-Through gradient estimator (Hubara et al., 2016; Zhao et al., 2020) and use  $\frac{\partial\mathcal{L}_r}{\partial\hat{m}_{i,j}^l}$  as a proxy of  $\frac{\partial\mathcal{L}_r}{\partial g_{i,j}^l}$ . We elucidate the loss function  $\mathcal{L}_r$  w.r.t. relation  $r$  in the next section.
2.3.3 Training and Inference
The resultant subnetwork  $\mathcal{LM}_{\theta_r}$  is expected to behave like a specialized neural knowledge base. That is, given a prompt requiring knowledge about relation  $r$ ,  $\mathcal{LM}_{\theta_r}$  should be able to fill in the missing object more accurately than  $\mathcal{LM}_{\theta}$ . Hence the learning objective for pruning mask generator  $\{G_r^l\}_{l_b \leq l \leq l_t}$ , where  $l_b$  and  $l_t$  indicate the range of transformer layers, is to find the subnetwork  $\mathcal{LM}_{\theta_r}$  that minimizes:
$$
\mathcal {L} _ {r} = - \mathbb {E} _ {(s u b j, T _ {r}, o b j) \in D _ {r}} [ \log P _ {\mathcal {L M} _ {\theta_ {r}}} (o b j | s u b j, T _ {r}) ] \tag {7}
$$
where  $D_r$  is the collection of prompts under relation  $r$ . The training procedure is conducted for
each relation  $r\in \mathcal{R}$  of interest. Finally, we obtain a set of trained  $\{\pmb {G}_r\}_{r\in \mathcal{R}}$  for the designated pretrained model  $\mathcal{LM}_{\theta}$
During inference, for deterministic pruning,  $M_r$  is obtained from  $G_r$  by Eq. (5). For stochastic pruning,  $M_r$  is obtained by taking the expectation value (i.e.,  $\sigma(G_r)$ ) of Bernoulli variables.
We present our pruning setup and detailed analysis in Section 3.1. Then we compare the knowledge transfer ability of pruned subnetworks against original PLMs in scenarios requiring knowledge of single or multiple relations for reasoning in Section 3.2 and Section 3.3 respectively.
3.1 Pruning & Analysis
Data Split	# Relations	# Prompts
Train	16	20,841
Validation	16	5,955
Test	16	2,978
Table 1: Statistics of C-LAMA.
Dataset. We use the ConceptNet (Speer and Havasi, 2012) subset of LAMA benchmark for both pruning and evaluation, denoted as C-LAMA. C-LAMA contains commonsense facts from the English part of ConceptNet that has single-token objects covering 16 relations. These facts are extracted from Open Mind Common Sense (OMCS). Since C-LAMA has no official data splits, we construct the train/validation/test splits with a ratio of 7:2:1. Detailed statistics are listed in Table 1.
Models. We experiment with the DistilBERT-base (Sanh, 2019), BERT-base, RoBERTa-base (Liu et al., 2019a), and the more recent MPNet-base (Song et al., 2020) as the choices of  $\mathcal{LM}$ . After pruning, each  $\mathcal{LM}$  will have 16 pruned subnetworks corresponding to the 16 commonsense relations. As a straightforward comparison, for each  $\mathcal{LM}$  we also obtain 16 fine-tuned models corresponding to the same 16 relations. Precision  $\mathrm{P@K}$  averaged across all relations is used to evaluate the prompt filling performance.
Setup. The prior distribution  $\phi (\cdot)$  is a Gaussian  $\mathcal{N}(\mu ,1)$  where  $\mu$  is the mean that controls the initial sparsity of the pruned model (e.g.,  $\mu = 0$  indicates  $50\%$  initial sparsity). We set  $l_{t}$  to be the top layer of a given model and choose  $l_{b}$  from  $\{3,4\}$  for DistilBERT,  $\{6,7,8,9\}$  for BERT, RoBERTa,
Figure 2: Ablation on the pruning masks (left) and effect of initial sparsity and pruned layers (right).
and MPNet. The temperature  $\tau$  is fixed as 0.1. The threshold  $t$  is fixed as 0.5. We use Adam (Kingma and Ba, 2015) with a batch size of 32 and a linear warm-up scheduler with 0.1 warm-up ratio for training the mask up to 6 epochs. The learning rate is fixed as  $3 \times 10^{-4}$ . We run all experiments with three different random seeds and report the averaged results. All experiments are conducted on a GTX 1080 Ti GPU with 11GB RAM.
Factors impacting performance. To investigate how  $\mu$  and  $l_{b}$  affect the performance, we perform a preliminary experiment by applying deterministic pruning on BERT-base with  $l_{b}$  in  $\{6,7,8,9\}$  and the initial sparsity in  $\{50\%, 54\%, 58\%, 62\%\}$ . Figure 2 (right) shows that (i) increasing the number of pruned layers helps distill more knowledge; and (ii) larger initial sparsity is more likely to prune away weights that are important to certain knowledge and cannot be recovered in the later training process. In general, we find an initial sparsity around  $50\%$  yields a decent performance both in probing and downstream applications. We adopt this setting in the rest of this paper unless otherwise stated.
How specialized are these subnetworks? Ideally, specialized subnetworks are expected to perform poorly on relations other than their associated ones. We verify this by instantiating the pruning mask upon BERT-base with a set of mismatched masks. Specifically, we corrupt the correspondence of relation between masks and prompts by shuffling the masks 15 times, as there are 16 relations in total. Then we calculate the micro-averaged  $\mathrm{P@K}$  for each shift and average the results. As indicated by the green curve in the left part of Figure 2, if we apply the mismatched masks from other relations, the  $\mathrm{P@1}$  score significantly drops to 3.8 from 43.8, a performance even worse than the original model. It suggests that the representation spaces for different commonsense relations modeled by these subnetworks are highly disentangled.
Model	P@1 (%)	P@2 (%)	P@3 (%)	Sparsity	\(l_b-l_t\)	#Param.
DistilBERT-base	11.4	16.6	19.9	0%	-	66M
DistilBERT-base w/ fine-tuning	27.9	36.3	41.4	0%	-	66M
DistilBERT-base w/ stochastic pruning	14.8	21.5	26.3	~30%	4-6	66M
DistilBERT-base w/ deterministic pruning	34.0	41.8	46.0	~50%	4-6	56M
BERT-base	12.9	18.4	21.8	0%	-	110M
BERT-base w/ fine-tuning	29.2	37.4	41.3	0%	-	110M
BERT-base w/ stochastic pruning	17.2	25.1	29.6	~30%	7-12	110M
BERT-base w/ deterministic pruning	43.8	49.5	52.2	~50%	7-12	88M
RoBERTa-base	15.4	21.2	24.6	0%	-	125M
RoBERTa-base w fine-tuning	11.7	14.4	16.4	0%	-	125M
RoBERTa-base w/ stochastic pruning	16.6	22.2	25.8	~30%	7-12	125M
RoBERTa-base w/ deterministic pruning	38.3	42.8	44.6	~50%	7-12	100M
MPNet-base	14.8	20.7	24.0	0%	-	110M
MPNet-base w/ fine-tuning	23.8	30.9	36.3	0%	-	110M
MPNet-base w/ stochastic pruning	19.8	27.9	33.2	~30%	7-12	110M
MPNet-base w/ deterministic pruning	47.9	52.8	55.6	~50%	7-12	88M
Table 2: Relational knowledge probing results on C-LAMA test set.
Non-triviality of subnetworks. We also examine the non-triviality of subnetworks by initializing the masks with a Bernoulli distribution  $B(0.5)$  and averaging the results from 5 different random seeds. If we apply such random masks with sparsity comparable to the learned ones, the P@1 drops drastically to 0.4 (red curve in the left part of Figure 2). This notable difference proves that the effective subnetworks cannot be trivially identified through random weights pruning.
Test set results. Table 10 summarizes the test set results. Among all original PLMs, RoBERTa achieves the highest  $\mathrm{P@1}$  score of 15.4 while DistilBERT gets the lowest 11.4. It indicates that while PLMs are shown to be helpful for downstream learning, they cannot accurately complete cloze-like prompts that require commonsense relation knowledge. This observation also coincides with previous finding (Zhang and Hashimoto, 2021) that the uniform masking adopted by PLMs is biased towards extracting statistical dependencies. Comparing the results for each pair of original and pruned models, we consistently observe a significant increase in pruned models, especially for deterministically pruned ones (27.4% on average). The pruned models also surpass their fine-tuned counterparts, which is likely due to that fine-tuning makes aggressive updates to parameters and overfits to the training set by memorizing spurious features.
To sum up, the substantial performance gains provide new evidences for the existence of sparse latent relational knowledge structures in PLMs. These structures are previously weakened by other pretrained weights reserved for more general-purpose use and are exposed by the proposed pruning method. It is worth noting that determinis
tic pruning excels by a big margin compared to stochastic one. It implies that successful extraction of relation-specific representation space relies on ignoring the information in the input that is irrelevant to the relation. Therefore, we focus our analysis on deterministically pruned PLMs and denote them by pruned in the rest of this paper.
Visualization of attention and representations. To explain how the subnetworks accommodate more accurate commonsense knowledge despite having far fewer weights than the full-scale models, we randomly sample several prompts that the subnetworks correctly answered but the full-scale model (BERT-base) failed to, and visualize the attention patterns in the last layer.
Figure 3: Attention weight visualization. At Location/PartOf is required in the left/right column.
Specifically, we focus on the attention weights between [MASK] and other tokens in the prompt. A first glance of change of attention pattern was given earlier in Figure 1, and now we show more examples of other ConcetpNet relations in Figure 3.
Method	Development Set	Test Set
MRR (%)	P@1 (%)	P@2 (%)	P@3 (%)	MRR (%)	P@1 (%)	P@2 (%)	P@3 (%)
Supervised KB completion models
DistMult (Yang et al., 2015)	8.5	4.2	6.6	8.3	10.5	5.4	8.4	10.9
ComplEx (Trouillon et al., 2016)	10.7	6.5	9.0	11.0	13.6	8.2	12.4	15.7
ConvE (Dettmers et al., 2018)	18.9	11.5	16.6	19.0	21.9	13.5	18.9	24.0
TuckER (Balažević et al., 2019)	17.3	10.9	14.8	18.8	21.6	14.0	20.4	24.0
SACN (Shang et al., 2019)	21.2	13.2	19.8	23.2	24.2	14.4	22.1	28.0
InteractE (Vashisth et al., 2020)	19.8	11.2	17.3	21.2	23.3	14.9	21.9	26.5
Unsupervised PLMs
DistilBERT-base	9.0	3.1	6.9	10.3	10.8	5.8	9.6	11.2
BERT-base	12.4	7.2	10.0	13.7	14.3	8.3	13.7	16.6
RoBERTa-base	8.3	4.2	6.0	7.1	9.4	5.1	7.1	9.3
MPNet-base	11.7	7.2	9.4	11.1	11.1	6.0	9.9	11.7
DistilBERT-base (pruned)	24.1	15.8	24.1	26.4	23.4	14.8	22.2	26.5
BERT-base (pruned)	23.7	15.5	22.1	27.0	22.8	14.3	20.9	26.0
RoBERTa-base (pruned)	9.0	4.9	7.1	8.9	9.5	6.1	7.6	11.4
MPNet-base (pruned)	22.1	12.9	21.2	25.5	20.0	11.4	18.8	22.9
Table 3: One-hop link prediction results. Best results are marked with bold font and second best with underline.
We observe that while the original pretrained model tends to attend to special tokens like period and [SEP], the subnetwork successfully grasps the relevant concepts (i.e., apple, worms, and basement) in the prompt hence produces the correct object. We also use t-SNE (van der Maaten and Hinton, 2008) to visualize the last layer's representation of [CLS] for each prompt. From Figure 4, the representations computed by original pretrained model are hardly separable as different types of knowledge are mixed together. In contrast, the pruned subnetwork can extract meaningful and disentangled representations for different commonsense relations.
Figure 4: t-SNE visualization of [CLS] representation from original (left) and pruned (right) BERT-base.
3.2 Single-Relation Scenario
In this section, we compare the transfer ability of pruned subnetworks against original PLMs in a scenario that explicitly requires knowledge of single commonsense relation, i.e., commonsense knowledge base completion (CKBC), which aims at populating a CKB with valid relational triples. We use the ConceptNet-100K benchmark provided by Li et al. (2016). To ensure a fair evaluation, we manually create a subset of ConceptNet-100K consisting of triples with single-token subject/object. We also
ensure that its dev/test set has no overlap with C-LAMA. Each relation is associated with a sentence template (provided in Appendix B) of which the wording is distinct from those in C-LAMA. The resulting dataset contains 17,891 training instances, 349 development instances, and 446 test instances.
One-hop link prediction. We first formulate CKBC as a link prediction task, i.e., predicting the missing object given the subject and relation. We regard the pruned subnetworks and the original PLMs as the unsupervised off-the-shelf neural knowledge base and include the results of several strong KB completion methods for reference.