MIT-Student implementiert Vision Transformer von Grund auf

Von Pixeln zu Vorhersagen: Der Bau eines Transformers für Bilder Seit Beginn der Computer-Vision-Revolution standen Faltungsneuronale Netze (CNNs) im Mittelpunkt fast jeder wichtigen Erkenntnis. Doch was wäre, wenn sie uns die ganze Zeit zurückhielten? Im Jahr 2020 stellte ein Forscherteam von Google die Vision Transformer (ViT) vor, eine bahnbrechende Architektur, die auf die Idee des Selbstaufmerksamkeitsmechanismus (Self-Attention) basiert und die Faltung ganz fallen lässt. Dieser Schritt markierte eine neue Ära im tiefen Lernen. 1. Hintergrund und Intuition 1.1 Von rekurrenten Modellen zur Aufmerksamkeit in der NLP Bis 2017 dominierten rekurrente neuronale Netze (RNNs) und Long Short-Term Memory-Netze (LSTMs) das Gebiet der natürlichen Sprachverarbeitung (NLP). Ob bei der Übersetzung oder dem Sprachmodellieren, diese Modelle waren erfolgreich, hatten aber fundamentale Einschränkungen. Sie verarbeiteten Sequenzen tokenweise, was die Parallelisierung des Trainings erschwerte, und verloren Informationen, je länger die Sequenzen wurden. Diese Engpässe erschwerten die Skalierung, insbesondere für Aufgaben, die eine globale Verständnis von Sprache erforderten. Im Jahr 2017 veröffentlichte Google in ihrem Paper "Attention Is All You Need" eine neue Architektur: den Transformer. Anstelle tokenweise zu verarbeiten, ermöglicht der Selbstaufmerksamkeitsmechanismus es jedem Token, direkt auf jeden anderen Token in der Sequenz zu achten. So lernen Worte, Fragen zu stellen (Queries), was sie hören sollen (Keys), und relevante Informationen zu sammeln (Values). Dieser Mechanismus beseitigte die Notwendigkeit von Rekurrenz und fixen Schrittfolgen, wodurch die Hauptprobleme der RNNs überwunden wurden. Innerhalb von zwei Jahren Revolutionierte die Transformer-Architektur die NLP vollständig. Sie erwies sich als effizienter, leichter skaliert und besser in der Modellierung langfristiger Abhängigkeiten. Große Durchbruchsmodule wie BERT, GPT und T5 basieren auf diesem Prinzip. 1.2 Kann Aufmerksamkeit die Faltung ersetzen? Der Wechsel zur Vision Die Frage, ob der Aufmerksamkeitsmechanismus auch die Faltung in der Computer-Vision ersetzen könnte, lag nahe. Dosovitskiy et al. stellten 2020 in ihrem Paper "An Image Is Worth 16×16 Words" die Vision Transformer vor. Sie schlugen vor, ein Bild wie einen Satz zu behandeln. Anstelle von Faltungsfiltern teilten sie Bilder in Patches und fütterten diese in einen standardmäßigen Transformer. Obwohl early ViTs große Datensätze benötigten, um mit CNNs konkurrieren zu können, zeigte dieser Ansatz, dass aufmerksamkeitsbasierte Modelle auch in der Vision funktionieren können. Seitdem hat die Vision Transformer eine Flut von Verbesserungen ausgelöst und ist zu einem zentralen Baustein der modernen Computer-Vision geworden. 2. Wie Vision Transformers funktionieren 2.1 Patch-Verarbeitung Transformers wurden ursprünglich für die Verarbeitung von Sequenzen wie Sätzen entwickelt. Bilder sind jedoch zweidimensionale Pixelgitter. Daher müssen Bilder in Patches unterteilt werden, die dann in 1D-Vektoren umgewandelt und in eine feste Größe projiziert werden. Zum Beispiel: Ein Bild wird in quadratische, nicht überlappende Patches (z.B. 16x16 Pixel) aufgeteilt. Jeder Patch wird in einen 1D-Vektor geflattet und durch eine lineare Projektion in eine feste Größe eingebettet. 2.2 Klassen-Token und positionale Einbettungen Transformers benötigen zwei zusätzliche Komponenten, um mit Bildsequenzen korrekt zu arbeiten: einen [CLS]-Token und positionale Einbettungen. Der [CLS]-Token wird am Anfang der Sequenz hinzugefügt und aggregiert während der Selbstaufmerksamkeit globale Informationen. Die positionalen Einbettungen kodieren die räumliche Struktur, da Transformers nicht von Natur aus die Reihenfolge verstehen. 2.3 Multi-Head Selbstaufmerksamkeit (MHSA) Das Herzstück des Vision Transformers ist die Multi-Head Selbstaufmerksamkeit. Statt einer einzelnen Aufmerksamkeitsfunktion teilt MHSA die Eingabe in mehrere "Heads". Jeder Head lernt, sich auf verschiedene Aspekte der Eingabe zu konzentrieren — einige auf Kanten, andere auf Textur, wieder andere auf räumliche Anordnung. Die Ausgaben der Heads werden dann konkateniert und zurück in den ursprünglichen Einbettungsraum projiziert. 2.4 Transformer-Encoder Einmal die Selbstaufmerksamkeit implementiert, wird sie in einem größeren Modul, dem Transformer-Block, eingepackt. Dieser Block kombiniert Selbstaufmerksamkeit und Multilayer-Perceptrons (MLPs) mit Residualverbindungen. Das Design ermöglicht es dem Modell, global zu denken (durch Selbstaufmerksamkeit) und diese Darstellungen zu transformieren (durch MLPs), wobei es durch Skip-Verbindungen stabil bleibt. 2.5 Klassifikationskopf Nach der Verarbeitung der Eingabe durch mehrere Transformer-Blöcke muss das Modell eine finale Vorhersage produzieren. Dies übernimmt der Klassifikationskopf. Der [CLS]-Token, der während der Einbettung hinzugefügt wurde, aggregiert Informationen von allen Bildpatches. Seine finale Einbettung wird dann durch eine einfache lineare Schicht in Klassenlogits umgewandelt. 3. Implementierungschritte 3.1 Patch-Verarbeitung Um Bildpatches in eine Sequenz von Einbettungen zu konvertieren, verwenden wir einen Trick: eine Conv2d-Schicht. Diese extrahiert nicht überlappende Patches und wendet eine gelernte lineare Projektion an — alles in einer Operation. Dies ist sauber, effizient und leicht zu backpropagieren. ```python class PatchEmbed(nn.Module): def init(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().init() self.img_size = img_size self.patch_size = patch_size self.num_patches = (img_size // patch_size) ** 2 self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): x = self.proj(x) x = x.flatten(2) x = x.transpose(1, 2) return x ``` 3.2 Klassen-Token und positionale Einbettungen Hier definieren wir ein ViTEmbed-Modul, das den [CLS]-Token und die positionalen Einbettungen hinzufügt. Dies erzeugt eine Sequenz der Form [B, num_patches + 1, embed_dim], die bereit ist für den Transformer-Encoder. ```python class ViTEmbed(nn.Module): def init(self, num_patches, embed_dim): super().init() self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) def forward(self, x): batch_size = x.shape[0] cls_tokens = self.cls_token.expand(batch_size, -1, -1) x = torch.cat((cls_tokens, x), dim=1) x = x + self.pos_embed return x ``` 3.3 Multi-Head Selbstaufmerksamkeit Jetzt implementieren wir die Multi-Head Selbstaufmerksamkeit, einen der wichtigsten Teile des Vision Transformers. Jeder Eingabetoken wird linear in Query (Q), Key (K) und Value (V)-Vektoren projiziert. Die Aufmerksamkeit wird parallel über mehrere Heads berechnet, dann konkateniert und zurück in die ursprüngliche Einbettungsdimension projiziert. ```python class MyMultiheadAttention(nn.Module): def init(self, embed_dim, num_heads): super().init() assert embed_dim % num_heads == 0, "embed_dim muss durch num_heads teilbar sein" self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.q_proj = nn.Linear(embed_dim, embed_dim) self.k_proj = nn.Linear(embed_dim, embed_dim) self.v_proj = nn.Linear(embed_dim, embed_dim) self.out_proj = nn.Linear(embed_dim, embed_dim) def forward(self, x): B, T, C = x.shape Q = self.q_proj(x) K = self.k_proj(x) V = self.v_proj(x) def split_heads(tensor): return tensor.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) Q = split_heads(Q) K = split_heads(K) V = split_heads(V) scores = torch.matmul(Q, K.transpose(-2, -1)) scores /= self.head_dim ** 0.5 attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, V) out = out.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(out) ``` 3.4 Transformer-Encoder Wir kombinieren nun alle Elemente in einem Transformer-Block, der Selbstaufmerksamkeit und MLP-Schichten mit Residualverbindungen stapelt. Dieses Design ermöglicht es dem Modell, global zu denken und Darstellungen zu transformieren, während es durch Skip-Verbindungen stabil bleibt. ```python class TransformerBlock(nn.Module): def init(self, embed_dim, num_heads, mlp_ratio=4.0): super().init() self.norm1 = nn.LayerNorm(embed_dim) self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) self.norm2 = nn.LayerNorm(embed_dim) self.mlp = nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Linear(int(embed_dim * mlp_ratio), embed_dim) ) def forward(self, x): x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x = x + self.mlp(self.norm2(x)) return x ``` 3.5 Alles zusammen Nun, da wir alle wesentlichen Komponenten des Vision Transformers gebaut haben — Patch-Verarbeitung, positionale Einbettungen, Multi-Head Selbstaufmerksamkeit, Transformer-Blöcke und den [CLS]-Token — ist es an der Zeit, alles in ein vollständiges Modell zu integrieren. ```python class SimpleViT(nn.Module): def init(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12, num_classes=1000): super().init() self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim) num_patches = (img_size // patch_size) ** 2 self.vit_embed = ViTEmbed(num_patches, embed_dim) self.blocks = nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) for _ in range(depth) ]) self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) def forward(self, x): x = self.patch_embed(x) x = self.vit_embed(x) x = self.blocks(x) x = self.norm(x) return self.head(x[:, 0]) ``` 4. Training des ViT 4.1 Datensatz: CIFAR-10 Wir trainierten unseren Vision Transformer auf CIFAR-10, einem bekannten Benchmark-Datensatz mit 60.000 Bildern in 10 Klassen (z.B. Flugzeuge, Katzen, Schiffe). Jedes Bild ist nur 32x32 Pixel groß, was CIFAR-10 zu einem idealen Testfall für einfache Modelle macht. 4.2 Modell-Setup: Anpassen des ViT für CIFAR-10 ViTs wurden ursprünglich für große Datensätze wie ImageNet entwickelt. Um das Training auf CIFAR-10 mit begrenzten Ressourcen möglich zu machen, passten wir das Modell an: python model = SimpleViT( img_size=32, # CIFAR-10-Bilder sind 32x32 Pixel patch_size=4, # 4x4 Patches → 64 Tokens in_chans=3, embed_dim=192, # Kleinere Einbettungsgröße depth=6, # Weniger Transformer-Blöcke num_heads=3, # Teilbar in 192 num_classes=10 # Für CIFAR-10 ).to(device) 4.3 Training-Setup Das Modell wurde mit folgenden Parametern trainiert: Optimierer: AdamW Learning Rate Scheduler: Cosine-Annealing Batch Size: 64 Trainingsdauer: Etwa 30 Sekunden pro Epoche dank der Effizienz des Transfomers. 4.4 Ergebnisse Wir trainierten den Vision Transformer für 30 Epochen, was insgesamt etwa 15 Minuten auf einer GPU dauerte. Am Ende des Trainings erreichte das Modell eine Genauigkeit von etwa 60% auf der Testmenge von CIFAR-10 — eine gute Grundlage, angesichts der Einfachheit des Modells und der relativ kleinen Datensatzgröße. Lernfortschritt: Wie die Trainingsplots unten gezeigt, verbesserte sich das Modell kontinuierlich. Es identifizierte viele Samples korrekt (z.B. Katzen und Frösche), hatte aber Schwierigkeiten mit visuell ähnlichen Klassen (z.B. Schiff als Flugzeug). Leistung nach Klasse: Das Balkendiagramm zeigt, wie gut das Modell in den 10 Klassen performte. Besonders bemerkenswert ist die Leistung auf Klassen, dievisuell stark voneinander abweichen. 5. Bewertung durch Branchenexperten Die Vision Transformer hat erhebliche Fortschritte im Bereich der Computer-Vision ermöglicht. Branchenexperten loben die Flexibilität und Effizienz des Modells, insbesondere seine Fähigkeit, langfristige Abhängigkeiten zu erfassen. Obwohl sie noch Verbesserungspotential sehen, insbesondere bei der Reduktion des Speicherverbrauchs und der Verarbeitungsgeschwindigkeit, sind sie von der Zukunftsperspektive begeistert. 6. Unternehmen und Profile Google, das Unternehmen hinter der Vision Transformer, ist ein führender Akteur im Bereich künstlicher Intelligenz und Maschinelles Lernen. Mit seiner starken Forschung und Entwicklungssparte setzt Google neue Standards in der Technologieindustrie. Die Vision Transformer ist nur einer der vielen Durchbrüche, die Google in den letzten Jahren erzielt hat. Die Vision Transformer hat auch andere Unternehmen und Forscher inspiriert, ähnliche Modelle und Architekturen zu entwickeln. Dies führt zu kontinuierlichen Innovationen und einer dynamischen Entwicklung im Bereich der Computer-Vision.