Forscher von Google Deepmind zeigen "Mixture-of-Depths", eine Methode, um die Rechenleistung von Transformer-Modellen effizienter zu nutzen.
Traditionelle Transformer-Modelle verwenden für jedes Token in einer Sequenz die gleiche Menge an Rechenleistung. Im Gegensatz dazu erlaubt Google Deepminds "Mixture-of-Depths" (MoD) dem Modell, diese Rechenleistung flexibel und gezielt auf die Token zu verteilen, die sie am meisten benötigen.
Dazu wird eine feste Obergrenze für die Rechenleistung pro Durchlauf festgelegt, z. B. dürfen maximal 50 % der Token die rechenintensiven Berechnungen durchlaufen. In jedem Block gibt es einen "Router", der für jeden Token einen Gewichtungswert berechnet. Token mit hohem Gewicht werden für die Berechnungen ausgewählt, Token mit niedrigem Gewicht werden übersprungen. Während des Trainings wird eine Anzahl Token mit den höchsten Gewichten für die Berechnungen ausgewählt.
Die restlichen Token werden einfach unverändert weitergereicht. Auf diese Weise können rechenintensive Schritte für nicht benötigte Token übersprungen werden. Das Modell lernt selbstständig, welche Token mehr und welche weniger Rechenleistung benötigen.
Mixture-of-Depths erreicht Leistung der Basismodelle
Trotz der signifikanten Reduktion der benötigten FLOPs pro Vorhersage konnten die MoD-Modelle nach dem Training die Leistung von Baseline-Modellen beibehalten oder sogar übertreffen. Letzteres deutet nach Ansicht des Teams darauf hin, dass die traditionelle gleichmäßige Verteilung der Rechenressourcen in Transformer-Modellen nicht immer optimal ist und dass eine gezieltere Verteilung der Berechnungen die Modellleistung steigern kann. Eine Anfrage an das fertig trainierte KI-Modell benötigt nur noch einen Bruchteil der Rechenleistung und kann bis zu 50 Prozent schneller sein.
Die Methode lässt sich zudem mit der inzwischen weit verbreiteten Mixture-of-Experts-Architektur kombinieren. MoD kann als Ergänzung zu MoE betrachtet werden, da beide Ansätze auf die Optimierung verschiedener Dimensionen des Modells abzielen: MoE auf die Breite (Expertenvielfalt) und MoD auf die Tiefe (Verarbeitungstiefe).
Die Möglichkeit, Rechenleistung dynamisch zuzuweisen und damit effizienter zu nutzen, könnte insbesondere in Anwendungsbereichen mit hohen Anforderungen an Rechenzeit und -ressourcen von großem Wert sein - die zusätzlich verfügbaren FLOPs könnten aber auch für das Training größerer Modelle genutzt werden. Für letztere ergeben sich nach Angaben der Forscher auch Speichereinsparungen. Einige MoD-Varianten benötigen weniger Beschleuniger, was darauf hindeutet, dass diese Effizienz bei der Skalierung auf größere Modelle von Bedeutung sein könnte.