Learning Decoupled Retrieval Representation for
Nearest Neighbour Neural Machine Translation

Given a source sentence ${\bm{x}}=\{x_{1},x_{2},\ldots,x_{|{\bm{x}}|}\}$ and a target prefix ${\bm{y}}_{<t}=\{y_{1},y_{2},\ldots,y_{t-1}\}$, the vanilla NMT predicts the next target word $y_{t}$ by:

$$p_{c}(y_{t}|{\bm{x}},{\bm{y}}_{<t})\propto\textrm{exp}\Big{(}q({\bm{h}}_{t})\Big{)}$$

(1)

where ${\bm{h}}_{t}=f_{\theta}({\bm{x}},{\bm{y}}_{<t})\in\mathcal{R}^{d}$ is the context vector at step $t$ with respect to ${\bm{x}}$ and ${\bm{y}}_{<t}$; $f_{\theta}(\cdot)$ can be arbitrary encoder-decoder network with parameters $\theta$, such as Transformer Vaswani et al. (2017); $q(\cdot)$ linearly projects ${\bm{h}}_{t}$ to target vocabulary size.

kNN-MT hypothesizes that the same target words have similar representations. To dynamically incorporate external sentence pairs $\mathcal{D}=\{({\bm{x}}^{(i)},{\bm{y}}^{(i)})\}_{i=1}^{|\mathcal{D}|}$, kNN-MT extends Eq. 1 by interpolating a retrieval-based probability $p_{r}$:

$$p_{knn}=(1-\lambda)\times p_{c}+\lambda\times p_{r}$$

(2)

where $\lambda$ is the interpolation coefficient as a hyper-parameter.

Specifically, kNN-MT first uses a pre-trained NMT model to force decoding each sentence pair $({\bm{x}}^{(i)},{\bm{y}}^{(i)})$ to build a key-value datastore $\mathcal{H}$:

$$\mathcal{H}=\bigcup_{i=1}^{|\mathcal{D}|}\bigcup_{t=1}^{|{\bm{y}}^{(i)}|}\Big{\{}({\bm{h}}^{(i)}_{t},y_{t}^{(i)})\Big{\}}$$

(3)

The key is the word-level context representation ${\bm{h}}^{(i)}_{t}$ and the value is the gold target word $y^{(i)}_{t}$. Then, given $\mathcal{H}$ and predicted target prefix $\hat{{\bm{y}}}_{<t}$ at test time, kNN-MT models $p_{r}(\hat{y_{t}}|{\bm{x}},\hat{{\bm{y}}}_{<t})$ by measuring the distance between query $\hat{{\bm{h}}}_{t}=f_{\theta}({\bm{x}},\hat{{\bm{y}}}_{<t})$ and its k-nearest representations $\{(\tilde{{\bm{h}}}_{i},\tilde{v}_{i})\}_{i=1}^{k}$ in $\mathcal{H}$:

$$p_{r}(\hat{y}_{t}|{\bm{x}},\hat{{\bm{y}}}_{<t})\propto\sum_{i=1}^{k}\mathbbm{1}_{\hat{y}_{t}=\tilde{v}_{i}}\textrm{exp}\Big{(}\frac{-d(\tilde{{\bm{h}}}_{i},\hat{{\bm{h}}}_{t})}{T}\Big{)},$$

(4)

where $d(\cdot)$ is $L_{2}$ distance; $T$ is temperature hyper-parameter; $\mathbbm{1}$ is the indicator function.

According to Eq. 1-4, we can see that the context representation ${\bm{h}}$ simultaneously plays two roles in kNN-MT: (1) the semantic vector for $p_{c}$; (2) the retrieval vector for $p_{r}$. We note that coupling the same ${\bm{h}}$ in the two scenes is sub-optimal. Recall that ${\bm{h}}$ in the translation model is generally learned through cross-entropy loss, which only pays attention to the gold target token and ignores others.¹¹1In practice, we often use its label-smooth variant, which evenly assigns a small probability mass to all non-gold labels without distinction. However, a good retrieval vector should be able to distinguish between different tokens, especially those owning similar representations. Therefore, we attempt to derive a new retrieval vector ${\bm{z}}$ from ${\bm{h}}$ for better retrieval performance.

We use a simple feedforward network as an adapter to transform the original representation ${\bm{h}}$ to desired retrieval representation ${\bm{z}}$:

$${\bm{z}}=\textrm{FFN}({\bm{h}})=\textrm{ReLU}({\bm{h}}{\bm{W}}_{1}+{\bm{b}}_{1}){\bm{W}}_{2}+{\bm{b}}_{2},$$

(5)

