Neural Search Space in Gboard Decoder (2024)

Abstract: Gboard Decoder produces suggestions by looking for paths that best match input touch points on the context aware search space, which is backed by the language Finite State Transducers (FST). The language FST is currently an N-gram language model (LM). However, N-gram LMs, limited in context length, are known to have sparsity problem under device model size constraint. In this paper, we propose Neural Search Space which substitutes the N-gram LM with a Neural Network LM (NN-LM) and dynamically constructs the search space during decoding. Specifically, we integrate the long range context awareness of NN-LM into the search space by converting its outputs given context, into the language FST at runtime. This involves language FST structure redesign, pruning strategies tuning, and data structure optimizations. Online experiments demonstrate improved quality results, reducing Words Modified Ratio by [0.26%, 1.19%] on various locales with acceptable latency increases. This work opens new avenues for further improving keyboard decoding quality by enhancing neural LM more directly.

## introduction gboard is a statistical-decoding-based keyboard on mobile devices developed by google. statistical decoding is far more necessary than one might think due to the error-prone process of "fat finger" touch input on small screens. according to @xcite , the per-letter error rate is around 8%-9% without decoding. with decoding, typos such as substitutions (due to the proximity of two keys or cognitive misspellings), omissions, insertions, and transpositions could be automatically corrected by the key-correction and (word) autocorrection functions in the gboard decoder, leading to an error-tolerant user experience. powered by language models (lm), the gboard decoder also * equal contribution. the decoding process involves two phases: building search space (decoder graph), and performing beam search within the space based on user touch inputs. gboard decoder utilizes context, a lexicon and language transducers -the familiar c • l • g composition @xcite -to construct the search space. c is a bi-key key to key transducer while l is a key to word transducer, c and l are statically composed together offline since the size is small. fig. 1-a illustrates how gesture typing and tap typing inputs are converted into bi-keys and fig. 1-b illustrates a composed c • l targeting four words. before this work, g is a n-gram language fst containing 64k words for n-grams and 170k words for uni-grams. composition between (c • l) and g are dynamically conducted due to the large size of g. fig. 1-c shows a simple g containing only four words, and fig. 1 -d illustrates a composed (c • l) • g, which is similar to (c • l) but with weights achieved by using the look-ahead composition filters proposed by @xcite . in practice, the whole search space of (c •l)•g like fig. 1 -d can't be fully expanded due to its huge size. only states which are close to users' bi-key inputs will be expanded. specifically, the states are pruned based on the combination of lm scores and spatial scores in the decoder graph while the user is typing. in this work, the n-gram lm is replaced by the nn-lm. due to the framework complexity brought by the rich functionalities, we propose a runtime conversion solution to minimize the framework change. we call the search space built on the nn-lm the neural search space (nss), and the original one the n-gram search space. algorithm 1 is called before users start to type a new word, for example, when users open the keyboard for an editor box or when users commit a word by tapping on space. algorithm 2 is called when users are typing a word; for tap-typing, it will be called upon each key tap. nss has minor changes in both algorithms, in algorithm 1, rather than finding a specific state in a static n-gram lm given context, u pdatelm inserts the next words and corresponding probabilities given context at the start state of g n lm as arcs, thus the start state is always 0. in al-gorithm 2, extendlm additionally extends the g n lm which aims to handle the multi-word problem discussed later. ## background recent advances of neural network lms(nn-lm), notably projects such as gpt-4 (openai, 2023), palm 2 @xcite , demonstrate their superior performance compared to n-gram lms, particularly in capturing longer context. federated learning (fl) @xcite with differential privacy (dp) @xcite enables gboard to improve lm quality with user data while preserving user privacy by distributing model training across user devices instead of collecting data centrally. prior work employed fl to train lms for next word prediction, smart compose, and on-the-fly rescoring in gboard following @xcite ; @xcite . however, these applications either operate on first pass decoding results produced by n-gram lms, or do not affect decoding suggestion which has the largest impact on typing experience. to benefit from fl of nn-lm and retain decoding efficiency, previous research has explored projecting or approximating nn-lms onto n-gram lms @xcite @xcite , and making the fst differentiable @xcite . however, such conversions inevitably incur losses due to limited context and back-off smoothing necessitated by sparsity @xcite . in this work, we replace the n-gram lm within the search space with an nn-lm trained via fl, enhancing long context capabilities. the deployed nn-lm is an lstm / cifg model similar to those in @xcite ; @xcite . ## challenges ideally, an nn-lm would score all known words for optimal coverage. however, vocabulary size is limited due to the high computational cost of the final dense layer. our deployed nn-lm has a 30k-word vocabulary (top words from federated counting), while the full lexicon contains 170k words. scoring the remaining 140k words in our generated fst poses a key challenge. missing the space key and mistyping it with the "cvbn" keys are the two common and consequential mistakes in mobile typing, turning multiple words into one single string (see appendix a.1 for demo cases). @xmath0, @xmath1, with subsequent words receiving contextless unigram scores. this penalizes and perhaps suppresses multi-word candidates. we address this using dynamic inference in section 4.3. gboard operates under strict latency constraints. key presses should trigger visible feedback within 20ms as highlighted in @xcite . nss inevitably increases latency due to nn-lm inference and fst conversion. dynamic inference, employed to address space substitution issues, significantly expands the search space by hypothesizing word separations at each frame, further exacerbates this challenge. ## methods we detail the u pdatelm and extendlm described in algorithm 1 and algorithm 2 respectively below. ## discussions this work successfully bridged the gap between the long range context awareness power of nn-lms and the efficiency requirements of gboard's decoder. by creatively adapting nn-lms to an fst structure and implementing latency optimizations, we deployed the neural search space in production. experiments demonstrated that nss significantly improves decoding quality, particularly on higherend devices, with an acceptable latency trade-off. this direct integration unlocks potential for future enhancements driven by nn-lm advancements, promising further gains in keyboard decoding and overall user experience. building upon nss, we identify several promising directions for future research: beyond enhancing nss, another avenue for exploration is replacing the fst-based decoder with a neural decoder. while we have investigated this approach, certain challenges hinder its immediate adoption as a full replacement: