Escalando Capas de Incrustación en Modelos de Lenguaje

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

Resumen

Proponemos SCONE (Incrustación de N-gramos Escalable, Contextualizada, Desplazada), un método para extender las capas de incrustación de entrada con el fin de mejorar el rendimiento del modelo de lenguaje a medida que aumenta el tamaño de la capa. Para evitar un aumento en los costos de decodificación, SCONE conserva el vocabulario original al mismo tiempo que introduce incrustaciones para un conjunto de n-gramos frecuentes. Estas incrustaciones proporcionan una representación contextualizada para cada token de entrada y se aprenden con un modelo separado durante el entrenamiento. Durante la inferencia, se precalculan y almacenan en memoria fuera del acelerador con un impacto mínimo en la velocidad de inferencia. SCONE permite dos nuevas estrategias de escalado: aumentar el número de incrustaciones de n-gramos en caché y escalar el modelo utilizado para aprenderlas, todo ello manteniendo una cantidad fija de FLOPS en tiempo de inferencia. Mostramos que al escalar ambos aspectos, SCONE supera a una línea base de 1.9 mil millones de parámetros en diversos corpus, utilizando solo la mitad de los FLOPS en tiempo de inferencia.

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.