Sliding Window Attention

Aufmerksamkeit mit gleitendem Fenster (SWA) beschränkt jedes Token auf die Aufmerksamkeit für eine feste Anzahl benachbarter Tokens – in einem kausalen Sprachmodell auf die w jüngsten –, statt wie bei vollständiger Selbstaufmerksamkeit die gesamte Sequenz einzubeziehen. Der Rechenaufwand für Aufmerksamkeit wächst dadurch linear statt quadratisch mit der Sequenzlänge. Ein generierendes Modell muss in seinem KV-Cache nur Schlüssel und Werte der letzten w Tokens behalten. Longformer machte das gleitende Fenster 2020 zum Kern seiner Aufmerksamkeitsmechanik. Mistral AIs Mistral 7B brachte es 2023 in ein weithin verwendetes offenes großes Sprachmodell (LLM), mit einem Fenster von 4.096 Tokens.

Wie das Fenster mit der Tiefe wächst#

Eine SWA-Schicht betrachtet nur w Tokens, doch Transformer-Schichten werden gestapelt. In Schicht 2 liest ein Token verborgene Zustände, die bereits jeweils w Tokens aus Schicht 1 zusammenfassen. So können Informationen mit jeder Schicht etwa w − 1 weitere Positionen zurücklegen, ähnlich wie gestapelte Faltungen das rezeptive Feld eines Faltungsneuronalen Netzes (CNN) vergrößern. In PyTorch wird das Fenster durch eine boolesche Aufmerksamkeitsmaske dargestellt. Durch Multiplikation der Maske mit sich selbst lässt sich verfolgen, wie weit Informationen gelangen:

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

Für 16 Tokens berechnet das Fenster pro Schicht 58 statt 136 Werte. Die Reichweite des letzten Tokens nimmt pro Schicht um 3 Tokens zu: von 4 auf 7, 10 und 13. Dieselbe Maske kann PyTorchs scaled_dot_product_attention als attn_mask übergeben werden. Bei Mistral 7B ergeben 32 Schichten mit einem Fenster von 4.096 Tokens eine theoretische Aufmerksamkeitsspanne von etwa 131.000 Tokens. Informationen aus weiter zurückliegenden Positionen müssen jedoch viele Zwischenschritte überstehen, um anzukommen.

Mistral 7B speichert seine Schlüssel und Werte außerdem in einem rollierenden Puffer mit w Speicherplätzen. Dabei wird Position i in den Platz i mod w geschrieben und der älteste Eintrag überschrieben. Bei Sequenzen mit 32.000 Tokens verringerte sich der Cache-Speicher dadurch auf ein Achtel, ohne die Qualität zu beeinträchtigen. Dank der festen Cache-Größe bleibt der Speicherverbrauch unabhängig von der Generierungslänge konstant.

Lokale und globale Schichten abwechseln#

Reine Fenster können in einem Schritt keine weit zurückliegende Information kopieren, daher kombinieren aktuelle Modelle SWA-Schichten mit vollständig globalen Schichten. Googles Gemma 2 wechselte sie im Verhältnis eins zu eins mit einem Fenster von 4.096 Tokens ab, und Gemma 3 verwendet fünf lokale Schichten pro globaler Schicht mit einem Fenster von nur 1.024 Tokens. Dadurch bleibt der KV-Cache über ein Kontextfenster von bis zu 128K Tokens klein. Nur die globalen Schichten speichern Schlüssel und Werte für den gesamten Kontext, sodass fünf von sechs Schichten jeweils nur 1.024 Tokens enthalten.

Vision-Transformer verwenden dieselbe Kombination für hochauflösende Bilder, bei denen jedes Bildsegment ein Token ist. Metas SAM 3, das im Ultralytics-Paket für die Bildsegmentierung verfügbar ist, verarbeitet in 28 der 32 Blöcke seines Bildcodierers jeweils nicht überlappende Fenster aus 24 × 24 Bildsegmenten und in den Blöcken 7, 15, 23 und 31 das gesamte Raster aus 72 × 72 Bildsegmenten eines 1.008 Pixel großen Bildes global.

SWA im Vergleich zu verwandten Begriffen#

• Swin Transformer: Teilt das Bild in feste, nicht überlappende Fenster auf, die von allen Tokens im jeweiligen Fenster gemeinsam genutzt werden, und verschiebt das Raster zwischen den Schichten, damit benachbarte Fenster Informationen austauschen.
• Sparse Attention: Die übergeordnete Familie fester Muster. SWA ist ihr einfachstes Mitglied; andere Varianten ergänzen dilatierte Fenster, globale Tokens oder zufällige Verbindungen.
• Attention Sinks: Die ersten Tokens einer Sequenz, die überschüssiges Attention-Gewicht aufnehmen. Wird bei der Inferenz ein Fenster auf ein Modell angewendet, das mit vollständiger Attention trainiert wurde, stört das Entfernen dieser Tokens die Textgenerierung. StreamingLLM behält sie deshalb neben dem Fenster im Cache.
• Deformable Attention: Jede Abfrage berücksichtigt einige Abtastpunkte, deren Versätze sie erlernt. Dadurch hängt die Nachbarschaft vom Inhalt ab.
• Flash Attention: Ein GPU-Kernel für exakte Attention. Die Implementierungen akzeptieren eine Fenstergröße, sodass SWA-Modelle die Blöcke außerhalb des Fensters vollständig überspringen.