Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Spelling Bee Embeddings for Language Modeling

Markus N. Rabe, Judith Clymo, Zheren Dong

TL;DR

This work tackles tokenization-induced inefficiencies in language models by introducing spelling bee embeddings, which augment standard token embeddings with a compact byte-level spelling signal derived from the token's UTF-8 representation. The approach preserves the original architecture and incurs negligible overhead, yielding consistent improvements across model sizes from 40M to 800M parameters and achieving compute efficiency comparable to about an 8% reduction in FLOPs or data for the same test loss. Empirical results on standard benchmarks show broad gains, with notable improvements in spelling-related tasks such as counting Rs in strawberry and the Count/Index tasks, though some limits exist (e.g., reverse spelling). Ablation analyses confirm the necessity of the combined token-and-byte embedding and reveal the approach's robustness to data scale, suggesting practical adoption and prompting reexamination of tokenization and vocabulary choices in future models.

Abstract

We introduce a simple modification to the embedding layer. The key change is to infuse token embeddings with information about their spelling. Models trained with these embeddings improve not only on spelling, but also across standard benchmarks. We conduct scaling studies for models with 40M to 800M parameters, which suggest that the improvements are equivalent to needing about 8% less compute and data to achieve the same test loss.

Spelling Bee Embeddings for Language Modeling

TL;DR

This work tackles tokenization-induced inefficiencies in language models by introducing spelling bee embeddings, which augment standard token embeddings with a compact byte-level spelling signal derived from the token's UTF-8 representation. The approach preserves the original architecture and incurs negligible overhead, yielding consistent improvements across model sizes from 40M to 800M parameters and achieving compute efficiency comparable to about an 8% reduction in FLOPs or data for the same test loss. Empirical results on standard benchmarks show broad gains, with notable improvements in spelling-related tasks such as counting Rs in strawberry and the Count/Index tasks, though some limits exist (e.g., reverse spelling). Ablation analyses confirm the necessity of the combined token-and-byte embedding and reveal the approach's robustness to data scale, suggesting practical adoption and prompting reexamination of tokenization and vocabulary choices in future models.

Abstract

We introduce a simple modification to the embedding layer. The key change is to infuse token embeddings with information about their spelling. Models trained with these embeddings improve not only on spelling, but also across standard benchmarks. We conduct scaling studies for models with 40M to 800M parameters, which suggest that the improvements are equivalent to needing about 8% less compute and data to achieve the same test loss.
Paper Structure (33 sections, 2 equations, 3 figures, 5 tables)

This paper contains 33 sections, 2 equations, 3 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Spelling Bee Embeddings
  4. Experimental Setup
  5. Datasets.
  6. Architecture.
  7. Notes on efficiency.
  8. Hyperparameters.
  9. Chinchilla optimal training.
  10. Optimizer.
  11. Learning rate schedule.
  12. Initialization.
  13. Counting FLOPs.
  14. Results
  15. Counting Rs
  16. ...and 18 more sections

Figures (3)

  • Figure 1: Spelling bee embeddings augment a standard token embedding $e_{\text{tok}}$ with a normalized sum of byte-level embeddings for the first 16 bytes of the token. Byte embeddings are position-encoded with RoPE using character positions inside the token.
  • Figure 2: Optimal learning rate as function of model size. Model sizes indicate the configurations listed in the Chinchilla paper hoffmann2022training. Actual parameter counts deviate due to changes in vocabulary size and the use of GQA.
  • Figure 3: Scaling study using the test loss. We trained a series of models with growing parameter counts, with and without spelling bee embeddings. Each dot represents a different training run according to the methodology described in Section \ref{['sec:setup']}. The lines are fitted to the data points with a shifted power law.