Technische Samenvatting: Normalisatie-Equivariantie voor Willekeurige Ruggegraten

Probleemstelling

Bij beeld-naar-beeld voorspellingstaken, met name bij blind beeldontruisen, lijden neurale netwerken vaak onder distributieveranderingen wanneer het ruisniveau op het testmoment afwijkt van het ruisniveau tijdens het trainen. Hoewel Schaalequivariantie (SE) ($f(ay) = af(y)$) en de strikt sterkere Normalisatieequivariantie (NE) ($f(ay + b\mathbf{1}) = af(y) + b\mathbf{1}$) bekend staan om hun verbeterde robuustheid tegen dergelijke verschuivingen, vormt het afdwingen van deze eigenschappen in moderne deep learning-architecturen aanzienlijke uitdagingen.

Bestaande methoden dwingen NE af door interne netwerklagen te beperken tot specifieke NE-compatibele families (zoals affiene-beperkte convoluties, sorteerende niet-lineariteiten en affiene residuen). Deze aanpak, hoewel theoretisch onderbouwd, kent twee grote nadelen:

1. Architecturale Incompatibiliteit: Het sluit standaardcomponenten uit die alomtegenwoordig zijn in moderne ruggegraten, zoals attentiemechanismen (softmax) en LayerNorm, die niet commuteren met affiene transformaties.
2. Runtime-Overhead: Het afdwingen van equivariantie via interne laagbeperkingen introduceert aanzienlijke rekentijd, waarbij sommige architecturale basismodellen tot een 1,6× vertraging vertonen ten opzichte van standaardmodellen.

Methodologie: De Normalisatieequivariante Wrapper (WNE)

De auteurs stellen een methode voor die NE afdwingt op het invoer-uitvoerniveau in plaats van op het niveau van de interne lagen. Deze aanpak berust op een volledige karakterisering van de NE-functieklasse.

Theoretische Karakterisering

Het artikel stelt Karakterisering 1 vast, waarin wordt bewezen dat een functie $f: \mathbb{R}^d \to \mathbb{R}^d$ normalisatieequivariant is dan en slechts dan als deze toelaat tot een normaliseren-processeren-hernormaliseren factorisatie. Specifiek, voor elke invoer $y$ met standaardafwijking $\text{std}(y) > 0$:
$$f(y) = \text{std}(y) \cdot g(T_{NE}(y)) + \mu(y)\mathbf{1}$$
Waarbij:

• $\mu(y)$ en $\text{std}(y)$ respectievelijk het instantiemiddelpunt en de standaardafwijking zijn.
• $T_{NE}(y) = \frac{y - \mu(y)\mathbf{1}}{\text{std}(y)}$ de normalisatiekaart is die $y$ projecteert op de genormaliseerde variëteit $\mathcal{M}$ (middelpunt nul, eenheid standaardafwijking).
• $g: \mathcal{M} \to \mathbb{R}^d$ een willekeurige functie is die werkt op de genormaliseerde invoer.

Deze karakterisering impliceert dat het afdwingen van NE geen beperking van de interne architectuur vereist; het is een equivalente herparameterisatie van de functieklasse.

De WNE-wrapper

