Specializing Word Embeddings (for Parsing) by Information Bottleneck (Extended Abstract)

1 Introduction

Word embedding systems like BERT and ELMo use spelling and context to obtain contextual embeddings of word tokens. These systems are trained on large corpora in a taskindependent way. The resulting embeddings have proved to then be useful for both syntactic and semantic tasks, with different layers of ELMo or BERT being somewhat specialized to different kinds of tasks [Peters et al., 2018b; Goldberg, 2019]. State-of-the-art performance on many NLP tasks can be obtained by fine-tuning, i.e., back-propagating task loss all the way back into the embedding function [Peters et al., 2018a; Devlin et al., 2018].

In this paper, we explore what task-specific information appears in the embeddings before fine-tuning takes place. We focus on the task of dependency parsing, but our method can be easily extended to other syntactic or semantic tasks. Our method compresses the embeddings by extracting just their syntactic properties-specifically, the information needed to reconstruct parse trees (because that is our task). Our nonlinear, stochastic compression function is explicitly trained by variational information bottleneck (VIB) to forget taskirrelevant information. This is reminiscent of canonical correspondence analysis [Anderson, 2003], a method for reducing the dimensionality of an input vector so that it remains predictive of an output vector, although we are predicting an output tree instead. However, VIB goes beyond mere dimensionality reduction to a fixed lower dimensionality, since it also avoids unnecessary use of the dimensions that are available in the compressed representation, blurring unneeded capacity via randomness. The effective number of dimensions may therefore vary from token to token. For example, a parser may be content to know about an adjective token only that it is adjectival, whereas to find the dependents of a verb token, it may need to know the verb's number and transitivity.

� Originally published in Proceedings of the 2019 Conference on EMNLP-IJCNLP [Li and Eisner, 2019]

We try compressing to both discrete and continuous taskspecific representations. Discrete representations yield an interpretable clustering of words. We also extend information bottleneck to control the contextual specificity of the token embeddings, making them more like type embeddings.

This specialization method is complementary to the previous fine-tuning approach. Fine-tuning introduces new information into word embeddings by backpropagating the loss, whereas the VIB method learns to exploit the existing information found by the pretrained language model. VIB also has less danger of overfitting, since it fits fewer parameters than fine-tuning. VIB is also very fast to train on a single GPU.

We discover that our syntactically specialized embeddings are predictive of the gold POS tags, validating the intuition that a POS tag summarizes a word token's syntactic properties. However, our representations are tuned explicitly for discriminative parsing, so our discrete tags prove to be even more useful for this task than POS tags, even at the same level of granularity. Our continuous tags are also more useful than the uncompressed ELMo representations, when it comes to generalizing to test data.

2 Formal Model

The information bottleneck (IB) method originated in information theory and has been adopted by the machine learning community as a training objective [Tishby et al., 2000; Tishby and Zaslavsky, 2015]

Let X represent an "input" random variable such as a sentence, and Y represent a correlated "output" random variable such as a parse. Suppose we know the joint distribution p"X;Y". (In practice, we will use the empirical distribution over a sample of "x; y" pairs.) Our goal is to learn a stochastic map p�"t j x" from X to some compressed representation T, which in our setting will be something like a tag sequence. IB seeks to minimize

where I"�; �" is the mutual information. A low loss means that T does not retain very much information about X (the second term), while still retaining enough information to predict Y. 1 The balance between the two MI terms is controlled by a Lagrange multiplier �. By increasing �, we increase the pressure to keep I"X; T" small, which "narrows the bottleneck" by favoring compression over predictive accuracy I"Y; T".

We extend the original IB objective (1) and add terms I"Ti; Xj ˆXi" to control the context-sensitivity of the extracted tags. Here Ti is the tag associated with the ith word, Xi is the ELMo token embedding of the ith word, and ˆXi is the same word's ELMo type embedding (before incorporating context).

