\documentclass[english]{jnlp_3.7}
\usepackage{jnlpbbl_2.2}
\usepackage[dvipdfmx]{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{url}
\usepackage{hangcaption_jnlp}
\usepackage{array}
\allowdisplaybreaks


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% author's definition
\usepackage{algorithm}
\usepackage{algpseudocode}
%%%%\usepackage{amsmath}
%%%%\usepackage{amssymb}
%%%%\usepackage{amsfonts}
\usepackage{bm}
\usepackage{bbm}
\usepackage{booktabs}
%%%%\usepackage{arydshln}
\usepackage{multirow}
%%%%\usepackage[subrefformat=parens]{subcaption}
\makeatletter
\let\MYcaption\@makecaption
\makeatother
\usepackage[subrefformat=parens]{subcaption}
\captionsetup{compatibility=false}
\makeatletter
\let\@makecaption\MYcaption
\makeatother
\setlength{\captionwidth}{\textwidth}
\usepackage{tabularx}
\usepackage{xspace}
%%%%\usepackage{udline}
%%%%\newcommand\blfootnote[1]{%
%%%%  \begingroup
%%%%  \renewcommand\thefootnote{}\footnote{#1}%
%%%%  \addtocounter{footnote}{-1}%
%%%%  \endgroup
%%%%}
\DeclareMathOperator*{\argmax}{argmax}
\DeclareMathOperator*{\argmin}{argmin}
\newcommand{\knn}{$k$NN\xspace}
\newcommand{\knnmt}{$k$NN-MT\xspace}
\newcommand{\pknn}{p_{k\mathrm{NN}}}
\newcommand{\pmt}{p_\mathrm{MT}}
\newcommand{\knearestneighbor}{$k$ nearest neighbor\xspace}
\newcommand{\nnearestneighbor}{$n$ nearest neighbor\xspace}
\newcolumntype{C}{>{\centering\arraybackslash}X}
\newcolumntype{R}{>{\raggedright\arraybackslash}X}
\newcolumntype{L}{>{\raggedleft\arraybackslash}X}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


\Volume{31}
\Number{2}
\Month{June}
\Year{2024}
%%%%\Issuetitle{Special Issue: XXXXXXXX}

\received{2023}{9}{1}
\revised{2023}{12}{1}
\accepted{2024}{1}{10}

\setcounter{page}{374}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\esubject{General Paper}
\etitle{Subset Retrieval Nearest Neighbor Machine Translation}
\eauthor{Hiroyuki Deguchi\affiref{NAIST}\affiref{NICT} \and Taro Watanabe\affiref{NAIST} \and Yusuke Matsui\affiref{UT} \and Masao Utiyama\affiref{NICT} \and Hideki Tanaka\affiref{NICT} \and Eiichiro Sumita\affiref{NICT}} 
\eabstract{$k$ nearest neighbor machine translation ($k$NN-MT) (Khandelwal et~al. 2021) boosts the translation quality of trained neural machine translation (NMT) models by incorporating an example search into the decoding algorithm. However, decoding is seriously time-consuming, that is, roughly 100 to 1,000 times slower than that of standard NMT, because neighbor tokens are retrieved from all the target tokens of parallel data in each timestep. In this paper, we propose ``Subset $k$NN-MT'', which improves the decoding speed of $k$NN-MT using two methods: (1) retrieving neighbor target tokens from a subset that is the set of neighbor sentences of the input sentence, not from all sentences, and (2) efficient distance computation technique suitable for subset neighbor search using a look-up table. Our subset $k$NN-MT achieved a speed-up of up to 134.2 times and an improvement in the BLEU score of up to 1.6 compared with those of $k$NN-MT in the WMT'19 De-En translation task, domain adaptation tasks in De-En and En-Ja translations, and the Flores101 multilingual translation task.}
\ekeywords{Neural Machine Translation, $k$ Nearest Neighbor Search}

\headauthor{Deguchi et al.}
\headtitle{Subset Retrieval Nearest Neighbor Machine Translation}

\affilabel{NAIST}{}{Nara Institute of Science and Technology (NAIST)}
\affilabel{NICT}{}{National Institute of Information and Communications Technology (NICT)}
\affilabel{UT}{}{The University of Tokyo}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{document}

\maketitle

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Introduction}

Neural machine translation (NMT) \cite{sutskever-etal-2014-sequence,bahdanau-etal-2015-neural,luong-etal-2015-effective,wu-etal-2016-google,vaswani-etal-2017-attention} has achieved a state-of-the-art performance and has become the focus of several studies.
Recently, \knearestneighbor machine translation (\knnmt) \cite{khandelwal-2021-nearest} has been proposed, which addresses the problem of the degradation of translation quality in out-of-domain data by incorporating an example search into the decoding algorithm.
\knnmt stores translation examples as a set of key--value pairs called ``datastore'' and retrieves \knearestneighbor target tokens in decoding.
This method improves the translation quality of NMT models without additional training.
However, decoding is seriously time-consuming, that is, roughly 100 to 1,000 times slower than that of standard NMT, because neighbor tokens are retrieved from all the target tokens of parallel data in each timestep.
In particular, in a realistic open-domain setting, \knnmt may be significantly slower because it needs to retrieve neighbor tokens from a large datastore that covers various domains.

We propose ``subset \knnmt'', which improves the decoding speed of \knnmt through two methods:
(1) retrieving neighbor target tokens from a subset that is the set of neighbor sentences of the input sentence, not from all sentences, and
(2) efficient distance computation technique suitable for subset neighbor search using a look-up table.
When retrieving neighboring sentences for a given input, we can employ arbitrary sentence representations, e.g., pre-trained neural encoders or term frequency--inverse document frequency (TF-IDF) vectors, to reduce the \knn search space.
When retrieving target tokens in each decoding step, the search space in the subset \knnmt varies depending on the input sentence; therefore, the clustering-based search methods used in the original \knnmt cannot be used.
Accordingly, we use asymmetric distance computation (ADC) \cite{jegou-etal-2011-product} in a subset neighbor search.

Our subset \knnmt achieved a speed-up of up to 134.2 times and an improvement in the BLEU score of up to 1.6 compared with those of the \knnmt in the WMT'19 German-to-English general domain translation task, domain adaptation tasks in German-to-English and English-to-Japanese translations
with open-domain settings, and the Flores101 multilingual translation task.\footnote{This is an extension of our prior work \cite{deguchi-etal-2023-subset}, and includes additional sections: implementation details, experiments for multilingual translation, and additional analyses.}
Our implementation, including both \knnmt and subset \knnmt, is available on GitHub.\footnote{\url{https://github.com/naist-nlp/knn-seq}}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{\knnmt}

\knnmt \cite{khandelwal-2021-nearest} retrieves the \knearestneighbor target tokens in each timestep, computes the \knn probability from the distances of the retrieved tokens, and interpolates the probability with the model prediction probability.
This method consists of two steps: (1) datastore creation, which creates key--value translation memory, and (2) generation, which calculates the output probability according to the nearest neighbors of the cached translation memory.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Datastore Construction}

A typical NMT model is composed of an encoder that encodes a source sentence $\bm{x} = (x_{1}, x_{2}, \ldots, x_{|\bm{x}|}) \in \mathcal{V}_{X}^{|\bm{x}|}$ and a decoder that generates target tokens $\bm{y} = (y_{1}, y_{2}, \ldots, y_{|\bm{y}|}) \in \mathcal{V}_{Y}^{|\bm{y}|}$, where $|\bm{x}|$ and $|\bm{y}|$ are the lengths of sentences $\bm{x}$ and $\bm{y}$, respectively, and $\mathcal{V}_{X}$ and $\mathcal{V}_{Y}$ are the vocabularies of the source and target languages, respectively.
The $t$-th target token $y_{t}$ is generated according to its output probability $P(y_{t} | \bm{x}, \bm{y}_{<t})$ over the target vocabulary, calculated from the source sentence $\bm{x}$ and generated target tokens $\bm{y}_{<t}$.
The \knnmt stores pairs of $D$-dimensional vectors and tokens in a datastore, represented as key--value memory $\mathcal{M} \subseteq \mathbb{R}^D \times \mathcal{V}_{Y}$.
The key ($\in \mathbb{R}^D$) is an intermediate representation of the final decoder layer obtained by teacher forcing a parallel sentence pair $(\bm{x}, \bm{y})$ to the NMT model,  and the value is a ground-truth target token $y_{t}$.
The datastore is formally defined as a set of tuples as follows:
\begin{equation}
    \mathcal{M} = \{ (f(\bm{x}, \bm{y}_{<t}), y_{t}) \mid (\bm{x}, \bm{y}) \in \mathcal{D}, 1 \leq t \leq |\bm{y}| \},
\end{equation}
where $\mathcal{D} = \{ (\bm{x}^i, \bm{y}^i) \}_{i=1}^{|\mathcal{D}|}$ denotes parallel data and $f: \mathcal{V}_{X}^{|\bm{x}|} \times \mathcal{V}_{Y}^{t-1} \to \mathbb{R}^D$ is a function that returns the $D$-dimensional intermediate representation of the final decoder layer from the source sentence and generated target tokens.
In our model, as in \shortciteA{khandelwal-2021-nearest}, the key is the intermediate representation before it is passed to the final feed-forward network.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Decoding}
During decoding, the \knnmt generates output probabilities by computing the linear interpolation between the \knn and MT probabilities, $\pknn$ and $\pmt$, respectively, as follows:
\begin{equation}
    P(y_t | \bm{x}, \bm{y}_{<t}) = \lambda \pknn(y_t | \bm{x}, \bm{y}_{<t}) + (1-\lambda) \pmt(y_t | \bm{x}, \bm{y}_{<t}),
    \label{eq:interpolation}
\end{equation}
where $\lambda$ is a hyperparameter for weighting the \knn probability.
Let $f(\bm{x},\bm{y}_{<t})$ be a query vector at timestep $t$.
The \knearestneighbor tokens of the query are searched from the datastore, and the top-$k$ key--value pairs are obtained.
The top $i$-th key and value in the \knearestneighbor tokens are $\bm{k}_i \in \mathbb{R}^D$ and $v_i \in \mathcal{V}_{Y}$, respectively.
Then $\pknn$ is defined as follows:
\begin{align}
    \pknn(y_t | \bm{x}, \bm{y}_{<t}) &= \frac{1}{Z} \sum_{i=1}^{k} \mathbbm{1}_{y_t = v_i} \exp{\left( \frac{-\lVert\bm{k}_i - f(\bm{x}, \bm{y}_{<t})\rVert_2^2}{\tau} \right)}, \label{eq:pknn-numerator} \\
    Z &= \sum_{i=1}^{k} \exp{\left( \frac{-\lVert\bm{k}_i - f(\bm{x}, \bm{y}_{<t})\rVert_2^2}{\tau} \right)},
    \label{eq:pknn-denominator}
\end{align}
where $\tau$ is the temperature for $\pknn$, and we set $\tau = 100$.

Note that this \knn search is seriously time-consuming\footnote{In our experiments on the WMT'19 German-to-English translation task, the datastore had 862M tokens, the vocabulary size was 42k, and the batch size was set to 12,000 tokens. While a normal Transformer translates 2,000 sentences in 7.5 s, and the \knnmt takes 2446.0~s. } \cite{khandelwal-2021-nearest}.
This is because \knn tokens are searched at each timestep when generating a target sentence.
For example, let $|\mathcal{M}|$ be one billion and $|\bm{y}|$ be 30.
If we naively search the \knn tokens, we need to calculate the distance between the query and key $|\mathcal{M}| \times |\bm{y}| = 30 \text{ billion}$ times, i.e., the time complexity is $\mathcal{O}(|\mathcal{M}| |\bm{y}|)$.
In other words, the speed problem of the \knnmt is caused by the large search space $|\mathcal{M}|$.\footnote{The \knnmt usually uses an approximate nearest neighbor algorithm instead of an exact search algorithm, but it is still time-consuming.}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 1
\begin{figure}[b]
\begin{center}
\includegraphics{31-2ia3f1.pdf}
\end{center}
\hangcaption{Illustration of our subset \knnmt, which reduces the search space of the \knnmt by retrieving the neighboring sentences of the input sentence. The green boxes in the sentence datastore indicate the neighboring examples of the input sentence. The search space is reduced from $|\mathcal{D}|$ sentences to $n (\ll |\mathcal{D}|)$ neighboring sentences.}
\label{fig:subset-knnmt}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Proposed Model: Subset \knnmt}

Our subset \knnmt drastically accelerates the vanilla \knnmt by reducing the \knn search space using sentence information.
In particular, subset retrieval (Section~\ref{sec:subset-knnmt}) retrieves neighboring sentences of the given input sentence and reduces the search space from all sentences to only neighboring sentences, as shown in Figure~\ref{fig:subset-knnmt}.
Additionally, we employ an efficient method to compute the distance between a query and keys using a look-up table when retrieving the \knn tokens from the reduced datastore (Section~\ref{sec:adc}).
While the original \knnmt employs a data structure and an algorithm optimized for the fixed search space, i.e., the full set of the datastore, our subset \knnmt employs an algorithm that efficiently searches the subset datastore which varies dynamically depending on the input sentence.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Subset Retrieval}
\label{sec:subset-knnmt}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Sentence Datastore Construction}
We construct a sentence datastore that stores pairs comprising a vector representation of a source sentence and key--value pairs of target tokens.
Concretely, a sentence datastore $\mathcal{S}$ is defined as follows:
\begin{align}
    \mathcal{S} &= \{ (h(\bm{x}^s), \mathcal{M}^{(s)}) \mid (\bm{x}^s, \bm{y}^s) \in \mathcal{D} \}, \\
    \mathcal{M}^{(s)} &= \{ (f(\bm{x}^s, \bm{y}^s_{<t}), y^s_{t}) \mid 1 \leq t \leq |\bm{y}^s| \}
