Fast Word Predictor for On-Device Application (2020)

Abstract: Learning on large text corpora, deep neural networks achieve promising results in the next word prediction task. However, deploying these huge models on devices has to deal with constraints of low latency and a small binary size. To address these challenges, we propose a fast word predictor performing efficiently on mobile devices. Compared with a standard neural network which has a similar word prediction rate, the proposed model obtains 60% reduction in memory size and 100X faster inference time on a middle-end mobile device. The method is developed as a feature for a chat application which serves more than 100 million users.

## introduction as a self-supervised learning task, next word prediction based on deep neural networks obtains robust performance by learning on large text corpora. given a textual context as input, these models shown in figure 1 @xcite encode the text to generate a probability distribution over a vocabulary for the next word. although various neural networks have been developed for efficient computation and performance, the word embedding layer and softmax layer are still essential parts in these architectures. however, this approach faces a bottleneck for deploying on mobile devices that is the huge number of parameters in the word embedding matrix and softmax layer. for a vocabulary of n words and a word embedding of dimension d, the word embedding matrix takes n × d parameters and the softmax layer takes h × n where h is the dimension of encoded text. @xcite , for instance, proposed a recurrent neural network with @xmath0, and @xmath1, so the word embedding matrix and softmax layer have 18m parameters in total. this cost has limited the applicability of deep neural networks on mobile devices for word prediction, word completion, and error correction tasks. various approaches have been proposed for deep neural compression. matrix factorization @xcite @xcite is applied to weight matrices to reduce model parameters while weight tying @xcite shares the parameters of the embedding layer with those of the output classifier. in addition, network pruning and sharing are also efficient for reducing network complexity. for instance, @xcite and @xcite explored the redundancy among neurons and propose to keep only the most the relevant parameters. recently, @xcite proposed an on-device word predictor employing shared matrix factorization and distillation to optimize memory and computation. these methods achieve promising results in memory reduction but limited efficiency in time constraint, especially on low-resource devices. as we analysed the bottleneck, computation complexity primarily depends on vocabulary's size n , word embedding dimension d, and hidden dimension h related to the embedding matrix and softmax layer. designed as a lookup table, the former have constant complexity, while the latter's operation is fully matrix multiplication. therefore, our proposed approach is to reduce this softmax layer's computation. according to our analysis on 150 million tokens from vietnamese wikipedia text and social text listed in in this work, we propose an approach which takes a contextual text and list of candidate words as input, and output the probability distribution over these words via the attention mechanism @xcite . this architecture shown in figure 1 (b) decreases significantly the number of softmax layer's parameters and performs faster due to the attention operation's parallelization. in addition, our proposed model keeps the contextual words and potential words in the same embedding space which is recommended for learning language models (press and wolf, 2017). on middle-end devices (e.g., samsung galaxy a8 -2016), the proposed model takes a smaller binary memory of 2.1mb (vs. 4.9mb) and faster inference time of 0.4ms (vs 41.2ms) while achieving a similar word prediction rate compared with a conventional neural model. the next section describes in detail the model's architecture, how to obtain a list of candidate words for a given context, and the mechanism to construct training samples for efficiently learning and diminishing bias. ## experiment according to the statistic shown in ## conclusion we have proposed an efficient approach for on-device word prediction. by making use of candidate words, ta distillation and matrix factorization, the model requires less parameters and computation while achieving a competitive performance. the model takes 2.1mb in memory and outputs a result in 0.4ms on the middle-end device. to optimize the model and keep it from being bias, we build a list of candidate words including potential words and frequent words. this approach of using candidate words is promising for enhancing user experience via personalizing the candidate list.