Seeing Both the Forest and the Trees: Multi-head Attention for Joint Classification on Different Compositional Levels

Our model is based on a bidirectional long short-term memory (BiLSTM) that builds contextual vector representations for each word. These vectors are then passed through a multi-head attention mechanism, which predicts label distributions for both individual words and the whole sentence. Each attention head is incentivized to be predictive of a particular label, allowing the system to also assign labels to individual words while composing a sentence-level representation for sentence classification.

The network takes as input a tokenized sentence of length $N$ and maps it to a sequence of vectors $[x_{1},x_{2},...,x_{N}]$. Each vector $x_{i}$, corresponding to the $i^{th}$ token in a sentence, is the concatenation of its pre-trained GloVe word embedding $w_{i}$  [Pennington et al. (2014] with its character-level representation $c_{i}$, similar to ?). Passing each vector $x_{i}$ to a BiLSTM  [Graves and Schmidhuber (2005], we obtain compact token representations $z_{i}$ by concatenating the hidden states from each direction at every time step and projecting these onto a joint feature space using a tanh activation (Equations 1-3).

This is followed by a multi-head attention mechanism with $H$ heads  [Vaswani et al. (2017]. By setting $H$ equal to the size of our token-level tagset, we can create a direct one-to-one correspondence between attention heads and possible token labels – attention head $h\in\{1,2,...,H\}$ gets assigned to the $h$-th token-level label. For each attention head we calculate keys, queries and values at every word position through a non-linear projection of $z_{i}$ (Equations 4-6). All the queries for a given attention head are then combined into a single vector through averaging, which will represent a query for the corresponding token-level label in the context of the given sentence (Equation 7).

The unnormalized attention scores $a_{ih}\in\mathbb{R}^{1}$ are then calculated through a dot product between the query and the associated key for a particular token in position $i$ (Equation 8). Given the established correspondence between attention heads and token labels, this score now represents the model confidence that the token in position $i$ has label $h$. Therefore, we can predict the probability distribution over the token-level labels by normalizing $a_{ih}$ with a softmax function (Equation 9). By concatenating the scores and normalizing them across the heads $h$, we get $\tilde{t}_{i}\in\mathbb{R}^{H}$, which we use as the token-level output from the model, for both optimization and evaluation as a token-level tagger.

Next, we use the same attention scores $a_{ih}$ to construct sentence-level representations of the input. We apply a sigmoid activation ($\sigma$) and normalization to produce the normalized attention weights $\alpha_{ih}\in[0,1]$ (Equation 10). Standard attention functions use a softmax activation, which is best suited for assigning most of the attention weight to a single token and effectively ignoring the rest. However, it is often necessary that higher-level representations pay attention to many different locations, or, in our case, to multiple tokens in a given sentence. By using the sigmoid instead of softmax, similar to ?), the model will need to make separate decisions for each token, which in turn encourages the attention scores to behave more similarly to sequence labeling predictions.

A sentence-level representation $s_{h}$ is obtained as the weighted sum over all the value vectors in the sentence (Equation 11). This is followed by two feed-forward layers: the first one is non-linear and projects the sentence representation onto a smaller feature space, while the second one is linear and outputs a scalar sentence-level score $o_{h}\in\mathbb{R}$ for each head $h$ (Equation 12).

To make a sentence prediction, the sentence scores need to be collected across all heads. The challenge arises as these $H$ scores (equal to the number of token labels) have to map to the number of sentence labels $S$, which are not always in direct correspondence. To solve this, we use the fact that many text classification tasks have a default label that is common between the token and the sentence label sets – for example, the neutral label for sentiment analysis or the no-named-entity label for NER. In our datasets, only two situations arise:

1. 1.

   $\boldsymbol{H=S\text{: }}$ Each sentence label has a corresponding word-tag, and thus one head associated with it. Therefore, we can directly concatenate the sentence scores across all heads into a vector $\tilde{o}=[o_{1};o_{2};...;o_{H}]$. An example of such a task is sentiment analysis, as the possible labels (positive, negative, and neutral) are the same for both sentences and tokens.

2. 2.

   $\boldsymbol{H\neq S\text{ and }S=2}\text{: }$ The sentence labels are binary, while the token labels are multi-class, therefore an appropriate correspondence between the heads and the two sentence labels needs to be found. We concatenate the score obtained for the default head $o_{d}$ (corresponding to the default label) with the maximum score obtained for the non-default heads $o_{nd}$: $\tilde{o}=[o_{d};o_{nd}]$, where $d$ and $nd$ are the indices of the default and non-default heads, respectively, and $o_{nd}=\max\limits_{h\neq d}(o_{h})$. Named entity recognition can be an example of such a task – while there are many possible tags on the token level, we only detect the binary presence of any named entities on the sentence level.

A probability distribution $\tilde{y}\in\mathbb{R}^{S}$ over the sentence labels is obtained by applying a softmax on the extracted scores: $\tilde{y}=softmax(\tilde{o})$. The most probable label is returned as the sentence-level prediction.

In our model, the sentence-level scores are directly intertwined with the token-level predictions in order to improve performance on both levels. The attention weights are explicitly using the same predicted scores as the token-level output. Therefore, when the model learns to detect specific types of tokens, it will also assign more importance to those tokens in the corresponding attention heads. At the same time, when the model learns to attend to particular tokens for sentence classification, this will also help in identifying the correct labels for the token-level tagging task. By joining the two tasks, the architecture is able to share the information on both levels of granularity and achieve better results. In addition, this allows us to explicitly teach the model to focus on the same evidence as humans when performing text classification, leading to more explainable systems.

In Figure 1, we illustrate how this architecture – to which we refer to as the multi-head attention labeler (MHAL) – is applied on one input word to compute one attention head.

Our model can be optimized both as a sentence classifier and as a sequence labeler using a cross-entropy loss. Both $L_{sent}$ and $L_{tok}$ minimize the summation over the negative log likelihood between the predicted sentence (or token) label distribution and the gold annotation:

where $y^{(s)}_{j}$ and $t^{(s)}_{ij}$ are binary indicator variables specifying whether sentence $s$ truly is a sentence of label $j$ and token $t$ at position $i$ in sentence $s$ truly is a token of tag type $j$, respectively.

Recall that the sentence label distribution is based on the attention evidence scores, which represent, in turn, the token scores used for word-level classifications. If we train our model solely as a sentence classifier (by providing only sentence-level annotations), the network will also learn to label individual tokens. As all the parameters used by the token labeling component are also part of the sentence classifier, they will be optimized during the sentence-level training. Moreover, the network will learn the important areas of a sentence, combining the scores from individual words to determine the overall sentence label. In this way, our model performs zero-shot sequence labeling, a type of transductive transfer learning  [Ruder (2017]. In addition, when both levels receive supervision, the token signal encourages the network to put more weight on the attention heads indicative of the correct labels.

We include an auxiliary attention loss objective, based on ?), which encourages the model to more closely connect the two labeling tasks on different granularity levels. In its original formulation, the loss could only operate over binary labels, whereas we extend it for general multi-class classification by imposing two conditions on the attention heads:

1. 1.

   There should be at least one word of the same label as the ground-truth sentence. Intuitively, most of the focus should be on the words indicative of the sentence type.

2. 2.

   There should be at least one word that has a default label. Even if the sentence has a non-default class, it should still contain at least one default word.

While these conditions are not true for every text classification task, they are applicable in many settings and hold true for all the datasets that we experimented with. The two conditions can be formulated as a loss function and then optimized during training:

Next, we propose a custom regularization term for the multi-head attention mechanism to motivate the network to learn a truly distinct representation sub-space for each of the query vectors $q_{h}$. As opposed to the keys and values, which are associated with different words, the queries $q_{h}$ encapsulate the essence of a certain tag. Therefore, these vectors need to capture the distinctive features of a particular label and how it is different from other labels. To push the network towards this goal, we introduce the term $R_{q}$ and calculate it as the average cosine similarity between every pair of queries $q_{h}$ and $q_{i}$, with $h\neq i$ (equation 16). $R_{q}$ penalizes high similarity between any two query vectors and motivates the model to push them apart. Thus, this technique imposes a wider angle between the queries, encouraging the model to learn unique, diverse, and meaningful vector representations for the tags.

Lastly, we include an auxiliary objective for language modeling (LM) operating over characters and words, following the settings proposed by ?). The hidden representations from the forward and backward LSTMs are mapped to a new, non-linear space and used to predict the next word in the sequence, from a fixed smaller vocabulary. Recently, many NLP systems using multi-task learning include LM objectives along the core task to inject corpus-specific information into the model, as well as syntactic and semantic patterns  [Dai and Le (2015, Peters et al. (2017, Akbik et al. (2018, Marvin and Linzen (2018]. In our case, we include an LM loss to help the model learn general language features. While performing well on language modelling itself is not an objective, we expect it to provide improved biases and language-specific knowledge that would benefit performance.

The final loss function $L$ is a weighted sum of all the objectives described above. Setting particular coefficients $\lambda$ allows us to investigate the effect of the different components as well as controlling the flow of the supervision signal and the importance of each auxiliary task: $L=\lambda_{sent}L_{sent}+\lambda_{tok}L_{tok}+\lambda_{attn}L_{attn}+\lambda_{R_{q}}R_{q}+\lambda_{LM}L_{LM}$.

To evaluate all the different properties of the model, we focus on three text classification datasets where annotations are available either for both individual words and full sentences or only for the words, but we can infer the ones for the sentences. We show some concrete examples in Table 1.

SST: The Stanford Sentiment Treebank  [Socher et al. (2013] is a dataset of human-annotated movie reviews used for sentiment analysis. It contains not only sentence annotations for positive (P), negative (N), and neutral (O) reviews but also phrase-level annotations, which we converted to token labels by accounting for the minimum spans of tokens (up to length three) of a certain sentiment. Therefore, on SST we have three labels both at the sentence and at the word level.

CoNLL03: The CoNLL-2003 dataset  [Tjong Kim Sang and De Meulder (2003] for named entity recognition (NER), contains five possible word-level tags, for person, organization, location, miscellaneous, or other, which is used for non-named entities. At the sentence-level, binary classification labels can be inferred based on whether the sentence contains any entities (annotated with label $\overline{\text{O}}$) or not (annotated with O).

FCE: The First Certificate in English  [Yannakoudakis et al. (2011] is a dataset for fine-grained grammatical error detection. Ungrammatical words can contain five possible mistakes: in content, form, function, orthography, or other. There is also a sixth label for grammatical words. A sentence that contains at least one world-level mistake is ungrammatical overall. Therefore, a binary sentence classification task naturally occurs, as sentences can be grammatical (annotated with O) or not (annotated with $\overline{\text{O}}$).

All datasets are already tokenized and split into training, dev, and test sets. In Appendix A, we provide some statistics on the corpus (Table 4) and on the annotations available per split and label (Table 5).

We chose the best values for our hyperparameters based on the performance on the development set (see Table 6 in Appendix A). We perform each experiment with five different random seeds and report the average results. Following ?), we also applied label smoothing [Szegedy et al. (2016] with $\epsilon=0.15$ to increase the robustness to noise and regularize the label predictions during training. As evaluation metrics, we report (based on the task) the precision (P), accuracy (Acc), and micro-averaged F₁ score of all the labels and of all the non-default labels (denoted by a superscript $*$), as it is common in the multi-task learning literature  [Changpinyo et al. (2018, Martínez Alonso and Plank (2017]. For CoNLL03, we use the dedicated CoNLL evaluation script, which calculates F₁ on the entity level.

We can optimize different variations of the architecture by changing the $\lambda$ weights in the loss and thereby choosing which components are active. We experiment with the following variations of the model:

• •

  MHAL-joint: Corresponds to the fully supervised experiment, and is optimized both as a sentence classifier and a sequence labeler by setting $\lambda_{sent}=\lambda_{tok}=1.0$, while all the other $\lambda$ values are $0.0$.

• •

  MHAL-joint+: In addition to training on both sentences and tokens ($\lambda_{sent}=\lambda_{tok}=1.0$), the auxiliary objectives are also activated ($\lambda_{attn}=0.01$, $\lambda_{LM}=0.1$, and $R_{q}=0.5$).

• •

  MHAL-sent: The model receives only sentence-level supervision ($\lambda_{sent}=1.0$) while all the other $\lambda$ values are set to $0.0$. No supervision is provided on the token level, which means this model performs zero-shot sequence labeling.

• •

  MHAL-sent+: Receives supervision on the sentence level ($\lambda_{sent}=1.0$), with the auxiliary objectives also activated ($\lambda_{attn}=0.01$, $\lambda_{LM}=0.1$, and $R_{q}=0.5$). No supervision is provided on the token level ($\lambda_{tok}=0.0$), so this model also performs zero-shot sequence labeling.

For comparison, we also evaluate two baseline models that do not connect the different hierarchical levels and specialize only on sentence classification or sequence labeling.

• •

  BiLSTM-sent: Following the description and implementation of ?), we built one of the strongest neural sentence classifiers based on BiLSTMs and soft attention; we tuned the hyper-parameters based on the development set to achieve the best performance on our tasks.

• •

  BiLSTM-tok: Widely-used bidirectional LSTM architecture for sequence labeling, which has been applied to many tasks including part-of-speech tagging  [Plank et al. (2016] and named entity recognition  [Panchendrarajan and Amaresan (2018]. We also tuned the hyperparameters based on the development set in order to achieve the best results on each of the evaluation datasets.

Fully-supervised: In this setting, we investigate whether training a joint model to solve the task on multiple levels provides a performance improvement over focusing only on one level. Table 2 compares the MHAL joint text classification performances to a BiLSTM attention-based sentence classifier and a BiLSTM sequence labeler. The results show that the multi-task models systematically outperform the single-task models across all tasks and datasets, emphasizing the effectiveness of sharing information between hierarchical levels. While additional annotation is required to train the multi-task models, the same input sentences are used in all cases, indicating that the benefits are coming directly from the model solving the task on multiple levels, as opposed to just from seeing more data examples. Despite the sentence-level labels for CoNLL03 and FCE having been derived automatically from the existing token-level annotation, they still provide a performance improvement for sequence labeling, further showing the benefit of the multi-level architecture. By teaching the model where to focus at the token level, the architecture is able to make better decisions on the sentence-level classification task. In addition, the sentence-level objective acts as a contextual regularizer and encourages the model to learn better compositional representations, thereby improving performance also on the token-level labeling task.

Comparing MHAL-joint against MHAL-joint+ with auxiliary objectives shows further improvements. The attention loss optimizes the model to make matching predictions between both levels, while the language modeling objective encourages the network to learn more informative word representations and composition functions. We also separately evaluated the regularization term, $R_{q}$, and found that it helps more on the sequence labeling tasks. This implies that using the intermediate per-head queries, the model indeed learns unique sub-space representations that help it assess the uncertainty of each word-tag pair and strengthen its labeling decision.

Semi-supervised: We further experiment with MHAL-joint+, using the supervision signal of all sentences but varying the percentage p of the word-level annotations. In Figure 2, we present the sequence labeling results of the multi-task model, in comparison to the single-task BiLSTM-tok, gradually increasing p to allow more tokens to guide learning. We observe that using only 30%-50% of the token-annotated data, our model already approaches the fully-supervised performance of the regular model (BiLSTM-tok-100%), suggesting that the two tasks are positively influencing each other. General, abstract knowledge about the entire sentence meaning fluidly flows down to the words, while fine-grained word-level information is propagated up to the sentence, showing a beneficial transfer in both directions.

Unsupervised: In Table 3, we evaluate the architecture as a zero-shot sequence labeler, trained without any token-level supervision signal. In this setting, the model learns to label individual tokens by seeing only examples of sentence-level annotations. Because the multi-level attention is directly wired together with the sequence labeling output, the model is still able to learn in this difficult setting. This experiment also illustrates that the information does indeed flow from the sentence level down to individual tokens in this architecture.

As no other model can operate in this setting, we can only compare against a random baseline, in which the labels are assigned uniformly at random from the available set of labels. For datasets where the label distribution is very skewed, this can still be a difficult baseline to beat. While MHAL-sent only outperforms this baseline on the CoNLL03 dataset, MHAL-sent+ outperforms both on all datasets and metrics. These results show that our auxiliary losses introduce a necessary inductive bias and allow for better transfer of information from the sentence level to the tokens. We visualized the decisions computed inside the attention heads for different example sentences and provide them in Appendix A (Figures 3 and 4, respectively). We also observed that the choice of the metric based on which the stopping criterion is selected plays an important and interesting role in our zero-shot experiments (see Appendix B for details).