In this section, we justify the additional term and sketch how to efficiently estimate variational bounds on all terms (lower bound for I"Y;T" and upper bound for the rest).

We instantiate the variational IB (VIB) estimation method [Alemi et al., 2016] on our dependency parsing task, as illustrated in Figure 1. We compress a sentence's word embeddings Xi into continuous vector-valued tags or discrete tags Ti ("encoding") such that the tag sequence T retains maximum ability to predict the dependency parse Y ("decoding"). Our chosen architecture compresses each Xi independently using the same stochastic, information-losing transformation.

The IB method introduces the new random variable T, the tag sequence that compresses X, by defining the conditional distribution p�"t j x". In our setting, p� is a stochastic tagger, for which we will adopt a parametric form (§ 2.1 below). Its parameters � are chosen to minimize the IB objective (2). By IB's independence assumption, 1 the joint probability can be factored as p�"x; y; t" = p"x" � p"y j x" � p�"t j x".

1 Since T is a stochastic function of X with no access to Y, it obviously cannot convey more information about Y than the uncompressed input X does. As a result, Y is independent of T given X, as in the graphical model T ! X ! Y .

2.1 I"X;T" - the Token Encoder p�"t j x"

Under this distribution, I"X;T" def =Ex »Et�p� "t jx" »log p� "t jx" p� "t" ……. Making this term small yields a representation T that, on average, retains little information about X. The outer expectation is over the true distribution of sentences x; we use an empirical estimate, averaging over the unparsed sentences in a dependency treebank. To estimate the inner expectation, we could sample, drawing taggings t from p�"t j x".

We must also compute the quantities within the inner brackets. The p�"t j x" term is defined by our parametric form. The troublesome term is p�"t" = Ex 0 »p�"t j x 0 "…, since even estimating it from a treebank requires an inner loop over treebank sentences x 0 . To avoid this, variational IB replaces p�"t" with a simpler distribution to obtain an upper bound on I"X;T". See [Li and Eisner, 2019] for the form of this distribution and how we optimize its parameters to tighten the uppper bound, jointly with optimizing �.

2.2 Two Token Encoder Architectures

We choose to define p�"t j x" = Q n i=1 p�"ti j xi". That is, our stochastic encoder will compress each word xi individually (although xi is itself a representation that depends on context): see Figure 1. We make this choice not for computational reasons-our method would remain tractable even without this-but because our goal in this paper is to find the syntactic information in each individual ELMo token embedding (a goal we will further pursue in § 2.3 below).

To obtain continuous tags, define p�"ti j xi" such that ti 2 R d is Gaussian-distributed with mean vector and diagonal covariance matrix computed from the ELMo word vector xi via a feedforward neural network. Alternatively, to obtain discrete tags, define p�"ti j xi" such that ti 2 f1; : : : ; kg follows a softmax distribution, where the k softmax parameters are similarly computed by a feedforward network. See [Li and Eisner, 2019] for architecture details.

2.3 I"Ti; Xj ˆXi" - the Type Encoder s�"ti j ˆ xi "

While the IB objective (1) asks each tag ti to be informative about the parse Y, we were concerned that it might not be interpretable as a tag of word i specifically. Given ELMo or any other black-box conversion of a length-n sentence to a sequence of contextual vectors x1; : : : ; xn, it is possible that xi contains not only information about word i but also information describing word i + 1, say, or the syntactic constructions in the vicinity of word i. Thus, while p�"ti j xi" might extract some information from xi that is very useful for parsing, there is no guarantee that this information came from word i and not its neighbors. Although we do want tag ti to consider contexte.g., to distinguish between noun and verb uses of word i-we want "most" of ti's information to come from word i itself. Specifically, it should come from ELMo's level-0 embedding of word i, denoted by ˆ xi-a word type embedding that does not depend on context.