\end{align}
where $h: \mathcal{V}_{X}^{|\bm{x}|} \to \mathbb{R}^{D'}$ represents a sentence encoder, which is a function that returns a $D'$-dimensional vector representation of a source sentence, and
$s \in \{ 1, \ldots, |\mathcal{D}|\}$ denotes the identifier of the $s$-th sentence pair in the parallel data.
Note that $\mathcal{M}^{(s)} \subset \mathcal{M}$ is the datastore created from only the $s$-th sentence pair.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Decoding}
At the beginning of decoding, the model retrieves the \nnearestneighbor sentences of the given input sentence by calculating the distances between the input sentence vector and the source sentence vectors of the sentence datastore $\mathcal{S}$.
Let $\mathcal{N} \subset \{ 1, \ldots, |\mathcal{D}|\}$ be the set of sentence numbers of the \nnearestneighbor sentences.
The nearest neighbor search space for the target tokens in the \knnmt is then drastically reduced by obtaining the datastore, as follows:
\begin{equation}
    \hat{\mathcal{M}} = \bigcup_{s \in \mathcal{N}} \mathcal{M}^{(s)} = \{ (f(\bm{x}^s, \bm{y}^s_{<t}), y^s_{t}) \mid s \in \mathcal{N}, 1 \leq t \leq |\bm{y}^s| \},
