Technische Samenvatting: Token Radius Attention voor Efficiënte Videogeneratie

1. Probleemstelling

Video Diffusion Transformers (VDiTs) zijn de dominante architectuur geworden voor hoogwaardige videogeneratie, waarbij ze vertrouwen op dichte 3D zelf-aandacht (self-attention) om complexe ruimtelijke en temporele afhankelijkheden te vangen. Echter, dichte aandacht brengt een kwadratische computationele kostenpost met zich mee ($O(N^2)$) ten opzichte van het aantal video-tokens $N$, wat een aanzienlijke inferentie-bottleneck vormt voor schaling en implementatie.

Bestaande benaderingen om deze kosten te beperken maken gebruik van sparse attention-mechanismen, die primair worden gecategoriseerd in head-level en block-level methoden.

• Head-level methoden wijzen gestructureerde ruimtelijke of temporele patronen toe aan individuele attention heads.
• Block-level methoden selecteren belangrijke query–key block-paren.

De kernbeperking: Beide paradigma's delen doorgaans een enkel computationeel budget (bijv. een vaste retentie-dichtheid of radius) over alle query-tokens binnen een head of block. Deze "gedeelde-allocatie"-aanname conflicteert met de intrinsieke aard van zelf-aandacht, waarbij de Softmax-normalisatie onafhankelijk voor elke query wordt uitgevoerd. Bijgevolg verspillen queries met geconcentreerde aandachtverdelingen (lage entropie) rekenkracht als ze gedwongen worden naar veel keys te kijken, terwijl queries met diffuse verdelingen (hoge entropie) kritieke interacties kunnen verliezen als het gedeelde budget te klein is. Deze mismatch verslechtert de generatiekwaliteit of faalt in het bereiken van optimale efficiëntie.

2. Methodologie: Token Radius Attention (TRA)

De auteurs stellen Token Radius Attention (TRA) voor, een training-vrij framework dat token-specifieke gestructureerde sparsiteit realiseert zonder expliciete per-query key-ranking. TRA werkt op basis van twee belangrijke empirische inzichten afgeleid van de analyse van VDiT-aandachtspatronen.

Inzicht I: Token-Specifieke Aandachtssparsiteit

De auteurs observeren dat de "gereteneerde dichtheid" (het fractie van de top-gerangschikte keys die nodig is om een doel-aandachtsmassa te behouden) met ordes van grootte varieert tussen queries binnen dezelfde laag en head. Cruciaal is dat zij een sterke log-lineaire relatie vinden tussen deze gereteneerde dichtheid en de aandachtsentropie.

• Mechanisme: Hoge entropie duidt op een diffuse aandachtverdeling die een groter token-budget vereist, terwijl lage entropie duidt op een geconcentreerde verdeling die minder tokens vereist.
• Implementatie: TRA gebruikt een analytische benadering om de query-entropie direct te mappen naar een token-specifiek budget ($\hat{B} \approx \tau \exp(H_i) / N$), waardoor de noodzaak voor dure per-query sortering of ranking wordt geëlimineerd.

Inzicht II: Token Radius Attention-patroon

De auteurs observeren dat dominante interacties voor een query een query-gecentreerde circulaire buurt vormen in het ruimtelijke domein. De genormaliseerde aandachtskans neemt ongeveer exponentieel af met de 2D ruimtelijke afstand.

• Mechanisme: In plaats van willekeurige top-$k$ keys te selecteren, zet TRA het scalaire token-budget om in een ruimtelijke radius.
• Temporele Decay: Om videosequenties te verwerken, wordt de ruimtelijke radius gemoduleerd door een temporele afstand-decay regel ($\phi(\delta) = \exp(-\gamma \delta)$), wat een regelmatige spatiotemporele ondersteuningsstructuur creëert.

De TRA-pipeline

