Масштабирование слоев встраивания в языковых моделях

Scaling Embedding Layers in Language Models

February 3, 2025

Авторы: Da Yu, Edith Cohen, Badih Ghazi, Yangsibo Huang, Pritish Kamath, Ravi Kumar, Daogao Liu, Chiyuan Zhang

cs.AI

Аннотация

Мы предлагаем SCONE (Масштабируемое, контекстуализированное, вынесенное, встраивание n-грамм), метод расширения слоев встраивания входных данных для улучшения производительности языковой модели при увеличении размера слоя. Чтобы избежать увеличения затрат на декодирование, SCONE сохраняет исходный словарь, в то время как вводит встраивания для набора часто встречающихся n-грамм. Эти встраивания обеспечивают контекстуализированное представление для каждого входного токена и обучаются с помощью отдельной модели во время обучения. Во время вывода они предварительно вычисляются и хранятся в памяти вне ускорителя с минимальным воздействием на скорость вывода. SCONE позволяет две новые стратегии масштабирования: увеличение количества кэшированных встраиваний n-грамм и масштабирование модели, используемой для их обучения, при этом сохраняя постоянное количество операций с плавающей запятой во время вывода. Мы показываем, что масштабирование обоих аспектов позволяет SCONE превзойти базовую модель с 1,9 миллиарда параметров на разнообразных корпусах, используя только половину операций с плавающей запятой во время вывода.

English

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.