Mise à l'échelle des couches d'incorporation dans les modèles de langage

Scaling Embedding Layers in Language Models

February 3, 2025

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

cs.AI

Résumé

Nous proposons SCONE (Scalable, Contextualisé, Déchargé, Incrustation d'N-grammes), une méthode pour étendre les couches d'incrustation d'entrée afin d'améliorer les performances des modèles de langage à mesure que la taille de la couche augmente. Pour éviter une augmentation des coûts de décodage, SCONE conserve le vocabulaire d'origine tout en introduisant des incrustations pour un ensemble de n-grammes fréquents. Ces incrustations fournissent une représentation contextualisée pour chaque jeton d'entrée et sont apprises avec un modèle séparé lors de l'entraînement. Lors de l'inférence, elles sont précalculées et stockées dans la mémoire hors accélérateur avec un impact minimal sur la vitesse d'inférence. SCONE permet deux nouvelles stratégies de mise à l'échelle : augmenter le nombre d'incrustations d'n-grammes mises en cache et mettre à l'échelle le modèle utilisé pour les apprendre, tout en maintenant des FLOPS fixes au moment de l'inférence. Nous montrons que le fait de mettre à l'échelle ces deux aspects permet à SCONE de surpasser une référence de 1,9 milliard de paramètres à travers divers corpus, tout en utilisant seulement la moitié des FLOPS au moment de l'inférence.

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.