Word-Level Representation From Bytes For Language Modeling

Byte-level and word-level tokenization are two extremes when it comes to representing text into indices of a vocabulary. The simplest method, character-level tokenization is to let the vocabulary $V$ be the alphabet. This method is hard to practice on multilingual text as the total size of alphabets of multiple languages can be up to a million. One way to tackle this problem is switching to token-free byte-level input because there are only 256 bytes, which can represent text in the smallest level of granularity. However, this tends to yield very long sequences with sparse information. Word-level tokenization lies on the other end of the spectrum. Word-level tokenization tends to require a very large vocabulary to cover a diverse amount of text. The copra BERT trained on contains more than 10 millions unique words. This leads to a large embedding layer and intense softmax computation when predicting. Tailored methods such as negative sampling are used in Feng et al. (2022) to train a word-level model successfully. Transformers mainly use subword tokenization right when it was first introduced. An easy way to explain subword tokenization is it uses a vocabulary containing a set of commonly occurring word segments like ’ing’ and ’ly’. However, Gowda and May (2020) and Xu et al. (2020) show that hyperparameters of subword methods, especially the size of the dictionary, affect final performance. Thus, some models Brown et al. (2020) allow a certain amount of redundancy for tokenization, resulting in varying granularities of token in its vocabulary and indirect learning.

This paper chooses to combine byte-level and word-level methods as proposed in Kim et al. (2015). However, to truly achieve token-free, we do not want to maintain a word-level vocabulary, therefore we need to come up with some rules or methods to split text sentences into words to provide bytes and word boundaries to the model. For splitting words in a sentence, the boundary of a word is well-defined in English. We split phrases into words simply by white space, punctuation, camel case, and snake case after following the preprocessing procedure used in Devlin et al. (2018). However, this is not the optimal choice for some non-English corpora, we leave this issue for future work.

After splitting a sentence to words, in order to construct word-level embedding from bytes for one of these words, we use cross attention mechanism proposed in Vaswani et al. (2017), which is an effective and more general pooling mechanism compared to character-aware convolutions. Given a sentence and let the $i$th word of that sentence contains $L$ bytes, we use the corresponding byte-level embeddings $(e_{i1}^{B},e_{i2}^{B},...,e_{iL}^{B})$ from a byte-level embedding lookup table $E\in R^{256\times d}$ to build the word embedding $e^{W}_{i}$ with a embedding size of $d$. We pack these byte-level embeddings together as key matrix $K_{i}$ and value matrix $V_{i}$ for cross-attention. Specifically, we have

$$K_{i}=\textrm{Concat}(e_{i1}^{B},e_{i2}^{B},...,e_{iL}^{B})\cdot W_{k}\\$$

(1)

$$V_{i}=\textrm{Concat}(e_{i1}^{B},e_{i2}^{B},...,e_{iL}^{B})\cdot W_{v}\\$$

(2)

$$e^{W}_{i}=\textrm{softmax}(\frac{Q_{i}K_{i}^{T}}{\sqrt{d}})V_{i}\\$$

(3)

Here, $Q_{i}\in R^{1\times d}$ denotes $i$th learnable word-query for position $i$ and $W_{k}\in R^{d\times d},W_{v}\in R^{d\times d}$ are matrices for key and value projections. We find having multiple word-queries for different positions helps convergence. After pooling by cross-attention, word-level embeddings are processed by a position-wise feed-forward layer as the standard procedure. Our word-level embeddings $e^{W}_{i}$ is then added by residual connection and other types of embeddings. Followed by a layer normalization Ba et al. (2016) and a linear projection, we obtain our final word-level hidden states $H^{W}_{i}$ for our encoder. This linear projection tactic resembles the Factorized embedding parametrization method introduced in Lan et al. (2019) and the rebalanced embeddings in Chung et al. (2020).

$$H^{W}_{i}=\textrm{LN}((\textrm{FFN}(e^{W}_{i})+e^{W}_{i}+e^{\textrm{pos}}_{i}+e^{\textrm{type}}_{i}))\cdot W_{\textrm{proj}}\\$$

(4)

where $e^{\textrm{pos}}$ is positional embedding, $e^{\textrm{type}}$ is token type embedding, $W_{\textrm{proj}}\in R^{d\times d_{\textrm{encoder}}}$ projects the word-level embedding to match the hidden size of the backbone encoder.

