Enhancing Layer Interaction Using Key-Correlated Layer Attention

Technische Samenvatting: Verbeteren van Laaginteractie via Key-Geë correlateerde Laag-Attention

Probleemstelling
Recente vooruitgang in het ontwerp van diepe neurale netwerken (DNN's) heeft gebruikgemaakt van Laag-Attention (Layer Attention) om interacties tussen lagen te verbeteren door de outputs van lagen als sequentie-tokens te behandelen. Echter, standaard laag-attention lijdt onder een kwadratische computationele complexiteit ($O(D^2)$) ten opzichte van de netwerkdiepte $D$, naast een significante groei in geheugengebruik door de noodzaak om outputs van alle voorgaande lagen te bufferen.

Om deze problemen aan te pakken, introduceerden eerdere werken Recurrent Layer Attention (RLA) en lineaire benaderingen zoals MRLA-L. Hoewel RLA de computationele overhead vermindert door Key ($K$) en Value ($V$) representaties te bufferen, fungeert het als een statisch mechanisme waarbij kenmerken van vroege lagen vast komen te liggen zodra ze zijn gegenereerd, wat dynamische informatie-updates beperkt. Bovendien proberen lineaire benaderingen zoals MRLA-L de complexiteit te verminderen door de attention-output lineair te benaderen. De auteurs observeren dat deze aanpak de bijdragen van eerdere lagen exponentieel verzwakt (zoals getoond in hun analyse van de vervalcoëfficiënten $\beta_l$), waardoor het moeilijk is om effectieve langetermijn afhankelijkheden tussen lagen (long-range cross-layer dependencies) te vestigen. Daarnaast slagen bestaande linearisatiemethoden vaak niet in het aanpakken van de geheugenbottleneck van het bufferen van voorgaande laag-outputs of de specifieke uitdagingen van het toepassen van lineaire attention op hiërarchische DNN-structuren.

Methodologie: Key-Correlated Layer Attention (KCLA)
De voorgestelde Key-Correlated Layer Attention (KCLA) is een nieuwe gelinieerde laag-attention methode, ontworpen om een lineaire computationele complexiteit te bereiken terwijl dynamische updates en het modelleren van langetermijnafhankelijkheden behouden blijven. De methodologie is gebouwd op drie kerninzichten en mechanismen:

1. Observatie van Key-correlatie: Empirische analyse onthult dat Key-representaties over verschillende lagen in standaard laag-attention een hoge cosinus-gelijkenis vertonen. Door hiervan gebruik te maken, benadert KCLA de standaard laag-attention output door de lineaire transformatie van de Value ($V$) te extraheren en de Key-representaties te benaderen als zijnde directioneel gecorreleerd. Dit maakt het mogelijk om de outputs van lagen te aggregeren in een recurrente staat ($U^t$) zonder de volledige geschiedenis van $K$- en $V$-matrices op te slaan.
2. Dynamische Softmax-restauratie: Een cruciale uitdaging bij het lineariseren van attention is het verlies van de softmax operatie, die essentieel is voor het normaliseren van attention-scores en het mogelijk maken van dynamische distributie-aanpassingen. KCLA introduceert een gemodificeerd multi-head attention mechanisme om deze capaciteit te herstellen:
   • Het maakt gebruik van meerdere heads met verschillende schalingfactoren (temperaturen) afgeleid van een logaritmisch uniforme sequentie.
   • Het gebruikt een leerbare lineaire transformatie $\Phi$ om de norm van de Query ($\|Q^t\|$) en de correlatiehoek ($\cos \theta$) naar head-specifieke gewichten te mappen.
   • Dit stelt het model in staat om de attention-distributie dynamisch aan te passen bij elke laag, waardoor de rol van de temperatuurparameter in de softmax effectief wordt gesimuleerd zonder de volledige $N \times N$ matrix te berekenen.
3. Cross-Stage Adaptatie: Om DNN's te accommoderen met variërende feature map-groottes en kanaaldimensies over verschillende stadia (stages), implementeert KCLA een specifieke adaptatieprocedure (Algoritme 1). Dit omvat het normaliseren van geaccumuleerde kenmerken, het toepassen van een $1 \times 1$ convolutie met een stride voor downsampling, en het herberekenen van normalisatiefactoren om de dimensionale verschillen tussen stadia te overbruggen.

Belangrijkste Bijdragen

• Observatie en Bewijs: De auteurs demonstreren dat Key-representaties over verschillende lagen in standaard laag-attention sterk gecorreleerd zijn in richting, wat de theoretische basis biedt voor de lineaire benadering.
• Nieuwe Linearisatie: KCLA wordt voorgesteld als een gelinieerde laag-attention methode die zowel de statische aard van RLA als de exponentiële verzwakking van MRLA-L vermijdt. Het bereikt een lineaire computationele complexiteit en een vaste ruimtelijke complexiteit die onafhankelijk is van de netwerkdiepte.
• Dynamische Softmax-benadering: Een nieuwe methode wordt geïntroduceerd om de softmax-operatie te benaderen met behulp van multi-head schaling en leerbare transformaties, wat een dynamische en adaptieve aanpassing van de attention-distributies bij elke laag mogelijk maakt.

Experimentele Resultaten
De auteurs evalueerden KCLA op drie fundamentele visie-taken: beeldclassificatie, objectdetectie en medische beeldsegmentatie.

• Beeldclassificatie: Op CIFAR-100 behaalde KCLA de beste prestaties over ResNet-dieptes van 20, 56 en 110, waarbij het MRLA-L, DLA-L en standaard laag-attention overtrof. Op ImageNet-1K verbeterde de integratie van KCLA in ResNet-50 en ResNet-101 de Top-1 nauwkeurigheid respectievelijk met +1,7% en +1,4% ten opzichte van de baseline. Het behield een licht maar consistent voordeel ten opzichte van MRLA-L, terwijl het aanzienlijk minder parameters gebruikte dan de dynamische DLA-L methode.
• Objectdetectie: Op de COCO2017 dataset met gebruik van Faster R-CNN en Mask R-CNN toonde KCLA competitieve prestaties. Met ResNet-101 overtrof het de complexere DLA-L architectuur op verschillende metrieken (AP$_{75}$, AP$_S$, AP$_M$), terwijl het 2,2 miljoen parameters minder gebruikte.
• Medische Beeldsegmentatie: Op de ISIC2017, ISIC2018 en Kvasir-SEG datasets presteerde KCLA consequent beter dan MRLA-L en de baseline ResNet-UNet, waarbij het superieure mIoU en DSC scores behaalde met minimale parameter-overhead.
• Complexiteitsanalyse: In tegen tegenstelling tot MRLA-B, waarvan geheugen en FLOPs kwadratisch groeien met de diepte, vertoont KCLA een lineaire groei in zowel de computationele last als de geheugenkosten.

Betekenis en Claims
Het artikel claimt dat KCLA succesvol de kloof overbrugt tussen de efficiëntie van lineaire attention en de representatiekracht van standaard attention. Door de standaard laag-attention mechanisme direct te benaderen op basis van Key-correlaties, vermijdt KCLA de problemen van statische feature-representatie van RLA en de degradatie van langetermijn afhankelijkheden van MRLA-L.

De auteurs benadrukken dat hun methode het vermogen behoudt om efficiënt langetermijn afhankelijkheden tussen lagen te modelleren, wat een wenselijke balans biedt tussen efficiëntie en prestaties. Zij merken op dat de voorgestelde methode voor het herstellen van de softmax-functie een kritieke factor is in het bereiken van deze resultaten. Hoewel de huidige implementatie vertrouwt op vooraf ingestelde temperatuurparameters, identificeren de auteurs het ontwerp van een volledig zelf-adaptieve, dynamische temperatuur multi-head lineaire laag-attention als een richting voor toekomstige optimalisatie. De code is publiekelijk beschikbaar om reproduceerbaarheid te faciliteren.