Technische Samenvatting: NanoFLUX

Probleemstelling

Grootschalige Text-to-Image (T2I) diffusiemodellen, zoals de FLUX-familie (opschalend tot 17B parameters) en Qwen-Image (20B parameters), hebben een staat van de kunst bereikt op het gebied van visuele kwaliteit. Echter, hun enorme omvang creëert een significante barrière voor on-device implementatie, wat kostbare GPU-infrastructuur of cloud-gebaseerde inferentie vereist. Bestaande on-device oplossingen vertrouwen vaak op kleinere modellen die niet over een vergelijkbare kwaliteit beschikken, of maken gebruik van compressietechnieken (pruning, kwantisatie, reductie van stappen) die de architecturale redundantie van grootschalige Diffusion Transformers (DiTs) en hun bijbehorende tekstencoders niet volledig aanpakken. De uitdaging ligt in het comprimeren van deze massieve modellen voor mobiele apparaten terwijl de generatiekwaliteit behouden blijft en de inferentie-latentie wordt geminimaliseerd.

Methodologie: NanoFLUX

De auteurs introduceren NanoFLUX, een 2.4B-parameter T2I-model dat is gedestilleerd van de 17B FLUX.1-Schnell teacher. De compressie wordt bereikt via een progressieve, door distillatie gedreven pipeline die zich richt op de twee zwaarste componenten: de 12B Diffusion Transformer (DiT) en de 5B T5-XXL tekstencoder.

1. Transformer Compressie (12B $\to$ 2B)

De auteurs hanteren een vierfasige progressieve compressiestrategie om de omvang van de DiT te verminderen:

• Fase C1: Attention Head Pruning (12B $\to$ 5B):

  • Analyse: Singular Value Decomposition (SVD) op de outputs van de attention heads onthulde dat niet alle heads onafhankelijk zijn; veel bevatten redundante informatie.
  • Actie: Het aantal attention heads ($H$) werd uniform verminderd van 24 naar 16.
  • Distillatie: Er werd een feature-level loss geïntroduceerd om ervoor te zorgen dat de resterende heads collectief de informatie van de geprunede heads behouden. De student werd geïnitialiseerd met de eerste 16 heads van de teacher.

• Fase C2: Head Dimension Reduction (5B $\to$ 3B):

  • Analyse: SVD werd opnieuw toegepast op de features van het 5B model.
  • Actie: De head dimensie ($d_H$) werd verminderd van 128 naar 96, waarbij de belangrijkste singuliere componenten werden behouden. Dit leverde een 3B parameter model op.

• Fase C3: Depth Pruning via Block Merging (3B $\to$ 2.5B):

  • Analyse: De input-output cosine similarity werd berekend voor de transformer blocks. Single-Stream (SS) blocks vertoonden een hoge gelijkenis (>0.85), wat duidt op functionele redundantie.
  • Actie: In plaats van deze blocks te verwijderen, hebben de auteurs een keten van 15 redundante SS blocks samengevoegd tot één enkele block door hun parameters te middelen. Dit verminderde de transformer diepte van 38 naar 24 blocks.

• Fase C4: Static Layer Normalization (2.5B $\to$ 1.8B):

  • Analyse: De door het model voorspelde Adaptive Layer Normalization (AdaLN) coëfficiënten vertoonden een lage variantie over samples heen, wat suggereert dat de dynamische conditionering van CLIP projecties redundant was.
  • Actie: Dynamische AdaLN coëfficiënten werden vervangen door vooraf berekende gemiddelden (Static-LN). Dit verwijderde ongeveer 0.7B parameters. Het model werd kortstondig gefinetuned met bevroren coëfficiënten om de kwaliteit te herstellen.

2. Progressieve Token Downsampling

Om de kwadratische complexiteit van self-attention aan te pakken zonder detail op te offeren:

