Technische Samenvatting: Outlier-tokens in Diffusie-Transformers temmen

Probleemstelling

Het artikel onderzoekt het fenomeen van outlier-tokens binnen Diffusie-Transformers (DiTs) die worden gebruikt voor beeldgeneratie. Hoewel eerdere werken (bijvoorbeeld Darcet et al.) hebben vastgesteld dat Vision Transformers (ViTs) in herkenningsopdrachten vaak een kleine subset van tokens produceren met abnormaal grote feature-normen die een disproportionele hoeveelheid aandacht absorberen terwijl ze beperkte lokale informatie dragen, bleef de rol van deze tokens in generatieve modellen ononderzocht.

De auteurs identificeren twee kritieke gaten:

1. Alomtegenwoordigheid in generatieve pijplijnen: Outlier-tokens komen niet alleen voor in de vooraf getrainde ViT-encoders (tokenizers) van moderne Representation Autoencoder (RAE)-pijplijnen, maar ook intern binnen de Diffusie-Transformers (denoisers) zelf.
2. Aard van het defect: De degradatie veroorzaakt door deze tokens is niet louter het gevolg van een paar extreme verlieswaarden. In plaats daarvan hypotheseren de auteurs dat outlier-tokens symptomen zijn van gecorrumpeerde lokale patch-semantiek. Het simpelweg maskeren van tokens met hoge normen tijdens het trainingstraject verbetert de prestaties niet, wat suggereert dat het probleem ligt in het verlies van ruimtelijke structuur en lokale informatie, en niet slechts in numerieke instabiliteit.

Methodologie: Dual-Stage Registers (DSR)

Om deze problemen aan te pakken, stellen de auteurs Dual-Stage Registers (DSR) voor, een uniforme interventiestrategie die register-tokens toepast op beide componenten van de RAE-DiT-pijplijn: de representatie-encoder en de diffusiegenerator.

1. Registers in Vision-encoders

De auteurs adresseren outlier-tokens in de encoderfase (de tokenizer) met twee benaderingen, afhankelijk van de beschikbaarheid van vooraf getrainde registers:

• Getrainde Registers: Voor encoders zoals DINOv2 die registers ondersteunen, gebruiken de auteurs vooraf getrainde register-tokens om artefacten met hoge normen in de feature-kaarten van de encoder te onderdrukken.
• Recursieve Test-Time Registers (TTR): Voor encoders zonder ingebouwde registers (bijvoorbeeld SigLIP2) voegen de auteurs tijdens de inferentie extra niet-getrainde tokens toe aan de invoersequentie. In gevallen waarin een tweede fase van outliers wordt gedetecteerd in de resulterende representatie, passen ze TTR recursief toe om de output te stabiliseren.

2. Registers in Diffusie-Transformers

De auteurs observeren dat zelfs met een gestabiliseerde encoder, de Diffusie-Transformer zelf interne outlier-tokens ontwikkelt, voornamelijk in intermediaire lagen (in tegenstelling tot standaard ViTs waar outliers vaak in de laatste lagen verschijnen).

• Diffusie-Registers: Ze introduceren een kleine set leerbare register-tokens direct in de Diffusie-Transformer. Deze tokens worden samen met de generator getraind om attentie-zuigers te absorberen en de representatieruimte te stabiliseren.
• Inferentieprotocol: Tijdens de inferentie worden de outputs van alle register-tokens (zowel van de encoder als van het diffusiemodel) verworpen, waarbij alleen de beeld-token-outputs worden behouden.

Belangrijkste Bijdragen

• Ontdekking van Interne Outliers: Het artikel toont aan dat outlier-tokens wijdverbreid zijn in DiTs, en voorkomen in zowel de encoder als de denoiser, met een duidelijk patroon van concentratie in de intermediaire lagen van het diffusiemodel.
• Afwijzing van de "Extreme Waarde"-hypothese: Door middel van verliesmaskeringsexperimenten tonen de auteurs aan dat het simpelweg uitsluiten van tokens met hoge normen uit het trainingsverlies de kwaliteit van de generatie niet verbetert. Dit leidt tot de conclusie dat de kern van het probleem de corruptie van lokale patch-semantiek is, en niet louter de aanwezigheid van extreme waarden.
• Dual-Stage Interventie: Het voorstellen van DSR biedt een eenvoudige, lichtgewicht mechanisme om de gehele generatiepijplijn te stabiliseren door outliers aan de bron (encoder) en tijdens het generatieproces (denoiser) aan te pakken.
• Generaliseerbaarheid: De methode blijkt effectief te zijn in verschillende invoerruimten (pixelruimte, VAE-latenten en representatieruimten) en diverse architecturen (SiT, JiT, RAE-gebaseerde ontwerpen).

Experimentele Resultaten

De auteurs evalueren DSR op ImageNet-1K class-conditional generatie (256×256) en grootschalige tekst-naar-beeld generatie.

• ImageNet-1K Prestaties:
  • Met SigLIP2-B als encoder verlaagt DSR de FID van 5,89 naar 4,58 (met 800 epochs) en verbetert de GenEval-scores voor tekst-naar-beeld taken van 0,426 naar 0,466.
  • In combinatie met DINOv2-B-encoders verbetert de FID van 4,16 naar 3,95.
  • DSR versnelt de convergentie en bereikt vergelijkbare kwaliteit met 4× minder trainings-epochs in vergelijking met de baseline.
• Schaalbaarheid: De verbeteringen zijn consistent over modelgroottes (DiT-S, B, XL), met slechts een marginale toename in rekentijd (ongeveer 10% toename in GFLOPs).
• Ablatiestudies:
  • Invoerdiepte: Registers zijn het meest effectief wanneer ze worden geïntroduceerd in de vroege tot middelste lagen (bijvoorbeeld beginnend bij blok 8). Het te vroeg of te laat introduceren ervan levert afnemende meeropbrengst op.
  • Aantal Registers: Een gematigd aantal registers (bijvoorbeeld 36) levert optimale prestaties op; te weinig (1) of te veel (100) verslechtert de prestaties.
• Vergelijking met Alternatieven: DSR presteert beter dan "in-context conditioning" (die invoer-afhankelijke tokens gebruikt) naarmate de invoerrepresentatie semantischer wordt, wat suggereert dat registers een robuuster mechanisme bieden voor globale capaciteit.

Betekenis en Claims

Het artikel beweert dat controle van outlier-tokens een kritiek, maar tot nu toe over het hoofd gezien, onderdeel is in het bouwen van robuuste Transformer-gebaseerde diffusiepijplijnen. Door aan te tonen dat deze artefacten lokale patch-semantiek degraderen en dat register-gebaseerde interventies deze effectief kunnen mitigeren, positioneren de auteurs DSR als een fundamentele verbetering voor zowel encoder- als generatorstabiliteit.

Het werk benadrukt dat de flexibiliteit van Transformers, hoewel gunstig voor globale interacties, specifieke artefacten (outlier-tokens) introduceert die de ruimtelijke fideliteit in generatie belemmeren. De voorgestelde oplossing wordt gepresenteerd als een eenvoudige, plug-and-play interventie die consistent de kwaliteit van de generatie verbetert in diverse settings, zonder complexe architecturale wijzigingen of het opnieuw trainen van de hele pijplijn vanaf nul te vereisen.