Discrete Cosine Transform Based Decorrelated Attention for Vision Transformers

1. Probleemstelling

Vision Transformers (ViTs) hebben de computervisie revolutionair veranderd door zelf-attentie-mechanismen toe te passen op afbeeldingspatches. Er blijven echter twee kritieke uitdagingen bestaan:

• Gevoeligheid voor Initialisatie: De standaardbenadering initialiseert projectiematrices voor query ($W_Q$), key ($W_K$) en value ($W_V$) met willekeurige getallen (bijvoorbeeld Xavier- of He-initialisatie). Dit behandelt gewichten als "witte ruis", die het spectrum uniform bestrijken maar structurele voorkennis missen. Dit kan leiden tot instabiele trainingsdynamiek, vooral op kleine datasets, en vereist aanzienlijke rekenkracht om te convergeren.
• Rekenkundige Efficiëntie: Zelf-attentie-mechanismen zijn rekenkundig duur ($O(N^2)$ complexiteit met betrekking tot sequentielengte en hoge parameteraantallen). Bestaande compressiemethoden offeren vaak nauwkeurigheid op of slagen er niet in om de specifieke correlatiestructuren die inherent zijn aan visuele data en attentiematrices, optimaal te benutten.

2. Methodologie

De auteurs stellen twee complementaire methoden voor die gebruikmaken van de Discrete Cosine Transform (DCT) om deze problemen aan te pakken:

A. DCT-gebaseerde Initialisatie

In plaats van willekeurige initialisatie, stellen de auteurs voor om één van de projectiegewichtsmatrices ($W_Q$, $W_K$ of $W_V$) te initialiseren met de orthogonale Type-II DCT-matrix.

• Mechanisme: De DCT-matrix dient als een gestructureerde feature-extractor. In tegenstelling tot willekeurige gewichten die ongecorreleerde ruis zijn, bestrijken DCT-basisvectoren het volledige frequentiespectrum systematisch.
• Redenering:
  • Spectrumdekking: Willekeurige initialisatie start alle gewichtvectoren met dezelfde spectrale verdeling. DCT-initialisatie zorgt ervoor dat elke gewichtvector een ander bandbreedtegebied bestrijkt (van lage naar hoge frequentie), waardoor een robuuster startpunt wordt geboden voor het leren van diverse features.
  • Decorrelatie: De DCT-basisvectoren benaderen de Karhunen-Loève-transformatie (KLT) voor sterk gecorreleerde signalen (zoals natuurlijke afbeeldingen), waardoor de inputfeatures aan het begin van de training effectief worden gedecorreleerd.
• Implementatie: Slechts één matrix wordt met DCT geïnitieerd om problemen met gewichtssymmetrie te voorkomen (waarbij identieke initialisatie van meerdere matrices leidt tot slecht leren). De anderen blijven willekeurig geïnitieerd.

B. DCT-gebaseerde Gecomprimeerde Attentie

Deze methode maakt gebruik van het frequentiedomein om dimensionaliteit en rekenkosten te verminderen.

• Mechanisme:
  1. Transformatie: Inputpatches worden met DCT langs de kanaalas naar het frequentiedomein getransformeerd.
  2. Truncatie: Hoge-frequentiecomponenten (die doorgaans overeenkomen met ruis en minder kritieke details) worden afgesneden. Alleen de bovenste $\tau$ fractie van de coëfficiënten (bijvoorbeeld 25%, 50%, 75%) wordt behouden.
  3. Projectie: De projecties voor Query, Key en Value worden berekend op deze tensor met gereduceerde dimensie ($\tau C$ in plaats van $C$).
  4. Reconstructie: De attentie-uitvoer wordt opgepadded naar de oorspronkelijke dimensie en via Inverse DCT (IDCT) teruggetransformeerd.
• Optimalisatie: De IDCT en de laatste lineaire outputlaag ($L_O$) worden samengevoegd tot één enkele lineaire laag om redundante bewerkingen te elimineren, aangezien er geen niet-lineariteit tussen hen bestaat.
• Efficiëntie: Dit verlaagt de rekenkundige complexiteit van de lineaire lagen van $O(C^3)$ naar $O(\tau^2 C^3)$ en verkleint de Softmax-matrix aanzienlijk.

