Learning Nearest Neighbour Informed Latent Word Embeddings to Improve Zero-Shot Machine Translation
Year: 2023

Abstract:
Multilingual neural translation models exploit cross-lingual transfer to perform zero-shot translation between unseen language pairs. Past efforts to improve cross-lingual transfer have focused on aligning contextual sentence-level representations. This paper introduces three novel contributions to allow exploiting nearest neighbours at the token level during training, including: (i) an efficient, gradient-friendly way to share representations between neighboring tokens; (ii) an attentional semantic layer which extracts latent features from shared embeddings; and (iii) an agreement loss to harmonize predictions across different sentence representations. Experiments on two multilingual datasets demonstrate consistent gains in zero shot translation over strong baselines.

heading: 1 Introduction
text: Many-to-many multilingual neural translation models (Firat et al., 2016; Johnson et al., 2017; Khandelwal et al., 2020; Fan et al., 2022) share a single representation space across multiple language pairs, which enables them to perform zero-shot translations between unseen pairs (Ha et al., 2017; Chen et al., 2022; Wu et al., 2022). Prior work on zeroshot translation has focused on aligning contextual, sentence-level representations from multiple languages (Ji et al., 2020; Pan et al., 2021a), to make these more ‘universal’ or language-agnostic (Gu et al., 2018; Gu and Feng, 2022). Non-contextual, token-level representations offer another space in which this kind of alignment could be pursued, but this space has not been thoroughly explored in prior work. Even lexicon-based methods (Conneau et al., 2020; Reid and Artetxe, 2022), which exploit tokenlevel anchors from multilingual dictionaries (Duan et al., 2020), still use these to align representations at the sentence level.
In this work, we explore a novel technique for sharing information across languages at the token
level, which exploits nearest neighbours (NNs) to aggregate information from subwords across multiple languages. When analysing embedding spaces, many authors speak in terms of “neighborhoods” or “subspaces” which group together tokens from a particular semantic field or other natural class. These neighborhoods form implicitly as a model learns similarities between embedded words or subwords. We propose to make this neighborhood structure explicit by forcing a model to consider a token’s neighbors when learning its embedding. Specifically, we dynamically perturb a translation model’s token embeddings at training time by averaging them with their nearest neighbors; thus a token like soccer may end up mixed together with related tokens such as football, fußball, or futbol from potentially distinct languages (Figure 1). This encourages the model to organize its subword embeddings in such a way that nearby tokens convey similar information to one another. We hypothesize that this process will produce a more structured embedding space which will in turn enable more fluent outputs. This process only uses the model’s embedding layer, and does not require any offline dictionaries or additional data.
Our experiments and ablations show that this simple technique significantly increases the effectiveness of translation models on the IWSLT17 and TED59 massively multilingual datasets. Concretely, our contributions include: (i) an efficient, gradient-friendly, soft representation-mixing technique which exploits token-level neighbors without changing the Transformer architecture; (ii) an attentional semantic layer which extracts features from
291
mixed representations to give neighbour-informed latent word embeddings, and which is a drop-in replacement for a conventional embedding layer; and (iii) an agreement loss which harmonizes predictions with and without neighbor-informed embeddings.

