Technische Samenvatting: Registers zijn van Belang voor Diffusietransformatoren in Pixelspace

Probleemstelling

Vision Transformers (ViTs) staan bekend om het vertonen van patch-token uitbijters met hoge normen, wat de kwaliteit van feature maps verslechtert en artefacten in attentiekarten veroorzaakt. Dit probleem wordt bij ViTs effectief verholpen door "register tokens" in te voeren, die fungeren als zinken voor deze uitbijters. Aangezien Diffusion Transformers (DiTs) steeds vaker transformer-architecturen adopteren en verschuiven naar training direct in pixelspace (weg van latent space), lijken ze structureel meer op ViTs. Dit roept twee kritieke vragen op:

1. Erven DiTs de uitbijters met hoge normen van patch-tokens over die bij ViTs worden waargenomen?
2. Zijn register tokens nuttig voor DiTs, met name in trainingsregimes in pixelspace?

Hoewel attentiezinken en het gedrag van speciale tokens zijn onderzocht bij Large Language Models (LLMs) en diffusietalenmodellen, blijven hun aanwezigheid en functionele rol in diffusietransformatoren voor afbeeldingen onderbelicht.

Methodologie

De auteurs voeren een systematische analyse uit van register tokens over verschillende diffusiemodelruimten (pixelspace, VAE-space en RAE-space) en vergelijken deze met standaard ViTs (specifiek DINOv2).

1. Empirische Analyse van Uitbijters: De studie onderzoekt token-voor-token feature-normen en attentiekarten in DiTs (specifiek pDiTs getraind met flow matching op ImageNet) versus ViTs. Ze analyseren of patch-tokens uitbijters met hoge normen vertonen bij afwezigheid van registers.
2. Integratie van Register Tokens: Register tokens worden geïmplementeerd als extra leerbare tokens die worden toegevoegd aan de patch-tokensequentie, volgens de aanpak die bij ViTs wordt gebruikt. Deze tokens worden uitgesloten van het trainingsdoel (verliesfunctie).
3. Analyse van Interne Representaties: De auteurs analyseren tussentijdse feature maps met behulp van Total Variation (TV) om ruimtelijke gladheid te meten en lineair proppen om de semantische inhoud van register tokens te evalueren. Ze onderzoeken hoe registers feature-normen en de coherentie van representaties op verschillende diffusietijdstappen (ruisniveaus) beïnvloeden.
4. Architecturale Verkenning: Gemotiveerd door de onderscheiden rollen van registers en patch-tokens, stelt het artikel een parameter-efficiënte dual-stream architectuur voor. Dit ontwerp verwerkt register- en patch-tokens via gespecialiseerde paden (specifiek het ontkoppelen van RMSNorm, MLP en adaLN-lagen) terwijl het attentiemechanismen en de meeste transformer-parameters deelt.
5. Ablatiestudies: De studie ablateert het aantal registers, de diepte waarop ze worden geïntroduceerd, en de specifieke componenten die gedualiseerd moeten worden om optimale configuraties te bepalen.

Belangrijkste Bevindingen en Bijdragen

1. DiTs verschillen fundamenteel van ViTs wat betreft Uitbijters

In tegenstelling tot ViTs vertonen standaard DiTs (zowel in latent als pixelspace) geen uitbijters met hoge normen van patch-tokens. Bij ViTs zijn uitbijters gelokaliseerd tot specifieke tokens (vaak in gebieden met weinig informatie), terwijl bij DiTs de normen van patch-tokens relatief uniform blijven en attentiekarten de uitbijters met weinig informatie missen die bij ViTs worden gezien.

2. Registers verbeteren Pixelspace DiTs ondanks Afwezigheid van Bestaande Uitbijters

Verrassend genoeg verbetert het toevoegen van register tokens aan DiTs significant de convergentie en de generatiekwaliteit, met name bij pixelspace-modellen.

• Mechanisme: In tegenstelling tot ViTs, waar registers bestaande uitbijters absorberen, creëren registers in DiTs zelf uitbijters met hoge normen.
• Effect: De aanwezigheid van registers verlaagt consistent de feature-normen van alle patch-tokens. Dit leidt tot schonere, gladdere tussentijdse feature maps, vooral bij hoge ruisniveaus (vroege diffusietijdstappen).
• Dubbele Rol: Register tokens specialiseren zich in twee onderscheiden rollen:
  • Norm-zinken: Tokens met extreem hoge normen die grootte absorberen van patch-tokens.
  • Dragers van Globale Informatie: Tokens met gematigde normen die diverse semantische informatie over de afbeelding coderen (bijvoorbeeld voorgrondobjecten, achtergrondcontext), geverifieerd via lineair proppen en visualisatie van attentiekarten.

