Reformer

Der Reformer ist ein speichereffizienter Transformer, den Nikita Kitaev, Łukasz Kaiser und Anselm Levskaya von UC Berkeley und Google Research in Reformer: Der effiziente Transformer auf der ICLR 2020 vorstellten. Er geht zwei Kostenfaktoren an, die die Sequenzlänge begrenzen. Die Attention mit Locality-Sensitive Hashing (LSH) ersetzt die vollständige n × n-Attention-Matrix und verursacht bei einer Sequenz der Länge L Kosten von etwa L log L. Umkehrbare Residual-Layer speichern die für die Backpropagation benötigten Aktivierungen einmal für das gesamte Netzwerk statt einmal pro Layer. Google berichtete, dass der Reformer Kontextfenster mit bis zu 1 Million Wörtern auf einem einzelnen Beschleuniger mit 16 GB Speicher verarbeiten kann. Außerdem verarbeitete er den gesamten Roman Schuld und Sühne als einzelnes Trainingsbeispiel.

LSH-Attention#

Nach der Softmax-Funktion wird die Attention jeder Query von den wenigen Keys dominiert, die ihr am ähnlichsten sind. Deshalb sucht der Reformer nur nach diesen Keys. Er verwendet eine gemeinsame Projektion für Queries und Keys und hasht jeden Vektor mithilfe von winkelbasiertem Locality-Sensitive Hashing: Er multipliziert ihn mit einer Zufallsmatrix R und nimmt die Position des größten Werts in [xR, −xR] als Bucket. Vektoren, die in ähnliche Richtungen zeigen, landen dadurch meist im selben Bucket. Anschließend werden die Tokens nach Buckets sortiert, die sortierte Sequenz in gleich große Abschnitte aufgeteilt und jedes Token berücksichtigt nur den eigenen Abschnitt sowie den vorherigen. Mehrere Hashing-Runden laufen parallel, um ähnliche Paare zu erfassen, die in einer einzelnen Runde getrennt werden.

LSH ist ein klassisches Verfahren für die approximative Vektorsuche nach nächsten Nachbarn: Es findet die meisten k-nächsten Nachbarn einer Query zu einem Bruchteil der Kosten einer exakten Suche. Der Bucketing-Schritt ist kurz genug, um ihn allein auszuführen:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

Der nahe Vektor mit einer Kosinusähnlichkeit von 0,96 zur Query liegt in 72 % der 1.000 Runden im selben Bucket. Der nicht verwandte Vektor mit einem Wert von 0,04 landet in 9 % der Runden im selben Bucket. Auf einen Bucket beschränkte Attention behält daher die meisten relevanten Paare bei.

Umkehrbare Layer und Aufteilung in Abschnitte#

Ein standardmäßiges neuronales Netz speichert für den Rückwärtsdurchlauf die Aktivierungen jedes Layers. Daher wächst der Speicherbedarf sowohl mit der Tiefe als auch mit der Länge. Der Reformer passt umkehrbare Residualnetzwerke (RevNets) an: Er teilt die Aktivierungen in zwei Hälften auf und aktualisiert jede Hälfte mithilfe einer Funktion der jeweils anderen.

Während des Rückwärtsdurchlaufs stellt jeder Layer seine Eingaben anhand seiner Ausgaben wieder her: x2 = y2 − FeedForward(y1) und anschließend x1 = y1 − Attention(x2). Dadurch werden nur die abschließenden Aktivierungen gespeichert und jeder Layer wird einmal neu berechnet. Gradienten-Checkpointing geht ähnlich vor: Es speichert Aktivierungen an einigen Checkpoints und berechnet die Layer dazwischen erneut. Der Reformer führt seine Feed-Forward-Layer außerdem abschnittsweise über Positionen aus, da ihr verborgener Layer um ein Mehrfaches breiter ist als das Modell. Der Aktivierungsspeicher begrenzt die Batchgröße beim Training von Ultralytics YOLO26: Mit batch=-1 misst AutoBatch den GPU-Speicherbedarf bei mehreren Batchgrößen, passt eine Gerade an und wählt die Batchgröße, für die eine Auslastung von 60 % des freien Speichers vorhergesagt wird.

Ergebnisse und Implementierungen#

Der Reformer erreichte bei Zeichentexten aus enwik8 und bei der pixelweisen Bildgenerierung mit ImageNet-64 dieselbe Leistung wie standardmäßige Transformer. Dabei benötigte er deutlich weniger Speicher und war bei langen Sequenzen schneller. Hugging Face Transformers implementiert ihn als ReformerModel und stellt Checkpoints bereit, die mit enwik8 und Schuld und Sühne trainiert wurden.

Reformer und verwandte Modelle#

• Longformer: Legt das Attention-Muster im Voraus fest: ein gleitendes Fenster um jedes Token sowie einige globale Tokens.
• Sparse Attention: Die Familie von Verfahren, die nur einige Token-Paare bewerten. LSH-Attention gehört zu den inhaltsbasierten Verfahren dieser Familie.
• Lineare Attention: Ersetzt die Softmax-Funktion durch eine Kernel-Merkmalsabbildung, sodass sich Attention in linearer Zeit berechnen lässt.
• Flash Attention: Berechnet exakte Attention in Kacheln des GPU-On-Chip-Speichers, ohne die Bewertungsmatrix zu speichern. Dadurch wächst der Speicherbedarf linear, während die Anzahl der Bewertungen quadratisch bleibt.