Vocabulary Learning via Optimal Transport for Neural Machine Translation (2021)

Abstract: The choice of token vocabulary affects the performance of machine translation. This paper aims to figure out what is a good vocabulary and whether we can find the optimal vocabulary without trial training. To answer these questions, we first provide an alternative understanding of vocabulary from the perspective of information theory. It motivates us to formulate the quest of vocabularization – finding the best token dictionary with a proper size – as an optimal transport (OT) problem. We propose VOLT, a simple and efficient solution without trial training. Empirical results show that VOLT beats widely-used vocabularies in diverse scenarios, including WMT-14 English-German translation, TED bilingual translation, and TED multilingual translation. For example, VOLT achieves 70% vocabulary size reduction and 0.5 BLEU gain on English-German translation. Also, compared to BPE-search, VOLT reduces the search time from 384 GPU hours to 30 GPU hours on English-German translation. Codes are available at https://github.com/Jingjing-NLP/VOLT .

## introduction due to the discreteness of text, vocabulary construction ( vocabularization for short) is a prerequisite for neural machine translation (nmt) and many other natural language processing (nlp) tasks using neural networks @xcite @xcite @xcite . currently, sub-word approaches like byte-pair encoding (bpe) are widely used in the community @xcite @xcite , and achieve quite promising results in practice @xcite @xcite @xcite . the key idea of these approaches is selecting the most frequent sub-words (or word pieces with higher probabilities) as the vocabulary tokens. in information theory, these frequency-based approaches are simple forms of data compression to reduce entropy @xcite , which makes the resulting corpus easy to learn and predict @xcite . however, the effects of vocabulary size are not sufficiently taken into account since current approaches only consider frequency (or entropy) as the main criteria. many previous studies @xcite @xcite show that vocabulary size also affects downstream performances, especially on low-resource tasks. due to the lack of appropriate inductive bias about size, trial training (namely traversing all possible sizes) is usually required to search for the optimal size, which takes high computation costs. for convenience, most existing studies only adopt the widely-used settings in implementation. for example, 30k-40k is the most popular size setting in all 42 papers of conference of machine translation (wmt) through 2017 and 2018 @xcite . in this paper, we propose to explore automatic vocabularization by simultaneously considering entropy and vocabulary size without expensive trial training. designing such a vocabularization approach is non-trivial for two main reasons. first, it is challenging to find an appropriate objective function to optimize them at the same time. roughly speaking, the corpus entropy decreases with the increase of vocabulary size, which benefits model learning @xcite . on the other side, too many tokens cause token sparsity, which hurts model learning @xcite . second, supposing that an appropriate measurement is given, it is still challenging to 3000 4000 5000 6000 7000 8000 solve such a discrete optimization problem due to the exponential search space. to address the above problems, we propose a vocabulary learning approach via optimal transport, volt for short. it can give an appropriate vocabulary in polynomial time by considering corpus entropy and vocabulary size. specifically, given the above insight of contradiction between entropy and size, we first borrow the concept of marginal utility in economics @xcite and propose to use marginal utility of vocabularization (muv) as the measurement. the insight is quite simple: in economics, marginal utility is used to balance the benefit and the cost and we use muv to balance the entropy (benefit) and vocabulary size (cost). higher muv is expected for pareto optimality. formally, muv is defined as the negative derivative of entropy to vocabulary size. figure 1 gives an example about marginal utility. preliminary results verify that muv correlates with the downstream performances on two-thirds of tasks (see figure 2 ). then our goal turns to maximize muv in tractable time complexity. we reformulate our discrete optimization objective into an optimal transport problem @xcite that can be solved in polynomial time by linear programming. intuitively, the vocabularization process can be regarded as finding the optimal transport matrix from the character distribution to the vocabulary token distribution. finally, our proposed volt will yield a vocabulary from the optimal transport matrix. we evaluate our approach on multiple machine translation tasks, including wmt-14 english-german translation, ted bilingual translation, and ted multilingual translation. empirical results show that volt beats widely-used vocabularies in diverse scenarios. furthermore, volt is a lightweight solution and does not require expensive computation resources. on english-german translation, volt only takes 30 gpu hours to find vocabularies, while the traditional bpe-search solution takes 384 gpu hours. ## related work initially, most neural models were built upon word-level vocabularies @xcite @xcite . while achieving promising results, it is a common constraint that word-level vocabularies fail on handling rare words under limited vocabulary sizes. researchers recently have proposed several advanced vocabularization approaches, like bytelevel approaches @xcite , characterlevel approaches @xcite @xcite , and sub-word approaches @xcite . byte-pair encoding (bpe) @xcite is proposed to get subword-level vocabularies. the general idea is to merge pairs of frequent character sequences to create sub-word units. sub-word vocabularies can be regarded as a trade-off between character-level vocabularies and word-level vocabularies. compared to word-level vocabularies, it can decrease the sparsity of tokens and increase the shared features between similar words, which probably have similar semantic meanings, like "happy" and "happier". compared to character-level vocabularies, it has shorter sentence lengths without rare words. following bpe, some variants recently have been proposed, like bpe-dropout @xcite , sentencepiece @xcite , and so on. despite promising results, most existing subword approaches only consider frequency while the effects of vocabulary size is neglected. thus, trial training is required to find the optimal size, which brings high computation costs. more recently, some studies notice this problem and propose some practical solutions @xcite @xcite . ## marginal utility of vocabularization in this section, we propose to find a good vocabulary measurement by considering entropy and size. as introduced in section 1, it is non-trivial to find an appropriate objective function to optimize them simultaneously. on one side, with the increase of vocabulary size, the corpus entropy is decreased, which benefits model learning @xcite . on the other side, a large vocabulary causes parameter explosion and token sparsity problems, which hurts model learning @xcite . to address this problem, we borrow the concept of marginal utility in economics @xcite and propose to use marginal utility of vocabularization (muv) as the optimization objective. muv evaluates the benefits (entropy) a corpus can get from an increase of cost (size). higher muv is expected for higher benefit-cost ratio. preliminary results verify that muv correlates with downstream performances on two-thirds of translation tasks (see figure 2 ). according to this feature, our goal turns to maximize muv in tractable time complexity. definition of muv formally, muv represents the negative derivation of entropy to size. for simplification, we leverage a smaller vocabulary to estimate muv in implementation. specially, muv is calculated as: ## maximizing muv via optimal transport this section describes the details of the proposed approach. we first show the general idea of volt in section 4.1, then describe the optimal transport solution in section 4.2, followed by the implementation details in section 4.3. ## experiments to evaluate the performance of volt, we conduct experiments on three datasets, including wmt-14 english-german translation, ted bilingual translation, and ted multilingual translation. ## conclusion in this work, we propose a new vocabulary search approach without trail training. the whole framework starts from an informtaion-therotic understanding. according to this understanding, we formulate vocabularization as a two-step discrete optimization objective and propose a principled optimal transport solution volt. experiments show that volt can effectively find a well-performing vocabulary in diverse settings.