Vocabulary Learning via Optimal Transport for Neural ...

1 Vocabulary Learning via Optimal Transport for Neural Machine Translation Jingjing Xu1 , Hao Zhou1 , Chun Gan1,2 , Zaixiang Zheng1,3 , Lei Li1. 1. ByteDance AI Lab 2. Math Department, University of Wisconsin Madison 3. Nanjing University Abstract 2018; Al-Rfou et al., 2019; Wang et al., 2020). The choice of token Vocabulary affects the per- The key idea of these approaches is selecting the formance of machine translation. This paper most frequent sub-words (or word pieces with aims to figure out what is a good Vocabulary higher probabilities) as the Vocabulary tokens. and whether one can find the Optimal vocab- In information theory, these frequency-based ap- ulary without trial training. To answer these proaches are simple forms of data compression to questions, we first provide an alternative un- reduce entropy (Gage, 1994), which makes the re- derstanding of the role of Vocabulary from the sulting corpus easy to learn and predict (Martin perspective of information theory.

2 Motivated by this, we formulate the quest of vocabular- and England, 2011; Bentz and Alikaniotis, 2016). ization finding the best token dictionary with However, the effects of Vocabulary size are not a proper size as an Optimal Transport (OT) sufficiently taken into account since current ap- problem. We propose VOLT, a simple and proaches only consider frequency (or entropy) as efficient solution without trial training. Em- the main criteria. Many previous studies (Sennrich pirical results show that VOLT outperforms and Zhang, 2019; Ding et al., 2019; Provilkov widely-used vocabularies in diverse scenar- et al., 2020; Salesky et al., 2020) show that vocab- ios, including WMT-14 English-German and TED's 52 translation directions. For example, ulary size also affects downstream performances, VOLT achieves 70% Vocabulary size reduc- especially on low-resource tasks. Due to the lack tion and BLEU gain on English-German of appropriate inductive bias about size, trial train- translation.

3 Also, compared to BPE-search, ing (namely traversing all possible sizes) is usually VOLT reduces the search time from 384 GPU required to search for the Optimal size, which takes hours to 30 GPU hours on English-German high computation costs. For convenience, most translation. Codes are available at https: existing studies only adopt the widely-used set- tings in implementation. For example, 30K-40K. is the most popular size setting in all 42 papers 1 Introduction of Conference of Machine Translation (WMT). Due to the discreteness of text, Vocabulary con- through 2017 and 2018 (Ding et al., 2019). struction ( vocabularization for short) is a prereq- In this paper, we propose to explore auto- uisite for Neural machine translation (NMT) and matic vocabularization by simultaneously consid- many other natural language processing (NLP) ering entropy and Vocabulary size without expen- tasks using Neural networks (Mikolov et al., 2013; sive trial training.)

4 Designing such a vocabulariza- Vaswani et al., 2017; Gehrmann et al., 2018; tion approach is non-trivial for two main reasons. Zhang et al., 2018; Devlin et al., 2019). Cur- First, it is challenging to find an appropriate objec- rently, sub-word approaches like Byte-Pair En- tive function to optimize them at the same time. coding (BPE) are widely used in the commu- Roughly speaking, the corpus entropy decreases nity (Ott et al., 2018; Ding et al., 2019; Liu et al., with the increase of Vocabulary size, which bene- 2020), and achieve quite promising results in prac- fits model Learning (Martin and England, 2011). tice (Sennrich et al., 2016; Costa-jussa and Fonol- On the other side, too many tokens cause to- losa, 2016; Lee et al., 2017; Kudo and Richardson, ken sparsity, which hurts model Learning (Allison . This work is done during the internship at ByteDance AI et al., 2006). Second, supposing that an appropri- Lab. ate measurement is given, it is still challenging to 7361.

5 Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, pages 7361 7373. August 1 6, 2021. 2021 Association for Computational Linguistics Eo translation tasks, including WMT-14 English- Entropy German translation, TED bilingual translation, BLEU and TED multilingual translation. Empirical re- sults show that VOLT beats widely-used vocabu- Entropy BLEU. laries in diverse scenarios. Furthermore, VOLT is a lightweight solution and does not require expen- sive computation resources. On English-German translation, VOLT only takes 30 GPU hours to find 3000 4000 5000 6000 7000 8000 9000 Size vocabularies, while the traditional BPE-Search so- lution takes 384 GPU hours. Figure 1: An illustration of marginal utility. We sample BPE-generated vocabularies with different sizes from Eo-En translation and draw their entropy (See ). 2 Related Work and BLEU lines.

6 Star represents the Vocabulary with Initially, most Neural models were built upon the maximum marginal utility. Marginal utility (See ) evaluates the increase of benefit (entropy de- word-level vocabularies (Costa-jussa and Fonol- crease) from an increase of cost (size). losa, 2016; Vaswani et al., 2017; Zhao et al., 2019). While achieving promising results, it is a common constraint that word-level vocabularies solve such a discrete optimization problem due to fail on handling rare words under limited vocabu- the exponential search space. lary sizes. To address the above problems, we propose Researchers recently have proposed several ad- a Vocabulary Learning approach via Optimal vanced vocabularization approaches, like byte- Transport , VOLT for short. It can give an appro- level approaches (Wang et al., 2020), character- priate Vocabulary in polynomial time by consider- level approaches (Costa-jussa and Fonollosa, ing corpus entropy and Vocabulary size.)

7 Specifi- 2016; Lee et al., 2017; Al-Rfou et al., 2019), cally, given the above insight of contradiction be- and sub-word approaches (Sennrich et al., 2016;. tween entropy and size, we first borrow the con- Kudo and Richardson, 2018). Byte-Pair Encoding cept of Marginal Utility in economics (Samuelson, (BPE) (Sennrich et al., 2016) is proposed to get 1937) and propose to use Marginal Utility of Vo- subword-level vocabularies. The general idea is cabularization (MUV) as the measurement. The to merge pairs of frequent character sequences to insight is quite simple: in economics, marginal create sub-word units. Sub-word vocabularies can utility is used to balance the benefit and the cost be regarded as a trade-off between character-level and we use MUV to balance the entropy (bene- vocabularies and word-level vocabularies. Com- fit) and Vocabulary size (cost). Higher MUV is pared to word-level vocabularies, it can decrease expected for Pareto optimality.

