Source-linked AI summary
Low-Complexity Probing via Finding Subnetworks
Steven Cao, Victor Sanh, Alexander M. Rush
TL;DR
Probing with learned classifiers leaves unclear how much task performance reflects information in the model representations versus what the probe learns. This paper searches for task-performing subnetworks through pruning, and reports that these probes outperform MLP probes on accuracy-complexity tradeoffs while better distinguishing pre-trained from random models.
Problem
Learned probes make it unclear how much task performance reflects properties encoded in representations versus what the probe learns.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
Subnetwork probing offers a way to assess what a model captures while also enabling location-based analysis of resulting subnetworks.
Source: §6 · p. 4, §7 · p. 5The reset-encoder model retains non-contextual word-embedding information, so it may have better probe accuracy on tasks based mainly on word type.
Source: §5 · p. 3
Abstract
from arXiv · showhide
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.
1 Introduction
Probing uses a shallow classifier on a model’s internal representations to test whether linguistic properties are encoded, but it remains unclear how much the classifier learns itself. The paper proposes searching for a pruned subnetwork as a lower-complexity probe and reports stronger accuracy-complexity results than MLP probes.
1 Introduction
- A shallow classifier trained on internal representations is used to test whether a model encodes a linguistic property.Its resulting accuracy is roughly indicative of whether that property is represented.Source: §1 · p. 1
- It remains unclear how much of a probe’s success comes from the learned classifier versus what the model representations already capture.This distinction has been the subject of recent debate.Source: §1 · p. 1
- Linear probes may miss nonlinear patterns, while deeper MLPs may learn the task themselves, creating an accuracy-complexity tradeoff.The stated goal is to find only the properties captured by the model.Source: §1 · p. 1
- The proposed probe searches for a task-performing subnetwork by setting a subset of model weights to zero.The search uses gradient descent on a continuous relaxation and yields fewer free parameters than MLP probes.Source: §1 · p. 1
- Subnetwork probes Pareto-dominate MLP probes, achieving higher accuracy at any desired probe complexity across the evaluated linguistic tasks.They also achieve higher accuracy on pre-trained models and lower accuracy on random models.Source: §1 · p. 1
2 Related Work
Prior probing work uses learned classifiers to study linguistic properties in model representations, including where tasks are captured across BERT’s layers. Because learned probes complicate interpretation, subsequent work has proposed random baselines and complexity-based evaluation, which this paper applies to pruning-based probes.
2 Related Work
- Probing commonly trains a shallow classifier on frozen representations, and BERT studies found lower-level tasks captured in lower layers, mirroring the classical NLP pipeline.The reported task sequence includes part-of-speech, parsing, named entity recognition, semantic roles, and coreference.Source: §2 · p. 2
- Learned classifiers make probing results difficult to interpret, motivating random-model or random-task baselines and information-theoretic measures of probe complexity.Prior proposals include mutual information and accuracy-versus-complexity Pareto curves.Source: §2 · p. 2
- Pruning has been used to learn new tasks through weight masks, but this paper uses pruning to reveal what a pre-trained model already captures.Prior work reported this approach for ImageNet and, in NLP, as an alternative to fine-tuning.Source: §2 · p. 2
3 Subnetwork Probing
The method finds a task-performing subnetwork by learning masks over a pre-trained encoder’s weights. It optimizes task loss while encouraging sparse masks, using a continuous relaxation that permits gradient-based search.
3 Subnetwork Probing
- A subnetwork is found by masking a subset of encoder weights, with masks learned by supervised gradient descent using a continuous relaxation.The approach can mask weights, neurons, or layers.Source: §3 · p. 2
- Each model weight is associated with a parameter that defines a random mask value following a hard-concrete distribution.The distribution has temperature and location parameters.Source: §3 · p. 2
- The hard-concrete mask is a soft Bernoulli-like variable stretched beyond [0, 1] and clamped back to [0, 1].The passage gives γ = −0.1 and ζ = 1.1 for the stretch interval.Source: §3 · p. 2
- Mask parameters are optimized to minimize expected task loss over sampled masks.The expectation is estimated by sampling U and taking a gradient; sparsity is encouraged with a penalty based on the probability a mask is nonzero.Source: §3 · p. 2
- The objective adds a regularization term to the data loss, with λ denoting regularization strength.The regularizer is based on expected mask sparsity.Source: §3 · p. 2
4 Probe Evaluation
Probe evaluation compares accuracy on pre-trained and reset models, then examines accuracy-complexity curves across probe types. Complexity is measured in bits needed to transmit probe parameters, with different encodings for subnetworks and MLPs.
4 Probe Evaluation
- The regularization strength λ is scheduled to stay at zero for the first 25% of training, rise linearly over the next 50%, and then remain fixed.The evaluation uses λmax = 1.Source: §4 · p. 2
- The random-model test compares probe accuracy on the pre-trained model, a reset encoder, and a model with both encoder and embeddings reset.An ideal probe has high pre-trained-model accuracy and low accuracy on reset models.Source: §4 · p. 3
- The parametric test varies probe complexity and plots accuracy-complexity curves for subnetworks and MLPs.Subnetwork complexity is varied by sharing masks across encoder weights, while MLP complexity is varied by restricting hidden-layer rank.Source: §4 · p. 3
- Probe complexity is measured as the bits needed to transmit probe parameters, using one bit per subnetwork mask parameter and a range-and-precision-based count for MLP parameters.The evaluation plots MLP bounds of 1 and 32 bits per parameter to avoid dependence on a single precision choice.Source: §4 · p. 3
5 Experimental Setup
Experiments use bert-base-uncased on part-of-speech tagging, dependency parsing, and named entity recognition, comparing neuron subnetworks with an MLP-1 probe across model states and probe complexities.
5 Experimental Setup
- The study probes bert-base-uncased on part-of-speech tagging, dependency parsing, and named entity recognition using task-specific datasets and classification heads.Part-of-speech tagging and NER use a linear classifier; dependency parsing uses a biaffine head and reports macro-averaged labeled attachment score.Source: §5 · p. 3
- Table 1 compares probe accuracy on the pretrained model, an encoder-reset model retaining embeddings, and a fully reset model, with fine-tuning as a reference.The reset encoder retains non-contextual word-embedding information but no context modeling.Source: Table 1 · p. 3, §5 · p. 3
- The primary baseline is MLP-1, while the subnetwork probe masks encoder weights; other experiments default to neuron masks and full-rank MLP-1.The probes apply the classification head to either the masked encoder or an MLP-1 applied to encoder representations.Source: §5 · p. 3
- Figure 1 plots pretrained-model subnetwork and MLP-1 probe accuracy versus complexity in ln(bits); the subnetwork achieves higher accuracy at any desired complexity, even against MLP-1’s optimistic 1-bit-per-parameter lower bound (upper bound: 32).The MLP-1 rank settings were 1, 2, 5, 10, 25, 50, 125, 250, and 768.Source: §5 · p. 3, Figure 1 · p. 4
- The layerwise analysis compares the percentage of nonzero weights in pretrained-model and reset-encoder subnetworks.The figure caption identifies the compared model conditions and layerwise measure.Source: Figure 2 · p. 4
6 Results
Subnetwork probes outperform MLP-1 on pretrained models and show lower accuracy on random models for dependency parsing and NER. They also dominate across complexity budgets, while the retained weights localize by task in pretrained models.
Accuracy-Complexity Tradeoff
- Subnetwork probes achieve much higher accuracy than MLP-1 across tasks on pretrained models and lower accuracy on random models for dependency parsing and NER.The authors interpret this pattern as indicating that subnetworks find properties when present and are less able to learn tasks on their own.Source: §6 · p. 4
- Subnetwork probing Pareto-dominates MLP-1, achieving higher accuracy at any complexity, including against an optimistic 1-bit-per-parameter MLP-1 bound.For part-of-speech and dependency parsing, subnetworks retain high accuracy at 72 bits while MLP-1 accuracy falls heavily around 20K bits.Source: §6 · p. 4
Subnetwork Analysis
- Pretrained subnetworks localize by task in the order part-of-speech, dependencies, then NER, whereas reset-encoder subnetworks are distributed uniformly across layers.This reproduces Tenney et al. (2019)'s ordering; here layer location is identified directly through pruning rather than classifiers trained at each layer.Source: §6 · p. 4
7 Conclusion
The authors conclude that subnetwork probing is more faithful to the model and supports richer analysis than existing probing approaches, while identifying further analysis directions.
7 Conclusion
- Subnetwork probing offers a more faithful view of model properties and richer analysis than existing probing approaches.The paper explores accuracy and location-based analysis and points to neuron explainability techniques as another possible direction.Source: §7 · p. 5
8 Ethical Considerations
The paper notes that pretrained models can exhibit biases present in their pretraining corpora, creating a risk of reinforcing social biases when deployed.
8 Ethical Considerations
- The authors suggest applying subnetwork analysis to bias-detection tasks such as CrowS-Pairs or StereoSet, although this paper examines linguistic tasks.They frame probing as a tool to understand and potentially mitigate biases in pretrained models.Source: §8 · p. 5
A.1 Hyperparameters
This appendix specifies training and reproducibility settings and reports that the subnetwork probe is generally more selective than MLP-1 across regularization strengths, except at λmax = 125.
A.1 Hyperparameters
- The mask, classification head, MLP-1, and fine-tuning baselines were trained for 30 epochs with Adam and linear warmup, using learning rates of 0.2 for masks and 5×10−5 for the others.Warmup covered the first 10% of the data; all optimizer settings are specified in the passage.Source: §A.1 · p. 7
A.2 Varying Regularization Strength
- The subnetwork probe is consistently more selective than MLP-1 across the tested λmax values, except at λmax = 125, which appears to impose too much sparsity.Table 2 reports subnetwork probing accuracies while regularization strength varies.Source: §A.2 · p. 7, Table 2 · p. 7
A.3 Reproducibility Checklist
- Experiments used BERT-base-uncased and English Universal Dependencies and CoNLL 2003 data, with one 12GB NVIDIA Tesla K80 GPU.The data passage specifies the splits and example counts; a fine-tuning run took about half an hour per task.Source: §A.3 · p. 7