3. Belangrijkste Bijdragen

1. Nieuwe Initialisatiestrategie: Introductie van de eerste DCT-gebaseerde initialisatie voor ViT-attentiegewichten, waarbij wordt aangetoond dat gestructureerde, frequentiedomein-priors superieur zijn aan willekeurige initialisatie.
2. Efficiënte Compressie: Voorstellen van een DCT-gebaseerde compressietechniek die parameters en FLOPs vermindert door hoge-frequentieruis te verwijderen, terwijl essentiële lage-frequentie semantische informatie behouden blijft.
3. Theoretische Rechtvaardiging: Verschaffen van een wiskundig verband tussen DCT en de decorrelatie van sterk gecorreleerde afbeeldingspatches (benadering van KLT), wat uitlegt waarom DCT superieur is aan willekeurige ruis of eenvoudige pooling voor attentiecompressie.
4. Uitgebreide Evaluatie: Valideren van de methoden over meerdere architecturen (Swin-T, Swin-S, ViT-B) en taken (Classificatie, Detectie, Segmentatie).

4. Experimentele Resultaten

De methoden werden geëvalueerd op CIFAR-10, ImageNet-1K en COCO 2017-benchmarks met Swin Transformers als ruggengraat.

• Afbeeldingsclassificatie (CIFAR-10 & ImageNet-1K):

  • Initialisatie: Het initialiseren van $W_K$ (Keys) met DCT verbeterde consistent de Top-1 nauwkeurigheid. Voor Swin-T op ImageNet-1K steeg de nauwkeurigheid van 81,47% naar 81,73%.
  • Compressie: Het behouden van 75% van de DCT-coëfficiënten ($\tau=0,75$) verlaagde het aantal parameters met ongeveer 13% (Swin-T) en 10% (Swin-S), terwijl de nauwkeurigheid ten opzichte van de baseline behouden bleef of licht verbeterde.
  • Vergelijking: De voorgestelde methode presteerde beter dan DCFormer (een eerdere DCT-gebaseerde methode) door lage-frequentiecomponenten nauwkeurig te behouden in plaats van ze te middelen, waardoor discriminatieve features behouden blijven.

• Objectdetectie & Segmentatie (COCO 2017):

  • Met DCT-initialisatie ($W_K$) en 75% compressie behaalden de modellen hogere Average Precision (AP) scores dan standaard Swin Transformers, terwijl het aantal parameters met 3 miljoen (Swin-T) en 5 miljoen (Swin-S) werd gereduceerd.

• Ablatiestudies:

  • Enkele versus Meerdere Initialisatie: Het initialiseren van meerdere matrices (bijvoorbeeld zowel $W_Q$ als $W_K$) met DCT verslechterde de prestaties door problemen met symmetriebreking; het initialiseren van één enkele matrix leverde de beste resultaten op.
  • Noodzaak van DCT: Het verwijderen van de DCT-transformatie uit de compressiepijplijn (waardoor alleen ruimtelijke pruning overbleef) veroorzaakte aanzienlijke dalingen in nauwkeurigheid, wat bevestigt dat frequentiedomein-decorrelatie essentieel is voor het succes van de methode.

5. Betekenis

Dit artikel toont aan dat frequentiedomein-analyse een krachtig hulpmiddel is voor het optimaliseren van Vision Transformers.

• Trainingsstabiliteit: Het biedt een principiële, niet-willekeurige manier om diepe attentienetwerken te initialiseren, wat mogelijk de trainingstijd verkort en de convergentie verbetert bij taken met een tekort aan data.
• Efficiëntie zonder Opoffering: Het bewijst dat hoge-frequentieruis veilig kan worden verwijderd uit attentiemechanismen, waardoor lichtere modellen met minder parameters en lagere FLOPs mogelijk zijn zonder prestaties te compromitteren (en soms zelfs te verbeteren).
• Generaliseerbaarheid: De aanpak is architectuur-agnostisch (getest op Swin en ViT) en taak-agnostisch (werkt voor classificatie, detectie en segmentatie), waardoor het een veelzijdige aanvulling is op de ViT-toolkit.

De code voor de implementatie is publiek beschikbaar, wat verdere research naar frequentie-gebaseerde attentiemechanismen faciliteert.