\end{equation}
where $\hat{\mathcal{M}} \subset \mathcal{M}$ is the reduced datastore for the translation examples from the \nnearestneighbor sentences.
During decoding, the model uses the same algorithm as \knnmt, except that $\hat{\mathcal{M}}$ is used as the datastore instead of $\mathcal{M}$.
The proposed method reduces the size of the nearest neighbor search space for the target tokens from $|\mathcal{D}|$ to $n$ ($\ll |\mathcal{D}|$) sentences.

Note that our method needs to store sentence representations in addition to the original datastore which stores token representations.
However, the number of sentences is usually one or two orders of magnitude smaller than the number of tokens, i.e., $|\mathcal{D}| \ll |\mathcal{M}|$; thus, the memory and storage usages will not be increased significantly.
Additionally, the subset \knnmt requires a neighboring sentence search, but it is not time-consuming because it only searches once before the decoding iterations, and the search space of the neighboring sentence search is smaller than that of all the target tokens, as mentioned above.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Efficient Distance Computation Using Look-up Table}
\label{sec:adc}

The subset \knnmt retrieves the \knearestneighbor target tokens using an efficient distance computation method that uses a look-up table.
In the original \knnmt, the inverted file index (IVF) is used for retrieving \knn tokens.
The IVF divides the search space into $N_\mathrm{list}$ clusters and retrieves tokens from neighboring clusters.
In contrast, in the subset \knnmt, the search space varies dynamically depending on the input sentence.
Thus, clustering-based search methods cannot be used; instead, we need to calculate the distance between a query and each key in the subset.
Accordingly, we use asymmetric distance computation (ADC) \cite{jegou-etal-2011-product} instead of the usual distance computation between floating-point vectors.
In ADC, the number of table look-ups is linearly proportional to the number of keys $N$ in the subset.
Therefore, it is not suitable for searching in a large datastore $\mathcal{M}$, but in a small subset $\hat{\mathcal{M}}$, the search is faster than the direct calculation of the L2 distance.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Product Quantization (PQ)}
The \knnmt datastore $\mathcal{M}$ may become too large because it stores the high-dimensional intermediate representations of all the target tokens of the parallel data.
For instance, the WMT'19 German-to-English parallel data, used in our experiments, contains 862 million tokens on the target side.
Therefore, if vectors are stored directly, the datastore would occupy 3.2 TiB when key vectors have 1024-dimensions,\footnote{$3.2~\text{TiB} \simeq 862.6\text{million tokens} \times 1024~\text{dimension} \times 32~\text{bit (float)} / 8~\text{bit per byte} / 1024^{4}$.}
 and this datastore would be difficult to load into the RAM.
To solve this memory problem, product quantization (PQ) \cite{jegou-etal-2011-product} is used in both \knnmt and our subset \knnmt, which includes both source sentence and target token search.

PQ splits a $D$-dimensional vector into $M$ sub-vectors and quantizes each $\frac{D}{M}$-dimensional sub-vector.
Codebooks are learned by the k-means clustering of key vectors in each subspace.
They are computed iteratively by (1) assigning the code of a key to its nearest neighbor centroid, and (2) updating the centroid of keys assigned to the code.
The $m$-th sub-space codebook $\mathcal{C}^{m}$ is formulated as
\begin{equation}
    \mathcal{C}^{m} = \{ \bm{c}_{1}^{m}, \ldots, \bm{c}_{L}^{m}\},\quad \bm{c}_{l}^{m} \in \mathbb{R}^{\frac{D}{M}}.
\end{equation}
In this study, the size of each codebook is set to $L=256$.
A vector $\bm{q} \in \mathbb{R}^D$ is quantized and its code vector $\bar{\bm{q}}$ is calculated as follows:
\begin{align}
    \bm{q} &= \left[{\bm{q}^1}^\top, \ldots, {\bm{q}^M}^\top \right]^\top,\quad \bm{q}^{m} \in \mathbb{R}^{\frac{D}{M}}, \\
    \bar{q}^{m} &= \argmin_{l} \| \bm{q}^m - \bm{c}_{l}^{m} \|_{2}^{2},\\
    \bar{\bm{q}} &= [\bar{q}^{1}, \ldots, \bar{q}^{M}]^{\top} \in \{1, \ldots, L\}^{M}.
\end{align}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 2
\begin{figure}[b]
\begin{center}
\includegraphics{31-2ia3f2.pdf}
\end{center}
\caption{Distance computation using asymmetric distance computation (ADC).}
\label{fig:adc}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Asymmetric Distance Computation (ADC)}

Our method efficiently computes the squared Euclidean distance between a query vector and quantized key vectors using ADC \cite{jegou-etal-2011-product} (Figure \ref{fig:adc} and Algorithm \ref{alg:adc-lookup}).
The ADC computes the distance between a query vector $\bm{q} \in \mathbb{R}^D$ and $N$ key codes
$\bar{\mathcal{K}} = \{ \bar{\bm{k}}_i \}_{i=1}^{N} \subseteq \{ 1, \ldots, L \}^M$.
First, the distance look-up table $\bm{A}^m \in \mathbb{R}^{L}$ is computed by calculating the squared Euclidean distance between a query $\bm{q}^m$ and the codes $\bm{c}_l^m \in \mathcal{C}^m$ in each sub-space $m$, as follows:
\begin{equation}
    \label{eq:adc-dt}
    A^m_l = \| \bm{q}^m - \bm{c}_{l}^{m} \|_{2}^{2}.
\end{equation}
Second, the distance between a query and each key $d(\bm{q}, \bar{\bm{k}}_i)$ is obtained by looking up the distance table as follows:
\begin{equation}
    \label{eq:adc-lu}
    d(\bm{q}, \bar{\bm{k}}_i) = \sum_{m=1}^{M} d_m(\bm{q}^m, \bar{k}_i^m) = \sum_{m=1}^M A^m_{\bar{k}_i^m}.
\end{equation}
A look-up table in each subspace, $\bm{A}^m \in \mathbb{R}^L$, consists of the distance between a query and codes.
The number of codes in each subspace is $L$, and the distance is scalar; therefore, $\bm{A}^m$ has $L$ distances.
The table look-up key is the code of the key itself, i.e., if the code of a key in the $m$-th subspace is $5$, the ADC looks-up $A^m_{5}$.
Using the ADC, the distance is computed only once\footnote{The direct distance computation requires $N$ times calculations according to $\lVert \bm{q} - \bm{k}\rVert^2$, i.e., the time complexity is $\mathcal{O}(ND)$. The ADC computes the distance only $L \ll N$ times, i.e., the time complexity for creating the distance table is $\mathcal{O}(L \times \frac{D}{M} \times M) = \mathcal{O}(LD)$, and it just looks-up the table $MN$ times in the constant time $\mathcal{O}(1)$. Therefore, the complexity is reduced from $\mathcal{O}(ND)$ to $\mathcal{O}(LD)$.} (Equation \ref{eq:adc-dt}).
The ADC does not decode PQ codes into $D$-dimensional key vectors;
therefore, it can compute the distance while maintaining the key in the quantization code, and the \knearestneighbor tokens are efficiently retrieved from  $\hat{\mathcal{M}}$.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% algo. 1
\begin{algorithm}[t]
\input{03algo01.tex}
\caption{~ADC look-up}
\label{alg:adc-lookup}
\end{algorithm}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Sentence Encoder}

In our subset \knnmt, various sentence encoder models can be employed.
The more similar sentences are extracted from $\mathcal{M}$, the more likely the subset $\hat{\mathcal{M}}$ comprises the target tokens useful for translation.
Hence, we need sentence encoders that compute vector representations whose distances are close for similar sentences.

