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
Published: 2025-02-03 18:59:32

Abstract: 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.