Cascade Token Selection for Transformer Attention Acceleration

Technische Samenvatting: Cascade Token-selectie voor Versnelling van Transformer-attention

Probleemstelling
Transformer-modellen ondervinden een aanzienlijke computationele knelpunt in het self-attention-mechanisme, waarbij de kosten voor het berekenen van het $QK^\top$-product en de daaropvolgende met softmax gewogen som kwadratisch schalen met de sequentielengte ($O(T^2 d_h)$). Hoewel recente methoden zoals FlashAttention de geheugenbandbreedte hebben geoptimaliseerd, blijft de algoritmische complexiteit een barrière voor lange sequenties.

Het begeleidende artikel, Activation Decorrelation Attention (ADA), stelde een oplossing voor om deze complexiteit te reduceren van $O(T^2)$ naar $O(r^2)$ door een kleine set van $r \ll T$ "representatieve" tokens te selecteren op elke laag. ADA identificeert deze tokens met behulp van een Gram-drempel: een token wordt als representatief beschouwd als de maximale cosinuscorrelatie met alle voorgaande tokens onder een drempel van $1 - \tau^2$ ligt. Echter vereist het selectieproces zelf het berekenen van een volledige $T \times T$ Gram-matrix op elke laag, wat kosten oplevert van $O(T^2 d)$. Voor gematigde sequentielengten (bijvoorbeeld $T=2048$) is deze selectie-overhead vergelijkbaar met de attention-kosten die het beoogt te reduceren, waardoor de potentiële snelheidswinst wordt tenietgedaan. Het kernprobleem dat door dit artikel wordt aangepakt, is de hoge computationele kosten van de stap voor het selecteren van representatieve tokens in ADA.

Methodologie: Het Cascade-mechanisme
Het artikel introduceert een "cascade"-mechanisme dat gebruikmaakt van de diepte-coherentie van de set representatieve tokens. De auteurs observeren dat de set informatieve tokens op laag $l$ sterk voorspellend is voor de set op laag $l+1$. Dit is theoretisch onderbouwd in de Lipschitz-continuïteit van transformer-blokken: tokens die op de ene laag bijna-duplicaten (redundant) of gedecorreleerd (representatief) zijn, behouden die relatie doorgaans op de volgende laag, met slechts kleine verstoringen naarmate contextuele informatie wordt toegevoegd.

In plaats van op elke laag de volledige $T \times T$ Gram-matrix opnieuw te berekenen, doet het cascade-mechanisme het volgende:

1. Erft de representatieve set $R^{(l)}$ af van de vorige laag.
2. Valideert de geërfde set door twee kleinere Gram-producten te berekenen:
   • Een $r \times r$ product onder de geërfde representanten om te controleren of er redundantie is ontstaan.
   • Een $(T-r) \times r$ kruisproduct tussen niet-representanten en de geldige geërfde representanten om te controleren of er niet-representanten gedecorreleerd zijn geraakt (d.w.z. nieuwe representanten).
3. Bijwerkt de set met een klein aantal toevoegingen en verwijderingen op basis van deze controles.

Deze aanpak reduceert de selectiecomplexiteit van $O(T^2 d)$ naar $O(Trd)$ per laag, een snelheidswinstfactor van $T/r$.

Belangrijkste Bijdragen

• Algoritmische Innovatie: Het voorstellen van een cascade token-selectie-algoritme dat de berekening van de volledige Gram-matrix vervangt door incrementele validatie- en update-stappen.
• Theoretische Rechtvaardiging: Een bewijs (Propositie 3) dat aantoont dat de diepte-coherentie van de representatieve set een gevolg is van de Lipschitz-continuïteit van transformer-blokken en het additieve karakter van residu-verbindingen.
• Conservatieve Garantie: Een bewijs (Propositie 4) dat aantoont dat de cascade-methode conservatief is; het produceert een super-set van de onafhankelijke representatieve set. Het mist nooit een ware representant, waardoor de attention-kwaliteit niet verslechtert, hoewel het het aantal representanten ($r$) licht kan verhogen.
• Orthogonale Versnelling: Het artikel positioneert deze methode als het token-dimensionale analogon van Gated Subspace Inference (GSI), die lineaire lagen versnelt via diepte-coherente basis-erfenis. Samen dekken ze de volledige transformer-forward pass (zowel de $d$-as als de $T$-as).

Experimentele Resultaten
De methode werd gevalideerd op drie model-families (GPT-2 124M, GPT-J 6B, OPT 6.7B) met AMD MI300X-hardware, met een sequentielengte van $T=512$ en een Gram-drempel $\tau=0.30$.

• Diepte Coherentie: De Jaccard-index tussen representatieve sets op opeenvolgende lagen is hoog (0,83 tot 0,94) in diepe lagen, wat de stabiliteit van de token-set bevestigt.
• Gram-bewerkingen Besparing: Het cascade-mechanisme verlaagde de totale Gram-bewerkingen met:
  • 22% voor GPT-2 (12 lagen)
  • 52% voor GPT-J (28 lagen)
  • 63% voor OPT (32 lagen)
    De besparingen nemen toe met de modeldiepte omdat de initiële volledige-Gram-kosten op laag 0 worden afgeschreven over meer lagen.
• Turnover Rate: In diepe lagen daalt de turnover rate (fractie van de set die wordt gewijzigd) tot 3–18%, wat aangeeft dat de geërfde set minimale aanpassing vereist.
• Schaalvoorspellingen: De auteurs voorspellen dat voor langere sequenties (bijvoorbeeld $T=8192$), waarbij $r$ sublineair groeit, de selectiesnelheidswinst $18\times$ kan bereiken, waardoor selectiekosten worden gereduceerd van een dominante term naar een verwaarloosbare.

Betekenis en Claims
Het artikel beweert dat het cascade-mechanisme een structurele eigenschap van getrainde transformers blootlegt: de set informatieve tokens is een coherente structurele eigenschap die zich door de netwerkdiepte voortplant. Door deze coherentie te exploiteren, maakt de methode de selectie-overhead praktisch toepasbaar voor ADA, waardoor de attention-compressie haar theoretische snelheidswinst kan realiseren zonder geblokkeerd te worden door de selectiestap.

De auteurs benadrukken dat dit een complementaire versnelling is voor bestaande methoden zoals FlashAttention (die geheugenbandbreedte optimaliseert) en sparse attention (die vaste patronen gebruikt). In tegenstelling tot token-pruning of -samenvoegingsmethoden die representaties permanent veranderen, selecteert cascade ADA dynamisch representanten op basis van inputstructuur zonder architecturale wijzigingen of hertraining. De combinatie van GSI en cascade ADA biedt een multiplicatieve reductie in inferentiekosten, waardoor de complexiteit van de forward pass effectief wordt gereduceerd tot de "effectieve complexiteit" van de input in plaats van de nominale modelgrootte.