Die jüngsten Forschungsergebnisse des Mamba-Teams sind bemerkenswert: Forscher von Universitäten wie Cornell und Princeton haben erfolgreich das große Transformer-Modell Llama zu Mamba „destilliert“ und einen neuartigen Inferenz-Decodierungsalgorithmus entwickelt, der die Inferenzgeschwindigkeit des Modells deutlich verbessert.
Das Ziel der Forscher war es, Llama in Mamba zu verwandeln. Warum? Weil die Entwicklung eines großen Modells von Grund auf teuer ist, und obwohl Mamba seit seiner Einführung große Aufmerksamkeit erregt hat, trainieren nur wenige Teams selbst große Mamba-Modelle. Es gibt zwar einige bekannte Varianten wie Jamba von AI21 und Hybrid Mamba2 von NVIDIA, aber viele erfolgreiche Transformer-Modelle enthalten ein umfangreiches Wissen. Wenn wir dieses Wissen nutzen und gleichzeitig den Transformer zu Mamba feinabstimmen können, ist das Problem gelöst.
Das Forschungsteam kombinierte verschiedene Methoden wie progressive Destillation, überwachtes Feintuning und zielgerichtete Präferenzoptimierung, um dieses Ziel zu erreichen. Dabei war neben der Leistung auch die Geschwindigkeit entscheidend. Mamba hat klare Vorteile bei der Inferenz langer Sequenzen, und auch für Transformer gibt es Inferenzbeschleunigungsmethoden wie die vorausschauende Decodierung. Da die einzigartige Struktur von Mamba diese Methoden nicht direkt unterstützt, entwickelten die Forscher einen völlig neuen Algorithmus und nutzten Hardware-Eigenschaften für eine Mamba-basierte vorausschauende Decodierung.
Schließlich wurden Zephyr-7B und Llama-38B erfolgreich in lineare RNN-Modelle umgewandelt, wobei die Leistung mit den Standardmodellen vor der Destillation vergleichbar war. Der gesamte Trainingsprozess benötigte nur 20 Milliarden Token und lieferte Ergebnisse, die mit dem von Grund auf trainierten Mamba7B-Modell (1,2 Billionen Token) und dem NVIDIA Hybrid Mamba2-Modell (3,5 Billionen Token) vergleichbar sind.
In Bezug auf die technischen Details sind lineare RNNs und lineare Aufmerksamkeit verwandt, daher konnten die Forscher die Projektionsmatrizen im Aufmerksamkeitsmechanismus direkt wiederverwenden und das Modell durch Parameterinitialisierung erstellen. Darüber hinaus hat das Forschungsteam die Parameter der MLP-Schichten im Transformer eingefroren und die Attention-Heads schrittweise durch lineare RNN-Schichten (d. h. Mamba) ersetzt. Die gruppenweise Abfrage-Aufmerksamkeit mit gemeinsam genutzten Schlüsseln und Werten wurde ebenfalls bearbeitet.
Während des Destillationsprozesses wurde eine schrittweise Ersetzung der Attention-Schichten verwendet. Das überwachte Feintuning umfasste zwei Hauptmethoden: eine basiert auf der KL-Divergenz auf Wortebene, die andere auf der Wissensdestillation auf Sequenzebene. Für die Optimierung der Benutzerpräferenzen verwendete das Team die Methode der direkten Präferenzoptimierung (DPO), wobei der Vergleich mit der Ausgabe des Lehrermodells gewährleistet, dass das Modell bei der Generierung von Inhalten die Erwartungen der Benutzer besser erfüllt.
Als Nächstes begannen die Forscher, die vorausschauende Decodierung des Transformers auf das Mamba-Modell anzuwenden. Vorausschauende Decodierung kann einfach als die Generierung mehrerer Ausgaben durch ein kleines Modell und die anschließende Überprüfung dieser Ausgaben durch ein großes Modell verstanden werden. Das kleine Modell läuft schnell und kann schnell mehrere Ausgabevektoren generieren, während das große Modell die Genauigkeit dieser Ausgaben bewertet und so die Gesamtgeschwindigkeit der Inferenz erhöht.
Um diesen Prozess zu realisieren, entwickelten die Forscher einen Algorithmus, bei dem das kleine Modell jeweils K Roh-Ausgaben generiert, und das große Modell gibt dann die endgültige Ausgabe und den Cache des Zwischenzustands zurück. Diese Methode funktionierte gut auf der GPU, wobei Mamba2.8B eine 1,5-fache Beschleunigung der Inferenz erreichte und eine Akzeptanzrate von 60 % erzielte. Obwohl die Ergebnisse auf verschiedenen GPU-Architekturen unterschiedlich waren, optimierte das Forschungsteam die Ergebnisse durch die Fusion von Kernen und die Anpassung der Implementierung weiter, um letztendlich die gewünschte Beschleunigung zu erreichen.
In der Experimentphase verwendeten die Forscher Zephyr-7B und Llama-3Instruct8B für ein dreistufiges Destillationstraining. Letztendlich konnten sie die Forschungsergebnisse innerhalb von 3 bis 4 Tagen auf 8 Karten mit 80 GB A100 reproduzieren. Diese Forschung zeigt nicht nur den Weg von Transformer zu Mamba, sondern bietet auch neue Ansätze zur Verbesserung der Inferenzgeschwindigkeit und Leistung zukünftiger Modelle.
论文地址:https://arxiv.org/pdf/2408.15237
