Token Merging (ToMe)

Token-Merging (ToMe) beschleunigt Transformer, insbesondere den Vision Transformer (ViT), indem ähnliche Tokens innerhalb des Netzwerks zusammengeführt werden, sodass jeder Block weniger davon verarbeitet. Bilder enthalten viele redundante Bereiche, etwa Himmel- oder Straßenflächen, die nahezu gleich aussehen. Werden solche Tokens gemittelt, sinkt der Rechenaufwand, während ihre Informationen erhalten bleiben. ToMe funktioniert mit bereits trainierten Modellen, ohne dass ein erneutes Training nötig ist, und lässt sich auch während des Trainings anwenden. Daniel Bolya und seine Kolleginnen und Kollegen bei Georgia Tech und Meta AI stellten die Methode in Token Merging: Your ViT But Faster vor, einem mündlich präsentierten Beitrag auf der ICLR 2023. Direkt einsatzbereit verdoppelte ToMe den Durchsatz von ViT-L bei 512 Pixeln und ViT-H bei 518 Pixeln, bei einem Genauigkeitsverlust von nur 0,2 % bis 0,3 %. Ein ViT-L-Videomodell wurde bei ähnlichem Genauigkeitsverlust um den Faktor 2,2 beschleunigt.

So funktioniert Token-Merging#

Ein ViT teilt ein Bild in Bildsegmente auf, aus denen jeweils ein Token entsteht. Jeder Block führt Self-Attention aus, deren Aufwand mit dem Quadrat der Token-Anzahl wächst, gefolgt von einem MLP, dessen Aufwand linear wächst. ToMe fügt zwischen Attention und MLP jedes Blocks einen Zusammenführungsschritt ein und entfernt jedes Mal eine feste Anzahl r von Tokens. Ein Modell mit 24 Blöcken hat somit am Ende 192 Tokens weniger als zu Beginn, wenn r = 8 gilt. Für die Zusammenführung wird bipartites Soft Matching verwendet, weil es schnell ist und kein iteratives Clustering erfordert:

1. Aufteilen: Teile die Tokens abwechselnd in zwei Mengen A und B auf.
2. Zuordnen: Verbinde jedes Token in A mit dem ähnlichsten Token in B. Gemessen wird die Kosinusähnlichkeit der über die Köpfe gemittelten Attention-Schlüssel.
3. Die besten Kanten beibehalten: Behalte die r ähnlichsten Paare.
4. Zusammenführen: Bilde aus jedem ausgewählten Paar ein gemeinsames Token. Gewichte dabei nach der Anzahl der ursprünglichen Bildsegmente, für die jedes Token bereits steht, und füge die beiden Mengen anschließend wieder zusammen.

Da zusammengeführte Tokens für mehrere Bildsegmente stehen, addiert ToMe den Logarithmus der Größe jedes Tokens zu den Attention-Werten. Diese Methode heißt proportionale Attention und sorgt dafür, dass ein Token aus vielen Bildsegmenten seinen entsprechenden Anteil an der Attention erhält. Das gesamte Verfahren ist differenzierbar, sodass sich Modelle mit ToMe trainieren lassen, um sie zusätzlich zu beschleunigen.

Dieses PyTorch-Beispiel wendet bipartites Soft Matching auf Metas DINOv2-small ViT aus Hugging Face Transformers an. Als Eingabe dient ein 448 x 448 großes Busbild. Nach jedem der zwölf Blöcke werden r Tokens zusammengeführt, und anschließend wird die resultierende Bildeinbettung mit der des Modells ohne Zusammenführung verglichen:

from urllib.request import urlopen

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
image = Image.open(urlopen("https://ultralytics.com/images/bus.jpg")).resize((448, 448))
pixels = processor(images=image, do_resize=False, do_center_crop=False, return_tensors="pt").pixel_values


def merge(x, r):
    """Bipartite soft matching: merge the r most similar pairs between alternating token sets A and B."""
    cls, a, b = x[:, :1], x[:, 1::2], x[:, 2::2]
    similarity = F.normalize(a[0], dim=-1) @ F.normalize(b[0], dim=-1).T
    best, match = similarity.max(dim=-1)  # each A token's most similar B token
    order = best.argsort(descending=True)
    merged, kept = order[:r], order[r:]
    index = match[merged][None, :, None].expand(1, r, x.shape[-1])
    b = b.scatter_reduce(1, index, a[:, merged], reduce="mean")  # average merged A tokens into their B matches
    return torch.cat([cls, a[:, kept], b], dim=1)


