Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling

We appreciate the time and effort you have taken to review our manuscript, and are truly thankful to the reviewers for the insightful comments. As a response, we address each point individually.

Analysis of the training & inference speed

To illustrate training efficiency, we show training throughputs in OLMoE experiments in the following table, where we run OE and baseline under the same hardware configurations. OLMoE-7B yields more overhead as we did not carefully optimize engineering configurations.

Table. Training throughputs for OE and baseline. We report average tokens per second in millions.

Theoretically, the additional flops introduced by OE is less than 0.5% (as shown in the table below). The overhead measured in the above experiments should mainly be introduced by the all-to-all communication (which is proportional to the size of data parallel). And these communication overheads can be further optimized through engineering techniques in the future.

Table. FLOPs per token in the forward pass.

For inference speed, we tested the prefill and decoding throughput on a single A100 GPU using the transformers library. For the OE models, the additional embedding parameters are offloaded to the CPU, incurring no GPU memory overhead. The numeric results are shown in the following table. The impact of OE on inference throughput is negligible, especially for larger models or larger batch sizes. In contrast, the sparse parameters introduced by MoE face severe memory access bottlenecks during inference. A very large batch size is required for the MoE model to achieve the same throughput as a dense model with the same activated parameters. Considering that the model inference might be carried out on more cost-effective but less computationally powerful inference GPUs, the relative overhead of OE could be further reduced.

Table. Inference speed for OE and baseline. The sequence length is fixed to 2048, and we report tokens per second for prefilling and decoding separately under various batchsizes. Settings that cause OOM are leave blank.

These analysis will be added to our paper in camera-ready version.

About Inference Diversity

This is an interesting perspective. We believe that theoretically there should be no difference, and we haven't observed such a phenomenon in practice either. In fact, examples of synthetic data can well answer this question. In the CFG task, if the problem of decreased output diversity exists, the predicted next token probability should have a larger difference from the groundtruth distribution that is calculated according to the grammatical rules. We actually calculated the KL divergence in our experiments, and the OE also shows better performance compared to the baseline, indicating that OE can model the grammar better and generate sequences as diverse as the language itself can be.

About Questions on Validation

Regarding the experimental setup, we mainly followed the experiments of OLMo and used its training data and evaluation protocols. Specifically, the evaluations include: the validation sets of public text datasets (such as C4-en-validation), on which we calculate the next token prediction loss/perplexity as the evaluation metrics; and open benchmarks (such as hellaswag), on which we calculate the zero-shot accuracy.

Figure 4 in the paper shows some of the evaluation set metrics that we are mainly concerned (explained in line 220~230). You can find comprehensive evaluation results on Figure 8 in the appendix, where we show the eval losses could has consistent gains.

Writing Issues

We appreciate the suggestions on paper writing. We'll improve the paper in the camera-ready version.

We appreciate the time and effort you have taken to review our manuscript, and are truly thankful to the reviewers for the insightful comments. As a response, we address each point individually.

The choice of using n-gram tokens.

Continuing to train BPE to obtain a larger vocabulary is indeed a more natural approach, and we have also attempted such a practice in our early experiments, where we keep frequent bigrams to construct a large one-to-one embedding table and discard unfrequent bigrams. However, the experimental results show that this approach is not better than simply hashing. This may be because the frequencies of the expanded embeddings differ too greatly (the hottest and the coldest can differ by a factor of 10,000), resulting in the parameters of the expanded embeddings not being sufficiently trained. In contrast, the hashing-based method can ensure that all expanded embeddings have an equal access frequency, thereby ensuring sufficient training and better scalability. Of course, we also encourage future work to continue exploring in this direction and to achieve end-to-end encoding of multi-granularity tokens directly from the tokenizer.

The slope of the scaling line is too small.