heading: 2 Translation with Nearest Neighbour Augmented Embeddings
text: We describe our model for nearest-neighbour informed token level embeddings (Figure 2) of subwords from multiple source languages.
Nearest Neighbor Retrieval Let Lemb be a word embedding layer that performs a lookup EMBp¨q using weights Wemb P R|V|ˆD, where V is a joint subword vocabulary over all languages and D is a fixed embedding dimension. Given the embedding q “ EMBpwq P R1ˆD of a subword w, we wish to find q’s nearest neighbour n (or neighbors n1, ..., nk) using maximum inner product search (MIPS) over the weight matrix Wemb:
n “ argmin xĂWemb ||q ´ x||22 “ argmin
xĂWemb p||x||22 ´ 2qTxq
(1)
Approximate solutions to (1) can be efficiently computed on-the-fly using anisotropic vector quantization (Guo et al., 2020).1
Given the approximate nearest neighbors (ANNs) n1, ..., nk of subword w, we compute a weighted average over these tokens’ embeddings
1Exact and approximate solutions yield similar results, but approximation gives significant gains in training speed.
with a weighting term λ:
EMBµpwq “ λ1 k
kÿ i“1 pEMBpniqq ` p1 ´ λqEMBpwq
(2) EMBµp¨q is computed directly from Wemb, which ensures that our technique remains gradientfriendly2 and does not need a separate warm-up step. Previous NN-based proposals for translation (Khandelwal et al., 2020) and language modeling (Khandelwal et al., 2019) have only explored NNs of contextualized representations, strictly for generation, and using neighbors from an offline frozen datastore of pretrained candidates. Their method proved effective for MT domain adaptation, rather than zero-shot translation which is the focus of this work. The ability to propagate gradients to a subword’s neighbors during training is novel and unique compared to previous NN-based techniques.
Attentional Semantic Representation To extract contextually-salient information from EMBµpwq, which combines information from many subwords in potentially disparate languages, we use a shared semantic embedding inspired by Gu et al. 2018; Wang et al. 2018 that shows a similar effect as topical modelling.
We introduce Wsem P RNˆD, where each of the N rows is taken to be a language-agnostic semantic representation. Wsem is shared across all languages. We use attention (Luong et al., 2015) to compute a latent embedding EMBlatentpwq using
2In Section 3.3 we introduce a caching heuristic which is not gradient-friendly; however, this is simply an implementation detail to speed up training, and the gradient-friendly presentation in this section achieves equivalent performance.
the averaged embedding EMBµpwq as query:
EMBlatentpwq “ SoftmaxpEMBµpwq.WTsemqWsem (3)
A residual connection from EMBµpwq gives the final NN-informed word embedding:
EMBknnpwq “ EMBlatentpwq ` EMBµpwq (4)
EMBknnpwq is a drop-in replacement for a conventional word embedding EMBpwq. Modelling Prediction Consistency Given a source sentence represented using conventional word embeddings and using NN-informed embeddings, following Kambhatla et al. (2022b) we model the loss with respect to target sentence yi as:
Li “ α1 LiNLLp pΘp yi|xiq qlooooooooooooomooooooooooooon source x-entropy
` α2 LiNLLp pΘp yi| kNNpxiqq qlooooooooooooooooomooooooooooooooooon k-NN embeds. source x-entropy ` β Lidistp pΘp yi|xiq, pΘp yi|kNNpxiqq qloooooooooooooooooooooooomoooooooooooooooooooooooon agreement loss (5)
where kNNpxiq denotes the set of k-nearest neighbors to token xi. This loss combines three terms: the first two are conventional negative log-likelihoods, while the third is an agreement loss measuring pairwise symmetric KL divergence between the output distributions for xi and kNNpxiq. This agreement-loss term performs coregularization by allowing explicit interactions between source sentences with and without NNinformed embeddings.

heading: 3 Experiments

heading: 3.1 Datasets
text: We conduct experiments on 2 multilingual datasets, each with BPE (Sennrich et al., 2016) vocabulary size of 32k subwords:
IWSLT17 (Cettolo et al., 2012) is an Englishcentric dataset3 totalling 1.8M parallel sentences. It has 8 supervised directions to and from German, Italian, Dutch and Romanian, each with about 220,000 parallel sentences, and 12 zero-shot directions. We use the official validation and test sets.
Ted59 (Qi et al., 2018) is a massively multilingual English-centric dataset4 with 116 translation directions totalling 10.8M parallel sentences. The imbalanced data—from 0.25M to just 2000 parallel samples for some language pairs—makes it ideal to study the effects of our method. Following (Aharoni et al., 2019; Raganato et al., 2021) we evaluate on 16 supervised pairs and 4 zero-shot (Arabic Ø French, Ukranian Ø Russian).

