Scaling Embedding Layers in Language Models
Authors: name: Da Yu, name: Edith Cohen, name: Badih Ghazi, name: Yangsibo Huang, name: Pritish Kamath, name: Ravi Kumar, name: Daogao Liu, name: Chiyuan Zhang

We propose SCONE (Scalable, Contextualized,
Offloaded, N-gram Embedding), a method for
extending input embedding layers to enhance language model performance as layer
size scales. To avoid increased decoding costs, SCONE retains the original
vocabulary while introducing embeddings for a set of frequent n-grams. These
embeddings provide contextualized representation for each input token and are
learned with a separate model during training. During inference, they are
precomputed and stored in off-accelerator memory with minimal impact on
inference speed. SCONE enables two new scaling strategies: increasing the
number of cached n-gram embeddings and scaling the model used to learn them,
all while maintaining fixed inference-time FLOPS. We show that scaling both
aspects allows SCONE to outperform a 1.9B parameter baseline across diverse
corpora, while using only half the inference-time FLOPS.