Gebaseerd op deze theorie introduceren de auteurs WNE, een parameterloze wrapper die kan worden toegepast op elke bestaande ruggegraat $g_\theta$ (inclusief CNN's en Transformers) zonder de interne gewichten of lagen te wijzigen. De forward-pass wordt gedefinieerd als:
$$f^{WNE}_\theta(y) = \text{std}(y) \cdot g_\theta(T_{NE}(y)) + \mu(y)\mathbf{1}$$

• Invoer: Ruzig beeld $y$.
• Normalisatie: Bereken instantiestatistieken $\mu(y), \text{std}(y)$ en normaliseer naar $\tilde{y} = T_{NE}(y)$.
• Verwerking: Geef $\tilde{y}$ door aan de willekeurige ruggegraat $g_\theta$.
• Hernormalisatie: Reconstrueer de uitvoer door schaling en verschuiving: $\hat{x} = \text{std}(y)g_\theta(\tilde{y}) + \mu(y)\mathbf{1}$.

Deze constructie garandeert exacte NE, ongeacht de interne componenten van de ruggegraat (zoals attentie, LayerNorm, BatchNorm). Voor numerieke stabiliteit wordt een kleine $\epsilon$ toegevoegd aan de standaardafwijking wanneer deze dicht bij nul ligt.

Belangrijkste Bijdragen

1. Methode: Een parameterloze, plug-and-play wrapper (WNE) die exacte invoer-uitvoer normalisatieequivariantie afdwingt op elke ruggegraat, inclusief transformers, zonder interne lagen te wijzigen.
2. Theorie: Een volledige karakterisering die aantoont dat de normaliseren-processeren-hernormaliseren factorisatie zowel noodzakelijk als voldoende is voor NE, waardoor het probleem wordt getransformeerd van een interne architecturale beperking naar een invoer-uitvoer parameterisatie.
3. Experimenten & Analyse:
   • Demonstratie van verbeterde robuustheid tegen ruisniveau-mismatch bij blind ontruimen voor zowel CNN- (DnCNN) als Transformer- (SwinIR, Restormer) ruggegraten.
   • Analyse die een verwaarloosbare GPU-overhead toont ten opzichte van architecturale basismodellen.
   • Een analyse op basis van genormaliseerde coördinaten die het robuustheidsmechanisme verklaart: onder NE verplaatst een ruisniveau-mismatch voornamelijk de distributie van "moeilijkheidsgraad" (afstand in genormaliseerde coördinaten) in plaats van het regressieprobleem zelf te veranderen, waardoor het model kan opereren in een interpolatieregime over ruisniveaus.

Experimentele Resultaten

De auteurs evalueerden WNE met behulp van een diagnostiek voor één ruisniveau-mismatch, waarbij modellen werden getraind op een vast ruisniveau ($\sigma_{train} \in \{10, 25, 50\}$) en getest over een breed scala aan $\sigma_{test}$.

• Robuustheid: WNE presteert aanzienlijk beter dan ongewikkelde basismodellen onder ruisniveau-mismatch. Voor SwinIR degradeert de baseline ernstig wanneer $\sigma_{test} \neq \sigma_{train}$, terwijl WNE een stabiele prestatie handhaaft. WNE komt de robuustheidstrends van architecturale NE-basismodellen (NE-arch) die uitgebreide interne aanpassingen vereisen, nauwkeurig overeen.
• Prestaties bij Gelijk Ruisniveau: Bij gelijk ruisniveau ($\sigma_{test} = \sigma_{train}$) blijft WNE concurrerend met de ongewikkelde baseline, zonder significante degradatie in piek-PSNR.
• Runtime-efficiëntie: Op de GPU voegt WNE verwaarloosbare overhead toe (ratio $\approx 1,00\times$) ten opzichte van de baseline. Daarentegen ondervinden architecturale NE-basismodellen (NE-arch) tot een 1,6× vertraging door de kosten van het afdwingen van equivariantie in interne lagen.
• Generalisatie: De voordelen strekken zich uit tot:
  • Verschillende ruggegraten (Restormer).
  • Kleurebeelden (3-kanaals).
  • Reële sensorruis (SIDD-dataset).
  • Niet-Gaussische corrupties (Uniform, Laplace, Rayleigh, JPEG).
  • Iteratief hergebruik in inverse problemen (Noise2Noise-getrainde onruisers gebruikt in beperkte inpainting), waarbij WNE stabiele reconstructie mogelijk maakt vanuit hoog-ruisige initialisaties die ongewikkelde modellen laten falen.

Betekenis en Claims

Het artikel stelt dat het afdwingen van normalisatieequivariantie niet slechts een benadering of een beperking van de functieruimte is, maar een equivalente herparameterisatie die kan worden bereikt zonder architecturale beperkingen.

De primaire betekenis ligt in het ontkoppelen van de robuustheidsvoordelen van NE van de architecturale beperkingen van eerdere methoden. Door het afdwingen van equivariantie naar de invoer-uitvoer grens te verplaatsen via een wrapper, stellen de auteurs het gebruik van standaard, hoogpresterende componenten (zoals attentie en LayerNorm) mogelijk, terwijl ze de theoretische garanties van NE behouden. Deze aanpak biedt een praktische oplossing voor het verbeteren van robuustheid in blind ontruimen en gerelateerde inverse problemen, zonder de rekenkosten die gepaard gaan met beperkingen op interne lagen. De auteurs merken op dat dit perspectief ook een duidelijke geometrische verklaring biedt voor robuustheid: het stabiliseert het regressieprobleem in genormaliseerde coördinaten, waardoor de overlap tussen trainen en testen in "moeilijkheidsgraad" behouden blijft, zelfs onder aanzienlijke verschuivingen in ruisniveau.