Analyzing the Training Dynamics of Image Restoration Transformers: A Revisit to Layer Normalization

Deze paper analyseert de trainingsdynamiek van Transformers voor beeldherstel, identificeert de beperkingen van conventionele Layer Normalization, en stelt een nieuwe, aanpassingsvaste variant genaamd i-LN voor die de prestaties aanzienlijk verbetert.

De kern

De Titel: Waarom beelden herstellen met AI soms "uit elkaar valt" (en hoe we het oplossen).

Stel je voor dat je een oude, beschadigde foto probeert te repareren. Je gebruikt een slimme computer (een "Transformer") die als een super-artist werkt. Maar in dit paper ontdekten de onderzoekers dat deze kunstenaar een vreemd gedrag vertoont: hij begint steeds wilder te tekenen, totdat de kleuren zo fel worden dat ze de foto volledig overspoelen en de details verdwijnen.

Hier is wat er aan de hand is, vertaald in simpele taal:

1. Het Probleem: De "Perfekte" Regel die faalt

In de wereld van AI gebruiken ze vaak een regel genaamd LayerNorm (Laagnormalisatie). Je kunt dit zien als een strenge leraar die elke leerling (elk stukje van de foto) apart controleert.

• Hoe het werkt: De leraar zegt: "Jij, pixel A, je moet precies op een gemiddelde score zitten. En jij, pixel B, jij ook."
• Het probleem: In een foto zijn pixels echter geen losse leerlingen; ze zijn een team. Een wolkenrand moet samenwerken met de lucht eromheen. Als de leraar elke pixel apart forceert om "normaal" te zijn, vergeten ze hun relatie met elkaar.
• Het gevolg: De AI probeert de regels van de leraar te omzeilen. In plaats van rustig te tekenen, begint hij met "schreeuwen". De getallen in de computer worden gigantisch groot (miljoenen!) en de foto verliest al zijn nuance. Het is alsof de kunstenaar in paniek raakt en alles in één felwitte kleur schildert.

2. De Oplossing: i-LN (De Nieuwe Coach)

De onderzoekers bedachten een nieuwe methode, genaamd i-LN. In plaats van een strenge leraar die iedereen apart afstraft, is dit een slimme coach die het hele team begrijpt.

Ze hebben twee dingen veranderd:

• Verandering 1: Kijk naar het hele team, niet naar één speler.
  • De oude manier: De leraar keek alleen naar één pixel en zei: "Jij moet gemiddeld zijn."
  • De nieuwe manier (i-LN): De coach kijkt naar de hele foto tegelijk. Hij zegt: "Kijk hoe de wolken en de lucht samenwerken. Laten we de verhoudingen behouden." Hierdoor blijven de mooie details en randen van de foto intact.
• Verandering 2: Pas je aan aan de situatie.
  • De oude manier: De leraar gebruikte altijd dezelfde regels, of het nu een zonnige dag of een storm was.
  • De nieuwe manier (i-LN): De coach past de regels aan op basis van de foto. Als de foto donker is, maakt hij de regels iets lichter. Als de foto veel details heeft, houdt hij die vast. Dit zorgt ervoor dat de AI niet vastloopt in een starre wereld.

3. Wat levert dit op?

Door deze simpele aanpassing gebeurt er magie:

• Stabiliteit: De AI schreeuwt niet meer. De getallen blijven rustig en beheersbaar.
• Beter resultaat: De herstelde foto's zijn scherper, hebben minder ruis en zien er natuurlijker uit.
• Robuustheid: Zelfs als je de computer op een minder krachtige manier laat werken (bijvoorbeeld op een mobiele telefoon met minder geheugen), blijft de nieuwe methode werken. De oude methode crashte dan vaak.

De Grootte Analogie: Het Orkest

Stel je een orkest voor dat een mooi liedje speelt (het herstellen van een foto).

• Met de oude methode (LayerNorm) krijgt elke muzikant een strikte notitie: "Speel precies op 60 decibel, ongeacht wat de buren doen." De violist probeert dit, de fluitist probeert dit. Omdat ze niet naar elkaar luisteren, ontstaat er chaos. Sommigen schreeuwen zo hard dat hun instrumenten breken (de "miljoen-schaal" divergentie).
• Met de nieuwe methode (i-LN) kijkt de dirigent naar het hele orkest. Hij zegt: "Speel in harmonie met elkaar." Hij past het volume aan op basis van het stuk dat gespeeld wordt. Het resultaat? Een prachtig, helder liedje zonder dat iemand zijn instrument kapot speelt.

Kortom: De onderzoekers hebben ontdekt dat de standaard-regels voor AI-fotoherstel te star waren. Door de regels iets flexibeler en meer "teamgericht" te maken, kunnen AI-modellen nu veel betere en stabielere foto's herstellen.

Technische samenvatting

Probleemstelling: Abnormale Trainingsdynamiek in IR-Transformers

De auteurs identificeren een kritiek, maar tot nu toe over het hoofd gezien probleem bij het trainen van Transformers voor Beeldherstel (Image Restoration - IR), zoals super-resolutie, ontruisen en regenverwijdering. Hoewel Transformer-architecturen met Layer Normalization (LN) de standaard zijn geworden, vertonen ze tijdens het trainen een ernstige instabiliteit:

1. Divergentie van Feature-magnitudes: De grootte van de interne features (activaties) divergeert dramatisch en bereikt schalen van miljoenen (bijv. $10^7$).
2. Ineenstorting van Entropie: De kanaal-gerichte entropie daalt scherp, wat aangeeft dat de activaties extreem geconcentreerd raken in specifieke kanalen in plaats van gelijkmatig verdeeld te zijn.

De auteurs stellen de hypothese dat dit fenomeen ontstaat doordat het netwerk probeert de beperkingen van conventionele per-token LayerNorm te omzeilen. Deze beperkingen staan haaks op de unieke eisen van IR-taken:

• Verlies van ruimtelijke correlaties: Conventionele LN normaliseert elk token (pixel) onafhankelijk, waardoor de relatieve verschillen tussen pixels (inter-pixel structuur) worden verstoord. Dit is funest voor het behoud van lage-niveau beelddetails.
• Gebrek aan input-afhankelijke statistieken: LN mapt features naar een uniforme ruimte, waardoor de specifieke statistische variabiliteit van de invoer (zoals de grootte van details of ruis) verloren gaat.

Methodologie: i-LN (Image Restoration Transformer Tailored Layer Normalization)

Om deze misalignementen op te lossen, stellen de auteurs i-LN voor. Dit is een eenvoudige "drop-in" vervanging voor conventionele LayerNorm die bestaat uit twee kerncomponenten:

1. Ruimtelijk Holistische Normalisatie (LN):*

   • In plaats van normalisatie per token (alleen over de kanaal-dimensie), berekent i-LN de gemiddelde waarde ($\mu$) en variantie ($\sigma^2$) over alle tokens en kanalen tegelijkertijd (de volledige spatio-kanaal dimensie).
   • Theoretisch voordeel: Dit behoudt de inter-pixel structuur. Wiskundig gezien is deze transformatie een homothety (een uniforme schaling en verschuiving), wat betekent dat hoekverhoudingen en afstanden tussen tokens behouden blijven, in tegenstelling tot conventionele LN.

2. Input-Adaptieve Rescaling:

   • Na de normalisatie worden de features expliciet opnieuw geschaald op basis van de berekende standaardafwijking ($\sigma$) van de invoer.
   • Doel: Dit herstelt de "globale schaal" die verloren ging door de normalisatie en zorgt ervoor dat het netwerk de input-afhankelijke statistieken (zoals de dynamische bereik van de features) kan behouden. Dit voorkomt dat het netwerk gedwongen wordt om extreme waarden te genereren om binnen de beperkingen van de normalisatie te blijven.

