Accelerating Vision Transformers on Brain Processing Unit

Technische Samenvatting: Versnelling van Vision Transformers op de Brain Processing Unit

Probleemstelling
Het artikel behandelt een kritieke architecturale mismatch tussen gespecialiseerde neurale verwerkingshardware en moderne Vision Transformer (ViT) modellen. Specifiek is de Horizon Brain Processing Unit (BPU) geoptimaliseerd voor vierdimensionale convolutie-operaties (INT8 gekwantiseerd) die typerend zijn voor Convolutional Neural Networks (CNN's). Echter, ViT-modellen, zoals de Data-efficient Image Transformer (DeiT), leunen zwaar op lineaire lagen en Layer Normalization (LayerNorm) die opereren op driedimensionale datastromen. Deze discrepantie voorkomt dat de BPU de kernberekeningen van ViTs versnelt, waardoor dergelijke modellen gedwongen worden op algemene CPU's te draaien en de acceleratiecapaciteiten van de hardware onbenut blijven.

Methodologie
Om deze incompatibiliteit op te lossen zonder dat hertraining van het model vereist is, stellen de auteurs een structurele hervorming van de DeiT-architectuur voor. De kernmethodologie behelst het vervangen van de niet-convolutionale operatoren door mathematisch equivalente convolutionale operatoren die de BPU kan versnellen:

1. Transformatie van Lineaire Lagen: Lineaire lagen, die doorgaans 3D-data $(B, N, C)$ verwerken, worden getransformeerd naar point-wise convoluties. Dit wordt bereikt door de gewichtsmatrix $W \in \mathbb{R}^{O \times I}$ te transformeren naar een 4D-tensor $W' \in \mathbb{R}^{O \times I \times 1 \times 1}$ via unsqueeze-operaties. De input-tensor wordt eveneens getransponeerd van $(B, N, C)$ naar $(B, C, 1, N)$ om aan te sluiten bij de eisen van de 4D-datastroom van de BPU.
2. Herformulering van LayerNorm: De Layer Normalization-operatie wordt gedecomposeerd in een sequentie van $1 \times 1$ convolutionale operaties. Het gemiddelde wordt berekend met een convolutie met gewichten ingesteld op $1/H$ (waarbij $H$ de verborgen dimensie is), en de variantie wordt berekend door het gemiddelde af te trekken, het resultaat te kwadrateren en een tweede convolutie toe te passen. De leerbare parameters $\gamma$ en $\beta$ worden gehervormd om overeen te komen met de inputdimensies.
3. Gewichtsovererving: Een cruciaal aspect van deze aanpak is dat het gestructureerde model geen hertraining of fine-tuning vereist. De originele pre-trained gewichten van standaard DeiT-checkpoints worden direct overgenomen door de gestructureerde convolutionale lagen.
4. Kwantisering: De aangepaste modellen worden verwerkt door de Horizon BPU algoritme-toolchain, waarbij gebruik wordt gemaakt van Post-Training Quantization (PTQ) met een min-max strategie en een kalibratiedataset om de modellen om te zetten naar INT8-precisie.

Belangrijkste Bijdragen

• Eerste Succesvolle Implementatie: De auteurs beweren dat dit de eerste succesvolle implementatie is van Vision Transformers op het BPU-platform die de hardware-acceleratiecapaciteiten volledig benut.
• Architecturale Adaptatie: Ze stellen een nieuwe methode voor om DeiT te herstructureren door lineaire en LayerNorm-lagen te vervangen door equivalente convolutionale operatoren, wat INT8-acceleratie mogelijk maakt zonder architecturaal verlies.
• Implementatie zonder Hertraining: De methode maakt directe overerving van pre-trained gewichten mogelijk, waardoor de computationele kosten en complexiteit van het hertrainen of fine-tunen van het volledige model voor de nieuwe operatorstructuur worden geëlimineerd.
• Uitgebreide Validatie: De aanpak wordt gevalideerd op twee verschillende datasets: de grootschalige ImageNet-dataset en een kleinere dataset voor bloemenclassificatie.

Experimentele Resultaten
Het artikel rapporteert uitgebreide experimenten waarbij de originele DeiT-modellen worden vergeleken met de voorgestelde BPU-geoptimaliseerde, gekwantiseerde versies:

• ImageNet Prestaties:
  • Nauwkeurigheid: Het gekwantiseerde DeiT-Base model bereikt een 80,4% Top-1 nauwkeurigheid op ImageNet, een daling van slechts 1,4% ten opzichte van de originele 81,8%. Gedistilleerde varianten tonen een nog grotere veerkracht, waarbij de DeiT-Base Dist. een nauwkeurigheid van 81,7% behoudt (een daling van 1,7%).
  • Anomaliedetectie: De auteurs merkten op dat zowel de originele als de gekwantiseerde modellen correcte labels voorspelden voor verschillende afbeeldingen in de ImageNet-validatieset die bekend stonden als foutief gelabeld in de officiële annotaties, wat suggereert dat de werkelijke nauwkeurigheid van de modellen iets hoger kan liggen dan de gerapporteerde metrieken.
• Bloemenclassificatie Prestaties:
  • Op een dataset met 5 bloemklassen bereikte het gekwantiseerde DeiT-Small een nauwkeurigheid van 97,4%, met slechts een daling van 0,3% ten opzichte van de full-precision versie.
• Inference Snelheidswinst:
  • De gekwantiseerde modellen vertoonden significante inference-versnellingen op de BPU-hardware vergeleken met hun floating-point tegenhangers.
  • DeiT-Base: Bereikte een versnelling van 3,82×.
  • DeiT-Small: Bereikte een versnelling van 2,16×.
  • DeiT-Tiny: Bereikte een versnelling van 1,30×.
  • De resultaten geven aan dat de voordelen van kwantisering en hardware-acceleratie meer uitgesproken zijn bij grotere modelarchitecturen.

Betekenis en Claims
Het artikel positioneert dit werk als een cruciale stap in het overbruggen van de kloof tussen Transformer-gebaseerde computer vision-modellen en gespecialiseerde embedded hardware. Door aan te tonen dat ViTs kunnen worden ingezet op BPU's met minimale nauwkeurigheidsdegradatie en substantiële snelheidsverbeteringen, stellen de auteurs dat hun methode een praktische oplossing biedt voor real-world edge deployment scenario's. Het vermogen om pre-trained gewichten te gebruiken zonder hertraining wordt benadrukt als een grote efficiëntiewinst, wat de overgang naar hardware-geaccelereerde Transformers haalbaar maakt voor omgevingen met beperkte middelen. De auteurs claimen bescheiden dat dit de eerste succesvolle poging vertegenwoordigt om de BPU-acceleratie voor DeiT-modellen volledig te benutten.