8 Formally, MUV the sparsity of tokens and increase the shared is defined as the negative derivative of entropy features between similar words, which probably to Vocabulary size. Figure 1 gives an example have similar semantic meanings, like happy and about marginal utility. Preliminary results verify happier . Compared to character-level vocabu- that MUV correlates with the downstream perfor- laries, it has shorter sentence lengths without rare mances on two-thirds of tasks (See Figure 2). words. Following BPE, some variants recently Then our goal turns to maximize MUV in have been proposed, like BPE-dropout (Provilkov tractable time complexity. We reformulate our dis- et al., 2020), SentencePiece (Kudo and Richard- crete optimization objective into an Optimal trans- son, 2018), and so on. port problem (Cuturi, 2013) that can be solved Despite promising results, most existing sub- in polynomial time by linear programming. In- word approaches only consider frequency while tuitively, the vocabularization process can be re- the effects of Vocabulary size is neglected.

9 Thus, garded as finding the Optimal Transport matrix trial training is required to find the Optimal size, from the character distribution to the Vocabulary which brings high computation costs. More token distribution. Finally, our proposed VOLT recently, some studies notice this problem and will yield a Vocabulary from the Optimal Transport propose some practical solutions (Kreutzer and matrix. Sokolov, 2018; Cherry et al., 2018; Chen et al., We evaluate our approach on multiple machine 2019; Salesky et al., 2020). 7362. 12 defined by the sum of token entropy. To avoid the 10 effects of token length, here we normalize entropy with the average length of tokens and the final en- 8. Count tropy is defined as: 6. 4. 1 X. Hv = P (i) log P (i), (2). lv i v 2. 0 where P (i) is the relative frequency of token i Spearman Score from the training corpus and lv is the average length of tokens in Vocabulary v. Figure 2: MUV and downstream performance are pos- itively correlated on two-thirds of tasks.

10 X-axis clas- Preliminary Results To verify the effectiveness sifies Spearman scores into different groups. Y-axis shows the number of tasks in each group. The middle of MUV as the Vocabulary measurement, we con- Spearman score is duct experiments on 45 language pairs from TED. and calculate the Spearman correlation score be- tween MUV and BLEU scores. We adopt the same 3 Marginal Utility of Vocabularization and widely-used settings to avoid the effects of In this section, we propose to find a good vocabu- other attributes on BLEU scores, such as model lary measurement by considering entropy and size. hyper-parameters and training hyper-parameters. As introduced in Section 1, it is non-trivial to find We generate a sequence of vocabularies with in- an appropriate objective function to optimize them cremental sizes via BPE. All experiments use the simultaneously. On one side, with the increase of same hyper-parameters. Two-thirds of pairs show Vocabulary size, the corpus entropy is decreased, positive correlations as shown in Figure 2.