To penalize Ti for capturing "too much" contextual information, our modified objective (2) adds a penalty term � � I"Ti; Xj ˆXi", which measures the amount of information about Ti given by the sentence X as a whole, beyond what is given by ˆXi: I"Ti; X j ˆXi" def = Ex »Eti�p� "ti jx" »log p� "ti jx" p� "ti j ˆ xi " ……. Setting � > 0 will reduce this contextual information.

We can derive an upper bound on I"Ti; Xj ˆXi" by approximating the conditional distribution p�"ti j ˆ xi" with a variational distribution s�"ti j ˆ xi". See details in [Li and Eisner, 2019].

Type Encoder Architectures Notice that s�"ti j ˆ xi " may be regarded as a type encoder, with the same neural architecture as p�"ti j xi" except that p� takes a token vector as input whereas s� takes a context-independent type vector. s� is not used at test time, but only as part of our training objective.

2.4 I"Y;T" - the Decoder q�"y j t"

Finally, I"Y;T" def = Ey; t�p� »log p� "y jt" p"y" …. The p"y" can be omitted during optimization as it does not depend on �. Thus, making I"Y;T" large tries to obtain a high log-probability p�"y j t" for the true parse y when reconstructing it from t alone.

But how do we compute p�"y j t"? This quantity effectively marginalizes over possible sentences x that could have explained t. To avoid this expensive computation, we obtain a lower bound of I"Y;T" by replacing p�"y j t" with a variational approximation q�"y j t". Here q� is a tractable conditional distribution, and may be regarded as a stochastic parser that runs on a compressed tag sequence t.

In short, when t is a stochastic compression of a treebank sentence x, we would like our variational parser on average to assign high log-probability q�"y j t" to its treebank parse y.

3 Training and Inference

With the variational approximations in §2, our final minimization objective is an upper bound on (2). Training involves computing two tractable KL divergence terms and a stochastic gradient step. See [Li and Eisner, 2019] for algorithmic and mathematical details.

To evaluate our trained model's ability to parse a sentence x from compressed tags, we obtain a parse as argmaxy q�"y j t", where t � p�"� j x" is a single sample. A better parser would instead estimate argmaxy Et »q�"y j t"… where Et averages over many samples t, but this is computationally hard.

4 Experimental Setup

Data Throughout §§5-6, we will examine our compressed tags on a subset of Universal Dependencies [Nivre et al., 2018], or UD, a collection of dependency treebanks. We experiment on 9 languages: Arabic, Hindi, English, French, Spanish, Portuguese, Russian, Italian, and Chinese - languages with different syntactic properties like word order. For each sentence, x is obtained by running the standard pre-trained ELMo [Gardner et al., 2017; Che et al., 2018] on the UD token sequence, and y is the labeled UD dependency parse without any part-of-speech (POS) tags.

Optimization We optimize with Adam [Kingma and Ba, 2014], a variant of stochastic gradient descent. We train for 50 epochs with minibatches of size 20 and L2 regularization. For each training sentence, we average over 5 i.i.d. samples of T to reduce the variance of the stochastic gradient. We experiment with different dimensionalities d 2 f5; 32; 256; 512g for the continuous tags, and different cardinalities k 2 f32; 64; 128g for the discrete tag set. We also tried different values �; � 2 f10 �6 ; 10 �5 ; � � � ; 10 1 g of the compression tradeoff parameter.

5 Scientific Evaluation

In this section, we study what information about words is retained by our automatically constructed tagging schemes. First, we show the relationship between I"Y;T" and I"X;T" on English as we reduce � to capture more information in our tags. 2 Second, across 9 languages, we study how our automatic tags correlate with gold part-of-speech tags (and in English, with other syntactic properties). See [Li and Eisner, 2019] for full tables and experimental details.

5.1 Tradeoff Curves

