Skalierung von Einbettungsschichten in Sprachmodellen

Scaling Embedding Layers in Language Models

February 3, 2025

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

cs.AI

Zusammenfassung

Wir schlagen SCONE (Skalierbare, Kontextualisierte, Ausgelagerte, N-Gramm Einbettung) vor, eine Methode zur Erweiterung von Eingabe-Einbettungsschichten zur Verbesserung der Leistung von Sprachmodellen, wenn die Schichtgröße skaliert wird. Um erhöhte Dekodierungskosten zu vermeiden, behält SCONE den ursprünglichen Wortschatz bei und führt Einbettungen für eine Reihe häufiger n-Gramme ein. Diese Einbettungen bieten eine kontextualisierte Darstellung für jedes Eingabe-Token und werden mit einem separaten Modell während des Trainings erlernt. Während der Inferenz werden sie vorberechnet und in einem Speicher außerhalb des Beschleunigers mit minimalem Einfluss auf die Inferenzgeschwindigkeit gespeichert. SCONE ermöglicht zwei neue Skalierungsstrategien: die Erhöhung der Anzahl der zwischengespeicherten n-Gramm-Einbettungen und die Skalierung des Modells, das zu ihrer Erstellung verwendet wird, während gleichzeitig die festgelegten FLOPS zur Inferenzzeit beibehalten werden. Wir zeigen, dass die Skalierung beider Aspekte es SCONE ermöglicht, eine 1,9 Milliarden Parameter umfassende Baseline über verschiedene Korpora hinweg zu übertreffen, während nur die Hälfte der FLOPS zur Inferenzzeit verwendet werden.

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.