SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Probleemstelling

Sparse-Linear Attention (SLA) is een eerdere methode die is ontwikkeld om diffusiemodellen (vooral voor video-generatie) te versnellen door een combinatie van sparse (verspreide) en lineaire attention te gebruiken. Hoewel SLA succesvol is, identified de auteurs twee fundamentele beperkingen:

1. Heuristische Routing: SLA gebruikt een vaste, heuristische strategie om te beslissen welke attention-berekeningen naar de sparse tak en welke naar de lineaire tak gaan. Dit gebeurt op basis van de grootte van de attention-weights (grote weights gaan naar sparse, kleine naar lineair). De auteurs stellen dat deze splitsing suboptimaal is en niet noodzakelijk leidt tot de beste balans tussen sparsiteit en kwaliteit.
2. Mismatch in Decompositie: Na een formele analyse van de attention-fout, ontdekten de auteurs dat de output van SLA niet perfect overeenkomt met de oorspronkelijke wiskundige motivatie (de decompositie van de full attention-matrix $P$ in een sparse deel $P_1$ en een lineair deel $P_2$). SLA introduceert een extra projectie-laag om deze fout te compenseren, maar dit leidt tot een complexe en minder efficiënte leeropdracht voor het model.

Daarnaast is er potentie voor verdere versnelling door het gebruik van low-bit kwantisatie, maar dit introduceert vaak kwantisatiefouten die de generatiekwaliteit verminderen.

Methodologie: SLA2

De auteurs stellen SLA2 voor, een verbeterde versie die drie kerninnovaties introduceert om de bovenstaande problemen op te lossen:

1. Leerbaar Routeringsmechanisme (Learnable Router)

In plaats van een vaste heuristiek, introduceert SLA2 een leerbare router ($R$).

• Functie: Deze router dynamisch bepaalt voor elke attention-berekening of deze via de sparse of de lineaire tak moet gaan.
• Implementatie: De router neemt de query ($Q$) en key ($K$) matrices als input (maar niet $V$, aangezien $V$ de attention-verdeling niet bepaalt). Om de rekentijd te beperken bij lange sequenties, worden $Q$ en $K$ eerst gemiddeld (pooling) over blokken van tokens. Vervolgens worden deze geprojecteerd via leerbare lineaire lagen ($proj_q, proj_k$) en wordt een mask ($M$) gegenereerd door de top-$k$ scores te selecteren.
• Training: Tijdens het trainen wordt een differentieerbare "SoftTop-k" operator gebruikt om gradiënten door de router te laten stromen, zodat het model de optimale splitsing kan leren. Tijdens inferentie wordt overgeschakeld op harde Top-k voor snelheid.

2. Faithful Sparse-Linear Formule

SLA2 herschrijft de formule voor de combinatie van de twee takken om directer overeen te komen met de theoretische decompositie.

• Oude aanpak (SLA): $O = O_s + \text{proj}(O_l)$. Hier moet de lineaire tak ook de schalingsfout van de sparse tak compenseren.
• Nieuwe aanpak (SLA2): $O = \alpha \odot O_s + (1 - \alpha) \odot O_l$.
  • Hierbij is $\alpha$ een leerbaar vector (met waarden tussen 0 en 1) die de verhouding bepaalt tussen de sparse en lineaire output.
  • Dit elimineert de noodzaak voor een extra projectie-laag en zorgt ervoor dat de schaling van de sparse output ($O_s$) correct wordt gecombineerd met de lineaire output ($O_l$), wat de leerbaarheid en stabiliteit verbetert.

3. Quantization-Aware Training (QAT) voor Low-Bit Attention

Om de snelheid verder te verhogen, wordt de sparse attention-tak ($O_s$) uitgevoerd met low-bit kwantisatie (bijv. INT8 of FP8).

• QAT-integratie: In plaats van kwantisatie pas na het trainen toe te passen (Post-Training Quantization), wordt de kwantisatie geïntegreerd in het trainingsproces.
• Werkwijze: Tijdens de forward pass wordt gebruik gemaakt van low-bit berekeningen, maar de backward pass (gradienten) blijft volledig in FP16. Dit stelt het model in staat om zijn parameters aan te passen aan de kwantisatiefouten, waardoor de nauwkeurigheid tijdens inferentie behouden blijft.

Belangrijkste Bijdragen

1. Analyse en Reformulering: Een diepgaande analyse van de fouten in de originele SLA-methode, wat leidt tot een wiskundig zuivere en directere formule voor sparse-linear attention.
2. Dynamische Routing: De introductie van een leerbare router die de splitsing tussen sparse en lineaire attention optimaliseert op basis van een duidelijk leerdoel, in plaats van een statische heuristiek.
3. Kwaliteitsbehoud bij Extreme Sparsiteit: Het succesvol combineren van QAT met sparse attention, waardoor extreme snelheidswinst mogelijk is zonder in te leveren op de visuele kwaliteit.
4. State-of-the-art Resultaten: SLA2 presteert beter dan bestaande methoden (zoals VMoBA, VSA en de originele SLA) op video-generatiemodellen.

Resultaten

De auteurs hebben SLA2 getest op de Wan2.1 video-generatiemodellen (1.3B en 14B parameters).

• Snelheid en Sparsiteit:
  • SLA2 bereikt een 97% attention-sparsiteit.
  • Dit resulteert in een 18,6x versnelling in attention-runtime ten opzichte van volledige attention (Full Attention).
  • Zelfs bij 97% sparsiteit is de berekening 11,7x sneller dan concurrenten die werken op 95% sparsiteit.
• Kwaliteit:
  • Opvallend genoeg presteert SLA2 bij 97% sparsiteit beter dan de baseline-methoden bij 90% sparsiteit op alle kwaliteitsmetrieken (zoals Imaging Quality, Motion Smoothness en Subject Consistency).
  • In sommige gevallen overtreft SLA2 zelfs de volledige attention (0% sparsiteit), wat wordt toegeschreven aan de hogere kwaliteit van het fine-tuning dataset.
• Efficiëntie:
  • De end-to-end latentie voor video-generatie wordt met een factor 2,3x (bij 1.3B model) tot 4,35x (bij 14B model) verlaagd.

Betekenis en Impact

SLA2 biedt een praktische en effectieve oplossing voor het "bottleneck" probleem van attention-berekeningen in grote diffusiemodellen voor video-generatie. Door de combinatie van een leerbare router, een wiskundig correcte formule en geavanceerde kwantisatietechnieken, maakt het mogelijk om video's te genereren met een veel lagere rekenkost (FLOPs) en snellere inferentietijden, zonder dat dit ten koste gaat van de visuele kwaliteit. Dit opent de deur voor het gebruik van grotere en complexere video-modellen op beperktere hardware, wat essentieel is voor de bredere adoptie van AI-video-generatie.