Dimensionamento de Camadas de Incorporação em Modelos de Linguagem

Scaling Embedding Layers in Language Models

February 3, 2025

Autores: Da Yu, Edith Cohen, Badih Ghazi, Yangsibo Huang, Pritish Kamath, Ravi Kumar, Daogao Liu, Chiyuan Zhang

cs.AI

Resumo

Propomos o SCONE (Incorporação de N-grama Escalonável e Contextualizada), um método para estender camadas de incorporação de entrada para aprimorar o desempenho do modelo de linguagem à medida que o tamanho da camada aumenta. Para evitar o aumento dos custos de decodificação, o SCONE mantém o vocabulário original enquanto introduz incorporações para um conjunto de n-gramas frequentes. Essas incorporações fornecem uma representação contextualizada para cada token de entrada e são aprendidas com um modelo separado durante o treinamento. Durante a inferência, elas são pré-calculadas e armazenadas na memória fora do acelerador com impacto mínimo na velocidade de inferência. O SCONE possibilita duas novas estratégias de escalonamento: aumentar o número de incorporações de n-gramas em cache e escalar o modelo usado para aprendê-las, mantendo os FLOPS fixos no tempo de inferência. Mostramos que escalando ambos os aspectos, o SCONE supera um modelo de referência com 1,9 bilhão de parâmetros em diversos corpora, utilizando apenas metade dos FLOPS no tempo de inferência.

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.