Angular Dispersion Accelerates k -Nearest Neighbors Machine Translation (2025)

Abstract: Augmenting neural machine translation with external memory at decoding time, in the form of k -nearest neighbors machine translation ( k -NN MT), is a well-established strategy for increasing translation performance. k -NN MT retrieves a set of tokens that occurred in the most similar contexts recorded in a prepared data store, using hidden state representations of translation contexts as vector lookup keys. One of the main disadvantages of this method is the high computational cost and memory requirements. Since an exhaustive search is not feasible in large data stores practitioners commonly use approximate k -NN lookup, yet even such algorithms are a bottleneck. In contrast to research directions seeking to accelerate k -NN MT by reducing data store size or the number of lookup calls, we pursue an orthogonal direction based on the performance properties of approximate k -NN lookup data structures. In particular, we propose encouraging angular dispersion of the neural hidden representations of contexts. We show that improving dispersion leads to better balance in the retrieval data structures, accelerating retrieval and slightly improving translations.

## key representation with dispersion dispersion on the sphere is well-established in the literature as discussed in §3, and it is common to apply dispersion of the angles for learning hyperspherical representations, i.e., magnitudes of the vectors are discarded. however, previous works show that the norm of the keys, in fact, might contain important information @xcite . to avoid discarding norms, contrary to other works, we apply dispersion only to the angles of the model's outputs while keeping their magnitudes intact and do gradient updates in euclidean space. figure 1 shows how angular dispersion operates in 2d. ## experimental setup 4.1 synthetic data to more directly control the effect of dispersion, we generate a family of synthetic data stores. specifically, we sample 10m random vectors of dimensionality 128 from the mixture of 5 power spherical distributions @xcite . the scale parameter κ for the power spherical distribution defines the concentration of the data points, where higher κ means higher concentration. we build a data store using the keys generated with κ = {1, 10, 50, 100, 1000} with 2048 centroids while picking a random length for each key in the [1, 100] range. ## conclusion in this work, we show that angular dispersion of the data store benefits the retrieval speed and performance in k-nn mt for both in-domain and domain adaptation scenarios. our analysis of cluster properties and synthetic data stores further indicates that dispersion of the keys balances the data store clustering, accelerating the approximate search, especially for larger data store sizes. our findings indicate that fine-tuning with dispersion can be a first efficient k-nn mt component. since there are no changes in architecture or k-nn data store construction, we hypothesize that other efficient methods that rely on distance metrics can be easily applied on top, which we leave for future work.