Technische Samenvatting: SparsePR – Training-vrije Sparse Attention voor Videogeneratie en Wereldmodellen

1. Probleemstelling

Videogeneratiemodellen (bijv. DiTs) en video-wereldmodellen verwerken lange spatiotemporele tokensequenties, waardoor kwadratische zelf-attention een primaire bottleneck vormt tijdens de inferentie. Hoewel training-vrije block-sparse attention een pad naar versnelling biedt, kampen bestaande benaderingen met twee fundamentele beperkingen wanneer sparsity wordt behandeld als een rij-gebaseerde eigenschap:

1. Ambiguïteit in Shared-Route Support: Rij-gebaseerde concentratie van attention bepaalt niet de benodigde uitvoerbare support wanneer meerdere queries een enkele routing-beslissing delen. Queries binnen een groep kunnen slecht overlappende key-supports hebben, waardoor de gecombineerde support (de unie van de keys die nodig zijn voor de groep) aanzienlijk groter kan zijn dan de som van de individuele supports.
2. Post-Softmax Foutmiscalculatie: Het behouden van de aandacht-massa alleen bepaalt niet de fout die wordt geïntroduceerd door overgeslagen interacties. Twee routes met een vergelijkbare behouden massa kunnen zeer verschillende output-fouten produceren, afhankelijk van het verschil tussen de outputs veroorzaakt door de weggelaten versus de behouden supports.

Huidige methoden houden vaak geen rekening met deze afhankelijkheden, wat leidt tot ofwel excessieve dichtheid (wat de snelheidswinst vermindert) of aanzienlijke kwaliteitsdegradatie (wat de reconstructiefout verhoogt).

2. Methodologie: SparsePR

De auteurs stellen SparsePR voor, een training-vrij framework dat Response-Coupled Partitioning combineert met Probe-Fitted Residual Reconstruction. De methode werkt volledig online zonder offline training.

A. Response-Coupled Partitioning

In tegenstelling tot eerdere methoden die queries en keys onafhankelijk clusteren of alternerende verfijning gebruiken, construeert SparsePR executable partities in een enkele asymmetrische pass op basis van de respons-geometrie afgeleid van de huidige attention-call:

• Key-Response Coördinaten: Gesamplede queries ($Q_s$) definiëren een metriek $M_K = Q_s^\top Q_s / M_s$. Keys worden geklusterd op basis van hun respons-profielen onder deze gesamplede queries (via $k$-means op eigenvectoren van $M_K$). Dit vormt gepaarde Key/Value (K/V) groepen.
• Query-Response Coördinaten: De centroids van de resulterende K/V-groepen definiëren een nieuwe metriek-ruimte voor queries. Queries worden geklusterd op basis van de gelijkenis van hun pre-softmax logit-profielen over deze K/V-groepen.
• Resultaat: Deze koppeling zorgt ervoor dat queries die een route delen, vergelijkbare support-vereisten hebben, waardoor de expansie van de "pooled support" die bij semantische clustering wordt waargenomen, wordt geminimaliseerd.

B. Probe-Fitted Residual Reconstruction

Om de fout door overgeslagen interacties aan te pakken, vertrouwt SparsePR niet op block-niveau benaderingen of temporele hergebruik. In plaats daarvan:

• Probe Selectie: Een kleine set query-rijen ($M \ll N_q$) wordt geselecteerd (gestratificeerd over query-groepen) en exact geëvalueerd om de dense residu te berekenen ($R = O_{dense} - O_{sparse}$).
• Affine Calibratie: De methode modelleert het residu als een affine functie van de sparse output ($R \approx b + xB$). Het past een call-specifieke affine map toe met behulp van de probe-residuen, geregulariseerd door ridge regressie.
• Subspace Restrictie: Om extrapolatie naar niet-geobserveerde output-richtingen te voorkomen, wordt de correctie geprojecteerd op de subspace gespannen door de leidende singuliere vectoren van de probe-residuen.
• Correctie: De gefitte map schat het residu voor niet-geprobeerde rijen, dat wordt toegevoegd aan de sparse output om de dense benadering te herstellen.

C. Sparse Executie

Het framework permuteert tokens naar aaneengesloten group-major lay-outs en gebruikt de FlashInfer variabele-block sparse attention primitief. De totale geëxecuteerde paar-dichtheid ($\rho_{exec}$) omvat zowel de gerouteerde sparse cellen als de exacte probe-rijen.

3. Belangrijkste Bijdragen

1. Structurele Analyse: Het artikel identificeert drie kritieke afhankelijkheden in executable block-sparse attention:
   • Per-query sparsity bepaalt niet de gedeelde-route support (pooled support kan 3–4$\times$ hoger zijn).
   • Behouden aandacht-massa garandeert geen lage post-softmax fout.
   • Partitie-geometrie beïnvloedt de "affine-verklaarbare" fractie van het residu aanzienlijk.
2. SparsePR Framework: Introductie van een training-vrije methode die response-coupled partitioning (om support-expansie te minimaliseren) combineert met probe-fitted residual reconstruction (om overgeslagen interacties te corrigeren).
3. Uitgebreide Evaluatie: Evaluatie over vier heterogene modellen (HunyuanVideo, Wan2.2, Cosmos-Predict2.5, Cosmos3-Nano) met gebruik van exact-pair accounting en end-to-end timing, wat consistente foutreductie en versnellingen aantoont.

4. Experimentele Resultaten

De auteurs evalueerden SparsePR op text-to-video, image-to-video en fysieke wereld-voorspellings-taken.

• Kwaliteitbehoud: SparsePR komt nauw overeen met de dense benchmark-kwaliteit. Bijvoorbeeld, op HunyuanVideo-13B bereikt het een PSNR van 31.844 (vs. de dense baseline) met een LPIPS van 0.087, waarmee het eerdere sparse methoden zoals SVG2 en SpargeAttn overtreft.
• Efficiëntie: De methode bereikt 1.48$\times$ tot 2.61$\times$ end-to-end versnellingen terwijl een gerealiseerde geëxecuteerde paar-dichtheid van slechts 22.0%–26.0% wordt behouden.
• Ablatie Inzichten:
  • Probe Fitting: Verantwoordelijk voor het merendeel van de foutreductie.
  • Response-Coupled Partitioning: Verlaagt aanzienlijk de "hard-drop" fout en verbetert de reconstructie onder eindige probe-budgetten door de orthogonale residuele energie te verminderen.
  • Overhead: Probe-fitted reconstructie voegt verwaarloosbare overhead toe (ongeveer 1.1% van de totale latentie), aangezien de primaire kostenbesparingen voortkomen uit de sparse routing.

5. Betekenis en Claims

Het artikel claimt dat effectieve sparse executie voor videomodellen afhangt van twee factoren die vaak apart worden behandeld: partitie-geometrie (die de support bepaalt die vereist is voor gedeelde routing) en de reconstructie van overgeslagen interacties (die de output-fout bepaalt).

SparsePR demonstreert dat door het partitieproces te koppelen aan de respons-geometrie van de huidige call en de residu-reconstructie te kalibreren met een kleine set exacte probes, het mogelijk is om bijna-dense getrouwheid te bereiken met substantiële computationele besparingen. De auteurs benadrukken dat deze aanpak training-vrij is, waardoor deze toepasbaar is op bestaande pretrained video DiTs en wereldmodellen zonder de noodzaak voor hertraining of distillatie. De resultaten suggereren dat eenvoudige rij-gebaseerde sparsity onvoldoende is voor videogeneratie, en dat gestructureerde, response-bewuste partitioning gecombineerd met precieze residu-correctie noodzakelijk is voor hoogwaardige versnelling.