Semantic Label Smoothing for Sequence to Sequence Problems

Label smoothing is a regularization technique commonly used in deep learning  (Szegedy et al., 2016; Chorowski and Jaitly, 2017; Vaswani et al., 2017; Zoph et al., 2018; Real et al., 2018; Huang et al., 2019), that improves calibration (Müller et al., 2019) and helps in label de-noising (Lukasik et al., 2020a). Here, one smooths labels by introducing a prior in the label space (often just a uniform distribution) in order to prevent overly confident predictions and achieve better model calibration, both of which lead to better generalization.

Given these benefits, it is natural to consider whether label smoothing can be applied to sequence-to-sequence (seq2seq) prediction tasks in Natural Language Processing. Here, inducing a label prior involves smoothing in sequence space. However, this is a challenging task because the output space is exponentially large for sequences, unlike the label space in standard classification. Previous works approached this challenge either by smoothing over individual tokens of the target sequence, or by sampling a few nearby targets according to Hamming distance or BLEU score Norouzi et al. (2016); Elbayad et al. (2018). These techniques however do not guarantee that the smoothed targets lie within the space of acceptable targets (i.e., the sampled new target may no longer be grammatically correct or even preserve semantic meaning).

In this work, we propose a label smoothing approach for seq2seq problems that overcomes this limitation. Given a large-scale corpus of valid sequences, our approach selects a subset of sequences that are not only semantically similar to the target sequence, but also well formed. We achieve this using a pre-trained model to find semantically similar sequences from the corpus, and then use BLEU scores to rerank the closest targets. We empirically show that this approach improves over competitive baselines on multiple machine translation tasks.

Sequence-to-sequence (seq2seq) learning involves learning a mapping from an input sequence $\mathbf{x}$ (e.g., a sentence in English) to an output sequence $\mathbf{y}$ (e.g., a sentence in French). Canonical applications include machine translation and question answering.

Formally, let $\mathscr{X}$ denote the space of input sequences (e.g., all possible English sentences), and $\mathscr{Y}$ the space of output sequences (e.g., all possible French sentences). We represent by $\mathbf{x}=[x_{1},x_{2},...x_{N}]$ an input sequence consisting of $N$ tokens, and similarly $\mathbf{y}=[y_{1},y_{2},...y_{N^{\prime}}]$ an output sequence with $N^{\prime}$ tokens. Our goal is to learn a function $f\colon\mathscr{X}\to\mathscr{Y}$ that, given an input sequence, generates a suitable target sequence.

To achieve this goal, we have a training set $\mathscr{S}\subseteq(\mathscr{X}\times\mathscr{Y})^{n}$ comprising pairs of input and output sequences. We then seek to minimise the objective

where $p_{\theta}(\cdot|\mathbf{x};\theta)$ is a parametrized distribution over all possible output sequences. Given such a distribution, we choose $f(\mathbf{x})=\operatorname{argmax}_{\mathbf{y}\in\mathscr{Y}}p_{\theta}(\mathbf{y}\mid\mathbf{x})$. Observe that one may implement (1) via a token-level decomposition,

This may be understood as a maximum likelihood objective, or equivalently the cross-entropy between $p_{\theta}(\cdot|\mathbf{x};\theta)$ and a one-hot distribution concentrated on $\mathbf{y}$.

Label smoothing (Szegedy et al., 2016; Pereyra et al., 2017; Müller et al., 2019) is a simple means of correcting this in classification settings. Smoothing involves simply adding a small reward to all possible incorrect labels, i.e., mixing the standard one-hot label with a uniform distribution over all labels. This regularizes the training and generally leads to better predictive performance as well as probabilistic calibration (Müller et al., 2019).

Given the success of label smoothing in classification settings, it is natural to explore its value in seq2seq problems. However, standard label smoothing is clearly inadmissible: it would require smoothing over all possible outputs $\mathbf{y}^{\prime}\in\mathscr{Y}$, which is typically an intractably large set. Nonetheless, we may follow the basic intuition of smoothing by adding a subset of related targets to the observed sequence $\mathbf{y}$, yielding a smoothed loss

Here, $\mathscr{R}(\mathbf{y})$ is a set of related sequences that are similar to the ground truth $\mathbf{y}$, and $\alpha>0$ is a tuning parameter that controls how much we rely on the observed versus related sequences.

The quality of $\mathscr{R}(\mathbf{y})$ is important for our task. Ideally, we would like an $\mathscr{R}(\mathbf{y})$ that: (i) is efficient to compute, and (ii) comprises sequences which meaningfully align with $\mathbf{x}$ (e.g., are plausible alternate translations). We now assess several options for constructing $\mathscr{R}(\mathbf{y})$ in light of the above.

Such random sequences contain general target language understanding (e.g., French grammar for an English to French translation task). However, these sequences are unlikely to have any semantic correlation with the true label.

While this approach increases the semantic similarity to $\mathbf{y}$, operating on a token level is limiting. For example, one may change the meaning of a factual sentence by changing even a few words. Further, operating at a per-token level limits the diversity of $\mathscr{R}(\mathbf{y})$, since, e.g., all sequences have the same number of tokens and structure as $\mathbf{y}$.

A key challenge is efficiently identifying semantically similar sequences to $\mathbf{y}$. To achieve this in a tractable manner, we propose the following procedure (see Algorithm 1). First, we assume the existence of an embedding space for output sequences. For example, this could be the result of BERT (Devlin et al., 2019), which embeds each sequence into a fixed vector representation. Given such an embedding space and a corpus $\mathscr{Y}_{\mathrm{ref}}$ of reference sequences, we may now efficiently compute the neighbors of $\mathbf{y}$, $\mathscr{N}(\mathbf{y})$, comprising the top-$k$ closest sequences in $\mathscr{Y}_{\mathrm{ref}}$ for the given $\mathbf{y}$ Indyk and Motwani (1998).¹¹1Alternatively, one could consider selecting highest scoring augmentations based on a pre-trained seq2seq model. However, the resulting quadratic computational complexity renders such an approach impractical.

The elements of $\mathscr{N}(\mathbf{y})$ can be expected to have high semantic similarity with $\mathbf{y}$, which is desirable. However, such sequences may not meaningfully align with the original input $\mathbf{x}$ (e.g., may not be sufficiently close translations). To account for this, we prune the elements from $\mathscr{N}(\mathbf{y})$ based on the BLEU score. Intuitively, this pruning retains sequences that are both semantically similar and have non-trivial token overlap with $\mathbf{y}$.

We use $\mathscr{Y}_{\mathrm{ref}}$ as all output sequences in the training set. In practice, one may however use any set of sequences that are valid for the domain in question. We find $k=100$ closest sequences in this space, and smooth over $k^{\prime}=5$ pruned sequences with the highest BLEU score to $\mathbf{y}$. In Table 1 we show example augmentations. Notice both the diversity of augmentations, as well as relatedness to the original targets.

We propose a novel label smoothing approach for sequence to sequence problems that selects a subset of sequences that are not only semantically similar to the target sequences, but are also well formed. We achieve this by using a pre-trained model to find semantically similar sequences from the training corpus, and then we use BLEU score to rerank the closest targets. Our method shows a consistent and significant improvement over state-of-the-art techniques across different datasets.

In future work, we plan to apply our semantic label smoothing technique to various sequence to sequence problems, including Text Summarization Zhang et al. (2019) and Text Segmentation Lukasik et al. (2020b). We also plan to study the relation between pretraining and data augmentation techniques.