• Mechanisme: Een ResNet-gebaseerd downsampling module ($D$) en upsampling module ($U$) werden in de transformer ingevoegd.
• Hybride Strategie: Het model werkt op een hybride resolutie schema. Vroege diffusie tijdstappen (die de globale structuur vastleggen) verwerken gedownsamplede tokens (reductie van het aantal tokens met 4x), terwijl latere tijdstappen (die details verfijnen) volledige resolutie tokens verwerken.
• Training: Een "Progressive Token Downsampling" (PTD) strategie werd gebruikt om de training te stabiliseren. Modules werden ingevoegd en incrementeel getraind, waarbij het downsampling punt iteratief eerder in het netwerk werd geplaatst.
• Resultaat: 23 van de 43 transformer blocks opereren op gedownsamplede tokens, wat de latentie aanzienlijk vermindert terwijl high-resolution verwerking waar nodig behouden blijft.

3. Tekstencoder Distillatie (5B $\to$ 330M)

• Doel: De 5B T5-XXL encoder distilleren naar een 330M T5-Large model.
• Methode: Een tweefasig distillatieproces:
  1. Embedding Matching: Een MLP wordt aan de student toegevoegd om de outputdimensie van de teacher te matchen, waardoor de MSE tussen prompt embeddings wordt geminimaliseerd.
  2. Block-wise Distillatie: De student wordt getraind om de prompt hidden states van de teacher op specifieke transformer layers te matchen. Cruciaal is dat de methode gebruikmaakt van visuele signalen uit de vroege lagen van de denoiser (via joint attention) om de tekstencoder te begeleiden, wat foutaccumulatie tegengaat.
  3. Stabiliteit: Gradiënten worden losgekoppeld voor tijdstappen met hoge ruis ($t < \hat{t}$) om onstabiele supervisie te voorkomen, waarbij de focus van de distillatie ligt op latere tijdstappen en de eerste drie transformer layers.

Belangrijkste Resultaten

• Modelgrootte: Het uiteindelijke NanoFLUX model is 2.4B parameters (2B DiT + 330M Text Encoder), een 7$\times$ reductie ten opzichte van de 17B FLUX.1-Schnell teacher.
• Prestaties:
  • Op mobiele apparaten genereert NanoFLUX 512 $\times$ 512 afbeeldingen in ongeveer 2.5 seconden.
  • Kwantitatieve benchmarks (One-IG, DPG, GenEval, HPDv3) laten zien dat NanoFLUX prestaties behoudt die vergelijkbaar zijn met het teacher model, met slechts een minimale degradatie in specifieke metrieken (bijv. HPSv3 daalt van 11.45 naar 10.41).
  • Kwalitatieve resultaten geven aan dat het model een beeldkwaliteit en prompt-getrouwheid behoudt die vergelijkbaar zijn met het originele 17B model.
• Efficiëntie: De progressieve token downsampling strategie balanceert succesvol tussen inferentie-latentie en generatiekwaliteit, waardoor het model efficiënt kan draaien op hardware met beperkte middelen.

Betekenis en Claims

Het artikel claimt dat NanoFLUX het eerste framework is dat uitgebreid onderzoek doet naar grootschalige modelcompressie voor T2I-generatie, specifiek gericht op modellen zoals FLUX. De betekenis ligt in:

1. Haalbaarheid van On-Device Generatie: Het aantonen dat hoogwaardige T2I-generatie haalbaar is op mobiele apparaten zonder afhankelijk te zijn van cloud-infrastructuur.
2. Architecturaal Inzicht: Het bieden van een praktisch pad voor het opschalen (naar beneden) van grote generatieve modellen door specifieke architecturale redundanties (attention heads, block depth, normalisatie dynamiek) te identificeren en te exploiteren, in plaats van enkel te vertrouwen op generieke kwantisatie of stapreductie.
3. Generaliseerbaarheid: Het bieden van inzichten en een progressieve distillatie pipeline die mogelijk generaliseerbaar zijn naar toekomstige grootschalige generatieve modellen.

De auteurs positioneren dit werk als een stap naar het verbreden van de toegankelijkheid van hoogwaardige generatieve AI voor gebruikers met beperkte computationele middelen.