Beyond Self-Attention: Sub-Quadratic Vision Transformers for Fast Image Captioning

Technische Samenvatting: Voorbij Self-Attention: Sub-kwadratische Vision Transformers voor Snelle Beeldtekstgeneratie

Probleemstelling
Beeldtekstgeneratie vereist het genereren van coherente, semantisch betekenisvolle tekstuele beschrijvingen vanuit visuele inhoud, waarbij een brug wordt geslagen tussen computer vision en natuurlijke taalverwerking. Hoewel Transformer-gebaseerde architecturen het veld aanzienlijk hebben geavanceerd door langetermijncontextuele afhankelijkheden te vangen via self-attention, kampen ze met twee primaire beperkingen. Ten eerste resulteren standaard self-attention mechanismen in het berekenen van paar interacties tussen alle beeldpatches, wat leidt tot een kwadratische computationele complexiteit ($O(n^2)$). Dit creëert aanzienlijke geheugen- en computationele overhead, vooral voor hogere resoluties en grootschalige toepassingen. Ten tweede lijden bestaande benaderingen vaak aan een gebrek aan rijke lokale feature-representaties en computationele redundantie, omdat ze alle patches gelijk verwerken, zelfs wanneer veel patches geen semantisch cruciale informatie bevatten.

Methodologie
De voorgestelde architectuur pakt deze beperkingen aan door de Vision Transformer (ViT) framework te herstructureren om het conventionele self-attention mechanisme te vervangen door een probabilistische clustering-benadering gebaseerd op een Gaussian Mixture Model (GMM). Het systeem werkt in twee hoofdfasen:

1. Beeldcodering (GMM-Transformer):

   • Patch Embedding: Invoerbeelden worden verdeeld in niet-overlappende patches, geprojecteerd in $D$-dimensionale embeddings en verrijkt met positional encodings.
   • Context-Gewogen GMM Clustering: In plaats van het berekenen van pairwise attention scores, behandelt het model patch-embeddings als continue willekeurige variabelen gegenereerd door een mengeling van $K$ multivariate Gaussische distributies ($K \ll N$).
   • Leerbare Contextgewichten: Een sub-laag wijst leerbare ruimtelijke gewichten toe aan patches, gemoduleerd door statistische eigenschappen (gemiddelde en variantie), waardoor het model semantisch significante regio's kan prioriteren.
   • EM-Algoritme: Het Expectation-Maximization (EM) algoritme wordt toegepast om soft-clustering uit te voeren. In de E-stap berekent het model "responsibilities" (soft assignment waarschijnlijkheden) die aangeven in welke mate elke patch bij een cluster behoort. In de M-stap worden de clusterparameters (gemiddelden, covarianties en mixing coëfficiënten) bijgewerkt.
   • Complexiteitsreductie: Dit proces vervangt de $N \times N$ attention matrix door een $N \times K$ responsibility matrix, waardoor de computationele complexiteit wordt gereduceerd van kwadratisch $O(n^2)$ naar bijna-lineair $O(nK)$.
   • Aggregatie: Verfijnde patch-features worden verwerkt door feed-forward netwerken met residual connections en layer normalization. Tot slot aggregeert een Global Average Pooling operatie de cluster-centroids tot een compact globaal feature vector.

2. Taaldecodering (Unified GPT-gebaseerde Decoder):

   • De decoder maakt gebruik van een autoregressieve GPT-stijl Transformer.
   • Unified Attention: In tegen tegenstelling tot traditionele encoder-decoder modellen die gemaskeerde self-attention en cross-attention scheiden, integreert deze architectuur ze in één enkele multi-head attention block.
   • Mechanisme: Tekstuele tokens dienen als query-vectoren, terwijl de gecombineerde sequentie van visuele tokens (cluster-centroids) en voorgaande tekstuele tokens dient als key en value vectoren.
   • Causale Masking: Een gespecialiseerde causale mask zorgt ervoor dat tekstuele tokens naar alle visuele tokens kijken, maar alleen naar eerder gegenereerde tekstuele tokens, waardoor de autoregressieve generatie-eigenschappen behouden blijven.
   • Generatie: De output wordt geprojecteerd via een Softmax-laag om de volgende token in de sequentie iteratief te voorspellen.

Belangrijkste Bijdragen

• Complexiteitsreductie: De architectuur reduceert de tijdcomplexiteit van de transformer encoder van kwadratisch naar bijna-lineair ($O(nK)$) door self-attention te vervangen door GMM-gebaseerde soft-clustering, terwijl het vermogen behouden blijft om contextueel coherente captions te genereren.
• Efficiënt Decoderontwerp: Een aangepaste GPT-gebaseerde decoder wordt ingezet die gemaskeerde self-attention en cross-attention verenigt in een enkele computationele block, wat multimodale interactie verbetert en redundantie vermindert.
• Empirische Validatie: Een uitgebreide evaluatie op de Flickr30k dataset demonstreert de competitieve prestaties van het model ten opzichte van bestaande state-of-the-art (SOTA) werken, waarbij een balans wordt gevonden tussen computationele efficiëntie en caption-kwaliteit.

Experimentele Resultaten
Het model werd geëvalueerd op de Flickr30k dataset (met de Karpathy split) met een configuratie van 8 transformer lagen, 4 attention heads en een verborgen dimensie van 768.

• Prestatie-metrieken: Het voorgestelde model behaalde top scores op geavanceerde semantische metrieken: BLEU-4 (0.364), METEOR (0.368) en SPICE (0.199). Deze resultaten duiden op superieure zinsstructuur en scene graph mapping vergeleken met verschillende SOTA modellen (bijv. TAKG, GNN-Visual, Scene Graph).
• Baseline Vergelijking: Hoewel het model achterblijft bij de ADA-baseline op basis van precisie-metrieken (BLEU-1: 0.733 vs 0.794; ROUGE: 0.568 vs 0.608; CIDEr: 0.723 vs 0.946), presteert het beter dan andere baselines zoals TopDown en GNN-Visual. De auteurs merken op dat de architectuur de voorkeur geeft aan diepe context boven eenvoudige replicatie.
• Efficiëntie: Het model opereert met een computationele kosten van ongeveer 21.17 GFLOPs en bevat ongeveer 32.8 miljoen trainbare parameters. Analyse van de confusion matrix toonde hoge true positive rates (45,580) en lage false positive/negative counts, wat wijst op een sterke classificatiecapaciteit.

Betekenis en Claims
Het artikel claimt dat het voorgestelde framework succesvol de computationele inefficiëntie aanpakt die inherent is aan standaard Vision Transformers. Door gebruik te maken van het Gaussian Mixture Model en het Expectation-Maximization algoritme, bereikt het model een significante reductie in computationele complexiteit van kwadratisch naar bijna-lineair, zonder het vermogen te verliezen om betekenisvolle feature-relaties te vangen. De auteurs stellen dat het model een sterke voorspellende capaciteit en schaalbare feature encoding vertoont, wat een verbeterde balans biedt tussen caption-kwaliteit, computationele efficiëntie en schaalbaarheid. Toekomstig werk is gericht op het integreren van geavanceerde multimodale pretraining, het aanpassen aan domeinspecifieke datasets en het implementeren van lichtgewicht modellen voor real-time edge toepassingen.