Low-Complexity Probing via Finding Subnetworks

TL;DR

The paper addresses how to probe linguistic properties learned by pretrained models without letting the probe itself learn the task. It introduces subnetwork probing, which searches for a sparse subnetwork within the encoder using gradient-based continuous relaxation of masks (HardConcrete) to perform targeted linguistic tasks, and compares it to traditional MLP probes. Results show that subnetworks achieve higher accuracy on pretrained models and are less capable on random models, with Pareto-dominance across complexity budgets; analysis also reveals that lower-level tasks are encoded in lower layers. This approach provides a more faithful, low-complexity probe and enables richer insights into where linguistic information is stored in pretrained transformers.

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.

Figures (2)

• Figure 1: Subnetwork probe and MLP-1 probe accuracy on the pre-trained model plotted versus probe complexity, measured in $\ln(\text{bits})$. For the MLP-1 probe, we plot lower and upper bounds on complexity of $1$ and $32$ bits per parameter. The subnetwork probe Pareto-dominates the MLP-1 probe in that it achieves higher accuracy for any desired complexity, even if we assume the optimistic lower bound on MLP-1 complexity of $1$ bit per parameter.
• Figure 2: The percentage of non-zero weights in each layer for subnetworks of the pre-trained model and the reset encoder model. While the reset encoder model's subnetworks are distributed uniformly across the layers, the pre-trained model's subnetworks are localized, with the order part-of-speech $\to$ dependencies $\to$ NER.