GNN-SL: Sequence Labeling Based on Nearest Examples via GNN

For a given word $w$, the output $h$ generated from the last layer of the vanilla SL model is used as its representation, where $h\in\mathbb{R}^{m}$. Then $h$ is fed into a multi-layer perceptron (MLP) to obtain the probability distribution $p_{\text{vanilla}}$ via a softmax layer:

$\displaystyle p_{\text{vanilla}}(y|w,x)=\text{softmax}(\text{MLP}(h))$

(1)

For each word $w$, its corresponding embedding $h$ is used to query $k$ nearest neighbors set $\mathcal{N}$ from the training set using $L^{2}$ Euclidean distance $d(h,\cdot)$ as similarity measure. The retrieved nearest neighbors set $\mathcal{N}$ is formulated as $(key,value)$ pairs, where $key$ represents the retrieved similar word and $value$ represents the corresponding SL label.

Then the retrieved examples are converted into a distribution over the label vocabulary based on an RBF kernel output Vert et al. (2004) of the distance to the original embedding $h$:

$\displaystyle p_{\text{kNN}}(y|w,x)\propto\sum_{(k,v)\in\mathcal{N}}\mathbbm{1}_{y=v}\exp(\frac{-d(k,h)}{T})$

(2)

where $T$ is a temperature parameter to flatten the distribution. Finally, the vanilla distribution $p_{\text{vanilla}}(y|w,x)$ is augmented with $p_{\text{kNN}}(y|w,x)$ generating the final distribution $p_{\text{final}}(y|w,x)$:

	$\displaystyle p_{\text{final}}(y|w,x)=$	$\displaystyle\lambda p_{\text{vanilla}}(y|w,x)+$		(3)
		$\displaystyle(1-\lambda)p_{\text{kNN}}(y|w,x)$

where $\lambda$ is adjustable to make a balance between $k$NN distribution and vanilla distribution.

In the constructed graph, we define three types of nodes $\mathcal{A}$:

(1) Input nodes, denoted by $a_{\text{input}}\in\mathcal{A}$, which correspond to words of the input sentence. In the example of Figure 2 Step 3, the input nodes are displayed with the word sequence $x=\{\text{Obama},\text{lives},\text{in},\text{Washington}\}$;

(2) Neighbor nodes, denoted by $a_{\text{neighbor}}\in\mathcal{A}$, which correspond to words in the retrieved neighbors. The context of nearest neighbors is also included (and thus treated as neighbor nodes) in an attempt to capture more abundant contextual information for the retrieved neighbors. For the example in Figure 2, for each input word with representation $h$, $k$ nearest neighbors are queried from the cached representations of all words in the training set with the $L2$ distance as the metric of similarity. Taking the input word $\{\text{Obama}\}$ as the example with $k=2$, we obtain two nearest neighbors $\{{\text{Obama},\text{Trump}}\}$ leveraging the $k$NN search. The contexts of each retrieved nearest neighbor are also considered by adding both left and right contexts around the retrieved nearest neighbor, where $\{{\text{Obama}}\}$ is expanded to $\{\text{[CLS]},\text{Obama},\text{is}\}$ and $\{{\text{Trump}}\}$ is expanded to $\{\text{[CLS]},\text{Trump},\text{is}\}\}$²²2[CLS] is a special token usually applied in pre-training language models (e.g., BERT) representing the beginning of a sentence.. The analysis of the size of the context is conducted in Section 6.

(3) Label nodes, denoted by $a_{\text{label}}\in\mathcal{A}$, since the labels of nearest neighbors provide important evidence for the input node to classify, we wish to pass the influence of neighbors’ labels to the input node along the graph. As will be shown in ablation studies in Section 6, the consideration of label nodes introduces a significant performance boost. Shown in Figure 2 Step 3, taking the input word $\{\text{Obama}\}$ as the example, the two retrieved nearest neighbors are $\{{\text{Obama}}\}$ and $\{{\text{Trump}}\}$, and both the corresponding node labels are $\{\text{B-PER}\}$.

With the above formulated, $\mathcal{A}$ can be rewritten as $\{a_{\text{input}},a_{\text{neighbor}},a_{\text{label}}\}$.

Given the three types of nodes $\mathcal{A}=\{a_{\text{input}},a_{\text{neighbor}},a_{\text{label}}\}$, we connect them using different types of edges to enable information passing.

We define four types of edges for $\mathcal{R}$: (1) edges within the input nodes $a_{\text{input}}$, notated by $r_{\text{input-input}}$; (2) edges between the neighbor nodes $a_{\text{neighbor}}$ and the input nodes $a_{\text{input}}$, notated by $r_{\text{neighbor-input}}$; (3) edges within the neighbor nodes $a_{\text{neighbor}}$, notated by $r_{\text{neighbor-neighbor}}$; and (4) edges between the label nodes $a_{\text{label}}$ and the neighbor nodes $a_{\text{neighbor}}$, denoted by $r_{\text{label-neighbor}}$. All types of edges are bi-directional which allows information passing on both sides. We use different colors to differentiate different relations in Figure 2 Step 3.

