A Non-Linear Structural Probe (2021)

Abstract: Probes are models devised to investigate the encoding of knowledge—e.g. syntactic structure—in contextual representations. Probes are often designed for simplicity, which has led to restrictions on probe design that may not allow for the full exploitation of the structure of encoded information; one such restriction is linearity. We examine the case of a structural probe (Hewitt and Manning, 2019), which aims to investigate the encoding of syntactic structure in contextual representations through learning only linear transformations. By observing that the structural probe learns a metric, we are able to kernelize it and develop a novel non-linear variant with an identical number of parameters. We test on 6 languages and find that the radial-basis function (RBF) kernel, in conjunction with regularization, achieves a statistically significant improvement over the baseline in all languages—implying that at least part of the syntactic knowledge is encoded non-linearly. We conclude by discussing how the RBF kernel resembles BERT’s self-attention layers and speculate that this resemblance leads to the RBF-based probe’s stronger performance.

## introduction probing has been widely used in an effort to better understand what linguistic knowledge may be encoded in contextual word representations such as bert @xcite and elmo @xcite . these probes tend to be designed with simplicity in mind and with the intent of revealing what linguistic structure is encoded in an embedding, rather than simply learning to perform an nlp task @xcite @xcite this preference for simplicity has often led researchers to place restrictions on probe designs that may not allow them to fully exploit the structure in which information is encoded @xcite . this preference has led many researchers to advocate the use of linear probes over non-linear ones @xcite . this paper treats and expands upon the structural probe of @xcite , who crafted a custom probe with the aim of investigating the encoding of sentence syntax in contextual representations. they treat probing for syntax as a distance learning problem: they learn a linear transformation that warps the space such that two words that are syntactically close to one another (in terms of distance in a dependency tree) should have contextual representations whose euclidean distance is small. this linear approach performs well, but the restriction to learning only linear transformations seems arbitrary. why should it be the case that this information would be encoded linearly within the representations? in this paper, we recast hewitt and manning (2019)'s structural probing framework as a general metric learning problem. this reduction allows us to take advantage of a wide variety of non-linear extensions-based on kernelization-proposed in the metric learning literature @xcite . these extensions lead to probes with the same number of parameters, but with an increased expressivity. by exploiting a kernelized extension, we are able to directly test whether a structural probe that is capable of learning non-linear transformations improves performance. empirically, we do find that non-linearity helps-a structural probe based on a radial-basis function (rbf) kernel improves performance significantly in all 6 languages tested over a linear structural probe. we then perform an analysis of bert's attention, asserting it is a rough approximation to an rbf kernel. as such, it is not surprising that the syntactic information in bert representations is more accessible with this specific non-linear transformation. we conclude that kernelization is a useful tool for analyzing contextual representations-enabling us to run controlled experiments and investigate the structure in which information is encoded. here, the goal of probing is to evaluate whether the contextual representations capture the syntax in a sentence. in the case of the structural probe, the goal is to see whether the syntactic distance between any two words can be approximated by a learned, linear distance function: this is simply minimizing the difference between the syntactic distances obtained from the dependency tree and the distance between the two vectors under our learned transformation. from the pairwise distances predicted by the probe, @xcite algorithm can be used to recover the onebest undirected dependency tree. ## kernelized metric learning the restriction to a linear transformation may hinder us from uncovering some of the syntactic structure encoded in the contextual representations. indeed, there is no reason a-priori to expect that bert encodes its knowledge in a fashion that is specifically accessible to a linear model. however, if we were to introduce non-linearity by using a neural probe, for example, we would have to pit a model with very few parameters (the linear model) against one with very many (the neural network); this comparison is not fair and also goes against the spirit of designing simple probes. to preclude the need for a neural probe, we instead turn to a kernelized probe. the key insight is that the structural probe reduces the problem of probing for linguistic structure to that of metric learning @xcite . this can be clearly seen in eq. ( 3 ), where the probe learns a distance metric between two representations in such a way that it matches the syntactic one. recognizing this relationship allows us to take advantage of established techniques from the metric learning literature to improve the performance of the probe without increasing its complexity, e.g. through kernelization. ## regularized metric learning we also take advantage of two common regularization techniques employed in the metric learning literature to further improve the transformations learned; both act on the matrix @xmath0. ( 3 ). the frobenius norm regularizer takes the form this is the matrix analogue of the l 2 squared regularizer. minimizing the frobenius norm of the learned matrix has the effect of keeping the values in the matrix small. it has been a popular choice for regularization in metric learning with adaptations to a variety of problems @xcite . we also consider the trace norm regularizer, which is of the form the trace norm regularizer is the matrix analogue of the l 1 regularizer and it encourages the matrix a to be low rank. as jain et al. ( 2010 ) point out, using a low-rank transformation in conjunction with a kernel corresponds to a supervised kernel dimensionality reduction method. ## experiments we experiment with @xcite probe on 6 typologically diverse languages, following the experimental design of hall maudslay et al. @xcite . our data comes from the universal dependency 2.4 treebank @xcite , providing sentences and their dependency trees, annotated using the universal dependencies annotation scheme. 3 for each sentence we calculate contextual representations using multilingual bert. for all languages, we took the first 12,000 sentences (or the maximum number thereof) in the train portion of the treebank and created new 80-10-10 train-test-dev splits. we present the results from our comparison of a re-implementation of @xcite linear structural probe and the non-linear kernelized probes in ## the rbf kernel and self-attention the rbf kernel produces improvements across all analyzed languages. this suggests that it is indeed the case that syntactic structure is encoded nonlinearly in bert. as such, analyzing this specific kernel may yield insights into what this structure is. indeed, none of the other kernels systematically improve over the linear baseline, implying this is not just an effect of the non-linearity introduced through use of a kernel-the specific structure of the rbf kernel must be responsible. in this section, we argue that the reason that the rbf kernel serves as such a boon to probing is that it resembles bert's attention mechanism; recall that bert's attention mechanism is defined as k projects vector h i into a key vector, while q projects h j into a query one. when the key and query vectors are similar (i.e. have a high dot product), the value of this equation is large and word j attends to word i. this bears a striking resemblance to the gaussian kernel. indeed, if we assume the linearly transformed representations have unit norm, i.e. the resemblance between these equations, taken together with the significant improvements in capturing syntactic distance, suggest that this encoded information indeed lives in an rbf-like space in bert. such information can then be used in its self-attention mechanism; allowing bert to pay attention to syntactically close words when solving the cloze language modeling task. being attentive to syntactically close words would also be supported by recent linguistic research, since words sharing syntactic dependencies have higher mutual information on average @xcite . the representations we analyze, though, are taken from bert's final layer; as such, they are not trained to be used in any self-attention layerso why should such a resemblance be relevant? bert's architecture is based on the transformer @xcite , and uses skip connections between each self-attention layer. such skip connections create an incentive for residual learning, i.e. only learning residual differences in each layer, while propagating the bulk of the information @xcite . as such, bert's final hidden representations should roughly live in the same manifold as its internal ones. it is interesting to note that the rbf kernel achieves the best performance in terms of uuas in all languages, but it only twice achieves the best performance in terms of dspr. this may be due to the fact that, as we can see by examination of eq. ( 8 ), the distance returned by the rbf kernel will not exceed 2, whereas syntactic distances in the tree will. further, the gradient of the rbf kernel contains an exponential term which will cause it to go to zero as distance increases (while an examination of the unkernelized loss function reveals the opposite behavior). this means that it will be less sensitive to the distances between syntactically distant words and focus more on words with small distances. this may partially explain its better performance on uuas, and comparably worse performance as measured by correlation (which counts pairwise differences between all words, not just those which are directly attached in the tree). furthermore, our probe's focus on nearby words resembles the general attentional bias towards syntactically close words @xcite . the direct resemblance between self-attention mechanisms and our proposed probe metric poses a new way of understanding results from more complex probes. while @xcite understood the euclidean-squared distance of hewitt and manning as an isometric tree embedding, their geometric interpretation did not factor in the rest of bert's architecture. such simplified contextless probes cannot tell us how linguistic properties are processed by a sequence of learned modules @xcite . however, we consider representations in the context of the model which is expected to employ them. from this perspective, simpler metrics may be rough approximations to our rbf kernel space, which is actually capable of measuring linguistic properties that can be easily extracted by an attention-based architecture. ## conclusion we find that the linear structural probe @xcite used to investigate the encoding of syntactic structure in contextual representations can be improved through kernelization, yielding a non-linear model. this kernelization does not introduce additional parameters and thus does not increase the complexity of the probe-at least if one treats the number of parameters as a good proxy for model complexity. at the same time, the rbf kernel improves probe performance in all languages under consideration. this suggests that syntactic information may be encoded non-linearly in the representations produced by bert. we hypothesize that this is true due to the similarity of the rbf kernel and bert's self-attention layers.