In this study, we employ two types of representations: \emph{neural} and \emph{non-neural}.
We can employ pre-trained neural sentence encoders.
Although they are required to support the source language, we expect that the retrieved sentences are more similar than those in other encoders because we can use models that have been trained to minimize the vector distance between similar sentences \cite{reimers-gurevych-2019-sentence}.
An NMT encoder can also be used as a sentence encoder by applying average pooling to its intermediate representations.
This does not require any external resources, but the encoder is not trained from the supervision of sentence representations, explicitly.
Alternatively, we can also use non-neural models, such as TF-IDF.
However, it is not clear whether TF-IDF based similarity is suitable for our method.
This is because even if sentences with close surface expressions are retrieved, they do not necessarily have similar meanings and may not yield the candidate tokens required for translation.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Implementation Details}

We used \textsc{fairseq} \cite{ott-etal-2019-fairseq} to implement the \knnmt and subset \knnmt, and \textsc{faiss} \cite{johnson-etal-2019-billion} to retrieve the \knn tokens in the \knnmt and for neighbor sentence search in the subset \knnmt.
Existing \knn libraries, including \textsc{faiss}, and algorithms, such as IVF, that are used in the original \knnmt are designed for full search but not for subset search \cite{matsui-etal-2018-reconfigurable}; therefore, we implemented the subset search and ADC look-up using \textsc{PyTorch}.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Subset Caching}
Quantized key codes and the value tokens of the subset are read at the beginning of decoding and cached during decoding.
Therefore, a large billion-scale array is accessed only once during decoding.
Note that the subset depends only on the input sentence, and the cache size does not change with the beam size.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Distance Look-up in Beam Search Decoding}
During decoding by beam search, the queries in the beams have different representations because a query vector is computed depending on the generated target tokens.
Let $B$ be the beam size, and $\mathbf{Q} \in \mathbb{R}^{B \times D}$ be the queries.
Note that we regard $\mathbf{Q}_i$ as the column vector by transposing the $i$-th row of the matrix $\mathbf{Q}$, i.e., $\mathbf{Q}_i \in \mathbb{R}^D$, and $\mathbf{Q}^m_i$ as the $m$-th subspace of $M$ subspaces in $\mathbf{Q}_i$, i.e., $\mathbf{Q}^m_i \in \mathbb{R}^{\frac{D}{M}}$.
The distance table of a subspace in PQ is computed from a query ($\in \mathbb{R}^{\frac{D}{M}}$) and codebook $\mathcal{C}^m$; thus, the table $\mathbf{A}^m \in \mathbb{R}^{L \times B}$ is computed for each beam.
\begin{equation}
    \mathbf{A}^m_l = \left[ \lVert \mathbf{Q}^m_{1} - \bm{c}^m_l \rVert_2^2, \lVert \mathbf{Q}^m_{2} - \bm{c}^m_l \rVert_2^2, \ldots, \lVert \mathbf{Q}^m_{B} - \bm{c}^m_l \rVert_2^2 \right]^\top.
\end{equation}
In contrast, the keys in the subset $\{ \bar{\bm{k}}_1, \ldots \bar{\bm{k}}_N \} \subseteq \{ 1, \ldots, L \}^M$ are the same across beams because they are not changed by the generated target tokens.
Then, the distances between a key and the queries for each beam $\mathbf{d}(\mathbf{Q}, \bar{\bm{k}}) \in \mathbb{R}^{B}$ are obtained as follows:
\begin{align}
    \mathbf{d}(\mathbf{Q}, \bar{\bm{k}}) &= \left[ d(\mathbf{Q}_1, \bar{\bm{k}}), \ldots d(\mathbf{Q}_B, \bar{\bm{k}}) \right]^\top, \label{eq:beam-dist}\\
    d(\mathbf{Q}_i, \bar{\bm{k}}) & = \sum_{m=1}^M d_m(\mathbf{Q}^m_i, \bar{k}^m) = \sum_{m=1}^M \mathbf{A}^m_{\bar{k}^m}. \label{eq:beam-dist-one}
\end{align}
From Equation~\ref{eq:beam-dist} and \ref{eq:beam-dist-one}, the distance table is looked up $M \times B$ times at each timestep during decoding.
This look-up can be parallelized using \texttt{torch.gather()} in PyTorch.
However, to perform a parallel look-up, the keys must be replicated for each beam resulting in multiple copies proportional to the beam size.
To avoid increasing the memory usage of key vectors, we designed not to allocate new memory by copying multiple instances for each beam, but only to create a new view of the tensor using \texttt{torch.expand()}.
The number of keys is the number of target tokens in neighboring sentences, for example, 10,000.
Therefore, this technique is helpful because it saves memory usage even if the beam size is increased.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Experiments}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Setup}

We compared the translation quality and speed of our subset \knnmt with those of the conventional \knnmt in open-domain settings that assume that the domain of an input sentence is unknown.
The translation quality was measured using \textsc{sacreBLEU}\footnote{Signature: \texttt{|nrefs:1|case:mixed|eff:no|tok:13a|smooth:exp|version:2.3.1}} \cite{post-2018-call} and \textsc{COMET} \cite{rei-etal-2020-comet}.
The decoding speed was evaluated based on the number of tokens generated per second (tok/s) on a single NVIDIA V100~GPU.
The time measurement includes all the processes since the source tokens are given until the output sequence is obtained by beam search; that is, in the \knnmt, it includes the time to search the \knearestneighbor tokens for each timestep in addition to the forward computation of the NMT model, and in the subset \knnmt, it includes the time to search the neighboring sentences, look-up the distance table, etc.
The speed, in tok/s, was calculated by dividing the number of generated tokens by the time required to translate the entire test set.
We varied the batch size settings:
either 12,000 tokens (B$_{\infty}$), to simulate the document translation scenario, or a single sentence (B$_1$), to simulate the online translation scenario.
The beam size was set to 5, and the length penalty was set to 1.0.
In the tables showing results, the best score is emphasized in bold font and the second-best score is underlined.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{$k$ Nearest Neighbor Search}
In the \knnmt, we set the number of nearest neighbor tokens to $k=16$.
We used the approximate distance computed from the quantized keys instead of the full-precision keys in Equation~\ref{eq:pknn-numerator}
and \ref{eq:pknn-denominator}
 following the original \knnmt \cite{khandelwal-2021-nearest} implementation.
Both the \knnmt datastore and our sentence datastore used IVF and optimized PQ (OPQ) \cite{ge-etal-2014-optimized}.
OPQ rotates the vectors to minimize the quantization error of PQ.
The subset \knnmt datastore is not applied clustering because we need to extract the subset tokens.
In this datastore, the 1024-dimensional vector representation, i.e., $D=1024$, 
was reduced in dimensionality to 256 by principal component analysis (PCA), and these vectors were then quantized by PQ.
At the search time, a query vector was pre-transformed into 256-dimensions by multiplying with the PCA matrix, and then, the \knn target tokens were searched using the ADC.
The subset of a datastore can be loaded into GPU memory, as it is significantly smaller than the original \knnmt datastore; hence, we retrieved \knearestneighbor tokens from a subset on a GPU.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Sentence Encoder}
We compared four sentence encoders: LaBSE, AvgEnc, TF-IDF, and BM25.
LaBSE \cite{feng-etal-2022-language} is a pre-trained sentence encoder, fine-tuned from multilingual BERT.
AvgEnc is the average pooled encoder hidden vector of the Transformer NMT model, which is also used for translation.
TF-IDF \cite{jones-1972-statistical} and BM25 \cite{jones-etal-2000-probabilistic} compute vectors that weight the important words in a sentence.
We used the raw token count as the term frequency and applied add-one smoothing to calculate the inverse document frequency, where a sentence was regarded as a document.
We set $k_1 = 2.0$ and $b = 0.75$ in BM25 \cite{jones-etal-2000-probabilistic}.
Both the TF-IDF and BM25 vectors were normalized by their L2-norm, and their dimensionality was reduced to 256-dimensions using singular value decomposition.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{In-Domain Translation}
\label{sec:exp-wmt19}

