Low-Complexity Probing via Finding Subnetworks (2021)

Abstract: The dominant approach in probing neural networks for linguistic properties is to train a new shallow multi-layer perceptron (MLP) on top of the model’s internal representations. This approach can detect properties encoded in the model, but at the cost of adding new parameters that may learn the task directly. We instead propose a subtractive pruning-based probe, where we find an existing subnetwork that performs the linguistic task of interest. Compared to an MLP, the subnetwork probe achieves both higher accuracy on pre-trained models and lower accuracy on random models, so it is both better at finding properties of interest and worse at learning on its own. Next, by varying the complexity of each probe, we show that subnetwork probing Pareto-dominates MLP probing in that it achieves higher accuracy given any budget of probe complexity. Finally, we analyze the resulting subnetworks across various tasks to locate where each task is encoded, and we find that lower-level tasks are captured in lower layers, reproducing similar findings in past work.

## introduction while pre-training has produced large gains for natural language tasks, it is unclear what a model learns during pre-training. research in probing investigates this question by training a shallow classifier on top of the pre-trained model's internal representations to predict some linguistic property @xcite @xcite . the resulting accuracy is then roughly indicative of the model encoding that property. however, it is unclear how much is learned by the probe versus already captured in the model representations. this question has been the subject of much recent debate @xcite @xcite . the code is available at https://github.com/ stevenxcao/subnetwork-probing . we would like the probe to find only and all properties captured by a model, leading to a tradeoff between accuracy and complexity: a linear probe is insufficient to find the non-linear patterns in neural models, but a deeper multi-layer perceptron (mlp) is complex enough to learn the task on its own. motivated by this tradeoff and the goal of lowcomplexity probes, we consider a different approach based on pruning. specifically, we search for a subnetwork -a version of the model with a subset of the weights set to zero -that performs the task of interest. as our search procedure, we build upon past work in pruning and perform gradient descent on a continuous relaxation of the search problem @xcite @xcite . the resulting probe has many fewer free parameters than mlp probes. our experiments evaluate the accuracycomplexity tradeoff compared to mlp probes on an array of linguistic tasks. first, we find that the neuron subnetwork probe has both higher accuracy on pre-trained models and lower accuracy on random models, so it is both better at finding properties of interest and less able to learn the tasks on its own. next, we measure complexity as the bits needed to transmit the probe parameters @xcite . varying the complexity of each probe, we find that subnetwork probing pareto-dominates mlp probing in that it achieves higher accuracy given any desired complexity. finally, we analyze the resulting subnetworks across various tasks and find that lower-level tasks are captured in lower layers, reproducing similar findings in past work @xcite . these results suggest that subnetwork probing is an effective new direction for improving our understanding of pre-training. however, these results are difficult to interpret due to the use of a learned classifier. one line of work suggests comparing the probe accuracy to random baselines, e.g. random models @xcite or random control tasks @xcite . other works take an informationtheoretic view: voita and titov (2020) measure the complexity of the probe in terms of the bits needed to transmit its parameters, while @xcite argue that probing should measure mutual information between the representation and the property. @xcite propose a pareto approach where they plot accuracy versus probe complexity, unifying several of these goals. we use these proposed metrics to compare our probing method to standard probing approaches. subnetworks. while pruning is widely used for model compression, some works have explored pruning as a technique for learning as well. @xcite found that a model trained on im-agenet could be used for new tasks by learning a binary mask over the weights. more recently, @xcite and @xcite showed the analogous result in nlp that weight pruning can be used as an alternative to finetuning for pre-trained models. our paper seeks to use pruning to reveal what the model already captures, rather than learn new tasks. ## subnetwork probing given a task and a pre-trained encoder model with a classification head, our goal is to find a subnetwork with high accuracy on that task, where a subnetwork is the model with a subset of the encoder weights masked, i.e. set to zero. we search for this subnetwork via supervised gradient descent on the head and a continuous relaxation of the mask. we also mask at several levels of granularity, including pruning weights, neurons, or layers. to learn the masks, we follow @xcite . letting φ ∈ r d denote the model weights, we associate the ith weight φ i with a real-valued parameter θ i , which parameterizes a random variable z @xmath0, 1] representing the mask. z i follows the hard concrete distribution hardconcrete(β, θ i ) with temperature β and location θ i , we will denote the mask as z @xmath1, θ i ) and the masked weights as φ * z. we can then optimize the mask parameters θ via gradient descent. specifically, let f (x; φ) denote the model. then, given a data point (x, y) and a loss function l, we can minimize the expectation of the loss, or we estimate the expectation via sampling: we sample a single u and take the gradient ∇ θ l(f (x; φ * z(u, θ)). to encourage sparsity, we penalize the mask based on the probability it is non-zero, or letting λ denote regularization strength, our objective becomes 1 |d| (x,y)∈d l(x, y, θ) + λr(θ). foot_1 ## probe evaluation to evaluate the accuracy-complexity tradeoff of a probe, we adapt methodology from recent work. first, we consider the non-parametric test of probing a random model @xcite . we check probe accuracy on the pre-trained model, the model with the encoder randomly reset (reset encoder), and the model with the encoder and embeddings reset (reset all). an ideal probe should achieve high accuracy on the pre-trained model and low accuracy on the reset models. 3 next, we consider a parametric test based on probe complexity. we first vary the complexity of each probe, where for subnetwork probing we associate multiple encoder weights with a single mask, 4 and for the mlp probe we restrict the rank of the hidden layer. we then plot the resulting accuracy-complexity curve @xcite . to plot this curve, we need a measure of complexity that can compare probes of different types. therefore, we measure complexity as the number of bits needed to transmit the probe parameters @xcite , where for simplicity we use a uniform encoding. in the subnetwork case, this encoding corresponds to using a single bit for each mask parameter. in the case of an mlp probe, each parameter is a real number, so the number of bits per parameter depends on its range and precision. for example, if each parameter lies in [a, b] and requires precision, then we need log( b-a ) bits per parameter. to avoid having the choice of precision impact results, we plot lower and upper bounds of 1 and 32 bits per parameter. ## experimental setup we probe bert-base-uncased @xcite for the following tasks: our primary probing baseline is the mlp probe with one hidden layer (mlp-1): with u, @xmath2. the choice of r restricts the rank of the hidden layer and thus its complexity. 5 then, if g(x; φ) is our pre-trained encoder and cls is the classification head, our two probes are f subnetwork (x) = cls(g(x; φ * z)) and while we vary the complexity of each probe to produce the accuracy-complexity plot, we default to neuron subnetwork probing and full rank mlp-1 probing in all other experiments. ## results accuracy-complexity tradeoff. figure 1 plots the results from the parametric experiments, where we vary the complexity of each probe, apply it to the pre-trained model, and plot the resulting accuracy-complexity curve. we find that the subnetwork probe pareto-dominates the mlp-1 probe in that it achieves higher accuracy for any complexity, even if we assume an overly optimistic mlp-1 lower bound of 1 bit per parameter. in particular, for part-of-speech and dependency parsing, the subnetwork probe achieves high accuracy even when given only 72 bits, while the mlp-1 probe falls off heavily at ∼20k bits. ## conclusion together, these results show that subnetwork probing is more faithful to the model and offers richer analysis than existing probing approaches. while this work explores accuracy and location-based analysis, there are other possible directions, e.g., applying neuron explainability techniques. therefore, we see subnetwork probing as a fruitful new direction for understanding pre-training.