An information theoretic view on selecting linguistic probes

There is increasing interest in assessing the linguistic knowledge encoded in neural representations. A popular approach is to attach a diagnostic classifier -- or"probe"-- to perform supervised classification from internal representations. However, how to select a good probe is in debate. Hewitt and Liang (2019) showed that a high performance on diagnostic classification itself is insufficient, because it can be attributed to either"the representation being rich in knowledge", or"the probe learning the task", which Pimentel et al. (2020) challenged. We show this dichotomy is valid information-theoretically. In addition, we find that the methods to construct and select good probes proposed by the two papers, *control task* (Hewitt and Liang, 2019) and *control function* (Pimentel et al., 2020), are equivalent -- the errors of their approaches are identical (modulo irrelevant terms). Empirically, these two selection criteria lead to results that highly agree with each other.


Introduction
Recently, neural networks have shown substantial progress in NLP tasks (Devlin et al., 2019;Radford et al., 2019). To understand and explain their behavior, a natural question emerge: how much linguistic knowledge is encoded in these neural network systems?
An efficient approach to reveal information encoded in internal representations uses diagnostic classifiers (Alain and Bengio, 2017). Referred to as "probes", diagnostic classifiers are trained on pre-computed intermediate representations of neural NLP systems. The performance on tasks they are trained to predict are used to evaluate the richness of the linguistic representation in encoding the probed tasks. Such tasks include probing syntax (Hewitt and Manning, 2019;Lin et al., 2019;Tenney et al., 2019a), semantics (Yaghoobzadeh et al., 2019), discourse features (Chen et al., 2019;Liu et al., 2019;Tenney et al., 2019b), and commonsense knowledge (Petroni et al., 2019;Poerner et al., 2019).
However, appropriate criteria for selecting a good probe is under debate. The traditional view that high-accuracy probes are better is challenged by Hewitt and Liang (2019), who proposed that the high accuracy could be attributed to either (1) that the representation contains rich linguistic knowledge, or (2) that the probe learns the task. To circumvent this ambiguity, they proposed to use the improvement of probing task performance against a control task (predicting random labels from the same representations), i.e., the "selectivity" criterion. Recently, , challenged this dichotomy from an information theoretic viewpoint. They proposed to use an "information gain" criterion, which empirically is the reduction in cross entropy from a "control function task" probing from randomized representation.
In this paper, we show the "non-exclusive-or" dichotomy raised by Hewitt and Liang (2019) is valid information-theoretically. There is a difference between the original NLP model learning the task and the probe learning the task.
In addition, we show that the "selectivity" criterion and the "control function" criterion are comparably accurate.  formulated their errors with the difference in a pair of KL divergences. We show that the error of the "selectivity" criterion (Hewitt and Liang, 2019), if measured from cross entropy loss, can be formulated in the difference in a pair of KL divergences as well. When randomizations are perfect, these two criteria differ by only constant terms.
Empirically, on a POS tag probing task on English, French and Spanish translations, we show that the "selectivity" and the "control function" criteria highly agree with each other. We rank experi-ments with over 10,000 different hyperparameter settings using these criteria. The Spearman correlation of the Hewitt and Liang (2019) vs.  criteria are on par with the correlations of "accuracy vs. cross entropy loss" -two very strong baselines.
Overall, we recommend using control mechanisms to select probes, instead of relying on merely the probing task performance. When randomization is done well, controlling the target or representations are equivalent.
The task of probing is usually formulated as a classifier problem, with the representations as input, and the features indicating information as output. A straightforward method to train a classifier is by minimizing cross entropy, which is the approach we follow. Note that Voita and Titov (2020) derived training objectives from minimum description lengths, resulting in cross entropy loss and some variants.

