Het schalen van insluitingslagen in taalmodellen

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

Samenvatting

We stellen SCONE (Schaalbaar, Gecontextualiseerd, Uitbesteed, N-gram Embedding) voor, een methode om input-embeddinglagen uit te breiden om de prestaties van taalmodellen te verbeteren naarmate de laaggrootte toeneemt. Om verhoogde decoderingskosten te vermijden behoudt SCONE de oorspronkelijke woordenschat terwijl het embeddings introduceert voor een set van veelvoorkomende n-grams. Deze embeddings bieden een gecontextualiseerde representatie voor elk invoertoken en worden geleerd met een apart model tijdens de training. Tijdens inferentie worden ze vooraf berekend en opgeslagen in off-accelerator geheugen met minimale invloed op de inferentiesnelheid. SCONE maakt twee nieuwe schalingsstrategieën mogelijk: het verhogen van het aantal gecachte n-gram embeddings en het schalen van het model dat wordt gebruikt om ze te leren, terwijl de FLOPS op inferentietijd constant blijven. We tonen aan dat het schalen van beide aspecten SCONE in staat stelt om een 1.9B parameter-baseline te overtreffen over diverse corpora, terwijl slechts de helft van de FLOPS op inferentietijd wordt gebruikt.

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.