For $r_{\text{input-input}}$ and $r_{\text{neighbor-neighbor}}$, they respectively mimic the attention mechanism to aggregate the context information within the input word sequence or the expanded nearest context, which are shown with the black and green color in Figure 2. For $r_{\text{neighbor-input}}$, it connects the retrieved neighbors and the query input word, transferring the neighbor information to the input word. For $r_{\text{label-neighbor}}$ colored with orange, information is passed from label nodes to neighbor nodes, which is ultimately transferred to input nodes.

For each edge $(s,e,n)$, the message transferred from the source node $s$ to the target node $n$ can be formulated as:

$\displaystyle\text{M}(s,e,n)=W^{v}_{\tau(s)}h^{l-1}_{s}W_{\phi(e)}$

(5)

where $d$ denotes the dimensionality of the vector, $W^{v}_{\tau(s)}\in\mathbb{R}^{d\times d}$ and $W_{\phi(e)}\in\mathbb{R}^{d\times d}$ are two learnable weight matrixes controling the outflow of node $s$ from the node side and the edge side respectively.

For the receiver node $n$, the importance of each neighbor $s$ with the relationship $e$ can be viewed as the attention mechanism in Vaswani et al. (2017), which relates different positions of a single sequence. To introduce the attention weights, we first map the source node $s$ to a key vector $K(s)$ and the target node $n$ to a query vector $Q(n)$:

$\displaystyle K(s)=W^{k}_{\tau(s)}h^{l-1}_{s},Q(n)=W^{q}_{\tau(n)}h^{l-1}_{n}$

(6)

where $W^{k}_{\tau(s)}\in\mathbb{R}^{d\times d}$ and $W^{q}_{\tau(n)}\in\mathbb{R}^{d\times d}$ are two learnable matrixes.

As we use different types of edges for node connections, we follow Hu et al. (2020) to keep a distinct edge-matrix $W_{\phi(e)}\in\mathbb{R}_{d\times d}$ for each edge type between the dot of $Q(n)$ and $K(s)$:

		$\displaystyle\text{P}(n,s)=K(s)W_{\phi(e)}Q(n)^{\mathsf{T}}\cdot\frac{\mu\langle\tau(s),\phi(e),\tau(n)\rangle}{\sqrt{d}},$		(7)
		$\displaystyle\text{A}(s,e,n)=\mathop{\text{softmax}}\limits_{s\in\mathcal{N}(n),e\in\phi(e)}(P(Q(n),K(s))),$

where $\mu\in\mathbb{R}^{|\mathcal{A}|\times|\mathcal{R}|\times|\mathcal{A}|}$ is a learnable matrix denoting the contribution of each edge with a different relationship. Similar to Vaswani et al. (2017), the final attention is the concatenation of heads which compute the attention weights independently:

		$\displaystyle\text{MultiHead}(s,e,n)=\text{Concat}(\text{head}_{1}\cdots\text{head}_{g})W^{O},$		(8)
		$\displaystyle\text{head}_{i}=\text{A}(s,e,n)\cdot\text{M}(s,e,n)$

where $g$ is the number of heads for multi-head attention (8 is used for our model) and $W^{O}$ is a learnable model parameter.

For each edge $(s,e,n)$, we now have the attention weight $\text{A}(s,e,n)$ and the information $\text{M}(s,e,n)$, the next step is to obtain the weighted-sum information from all neighboring nodes:

$\displaystyle\text{Aggregate}(\cdot)=W^{o}_{\tau(n)}(\mathop{\oplus}\limits_{\forall s\in\mathcal{N}(n)}\text{MultiHead}(s,e,n))$

(9)

where $\oplus$ is element-wise addition and $W^{o}_{\tau(n)}\in\mathcal{R}^{d\times d}$ is a learnable model parameter used as an activation function like a linear layer.

The aggregated representation for each input word is used as its final representation, passed to the softmax layer for classification. For all our experiments, the number of heads is 8.

As described in Section 4, we need the pre-trained vanilla SL model to extract features to initial the nodes of the constructed graph. For all our experiments, we choose the standard BERT-large Devlin et al. (2018) and RoBERTa-large Liu et al. (2019) for English tasks, as well as the standard BERT-large and ChineseBERT-large Sun et al. (2021) for Chinese tasks.

In the process of $k$NN retrieval, the number of nearest neighbors $k$ is set to 32, and the size of the nearest context window is set to 7 (setting both the left and right side of the window to 3). The two numbers are chosen according to the evaluation in Section 6 and perform best in our experiments. For the $k$ nearest search, we use the last layer output of the pre-trained vanilla SL model as the representation and the $L^{2}$ distance as the metric of similarity comparison.

The task of NER is normally treated as a char-level tagging task: outputting a NER tag for each character. We conduct experiments on CoNLL2003 Sang and De Meulder (2003) and OntoNotes5.0 Pradhan et al. (2013) for English, and MSRA Levow (2006), OntoNotes4.0 Pradhan (2011) for Chinese.

For the chosen baselines, we make a comparison with Cross-View Training CVT Clark et al. (2018), BERT-MRC Li et al. (2019a) for English datasets and Lattice-LSTM Zhang and Yang (2018), Glyce-BERT Meng et al. (2019), BERT-MRC Li et al. (2019a) for Chinese datasets.

