Learning Affine-Equivariant Proximal Operators

1. Il Problema

I problemi inversi nell'elaborazione dei segnali e nel machine learning (come il denoising di immagini, la deblurring e la compressione sensoriale) sono spesso mal posti e richiedono regolarizzazione per ottenere soluzioni uniche o stabili.
Un approccio moderno utilizza gli Operatori Prossimali ($\text{prox}_R$) come sottopassi in solutori iterativi. Recenti lavori hanno introdotto le Reti Prossimali Apprese (LPN), che apprendono operatori prossimali esatti per regolarizzatori basati sui dati, garantendo che la rete sia il gradiente di un potenziale convesso.

Tuttavia, esiste una limitazione fondamentale: le LPN standard non incorporano necessariamente strutture aggiuntive cruciali per molti compiti, come l'equivarianza rispetto a trasformazioni geometriche (spostamenti e scale). Quando i dati presentano simmetrie (es. livelli di rumore variabili o cambiamenti di luminosità/contrasto), imporre queste proprietà nel regolarizzatore può migliorare significativamente la generalizzazione e la robustezza.
Il problema centrale affrontato è: come progettare operatori prossimali appresi che siano sia esatti (garantiti matematicamente) sia equivarianti rispetto a trasformazioni affini (spostamento e scala)?

Le soluzioni esistenti falliscono su questo fronte:

• Il "trucco della normalizzazione" (normalizzare input/output) rompe la garanzia che la rete rappresenti un operatore prossimale (non è più il gradiente di una funzione convessa).
• Le architetture equivarianti generiche (es. rimozione di bias e uso di SortPool) non garantiscono che il gradiente della funzione appresa (necessario per essere un operatore prossimale) mantenga l'equivarianza.

2. Metodologia

Gli autori propongono le Affine-Equivariant Learned Proximal Networks (AE-LPN), un framework che costruisce operatori prossimali esatti e affini-equivarianti.

Fondamenti Teorici

L'approccio si basa sulla caratterizzazione degli operatori prossimali come gradienti di funzioni potenziali convesse ($\psi$). Per garantire l'equivarianza affine $f(ax + b) = af(x) + b$, gli autori derivano condizioni specifiche sulla funzione potenziale $\psi$:

1. Equivarianza di Scala: Se una funzione $\Psi$ è omogenea di grado 2 ($\Psi(ax) = a^2\Psi(x)$), allora il suo gradiente $\nabla\Psi$ è omogeneo di grado 1 (cioè equivariante alla scala: $\nabla\Psi(ax) = a\nabla\Psi(x)$).
2. Equivarianza di Spostamento: Per gestire lo spostamento, la funzione deve essere costruita decomponendo il segnale nella sua media (proiezione sul vettore di tutti 1, $P$) e nel suo complemento ortogonale ($I-P$).

Costruzione dell'Architettura AE-LPN

L'architettura modifica le LPN standard (basate su Input Convex Neural Networks - ICNN) come segue:

1. Rimozione dei Bias: Si eliminano tutti i termini di bias (bias terms) dai layer e dalle connessioni residue per preservare l'omogeneità.
2. Funzioni di Attivazione: Si utilizzano funzioni di attivazione equivarianti alla scala, come SortPool o ReLU.
3. Quadratura dell'Output: Si costruisce una rete ICNN $\Psi_\theta$ e si eleva al quadrato il suo output per ottenere $h_\theta(x) = \Psi_\theta(x)^2$. Questo garantisce che $h_\theta$ sia omogenea di grado 2.
4. Combinazione per l'Equivarianza Affine: La funzione potenziale finale $\tilde{h}(x)$ è definita come:
   $$\tilde{h}(x) = h_\theta((I - P)x) + \frac{1}{2}\|Px\|_2^2$$
   Dove $Px$ è il vore della media dell'input. Questa struttura garantisce che il gradiente $\nabla\tilde{h}(x)$ sia sia convesso (quindi un operatore prossimale esatto) sia affine-equivariante.

Addestramento

Le AE-LPN vengono addestrate utilizzando la Loss di Matching Prossimale (Proximal Matching Loss, $L_{PM}$), che promuove l'adattamento alla densità logaritmica dei dati (MAP estimate) in modo non supervisionato. Il processo include una fase di pre-addestramento con perdita $L_1$ seguita dall'ottimizzazione con $L_{PM}$, dove il parametro $\gamma$ viene gradualmente ridotto.

3. Risultati Chiave

Gli esperimenti sono stati condotti su dati sintetici e reali (dataset BSDS500 per il denoising di immagini).

• Apprendimento di Operatori Prossimali Esatti: Su una distribuzione normale divisa (split normal), solo le AE-LPN sono state in grado di apprendere l'operatore prossimale vero e proprio mantenendo l'equivarianza di scala. Le reti con il "trucco della normalizzazione" producevano denoiser equivarianti ma non garantivano di essere operatori prossimali esatti (non gradienti di funzioni convesse).
• Robustezza al Rumore (Out-of-Distribution): Nel denoising di immagini, le AE-LPN mostrano una robustezza superiore rispetto alle LPN standard quando testate su livelli di rumore diversi da quelli di addestramento. Mentre le LPN standard crollano in prestazioni se il rumore di test differisce da quello di training, le AE-LPN mantengono prestazioni elevate grazie all'invarianza di scala appresa.
• Invarianza Affine: Sperimentando trasformazioni affini (cambiamenti di luminosità/contrasto simulati da $g(x) = \alpha x + (1-\alpha)$), le AE-LPN hanno dimostrato di mantenere la struttura affine dell'immagine con errori trascurabili (PSNR altissimo tra $f(g(x))$ e $g(f(x))$), a differenza degli altri modelli che fallivano nel preservare tale struttura.

4. Contributi Principali

1. Definizione Teorica: Identificazione delle condizioni matematiche necessarie affinché il gradiente di una funzione convessa sia affine-equivariante (omogeneità di grado 2 della potenziale).
2. Architettura AE-LPN: Progettazione di una rete neurale che soddisfa queste condizioni, garantendo che l'output sia un operatore prossimale esatto e affine-equivariante.
3. Superamento dei Limiti Esistenti: Dimostrazione che il "trucco della normalizzazione" è incompatibile con la garanzia di essere un operatore prossimale, offrendo una soluzione architetturale intrinseca invece.
4. Validazione Empirica: Prove che l'equivarianza appresa migliora significativamente la generalizzazione in scenari di distribuzione diversa (out-of-distribution), specialmente per il denoising con livelli di rumore variabili.

5. Significato e Implicazioni

Questo lavoro colma un divario importante tra l'apprendimento profondo adattivo e le garanzie matematiche necessarie per la risoluzione stabile di problemi inversi.

• Affidabilità: Fornisce operatori prossimali che sono sia "intelligenti" (adattati ai dati) sia "robusti" (rispettano le simmetrie fisiche dei dati).
• Applicazioni Future: Oltre al denoising, il framework è promettente per applicazioni come la ricostruzione CT, dove variazioni di intensità e finestre (windowing) sono comuni, e per compiti che richiedono un campionamento più efficiente (sample complexity).
• Generalizzazione: Dimostra che incorporare vincoli di simmetria (equivarianza) direttamente nella struttura della rete, piuttosto che tramite normalizzazione post-hoc, porta a modelli più stabili e teoricamente fondati.