While with the method above, it is sufficient for language modeling. Still, predicting at word-level is computationally intense due to the large size of the lookup table. Feng et al. (2022) showcases a receipt to train a word-level BERT by utilizing negative sampling. But such a method is not direct to implement in our setting since we don’t even have a word-level vocabulary. Jozefowicz et al. (2016) proposes a contrastive learning based method that dynamically generates output embedding given the character-level embedding sequence of the target word. But contrastive learning can easily lead to model collapse and requires extra caring to train. There is another way to train a fully vocabulary-independent model, Jozefowicz et al. (2016); Xue et al. (2021); Tay et al. (2021) suggest predicting a word by auto-regressively decoding a sequence of characters given a word-level representation. However, due to efficiency constraints, Jozefowicz et al. (2016)’s character-level decoding is based on a pretrained then frozen word-level model. We try to avoid these complex designs and instead choose to predict on sub-word level. Chung et al. (2020) shows that decoupling the hidden sizes of input and output embeddings can enhance performance. We believe it is feasible to take a further step and decouple input and output text granularity. By predicting subwords we can save a lot of effort and keep our model behave similarly to whole-word masking models and make it comparable to BERT baseline if we reuse its subword vocabulary. Besides, this method is token-free anyway when transfer learning on non-generative tasks. Subword-level features can be constructed by up-sampling from word-level hidden states. To keep things simple, we use a positional up-sampling method. Say we have the representation $H_{i}$ of the $i$th word, that word contains multiple subwords and we want to want to predict the $j$th one, we have

$$H_{ij}=H_{i}+P_{j}$$

(5)

Where $P_{j}$ is a positional query for the $j$th subword of a word. We feed $H_{ij}$ into a prediction head for the final subword level prediction. We find that simply adding positional queries to the word-level representation respectively works better than more complex attention up-sampling Lee et al. (2018).

While Byte2Word makes more sense on multilingual data, we evaluate our method in English due to budget constraint. We follow Izsak et al. (2021) and pretrain our model on English Wikipedia and BookCorpus Zhu et al. (2015). This corpora consists  20GB of raw text, and  10 million unique words, which should be enough to test if our method can learn a diversity of words. We provide some statistical information about the corpora in Table 2. We follow the masking strategy in Devlin et al. (2018) but slightly modify the input format to use specific byte values for special tokens. Specifically we substitute "[PAD]" for HEX0 , "[UNK]" for HEX1, "[CLS]" for HEX2, "[SEP]" for HEX3 and "[MASK]" for HEX4. For instance, instead of "[CLS] a [MASK] sentence exmaple. [SEP]", Byte2Word takes the input as "\x00 a \x03 sentence example. \x02". In this way our model learns each special token from a single index, sparing extra hassle. Also to control memory usage, we limit the max amount of bytes a word can contain to 128. Similar to $\textrm{BERT}_{\textrm{Large}}$, we adopt a 24-layer transformer encoder with 16 attention heads and 1024 hidden sizes for our model backbone. To minimize computation overhead, we limit the size of Byte2Word embedding layer $d$ to 192. In this way, our model consume less than $0.2\%$ extra FLOPS per inference compared to $\textrm{BERT}_{\textrm{Large}}$. We also find that increasing the width of byte2word embedding slows convergence in ablation study in 5.4. To ensure pretraining efficiency we follow Liu et al. (2019); Izsak et al. (2021) and discard next-sentence prediction. We build our code base upon Izsak et al. (2021) and train our Byte2Word model for 230k steps on 8x NVIDIA T4 16 GB for roughly a month.

We test the performance of our model on GLUE benchmark Wang et al. (2018), the standard evaluation suite on language understanding for pretrained language model, and compare our result to subword-level baseline $\textrm{BERT}_{\textrm{Large}}$. During finetuning, instead of performing a grid search over sets of hyper-parameters, we use a fixed set of hyper-parameters across all tasks for each model. Table 1 shows test set results of the GLUE benchmark. Our Byte2Word model performs on par with $\textrm{BERT}_{\textrm{Large}}$ on MNLI, QNLI and QQP, outperforms it on SST-2, MRPC and STS-B. However, our model has lower results on COLA and RTE. Overall, this amounts to a <0.5% difference in the average score, proving that our method can reach on-par performance with sub-word level model, while being toke-free for transfer learning and can accept and adapt to new words.

To explore Byte2Word’s ability to handle noisy input, similar to ByT5 Xue et al. (2021), we test our method on learning synthetic noise injected during transfer learning. We experiment four synthetic noising schemes listed below

• •

  Random drop: 10% of characters will be dropped in a sentence

• •

  Repetition: 20% of characters will be repeated for 1-3 times(drawn uniformly).