where ${\bm{W}}_{1}\in\mathcal{R}^{d\times d_{f}}$, ${\bm{W}}_{2}\in\mathcal{R}^{d_{f}\times d_{o}}$, ${\bm{b}}_{1}\in\mathcal{R}^{d_{f}}$, and ${\bm{b}}_{2}\in\mathcal{R}^{d_{o}}$ are learnable parameters; $d_{f}$ and $d_{o}$ are the intermediate hidden size and output size of the adapter, respectively. When $d_{o}<d$, the adapter network can be regarded as a dimension reducer. As FFN is very lightweight compared to the calculation of ${\bm{h}}$, there is almost no latency in converting ${\bm{h}}$ to ${\bm{z}}$.For convenience, in the following description, we redefine ${\bm{h}}_{i}$ as the key of $i$-th key-value pair in the original datastore $\mathcal{H}$, and the corresponding value is denoted by $Y_{i}$ when there is no ambiguity. In this way, the new datastore $\mathcal{Z}$ can be denoted as $\mathcal{Z}=\{({\bm{z}}_{i},Y_{i})|i=1,\ldots,|\mathcal{H}|\}$, where ${\bm{z}}_{i}=\textrm{FFN}({\bm{h}}_{i})$.

In machine translation field, contrastive learning has been applied in multilingual translation Pan et al. (2021); Wei et al. (2021), cross-modal translation Ye et al. (2022), and learning robust representation for low-frequency word Zhang et al. (2021) etc. In this work, we use supervised contrastive learning Grill et al. (2020) with multiple positive and negative samples to learn the desired retrieval representation ${\bm{z}}$. Here, we regard the unique token $v$ in the target vocabulary $V$ as a natural supervision signal. We aim to make ${\bm{z}}$ more distinguishable, for example, pulling ${\bm{z}}$ of the same words together and pushing ${\bm{z}}$ of different words apart. Specifically, we first divide $\mathcal{Z}$ into $|V|$ clusters according to the token class label. E.g., $C_{v}=\{{\bm{z}}_{i}|i=1,\ldots,|\mathcal{Z}|,Y_{i}=v\}$, where $C_{v}$ is the context representation cluster of token $v$. Thus, given any context representation ${\bm{z}}\in\mathcal{Z}$ and its token label $v$, we can construct M positive samples ${\bm{z}}^{+}=\{{\bm{z}}_{1}^{+},\ldots,{\bm{z}}_{i}^{+},\ldots,{\bm{z}}_{M}^{+}\}$, where ${\bm{z}}_{i}^{+}$ is uniformly sampled from its owned cluster $C_{v}$ and ${\bm{z}}_{i}^{+}\neq{\bm{z}}$.²²2We use sampling with replacement when $|C_{v}|<M$. Likely, we further construct N negative samples ${\bm{z}}^{-}=\{{\bm{z}}_{1}^{-},\ldots,{\bm{z}}_{i}^{-},\ldots,{\bm{z}}_{N}^{-}\}$, where ${\bm{z}}_{i}^{-}\in\backslash C_{v}$, $\backslash C_{v}$ denotes other clusters except $C_{v}$. In the next part, we will describe how to build ${\bm{z}}^{-}$. Finally, given the anchor vector ${\bm{z}}$, its multiple positive samples ${\bm{z}}^{+}$ and multiple negative samples ${\bm{z}}^{-}$, we learn the adapter network through the following contrastive learning loss:

$$-\textrm{log}\frac{\sum\limits_{1\leq i\leq M}\textrm{exp}(s({\bm{z}},{\bm{z}}_{i}^{+}))}{\sum\limits_{1\leq i\leq M}\textrm{exp}(s({\bm{z}},{\bm{z}}_{i}^{+}))+\sum\limits_{1\leq j\leq N}\textrm{exp}(s({\bm{z}},{\bm{z}}_{j}^{-}))},$$

(6)

where $s(\cdot)$ is the score function implemented as cosine similarity with temperature $T^{\prime}$: $s({\bm{a}},{\bm{b}})=\frac{1}{T^{\prime}}\times\frac{{\bm{a}}^{T}{\bm{b}}}{\|{\bm{a}}\|\cdot\|{\bm{b}}\|}$. Note that $T^{\prime}$ is the temperature in training, which is different from the inference temperature $T$ in Eq. 4.

The key for Eq. 6 is the construction of negative samples ${\bm{z}}^{-}$. A trivial solution is randomly sampling from the entire space of $\backslash C_{v}$. However, this negative sample may be too easy to provide the effective learning signal Robinson et al. (2020). On the contrary, an extreme method for hard negative samples is to traverse $\backslash C_{v}$ to find the most similar negative samples for the anchor. The problem is that $|\backslash C_{v}|$ is close to $|\mathcal{Z}|$, with a scale of millions or more, resulting in enormous computational complexity. To solve it, we propose a fast and cheap approach to constructing hard negative samples. Specifically, we first collect the cluster centre $\bar{C}_{v}=\frac{1}{|C_{v}|}\sum_{i=1}^{|\mathcal{Z}|}\mathbbm{1}_{Y_{i}=v}{\bm{z}}_{i}$. We calculate the nearest K (K$>=$N) cluster centers w.r.t the anchor and randomly sample N clusters to make the source of the negative sample diverse. Then we randomly sample one point from the corresponding cluster as a negative sample. As the anchor vector only involves querying $|C|$ cluster centers and $|C|<<|\mathcal{Z}|$, our approach runs faster than the exact global search.