heading: 3.2 Baselines and Related Work
text: We compare against methods for encoder manifold alignment. These include strong baselines such as sentence representation alignment (SRA; Arivazhagan et al. 2019), softmax forcing (SF; Pham et al. 2019), the contrastive multilingual model (CL; Pan et al. 2021b), multilingual Transformer with disentagled positional embedding (DP; Liu et al. 2021), and latent variable based denoising (LV; Wang et al. 2021), along with the vanilla manyto-many zero-shot model (M2M). On TED59, we compare against CL and 3 explicit multilingual alignment techniques proposed by Raganato et al. (2021): word-alignment, language tag alignment, and the union of the two. We also implement and compare against Raganato et al.’s (2021) sparse 1.5entmax cross-attention variant.

heading: 3.3 Model and Implementation Details
text: All models use the configuration in Vaswani et al. 2017 using the fairseq toolkit (Ott et al., 2019). See reproducibility details in Appendix A.
3https://wit3.fbk.eu/2017-01 4github.com/neulab/word-embeddings-for-nmt
We use ScANN (Guo et al., 2020) for efficient ANN search 5 with k “ 3. To increase training speeds, we cache each subword’s ANNs for 400 iterations before recomputing them. We only (peridocally) cache subword IDs: the embedding EMBµp¨q is always computed directly from Wemb. We set λ “ 0.5, α1, α2 “ 1, and β “ 5. The attentional latent semantic representation layer has 512 dim (same as the embedding layer) and a size N of 1000 for IWSLT17 (smaller dataset) and 5000 for TED59 (larger dataset). We did not tune this hyperparameter and chose the values based on the size of the datasets. For evaluation, we report sacreBLEU (Post, 2018).

heading: 3.4 Results
text: Main Results. Tables 1 and 2 show our main results. On IWSLT17, our latent k-NN embedding model outperforms several strong baselines, including sentence-representation alignment and contrastive learning, by an average of 0.62 and 0.11 BLEU respectively across the 12 zero-shot pairs. Compared to the baseline many-to-many model, our method yields a 1.5 BLEU gain on average. Our method is able to improve zero-shot performance without deteriorating supervised performance.
On the TED59 dataset, we follow Raganato et al. (2021) in comparing against two multilingual model variants: the standard Transformer, and the Transformer with sparse entmax instead of standard softmax cross-attention. Our approach gains „3 BLEU points against the baseline, and 2 BLEU
5We use asymmetric hashing with 2-dimensional blocks and a quantization threshold of 0.2, and re-order the top 100 ANN candidates.
against the stronger contrastive model. Further, our model consistently outperforms strong, explicitly alignment-based methods.
Target-language Accuracy. To supplement the evaluation, we provide the accuracy score for target language identification6 in zero-shot scenarios, called Acc0. While the classical many-to-many NMT models (Johnson et al., 2017; Aharoni et al., 2019) enable zero-shot translations, several studies have shown that these models fail to reliably generalize to unseen language pairs, ending up with an off-target translation issue (Zhang et al., 2020). The model ignores the language label and the wrong target language is produced as a result. We observe significant improvements in target language accuracy, up to nearly 99% (absolute).

heading: 4 Analysis
text: Ablation Study. Table 3 reports ablations on the IWSLT17 test set. We find that kNN embeddings alone yield improvements over the baseline manyto-many model. By contrast, absent the other parts of our model, the attentional semantic layer deteriorates model performance. Only in combination with the agreement loss do we observe a benefit from this component.
Embedding Analysis. Figure 3 visualizes subword representations from models trained on IWSLT17. Each subword is colored according to the language in which it is most frequent. The overall layout of the two spaces is similar, although the
6We utilize FastText (Joulin et al., 2017) as a language identification tool to compare the translation language with the reference target language and keep count of the number of matches.

