Hungry Hungry Hippos: Hin zum Sprachmodellieren mit Zustandsraummodellen

Zustandsraummodelle (SSMs) haben in bestimmten Modaliäten Leistungen auf dem Stand der Technik für Sequenzmodellierung gezeigt, schneiden jedoch im Bereich der Sprachmodellierung unter den Aufmerksamkeitsmechanismen ab. Zudem sind SSMs trotz einer nahezu linearen Skalierung mit der Sequenzlänge – im Gegensatz zur quadratischen Skalierung von Transformers – weiterhin langsamer, da sie eine schlechte Hardware-Auslastung aufweisen. In diesem Paper machen wir Fortschritte bei der Erklärung der Ausdruckskraftlücke zwischen SSMs und Aufmerksamkeit im Bereich der Sprachmodellierung sowie bei der Verringerung der Hardware-Hürde zwischen SSMs und Aufmerksamkeit. Zunächst nutzen wir synthetische Sprachmodellierungsaufgaben, um die Lücke zwischen SSMs und Aufmerksamkeit besser zu verstehen. Wir stellen fest, dass bestehende SSMs Schwierigkeiten mit zwei Fähigkeiten haben: der Wiederaufnahme früherer Tokens in der Sequenz und dem Vergleich von Tokens innerhalb der Sequenz. Um den Einfluss dieser Schwächen auf die Sprachmodellierung besser zu verstehen, schlagen wir eine neue SSM-Schicht, H3, vor, die explizit für diese Fähigkeiten konzipiert ist. H3 erreicht bei den synthetischen Sprachen die Leistung von Aufmerksamkeit und liegt bei OpenWebText nur 0,4 PPL hinter Transformers. Darüber hinaus übertrifft ein hybrider 125-M-Parameter-H3-Aufmerksamkeits-Modell, das zwei Aufmerksamkeitsschichten beibehält, Transformers auf OpenWebText um 1,0 PPL. Um die Effizienz des Trainings von SSMs auf moderner Hardware zu verbessern, stellen wir FlashConv vor. FlashConv nutzt einen vereinten Block-FFT-Algorithmus, um die Effizienz für Sequenzen bis zu 8K zu steigern, und führt einen neuartigen Zustandsübertragungs-Algorithmus ein, der die rekurrenten Eigenschaften von SSMs ausnutzt, um auf längere Sequenzen zu skalieren. FlashConv erzielt eine Verdoppelung der Geschwindigkeit im Long-Range-Arena-Benchmark und ermöglicht es hybriden Sprachmodellen, Text 2,4-mal schneller zu generieren als Transformers. Mit FlashConv skalierten wir hybride H3-Aufmerksamkeits-Modelle bis zu 2,7 Milliarden Parametern auf dem Pile und erzielten vielversprechende erste Ergebnisse: Die Modelle erreichen eine niedrigere Perplexität als Transformers und übertrifft diese bei der Null- und Few-Shot-Lernung in der Mehrheit der Aufgaben des SuperGLUE-Benchmarks.