3. Ruimte-specifieke Effectiviteit

De voordelen van register tokens zijn sterk afhankelijk van de trainingsruimte:

• Pixelspace: Registers leveren de meest significante verbeteringen in generatiekwaliteit (FID). Pixelspace-modellen produceren inherent ruisigere tussentijdse representaties met hogere feature-normen, waardoor het regularisatie-effect van registers cruciaal is.
• Latent Space (VAE/RAE): Registers bieden gematigde winsten of, in het geval van op RAE gebaseerde modellen, degraderen ze de prestaties lichtjes. Latent ruimtes zijn meer gestructureerd en lager dimensionaal, wat de behoefte aan de sterke regularisatie die registers bieden, vermindert.

4. Impliciete Registers in Bestaande Architecturen

De auteurs observeren dat recente pixelspace DiT-architecturen (bijvoorbeeld JiT) die in-context class conditioning gebruiken (het toevoegen van class-embeddings aan de tokensequentie), impliciet register-achtige mechanismen incorporeren. Deze in-context tokens vertonen vergelijkbaar gedrag als expliciete registers: sommigen fungeren als norm-zinken, terwijl anderen globale semantische informatie coderen. Dit suggereert dat prestatiewinsten in dergelijke modellen grotendeels kunnen voortkomen uit dit register-achtige gedrag in plaats van de class-informatie zelf.

5. Parameter-efficiënte Dual-Stream Architectuur

Om deze inzichten te benutten, stellen de auteurs een compacte dual-stream architectuur voor die de verwerking van register tokens specialiseert.

• Ontwerp: Registers worden alleen geïntroduceerd in diepere lagen (beginnend bij blok 4), waar semantische structuur wordt vastgesteld. De architectuur gebruikt gedeelde attention maar ontkoppelt RMSNorm, MLP en adaLN voor registers en patches.
• Efficiëntie: Deze aanpak vergroot de modelgrootte met slechts ~14% (vergeleken met ~77% voor naïeve volledige duplicatie) terwijl het vergelijkbare inferentiekosten (GFLOPs) behoudt.
• Prestaties: Deze architectuur verbetert consistent de generatiekwaliteit over modelgroottes (B, L, H) en resoluties (256x256, 512x512) in vergelijking met single-stream baselines.

Resultaten

• Generatiekwaliteit: Op ImageNet 256x256 presteren pDiTs met registers consistent beter dan die zonder. Bijvoorbeeld, een pDiT-B/16-model verbeterde zijn FID van 7,39 naar 5,30 na 200 epoches.
• Feature-gladheid: Modellen met registers vertonen lagere Total Variation (TV)-ratio's in tussentijdse features bij hoge ruisniveaus, wat wijst op gladdere en coherenter representaties.
• Ablatie: Registers zijn het meest effectief wanneer ze worden geïntroduceerd na de eerste 4 lagen en wanneer een groter aantal tokens wordt gebruikt (bijvoorbeeld 32 registers voor pDiT-B, vergeleken met 4 voor ViTs).

Betekenis en Beweringen

Het artikel beweert dat register tokens een kritiek, maar eerder over het hoofd gezien, onderdeel zijn voor hoogwaardige diffusiegeneratie in pixelspace. De betekenis ligt in:

1. Hertekening van de Rol van Registers: Aantonen dat registers niet slechts een oplossing zijn voor ViT-specifieke uitbijters, maar een bredere regularisatiefunctie dienen in diffusiemodellen door feature-normen te beheren en ruimtelijke coherentie te versterken.
2. Verklaring van Recente Successen: Het bieden van een rationale voor de sterke empirische prestaties van recente pixelspace DiTs die in-context conditioning gebruiken, waarbij een deel van hun succes wordt toegeschreven aan impliciet register-gedrag.
3. Efficiënt Architecturaal Ontwerp: Het aanbieden van een parameter-efficiënt dual-stream ontwerp dat inspreekt op de onderscheiden rollen van registers en patches, en state-of-the-art generatiekwaliteit bereikt met verwaarloosbare runtime-overhead.

De auteurs concluderen dat hoewel hun ontwerpen de belofte van token-specialisatie vastleggen, het volledige potentieel van het gebruik van registers in diffusietransformatoren in pixelspace een open vraag blijft voor toekomstig onderzoek.