Scaling degli Strati di Embedding nei Modelli Linguistici

Scaling Embedding Layers in Language Models

February 3, 2025

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

cs.AI

Abstract

Proponiamo SCONE (Scalable, Contestualizzato, Spostato, Embedding di N-grammi), un metodo per estendere i livelli di incorporamento dell'input al fine di migliorare le prestazioni del modello linguistico all'aumentare delle dimensioni del livello. Per evitare un aumento dei costi di decodifica, SCONE mantiene il vocabolario originale introducendo al contempo embedding per un insieme di n-grammi frequenti. Questi embedding forniscono una rappresentazione contestualizzata per ciascun token di input e vengono appresi con un modello separato durante l'addestramento. Durante l'inferenza, vengono precomputati e memorizzati nella memoria esterna all'acceleratore con un impatto minimo sulla velocità di inferenza. SCONE consente due nuove strategie di scalabilità: aumentare il numero di embedding di n-grammi memorizzati nella cache e scalare il modello utilizzato per apprenderli, il tutto mantenendo i FLOPS fissi durante l'inferenza. Dimostriamo che scalare entrambi gli aspetti consente a SCONE di superare un modello di base con 1,9 miliardi di parametri su corpora diversi, utilizzando solo la metà dei FLOPS durante l'inferenza.

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.