We evaluated the translation quality and speed of the subset \knnmt in the WMT'19 De-En translation task (newstest2019; 2,000 sentences) and compared them with those of the original \knnmt \cite{khandelwal-2021-nearest} and other prior work \cite{meng-etal-2022-fast,martins-etal-2022-chunk}.
We used a trained Transformer big model implemented in \textsc{fairseq} \cite{ott-etal-2019-fairseq}
as the base machine translation (MT) model.\footnote{\url{https://dl.fbaipublicfiles.com/fairseq/models/wmt19.de-en.ffn8192.tar.gz}}
We constructed the datastore from the parallel data of the WMT'19 De-En news translation task.
We removed all the sentences of the parallel data shorter than one token and longer than 250 tokens.
We also removed all the sentences in which the sentence length in one language was more than 1.5 times longer than that in the other language; that is, the ratio of tokens between the source and target was $>1.5$.
The datastore contained 862.6 million target tokens obtained from 29.5 million sentence pairs.
The subset size was set to $n=512$.
For the fast \knnmt, we constructed additional source side datastores for each source token, and used \textsc{fast\_align} \cite{dyer-etal-2013-simple} to obtain source-to-target word alignment, following {\shortciteA{meng-etal-2022-fast}}.
We then retrieved the 512 nearest neighbor source tokens from the source side datastores for each input token in the decoding time of the fast \knnmt.
Note that the total size of the source side datastores is close to that of the \knnmt datastore; thus, the fast \knnmt consumes twice as much storage and memory as the original \knnmt.
In the chunk-based \knnmt, the chunk size was set to 16, and the hyperparameters that determine the interval of retrieval {$i_\mathrm{max}$} and {$i_\mathrm{min}$} were set to 2 and 16, respectively, following {\shortciteA{martins-etal-2022-chunk}}.
The details of the \knn indices are listed in Table~\ref{tab:hparam-index}.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 1
\begin{table}[b]
\input{03table01.tex}%
\caption{Details of \knn indexes. ``DS'' indicates ``Datastore''.}
\label{tab:hparam-index}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 2
\begin{table}[b]
\input{03table02.tex}%
\hangcaption{Results of translation quality and decoding speed in the WMT'19 De-En translation task. ``$h$:'' indicates the type of sentence encoder used. }
\label{tab:result-wmt19deen}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Table~\ref{tab:result-wmt19deen} presents the experimental results.
The table shows that, although \knnmt improved the BLEU score by 0.9 points from the base MT model without additional training, the decoding speed was 326.1 times and 51.7 times slower with the B$_{\infty}$ and B$_{1}$ settings, respectively.
In contrast, our subset \knnmt ($h$: LaBSE) was 111.8 times (with B$_{\infty}$) and 47.4 times (with B$_{1}$) faster than the \knnmt with no degradation in the BLEU score.
Te subset \knnmt ($h$: AvgEnc) achieved speed-ups of 92.7 times (with B$_{\infty}$) and 38.9 times (with B$_{1}$) with a slight quality degradation ($-$0.2 BLEU and $-$0.05 COMET), despite using no external models.
We also evaluated our subset \knnmt using non-neural sentence encoders ($h$: TF-IDF, BM25).
The results show that both TF-IDF and BM25 can generate translations with almost the same BLEU score and speed as neural sentence encoders.
The chunk-based \knnmt and fast \knnmt generated translations 4 and 15 times faster than the original \knnmt, respectively.
Note that the chunk-based \knnmt does not use any additional resources and the fast \knnmt needs to create additional source side datastores that consume large amounts of memory and storage, and requires a source-to-target word alignment tool.
The subset \knnmt uses a sentence encoder and creates the sentence datastore which has {$|\mathcal{D}| \ll |\mathcal{M}|$} sentence representations in addition to the \knnmt datastore.
In summary, this experiment showed that our subset \knnmt was two orders of magnitude faster than the \knnmt, and had the same translation quality.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 3
\begin{table}[b]
\input{03table03.tex}%
\caption{Datastore statistics in the domain adaptation task.}
\label{tab:stats-dom}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Domain Adaptation}
\label{sec:exp-domain}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{German-to-English}

We evaluated the subset \knnmt on out-of-domain translation in the IT, Koran, Law, Medical, and Subtitles domains \cite{koehn-knowles-2017-six,aharoni-goldberg-2020-unsupervised} with open-domain settings.
The datastore was constructed from parallel data by merging all the target domains and the general domain (WMT'19 De-En) assuming that the domain of the input sentences is unknown.
The datastore contained 895.9 million tokens obtained from the 30.8 million sentence pairs listed in Table {\ref{tab:stats-dom}\subref{tab:stats-dom-deen}}.
The NMT model was the same as that used in Section~\ref{sec:exp-wmt19}, trained from the WMT'19 De-En.
The subset size was set to $n=256$, and the batch size was set to 12,000 tokens, i.e., B$_\infty$.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 4
\begin{table}[b]
\input{03table04.tex}%
\hangcaption{Results of out-of-domain translation with open-domain settings. ``Avg'' denotes the average scores. ``BL'' and ``CM'' denote the BLEU and COMET scores, respectively.}
\label{tab:result-alldomain}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

The results are presented in Table~\ref{tab:result-alldomain}.
Compared with the base MT model, the \knnmt improved the translation quality in all the domains, but the decoding speed was much slower.
In contrast, the subset \knnmt generated translations faster than the \knnmt.
However, in the domain adaptation task, there were differences in translation quality between the use of neural sentence encoders and the use of non-neural sentence encoders.
The table shows that the use of non-neural sentence encoders (TF-IDF and BM25) caused a decrease in translation quality, whereas the use of neural sentence encoders (LaBSE and AvgEnc) did not.
In addition, compared with the \knnmt, our subset \knnmt with neural encoders achieved an improvement of up to 1.6 BLEU points on some datasets.
In summary, these results show that neural sentence encoders are effective in retrieving domain-specific nearest neighbor sentences from a large datastore.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{English-to-Japanese}

We also evaluated our model using English-to-Japanese translation.
We used a pre-trained Transformer big model trained using JParaCrawl v3 \cite{morishita-etal-2022-jparacrawl} and evaluated its translation quality on Asian Scientific Paper Excerpt Corpus (ASPEC) \cite{nakazawa-etal-2016-aspec} and Kyoto Free Translation Task (KFTT; created from Wikipedia's Kyoto articles) \cite{neubig-2011-kyoto}.
The datastore was constructed from the parallel data by merging ASPEC, KFTT, and the general domain (JParaCrawl v3), as shown in Table {\ref{tab:stats-dom}\subref{tab:stats-dom-enja}}.
Note that ASPEC contains 3 million sentence pairs, but we used only the first 2 million pairs for the datastore to remove noisy data, following \shortciteA{neubig-2014-forest}.
The datastore contained 735.9 million tokens obtained from 24.4 million sentence pairs.
The subset size was set to $n=512$, and the batch size was set to 12,000 tokens.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 5
\begin{table}[t]
\input{03table05.tex}%
\hangcaption{Results of out-of-domain translation in English-to-Japanese translation. The speed is measured with the B$_\infty$ setting.}
\label{tab:result-enja}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Table~\ref{tab:result-enja} presents the results.
These results show that the \knnmt improved the out-of-domain translation quality compared with the base MT model on language pairs other than German-to-English.
For English-to-Japanese, the subset \knnmt improved the decoding speed, but the subset \knnmt with TF-IDF and BM25 degraded the translation quality compared with the \knnmt.
However, the subset \knnmt still achieved higher BLEU scores than the base MT model without any additional training steps, and was two orders of magnitude faster than the \knnmt.

In summary, the subset \knnmt could achieve better translation quality than the base MT model in a tradeoff for a slowdown to approximately 40\% of the base MT model in open-domain settings, whereas the original \knnmt decreased the decoding speed to less than 1\% that of the base MT model.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Multilingual Translation}

We also evaluated the multilingual translation quality across 11 
\pagebreak
translation directions using the Flores-101 dataset \cite{goyal-etal-2022-flores}, which was created from English Wikipedia.
We used the Flores101-M2M100\footnote{\url{https://dl.fbaipublicfiles.com/flores101/pretrained_models/flores101_mm100_615M.tar.gz}} model with 615 million parameters, which was extended from M2M \cite{fan-etal-2021-beyond} to support languages included in Flores-101 by training from OPUS data.
The datastore of each language pair was constructed from  CCMatrix \cite{schwenk-etal-2021-ccmatrix} which is extracted from Common Crawl.
Note that each datastore was created from the parallel data of the language pair to be translated.
In this experiment, we employed LaBSE and AvgEnc for the sentence encoder.
We tuned the subset size $n$ to maximize the BLEU score among $\{256, 512, 1024, 2048\}$ in the validation sets of the En-Ja and Ja-En translations and set $n=2048$.
The batch size was set to 12,000 tokens, i.e., B$_\infty$.
We used the \texttt{flores101} tokenizer implemented in \textsc{sacreBLEU}\footnote{Signature: \texttt{nrefs:1|case:mixed|eff:no|tok:flores101|smooth:exp|version:2.3.1}} to calculate the BLEU scores.

Table~\ref{tab:result-multilingual} presents the results of the multilingual translation.
In the table, ``Avg'' denotes the average of the scores.
The results showed that both the \knnmt and subset \knnmt improved the translation quality in multilingual translation.
In Avg, the subset \knnmt degraded the translation quality compared with the \knnmt, but was 19 times faster measured in terms of tok/s.
Comparing the decoding speed for each language pair with that of the \knnmt, the subset \knnmt was 16.8 times faster in Uk-Ja with the smallest datastore and 134.2 times faster in En-Zh with the largest datastore.
In summary, this experiment shows that the subset \knnmt is more effective when the datastore is larger because a larger datastore will be reduced more examples from the search space by our subset retrieval.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 6
\begin{table}[t]
\input{03table06.tex}%
\hangcaption{Results of multilingual translation. The speed is evaluated with B$_\infty$. ``BL'' and ``CM'' denote the BLEU and COMET scores, respectively. The number next to a language name pair indicates the size of the datastore, i.e., the number of target tokens in the parallel data. 
}
\label{tab:result-multilingual}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 7
\begin{table}[t]
\input{03table07.tex}%
\caption{Translation examples in the medical domain.}
\label{tab:case-medical}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Discussion}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Case Study: Effects of Subset Search}
\label{sec:case}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Effective Cases of Subset \knnmt}
The translation examples in the medical domain are shown in Table \ref{tab:case-medical}, and the search results for the top-3 nearest neighbor sentences are shown in Table~\ref{tab:case-search-results-medical}.
In the table, the subset \knnmt results were obtained using the LaBSE encoder.
Table~\ref{tab:case-medical} shows that the subset \knnmt  correctly generated the medical term ``Co-administration''.
The results of the nearest neighbor sentence search (Table \ref{tab:case-search-results-medical}) show that ``Co-administration'' was included in the subset.
In detail, there were 30 cases of ``Co-administration'' and no case of ``A joint use'' in the entire subset consisting of $n=256$ neighboring sentences.
The base MT and kNN-MT had the subwords of ``Co-administration'' in the candidates; however, the subwords of ``A joint use'' had higher scores.
Table~\ref{tab:case-nll} presents the negative log-likelihood (NLL) of the first three tokens and their averages for each model.
The second token of subset \knnmt, ``-'' (hyphen), had a significantly lower NLL than the other tokens.
The numbers of ``joint'' and ``-'' in the subset were 0 and 101, respectively, and the \knearestneighbor tokens were all ``-'' in the subset \knnmt.
Therefore, the NLL was low because $\pknn(\text{``-''}) = 1.0$; hence, so the joint probability of a beam that generated the sequence ``Co-administration'' was higher than that for ``A joint use''.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 8
\begin{table}[b]
\input{03table08.tex}%
\hangcaption{Top-3 neighbor sentences of our subset \knnmt in Table \ref{tab:case-medical}. ``S-'' and ``T-'' denote the top-$n$ neighbor source sentences and their translations, respectively.}
\label{tab:case-search-results-medical}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 9
\begin{table}[b]
\input{03table09.tex}%
\hangcaption{Negative log-likelihood (NLL) of the first three tokens and their average in the case of Table~\ref{tab:case-medical}. Note that a smaller NLL indicates a higher probability.}
\label{tab:case-nll}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

In summary, the proposed method can retrieve more appropriate words by searching for a subset consisting only of neighboring cases when translation examples from the target domain are contained in the datastore.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 10
\begin{table}[b]
\input{03table10.tex}%
\caption{Japanese-to-English translation examples in the Flores-101 multilingual translation task.}
\label{tab:case-flores101}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 11
\begin{table}[b]
\input{03table11.tex}%
\hangcaption{Top-3 neighbor sentences of our subset \knnmt in Table \ref{tab:case-flores101}. ``S-'' and ``T-'' denote the top-$n$ neighbor source sentences and their translations, respectively.}
\label{tab:case-search-results-flores101}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 12
\begin{table}[b]
\input{03table12.tex}%
\hangcaption{Translation examples containing ``termites'' in the Japanese-to-English datastore constructed from CCMatrix.}
\label{tab:dstore-termites}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Ineffective Cases of Subset \knnmt}

Japanese-to-English translation examples in the Flores-101 multilingual translation task are shown in Table \ref{tab:case-flores101}, and the search results of the top-3 nearest neighbor sentences are shown in Table~\ref{tab:case-search-results-flores101}.
The subset \knnmt results listed in the table were obtained using the LaBSE encoder.
Table~\ref{tab:case-flores101} shows that the subset \knnmt incorrectly generated animal names, ``アリ'' $\rightarrow$ ``ants'' and ``シロアリ'' $\rightarrow$ ``termites''.
The results of the nearest neighbor sentence search (Table \ref{tab:case-search-results-flores101}) showed that neither word was included in the subset.
In detail, there were no cases of ``ants'' and ``termites'' in the entire subset consisting of $n=2048$ neighbor sentences.
Table~\ref{tab:dstore-termites} shows translation examples containing ``termites'' in the datastore.
Compared with the input sentences in Table~\ref{tab:case-flores101}, the topics of the sentences containing  ``termites'' were not matched.
By contrast, as the \knnmt searches on a token-by-token basis, it is also possible to retrieve target tokens from translation examples with different topics.
In summary, the subset \knnmt degrades the translation quality compared with the \knnmt when the neighboring sentences contain no correct words.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Diversity of Subset Sentences}

