Technische Samenvatting: CoReDiT

Probleemstelling

Diffusietransformatoren (DiTs) hebben state-of-the-art prestaties gevestigd in beeld- en videogeneratie door effectief lange-afstand afhankelijkheden te modelleren via attentiemechanismen. Hun rekenkosten schalen echter kwadratisch met het aantal tokens, waardoor hoogresolutie synthese en videogeneratie met meerdere frames voor op hulpbronnen beperkte mobiele platformen onbetaalbaar duur worden. Hoewel een significant deel van de tokens in de latente ruimte overeenkomt met visueel redundante gebieden (bijvoorbeeld uniforme achtergronden), staan bestaande token-pruningmethoden voor drie hoofduitdagingen wanneer ze worden toegepast op diffusie-inferentie:

1. Efficiënte Lokalisatie: Het identificeren van tokens met lage salientie zonder de kwadratische overhead van het berekenen van globale attentiescores.
2. Semantische Behoud: Het handhaven van visuele continuïteit en het voorkomen van artefacten (bijvoorbeeld textuurbreuken, vervormingen van grenzen) wanneer tokens worden overgeslagen, wat kritiek is voor generatieve taken in tegenstelling tot discriminatieve taken.
3. Adaptieve Planning: Het bepalen van een pruningstrategie die zich aanpast aan variërende redundantieniveaus over verschillende transformerblokken en ontdoeningstijdstappen.

Methodologie: CoReDiT

CoReDiT is een gestructureerd token-pruningframework dat is ontworpen om voorgeprogrammeerde DiTs te versnellen terwijl de outputtrouw wordt behouden. Het werkt via drie kerncomponenten:

1. Tokenselectie op basis van Ruimtelijke Coherentie

In plaats van te vertrouwen op globale attentiescores (die $O(N^2)$ zijn), maakt CoReDiT gebruik van een Ruimtelijke Coherentie (SC) score om lokale redundantie te schatten in lineaire tijd ($O(N)$).

• Mechanisme: Het tokenrooster wordt opgedeeld in niet-overlappende roosters. Voor elke token wordt de SC-score berekend als het gemiddelde van de cosinus-sequentie tussen de genormaliseerde embedding van de token en de gemiddelde genormaliseerde embedding van zijn rooster.
• Logica: Tokens met hoge SC-scores worden als redundant beschouwd omdat ze goed worden verklaard door hun ruimtelijke buren.
• Selectie: De top-$K$ tokens met de hoogste SC-scores worden geselecteerd om over te slaan ($X_S$), terwijl de overige tokens ($X_R$) doorgaan naar de Multi-Head Self-Attention (MHSA)-laag.

2. Coherentie-gestuurde Tokenreconstructie

Om ruimtelijke discontinuïteiten te voorkomen die worden veroorzaakt door het overslaan van tokens, reconstrueert CoReDiT de attentie-outputs voor overgeslagen tokens ($Y_S$) met behulp van de outputs van behouden buren ($Y_R$).

• Benadering: In plaats van paarsgewijze gelijkenissen te berekenen voor reconstructie (wat overhead toevoegt), benadert de methode de gelijkenisgewichten met behulp van de vooraf berekende SC-scores van de buren.
• Lokale Subroosters: Om oververzachting te voorkomen, wordt de reconstructie uitgevoerd binnen kleinere subroosters.
• Micro-ontwerpen:
  • Afwisselende Roostergroottes: Roosterafmetingen wisselen af tussen opeenvolgende blokken om grensartefacten te verschuiven en informatieverplaatsing over roostergrenzen mogelijk te maken.
  • m-Stride Behoud: Een deterministisch stride-patroon zorgt ervoor dat er in elk subrooster ten minste één token wordt behouden, wat noodzakelijke ankers biedt voor reconstructie.

3. Progressieve, Blok-Adaptieve Pruningplanning

Omdat redundantie varieert over transformerblokken en ontdoeningstappen, hanteert CoReDiT een progressieve pruningstrategie tijdens lichtgewicht fine-tuning.

