SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qian Huang
qhwang@cs.stanford.edu
Eric Zelikman
ezelikman@cs.stanford.edu
Sarah Li Chen
sachen@stanford.edu
Yuhuai  $\mathbf{W}\mathbf{u}^{12}$
yuhuai@cs.stanford.edu
Gregory Valiant
gvaliant@cs.stanford.edu
pliang@cs.stanford.edu
$^{1}$ Stanford University
Token embeddings, a mapping from discrete lexical symbols to continuous vectors, are at the heart of any language model (LM). However, lexical symbol meanings can also be determined and even redefined by their structural role in a long context. In this paper, we ask: is it possible for a language model to be performant without any fixed token embeddings? Such a language model would have to rely entirely on the co-occurrence and repetition of tokens in the context rather than the a priori identity of any token. To answer this, we study lexinvariant language models that are invariant to lexical symbols and therefore do not need fixed token embeddings in practice. First, we prove that we can construct a lexinvariant LM to converge to the true language model at a uniform rate that is polynomial in terms of the context length, with a constant factor that is sublinear in the vocabulary size. Second, to build a lexinvariant LM, we simply encode tokens using random Gaussian vectors, such that each token maps to the same representation within each sequence but different representations across sequences. Empirically, we demonstrate that it can indeed attain perplexity comparable to that of a standard language model, given a sufficiently long context. We further explore two properties of the lexinvariant language models: First, given text generated from a substitution cipher of English, it implicitly implements Bayesian in-context deciphering and infers the mapping to the underlying real tokens with high accuracy. Second, it has on average 4X better accuracy over synthetic in-context reasoning tasks. Finally, we discuss regularizing standard language models towards lexinvariance and potential practical applications.
All language processing systems rely on a stable lexicon, which assumes that a token (a word or subword such as tree) has a consistent contribution to the meaning of a text (though of course this meaning is mediated by context). In neural language models (LMs), this contribution is the token embedding, which stably maps each token into a continuous vector [21, 16, 17, 7, 6]. However, in real language, a token's contribution might be determined by its structural role; in math and code, novel variable names are arbitrarily defined to carry new meaning, and poems such as Jabberwocky exploit humans' lexical flexibility in interpreting novel words such as vorpal. Besides standard language understanding, this lexical flexibility also correlates with a stronger in-context reasoning performance. For example, GPT-3 [6] and other large language models that demonstrate high lexical flexibility show strong performance on tasks involving in-context reasoning over new concepts and rules.
Figure 1: Definition (a) and construction (b) of lexinvariant language model
Motivated by the above, we ask whether we can push this flexibility to the extreme: can we build a language model without any stable lexical mapping? To this end, we formulate and study such lex invariant language models. We define a lex invariant language model as a language model that assigns the same probability to all lexical permutations of a sequence. Formally, we define a lexical permutation  $\pi$  to be a one-to-one mapping of a set of lexical symbols  $^1$  onto itself. Then the lex invariant language model is defined as a language model over the symbol sequence  $x_{1},\ldots ,x_{n}$  with the following property:
$$
p \left(x _ {1}, \dots , x _ {n}\right) = p \left(\pi \left(x _ {1}\right), \dots , \pi \left(x _ {n}\right)\right) \forall \pi \tag {1}
$$
For example, a lex invariant language model (whose vocabulary is letters and space) should assign the same probability to the phrase "a big banana" as "e cop cekeke" because the two are the same up to the permutation  $\pi = \{a: e, b: c, i: o, n: k, g: p, \dots\}$  (Figure 1a).
The central question is: how well can lexinvariant language models predict the next token given an increasingly long context? We find the answer is almost as well as standard language models, both theoretically and empirically. This is rather surprising given that lexinvariance seems like a strong limitation (a model doesn't know what any individual symbol means!) However, the intuition is that given longer contexts, a lexinvariant model can both infer the latent permutation  $\pi$  (lazily) up to whatever ambiguity is present in the language model, and do the standard next word prediction task jointly.
Theoretically, we prove that a constructed lexinvariant language model can converge to the true language model as the context length increases—that is, the average L1 distance between the predictions of the two models decreases with a convergence rate of  $O\left(\left(\frac{d}{T}\right)^{\frac{1}{4}}\right)$ , where  $T$  is the length of the context and  $d$  is the vocabulary size, and where the big-O notation hides polylogarithmic factors of  $d$  and  $T$  and an absolute constant that is independent of the language model.
Empirically, we train a lexinvariant LM by replacing standard embeddings in a decoder-only Transformer [24] with per-sequence random Gaussian vectors, such that the same symbols get the same embedding within each sequence but get different embedding across sequences (Figure 1b). We indeed see that the perplexity gap between the lexinvariant LM and the standard LM shrinks as context length increases, as shown in Section 3.2. With a 150M parameters Transformer and a small character-level vocabulary (130 tokens), the average perplexity gap shrinks from 9X to less than 1X the average perplexity of a standard LM after observing 512 tokens over The Pile [9]. With a larger 32K vocabulary, the gap also shrinks, especially on the more structured text like GitHub code, albeit at a much slower rate.
We then explore two additional properties of the lexinvariant LM: in-context deciphering and symbol manipulation. First, we show that given a ciphertext generated by applying a substitution cipher to English text, the lexinvariant LM can be seen as implicitly approximating Bayesian inference of the lexical permutation, i.e., cipher key, in-context. To show this empirically, we train a small MLP probe on top of a frozen pretrained lexinvariant LM to predict the deciphered token corresponding to the last
<sup>1</sup>We specifically consider lexical symbols as tokens, not necessarily words or other linguistic units.
seen cipher token. We can then read out the inferred cipher key with each prefix of the sequence. We show that the accuracy of this inferred cipher key quickly improves as context length grows, reaching  $99.6\%$  average accuracy. We also show examples in Section 3.4 that visualize the uncertainties over different possible lexical mappings maintained by the lexinvariant LM when the cipher key is ambiguous and that the semantic meaning of a symbol with very rare occurrence can be inferred efficiently relative to other common symbols in context. Second, we show that lexinvariant models perform better than traditional models over synthetic pure in-context reasoning tasks that involve symbol manipulation. We observe a significant 4X improvement over a standard language model.
While the primary motivation of this paper is scientific exploration of a new idea, lexinvariance, we were also curious to see if it could help improve certain tasks, generalizing the performance gain we see on synthetic tasks. We stress that for most practical applications, lexinvariance is far too strong, so these experiments are intended to be illustrative rather than be a recipe for improving state-of-the-art. We discuss potential approaches to integrate the idea of lexinvariant LM into standard language modeling as a form of regularization, such that the LM assumes some form of partially stable symbol representations. The resulting LM can improve upon a standard language model over some BIG-bench tasks [23].
2 Lexinvariant Language Model
We define a language model as a probability distribution  $p(x_1, \ldots, x_n)$  over input token sequences  $x_1, \ldots, x_n \in \mathcal{V}^n$ , where  $\mathcal{V}$  is some vocabulary over symbols. A language model is lexinvariant if for all permutations  $\pi: \mathcal{V} \to \mathcal{V}$  and for all token sequences  $x_1, \ldots, x_n \in \mathcal{V}^n$ ,  $p(x_1, \ldots, x_n) = p(\pi(x_1), \ldots, \pi(x_n))$ . For example, if  $\mathcal{V} = \{\mathrm{a}, \mathrm{b}\}$ , then the model should assign the same probability to  $aab$  and  $bba$ . One example  $p$  that satisfies this could simply be
$$
p (x) = \left\{ \begin{array}{l l} 1 / 2 & x \in \{a a b, b b a \} \\ 0 & \text {o t h e r w i s e} \end{array} \right. \tag {2}
$$
Can such a lex invariant language model predict language well, even though it can only make next token predictions based on the structure of co-occurrence and repetition of input tokens in a single context?
2.1 Convergence on Language Modeling Performance
We show that we can construct a lexinvariant LM (as shown in Figure 2) to model the true language distribution faithfully, given a long enough context. The lexinvariant language model can essentially infer back the latent permutation  $\pi$  as it observes more symbols.
Figure 2: Probabilistic graphical model for the lexinvariant LM associated with the true language distribution  $p(x_{1},\ldots ,x_{n})$ .
As an intuitive example, suppose that  $\mathcal{V} = \{a,b\}$  and the true language only contains two sequences babbbb and ababab (and their prefixes) with even probability. When given only the first three letters, a lexinvariant model can't tell the latent permutation and can only assign the same probability to  $a$  and  $b$  for the next letter: Due to the lexinvariant property, it assigns the same probability to  $p(a|aba) = p(b|bab)$  as well as to  $p(b|aba) = p(a|bab)$ . Further,  $p(a|aba) = p(b|aba)$  because the permutations to prefixes aba and bab are equally probable. In contrast, when considering the prefix abab, the fourth letter resolves the ambiguity in possible permutations  $\pi$ . (Since bab is not in the true language distribution,  $\pi$  cannot map  $a$  to  $b$ .) Therefore, the model can correctly predict that  $p(a|abab) = 1$  and  $p(b|abab) = 0$ .
Formally, for a given language model  $p$ , we define the associated lexinvariant language model  $p'(x_1, \ldots, x_n)$  as  $\mathbb{E}_{\pi}[p(\pi^{-1}(x_1), \ldots, \pi^{-1}(x_n))]$ . Analyzing it, we have the following theorem:
Theorem 2.1. Let  $x_{1}, \ldots, x_{n}$  be any token sequence generated by an arbitrary language distribution  $p$  with an alphabet of size  $d$ . Let  $p'(x_{1}, \ldots, x_{n}) = \mathbb{E}_{\pi}[p(\pi^{-1}(x_{1}), \ldots, \pi^{-1}(x_{n}))]$ . Then, for any
$$
0 <   \epsilon , \delta <   1 / 2,
$$
$$
\frac {1}{T} \sum_ {t = 1} ^ {T} \| p (x _ {t} | x _ {1}, \dots , x _ {t - 1}) - p ^ {\prime} (x _ {t} | x _ {1}, \dots , x _ {t - 1}) \| _ {1} \leq \epsilon
$$
with probability greater than  $1 - \delta$ , when  $T \geq \frac{d}{\epsilon^4} \log(d, \frac{1}{\epsilon}, \frac{1}{\delta})$ , where the polylogarithmic term hides an absolute constant that is independent of  $p$ .
This theorem says that this associated lexinvariant language model converges to modeling the true language distribution fairly efficiently—with polynomial rate and near-linear dependence on vocabulary size  $d$ . Strikingly, this holds irrespective of the properties of the language distribution  $p^2$ . In other words, a language model can indeed infer the operational meaning of the tokens in context based solely on the structure of the symbols!
We give a complete proof of this theorem in Appendix A. At a high level, this convergence happens because at most timesteps  $t$ , the new observation  $x_{t}$  either provides new information about the permutation  $\pi$ , or  $x_{t}$  has similar likelihood under the permutations that are likely given  $x_{1},\ldots,x_{t-1}$ . In the simplest case, if the posterior  $\mathcal{P}(\pi|x_{1},\ldots,x_{n})$  concentrates on the correct  $\pi$ , then we converge to the standard LM. But even if it doesn't, that means the uncertainty about  $\pi$  should not matter for predicting the next token. We make this precise by interpreting  $p^{\prime}(x_{t}|x_{1},\ldots,x_{t-1})$  as performing a multiplicative weights algorithm with the Hedge strategy of Freund and Schapire [8], and then relate the regret bounds to the average KL divergence between the predictions of  $p$  and  $p^{\prime}$ , and ultimately the average  $\mathcal{L}_{1}$  distance between these predictions.
2.2 In-context Bayesian Deciphering
We can see the associated lexinvariant language model as implicitly learning to approximate an in-context Bayesian deciphering process, i.e. inferring a probability distribution over possible lexical permutations based on seen tokens, with the language modeling prior:
$$
\begin{array}{l} p ^ {\prime} (x _ {n + 1} | x _ {1}, \dots , x _ {n}) \\ = \sum_ {\pi} \frac {1}{d !} \frac {p \left(\pi^ {- 1} \left(x _ {1}\right) , \dots , \pi^ {- 1} \left(x _ {n + 1}\right)\right)}{p ^ {\prime} \left(x _ {1} , \dots , x _ {n}\right)} \\ = \sum_ {\pi} \frac {p \left(\pi^ {- 1} \left(x _ {1}\right) , \dots , \pi^ {- 1} \left(x _ {n + 1}\right)\right)}{p \left(\pi^ {- 1} \left(x _ {1}\right) , \dots , \pi^ {- 1} \left(x _ {n}\right)\right)} \frac {\frac {1}{d !} p \left(\pi^ {- 1} \left(x _ {1}\right) , \dots , \pi^ {- 1} \left(x _ {n}\right)\right)}{p ^ {\prime} \left(x _ {1} , \dots , x _ {n}\right)} \tag {3} \\ = \sum_ {\pi} \underbrace {p \left(\pi^ {- 1} \left(x _ {n + 1}\right) \mid \pi^ {- 1} \left(x _ {1}\right) , . . . , \pi^ {- 1} \left(x _ {n}\right)\right)} _ {\text {l a n g u a g e m o d e l i n g}} \underbrace {\mathcal {P} \left(\pi \mid x _ {1} , . . . , x _ {n}\right)} _ {\text {i n f e r r i n g l e x i c a l p e r m u t a t i o n}} \\ \end{array}
$$
As shown above,  $p'$  can be reduced to two parts, where the first part is normal language modeling and the second part is the probability distribution of lexical permutations based on seen tokens. So the lex invariant language model is implicitly learning to model  $\mathcal{P}(\pi | x_1, \dots, x_n)$ .
We can make this approximate in-context Bayesian deciphering explicit by training a small probe to predict  $\mathcal{P}(\pi |x_1,\dots,x_n)$  given the internal representation of the lexinvariant language model. We will show that this indeed recovers  $\pi$  reasonably accurately in the experiment section.
2.3 Constructing a Lexinvariant Language Model
We now consider how to construct a lexinvariant LM in practice. A typical neural language model, such as a Transformer, converts input tokens to continuous vectors using token embedding and then passes these vectors as input to the rest of the neural network. Thus, the language model  $p$  it parameterizes depends on the token embedding  $E: \mathcal{V} \to \mathbb{R}^d$ :
$$
p \left(x _ {1}, \dots , x _ {n}\right) = T \left(E \left(x _ {1}\right), \dots , E \left(x _ {n}\right)\right) \tag {4}
$$
To make a neural LM lexinvariant, we can replace the standard stable token embedding  $E$  with a randomized  $E$  and take the expectation over  $E$ . Each token  $x \in \mathcal{V}$  has an independent embedding
2The convergence rate could be better depending on the language distribution, such as on math and code, where the symbols should have clear functional meaning in context. We explore this empirically in the experiment section.
$E(x)\sim \mathcal{N}(0,\mathcal{I}_d)$  , and the language model becomes
$$
p \left(x _ {1}, \dots , x _ {n}\right) = \mathbb {E} \left[ T \left(E \left(x _ {1}\right), \dots , E \left(x _ {n}\right)\right) \right] \tag {5}
$$
Since  $E \stackrel{d}{=} E \circ \pi$ , the right-hand side is the same when  $x_{i}$  are applied with any permutation  $\pi$ , i.e., for any  $x_{1}, \ldots, x_{n}$ :
$$
\mathbb {E} [ T (E (x _ {1}), \dots , E (x _ {n})) ] = \mathbb {E} [ T (E (\pi (x _ {1})), \dots , E (\pi (x _ {n}))) ], \tag {6}
$$
showing that the Transformer with random  $E$  is lexinvariant as in Eq. 1. Now we can train this lexinvariant LM similarly to a standard LM. Concretely, we sample a new  $E$  for each training sequence and minimize the standard language modeling loss as in a standard neural LM. Here we are stochastically optimizing a variational lower bound of the standard language modeling loss with this randomized model by taking the expectation to the outside of the loss over log likelihood. Effectively, the same token gets the same random embedding within each training sequence, but different embedding across training sequences.
In practice, we focus on training decoder-only Transformers with a next token prediction objective in this work, where the model directly models  $p(x_{n+1} | x_1, \ldots, x_n)$  instead of the joint distribution. Our definitions and analysis above still hold in general. The only modification is that the final readout matrix also needs to be replaced with the same  $E$ , so that the Transformer can predict the embedding of the next token based on the embedding of input tokens.
Architecture. For all experiments, we use decoder-only Transformer architecture with T5 relative position bias [19]. We use models with 150M parameters, with 12 layers, 8 heads, head dimension 128, and MLP dimension 4096.
Training. We use the Adafactor optimizer [22], with a cosine decay learning rate schedule [13] from 0.01 to 0.001 based on preliminary experiments. We train the models from scratch for 250K steps on all the settings, with 512 sequence length and 64 batch size. We ran all of our experiments on 8 TPU cores. Our models are implemented in JAX [5].
Datasets. For datasets, we mainly use the Pile [9], a large open-source corpus that contains text collected from 22 diverse high-quality sources. We also run experiments on two additional datasets to explore their effects on the behavior of lexinvariant models: Wiki-40B [10], which contains high quality processed Wikipedia text in  $40+$  languages, and GitHub (subset of the Pile), which contains code from GitHub repositories with more than 100 stars and less than 1GB files.
3.2 Convergence to Standard Language Models
We first show empirically that lexinvariant LMs can mostly recover the next token prediction performance of standard LMs after a long enough context. As already discussed in section 2.1, the lexinvariant LM will theoretically converge to a standard LM as the context becomes long enough to resolve ambiguity. Here we verify this experimentally and show the variation of this convergence across corpora and tokenizations.
To show this, we train lex invariant and standard LMs with both character-level vocabulary (128 acci characters) and T5 default vocab (32k tokens) over the three datasets. For each model, we measure
Figure 3: Perplexity over the Pile with character-level vocabulary (left) and T5 default vocab (right).
the perplexity of each token in each sequence w.r.t. context length, smoothed by moving average within each sequence, i.e.  $P(x_{i},\dots,x_{i + k}|x_{1},\dots,x_{i})^{\frac{1}{k}}$  for context length  $i$ . We set the moving average window  $k = 100$ . We plot results over 100 sequences. As shown in figure 3, the perplexity gap between lexinvariant LM and standard LM gradually shrinks as the prefix becomes longer and longer, albeit much more slowly with a larger vocabulary. This makes intuitive sense since a larger vocabulary has more possibilities of permutations and requires many more prefix tokens to disambiguate. For the 32K vocabulary, the 512 context length will only allow the model to see a very small number of tokens, let alone to see tokens more than once. Nonetheless, the model still manages to show the trend of convergence, since even a small number of repetitions can form common patterns in grammar (such as the usage of spaces, punctuation, articles, etc). For the character-level vocabulary, the perplexity gap shrinks from 9X to less than 1X the average perplexity of the standard LM. With a context length of 511, the lexinvariant LM converges to perplexity 3.38, almost comparable to the perplexity of the standard LM of 2.00. Additionally, we observe that the gap shrinks significantly faster for models trained over Github than standard English text like Wiki-40B since code is more structured and it is easier to decipher the token permutation. We show the comparison across different datasets in Figure 7 in Appendix.
3.3 Recovering Substitution Ciphers
Here we show that lexinvariant LM is implicitly performing Bayesian in-context deciphering by testing its ability to recover cipher keys (e.g. Figure 4a) from character-level substitution ciphers, e.g. uC; kvR5W 4mfzd @f | Svcgn fw; m uCRmu; ;d ]%~ :fBn. For the lexinvariant LM, this cipher text is perceived as the same as the quick brown fox jumps over thirteen lazy dogs, due to the lexinvariant property. It will then proceed to complete the cipher text with %d: uC; @f| with the same probability as it will complete the normal text with and the fox.
Because of this, we cannot directly read out the distribution of possible cipher keys  $\mathcal{P}(\pi | x_1, \dots, x_{n-1})$  implicitly inferred by the lexinvariant LM. To do this, we train a small two-layer MLP probe on top of a frozen trained lexinvariant LM. For each training sequence, we first embed the input sequence with a randomly sampled token embedding  $E$  as described in section 2.3 and obtain the hidden activation of the final layer generated by the frozen lexinvariant LM. Then, we pass this activation through the two-layer MLP probe. Finally, instead of decoding the output activations to classification logits with the same  $E$  as in the lexinvariant LM, we instead use another learnable non-randomized token embedding matrix  $E'$  so that the probe can recover the deciphered token with stable token embeddings. Overall, we train the probe jointly with this embedding matrix  $E'$  to predict the current token. Effectively, we are training the probe to decipher the current token using the representation provided by the lexinvariant LM. We train the probe over the same corpus as the original lexinvariant LM for 10k steps. With this probe, we can directly visualize  $\mathcal{P}(\pi^{-1}(x_n) | x_1, \dots, x_n)$  inferred by the lexinvariant LM, which is effectively one row in the permutation matrix representing  $\pi$ .
Now we can use this probe to explicitly recover the cipher key. An example ground truth cipher key that we want to recover is shown in Figure 4a. Note that although the substitution cipher is only among lowercase letters, the character-level lex invariant model we use assumes that all permutations among the 128 characters are possible, making the deciphering even more challenging.
(b) Majority vote prediction
(c) Cipher key prediction accuracy
Figure 4: (a) (b): Cipher key matrix, where the vertical axis shows the cipher characters and the horizontal axis shows the deciphered letters. The highlighted entries show the correspondences between cipher characters and the actual letters, e.g. % deciphers to 1. (c): Cipher key prediction accuracy, averaged across 1000 input sequences. Context length denotes the start index of the window.
Figure 5: Predicted cipher key for windows of size 50, at indices 0, 50, 100, 200, and 400. Generated using temperature of  $T = 1$ .
Concretely, we first input ciphertext through the frozen lexinvariant LM with the probe to produce a deciphered sequence. We then select a window of size 100 in the middle of the sequence and perform a majority vote over the corresponding deciphered tokens of each cipher token seen in this window. This essentially produces a predicted cipher key matrix for each window, and we can measure its precision against the ground truth. As shown in Figure 4c, such a cipher key prediction generally has increasingly higher precision as the window is selected later in the context, and it becomes near-perfect by the end of the sequence. Specifically, the cipher key matrix produced by the last window has an average precision of  $99.6\%$  over 1000 input sequences.