We hypothesized that the noise introduced by sentence encoders would cause a difference in the accuracy.
For example, if a sentence search is insufficiently accurate, translation examples related to the input sentence cannot be retrieved.
In addition, we can expect consistency of translations by retrieving based not only on the semantic similarity of sentences but also on style and other aspects.
From the results of Section~{\ref{sec:case}} and Table~{\ref{tab:case-nll}}, one characteristic subword that frequently occurs in the \knn changed the order of the beams, which contributed to the improvement of the translation quality of the subset \knnmt.
Thus, if the subset includes only vocabulary that is more relevant to translation, the translation accuracy may be improved.

In this section, we investigate whether a better sentence encoder reduces the noise injected into a subset.
We investigated the relationship between the vocabulary diversity in this subset and the translation quality in the medical domain.
As an output sentence is affected by the subset, we measured the unique token ratio of both the source and target languages in the subset as the diversity as follows:
\begin{equation}
    \frac{\text{number of unique tokens}}{\text{number of subset tokens}}.
\end{equation}

Table~\ref{tab:diversity-h} presents the BLEU score and unique token ratio for the various sentence encoders, in which ``source'' and ``target'' indicate the diversity of the neighbor sentences on the source and target sides, respectively.
The results show that the more diverse the source side, the more diverse the target side.
Furthermore, the low the vocabulary diversity of both the source and target languages in the subset, the higher the BLEU score.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 13
\begin{table}[b]
\input{03table13.tex}%
\hangcaption{BLEU score and unique token ratio in the subset obtained using each sentence encoder in the medical domain.}
\label{tab:diversity-h}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