• •

  Uppercase: All characters will be converted to uppercase

• •

  Random case: All characters will be converted to uppercase or lowercase randomly. This pattern is simulating Alternating caps²²2https://en.wikipedia.org/wiki/Alternating_caps existing on the Internet

These types of noise are added to training and evaluation data. We compare our method to vanilla byte-level model ByT5 and case-sensitive subword-level model BERT Cased on MNLI. For ByT5, we adopt a method introduced in EncT5 Liu et al. (2021) to finetune ByT5’s encoder on MNLI but keep the training budget close to our baseline model. Table 3 shows the test performance of byte-level, subword-level and our hybrid method Byte2word on MNLI. Byte2Word has superior performance on all types of noise compared to subword-level model. On random dropping, ByT5 Encoder has less performance loss, but note that ByT5 is trained on mC4 Xue et al. (2020), a much larger multilingual dataset crawled from the Internet which contains noisy text. On Random case, it’s no surprise to see our method perform terribly, due to the strategy of splitting camel case in the input. This result also shows that word boundary is critical for learning language representations. Given the result of injecting random case noise after word segmentation degenerate minimal performance in Table 3, we presume the performance would be much better if we make fewer assumptions on data preprocessing. It’s also interesting to see that while Byte2Word model is pretrained using an uncased, subword-level vocabulary, this setting does not hinder finetuning on noisy text. Our Byte2word method can learn the noise pattern on the fly and has least degeneration on various noise types.

Unlike Xue et al. (2021), we evaluate cross-lingual transfer ability of our method by finetuning our English-only pretrained model on non-English downstream datasets. Since our model is not pretrained on multilingual corpora, it has to utilize cognates or shared grammar to facilitate learning. We test our model on XNLI Conneau et al. (2018) which is basically MNLI translated into multiple languages. Our results on a subset of XNLI are shown in Table 4. While the performance of our method is still behind multilingual language models, it has the better ability when adapting to new

We analyze our learned Byte2Word model by calculating the cosine similarity of the learned word-level representations. Compare to BERT’s sparse embedding space(Figure 3), we speculate representations learned by our Byte2Word method seems to contain less semantic meaning, as the cosine distance between "live" and "liver" is too close. Of course, hidden states usually are not as sparse as items of an embedding table. Based on these findings, we use word pairs from Kirov et al. (2018) and compute Spearman’s correlation between cosine similarity of the learned representations and various types of edit distance of these pairs. Results are in Table 5. We can see that our representations have a strong correlation with the jaro-winkler distance, which factors in matching characters and transpositions. We believe that a single layer of cross-attention with limited hidden size serves as an information bottleneck that makes representations only contain character-level, shallow abstraction and saves the high-level learning for the following encoder layers. Also, it’s interesting to point out that while no case information existed in supervised signal during pretraining, Byte2Word still can distinguish uppercase and lowercase in text.

To make sure that our choice of small embedding size $d$ for Byte2Word model does not degrade performance, we do an ablation study on the convergence rate of different embedding sizes. We use the previous setting in 4.1 to pretrain Byte2Word models with multiple embedding sizes of 96, 192, 512, and 1024 but limit the total training step to 23k, roughly 10% of our pertraining budget. The result is shown in Figure2. We can see that while all models reach similar perplexity after a certain amount of computing, 192 has the fastest convergence rate. With the mind of limiting the hidden size of Byte2Word module as small as we can, choosing 192 seems like a great choice that balances computation and performance.

To showcase our method is superior to previous, 1D convolution-based methods that consult the characters of a token to produce representation of a single word, we pretrain a Byte2Word model and CharacterBERT using the 10% of our pretraining budget in 4.1 and evaluate their performance on MNLI, the most representative task of the GLUE benchmark. Similar to Aguilar et al. (2021), we use BERT’s vocabulary as a pseudo word-level vocabulary to minimize architecture differences compared to subword-level baseline. Results in Table 6 show that attention based pooling method of Byte2Word performs better than 1D CNN pooling.

$\textrm{BERT}_{\textrm{Large}}$’s subword lookup table has a size of 30522 and contributes roughly 31M of parameters, nearly 10% of the total model size. On the contrary, our Byte2word method has a small size lookup table and a following cross-attention pooling, composing about 0.4M of parameters. While it has 2x more parameters than 1D CNN character-aware embedding method, compared to BERT it is negligible. However, embedding lookup albeit its size does not require computation, but our Byte2Word method has extra computation with a ceiling of 25 MFLOPS per inference if we set our byte-level hidden size to 192 and input sentence to 512 words.