heading: ID Component dev.2010 test.2010
text: baseline model (left) exhibits a clear ring-shaped gap dividing the embeddings into two groups. With ANN embeddings (right), this gap is eliminated and the layout of the embeddings appears more homogeneous. Quantitatively, the average distance from a subword to its neighbors exhibits a smaller variance in the ANN model than in the baseline, which further supports the reading that ANN training creates a more homogeneous representation space in which subwords are more uniformly distributed.
Table 4 shows nearest neighbors for a random sample of subwords (additional examples in Table 5 in Appendix B). With ANN training, a subword’s nearest neighbors are generally its synonyms (e.g. _wonderful, _large _tremendous, and _big as neighbors to _great) or derived forms (e.g. _încep, _începem, _început, _începe beside _înceap). In the baseline, it is more likely to find neighbors with no apparent relation, such as _erzählen ‘tell’ and _stemmen ‘hoist’ or ‘accomplish’ beside _America. This suggests that ANN embeddings help a model to better organize its subword embedding space into coherent, semanticallyrelated subspaces.
We quantify this trend by labeling each subword according to the language in which it is most frequently attested. In the baseline model, we find that on average only 2.7 of a subword’s 6 nearest neighbors come from the same language as that subword. This average rises to 3.6 in the ANN model, demonstrating that ANN training significantly increases the number of same-language neighbors on average.
In the ANN model, a few rare subwords (?, ž, ć) are disproportionately common among the nearest neighbors of many other subwords. We speculate that these tokens may act as pivots for information to flow between their many neighbours. Their high centrality means that these tokens provide avenues for information to flow between a large number of subwords, even those which never occur in sentences together. Because these tokens are rare, there is also very little penalty for the model to “corrupt” their representations with information from neighboring subwords.

heading: 5 Other Related Work
text: A vast body of work addresses zero-shot translation. Most methods focus on producing languageagnostic encoder outputs (Pham et al., 2019). Wei et al. (2021) introduce multilingual contrastive learning, while Yang et al. (2021) adopt auxiliary target language prediction. To enable the input tokens to be positioned without constraints, Liu et al. (2021) eliminate the residual connections within a middle layer of the encoder. Yang et al. (2022); Gu and Feng (2022) employ optimal transport to improve contextual cross-alignments, in contrast to our method which performs soft, non-contextual alignment between subwords in the continuouslyupdating embedding space. Other methods extend the training data using monolingual data (AlShedivat and Parikh, 2019) to pretrain the decoder (Gu et al., 2019), and random-online backtranslation (Zhang et al., 2020). Lin et al. (2021); Reid and Artetxe (2022) use dictionary based alignments to produce pseudo-cross-lingual sentences. Other approaches that enhance token level representations include multiple subword segmentations (Wu et al., 2020; Kambhatla et al., 2022a), enciphered source text (Kambhatla et al., 2022b) and stroke sequence modelling (Wang et al., 2022). While all these techniques rely on multilingual training paradigm for machine translation, they either rely on external data and use explicit augmentations. We do not
use any external data or explicit alignments and our model can be trained end-to-end like a regular multilingual model.

heading: 6 Conclusion
text: We described a novel approach to harness nearest neighbors at the token level and learn nearestneighbour informed word embeddings for every word in a source language for many-to-many multilingual translation. Our experiments show that this simple yet effective approach results in consistently better zero-shot translations across multiple multilingual datasets. Additionally, our model produces translations in the right target language with high accuracy. Our analysis shows that our model learns to organize subwords into semantically-related neighborhoods, and reduces the separation between monolingual subspaces in the embedding space.
Limitations
While our method is effective in zero-shot settings, we find that it has limited implications in supervised settings. This is because improving zero-shot translation presents a tug-of-war between language-agnostic and language-specific representations, each of which has a distinct effect on the model. Another major downside is reduced training speed relative to the baseline many-to-many model. We note that this is an artifact of the agreement loss (KLDiv.) which entails two forward-passes for each update. Finally, in the present work, we compute k-NNs for every source word in a sentence. Although this has yielded strong results, we would like to explore a more explainable setting where k-NNs can be applied to specific source words. We leave such explorations to future work.

heading: Acknowledgements
text: We would like to thank the anonymous reviewers for their helpful comments. The research was partially supported by the Natural Sciences and Engineering Research Council of Canada grants
NSERC RGPIN-2018-06437 and RGPAS-2018522574 and a Department of National Defence (DND) and NSERC grant DGDND-2018-00025 to the third author, and by an NSERC award CGSD3547773-2020 to the second author.

