Mischung von Tiefen: Dynamische Zuweisung von Berechnungen in auf Transformer basierenden
SprachmodellenMixture-of-Depths: Dynamically allocating compute in transformer-based
language models
Transformer-basierte Sprachmodelle verteilen FLOPs gleichmäßig über Eingabesequenzen. In dieser Arbeit zeigen wir, dass Transformer stattdessen lernen können, FLOPs (oder Berechnungen) dynamisch bestimmten Positionen in einer Sequenz zuzuweisen, wobei die Zuweisung entlang der Sequenz für verschiedene Schichten über die Modelltiefe optimiert wird. Unsere Methode setzt ein Gesamtberechnungsbudget durch Begrenzung der Anzahl von Tokens (k) um, die an der Self-Attention und MLP-Berechnungen in einer bestimmten Schicht teilnehmen können. Die zu verarbeitenden Tokens werden vom Netzwerk mithilfe eines Top-k-Routingmechanismus bestimmt. Da k a priori definiert ist, verwendet dieses einfache Verfahren einen statischen Berechnungsgraphen mit bekannten Tensorgrößen, im Gegensatz zu anderen bedingten Berechnungstechniken. Dennoch, da die Identitäten der k Tokens variabel sind, kann diese Methode FLOPs nicht gleichmäßig über die Zeit- und Modelltiefe-Dimensionen verteilen. Somit ist der Berechnungsaufwand insgesamt vollständig vorhersehbar, aber dynamisch und kontextsensitiv auf der Token-Ebene. Modelle, die auf diese Weise trainiert werden, lernen nicht nur, Berechnungen dynamisch zuzuweisen, sondern tun dies auch effizient. Diese Modelle erreichen die Baseline-Performance für äquivalente FLOPS und Trainingszeiten, benötigen jedoch nur einen Bruchteil der FLOPs pro Vorwärtspass und können während der Post-Training-Samplingphase um mehr als 50\% schneller sein.