As we lower � to retain more information about X, both I"X;T" and I"Y;T" rise (Figure 2). There are diminishing returns: after some point, the additional information retained in T does not contribute much to predicting Y. Also noteworthy is that at each level of I"X; T", very low-dimensional tags (d = 5) perform on par with high-dimensional ones (d = 256). (Note that the high-dimensional stochastic tags will be noisier to keep the same I"X; T".) The low-dimensional tags allow far faster CPU parsing. This indicates that VIB can achieve strong practical task-specific compression.

5.2 Learned Tags vs. Gold POS Tags

We investigate how our automatic tag Ti correlates with the gold POS tag Ai provided by UD.

Continuous Version In Figure 3, the first figure shows the original uncompressed level-1 ELMo embeddings of the tokens in test data. In the two-dimensional visualization, the POS tags are vaguely clustered but the boundaries merge together and some tags are diffuse. The second figure is when

2 We always set � = � to simplify the experimental design.

� = 10 �3 (moderate compression): our compressed embeddings show clear clusters that correspond well to gold POS tags. The third figure is when � = 1 (too much compression), when POS information is largely lost. An interesting observation is that the purple NOUN and blue PROPN distributions overlap in the middle distribution, meaning that it was unnecessary to distinguish common nouns from proper nouns for purposes of our parsing task 3 .

Beyond Gold POS Tags, we study other syntactic distinctions: tense, number, and transitivity. Even with moderate compression, ti achieves 0:87 classification accuracy in distinguishing between transitive and intransitive English verbs.

Discrete Version We also quantify how well our specialized discrete tags capture the traditional POS categories, by investigating I"A;T". In effect, we are doing transfer learning, fixing our trained IB encoder (p�) and now using it to predict A instead of Y, but otherwise following § 2.4. We experiment on all 9 languages, taking Ti at the moderate compression level � = 0:001, k = 64. Experimental details and results are in [Li and Eisner, 2019]. Averaging over the 9 languages, the reconstruction retains 71% of POS information. We can conclude that the information encoded in the specialized tags correlates with the gold POS tags, but does not perfectly predict the POS.

6 Engineering Evaluation

As we noted in §1, learning how to compress ELMo's tags for a given task is a fast alternative to fine-tuning all the ELMo parameters. We find that indeed, training a compression method to keep only the relevant information does improve our generalization performance on the dependency parsing task.

Baselines. The continuous baselines include: Iden, which leaves the ELMo embeddings uncompressed; MLP, which uses a multi-layer perceptron ([Dozat and Manning, 2016]) to reduce the dimensionality in a task-specific and nonlinear way. The discrete baseline: POS uses the k � 17 gold POS tags from the UD dataset.

3 Both can serve as arguments of verbs and prepositions. Both can be modified by determiners and adjectives.

Runtime. Our VIB approach is quite fast: it trains on 10,000 sentences in 100 seconds, per epoch, on a single GPU.

Analysis. In the continuous case, the VIB representation outperforms all three baselines in 8 of 9 languages, and is not significantly worse in the 9th language (Hindi). In short, our VIB joint training generalizes better to test data. This is because the training objective (2) includes terms that focus on the parsing task and also regularize the representations. In the discrete case, the VIB representation outperforms gold POS tags (at the same level of granularity) in 6 of 9 languages, and of the other 3, it is not significantly worse in 2. This suggests that our learned discrete tag set could be an improved alternative to gold POS tags [Klein and Manning, 2003] when a discrete tag set is needed for speed.

7 Conclusion

In this paper, we have proposed two ways to syntactically compress ELMo word token embeddings, using variational information bottleneck. We automatically induce stochastic discrete tags that correlate with gold POS tags but are as good or better for parsing. We also induce stochastic continuous token embeddings (each is a Gaussian distribution over R d ) that forget non-syntactic information captured by ELMo. These stochastic vectors yield improved parsing results, in a way that simpler dimensionality reduction methods do not. They also transfer to the problem of predicting gold POS tags, which were not used in training.

Acknowledgments

This work was supported by the National Science Foundation under Grant No. 1718846.