SALAD: Achieve High-Sparsity Attention via Efficient Linear Attention Tuning for Video Diffusion Transformer

Probleemstelling

Diffusion Transformers (DiT) hebben zich bewezen als krachtige modellen voor video-generatie, maar ze kampen met een significant efficiëntieprobleem bij het modelleren van lange sequenties (zoals hoge resolutie en lange duur video's). De standaard "full attention" mechanisme heeft een kwadratische complexiteit ($O(N^2)$), wat leidt tot enorme latentie en rekenkosten.

Om dit op te lossen, zijn er diverse sparse attention (verspreide aandacht) mechanismen ontwikkeld die de interacties tussen tokens beperken. Echter, deze benaderingen hebben twee grote beperkingen:

1. Training-vrije methoden: Deze kunnen slechts matige mate van verdunning (sparsity, vaak 40-60%) bereiken, wat leidt tot beperkte versnelling.
2. Training-gebaseerde methoden: Hoewel deze zeer hoge verdunning (80-95%) kunnen bereiken, vereisen ze enorme hoeveelheden data en rekenkracht voor fine-tuning. Bovendien faalt de gebruikelijke LoRA (Low-Rank Adaptation) vaak om de kwaliteitsdaling volledig te compenseren bij extreme verdunning. Het resultaat is vaak een verlies aan lange-afstand afhankelijkheden, wat zich manifesteert als artefacten (bijv. dubbele objecten of inconsistenties in tekst en tijd).

Methodologie: SALAD

De auteurs stellen SALAD voor, een efficiënt aandachtsmodule die een lichtgewicht lineaire attentie-tak parallel introduceert aan de bestaande sparse attention. Het doel is om de informatieverlies van de sparse attention te compenseren zonder de rekenkosten van full attention te herstellen.

De kerncomponenten van SALAD zijn:

1. Parallelle Architectuur:

   • De sparse attention (bijv. Sliding-Window of Top-K) blijft de hoofdrol spelen voor het modelleren van de meeste sequentie-informatie.
   • Een lineaire attention-tak (met $O(N)$ complexiteit) wordt parallel toegevoegd. Deze tak kan globale token-interacties modelleren en vult de informatie aan die door de hoge verdunning is verloren.
   • Beide takken delen dezelfde query, key en value projectiematrices (gefinetuned met LoRA), wat de parameter-efficiëntie maximaliseert.

2. Multi-level Static-Dynamic Scaling Strategy:
   De auteurs merken op dat een simpele sommatie van de twee takken suboptimaal is vanwege een "rank-ongelijkheid": de lineaire tak heeft een veel lagere rang (rank) dan de sparse tak en kan de informatieve structuren van de sparse tak onbedoeld onderdrukken. Om dit op te lossen, introduceren ze een schaalstrategie:

   • Layer-wise Static Scaling (LWSS): Een zero-geinitialiseerde projectielag die per transformer-laag de bijdrage van de lineaire tak aanpast. Dit zorgt voor een stabiele start van de training.
   • Input-Timestep-Aware Dynamic Scaling (ITADS): Een dynamische module die een globale schalingsfactor genereert op basis van de invoer en de denoising-tijdstap ($t$). Deze factor (beperkt tussen 0 en 1 via een sigmoid functie) reguleert de invloed van de lineaire tak adaptief.
   • Doel: De lineaire tak fungeert als een gecontroleerde, aanvullende "low-rank" component die cruciale globale informatie injecteert zonder de dominante sparse output te overweldigen.

3. Post-tuning Branch Dropping:
   Omdat de schalingsfactoren variëren per tijdstap en input, kunnen takken met een zeer lage schalingswaarde tijdens de inferentie worden verwijderd (gepruned) zonder kwaliteitsverlies, wat de efficiëntie verder verhoogt.

Belangrijkste Bijdragen

1. SALAD Architectuur: Een efficiënte hybridisatie van sparse en lineaire attention die hoge verdunning mogelijk maakt met minimale extra parameters (slechts ~4,99% extra t.o.v. het voorgeprogrammeerde model).
2. Schaalstrategie: De ontwikkeling van de Multi-level Static-Dynamic Scaling, die essentieel is om de rank-ongelijkheid tussen de takken op te lossen en te voorkomen dat de lineaire tak de sparse output domineert.
3. Efficiënte Finetuning: SALAD vereist slechts 2.000 video-voorbeelden en minder dan 1.600 trainingsstappen (ongeveer 20-30 GPU-uren), wat aanzienlijk minder is dan concurrerende methoden die duizenden tot miljoenen clips nodig hebben.

Resultaten

De auteurs evalueren SALAD op het Wan2.1 model (1.3B en 14B parameters) met verschillende token-lengtes (30k en 100k tokens).

• Verdunning en Snelheid: SALAD bereikt tot 90% verdunning (sparsity) en een inferentieversnelling van 1,52x tot 2,03x ten opzichte van full attention.
• Kwaliteit: In tegenstelling tot LoRA-only fine-tuning of training-vrije methoden, behoudt SALAD een generatiekwaliteit die vergelijkbaar is met, en soms zelfs beter is dan, de full-attention baseline. Dit wordt gemeten via VBench-metrics (Subject Consistency, Background Consistency, Image Quality, Text Consistency).
• Vergelijking:
  • Bij 90% verdunning presteert SALAD aanzienlijk beter dan LoRA-gefine-tuned sparse attention, die vaak artefacten vertoont.
  • SALAD overtreft training-vrije methoden (zoals SVG2 en PARO) in zowel snelheid als kwaliteit bij hoge verdunning.
  • De methode is schaalbaar naar grotere modellen (14B parameters) met vergelijkbare resultaten.

Betekenis en Impact

SALAD biedt een praktische en kosteneffectieve oplossing voor het probleem van lange sequenties in video-generatie. Door de afhankelijkheid van enorme datasets en rekenkracht voor het trainen van sparse attention te elimineren, maakt het high-sparsity video-generatie toegankelijker. De methode bewijst dat een zorgvuldig gebalanceerde combinatie van sparse en lineaire attention, gereguleerd door dynamische schaling, de kwaliteitsdaling bij extreme verdunning kan overbruggen. Dit opent de deur voor snellere, goedkopere en hogere kwaliteit video-generatie in toekomstige Diffusion Transformer-toepassingen.