1. Entropy-to-Budget: Tijdens een initiële dichte "warm-up" fase berekent TRA de aandacht-entropie voor elke query. Deze entropie wordt gemapt naar een specifiek token-budget.
2. Budget-to-Radius: Het budget wordt omgezet in een query-specifieke basisradius. Deze radius wordt vervolgens voor elk frame aangepast op basis van de temporele afstand om een 2D disk-mask te vormen.
3. Efficiënte Executie:
   • Tile-Major Layout: Video-tokens worden herschikt in een tile-major volgorde om ervoor te zorgen dat ruimtelijk nabijgelegen tokens (die binnen de radius vallen) contigu (aaneengesloten) zijn in de sequentie.
   • Fused CUDA Kernels: Een aangepaste kernel fuseert afstandsberekening, radiusvergelijking, block-pruning en token-voting om de onregelmatige token-niveau radius-masks om te zetten in regelmatige block-sparse masks die compatibel zijn met FlashInfer.
   • Cross-Step Reuse: Entropie wordt alleen berekend tijdens de warm-up fase en hergebruikt voor daaropvolgende sparse stappen, waarbij wordt geprofiteerd van de geobserveerde stabiliteit van entropiepatronen over de denoising-stappen heen.

3. Belangrijkste Bijdragen

1. Ontdekking van Token-Specifieke Sparsiteit: Het paper onthult dat de gereteneerde dichtheid drastisch varieert tussen queries maar een voorspelbare log-lineaire relatie volgt met de aandacht-entropie, en dat dominante interacties een query-gecentreerd radiuspatroon volgen.
2. Token Radius Attention (TRA): Een training-vrij framework dat query-specifieke computationele eisen transformeert naar regelmatige spatiotemporele ondersteuning via een entropy-to-budget-to-radius pipeline, waarbij expliciete key-ranking wordt vermeden.
3. Systeem Co-Design: De auteurs ontwerpen een radius-mask kernel en een gefuseerde entropie-kernel samen om overhead te minimaliseren, wat efficiënte executie op bestaande block-sparse backends mogelijk maakt.
4. Uitgebreide Validatie: TRA wordt gevalideerd over zeven configuraties van Wan2.1, Wan2.2 en HunyuanVideo (variërend van 1.3B tot 14B parameters) voor zowel Text-to-Video (T2V) als Image-to-Video (I2V) taken.

4. Experimentele Resultaten

TRA werd geëvalueerd tegen dichte aandacht en drie training-vrije sparse baselines (SVG, SVG2 en Radial) op de VBench benchmark suite.

• Efficiëntie: TRA behoudt slechts 9%–19% van de aandacht-interacties. Het bereikt een versnelling van 1.56× tot 2.05× ten opzichte van dichte modellen over alle geteste configuraties. Bijvoorbeeld, op HunyuanVideo-13B bereikt het een versnelling van 2.05× voor T2V en 2.02× voor I2V.
• Kwaliteit: TRA behoudt een concurrerende generatiekwaliteit en behaalt vaak de beste VBench Overall scores onder de sparse methoden. Op Wan2.1-14B T2V benadert het bijna de score van dichte aandacht terwijl het een 1.75× versnelling biedt.
• Getrouwheid (Fidelity): Kwalitatieve resultaten laten zien dat TRA de identiteit van onderwerpen, structurele integriteit en temporele consistentie beter behoudt vergeleken met andere sparse methoden, met minder artefacten zoals flikkering of vervorming.
• Ablatie-studies: Het verwijderen van de entropie-gestuurde budget-mapping of de radius-masking (vervangen door 1D afstand) verslechtert de prestaties aanzienlijk, wat de noodzaak bevestigt van zowel token-specifieke adaptiviteit als gestructureerde ruimtelijke ondersteuning.

5. Betekenis en Claims

Het paper claimt dat Token Radius Attention een fundamentele inefficiëntie in huidige sparse attention-paradigma's aanpakt: de mismatch tussen gedeelde allocatie-budgetten en token-specifieke aandachtseisen. Door gebruik te maken van de intrinsieke relatie tussen entropie en aandacht-sparsiteit, bereikt TRA een gunstige kwaliteit-efficiëntie balans zonder extra training of complexe online ranking-mechanismen te vereisen.

De auteurs positioneren TRA als een complementaire techniek voor andere acceleratiemethoden (zoals caching of kwantisatie) die zich specifiek richt op het verminderen van de kwadratische query–key interactiekosten binnen de uitgevoerde aandacht-lagen. Het werk demonstreert dat token-niveau adaptiviteit efficiënt gerealiseerd kan worden via gestructureerde, radius-gebaseerde masking, wat een schaalbaar pad biedt voor de implementatie van hoogwaardige videogeneratiemodellen.