Formulation
We adopt the information theoretic formulation of linguistic probes of , and briefly summarize as follows.
We want to probe true labels T from representations R. An ideal probe should accurately report the code-target mutual information I(T ; R), which is unfortunately intractable. We will write I(T ; R) in an alternative form.
Let p(T | R) be the unknown true conditional distribution, and a diagnostic probe, according to the setting in literature (Alain and Bengio, 2017;Hewitt and Manning, 2019;, is an approximation q θ (T | R) parameterized by θ, then: where p and q θ stand for p(T | R) and q θ (T | R) respectively. We also use H(p, q θ ) = −E p log q θ to represent the cross entropy for simplicity.

The source of probing error
A valid dichotomy Traditionally, people use the cross entropy loss of the diagnostic probe H(p, q θ ) to approximate I(T ; R). We can derive a source of error by rewinding the above formulations: The first term on RHS, H(T ), is independent of either R or θ. Therefore, a low cross entropy loss H(p, q θ ) can be caused by either of the two scenarios: • High code-target mutual information I(T ; R), indicating the representation R contains rich information about the target T .
• Low KL-divergence between p(T | R) and q θ (T | R), indicating the probe learns the task.
The two scenarios exactly correspond to the dichotomy of Hewitt and Liang (2019).
A good probe To get a good probe, we want H q θ (T | R) to approximate I(T ; R) as much as possible. This means a good probe should minimize KL(p q θ ), as proposed by .
However, empirically  (as well as many previous articles) used the cross entropy loss H q φ (T | R) to select good probes, which is insufficient, as described above. Alternatively, Hewitt and Liang (2019) and  proposed control tasks and control functions, respectively.

The control task
The control task (Hewitt and Liang, 2019) sets random targets for the probing task. Let us use c(T ) to indicate a "control function" applied on a token v that originally has label T . The control function could nullify the information of the input, if necessary.
If we measure the difference between cross entropy in the control task and probing task , we can derive a form of error margin in their measurements 1 . Let us use a short-hand notations H(p c , q θc ) − H(p, q θ ) for clarity. Now, what does the diff between cross entropy on control task and probing task actually contains?
We already knew that H(T ) = Const. According the definition of control function, the output c(T ) would be independent of R. Then: Therefore: where ∆ h is a short-hand notation for the measurement error in the control task criteria: (1) When the probe fits the true distribution to a similar extent on both the control task and probing task, the error ∆ h would be small. Unfortunately, both KL terms are intractable.

The control function
Control function  introduces a random processor c(·) on the representation R.
To measure the information gain, they used an "information gain" criterion: 1 Note that Hewitt and Liang (2019) used accuracy instead of cross entropy. We discuss cross entropy so as to compare the errors against the control function  Noticing that mutual information terms are intractable, they approximated the objective with the difference between cross entropy in the control function task (we refer to as "control function" henceforth) and the probing task: To compute the error of this approximation, they reformulated the terms as following: The two H(T ) terms cancel out, then: Where we abbreviate similarly as we did for the control task. Specifically, we write φ for the probe parameters of control function to tell apart from θ in the control task.  showed that the error of their approximation, ∆ p = G(T, R, c) − G(T, R, c), can be expressed as: Again, when the probe fits the true distribution to a similar extent on both the target labels distribution and the probing task, the ∆ p will be small. Unfortunately, both KL terms are intractable too.

Control task vs control function
From Equations 1 and 2, we showed that the selectivity criterion of Hewitt and Liang (2019) and the information gain criterion of , if both measured in cross entropy loss, have very similar errors in approximating information gains.
These errors, ∆ h and ∆ p respectively, appear in very similar forms 2 . The probes selected from these two criteria should be highly correlated to each other, and our experiments will confirm.  Table 1: Spearman correlations between t acc (the "selectivity" criterion (Hewitt and Liang, 2019)) and f ent (the "gain" criterion ) are on par with two "accuracy vs. cross entropy" correlations.

Setup
We use the same family of probes as Hewitt and Liang (2019) and , multiple layers perceptrons with ReLU activations, to show the correlations of their "good probe" criteria (control task and control function, respectively). Overall, we sweep the probe model hyperparameters with a unified training scheme on three tasks (probe, control task, control function). The control task (function) setting includes labels (embeddings) drawn from a uniform random sample once before all experiments. In each training, we follow the setting of (Hewitt and Liang, 2019). We save the model with the best dev loss, report the test set loss and accuracy, and average across 4 different random seeds.
Data We use the Universal Dependency (Zeman et al., 2019) dataset loaded with the Flair toolkit (Akbik et al., 2018). We examine three languages: English, French, and Spanish. For the probing task, we use POS with labels provided by SpaCy 3 . We use the embedding of multilingual BERT (mBERT) implemented by huggingface (Wolf et al., 2019). If a word is split into multiple word pieces, we average its representations.

The "good probes" are good for both
When measuring the qualities of probes using the "selectivity" (Hewitt and Liang, 2019) or "information gain"  criterion, we show that the rules-of-thumb for training good probes largely agree.
• Early stopping before 24,000 gradient steps (approximately 4 epochs) could inhibit probe quality, but longer training procedures do not improve the probe qualities considerably. • Smaller probes are better in general, but exceptions exist. For example, when weight decay 3 https://spacy.io is set to 0, probes with one hidden layer and 40 hidden neurons are better in both criteria. • A small weight decay is beneficial. We include more descriptions, including comprehensive experiment configurations and plots in the Supplementary Material.

The high correlation between criteria
In addition to the qualitative correlations shown above, we compute the correlations between the two criteria over a grid-search style hyperparameter sweep of over 10,000 configurations. For each "probe, control task, control function" experiment set, we record the following four criteria: • t acc: Difference between probing task and control task accuracy. This is the "selectivity" criterion of Hewitt and Liang (2019). • f ent: Difference between control function and probing task cross entropy. This is the "gain" criterion of . • t ent: Difference between the control task and the probing task cross entropy. • f acc: Difference between the probing task and control function accuracy. We collect all experiments of each language according to these criteria, and use Spearman correlation to test three pairs of correlations. As is reported in Table 1, the (t acc, f ent) correlations are comparable to two strong baselines, (t acc, t ent) and (f acc, f ent), the correlations between measurements in accuracy and cross entropy losses.

Conclusion
When selecting probes that better approximate I(T ; R), we recommend measuring with a control mechanism instead of relying on the traditional cross entropy on probing task. We show both information-theoretically and empirically, that controlling the targets and representations are equivalent, as long as the control mechanism is randomized.

A Difference between the two criteria
In Section 3 we show that the error of the two criteria can both be written as difference between a pair of KL divergence (modulo a constant term).
Here we further simplify the terms when we assume the control task and functions take perfectly random distributions (i.e., independent of the task and representations, respectively).
When we use the same hyperparameter setting, q θ (T | R) and q φ (T | R) should be able to approximate p(T | R) to the same extent, so KL(p q θ ) and KL(p q φ ) cancel out. Additionally, following the definitions of the control function and control tasks, we can simplify as follows: p(c(T )) q θc (c(T )) KL(p(T |c(R)) q φc (T |c(R))) = E p(T ) p(T ) q φc (T ) Therefore, the difference between errors of criteria of Hewitt and Liang (2019) and  are: In short, these two criteria differ by terms dependent only on the randomization functions and the inherent distributions of task labels, i.e., irrelevant terms. Hewitt and Liang (2019) proposed some rules-ofthumb to select good probes, including a "simple probe" suggestion. We sweep the hyper parameters to test whether these rules also apply when measuring probe qualities using the control function .
Early stopping could inhibit probe quality Early stopping, if stopped before 24,000 gradient steps (approximately 4 epochs) may inhibit the quality of probes. In addition, we Figure 1 shows high correlation between the "selectivity" (Hewitt and Liang, 2019) and "information gain"  criteria.
Small weight decay is beneficial We find that smaller weight decays (e.g., 0.01) are more beneficial for probes than larger weight decays. While the two criteria rank the capacity of probes similarly, the most simple probes tend to stand out more distinctively with the "selectivity" criterion , as are shown on Figures 2, 3, and 4.
Smaller probes are not necessarily better We find that while smaller probes have higher "selectivity" and "information gain" for mBERT representations, probes with one hidden layer and 40-80 hidden neurons are better than more simplistic probes, as shown in Figures 5, 6 and 7. The plots show consistency between the two criteria. For example, larger models and more layers do not necessary lead to better results. Neither are the smallest probes with 0 hidden layers.
Note that we also swept hyperparameters for FastText, where probes with less parameters do not always outperform more complex probes in either accuracy, loss, selectivity, or information gain. Figures 8 and 9 illustrate these observations.

C Reproducibility
On a T4 GPU card, training one epoch takes around 20 seconds. Without setting maximum gradient steps, 98.6% of experiments finish within 400 epochs. We will open source our codes. Figure 1: Max gradient step vs accuracy, t acc, f acc (blue) and loss, t ent, f ent (green) on English. The "t" refers to control task (Hewitt and Liang, 2019), and "h" refers to control function . In these set of experiments, we look for the best learning rate and zero weight decay in each configuration. Figure 2: The "difference of accuracy" (Hewitt and Liang, 2019) and the "difference of loss"  criteria against weight decay on model configurations, on UD English. For each configuration, the learning rate leading to the highest accuracy is selected. Figure 3: The "difference of accuracy" (Hewitt and Liang, 2019) and the "difference of loss"  criteria against weight decay on model configurations, on UD French. For each configuration, the learning rate leading to the highest accuracy is selected. Figure 4: The "difference of accuracy" (Hewitt and Liang, 2019) and the "difference of loss"  criteria against weight decay on model configurations, on UD Spanish. For each configuration, the learning rate leading to the highest accuracy is selected. Figure 5: The "difference of accuracy" (Hewitt and Liang, 2019) and the "difference of loss"  criteria with different learning rates on model configurations, on UD English. The weight decay is set to 0. Figure 6: The "difference of accuracy" (Hewitt and Liang, 2019) and the "difference of loss"  criteria with different learning rates on model configurations, on UD French. The weight decay is set to 0.
Figure 7: The "difference of accuracy" (Hewitt and Liang, 2019) and the "difference of loss"  criteria with different learning rates on model configurations, on UD Spanish. The weight decay is set to 0.    of probes on FastText. Probes with different capacities are ranked similarly using these two criteria.