After training, we use the well-trained FFN to rebuild the retrieval datastore $\mathcal{H}$ into $\mathcal{Z}$. To further reduce calculation cost at test time, we introduce PCA to reduce the dimension of the retrieval vector. We also add normalization after PCA to guarantee the numerical stability of the input to the inner product. Another difference with Eq. 4 is that we use the inner product instead of the L2 distance as distance metrics. The reason is that using consistent distance metrics in training and inference improves performance in primitive experiments. Concretely, we modify the original kNN-MT in Eq. 4 as:

$$p_{r}(\hat{{\bm{y}}}_{t}|{\bm{x}},\hat{{\bm{y}}}_{<t})\propto\sum_{i=1}^{k}\mathbbm{1}_{\hat{y}_{t}=\tilde{v}_{i}}\textrm{exp}\Big{(}\frac{g(\tilde{{\bm{z}}}_{i})\otimes g(\hat{{\bm{z}}}_{t})}{T}\Big{)},$$

(7)

where $g(x)=\textrm{Norm}(\textrm{PCA}(x))$, $\otimes$ denotes inner product operation, $\tilde{{\bm{z}}}_{i}$ is the $i$-th nearest neighbor in $\mathcal{Z}$ for the current retrieval representation $\hat{{\bm{z}}}_{t}$. As a bonus, since the numeric range of the normalized inner product is $[0,1]$, which can be seen as the confidence in retrieving.³³3L2 distance lacks this feature because its numeric range is too broad, e.g., 0~1000 in our observation. We leverage this nature to modify the interpolation coefficient $\lambda$ in Eq. 2 to be aware of retrieval confidence:

$$\lambda^{*}=\lambda\times\frac{\sum_{i=1}^{k}g(\tilde{{\bm{z}}}_{i})\otimes g(\hat{{\bm{z}}}_{t})}{k}.$$

(8)

$\lambda^{*}$ can be considered a simple adaptive coefficient like Zheng et al. (2021); Jiang et al. (2021); Wang et al. (2022), but does not require training.

The closest work with us is CKMT Wang et al. (2022). As illustrated in Figure 1, there are two major differences compared with CKMT: (1) Clknn uses multiple positive and negative samples, while CKMT only considers a single positive and negative sample, limiting the exploration of representation space. (2) CKMT requires to partition clusters through cost-expensive clustering in full-scale datastore, while Clknn predefines clusters based on vocabulary labels and only involves calculating cluster centers. In practice, we spent about 6 hours on the CPU to complete the cluster operation in CKMT, while Clknn only takes about 3 minutes.

Table 2 reports the SacreBLEU scores in five domains. We can see that: (1) Clknn is robust about training data: using out-domain or in-domain average improves 1+ points than our kNN-MT; (2) The gap between in-domain and out-domain is small (about 0.3 points), meaning that our approach does not rely on in-domain data and is more practical than Zheng et al. (2021); Jiang et al. (2021); (3) using proposed $\lambda^{*}$ slightly improve the performance across the board. These results show that learning independent retrieval representation is helpful for vanilla kNN-MT. Besides, we also compare the inference speed between Clknn and kNN-MT through running five times on IT test set. The results show that Clknn has a comparable speed (97%$\pm$2%) to that of kNN-MT because the adapter in Clknn is very lightweight.

One of the main differences between Wang et al. (2022) and us is that we use multiple positive and negative samples in our training objective. We vary the number of M and N and report the BLEU scores in Table 3. As we can see, increasing M and N is helpful for our method. However, large M cannot befit more than increasing N. We attribute it to positive samples that are too easy to learn because most are close in embedding space. On the contrary, negative samples from different clusters can provide a stronger learning signal. To further validate the effectiveness of multiple samples, we also conduct experiments on Medical. The results are similar to that of IT: using M=2, N=32 is 1.64 BLEU points higher than using M=1, N=1 (56.52 vs. 54.88). It indicates that using multiple positive and negative samples is necessary to achieve good performance for contrastive learning.

Intuitively, our approach can learn more accurate retrieval representation than vanilla kNN-MT. To validate this hypothesis, we use IT validation as the datastore and plot the retrieval accuracy on top-k in Figure 3. We can see that Clknn has more robust retrieval accuracy than kNN-MT no matter how k changes. It indicates that the performance improvement comes from our better retrieval representation.

We visually present the differences between baseline and Clknn on embedding space. Specifically, we split three categories according to the word frequency in IT training set: HIGH(the first 1%), Middle(40%-60%) and LOW(the last 1%) ⁶⁶6We filter words whose frequency is less than 10.. We uniformly sample 10 unique words in each category and randomly sample 10 unique vector representations from the training datastore. We use t-SNE to plot these representations, as illustrated in Figure 2. We can see that: (1) high-frequency words’ representations are prone to distinguish for both baseline and Clknn; (2) Clknn has more close distances in the same vocabulary than baseline; (3) Clknn has more robust accuracy for low-frequency words.