As for the slope of the embedding parameters' scaling curve, we are not proposing to replace the scaling law of the standard dense parameters, but rather to provide an additional, nearly COST-FREE, second growth curve. From this perspective, as long as this scaling benefit exists, it remains a better solution. Typically, you can keep scaling up the embedding parameters with negligible inference cost and barely any additional training costs other than the GPU memory usage (please refer to our response to Reviewer wJC2 for numberic results). To inference, the large embedding table can simply be offloaded to CPU. And in fact, the issue of GPU memory overhead during training can also potentially be addressed through CPU offload and prefetch (might be solved in future work). Under these circumstances, as long as we scale it up to a large enough extent, e.g., 128 times the original vocabulary size, we can achieve significant performance improvements, e.g., 400M OE loss is on par with the 1B baseline.

Typos in Equations.

Thanks for your careful reading. We apologize for the confusion on our formulations. Yes, there is a typo in equation (5). We leverage both 1-, 2-,.., n-gram tokens, so the equation is expected to be $`OE`(x)= \mathbb{E}^{V\times d}(x^{(-1)}) +\sum_{i=2}^{n} \mathbb{E}^{m\times \frac{d}{n}|k}(x^{(-i)})$ As for parameter k, it indicates a slicing factor. And flattening this equation should result in: $`OE`(x)= \mathbb{E}^{V\times d}(x^{(-1)}) +\sum_{i=2}^{n} \sum_{j=1}^{k}\mathbb{E}^{m\times \frac{d}{nk}}(x^{(-i)})W_{i,j},$ where $W_{i,j}\in \mathbb{R}^{\frac{d}{nk}\times d}$ We hope this could make the formulation clear. We'll revise the paper in the camera-ready version.

We appreciate the time and effort you have taken to review our manuscript, and are truly thankful to the reviewers for the insightful comments. As a response, we address each point individually.

Performance on Few-shot Tasks

We have conducted few-shot evaluations for in-house experiments. Our in-house baseline follows MoE architecture with 400M activated parameters and a total of 4B sparse parameters. And we implement OE to scale up embedding table m to 36M. Here, we share some of the results.

Reasoning Related Benchmarks:

Knowledge Related Benchmarks:

Math Related Benchmarks

We will add these results to our paper in the camera-ready version.

Analysis of the training & inference speed

We put thorough analysis on response to Reviewer wJC2. In conclusion, OE has 4% and 8% training overhead on OLMoE 1.3B and OLMoE 7B. We emphasize the larger overhead in OLMoE7B is owing to that we did not carefully optimize engineering configurations. Theoretically, OE only introduces less than 0.4% FLOPs. The overhead measured in the experiments should mainly be introduced by the all-to-all communication (which is proportional to the size of data parallel). And these communication overheads can be further optimized through engineering techniques in the future.

As for inference, we tested the prefill and decoding throughput on a single A100 GPU using the transformers library. For the OE models, the additional embedding parameters are offloaded to the CPU, incurring no GPU memory overhead. The impact of OE on inference throughput is negligible (around 2%).

How would one choose the sparsity tradeoff given fixed memory budget if combining both MoE and OE?

It's an interesting question. First of all, it must be admitted that the performance of OE with the same number of sparse parameters is not as good as that of MoE. However, we'd like to emphasize that scaling up the parameters of MoE is not free. During the inference, MoE increases the GPU memory overhead. Moreover, under small batch sizes, there are memory access bottlenecks, and the inference efficiency is usually far from reaching that of a dense model with the same activated parameters (as shown in the inference speed table on our response to Reviewer wJC2). In contrast, during the inference process, OE can be completely offloaded to the CPU, incurring no GPU memory overhead, and the reduction in efficiency is almost negligible. Our recommended approach is to first determine the size of MoE according to the inference requirements, and then use the remaining GPU memory to scale the embedding table.

In addition, considering that OE has a good property during the training phase: the input of OE only depends on the token id, it provides more room for engineering optimization. For example, during the training phase, OE can possibly be offloaded to the CPU and the embeddings of the next micro batch can be prefetched and overlap with current micro batch's forward. On this occasion, OE and MoE can operate independently without interfering with each other. We are also exploring different engineering solutions for this issue.

Related work & Equation Typos

