Das ist etwas, was jeder, der viel Zeit mit Transformatoren und Selbstaufmerksamkeit verbracht hat, schon hundertmal gehört hat. Beides ist absolut wahr, wir haben das alle schon erlebt: Wenn man versucht, die Kontextgröße seines Modells zu erhöhen, kommt plötzlich alles zum Stillstand. Doch gleichzeitig gibt es, so scheint es, praktisch jede Woche ein neues hochmodernes Modell mit einer neuen rekordverdächtigen Kontextlänge. (Gemini hat eine Kontextlänge von 2 Millionen Token!)
Es gibt viele ausgefeilte Methoden wie RingAttention, die das Coaching unglaublich langer Kontextlängen in großen verteilten Systemen ermöglichen, aber was mich heute interessiert, ist eine einfachere Frage.
Wie weit kommen wir allein mit linearer Aufmerksamkeit?
Dies wird eine Artwork Blitztour, aber haben Sie etwas Geduld, denn wir werden ein paar wichtige Punkte ansprechen, bevor wir uns mit den Ergebnissen befassen.
Wir können den traditionellen Aufmerksamkeitsmechanismus im Wesentlichen mit zwei Kernpunkten zusammenfassen:
- Zunächst nimmt der typische Softmax-Aufmerksamkeitsausdruck das Produkt aus Abfrage- und Schlüsselmatrizen, normalisiert es zur Gewährleistung der Stabilität und nimmt dann das Softmax (zeilenweise), um die Aufmerksamkeitswerte zwischen den einzelnen Elementen der Sequenz zu erhalten.
- Zweitens wird die Zeitkomplexität von den N² Skalarprodukten dominiert, und das Produkt innerhalb des Softmax ist der limitierende Faktor. Dort berechnen wir die Aufmerksamkeitswerte.
Dies wird in der traditionellen Type wie folgt ausgedrückt:
Es stellt sich heraus, dass wir etwas anders darüber nachdenken können, wenn wir unsere Mathematikerfreunde fragen. Man kann sich den Softmax als eine von vielen Möglichkeiten vorstellen, die Wahrscheinlichkeitsverteilung zu beschreiben, die Token miteinander in Beziehung setzt. Wir können jedes beliebige Ähnlichkeitsmaß verwenden (das Skalarprodukt ist eines der einfachsten), und solange wir es normalisieren, ist alles in Ordnung.
Es ist ein bisschen schlampig, das zu sagen Ist Aufmerksamkeit, denn tatsächlich ist es nur die Aufmerksamkeit, die wir kennen und lieben, wenn die Ähnlichkeitsfunktion die Exponentialfunktion des Skalarprodukts von Abfragen und Schlüsseln (siehe unten) ist, wie wir es im Softmax finden. Aber hier wird es interessant: Was wäre, wenn wir statt dieses Ausdrucks eine Näherung vornehmen könnten?
Wir können davon ausgehen, dass es eine Characteristic-Map gibt.phi” was uns ein Ergebnis gibt quick das gleiche wie die Exponentialfunktion des Skalarprodukts. Und entscheidend ist, dass wir den Ausdruck auf diese Weise schreiben können, um mit der Reihenfolge der Matrixmultiplikationsoperationen zu spielen.
Im Papier Sie schlagen die Exponential Linear Unit (ELU) aufgrund einer Reihe nützlicher Eigenschaften als Merkmalskarte vor:
- Für Werte über 0 liefert ELU(x) ein lineares Ergebnis, das zwar nicht mit der Exponentialfunktion identisch ist, aber die relative Reihenfolge zwischen den Wertungen beibehält.
- Bei Werten kleiner oder gleich 0 bleibt der Exponentialterm der kontinuierliche Charakter der Funktion erhalten und stellt sicher, dass die Gradienten nicht einfach verschwinden.
Wir werden hier nicht allzu viel Zeit darauf verwenden, aber es ist empirisch ziemlich intestine bestätigt, dass dies eine gute Annäherung an die Softmax-Funktion ist.
Dadurch können wir die Reihenfolge der Operationen ändern. Wir können zuerst das Produkt unserer Characteristic-Map von Okay mit V nehmen, um einen KV-Block zu erstellen, und dann das Produkt mit Q. Das quadratische Produkt wird über die Modelldimensionsgröße und nicht über die Sequenzlänge bestimmt.
Wenn wir dies alles in den linearen Aufmerksamkeitsausdruck einfügen, erhalten wir:
Wobei wir die Begriffe in den Klammern nur einmal professional Abfragezeile berechnen müssen.
(Wenn Sie wissen möchten, wie die gelegentliche Maskierung hier hineinpasst und wie die Farbverläufe berechnet werden, werfen Sie einen Blick in das Dokument. Oder schauen Sie in diesem Bereich nach einem zukünftigen Weblog.)
Die mathematischen Argumente sind überzeugend, aber persönlich bin ich immer etwas misstrauisch, bis ich einige Benchmarks gesehen habe.
Sehen wir uns zunächst die Codeausschnitte an, die jeden dieser Begriffe beschreiben. Die Softmax-Aufmerksamkeit wird Ihnen sehr vertraut vorkommen, wir machen hier nichts Ausgefallenes.
class TraditionalAttention(nn.Module):
def __init__(self, d_k):
tremendous(TraditionalAttention, self).__init__()
self.d_k = d_kdef ahead(self, Q, Okay, V):
Z = torch.sqrt(torch.tensor(self.d_k, machine=Q.machine, dtype=torch.float32))
scores = torch.matmul(Q, Okay.transpose(-2, -1)) / Z
attention_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attention_weights, V)
return output
Dann beginnen wir für die lineare Aufmerksamkeit mit dem Abrufen der Abfrage-, Schlüssel- und Wertematrizen und wenden dann die ELU(x)-Funktionszuordnung auf die Abfrage und die Schlüssel an. Dann verwenden wir die Einsummennotation, um die Multiplikationen durchzuführen.
class LinearAttention(nn.Module):
def __init__(self):
tremendous(LinearAttention, self).__init__()
self.eps = 1e-6def elu_feature_map(self, x):
return F.elu(x) + 1
def ahead(self, Q, Okay, V):
Q = self.elu_feature_map(Q)
Okay = self.elu_feature_map(Okay)
KV = torch.einsum("nsd,nsd->ns", Okay, V)
# Compute the normalizer
Z = 1/(torch.einsum("nld,nd->nl", Q, Okay.sum(dim=1))+self.eps)
# Lastly compute and return the brand new values
V = torch.einsum("nld,ns,nl->nd", Q, KV, Z)
return V.contiguous()
Dies in Code geschrieben zu sehen ist schön und intestine, aber was bedeutet es tatsächlich experimentell? Von welcher Leistungssteigerung sprechen wir hier? Es kann schwierig sein, den Grad der Beschleunigung zu erfassen, wenn man von einem quadratischen zu einem linearen Engpass wechselt, additionally habe ich das folgende Experiment durchgeführt.
Wir werden eine einzelne Aufmerksamkeitsebene mit einer festen d_k-Modelldimension von 64 verwenden und die Zeit für einen Vorwärtsdurchlauf eines Sequenzsatzes mit 32 Batchgrößen als Benchmark verwenden. Die einzige zu ändernde Variable ist die Sequenzlänge, die zwischen 128 und 6000 liegt (die GPT-3-Kontextlänge als Referenz beträgt 2048). Jeder Durchlauf wird 100 Mal durchgeführt, um einen Mittelwert und eine Standardabweichung zu erhalten, und die Experimente werden mit einer Nvidia T4 GPU durchgeführt.
Für ein so einfaches Experiment sind die Ergebnisse ziemlich bemerkenswert.
Die Ergebnisse zeigen, dass wir sogar bei einem unglaublich kleinen Spielzeugbeispiel eine bis zu 60-fache Beschleunigung erreichen.
Diskussion
Hieraus lassen sich einige offensichtliche Erkenntnisse gewinnen:
- Der Vorteil der linearen Aufmerksamkeit ist enorm – entweder in Bezug auf die Geschwindigkeit, ein höherer Durchsatz ist immer eine gute Sache. Oder in Bezug auf den Speicherbedarf zur Verarbeitung langer Sequenzen. In Umgebungen mit wenig Speicher kann dies ein großer Vorteil sein.
- Das Verhältnisdiagramm weist einen überraschenden Knick auf – was uns vermuten lässt, dass hier eine zusätzliche Optimierung auf niedrigerer Ebene stattfindet, was bedeutet, dass das erwartete Verhältnis nicht ganz eintritt. Wir müssen dieses Ergebnis additionally mit Vorsicht genießen.
Der Vollständigkeit halber sei hier nicht verwechselt mit „lineare Aufmerksamkeit ist bei kleinen Modellen 60-mal schneller“. In Wirklichkeit machen die Feedforward-Schichten oft einen größeren Teil der Parameter in einem Transformer aus, und die Kodierung/Dekodierung ist oft auch eine Komponente mit limitierender Größe. Aber bei diesem eng definierten Downside ist das ziemlich beeindruckend!