Technische Samenvatting: Heterogene Staptoewijzing (HSA) voor Efficiënte Videogeneratie

1. Probleemstelling

Diffusietransformatoren (DiT's) hebben state-of-the-art prestaties gevestigd in videogeneratie, maar lijden onder enorme rekenkosten. Standaard inferentieprotocollen passen een uniform aantal ontdoeningstappen ($T$) toe op elke spatiotemporale token in de sequentie. Deze monolithische benadering negeert de inherente asymmetrie van visuele data: achtergrondgebieden, statische patches en grove structuren zijn perceptueel eenvoudiger en vereisen minder ontdoeningstappen dan detailrijke voorgrondobjecten of gebieden met hoge beweging. Bovendien is het menselijk gezichtsvermogen minder gevoelig voor kwaliteitsdegradatie in sterk dynamische componenten vergeleken met statische componenten. Bijgevolg is het uniform toepassen van een volledig ontdoeningbudget op alle tokens rekenkundig verspillend en faalt het om aan te sluiten bij perceptuele prioriteiten.

Bestaande versnellingsmethoden, zoals stapniveau-features caching (bijv. TeaCache), nemen doorgaans globale beslissingen: bij een gegeven iteratie worden óf alle tokens berekend óf alle tokens hergebruiken opgeslagen uitvoer. Deze methoden maken geen gebruik van ruimtelijke of token-niveau heterogeniteit, en behandelen de sequentie als een monolithisch blok in plaats van een collectie tokens met variërende dynamische eisen.

2. Methodologie: Heterogene Staptoewijzing (HSA)

De auteurs stellen Heterogene Staptoewijzing (HSA) voor, een trainingsvrij, modelagnostisch inferentie-algoritme dat tokenspecifieke ontdoeningstrajecten toewijst binnen een voorgeöorde DiT. De kernidee is om de tokensequentie te partitioneren in groepen op basis van hun snelheidsdynamiek, waarbij minder stappen worden toegewezen aan langzamer evoluerende tokens terwijl een volledig schema wordt gehandhaafd voor een baselijn-groep.

2.1 Tokenpartitionering en Planning

De set van $N$ tokens wordt gepartitioneerd in $K$ disjuncte groepen ($G_1, \dots, G_K$). Elke groep $G_k$ krijgt een stapbudget $s_k$ toegewezen, dat een deler moet zijn van het totale aantal stappen $T$.

• Baselijn-groep ($G_1$): Krijgt het volledige budget $s_1 = T$. Deze tokens zijn actief bij elke iteratie.
• Gereduceerde groepen: Krijgen kleinere budgetten toegewezen (bijv. $T/2, T/4$). Deze tokens zijn alleen actief wanneer de globale iteratie $i$ voldoet aan $i \pmod{T/s_k} = 0$.
• Dynamische selectie: Het artikel evalueert vier strategieën voor het toewijzen van tokens aan groepen. De strategie Dynamische Tokenselectie, die tokens toewijst op basis van de per-token L1 relatieve verandering van snelheidsvoorspellingen (waarbij langzamer evoluerende gebieden zoals achtergronden worden geïdentificeerd), levert consistent de beste prestaties op.

2.2 KV-cache Synchronisatie

Een primaire uitdaging bij heterogene planning is dat tokens met verschillende stapaantallen op een gegeven wandkloktijd op verschillende ruisniveaus ($\sigma$) verkeren. Zelfaandacht vereist dat alle tokens naar elkaar verwijzen, wat een synchronisatieprobleem creëert. HSA lost dit op via een KV-cache synchronisatiemechanisme:

1. Actieve tokens: Alleen de subset van actieve tokens ($A_i$) voert een forward pass uit door de QKV-projecties.
2. Cache-update: Verse Key ($K$) en Value ($V$) projecties voor actieve tokens worden in een per-laag cache geschreven, waarbij verouderde vermeldingen worden overschreven.
3. Volledige context-aandacht: Actieve tokens berekenen queries tegen de volledige gecachete $K, V$-matrices (die zowel verse vermeldingen voor actieve tokens als verouderde maar geldige vermeldingen voor overgeslagen tokens bevatten).
4. Complexiteitsreductie: Dit reduceert de zelfaandachtscomplexiteit van $O(N^2)$ naar $O(|A_i| \cdot N)$, waarbij $|A_i|$ het aantal actieve tokens bij stap $i$ is. Overgeslagen tokens ondergaan geen queryberekening, cross-aandacht of feed-forward operaties tijdens inactieve iteraties.

2.3 Gecachete Euler-update

Om de latente toestanden van overgeslagen tokens te laten vorderen zonder extra modelbeoordelingen, hanteert HSA een gecachete Euler-update.

• Elke token slaat de snelheid ($\hat{v}_n$) op die bij de meest recente actieve stap is voorspeld.
• Bij elke globale iteratie worden alle tokens (actief en overgeslagen) voortgezet met dezelfde incrementele Euler-stap $(\sigma_{i+1} - \sigma_i)$.
• Actieve tokens gebruiken hun vers berekende snelheid; overgeslagen tokens hergebruiken hun gecachete snelheid.
• Deze update wordt geformuleerd als een enkele tensoroperatie over alle $N$ tokens, waardoor vertakkingen worden vermeden en GPU-efficiëntie behouden blijft.

2.4 Fasebewuste Cachingvenster

Omdat vroege en late ontdoeningfasen respectievelijk cruciaal zijn voor globale structuur en fijne details, beperkt HSA caching tot het middengedeelte van het traject. De eerste en laatste $m$ stappen voeren volledige, ongecachete forward passes uit voor alle tokens, zodat structurele integriteit en detailverfijning niet worden aangetast door verouderde benaderingen.

3. Belangrijkste Bijdragen

1. HSA-algoritme: Een volledig trainingsvrij, plug-and-play inferentieparadigma dat heterogene ontdoeningstrajecten toewijst aan spatiotemporale tokengroepen op basis van snelheidsdynamiek.
2. KV-cache Synchronisatie: Een mechanisme dat de volledige $N$-token aandachtscontext voor actieve tokens behoudt terwijl inactieve tokens volledig worden omzeild, waardoor de per-iteratie aandachtskosten worden gereduceerd tot $O(|A_i| \cdot N)$.
3. Gecachete Euler-update: Een eenvoudige tensoroperatie die overgeslagen tokens voortzet met behulp van gecachete snelheden, waarbij de generatiekwaliteit wordt gehandhaafd zonder extra modelbeoordelingen.
4. Uitgebreide Evaluatie: Een analyse van vier tokengroeperingsstrategieën en vier budgetvoorinstellingen, die robuustheid aantoont over verschillende inferentiebudgetten.

4. Experimentele Resultaten

De auteurs evalueerden HSA op Wan-2.1-1.3B en Wan-2.2-A14B (Tekst-naar-Video en Afbeelding-naar-Video) en LTX-2 (Audio-Video).

• Kwaliteit-Tijdverloop Trade-off: HSA bereikt een superieure Pareto-grens vergeleken met uniforme Flow Matching (FM) en state-of-the-art cachingmethoden zoals TeaCache.
• Aggressieve Versnelling: De voordelen zijn het meest uitgesproken bij agressieve versnellingregimes (50% en 25% van de runtime). Bij 25% runtime presteert HSA aanzienlijk beter dan baselines op de VBench-benchmark (Totale Score: 79,87% versus 75,68% voor FM en 78,33% voor TeaCache).
• Dimensionale Robuustheid: Waar baselines lijden aan "catastrofaal dimensioneel instorten" (ernstige degradatie in specifieke VBench-dimensies zoals bewegingsgladheid of temporele consistentie) onder strakke budgetten, volgt HSA robuust de referentie met volledig budget over alle 16 evaluatiedimensies.
• Afbeelding-naar-Video (I2V): HSA behoudt succesvol de afbeeldingsconditionering en visuele fideliteit in I2V-taken, terwijl uniforme FM-baselines bij verminderde stappen ($T=10$) ernstig structureel instorten vertonen.
• Geen Offline Profiling: In tegenstelling tot TeaCache, dat dure offline profiling vereist om fluctuaties in modeluitvoer te schatten, werkt HSA zonder enige offline tuning of fine-tuning.

5. Betekenis en Claims

Het artikel stelt dat HSA een fundamentele inefficiëntie in DiT-inferentie aanpakt door rekenkracht af te stemmen op de perceptuele en dynamische complexiteit van visuele data. Door tokens individueel in plaats van globaal te behandelen, bereikt HSA aanzienlijke snelheidswinsten zonder de kwaliteitsdegradatie die doorgaans gepaard gaat met agressieve versnelling.

De auteurs benadrukken dat HSA een flexibel kader is dat kan worden geïmplementeerd met diverse tokengroeperingsstrategieën. Hoewel dynamische selectie momenteel de sterkste resultaten oplevert, voorziet het kader in verdere verbeteringen met meer verfijnde groeperingsstrategieën. Het werk legt de basis voor een nieuwe klasse van inferentie-algoritmen die de inherente heterogeniteit van spatiotemporale data benutten om hoogwaardige videogeneratie te democratiseren.