• Blok-Adaptiviteit: Het pruningbudget wordt toegewezen aan blokken met de hoogst gemeten redundantie (geschat door de som van SC-scores van kandidaat-tokens).
• Tijdstap-afname: Pruning is agressiever in vroege ontdoeningstappen (waar ruis overheerst en structuur grof is) en wordt verminderd in latere stappen (waar fijne details worden verfijnd). Een twee-fase afname-schema wordt gebruikt om efficiëntie in evenwicht te brengen met de complexiteit van het beheren van stap-specifieke computatiegrafieken.
• Training: Het model wordt gefine-tuned met behulp van een distillatieverlies dat de oorspronkelijke modeloutputverlies en een pruning-blokverlies combineert om prestaties te herstellen.

Belangrijkste Bijdragen

1. Lineaire-Tijd Redundantie Schatting: De introductie van een ruimtelijke coherentiescore die redundante tokens identificeert in $O(N)$ tijd, waardoor de rekenkundige bottleneck van globale attentieanalyse wordt vermeden.
2. Behoud van Dicht Rooster: Een reconstructiemechanisme dat overgeslagen attentie-outputs synthetiseert vanuit buren, zodat het tokenrooster dicht blijft voor daaropvolgende lagen en visuele artefacten worden voorkomen.
3. Adaptieve Pruningstrategie: Een progressieve, blok-adaptieve planning die dynamisch pruningbudgetten toewijst op basis van empirische redundantiestatistieken over diepte en tijd, wat stabiele aanpassing faciliteert.

Experimentele Resultaten

CoReDiT werd geëvalueerd op state-of-the-art DiT-ruggengraten, waaronder PixArt-α, PixArt-Σ en MagicDrive-V2.

• Efficiëntiewinsten:
  • Bereikte tot 55% reductie in self-attention FLOPs.
  • Inferentieversnellingen: 1,33× op cloud-GPU's (Nvidia H100) en 1,72× op mobiele NPUs (Qualcomm Snapdragon 8 Elite).
  • Geheugenruimte: Maakte generatie mogelijk bij resoluties tot 1920 op mobiele apparaten, waarbij het basismodel Out-of-Memory (OOM)-fouten tegenkwam bij 1600.
• Kwaliteitsbehoud:
  • Op PixArt-α-1024 met 40% pruning behield CoReDiT een FID van 28,7 (tegenover 27,3 voor de baseline) en concurrerende CLIP/IS-scores. Met distillatie werd de FID-kloof verkleind tot 27,4.
  • Op PixArt-Σ-2048 verminderde een gemiddelde pruningratio van 23% de self-attention latentie met 22% met slechts een bescheiden daling in kwaliteitsmetrieken.
  • Op MagicDrive-V2 (video) behield een reductie van 39% in self-attention FLOPs de kwaliteit per frame (PSNR, LPIPS, SSIM) en conditionele uitlijning (mAP, mIoU), hoewel een lichte daling in temporele consistentie (FVD) werd genoteerd.
• Ablatiestudies:
  • Willekeurige tokenselectie resulteerde in catastrofale kwaliteitsdegradatie (FID 644,5), wat de noodzaak van coherentie-gebaseerde selectie bevestigt.
  • Het uitschakelen van reconstructie leidde tot aanzienlijk kwaliteitsverlies (FID 30,4 versus 28,5), wat het belang van de reconstructiestap bevestigt.
  • Uniforme pruning gevolgd door fine-tuning presteerde slechter dan het voorgestelde progressieve adaptieve schema.

Betekenis en Claims

De paper beweert dat CoReDiT een praktische oplossing biedt voor het implementeren van hoogtrouw diffusiemodellen op randapparaten door de specifieke uitdagingen van generatieve token-pruning aan te pakken. Zijn betekenis ligt in:

• Schaalbaarheid: Het mogelijk maken van hoogresolutie synthese en videogeneratie op mobiel hardware door zowel reken- als geheugenvoetafdruk te verminderen.
• Architecturale Compatibiliteit: Het behouden van compatibiliteit met bestaande DiT-architecturen zonder dat structurele wijzigingen of complexe routingmodules nodig zijn.
• Efficiëntie-Kwaliteit Trade-off: Het demonstreren dat gestructureerde, coherentie-gestuurde oversprongen de inferentiekosten drastisch kunnen verminderen (tot 55% self-attention FLOPs) terwijl de perceptuele kwaliteit vergelijkbaar blijft met het volledige model.

De auteurs merken op dat hoewel de huidige aanpak zich richt op ruimtelijke redundantie, toekomstig werk gericht is op het uitbreiden van efficiëntie naar FFN/MLP-componenten en het opnemen van cross-frame coherentiedoelstellingen om de temporele consistentie van video verder te verbeteren.