We also investigated the relationship between the sentence encoder representation and the BLEU scores.
In particular, we evaluated the translation quality when noise was injected into a subset by retrieving $n$ sentences from outside the nearest neighbor.
To test our hypothesis, we experimented with two artificially created subsets.
One is ``Bottom of $2n$'', in which the $n$ furthest sentences of the $2n$ neighbor sentences, which simulates the $n$ nearest neighbor sentences cannot be retrieved.
The other is ``Random of $2n$'', in which $n$ sentences are randomly selected from the $2n$ neighbor sentences, i.e., it can be regarded as a subset that mixes roughly half of the ``Top'' subset and noise examples from half of ``Bottom of $2n$''.
Thus, ``Random of $2n$'' uses more similar examples than ``Bottom of $2n$''.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 14
\begin{table}[t]
\input{03table14.tex}%
\hangcaption{BLEU score and unique token ratio in the subset obtained using different $n$-selection methods in the medical domain.}
\label{tab:diversity-nselection}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Table~\ref{tab:diversity-nselection} presents the results of various $n$-selection methods when LaBSE was used as the sentence encoder.
In the table, ``Top'' indicates the \nnearestneighbor sentences.
``Bottom of $2n$'' and ``Random of $2n$'' had higher diversity than ``Top'' on both the source and target sides, and the BLEU scores were correspondingly lower.
In addition, ``Random of $2n$'' achieved a higher BLEU score than ``Bottom of $2n$'' with a lower unique ratio.
 resulting in lower translation quality than ``Top''.
These experiments showed that a sentence encoder that appropriately calculates similarity can reduce noise and prevent the degradation of translation quality because the subset consists only of similar sentences.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Analysis of Decoding Speed}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Efficiency of ADC}
The subset \knnmt computes the distance between a query vector and key vectors using the ADC, as described in Section~\ref{sec:adc}.
The efficiency of the ADC in the WMT'19 De-En is presented in Table \ref{tab:result-adc}.
The results show that ``w/ ADC'' is roughly 4 to 5 times faster than ``w/o ADC''.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Effect of Parallelization}
The method and implementation of our subset \knnmt are designed for parallel computing.
We measured the translation speed for different batch sizes in WMT'19 De-En.
%%%%Figure~\ref{fig:speed}\subref{fig:speed-bsz} shows that the subset \knnmt ($h$: LaBSE) was two orders of magnitude faster than the \knnmt even when the batch size was increased.
    Figure~\ref{fig:speed}(a) shows that the subset \knnmt ($h$: LaBSE) was two orders of magnitude faster than the \knnmt even when the batch size was increased.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{Subset Size}
\label{sec:speed}
We measured the translation speed for different subset sizes, i.e., the number of \nnearestneighbor sentences in WMT'19 De-En.
%%%%Figure~\ref{fig:speed}\subref{fig:srck-speed} shows the translation speed of the subset \knnmt ($h$: LaBSE).
    Figure~\ref{fig:speed}(b) shows the translation speed of the subset \knnmt ($h$: LaBSE).