Results for the NER task are shown in Table 1, and from the results:

(1) We observe a significant performance boost brought by $k$NN, respectively +0.06, +1.62, +1.75 and +0.19 for English CoNLL 2003, English OntoNotes 5.0, Chinese OntoNotes 4.0 and Chinese MSRA. which proves the importance of incorporating the evidence of retrieved neighbors.

(2) We observe a significant performance boost for vanilla+GNN over both the vanilla model: respectively +2.14 on the BERT-Large and +1.78 on the RoBERTa-Large for English OntoNotes 5.0, and +1.10 on the BERT-Large and +0.40 on the ChineseBERT-Large for Chinese MSRA. Due to the fact that both vanilla+GNN and the vanilla model output the final layer representation to the softmax function to obtain final probability and that $p_{\text{kNN}}$ do not participate in the final probability interpolation for both, the performance boost over vanilla SL demonstrates that we are able to obtain better token-level representations using GNNs.

(3) When comparing with vanilla+$k$NN, we observe further improvements of +0.33, +2.14, +3.17 and +1.10 for English CoNLL 2003, English OntoNotes 5.0, Chinese OntoNotes 4.0 and Chinese MSRA dataset, respectively, showing that the proposed GNN model has the ability to filtrate the useful neighbors to augment the vanilla SL model;

(4) There is an inconspicuous improvement brought by interpolating both the $k$NN probability and GNN probability (vanilla + $k$NN + GNN) over the proposed GNN-SL (vanilla + GNN), e.g., +2.16 v.s. +2.14 on the BERT-Large for English OntoNotes 5.0, and +1.13 v.s. +1.10 on the BERT-Large for Chinese MSRA. This demonstrates that as the evidence of retrieved nearest neighbors (and their labels) has been assimilated through GNNs in the representation learning stage, the extra benefits brought by interpolating $p_{\text{knn}}$ in the final prediction stage is significantly narrowed.

The task of CWS is normally treated as a char-level tagging problem: assigning seg or not seg for each input word.

For evaluation, four Chinese datasets retrieved from SIGHAN 2005³³3The website of the 4-th Second International Chinese Word Segmentation Bakeoff (SIGHAN 2005) is: http://sighan.cs.uchicago.edu/bakeoff2005/ are used: PKU, MSR, CITYU, and AS, and benchmarks Multitask pretrain Yang et al. (2017), CRF-LSTM Huang et al. (2019) and Glyce+BERT Meng et al. (2019) are chosen for the comparison.

Results for the CWS task are shown in Table 2. From the results, same as the former Section 5.3, with different vanilla models as the backbone, we can observe obvious improvements by applying the $k$NN probability (vanilla + $k$NN) or the GNN model (vanilla + GNN), while keeping the same results between vanilla + GNN and vanilla + GNN + $k$NN, e.g., for PKU dataset +0.1 on BERT + $k$NN, +0.3 both on BERT + GNN and BERT + GNN + $k$NN. Notablely we achieve SOTA for all four datasets with the ChineseBERT: 96.9 (+0.2) on PKU, 98.3 (+0.3) on CITYU, 98.5 (+0.2) on MSR and 96.9 (+0.2) on AS.

The task of POS is normally formalized as a character-level sequence labeling task, assigning labels to each of the input word.

We use Wall Street Journal (WSJ) and Tweets Ritter et al. (2011) for English datasets, Chinese Treebank 5.0, Chinese Treebank 6.0, and UD1.4 Xue et al. (2005) for Chinese datasets.

For the comparisons, we choose Meta-BiLSTM Bohnet et al. (2018) for English datasets and Joint-POS Shao et al. (2017), Lattice-LSTM Zhang and Yang (2018) and Glyce-BERT Meng et al. (2019) for Chinese datasets.

Results for the POS task are shown in Table 3 for Chinese datasets and Table 4 for English datasets. As shown, with different vanilla models and datasets, the phenomenons are the same as the former Section 5.3 that improves largely based on the $k$NN probability and further on the GNN model. For the results of Chinese CTB5 as the example, +0.42 on BERT + $k$NN, further +0.70 on BERT + GNN, and +0.73 on BERT + GNN + $k$NN.

To illustrate the augment of our proposed GNN-SL, we visualize the retrieved $k$NN examples as well as the input sentence in Table 5. For the first example the long-tail case “Phoenix”, which is assigned with “LOCATION” by the vanilla SL model, is amended to “ORGANIZATION” by the nearest neighbors. Especially, we can observe that both 1-th and 8-th retrieved labels are “ORGANIZATION” while the 16-th retrieved label is “LOCATION” which is against the ground truth. That phenomenon proves that retrieved neighbors do relate to the input sentence in different ways: some are close to the original input sentence while others are just noise, and our proposed GNN-SL has the ability to better model the relationships between the retrieved nearest examples and the input word. For the second example, with the augment of the retrieved nearest neighbors, our proposed GNN-SL outputs the correct label “PERSON” for the word “Tom Moody”.