Thanks for your careful reading, equation 5 has some typos. The correct formula should be: $`OE`(x)= \mathbb{E}^{V\times d}(x^{(-1)}) +\sum_{i=2}^{n} \mathbb{E}^{m\times \frac{d}{n}|k}(x^{(-i)})$ And we appreciate the suggestions on paper writing and related work. We will improve the paper in the camera-ready version.

We appreciate the time and effort you have taken to review our manuscript, and are truly thankful to the reviewers for the insightful comments. As a response, we address each point individually.

About Questions on Claims And Evidence

As far as the over-encoding technique itself, indeed, it is typically a network module that improves performance and is irrelevant to the tokenizer. However, we believe the insight behind is worth noticing for the tokenizer design:

1. We show that the encoding and decoding in tokenization should be separately considered to maximize model performance. One should be careful in increasing decoding vocabulary.
2. Hierarchical designs improve OE, but can multi-granularity tokens be properly considered in tokenizer design itself? By introducing OE, we show that simply leveraging hashed n-gram tokens could yield significant improvements to LLM. We could expect this gain to be pushed further with delicate tokenizer designs.

As for the "more efficient LLM", we mainly consider our efficiency on the proper use of sparse embedding parameters, similar to MoE, which improves model performance with negligible training and inference overheads. It is efficient as you can obtain powerful models with possibly half training and inference budget.

About tease apart the gains.

The reviewer mentioned:

It would be nice to see experiments that teased apart if the gains are from having such a larger vocabulary or from the explicit modeling of n-gram composition. It would have been nice to see an ngram representation made of the sum of each token in the n-gram with shared embeddings

Actually, we tried such an approach in the early stage of our research. Summing n-gram tokens from a shared embedding doesn't work because there is no positional distinction for each token in the n-gram (addition is commutative). This makes the current input token unclear, which is detrimental to the model's performance. If we simply remove the condition of sharing embeddings and apply n different embedding tables to these n tokens, the model can achieve some benefits. Later, we found that the gains from this implementation are close to those predicted by the OE scaling curve with $m=(n - 1)\times V$. Essentially, this implementation is a product decomposition for the full n-gram embedding table. Alternatively, you can also view it as OE using a special hash function, and the choice of hash function isn't crucial for OE.

In addition, our ablation study has also ablated the effect of increasing vocabulary size and n-gram composition. In Figure 5, we fix n=2 and varied vocabulary size, showing the gains from scaling vocabulary size. In Table 3, we fix the embedding parameters and ablated gains introduced by hierarchical n-gram composition. We conclude that both vocabulary size and well-designed n-gram composition are important to OE.

About the under-trained tokens

It's an interesting perspective to consider under-trained tokens. However, we believe OE should not have such a problem. Under-trained tokens are potentially harmful mainly due to their under-trained embedding vectors. For OE, n-gram embeddings are equally visited during training owing to the many to one hashing. Though there might be some unknown n-gram tokens occur during inference, their embedding vectors are guaranteed to be frequently trained. So, we believe the unseen n-gram tokens are more like a normal spelling mistake and will not cause catastrophic consequences. Moreover, the token embeddings should contain at least one well-trained uni-gram embedding under the hierarchical design, which also improves robustness against unseen n-grams.

What kind of text benefits the most from over-encoding

Examining how n-gram embedding contributes to some specific token is kind of difficult. As the n-gram embeddings improve the keys and values in the attention layer simultaneously. As a result, it is difficult to tell where the improvements on specific tokens come from. However, from the evaluation results, we do notice that OE has significantly larger improvement in knowledge-related tasks (see few-shot results in our response to Reviewer ye13). We hypothesis that coarse-grained token embeddings help memorizing concepts of proper nouns, which are usually broken into several tokens under the base tokenizer.

About the Relevance of BLT

We are sorry for the word 'patching' that misleads the understanding. We typically want to talk about the n-gram embedding technique used in BLT. They apply n-gram byte embeddings to the byte-level sequence, which does not reduce byte sequence length as well. We'll revise the paper to make this more clear.

Related Work & Writing Issues

We appreciate the suggestions on related work and paper writing. We'll improve the paper in the camera-ready version.