[Übersetzung] MoD (Mixture-of-Depths): Ein Ansatz zur Optimierung der Rechenoperationen von Transformer-basierten Sprachmodellen, und MoDE (MoD+MoE)

PyTorchKR

• MoE-Methoden (Mixture-of-Experts), bei denen wie bei Jamba, Qwen1.5-MoE oder DBRX einer oder mehrere „Experten“-Blöcke ausgewählt werden, um die Eingabe-Token zu verarbeiten, werden in letzter Zeit häufig eingesetzt. Solche MoE-Ansätze zeigen bessere Leistung, indem sie Teile großer LLMs dynamisch steuern und Rechenressourcen effizient zuweisen.

• Wenn MoE auf diese Weise die „Breite (width)“ des Modells dynamisch steuert, indem es einige Experten aus mehreren auswählt, verfolgt das hier vorgestellte MoD (Mixture-of-Depths) einen Ansatz, der die „Tiefe (depth)“ des Modells dynamisch steuert. Das heißt: Für bestimmte Token werden nicht alle Layer des Modells berechnet, sondern einige Layer werden übersprungen. Schauen wir es uns gemeinsam an.

---

MoD (Mixture-of-Depths): Ein Ansatz zur Optimierung der Rechenoperationen von Transformer-basierten Sprachmodellen (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

Einführung

Traditionell weisen Transformer-Modelle allen Teilen einer Eingabesequenz die gleichen Rechenressourcen zu. Die hier vorgestellte MoD-Transformer-Arbeit (Mixture-of-Depths) versucht jedoch, Rechenressourcen dynamisch bestimmten Sequenzpositionen im Transformer zuzuweisen und diese Zuweisung über verschiedene Layer hinweg zu verändern. Um diese Mixture-of-Depths-(MoD)-Strategie zu nutzen, verwenden die Autoren einen Top-k-Routing-Mechanismus, der entscheidet, welche Token verarbeitet werden, und so den gesamten Rechenaufwand steuert. Diese Technik erreicht eine Leistung, die mit bestehenden Modellen vergleichbar oder besser ist, benötigt dabei jedoch deutlich weniger FLOPs und bietet beim Sampling nach dem Training eine Geschwindigkeitssteigerung von bis zu 50 %.

Transformer-Architektur, vorgestellt im Paper Attention is All You Need

Aktuelle Transformer-Modelle und LLMs (Large Language Models) wachsen weiterhin in Größe und Komplexität. Diese Modelle nutzen Milliarden von Parametern, um sich menschlichem Sprachverständnis und menschlicher Sprachgenerierung anzunähern, und zeigen dadurch beeindruckende Leistungen. Allerdings gehen diese Fortschritte mit erheblichen Rechenkosten und Energieverbrauch einher, und die Vergrößerung der Modelle erhöht Trainings- und Inferenzzeiten drastisch.

Das liegt daran, dass Transformer für alle Token einer Eingabesequenz die gleiche Menge an Rechenoperationen ausführen. Der Ansatz geht jedoch von der Idee aus, dass nicht alle Token gleich viele Informationen liefern: Manche spielen eine Schlüsselrolle für das Kontextverständnis, andere sind vergleichsweise weniger wichtig.

Der MoD Transformer weist jedem Token dynamisch Rechenaufwand zu, sodass das Modell seine Ressourcen dort konzentrieren kann, wo sie benötigt werden. Dadurch lassen sich die gesamten Rechenkosten senken und zugleich die Effizienz der Informationsverarbeitung erhöhen, indem wichtigen Token mehr Aufmerksamkeit geschenkt wird. Besonders bei großen Modellen wie LLMs kann der MoD-Ansatz ein wichtiger Weg sein, Zeit und Ressourcen sowohl beim Training als auch bei der Inferenz zu sparen.

Einführung in die MoD-(Mixture-of-Depths)-Technik

Die MoD-Technik integriert in jeden Transformer-Layer einen Top-k-Routing-Mechanismus, der dynamisch auswählt, welche Token verarbeitet werden. Dadurch durchläuft nur ein ausgewählter Teil der Token den vollständigen Rechenprozess, während andere Token Layer überspringen.

Abbildung 1 / links: Überblick über den MoD-(Mixture-of-Depths)-Transformer: Ähnlich wie bei der MoE-(Mixture-of-Experts)-Struktur wird über einen Router der Rechenpfad gewählt. Anders als bei MoE wird hier entschieden, ob der Standardblock (Self-Attention und MLP) berechnet wird oder nicht. Falls ein Eingabe-Token ($X_i$) gemäß der Router-Entscheidung nicht die vollständige Berechnung aktiviert, durchläuft es die Berechnung dieses Layers nicht.
Abbildung 1 / rechts: Ergebnis der Routing-Entscheidungen für eine kurze Sequenz mit 64 Token. (X-Achse: Sequenz, Y-Achse: Layer) Lila bedeutet, dass die Berechnung dieses Layers vollständig durchgeführt wurde, Orange bedeutet, dass die Berechnung dieses Layers übersprungen wurde.

Wie die obige Abbildung zeigt, implementiert MoD einen prädiktiven Router innerhalb der Transformer-Architektur, der für jedes Token entscheidet, ob die Berechnung eines Layers ausgeführt wird. Dadurch können Rechenressourcen auf relativ „wichtige“ Token konzentriert werden.

Training des MoD Transformers

Der Trainingsprozess des MoD Transformers ähnelt grundsätzlich dem eines Standard-Transformers. Die Eingabesequenz wird dem Modell zugeführt und beim Durchlaufen der einzelnen Layer werden die erforderlichen Berechnungen ausgeführt. Der MoD Transformer enthält jedoch einen zusätzlichen Schritt, bei dem der Router entscheidet, welchen Rechenpfad jedes Token nehmen soll.

1. Berechnung der Router-Gewichte: Der Router berechnet in jedem Layer des Modells für jedes Token ein Gewicht. Dieses Gewicht repräsentiert die Menge an Berechnung, die das Token erhalten soll, und basiert auf der Wichtigkeit des Tokens.

2. Auswahl der obersten Token: Auf Basis der berechneten Gewichte wählt der Router für jeden Layer die obersten Token aus. Diese Token werden auf einen Pfad geleitet, auf dem die vollständige Berechnung ausgeführt wird.

3. Zuweisung des Rechenpfads: Die ausgewählten Token werden einem Pfad zugewiesen, der die Standardberechnung (Self-Attention und MLP) ausführt. Die übrigen Token werden über Residual-Verbindungen weitergegeben und gelangen ohne zusätzliche Berechnung in den nächsten Layer.

4. Loss-Berechnung und Backpropagation: Die Ausgabe des Modells wird mit der finalen Loss-Funktion bewertet, und die Gradienten werden durch das Modell zurückpropagiert, um die Gewichte zu aktualisieren. Dabei wird auch der Router trainiert, sodass er im Laufe der Zeit effizientere Routing-Entscheidungen treffen kann.

Routing-Schemata des MoD Transformers

Die Routing-Implementierung ist ein zentraler Bestandteil des MoD Transformers. Die Autoren verglichen drei Routing-Varianten und stellten fest, dass Token-Routing mit Expert-Choice-MoD die beste Leistung zeigt. Werfen wir einen Blick auf den Vergleich mit den anderen vorgestellten Routing-Ansätzen:

Token-Choice-Routing

Beim Token-Choice-Routing wählt jedes Token selbst den Pfad, den es durchläuft. Das Modell erzeugt für jedes Token eine Wahrscheinlichkeitsverteilung über verschiedene Rechenpfade (z. B. unterschiedliche Layer oder Verarbeitungsblöcke) und leitet das Token entsprechend dieser Verteilung an den jeweiligen Pfad weiter.

Token-Choice-Routing ist sehr flexibel, weil jedes Token den für sich am besten geeigneten Rechenpfad wählen kann. Außerdem kann das Modell die Wichtigkeit oder den Kontext jedes Tokens feiner berücksichtigen und so auf einzelne Token optimieren.

Expert-Choice-Routing

Beim Expert-Choice-Routing wählt das Modell direkt aus, welche Token von welchem Pfad (oder „Experten“) verarbeitet werden. Dies geschieht mithilfe eines Top-k-Mechanismus, der auf Basis der Router-Gewichte der Token die obersten Token auswählt.

Durch dieses Expert-Choice-Routing können die Token gleichmäßiger auf die einzelnen Pfade verteilt werden, sodass die Rechenlast ausgewogener verteilt und die Rechenressourcen des Modells effizienter genutzt werden. Außerdem sind die obersten Token klar bestimmt, was die Vorhersagbarkeit und Stabilität des Modells verbessert.

Expert-Choice MoD

Expert-Choice MoD überträgt das Konzept des Expert-Choice-Routings auf die Struktur des MoD Transformers. Dabei wählt das Modell für jeden Layer oder Verarbeitungsschritt die obersten Token aus, und nur diese ausgewählten Token werden entlang eines Pfads verarbeitet, auf dem Berechnungen stattfinden. Die übrigen Token folgen über Residual-Verbindungen einem kostengünstigeren Pfad.

Da bei diesem Ansatz nur ausgewählte Token berechnet werden, sinkt der gesamte FLOP-Verbrauch des Modells deutlich. Gleichzeitig kann die Leistung erhalten oder verbessert werden, da für wichtige Token die Berechnung über vollständige Layer hinweg sichergestellt wird, während unnötige Berechnungen minimiert werden. Schließlich erlaubt das dynamische Routing über den Router, Rechenressourcen flexibel entsprechend den Eigenschaften der Eingabesequenz zuzuweisen.

Implementierung des Routings

Die Implementierung des Routings ist in den folgenden drei Hauptschritten möglich:

1. Token-Embeddings und Bestimmung der Router-Gewichte: Für jeden Layer des Modells gibt es eine Menge von Token-Embeddings $X^i$ der gegebenen Sequenz. Für jedes Token-Embedding $x^l_i$ berechnet der Router über eine lineare Transformation ein skalares Gewicht $r^l_i$, das ausdrückt, wie wahrscheinlich es ist, dass dieses Token an der Berechnung teilnimmt.

2. Bestimmung der obersten Gewichte: Aus der Menge der Router-Gewichte $R^l$ der Token wird das Perzentil $P_\beta(R^l)$ des $\beta$-ten Gewichts ermittelt. Dabei bestimmt $\beta$, welcher Anteil der Token vom Modell in jedem Layer verarbeitet wird. Token, deren Router-Gewicht größer als $P_\beta(R^l)$ ist, nehmen daher an der Berechnung teil.

3. Auswahl des Rechenpfads: Jedes Token wird mit seinem Router-Gewicht $r^l_i$ mit $P_\beta(R^l)$ verglichen. Ist es größer oder gleich, wird die Berechnung dieses Layers ausgeführt; ist es kleiner, wird das Token direkt an den nächsten Layer weitergegeben. Dadurch wird die Berechnung nur für ausgewählte Token durchgeführt, während die übrigen einen Pfad ohne Rechenkosten nehmen können.

Bei der Implementierung dieses Routings sollten insbesondere die folgenden Punkte berücksichtigt werden:

• Gewichtsbasierte dynamische Entscheidung: Der Routing-Prozess wird dynamisch auf Basis der Gewichte jedes Tokens entschieden, was die Wichtigkeit des Tokens widerspiegelt. Wichtige Token erhalten mehr Berechnung zur Optimierung der Modellleistung, während weniger wichtige Token über den Residual-Pfad Kosten sparen.

• Recheneffizienz und Leistungsoptimierung: Über diesen Routing-Mechanismus kann der MoD Transformer Rechenressourcen effizient nutzen und dabei die Modellleistung erhalten oder verbessern. Insbesondere kann der gesamte FLOP-Verbrauch reduziert werden, indem in rechenintensiven Layern nur wichtige Token verarbeitet werden.

• Einfachheit und Effizienz der Implementierung: Das Routing lässt sich über lineare Transformationen und Perzentil-Berechnung umsetzen und ist daher relativ einfach und effizient. Dadurch wird eine dynamische Zuweisung von Rechenressourcen ermöglicht, ohne die Rechenlast beim Training oder bei der Inferenz wesentlich zu erhöhen.

Beispiel für die Funktionsweise des MoD Transformers

Anhand des bisher Gesehenen erklären wir die Funktionsweise des MoD Transformers mit einem Beispiel.

Nehmen wir an, die Eingabesequenz enthält 100 Token und der Router hat auf 10 gesetzt. In diesem Fall wählt der Router in jedem Layer die 10 wichtigsten Token aus, und nur diese Token werden entlang des Standard-Rechenpfads verarbeitet.

Die übrigen 90 Token überspringen die Berechnung und werden über Residual-Verbindungen direkt an den nächsten Layer weitergegeben. Auf diese Weise reduziert der MoD Transformer den gesamten FLOP-Verbrauch und stellt dennoch sicher, dass wichtige Token die notwendige Berechnung erhalten.

Leistungsvergleich des MoD Transformers

Trainingsleistung

Hyperparameter-Optimierung: In Experimenten zum Training einer MoD-Transformer-Variante, bei der Routing-Blöcke mit 12,5 % Kapazität abwechselnd mit vollständigen Attention-Blöcken angeordnet wurden, wurden optimale Hyperparameter-Einstellungen gefunden. Diese Konfiguration lieferte in der isoFLOP-Analyse bessere Ergebnisse als die optimale Baseline.

• Recheneffizienz: Eine bestimmte MoD-Variante erreichte dieselbe Leistung wie bestehende Modelle mit weniger FLOPs und zeigte dabei eine um 66 % höhere Verarbeitungsgeschwindigkeit.

Routing-Analyse

Sparsamkeit der Routing-Entscheidungen: Durch die Verwendung von Routing mit 12,5 % Kapazität umgingen die meisten Token den Block, und nur wenige wichtige Token passierten ihn. Dies ließ sich auch an der Verteilung der Router-Gewichte erkennen; das Netzwerk bevorzugte beim Routing je nach Tiefe teilweise bestimmte Token.

Inferenzleistung

Autoregressive Inferenz: Beim Wechsel von einem nicht-kausalen Top-k-Routing-Schema im Training zu einem kausalen, vorhersagebasierten Ansatz in der Inferenz gab es nahezu keinen Leistungsverlust. Das scheint darauf zurückzuführen zu sein, dass dieses Vorhersageproblem leicht zu erlernen ist und bereits früh im Training eine Genauigkeit von über 97 % erreicht wurde.

Mixture-of-Depths-and-Experts (MoDE) = MoD + MoE

MoDE (Mixture-of-Depths-and-Experts) bezeichnet die Integration der MoD-(Mixture-of-Depths)-Technik und der MoE-(Mixture of Experts)-Technik. Ziel dieses Ansatzes ist es, die Vorteile beider Methoden zu kombinieren und sowohl Recheneffizienz als auch Leistung von Transformer-Modellen gleichzeitig zu optimieren. Während der MoD-Ansatz auswählt, ob für ein Token tiefer oder weniger tief gerechnet wird, verarbeitet der MoE-Ansatz ein Token über den am besten geeigneten „Experten“ aus mehreren verfügbaren Experten. MoDE integriert diese beiden Ansätze, sodass Token bestimmte Blöcke überspringen oder selektiv tiefere Berechnungen ausführen können und zugleich aus mehreren Experten die optimale Verarbeitungsweise gewählt werden kann.

Implementierung von MoDE

MoDE hat zwei wesentliche Implementierungsformen: stufenweises (Staged) MoDE und integriertes (Integrated) MoDE.

1. Staged MoDE (stufenweises MoDE): Beim stufenweisen MoDE wird zuerst die MoD-Technik angewendet, um den Rechenpfad jedes Tokens zu bestimmen. Danach wird die MoE-Technik angewendet, um die Token, die auf den ausgewählten Pfaden verarbeitet werden, auf die jeweiligen Expertenblöcke zu routen. So kann das Modell zunächst die Rechentiefe steuern und anschließend in jedem Schritt den optimalen Experten auswählen.

2. Integrated MoDE (integriertes MoDE): Beim integrierten MoDE werden MoD und MoE in einem einzigen Routing-Schritt zusammengeführt. Das heißt, beim Routing eines Tokens werden in einer einzigen Entscheidung gleichzeitig die Tiefe des Rechenpfads und der zu verwendende Experte ausgewählt. Dieser Ansatz kann einfacher und effizienter in der Implementierung sein und bietet den Vorteil, für jedes Token gleichzeitig den passendsten Rechenpfad und Experten zu bestimmen.

Vorteile von MoDE

• Recheneffizienz: MoDE steuert den gesamten Rechenaufwand des Modells effektiv, indem für jedes Token dynamisch die benötigte Rechenmenge und der zu verwendende Experte bestimmt werden.
• Leistungsoptimierung: Indem jedes Token dem am besten geeigneten Rechenpfad und Experten zugewiesen wird, maximiert MoDE die Modellleistung. Das kann besonders nützlich sein, wenn komplexe Probleme gelöst oder unterschiedliche Arten von Daten verarbeitet werden sollen.
• Flexibilität und Skalierbarkeit: MoDE bietet Flexibilität für den Einsatz in verschiedenen Modellarchitekturen und Tasks und liefert ein skalierbares Framework, das Recheneffizienz und Leistung gleichzeitig verbessern kann.

MoDE zeigt eine mögliche Richtung für die zukünftige Entwicklung von Transformer-Modellen auf und leistet einen wichtigen Beitrag zur Forschung, die Wege sucht, die Modellleistung bei minimalen Rechenkosten zu maximieren.

Weiterführende Lektüre

https://arxiv.org/abs/2404.02258

---

Dieser Text basiert auf einer mit einem GPT-Modell zusammengefassten Darstellung. Daher kann es sein, dass Inhalte anders als im Original oder nicht vollständig im Sinne des Originals wiedergegeben wurden. Wenn Sie das Thema interessiert, lesen Sie bitte auch den Originaltext. Falls Ihnen beim Lesen unnatürliche oder fehlerhafte Stellen auffallen, geben Sie bitte in einem Kommentar Bescheid.

⚠️Werbung⚠️: Fanden Sie diesen von der PyTorch Korea User Group zusammengestellten Beitrag nützlich? Wenn Sie Mitglied werden, erhalten Sie wichtige Beiträge per E-Mail! Standardmäßig wöchentlich, aber eine Umstellung auf täglich ist ebenfalls möglich.