Griffin
Griffin : Kombination von Gated Linear Recurrences mit lokaler Aufmerksamkeit für effiziente Sprachmodelle
arXiv
Modellarchitektur
Alle unsere Modelle enthalten die folgenden Komponenten: (i) einen Restblock, (ii) einen MLP-Block und (iii) einen zeitlich mischenden Block. Während (i) und (ii) in allen Modellen gleich sind, betrachten wir drei zeitliche Mischblöcke: globale Multi-Query Attention (MQA), lokale (Schiebefenster) MQA und unseren vorgeschlagenen wiederkehrenden Block. Als Teil des wiederkehrenden Blocks verwenden wir die Real-Gated Linear Recurrent Unit (RG-LRU) – eine neuartige wiederkehrende Schicht, die von der linearen wiederkehrenden Einheit Orvieto et al., 2023b, inspiriert ist.
Der Restblock, wie in Abbildung 2(a) dargestellt, definiert die globale Struktur unserer Modelle und ist von Transformern vor der Norm inspiriert (Xiong et al., 2020). Nachdem wir die Eingabesequenz eingebettet haben, leiten wir sie weiter $N$ solche Blöcke ( $N$ was die Modelltiefe angibt) und dann wenden wir RMSNorm Zhang und Sennrich, 2019 an, um die endgültigen Aktivierungen zu erzeugen. Um die Token-Wahrscheinlichkeiten zu berechnen, wenden wir eine letzte lineare Ebene gefolgt von einem Softmax an. Die Gewichte dieser Ebene werden mit der Eingabeeinbettungsebene geteilt.
Restblock
Abbildung 2: a) Das Hauptrückgrat unserer Modusarchitektur ist der Restblock, der gestapelt ist $N$ mal. b) Der von uns verwendete Gated-MLP-Block. c) Der wiederkehrende Block, den wir als Alternative zu Multi Query Attention (MQA) vorschlagen. Es verwendet unsere vorgeschlagene RG-LRU-Schicht, definiert in Abschnitt 2.4.
Der Restblock enthält zwei Komponenten, die der Reihe nach angewendet werden. Die erste Komponente nimmt den verborgenen Zustand ein $chi$ und wendet eine RMSNorm Zhang und Sennrich, 2019 an, gefolgt vom zeitlichen Mischblock. Anschließend führen wir die Ausgabe mit einer Skip-Verbindung von zusammen $chi$ durch Addition. In ähnlicher Weise wendet die zweite Komponente RMSNorm an, gefolgt vom MLP-Block, und führt dann seinen Ausgang mit einer Sprungverbindung vom Eingang des RMSNorm zusammen. Dieser Block ist in Abbildung 2 (a) dargestellt.
MLP-Block
Wir verwenden einen Gated-MLP-Block Dauphin et al., 2017 (dargestellt in Abbildung 2(b)), der aus seiner Dimensionseingabe zwei Zweige erstellt $D$ . Wir wenden eine lineare Ebene mit Ausgabedimension an $MD$ auf jedem Zweig, wo $M$ bezeichnet den Expansionsfaktor. Der Einfachheit halber verwenden wir $M=3$ während dieser Arbeit. Wir wenden eine GeLU-Nichtlinearität (Hendrycks und Gimpel, 2016) auf einen der Zweige an, bevor wir sie durch elementweise Multiplikation zusammenführen, ähnlich wie bei GeGeLU Shazeer, 2020. In unserem MLP-Block wenden wir jedoch eine letzte lineare Ebene mit Ausgabedimension an $D$ auf den Ausgaben der GeGeLU-Schicht.
Zeitlich mischende Blöcke
Der zeitliche Mischblock ist die Komponente unseres Modells, die Aktivierungen verborgener Schichten an verschiedenen zeitlichen Orten in der Sequenz aggregiert. Wir betrachten drei zeitlich mischende Blöcke: den globalen MQA Shazeer, 2019, den lokalen MQA Beltagy et al., 2020 und unseren vorgeschlagenen Recurrent-Block.
Globale Aufmerksamkeit bei mehreren Abfragen
Sofern nicht anders angegeben, verwenden wir MQA anstelle von MHA, um die Inferenzgeschwindigkeiten unserer Transformer-Basislinien Shazeer, 2019 zu verbessern. Wir verwenden eine feste Kopfabmessung $D_{Kopf}=128$ , und wir legen die Anzahl der Aufmerksamkeitsköpfe fest $H$ so dass $HD_{Kopf}=D$ . Hierzu ist die Modelldimension erforderlich $D$ ein Vielfaches von 128 sein. Wir verwenden keine absoluten Positionseinbettungen, sondern verwenden Rotary Position Embedding (RoPE) Su et al., 2021 als relative Positionseinbettung.
Aufmerksamkeit für lokale Schiebefenster
Einer der Hauptnachteile der Verwendung globaler Aufmerksamkeit besteht darin, dass ihre Rechenkomplexität quadratisch mit der Sequenzlänge zunimmt. Um dieses Problem anzugehen, haben mehrere Arbeiten damit begonnen, die lokale Aufmerksamkeit zu übernehmen, Beltagy et al., 2020, auch bekannt als Sliding Window Attention. Dadurch kann sich jede Position in der Vergangenheit nur um eine feste Anzahl von Token kümmern. Dadurch werden nicht nur die rechentechnischen FLOPs reduziert, sondern auch die Größe des KV-Cache auf die Größe des Fensters beschränkt, sodass die Sequenzlänge nicht mehr quadratisch ist. Alle anderen Details entsprechen denen des globalen MQA.
Wiederkehrender Block
Unser wiederkehrender Block (Abbildung 2(c)) ähnelt dem GSS-Block Mehta et al., 2022 und dem von Mamba Gu und Dao, 2023 verwendeten Block. Wir übernehmen die Eingabe der Dimension $D$ und wenden Sie zwei lineare Ebenen mit Ausgabedimension an $D_{RNN}$ parallel, wodurch zwei Zweige entstehen. Auf dem ersten Zweig wenden wir eine kleine trennbare Conv1D-Schicht an, inspiriert vom Shift-SSM in H3 Dao et al., 2022b, mit einer zeitlichen Filterdimension von 4. Beachten Sie, dass diese Conv1D-Schicht mit gerade sehr klein ist $4D$ Parameter. Wir folgen der Conv1D-Schicht mit unserer vorgeschlagenen RG-LRU-Schicht (unten definiert). Auf dem zweiten Zweig wenden wir eine GeLU-Nichtlinearität an und führen dann die Zweige durch elementweise Multiplikation zusammen. Anschließend wenden wir eine abschließende lineare Ebene mit Ausgabedimension an $D$ .
Real-Gated Linear Recurrent Unit (RG-LRU)
Unsere vorgeschlagene RG-LRU-Schicht verfügt über eine einfache Wiederholung, die von der Linear Recurrent Unit (LRU) von Orvieto et al., 2023b, inspiriert ist, enthält jedoch einen Gating-Mechanismus, der durch die Literatur zu nichtlinearen RNNs, insbesondere LSTMs Hochreiter und Schmidhuber, 1997, motiviert ist GRUs Chung et al., 2014. Die Gleichungen, die die Schicht beschreiben, lauten wie folgt:
$$begin{align} r_t &= sigma(W_{a} x_t + b_a), & text{Recurrence Gate} \ i_t &= sigma(W_{x} x_t + b_x), & text{ Eingabegatter} \ a_t &= a^{cr_t}, & text{} \ h_t &= a_t odot h_{t-1} + sqrt{1 - a_t^2} odot (i_t odot x_t). & text{} end{align}$$
Die Ausgabe der Ebene ist $y_t=h_t$ und die Nichtlinearität $sigma$ in den Gleichungen ist die Sigmoidfunktion. Das wiederkehrende Gewicht $a$ in Gleichung (4) ist diagonal. Daher sind alle Operationen elementweise. Wir parametrisieren $a$ in Gleichung (3) als $a=sigma(Lambda)$ , Wo $Lambda$ ist ein lernbarer Parameter. Das garantiert $0 <= a <= 1$ , um sicherzustellen, dass die Wiederholung stabil ist. Die Variable $c$ ist eine skalarwertige Konstante, die auf 8 gesetzt ist. Für die numerische Stabilität berechnen wir in der Praxis $a^{cr_t}$ im Log-Space (siehe Anhang A). Die Ebene verfügt über Tore an beiden Eingängen $x$ und das wiederkehrende Gewicht $a$ . Keines der Gatter hängt jedoch vom wiederkehrenden Zustand ab $h_{t-1}$ , wodurch sichergestellt wird, dass die Berechnung effizient auf dem Gerät ausgeführt werden kann. Wir initialisieren beide $W_{a}$ Und $W_{b}$ mit LeCun init LeCun et al., 2002. Wir initialisieren $Lambda$ so dass $a^c$ gleichmäßig verteilt ist 0,9 $ Und 0,999 $ zu Beginn der Ausbildung, ähnlich wie (Orvieto et al., 2023b.). Im Gegensatz zu vielen neueren Arbeiten in der SSM-Literatur verwendet die RG-LRU keine Initialisierung, die von der Theorie orthogonaler Polynome (Gu et al., 2020) inspiriert ist, und sie ist auch nicht als Diskretisierung eines zugrunde liegenden kontinuierlichen Systems definiert (Gu et al., 2021a. Im Gegensatz zur ursprünglichen LRU-Schicht verwenden wir bei der Wiederholung keine komplexe Algebra. Während die Verwendung komplexer Wiederholungen zu einer ausdrucksstärkeren Ebene führen würde (Orvieto et al., 2023a), stellten wir fest, dass komplexe Wiederholungen für die Sprachmodellierung in der Praxis nicht vorteilhaft waren, wie auch von Gu und Dao, 2023, beobachtet wurde. (siehe Anhang B)
Torverhalten
Das Eingangstor $i_t$ ähnelt dem in LSTM, das die Eingabe filtern (oder verkleinern) kann $x_t$ . Allerdings ist unseres Wissens unser Wiederholungstor $r_t$ unterscheidet sich von anderen Gating-Mechanismen in der Literatur. Beispielsweise ist der in Mamba Gu und Dao, 2023 vorgeschlagene Auswahlmechanismus vergleichbar mit dem Update-Gate von GRUs, das interpoliert $x_t$ . Seine Auswirkung auf den verborgenen Zustand ermöglicht es ihm, seinen Zustand zurückzusetzen und alle Informationen aus der Vergangenheit zu vergessen, ähnlich dem Vergessenstor im LSTM. Im Gegensatz dazu kann unser Wiederholungsgatter annähernd zwischen dem Standard-LRU-Update von Orvieto et al., 2023a und dem vorherigen verborgenen Zustand interpolieren, was es ihm ermöglicht, die Eingabe effektiv zu verwerfen und alle Informationen aus dem vorherigen Verlauf beizubehalten (weitere Einzelheiten finden Sie in Anhang A). ). Wir glauben, dass die Schlüsselrolle dieses Gatters darin besteht, dem Modell zu ermöglichen, ein superexponentielles Gedächtnis zu erreichen, indem der Einfluss nicht informativer Eingaben reduziert wird.