De formule voor een blok $B$ (Attention of FFN) met i-LN wordt:
$$B(x; f, \text{i-LN}) = x + \sqrt{\sigma^2 + \epsilon} \cdot f(\text{LN}^*(x))$$
Waarbij $\text{LN}^*$ de holistische normalisatie is en de term $\sqrt{\sigma^2 + \epsilon}$ de input-adaptieve rescaling vertegenwoordigt.

Belangrijkste Bijdragen

• Analyse van Trainingsdynamiek: Het paper biedt een diepgaande analyse van waarom IR-Transformers instabiel trainen met standaard LN, gekoppeld aan het verlies van ruimtelijke correlaties en input-statistieken.
• Theoretische Onderbouwing: De auteurs bewijzen wiskundig dat conventionele per-token LN de inter-pixel structuur niet behoudt, terwijl hun voorgestelde holistische normalisatie (LN*) dit wel doet (tot op een globale schaal).
• Eenvoudige Implementatie: i-LN vereist geen complexe architecturale wijzigingen en fungeert als een directe vervanging voor bestaande LN-lagen in modellen zoals HAT, SwinIR en DRCT.
• Robuustheid: Het paper toont aan dat i-LN niet alleen de prestaties verbetert, maar ook de stabiliteit verhoogt bij verminderde precisie (inferentie in FP16 of INT8), waar standaard LN vaak faalt door overloop (infinity values).

Resultaten

Uitgebreide experimenten op diverse benchmarks (Set5, Set14, Urban100, etc.) en taken (Super-Resolution, Denoising, Deraining, JPEG-artefactverwijdering) tonen aan:

• Prestatieverbetering: i-LN levert consistent hogere PSNR en SSIM scores op vergeleken met conventionele LN, LayerScale, RMSNorm en andere normalisatiemethoden. Bijvoorbeeld, bij 4x Super-Resolution met HAT1 stijgt de PSNR op Urban100 van 26.55 naar 27.17.
• Stabilisatie van Features: In tegenstelling tot de divergentie naar miljoenen bij standaard LN, blijven de feature-magnitudes met i-LN goed gebonden (rond de verwachte verdeling van een normale verdeling). De kanaal-entropie blijft hoog en stabiel, wat wijst op een gezonde verdeling van activaties.
• Verbeterde Ruimtelijke Modellering: Visualisaties van Relative Position Embeddings (RPE) tonen aan dat i-LN leidt tot gestructureerdere en minder ruisachtige patronen, wat aangeeft dat het netwerk ruimtelijke relaties beter leert modelleren.
• Robuustheid bij Lage Precisie: Bij inferentie met half-precision (FP16) of kwantisatie (INT8) faalt het model met standaard LN vaak (generatie van oneindige waarden/artefacten), terwijl i-LN stabiel blijft en bijna geen kwaliteitsverlies vertoont.

Significantie

Deze studie is significant omdat het een fundamenteel probleem in de huidige staat van de kunst voor Beeldherstel-Transformers aanpakt. Het toont aan dat de succesvolle toepassing van Transformers in andere domeinen (zoals NLP of classificatie) niet zomaar kan worden overgenomen voor lage-niveau visietaken zonder de normalisatiestrategie aan te passen.

Door de normalisatie aan te passen aan de specifieke eisen van IR (behoud van ruimtelijke correlaties en input-statistieken), biedt i-LN een eenvoudige maar krachtige oplossing die:

1. De trainingsstabiliteit drastisch verbetert.
2. De uiteindelijke reconstructiekwaliteit verhoogt.
3. De toepasbaarheid van IR-Transformers op resource-beperkte apparaten (edge devices) mogelijk maakt door betere stabiliteit bij lage precisie.

Kortom, i-LN positioneert zich als een nieuwe standaard voor normalisatie in Image Restoration Transformers.