The subset \knnmt was two orders of magnitude faster than  the \knnmt, even when the subset size was increased.
The results also show that the speed decreased from $n=256$ compared with that of the base MT model.
We also observed that 71.7\% of the time was spent searching for \knn tokens from the subset when $n=2048$.
Although the ADC look-up search is slow for a large datastore, it is fast for \knn search when the subset size $n$ is not large \cite{matsui-etal-2018-reconfigurable}, e.g., $n=512$.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 15
\begin{table}[t]
\input{03table15.tex}%
\hangcaption{Efficiency of ADC in WMT'19 De-En. The results show the number of tokens generated per second, i.e., $\uparrow$tok/s, with the B$_\infty$ setting.}
\label{tab:result-adc}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 3
\begin{figure}[t]
\begin{center}
\includegraphics{31-2ia3f3.pdf}
\end{center}
%%%%\label{fig:speed-bsz}
%%%%\label{fig:srck-speed}
%%%%\label{fig:srck-bleu}
\hangcaption{Translation speed for different batch sizes, and subset sizes and translation quality for different subset sizes in WMT'19 De-En.}
\label{fig:speed}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%Figure~\ref{fig:speed}\subref{fig:srck-bleu} shows the translation quality results on the development set (newstest2018).
    Figure~\ref{fig:speed}(c) shows the translation quality results on the development set (newstest2018).
The results show that a larger $n$ improved the BLEU up to $n=512$ but decreased the BLEU for greater values of $n$.
In terms of translation quality and speed, we set $n=512$ for WMT'19 De-En.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Relationship Between Neural/Non-neural Encoders and Translation Quality}

From Section~\ref{sec:exp-wmt19} and \ref{sec:exp-domain}, 
the translation quality of the non-neural sentence encoder 
\pagebreak
was almost the same as that of the neural sentence encoder in the WMT'19 translation task, whereas the non-neural sentence encoder degraded the translation quality compared with the neural sentence encoder in the domain adaptation task.
We hypothesized that one of the causes of this phenomenon is that calculating TF-IDF and BM25 for a sentence rather than for a document, would not extract sufficient statistics, particularly in short sentences.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 4
\begin{figure}[b]
\begin{center}
\includegraphics{31-2ia3f4.pdf}
\end{center}
\caption{Total difference in sentence BLEU for each length bucket.}
\label{fig:diff-bleu}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

To verify this, we measured the total difference in sentence BLEU when using LaBSE and TF-IDF for each length bucket of source sentences.
Note that the length bucket indicates the range from $i$ (inclusive) to $i+10$ (exclusive).
Figure~\ref{fig:diff-bleu} shows the results of the WMT'19 translation and medical domain adaptation tasks.
TF-IDF often degraded translation quality in short sentences, and this degradation was suppressed as the sentence length increased in both datasets.
As shown in Figure~\ref{fig:lengths}, the medical domain task had more short source sentences than the WMT'19 translation task.
Therefore, the score difference between TF-IDF and LaBSE in the medical domain could have been larger than that in the WMT'19 translation task because of sentence lengths.

To summarize, we observed that the non-neural encoder, TF-IDF, degraded the translation quality, particulary for short sentences, whereas the neural encoder, LaBSE, retrieved similar sentences robustly and prevented degradation even for short sentences.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Related Work}

The first type of example-based MT method was analogy-based MT \cite{nagao-1984-framework}.
Our idea of using source sentence similarity to obtain relevant translation examples for input sentences was inspired by \shortciteA{nagao-1984-framework}.
\shortciteA{zhang-etal-2018-guiding,gu-etal-2018-search} incorporated example-based methods into NMT models, which retrieved examples according to the edit distance.
\shortciteA{bulte-tezcan-2019-neural} and \shortciteA{xu-etal-2020-boosting} concatenated an input sentence and translations of similar sentences of the input.
Both the \knnmt and subset \knnmt retrieve \knn tokens according to the distance between the intermediate representations and interpolate the output probability.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 5
\begin{figure}[t]
\begin{center}
\includegraphics{31-2ia3f5.pdf}
\end{center}
\caption{Cumulative distribution of the lengths of source sentences.}
\label{fig:lengths}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

To improve the decoding speed of \knnmt, the fast \knnmt \cite{meng-etal-2022-fast} constructs additional datastores for each source token, and
 reduces the \knn search space using their datastores and word alignments.
The subset \knnmt requires a sentence datastore smaller than the source token datastores and does not require word alignment.
\shortciteA{martins-etal-2022-chunk} cached n-gram tokens adjacent to the retrieved tokens and reduced the number of queries for the entire datastore, leading to a speed-up of up to 4 times compared with \knnmt.
In contrast, the subset \knnmt does not search for the entire datastore during decoding.
\shortciteA{dai-etal-2023-simple} reduced the \knn search space by retrieving the neighboring sentences from the input sentence.
They searched for neighboring sentences using BM25 scores with \textsc{ElasticSearch};\footnote{\url{https://github.com/elastic/elasticsearch}} therefore, our subset \knnmt with BM25 can be regarded as an approximation of their method.
They also proposed ``adaptive lambda'', which dynamically computes the weights of the lambda of linear interpolation in Equation \ref{eq:interpolation} from the distance between the query and the nearest neighbor key vectors.
However, adaptive lambda requires an exact distance and cannot employ datastore quantization and ADC look-up.
To improve the translation quality of \knnmt, \shortciteA{zheng-etal-2021-adaptive} computed a weighted average of \knn probabilities $\pknn$ over multiple values of $k$.
Each weight is predicted by a ``meta-$k$ network'', trained to minimize the cross-entropy in the training data.
Their adaptive \knnmt only improved the translation quality, and its decoding speed was almost the same as that of the \knnmt \cite{zheng-etal-2021-adaptive}.
In contrast, we focused on improving the decoding speed.
Additionally, our subset \knnmt outperformed the \knnmt in some domain adaptation tasks as a positive side effect of subset retrieval.
For the other tasks, \knn-LM \cite{khandelwal-etal-2020-generalization}, efficient \knn-LM \cite{he-etal-2021-efficient}, and RETRO \cite{borgeaud-etal-2022-improving} used \knn search for language modeling (LM).
Our subset search method may be applied to LM regarding the prompt text as the query.
However, the method of constructing the sentence datastore from monolingual data is non-trivial, and we leave this issue for future work.

Some studies used sentence similarity to improve the translation quality of NMT models.
\shortciteA{wieting-etal-2019-beyond} showed that minimum risk training using the cosine similarity between the generated hypothesis and the reference translation improved the translation quality of NMT models.
Another approach used the sentence similarity between the output sentence and the reference as the reward for reinforcement learning \cite{yasui-etal-2019-using} to prevent excessive penalties owing to cross-entropy which did not take into account the semantics of the sentence.
Both methods used sentence similarity to place the semantics of the output sentence close to the reference translation, whereas our method used sentence similarity to search for translation examples.
These methods used sentence similarity on the target side, whereas our method used the similarity between the input sentence and the source sentences in the parallel data.


Quality estimation models and metric models, which use the similarity between the source sentence and the hypothesis, or the hypothesis and the reference, have been proposed to evaluate the translation quality \cite{rei-etal-2022-comet,rei-etal-2022-cometkiwi,sellam-etal-2020-bleurt}.
These models used the similarity on the target side, whereas our model used the similarity on the source side.
\shortciteA{sellam-etal-2020-bleurt} augmented the training data of the metric model by mask-filling with BERT \cite{devlin-etal-2019-bert}, back-translation, and dropping words to allow the model to capture various errors.
In our model, we may improve the accuracy of the similar sentence search by fine-tuning the sentence encoder to retrieve better subsets that improve the translation quality.

In the field of \knn searches, \shortciteA{matsui-etal-2018-reconfigurable} allowed a search in dynamically created subsets, whereas conventional search methods assume only a full search.
The subset \knnmt retrieves \knn tokens from a subset depending on the given input.
In our subset \knnmt, the decoding speed is slow when the subset size $n$ is large.
The bottleneck is the look-up in the distance table, and this can be improved by efficient look-up methods that use SIMD \cite{andre-etal-2015-cache,matsui-etal-2022-4bit}.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Conclusion}

In this paper, we proposed ``subset \knnmt'', which improves the decoding speed of the \knnmt using two methods:
(1) retrieving neighbor tokens from only the neighbor sentences of the input sentence, not from all sentences, and
(2) efficient distance computation technique suitable for subset neighbor search using a look-up table.
Our subset \knnmt achieved a speed-up of up to 134.2 times and an improvement in BLEU of up to 1.6 compared with those of the \knnmt in the WMT'19 De-En translation task, domain adaptation tasks in De-En and En-Ja translations, and the Flores101 multilingual translation task.
From the experiments, we observed that the translation quality varied depending on the sentence encoder.
For future work, we would like to compare the proposed method with other pre-trained models and also fine-tune sentence encoders maximizing the metrics.
In addition, we would like to apply our method to other text generation tasks, such as text simplification and speech-to-text translation.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\acknowledgment

This work was supported by JSPS KAKENHI Grant Numbers JP22J11279 and JP22KJ2286.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\bibliographystyle{jnlpbbl_1.7}
\bibliography{03refs}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{biography}

\bioauthor[:]{Hiroyuki Deguchi}{%
Hiroyuki Deguchi received B.E. and M.E. degrees from Ehime University in 2019 and 2021, respectively.
He is currently a Ph.D. student at Nara Institute of Science and Technology (NAIST) and a technical researcher at National Institute of Information and Communications Technology (NICT).
}

\bioauthor[:]{Taro Watanabe}{%
Taro Watanabe received his B.E. and M.E. degrees in information science from Kyoto University in 1994 and 1997, respectively, and obtained an M.S. degree in Language and Information Technologies from the School of Computer Science, Carnegie Mellon University in 2000.
In 2004, he received a Ph.D. in informatics from Kyoto University.
After working as a researcher at ATR, NTT and NICT, and as a software engineer at Google, he is a professor at the Nara Institute of Science and Technology starting in 2020.
His research interests include natural language processing, machine learning and machine translation.
}


\bioauthor[:]{Yusuke Matsui}{%
Yusuke Matsui received his Ph.D. degree in EECS from the University of Tokyo in 2016. He is currently a Lecturer (Assistant Professor) with the Department of Information and Communication Engineering, Graduate School of Information Science and Technology, the University of Tokyo, Japan. His research interests lie in computer vision and multimedia processing, with a particular interest in large-scale indexing. 
}

\bioauthor[:]{Masao Utiyama}{
Masao Utiyama is the director with Advanced Translation Technology Laboratory, the National Institute of Information and Communications Technology, Japan from 2023. He completed his doctoral dissertation with the University of Tsukuba, in 1997. His main research field is machine translation.

\bioauthor[:]{Hideki Tanaka}{
Hideki Tanaka is a Research Executive Director of the Advanced
Translation Technology Laboratory of the National Institute of
Information and Communications Technology in Japan. He received BS, MS,
and Ph.D. from Kyushu University in 1982, 1984, and 1995 respectively.
His research interest lies in Natural Language Processing and Machine
Learning.
}

\bioauthor[:]{Eiichiro Sumita}{
He is a fellow of NICT.
Before joining NICT, he worked with IBM and ATR.
His research field is mainly machine translation.
He also serves the president of AAMT.
He received a PhD in engineering from Kyoto University in 1999.
}
}
\end{biography}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\biodate


\end{document}

