Attention Sinks

Ein Attention Sink ist ein Token, das in vielen Köpfen und Schichten einen großen Anteil der Attention-Gewichte erhält, obwohl sein Inhalt für die Anfrage irrelevant ist. Meist handelt es sich um das erste Token der Sequenz. Guangxuan Xiao und Kollegen am MIT, bei Meta AI und an der Carnegie Mellon University benannten diesen Effekt in ihrer bei ICLR 2024 veröffentlichten Arbeit Efficient Streaming Language Models with Attention Sinks. Sie zeigten, dass große Sprachmodelle (LLMs) wie Llama 2, MPT, Falcon und Pythia versagen, wenn ein gleitender KV-Cache die ersten Tokens entfernt. Werden vier dieser Tokens beibehalten, ist die Generierung wieder stabil.

Warum Attention Sinks entstehen#

Self-Attention wandelt die Scores jeder Anfrage mit einer Softmax in Gewichte um, deren Summe immer eins beträgt. Ist kein Token im Kontext für einen Kopf relevant, müssen die Gewichte trotzdem irgendwohin fließen. In einem kausalen Sprachmodell sind die ersten Tokens für jedes spätere Token sichtbar. Beim Training werden sie deshalb zum Standardziel für Gewichtung, die ein Kopf nicht benötigt. In der StreamingLLM-Analyse von Llama 2 konzentrierten alle Schichten über den beiden untersten die Attention aller Köpfe auf die Anfangstokens. Ein zusätzliches Sink-Logit in der Softmax verdeutlicht diesen Effekt:

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Ohne Sink verteilt der Kopf sein gesamtes Gewicht auf sechs irrelevante Werte und gibt eine Ausgabe mit einer Norm von 0,83 zurück. Bei einem Sink-Logit von 3 erhält der Sink 77 % des Gewichts und die Ausgabenorm sinkt auf 0,19. Bei dieser Anfrage tut der Kopf also fast nichts. Die offenen Gewichtsmodelle gpt-oss von OpenAI enthalten diese Funktion. Für jeden Attention-Kopf wird ein lernbares Sink-Logit verwendet, das einen Platz für überschüssiges Gewicht bietet, ohne die Ausgabe zu beeinflussen.

Streaming-Generierung mit StreamingLLM#

Ein Modell, das unbegrenzt Text generiert, wächst über jedes Kontextfenster hinaus und kann nicht jedes Token in seinem KV-Cache behalten. Eine einfache Fenster-Attention, die nur die jüngsten Tokens speichert, versagt, sobald die ersten Tokens herausfallen. StreamingLLM behält Schlüssel und Werte von vier anfänglichen Tokens dauerhaft neben einem rollierenden Fenster der jüngsten Tokens. Positionen werden anhand des Platzes im Cache statt im ursprünglichen Text vergeben. Bei jedem Schritt werden die rotatorischen Positions-Embeddings (RoPE) erneut angewendet:

Damit blieb die Sprachmodellierung über 4 Millionen Tokens und mehr stabil und war bis zu 22,2-mal schneller als ein gleitendes Fenster, das den Cache für jedes neue Token erneut berechnet. Das Fenster begrenzt weiterhin, woran das Modell sich erinnern kann: Tokens, die aus der Mitte entfernt wurden, sind endgültig verloren. Ein Streaming-Chatbot oder KI-Agent erinnert sich daher nur an die jüngsten Gesprächsbeiträge. Die Autoren trainierten Modelle außerdem mit einem eigenen, lernbaren Platzhalter-Token am Anfang jeder Stichprobe, das die Sink-Funktion allein übernahm.

Attention Sinks und verwandte Begriffe#

• Gleitende Fenster-Attention: Beschränkt jedes Token auf jüngere Nachbartokens. Wird sie bei der Inferenz auf ein Modell angewendet, das mit vollständiger Attention trainiert wurde, versagt sie, sobald die Sinks aus dem Fenster fallen.
• Kontextverfall: Der Genauigkeitsverlust, der eintritt, wenn eine Eingabe länger wird, aber noch in das Kontextfenster passt.

Sinks in Bildverarbeitungsmodellen#

Vision-Transformer (ViTs) zeigen ein verwandtes Verhalten. In Vision Transformers Need Registers fanden Forschende von Meta und Inria Tokens mit ungewöhnlich großen Normen auf Hintergrundbereichen mit wenig Information in überwachten und selbstüberwachten ViTs. Diese Tokens wurden für interne Berechnungen zweckentfremdet. Einige zusätzliche Register-Tokens in der Eingabe übernahmen diese Funktion, entfernten Artefakte aus Merkmalskarten und Attention-Karten und verbesserten die dichte Vorhersage.

Meta-SAM 2, das im Ultralytics-Paket ausgeführt wird, plant seinen Videospeicher nach einem ähnlichen Muster aus Anker und Fenster: Bei jedem neuen Frame greift die Memory-Attention auf Einträge der Speicherbank der vom Nutzer vorgegebenen Frames im gesamten Video sowie auf die Einträge der sechs jüngsten Frames zu. So bleiben die Aufmerksamkeitskosten pro Frame unabhängig von der Videolänge konstant.