with torch.no_grad():
    for r in (0, 32, 64):
        x, counts = model.embeddings(pixels), []
        for layer in model.encoder.layer:
            counts.append(x.shape[1])
            x = merge(layer(x), r)  # merge r tokens after every transformer block
        embedding = model.layernorm(x)[:, 0]  # CLS token: the image embedding
        reference = embedding if r == 0 else reference
        similarity = F.cosine_similarity(embedding, reference).item()
        print(f"r={r}: mean tokens per block {sum(counts) / 12:.0f}, CLS similarity to unmerged {similarity:.3f}")

Das Modell ohne Zusammenführung verarbeitet in jedem Block 1025 Tokens: 1.024 Bildsegmente plus das CLS-Token. Werden pro Block 32 Tokens zusammengeführt, sinkt der Durchschnitt auf 849, während die Bildeinbettung eine Kosinusähnlichkeit von 0.983 zum Original beibehält. Bei 64 zusammengeführten Tokens pro Block sinkt der Durchschnitt auf 673 und die Ähnlichkeit beträgt 0.887. Das zeigt, wie r den Rechenaufwand gegen die Genauigkeit abwägt.

Einsatzbereiche von Token-Merging#

Token Merging for Fast Stable Diffusion wendet dieselbe Idee auf die Transformer-Blöcke des Entrauschungsnetzwerks von Stable Diffusion an. Dabei lassen sich bis zu 60 % der Tokens entfernen. Bilder werden so bis zu doppelt so schnell und mit bis zu 5,6-mal weniger Speicher erzeugt, ohne zusätzliches Training. Über Bilder hinaus verdoppelte das ToMe-Paper den Durchsatz eines mit Token-Merging trainierten Audio-ViT-B bei einem mAP-Verlust von 0,4 % und beschleunigte das Fine-Tuning eines maskierten Autoencoders für Videos um bis zu das Zweifache. Forschende wenden ähnliche Zusammenführungs- und Beschneidungsmethoden auf visuelle Tokens von Vision-Language-Modellen an. LLaVA-1.5 beansprucht beispielsweise für jedes Bild 576 Tokens des Sprachmodell-Kontexts.

Metas Segment Anything Model (SAM) im Ultralytics-Paket codiert jedes Bild mit 1024 x 1024 Pixeln als 4.096 Bildsegment-Tokens. Im MobileSAM-Vergleich benötigt sein ViT-H-Bildcodierer auf einer GPU 452 ms der insgesamt 456 ms, die SAM pro Bild benötigt. Genau auf diesen Rechenaufwand zielen Methoden zur Token-Reduktion ab.

Token-Merging im Vergleich zu verwandten Konzepten#

• Token-Pruning: Verfahren wie DynamicViT und EViT verwerfen als unwichtig eingestufte Tokens. Dabei gehen deren Informationen verloren, und in der Regel ist Training erforderlich. ToMe führt Tokens stattdessen zusammen, sodass ihre Informationen gemittelt werden, und funktioniert auch mit Modellen, die ohne ToMe trainiert wurden.
• Tokenisierung: Erstellt vor dem Durchlauf des Netzwerks die ursprünglichen Tokens aus Bildsegmenten oder Text. ToMe verringert ihre Anzahl innerhalb des Netzwerks.
• Flash Attention und Sparse Attention: Senken die Kosten oder Laufzeit der Attention bei einer bestimmten Anzahl von Tokens. ToMe reduziert die Anzahl der Tokens, wodurch auch die MLP-Schichten schneller werden, und lässt sich mit diesen Verfahren kombinieren.
• Mixture of Depths (MoD): Ein erlernter Router entscheidet, welche Tokens in Sprachmodellen einen Block überspringen. ToMe führt redundante Tokens dauerhaft anhand einer festen Ähnlichkeitsregel zusammen.
• Niedrigere Eingabeauflösung: Verringert die Token-Anzahl gleichmäßig über das ganze Bild. ToMe führt Tokens vor allem in redundanten Bereichen wie dem Hintergrund zusammen und erhält detailreiche Bereiche in voller Auflösung.