heading: A Reproducibility Details
text: A.1 Data
IWSLT17 (Cettolo et al., 2012) is an Englishcentric dataset7 totalling 1.8M parallel sentences. It has 8 supervised directions to and from German, Italian, Dutch and Romanian, each with about 220,000 parallel sentences, and 12 zero-shot directions. We use the official validation and test sets.
Ted59 (Qi et al., 2018) is a massively multilingual English-centric dataset8 with 116 translation directions totalling 10.8M parallel sentences. The imbalanced data—from 0.25M to just 2000 parallel samples for some language pairs—makes it ideal to study the effects of our method. Following (Aharoni et al., 2019; Raganato et al., 2021) we evaluate on 16 supervised pairs (Azerbaijani, Belarusian, Galician, Slovak, Arabic, German, Hebrew, and Italian to and from English) and 4 zero-shot (Arabic Ø French, Ukranian Ø Russian). Note that of these languages, Azerbaijani, Belarusian, Galician, and Slovak are low resource with only 5.9k, 4.5k, 10k and 61.5k paralle samples to/from English.
All settings and baselines use sentencepiece9 for subword tokenization using byte-pair encodings (BPEs; Sennrich et al. 2016) with 32000 merge operations.
A.2 Model and Hyperparameters
All models follow the basic configuration of Vaswani et al. (2017), using the fairseq toolkit (Ott et al., 2019) in PyTorch. This includes 6 layers of encoder and eecoder each with 512 dim and 2048 feed-forward dimension. The 512 dim word embedding layer has a vocabulary size of 32000. All word-embeddings in the model (encoder, decoder input/output) are shared, although the latent embedding layer alone is specific to encoder only. This implies that any updates to the actual embedding layer because of k-NN tokens also impacts the decoder.
The attentional latent semantic representation layer has 512 dim (same as the embedding layer) and a size N of 1000 for IWSLT17 (smaller dataset) and 5000 for TED59 (larger dataset). We did not tune this hyperparameter and chose the values based on the size of the datasets. This implies that this layer adds 0.5M trainable parameters to
7https://wit3.fbk.eu/2017-01 8github.com/neulab/word-embeddings-for-nmt 9https://github.com/google/sentencepiece
the IWSLT17 model and 2.5M parameters to the TED59 model. However, note that the total trainable parameters are still much lower than that of the baselines – this because our models have shared embedding layers.
We use the Adam optimizer with inverse square root learning scheduling and 6k warm steps, lr “ 0.0007 and dropout of 0.3 (IWSLT17), or 10k warmup steps, lr “ 0.005 and dropout of 0.2 (TED59). The batch size is 4096 tokens for each of four A100 GPUs.
We use ScANN (Guo et al., 2020) for efficient ANN search10 with k “ 3. To increase training speeds, we cache each subword’s ANNs for 400 iterations before recomputing them. We only (peridocally) cache subword IDs: the embedding EMBµp¨q is always computed directly from Wemb. The value of λ is set to 0.5 (Equation 1). We follow Kambhatla et al. (2022b) to set the values of α1, α2 to 1, and β to 5 (Equation 5).
Evaluation. For evaluation, all translations are generated with beam size 5. We report casesensitive BLEU scores (Papineni et al., 2002) using sacreBLEU11 (Post, 2018). We report detokenized BLEU for IWSLT17 and tokenized BLEU for TED59 for fair comparison with prior work (Aharoni et al., 2019; Raganato et al., 2021).

