Learning Affine-Equivariant Proximal Operators

1. Problématique

Les problèmes inverses en traitement du signal et en apprentissage automatique (comme le débruitage d'images, la défloutage ou la reconstruction tomographique) sont souvent mal posés et nécessitent une régularisation. Une approche moderne consiste à utiliser des opérateurs proximaux appris via des réseaux de neurones, appelés Learned Proximal Networks (LPN). Ces réseaux paramètrent des fonctions qui calculent exactement l'opérateur proximal d'un régularisateur, garantissant ainsi la convergence théorique des algorithmes itératifs (comme Plug-and-Play).

Cependant, deux limitations majeures persistent :

1. Manque de structure : Les LPN existants apprennent des régularisateurs génériques sans garantir qu'ils respectent des propriétés structurelles importantes des données, telles que l'équivariance aux transformations géométriques.
2. Robustesse aux changements de distribution : Les modèles appris sont souvent sensibles aux changements de distribution (par exemple, des niveaux de bruit différents entre l'entraînement et le déploiement).

L'objectif de cet article est de concevoir des opérateurs proximaux appris qui sont affinement équivariants (invariants aux translations et aux changements d'échelle), tout en conservant la garantie mathématique d'être de vrais opérateurs proximaux (c'est-à-dire les gradients de fonctions convexes).

2. Méthodologie

Les auteurs proposent une nouvelle architecture appelée AE-LPN (Affine-Equivariant Learned Proximal Networks). La méthodologie repose sur trois piliers théoriques et architecturaux :

A. Fondements Théoriques

Un opérateur proximal $prox_R(y)$ est défini comme le minimiseur d'une fonction objectif combinant une fidélité aux données et un régularisateur $R$. Selon le théorème de Moreau, un opérateur $f$ est un opérateur proximal si et seulement s'il est le gradient d'une fonction convexe $\psi$ ($f = \nabla \psi$).
Pour garantir l'équivariance affine ($f(ax + b) = af(x) + b$), les auteurs démontrent que :

• Si une fonction $\Psi$ est homogène de degré 2 ($\Psi(ax) = a^2\Psi(x)$), alors son gradient est homogène de degré 1 (équivariant à l'échelle).
• Pour gérer les translations (décalages), il faut décomposer le signal en sa moyenne (composante constante) et sa partie centrée.

B. Construction Architecturale

L'architecture AE-LPN est construite à partir d'un Input Convex Neural Network (ICNN) modifié :

1. Suppression des biais : Tous les termes de biais (biases) sont retirés des couches du réseau pour permettre l'équivariance à l'échelle.
2. Fonctions d'activation : Utilisation de fonctions d'activation équivariantes à l'échelle, telles que SortPool (basé sur le tri) ou ReLU, au lieu des activations standards.
3. Homogénéité de degré 2 : La sortie du réseau ICNN (qui est déjà une fonction convexe) est mise au carré. Si $\Psi_\theta$ est le potentiel ICNN, la nouvelle fonction potentielle est $h_\theta(x) = \Psi_\theta(x)^2$. Cela garantit que $h_\theta$ est homogène de degré 2.
4. Décomposition Affine : Pour obtenir l'équivariance complète (translation + échelle), le potentiel est défini comme suit :
   $$\tilde{h}(x) = h((I - P)x) + \frac{1}{2}\|Px\|_2^2$$
   Où $P$ est le projecteur sur le vecteur de tous les uns (calculant la moyenne du signal) et $(I-P)x$ est la version centrée du signal.
   • Le gradient de cette fonction $\nabla \tilde{h}(x)$ est prouvé être un opérateur proximal exact et affine-équivariant.

C. Entraînement

Les modèles sont entraînés avec la perte d'appariement proximal (Proximal Matching Loss - $L_{PM}$), introduite précédemment par Fang et al. (2024). Cette perte permet d'apprendre le régularisateur implicite (la densité de probabilité des données) de manière non supervisée en minimisant la distance entre la sortie débruitée et le signal vrai, en utilisant un paramètre $\gamma$ qui tend vers zéro.

3. Contributions Clés

1. Théorème de caractérisation : Les auteurs établissent les conditions nécessaires et suffisantes pour qu'un réseau neuronal convexe génère un gradient affine-équivariant (homogénéité de degré 2 + décomposition moyenne/centrée).
2. Architecture AE-LPN : Une construction pratique qui intègre ces contraintes dans un ICNN, garantissant que le réseau est à la fois un opérateur proximal exact et affine-équivariant.
3. Résolution du compromis : Contrairement aux "tricks" de normalisation (normalisation par la moyenne et l'écart-type) qui brisent la propriété de gradient convexe, ou aux réseaux équivariants génériques qui ne sont pas des opérateurs proximaux, l'AE-LPN préserve les deux propriétés simultanément.

4. Résultats Expérimentaux

Les auteurs évaluent leur méthode sur des données synthétiques et réelles (BSDS500) :

• Apprentissage d'opérateurs proximaux (Cas 1D) : Sur une distribution "normale divisée" (split normal), seul l'AE-LPN parvient à apprendre le véritable opérateur proximal tout en respectant l'équivariance à l'échelle. Les tentatives avec la "normalisation trick" échouent à produire un vrai opérateur proximal (le gradient n'est plus celui d'une fonction convexe).
• Débruitage robuste (Niveaux de bruit variables) :
  • Les modèles sont entraînés avec un bruit de niveau $\sigma=0.1$.
  • Lors du test sur des niveaux de bruit différents, les AE-LPNs montrent une robustesse supérieure par rapport aux LPN standards.
  • L'équivariance à l'échelle s'avère plus critique que l'équivariance aux translations pour le débruitage d'images.
• Équivariance Affine : L'AE-LPN est le seul modèle capable de maintenir la structure affine de l'image lors de transformations (changement de luminosité/contraste simulé par $g(x) = \alpha x + (1-\alpha)$), avec des erreurs de PSNR négligeables entre $f(g(x))$ et $g(f(x))$.

5. Signification et Impact

Ce travail est significatif car il comble le fossé entre l'apprentissage profond (flexibilité) et la théorie de l'optimisation convexe (garanties).

• Généralisation hors distribution : En intégrant l'équivariance affine, les modèles deviennent beaucoup plus robustes aux variations de conditions d'acquisition (ex: changements de luminosité ou d'échelle en imagerie médicale) sans nécessiter de réentraînement.
• Garanties théoriques : L'approche assure que les itérations de type "Plug-and-Play" convergent, ce qui est crucial pour les problèmes inverses critiques.
• Applications futures : Les auteurs suggèrent que cette méthode pourrait améliorer la complexité d'échantillonnage et être appliquée à des domaines comme la reconstruction CT, où les variations d'intensité et de fenêtrage sont omniprésentes.

En résumé, les AE-LPNs offrent un cadre rigoureux pour apprendre des régularisateurs profonds qui respectent les symétries physiques des données tout en conservant les propriétés mathématiques nécessaires à la résolution stable de problèmes inverses.