- Mixture-of-Depths (MoD) von Google DeepMind ist ein Verfahren, bei dem Transformer-Sprachmodelle nicht für alle Tokens die gleichen FLOPs verwenden, sondern pro Layer nur die Tokens an Attention- und MLP-Berechnungen teilnehmen lassen, die sie benötigen
- Ein Router pro Layer erzeugt für jedes Token ein skalares Gewicht, und nur Tokens innerhalb einer vorab festgelegten Top-k-Kapazität durchlaufen den Block; die übrigen umgehen ihn über eine Residual Connection
- Indem k im Voraus festgelegt wird, bleiben Rechengraph und Tensorgrößen statisch, während sich die Verteilung der Berechnung pro Token je nach Kontext dynamisch verändert
- In Experimenten erzielte MoD bei gleichem isoFLOP-Budget entweder bei derselben Trainingszeit einen niedrigeren Loss als der Basis-Transformer oder reduzierte bei gleicher Leistung die FLOPs pro Forward Pass, wodurch Training und Sampling-Schritte schneller wurden
- Die beste Konfiguration wandte Blöcke mit 12,5 % Kapazität auf jeden zweiten Block an; beim auto-regressiven Sampling verringerte ein prädiktorbasiertes Routing das Nicht-Kausalitätsproblem von Top-k
Auf welche Rechenverschwendung MoD abzielt
- Herkömmliche Transformer-Sprachmodelle verwenden in einem Forward Pass für alle Tokens die gleiche Rechenmenge
- MoD basiert auf der Annahme, dass nicht alle Tokens und Sequenzen für die Vorhersage gleich viel Zeit und Aufwand benötigen, und verteilt Berechnung auf die benötigten Tokens, während das gesamte Rechenbudget sinkt
- Bestehende bedingte Berechnung (conditional computation) kann dynamische Rechengraphen erzeugen und passt damit möglicherweise nicht zu den statischen Rechengraphen, die aktuelle Hardware bevorzugt
- Wenn der Nutzer vor dem Training ein gesamtes Rechenbudget festlegt, lernt das Modell innerhalb dieses Budgets, wo Berechnung pro Token und pro Layer eingesetzt werden soll
- Rechenaufwand, Speichernutzung und FLOPs pro Forward Pass lassen sich im Voraus vorhersagen
- Welche Tokens geroutet werden, hängt vom Eingabekontext ab
Routing-Struktur und wie Berechnung eingespart wird
- MoD verwendet wie ein MoE-Transformer einen Router, schickt Tokens aber nicht zu einem von mehreren Experts, sondern auf einen von zwei Pfaden
- Standard-Blockberechnung mit Attention und MLP
- Residual Connection, die den Wert unverändert weitergibt
- Die Residual Connection hat geringe Rechenkosten, und die Blockausgabe wird vom Eingabewert bestimmt
- Wenn die Kapazität eines Blocks kleiner als die Gesamtzahl der Tokens T gesetzt wird, sinken die FLOPs pro Forward Pass im Vergleich zu einem Basis-Transformer
- Wird die Kapazität beispielsweise auf T/2 reduziert, wird bei der Query-Key-Matrixmultiplikation der Attention aus T² nur noch (T/2)², also etwa 25 % der FLOPs
- Bei Kapazität 0 ist das Verfahren zwar sehr schnell, nutzt aber den Großteil der Transformer-Parameter nicht und verschlechtert dadurch die Leistung; bei Kapazität T entspricht es dem Basis-Transformer
- Ziel ist es, zwischen diesen beiden Extremen einen Punkt zu finden, an dem die Leistung dem Basis-Transformer entspricht oder besser ist und zugleich die Schrittgeschwindigkeit steigt
Top-k Expert-Choice-Routing
- In jedem Layer l erzeugt der Router für jedes Token-Embedding ein skalares Gewicht
- Entsprechend der vom Nutzer festgelegten Kapazität C nehmen unter den Sequenz-Tokens nur die Top-k-Tokens an Attention- und MLP-Berechnungen teil
- Nicht ausgewählte Tokens überspringen die Blockberechnung und werden über die Residual Connection an den nächsten Layer weitergegeben
- Dieses Verfahren basiert auf Expert-Choice-Routing
- Beim Token-Choice-Routing wählen Tokens ihren Pfad selbst, was zu Problemen beim Load Balancing führen kann
- Beim Expert-Choice-Routing wählt jeder Pfad seine Top-k-Tokens aus, sodass die Kapazität exakt gefüllt wird
- In MoD gibt es effektiv nur zwei Pfade, „Blockberechnung“ und „Umgehung“, sodass eine einzelne Top-k-Operation die Tokens in zwei Gruppen teilen kann
- Die Router-Gewichte werden mit der Blockausgabe multipliziert, liegen damit auf dem Gradient Path und werden beim Sprachmodelltraining durch Gradientenabstieg beeinflusst
Das Nicht-Kausalitätsproblem beim Sampling
- Die Top-k-Operation des Expert-Choice-Routings ist nicht-kausal (non-causal)
- Ob ein Token in den Top-k landet, lässt sich erst entscheiden, wenn auch die Router-Gewichte späterer Tokens bekannt sind
- Beim auto-regressiven Sampling können keine Informationen über zukünftige Tokens verwendet werden
- Um dies zu vermeiden, wurden zwei Methoden getestet
- Ein zusätzlicher Binary-Cross-Entropy-Loss, damit die Sigmoid-Ausgabe des Routers vorhersagt, ob ein Token für Top-k ausgewählt wird
- Ein kleiner zusätzlicher MLP-Prädiktor, der dieselbe Eingabe wie der Router erhält, aber mit Stop-Gradient arbeitet
- Die erste Methode beeinflusste das Hauptziel des Sprachmodelltrainings um etwa 0,2–0,3 %
- Die zweite, prädiktorbasierte Methode beeinflusste das Ziel des Sprachmodelltrainings nicht und hatte auch keinen signifikanten Einfluss auf die Schrittgeschwindigkeit
- Die zusätzliche Vorhersageaufgabe erreichte experimentell schnell eine Genauigkeit von 99 %
Versuchsergebnisse und Beobachtungen
- Die ersten Experimente suchten MoD-Hyperparameter innerhalb eines Gesamtbudgets von 6e18 FLOPs
- Die optimale MoD-Konfiguration wendete Routing auf jeden zweiten Block an und setzte Top-k auf 256
- Bei einer Sequenzlänge von 2048 durchlaufen damit 256 Tokens, also 12,5 %, Attention und MLP
- Die übrigen 1792 Tokens, also 87,5 %, umgehen den Block
- Eine MoD-Variante mit 220 Mio. Parametern lieferte leicht bessere Leistung als die isoFLOP-optimierte Baseline und war während des Trainings bei der Schrittgeschwindigkeit um mehr als 60 % schneller; auf derselben Hardware war die gesamte Wall-Clock-Trainingszeit jedoch nahezu gleich
- Auch in isoFLOP-Analysen bei 6e18, 2e19 und 1e20 FLOPs zeigten die FLOP-optimalen MoD-Modelle tendenziell mehr Parameter als die Baseline
- Es gibt MoD-Varianten, die bei gleicher Leistung oder niedrigerem Loss weniger FLOPs pro Forward Pass benötigen
- Die höhere Schrittgeschwindigkeit kommt aus zwei Faktoren
- Einige Tokens umgehen den Block, wodurch bei gleicher Modellgröße die FLOPs pro Parameter sinken
- Es kann eine kleinere MoD-Variante gewählt werden, die dieselbe Leistung wie die Baseline erreicht
- MoD mit stochastischem Routing zeigte deutlich schlechtere Leistung als der Basis-Transformer und normales MoD, was die Bedeutung von gelerntem Routing unterstreicht
- Für die auto-regressive Auswertung wurden 256.000 Sequenzen und 500 Mio. Tokens an Held-out-Daten verwendet; auch mit prädiktorbasiertem Routing blieb der Leistungsverlust gering
- MoD lässt sich auch mit MoE zu Mixture-of-Depths-and-Experts (MoDE) kombinieren
- Staged MoDE entscheidet vor der Attention, ob ein Token in den Block geschickt oder umgangen wird
- Integrated MoDE integriert MoD-Routing, indem zwischen bestehende MLP-Experts ein „No-op“-Expert eingefügt wird
- Die Routing-Analyse zeigte Muster, bei denen einige Tokens häufig mehrere Blöcke durchlaufen, während andere wann immer möglich umgangen werden
- Eine vorläufige Analyse deutete darauf hin, dass Tokens, die häufig durch Blöcke laufen, mit höherer Entropie bei der Ausgabevorhersage korrelieren, was darauf hindeutet, dass es sich um schwerer vorhersagbare Tokens handeln könnte
1 Kommentare
Hacker-News-Meinungen
Komplexeres Routing wird mit Sicherheit häufiger werden.
Ich glaube insbesondere, dass es eines Tages in Richtung rekursives Routing gehen wird, bei dem eine Expertengruppe erneut durchlaufen wird. Künftig dürfte die „Gedankenkette“ (chain-of-thought) rekursiv im Inneren des Modells stattfinden.
Näher an der Idee der „Experten“ ist ein Paper aus dem vergangenen Jahr zu Sparse Universal Transformers, das Universal Transformer mit einer Sparse Mixture of Experts kombiniert, sodass ein Gating-Mechanismus entscheidet, welche Transformer-Blöcke in welcher Reihenfolge genutzt werden.
Es ist nicht mein Fachgebiet, aber nach meinem Verständnis ist es schwierig, das sauber zu trainieren, und um ähnliche Ergebnisse wie ein normaler Transformer zu erzielen, braucht man bei der Inferenz insgesamt mehr Rechenaufwand. Trotzdem ist es eine interessante Richtung, und dass es eine Obergrenze für die Zahl der Rechenschritte pro Token gibt, halte ich für einen der großen Nachteile der klassischen Transformer-Architektur.
Wenn man eine zufällige Anzahl wählt oder mehrere Rekursionstiefen ausprobiert, wird die Ausgabe „unscharf“. Man weiß dann also nicht, ob die Ausgabe einer bestimmten Schicht wichtige Informationen für das Endergebnis liefern soll oder ob sie die bestmögliche Eingabe für die nächste Rekursionsrunde sein soll.
Tieferes Schlussfolgern dürfte eher aus Dynamik auf Token-Ebene kommen als aus Dynamik auf Schichtebene. Ein Beispiel ist das aktuelle Quiet-STaR-Paper, bei dem das Modell verworfene Begründungs-Tokens ausgibt: https://arxiv.org/abs/2403.09629
Könnte das wichtigste Paper des Jahres 2024 sein.
Die Idee, Modelle zu wollen, die nicht für jedes Token dieselbe Rechenmenge aufwenden müssen, gibt es schon lange, aber zum ersten Mal sehe ich dafür einen überzeugenden Mechanismus.
Eine einfache Einführungsversion:
Stell dir einen klugen Assistenten vor, der Sprache verstehen und verarbeiten kann. Normalerweise schenkt dieser Assistent jedem Wort die gleiche Aufmerksamkeit, egal wie wichtig das jeweilige Wort für die Gesamtbedeutung ist.
Nehmen wir nun an, man hätte eine Methode gefunden, dem Assistenten beizubringen, seine „Gehirnressourcen“ intelligenter einzusetzen. Statt jedem Wort die gleiche Aufmerksamkeit zu geben, konzentriert er sich stärker auf die Wörter, die für das Verständnis der Bedeutung am wichtigsten sind, und passt diese Fokussierung je nach Kontext spontan an.
Damit der Assistent nicht überlastet wird, setzt man außerdem eine Grenze für die gesamten „Gehirnressourcen“, die er zu einem bestimmten Zeitpunkt verwenden darf. Man gibt ihm gewissermaßen ein Budget und sagt: „Du kannst nur für eine bestimmte Anzahl von Wörtern gleichzeitig Ressourcen einsetzen.“ Dann muss der Assistent entscheiden, welche Wörter am wichtigsten sind.
Trotz dieser Grenze passt der Assistent seinen Ressourceneinsatz flexibel an. Für manche Wörter verwendet er mehr, für andere weniger, und priorisiert je nach Situation.
So trainiert arbeitet der Assistent effizient und richtet seine Aufmerksamkeit intelligent aus. Er versteht genauso gut wie ein Assistent, der jedem Wort gleich viel Aufmerksamkeit schenkt, verbraucht aber insgesamt weniger Gehirnressourcen, sodass Antworten und die Verarbeitung neuer Informationen deutlich schneller werden.
Ich habe auf Basis meines Verständnisses hier eine kleine Zusammenfassung geschrieben:
https://lifeinthesingularity.com/p/googles-breakthroughs-in-...
Das ist Mixture of Experts sehr ähnlich. Nur dass Tokens nicht an mehrere Experten geroutet werden, sondern „einem einzelnen Experten zugewiesen werden, der dynamisch übersprungen werden kann“.
“This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
Ist das die Methode, die ein Kontextfenster mit 10 Millionen Tokens möglich macht? Oder geht es einfach um künftig längere Kontextfenster?
Nachdem ich versucht hatte, einige RASP-Algorithmen zu verstehen und zu implementieren, kam ich zu dem Schluss, dass bestimmte Funktionen eine gewisse Anzahl von Transformer-Schichten brauchen, um zu funktionieren.
Folgt man dieser Logik, zeigt sich, dass die von Transformern gelernten Funktionen über mehrere Heads verteilt sein können. Solche Funktionen zu wiederholen, kann extrem wertvoll sein, um ein Problem zu verstehen und zu lösen, aber die aktuelle Inferenzmethode kann zusammenhängende Gruppen von Heads nicht wiederholt ausführen. Dieses Paper wirkt tatsächlich wie eine vielversprechende Richtung.
[1] https://arxiv.org/pdf/2106.06981.pdf
[2] https://www.youtube.com/watch?v=t5LjgczaS80
Der einzige Nachteil daran, dass LLMs sich so schnell entwickeln, ist vielleicht, dass Papers schneller erscheinen, als irgendjemand, zumindest außerhalb von Google, die Verbesserungen lernen und testen kann.
Als ich mit Deep Learning angefangen habe, waren ReLU und Dropout gerade groß, und man konnte auf einer Consumer-1080 ein oder zwei Codezeilen ändern und innerhalb weniger Stunden testen, ob es eine Verbesserung bringt. Heute muss man wohl ein paar Wochen warten, bis jemand wie Mistral es ausprobiert.
Ich konzentriere mich auf Quantisierungsansätze und teste mit alten GPUs der vorherigen Generation.
hu-po reviewt AI-Paper in ausführlichen Live-Streams.
Sehr empfehlenswert; hier geht es um das diskutierte mixture-of-depths-Paper: https://www.youtube.com/watch?v=Teru_qIdB8Y