heading: B Nearest Neighbor Examples
text: See Table 5.
10https://github.com/google-research/ google-research/tree/master/scann . We use asymmetric hashing with 2-dimensional blocks and a quantization threshold of 0.2, and re-order the top 100 ANN candidates.
11case:mixed|eff:no|tok:13a|smooth:exp|version:2.3.1
Su bw
or d
N ea
re st
N ei
gh bo
rs (B
as el
in e)
N ea
re st
N ei
gh bo
rs (O
ur s)
_F is
ch _fi
nd et
œ _c
he st
ii _N
et z
fis ch
m ög
lic hk
ei te
n er
ei fis
ch _p
es ca
_fi sh
_F is
ch ž
sc ha
ft ho
od er
un g
un ge
n ga
at _g
eh ab
t sc
ha ft
w ür
di g
lic hk
ei t
sh ip
ni ss
e äg
lic h
sc ha ft th e tje pe d ow n as ta by _s ol am en te te ch th er th by th e ? _t he is ce iz ăm _e rz äh le n ”& _g eh ab t oa ra _y ou r _t he ir _o ur _t he ć ž _M us ic m at _c ar to on ho od _c on ne ss io ne zi a _ş co al a _m us ic a _m us ic ž _M us ic di ne ć _p ic io r _s fâ rş it _p la at je _m es aj _t er en _a vi on _g eh ab t _c or pu l _b ra t , _p ag in ă _p ic io r ž ? er n ei en ie re er en er un g _t en ut o _g eh ab t ur en un ge n er t er en st er n er n _î nc ea pă _p op ol ar e _c on do tto _m iş că _b ek ijk en _c re as că _c re ez e _g ep ub lic ee rd _î nc ep _î nc ep em _î nc ep ut _î nc ep e m uo
vo no
_d em
oc ra
t ,ia an
al is
i _î
nt , el
eg e
_p op
ol ar
e iz
ăm _ş
co al
a _d
es ha
lb _t
er or
is m
m uo
vo no
_d em
oc ra
tic di
ne _b
io di
ve rs
ità ć
_p ur
e ri
sc he
_g io
va ne
_a pp
en a
_t ar
e \u
0e 22
_a ve
ss e
_s em
pl ic
em en
te _u
ni qu
e _t
ra gi
c _c
om pl
et e
_s ol
e _p
ur e
_g en
om ic
_fi na
nz ia
Â _p
op ol
ar e
_r ăs
pâ nd
i _g
en om
en m
ög lic
hk ei
te n
_e le
ct ro
ni c
_g en
om e
_r ob
ot ic
ž _g
en et
ic _g
en om ic _A bb ia m o _p er ci ò _g eh ab t _s pu ne m _c on do tto iz ăm _a ve ss e ab bi am o m m o ia m o _A bb ia m o _a bb ia m o ć iz ăr i am en to is ie re n ie ru ng iz za zi on e _r ăs pâ nd i iz ar e iz ar e ită ţi aţ ie iz ăr i m uo vo no ne lli _n eg at iv e _a lte le az io ni ie re _b un e _e no rm ou s oa se _i lle ga l _a lte rn at iv e _e vi l _p os iti ve _n eg at iv _n eg at iv e _t ak e _s ol am en te _g em ac ht _s pi ng e _a cc om pa gn a _p re so _t en ut o _t ak es _t ak en _t ak in g _t oo k ć _t ak e _m uz ie k _p er co rs o _b es ef _o nd er w ijs _e rz äh le n _v re ug de oa ra _m us ic m uo vo no _M us ik _m us ic a _m uz ie k ć _K ar te _B ib lio th ek _l av or i st ra te gi e _c he st ii _c if re ka ar t _W ei se _S pr ac he _c ar ta _m on ta gn a kj es _K ar te _f un ct , io na _m iş că _m un ci m at ig _r ea liz a _f un ct , ie _f un ct , io na _f un cţ io ne az ă _f un ct , io ne az ă _f un zi on a _f un cţ io na _f un ct , io na ć _n aţ io na l _p op ol ar e ie re _b az ă _c on do tto _e se nţ ia l _p ol iti c ju is t _r ur al äg lic h _N at io na l _n aţ io na l _n at io na l _A m er ic a ta te _e rz äh le n _s te m m en di ne _f ac ul ta te _c he st iu ne _U SA _A si a _A fr ic a _A m er ic an _A m er ic a ć
Ta bl
e 5:
A pp
ro xi
m at
e ne
ar es
tne
ig hb
or s
fo ra
sa m
pl e
of su
bw or
ds ,c
om pu
te d
w ith
(r ig
ht )a
nd w
ith ou
t( le
ft )A
N N
tr ai
ni ng
.

heading: Future Work
text: We leave such explorations to future work