Learning Affine-Equivariant Proximal Operators

Probleemstelling

Inversproblemen, zoals beeldontrussing (denoising) en super-resolutie, vereisen vaak regularisatie om stabiele oplossingen te vinden. Traditionele methoden gebruiken expliciete regularisatoren, terwijl moderne data-gedreven benaderingen vaak "Learned Proximal Networks" (LPNs) gebruiken. LPNs leren een proximal operator (een denoising-stap) door deze te modelleren als de gradiënt van een convexe potentiaalfunctie.

Hoewel LPNs waarborgen bieden dat de geleerde functie een echte proximal operator is, missen ze vaak de mogelijkheid om specifieke structurele eigenschappen, zoals equivariantie, in te bouwen.

• Equivariantie betekent dat een transformatie van de invoer (bijv. verschuiving of schaling) resulteert in een overeenkomstige transformatie van de uitvoer.
• In beeldverwerking is het cruciaal dat denoisers robuust zijn tegen veranderingen in ruisniveaus (schaling) en helderheid (verschuiving).
• Bestaande methoden om equivariantie te bereiken (zoals normalisatie-trucs of het verwijderen van bias-termen in netwerken) breken echter de garantie dat de functie een exacte proximal operator is (d.w.z. de gradiënt van een convexe functie). Er is dus een behoefte aan een methode die data-gedreven is, exacte proximal operators levert, en tegelijkertijd affine-equivariant (schaling en verschuiving) is.

Methodologie

De auteurs introduceren Affine-Equivariant Learned Proximal Networks (AE-LPNs). De kern van de methode ligt in het construeren van een convexe potentiaalfunctie $\tilde{h}(x)$ waarvan de gradiënt $\nabla\tilde{h}(x)$ zowel convex is (voor de proximal-garantie) als affine-equivariant.

De constructie verloopt in twee stappen:

1. Schalingsequivariantie (Scale-Equivariance):

   • Om schalingsequivariantie te garanderen, moet de gradiënt van de potentiaalfunctie homogeen van graad 1 zijn.
   • Volgens Lemma 1 is de gradiënt van een functie die homogeen is van graad 2, automatisch homogeen van graad 1.
   • De auteurs bouwen een Input Convex Neural Network (ICNN) zonder bias-termen en met schalingsequivariante activatiefuncties (zoals SortPool). De uitgang van dit netwerk wordt vervolgens gekwadrateerd ($h_\theta(x) = \Psi_\theta(x)^2$). Dit zorgt ervoor dat de functie homogeen van graad 2 is, waardoor de gradiënt schalingsequivariant is.

2. Verschuivingsequivariantie (Shift-Equivariance):

   • Om verschuivingsequivariantie toe te voegen, wordt het signaal ontleed in zijn gemiddelde (projectie op de vector van enen, $P$) en het gemiddelde-vrije deel ($I-P$).
   • De definitieve potentiaalfunctie wordt gedefinieerd als:
     $$\tilde{h}(x) = h((I - P)x) + \frac{1}{2}\|Px\|_2^2$$
     Hierbij is $h$ de eerder geconstrueerde schalingsequivariante functie. De term $\frac{1}{2}\|Px\|_2^2$ zorgt voor de correcte respons op verschuivingen.
   • De gradiënt van deze functie, $\nabla\tilde{h}(x)$, is bewezen zowel schaling- als verschuivingsequivariant te zijn, terwijl het nog steeds de gradiënt van een convexe functie blijft.

Training:
De netwerken worden getraind met de Proximal Matching Loss (een methode uit eerdere LPN-werk). Deze loss-functie zorgt ervoor dat de geleerde proximal operator overeenkomt met de log-dichtheid van de trainingsdata (MAP-schatting), zelfs zonder dat de onderliggende regularisator expliciet bekend is.

Belangrijkste Bijdragen

1. Theoretische Formulering: De auteurs bewijzen dat een convexe functie die homogeen is van graad 2 een schalingsequivariante gradiënt heeft, en combineren dit met een projectie-methode om volledige affine-equivariantie te bereiken.
2. AE-LPN Architectuur: Een nieuwe netwerkarchitectuur die de garanties van LPNs (exacte proximal operators) behoudt terwijl deze strikt voldoet aan affine-transformaties (schaling en verschuiving).
3. Oplossing voor het "Normalisatie-probleem": Ze tonen aan dat de gebruikelijke "normalisatie-truc" (input normaliseren, verwerken, dan denormaliseren) de garantie van een convex potentiaal breekt, en bieden een architecturale oplossing die dit probleem omzeilt.

Resultaten

De auteurs evalueren AE-LPNs in drie scenario's:

1. Synthetisch voorbeeld (Split Normal Distributie):

   • Op een 1D-synthese-experiment leert de AE-LPN de ware proximal operator en de onderliggende regularisator perfect.
   • Andere methoden (zoals LPN met normalisatie-truc) leveren weliswaar een equivariante denoiser, maar geen exacte proximal operator (geen gradiënt van een convexe functie).

2. Beelddenoising over verschillende ruisniveaus:

   • Getraind op een ruisniveau van $\sigma=0.1$, getest op BSDS500.
   • AE-LPNs tonen een aanzienlijk hogere robuustheid bij test-ruisniveaus die afwijken van het trainingsniveau (out-of-distribution), vergeleken met standaard LPNs.
   • Schalingsequivariantie bleek in dit scenario belangrijker dan verschuivingsequivariantie voor de prestaties.

3. Evaluatie van Affiene Transformaties:

   • De auteurs testen of de denoiser de structuur van het beeld behoudt bij transformaties $g(x) = \alpha x + (1-\alpha)$ (simulatie van helderheidsveranderingen).
   • Alleen de AE-LPN behoudt de equivariantie-eigenschap perfect (hoge PSNR tussen $f(g(x))$ en $g(f(x))$), terwijl andere modellen hierin falen of geen exacte proximal operator zijn.

Betekenis en Impact

Dit werk is significant omdat het de kloof overbrugt tussen theoretische garanties (dat een netwerk een geldige proximal operator is voor inversproblemen) en praktische robuustheid (equivariantie tegen veranderingen in data-distributie).

• Robuustheid: AE-LPNs zijn superieur in situaties waar het test-omgeving verschilt van het trainingsomgeving (bijv. andere ruisniveaus of helderheid), wat cruciaal is voor toepassingen zoals medische beeldvorming (CT-reconstructie) waar instellingen variëren.
• Algoritme-stabiliteit: Omdat AE-LPNs exacte proximal operators zijn, kunnen ze veilig worden gebruikt in iteratieve oplossingsalgoritmen (zoals Plug-and-Play), waar convergentie-eigenschappen afhankelijk zijn van de proximal-garantie.
• Toekomstperspectief: De methode opent de deur voor het leren van priors in complexe toepassingen waar schaling en verschuivingen inherent zijn, zonder in te leveren op de wiskundige eigenschappen die nodig zijn voor stabiele inversie.