Learning Affine-Equivariant Proximal Operators

1. Problemstellung

In der Signalverarbeitung und im maschinellen Lernen, insbesondere bei der Lösung schlecht gestellter inverser Probleme (z. B. Bildentrauschung, Super-Resolution), spielen Proximal-Operatoren eine fundamentale Rolle. Diese Operatoren werden oft als Denoising-Schritte in iterativen Algorithmen verwendet.

Das zentrale Problem liegt in der Balance zwischen Datenadaptivität und strukturellen Garantien:

• Gelernte Proximal-Netzwerke (LPNs): Neuere Ansätze (Fang et al., 2024) nutzen neuronale Netze, um Proximal-Operatoren zu lernen, die exakt als Gradienten einer konvexen Funktion definiert sind. Dies garantiert, dass sie echte Proximal-Schritte für ein implizites Regularisierungsterm darstellen.
• Fehlende Äquivarianz: Viele reale Daten zeigen Symmetrien, wie z. B. Verschiebungen (Shift) und Skalierungen (Scale). Es ist wünschenswert, dass die gelernten Regularizer und ihre Proximal-Operatoren diese Symmetrien respektieren (äquivariant sind).
• Das Dilemma:
  • Herkömmliche Methoden zur Erzeugung äquivarianter Funktionen (z. B. Normalisierungstricks wie Subtraktion des Mittelwerts und Division durch die Standardabweichung) zerstören die Eigenschaft, ein Gradient eines konvexen Potentials zu sein. Damit verlieren sie die Garantie, ein echter Proximal-Operator zu sein.
  • Architekturen, die intrinsisch äquivariant sind (z. B. durch Entfernen von Bias-Termen), garantieren zwar die Äquivarianz der Funktion selbst, aber nicht notwendigerweise die Äquivarianz ihres Gradienten. Da LPNs jedoch als Gradienten definiert sind, ist dies kritisch.

Die Fragestellung lautet: Wie kann man datengetriebene, gelernte Proximal-Operatoren entwerfen, die exakt als Gradienten einer konvexen Funktion definiert sind und gleichzeitig affin-äquivariant (insbesondere verschiebungs- und skalierungsäquivariant) sind?

2. Methodik

Die Autoren schlagen Affine-Equivariant Learned Proximal Networks (AE-LPNs) vor. Der Ansatz basiert auf der Konstruktion eines speziellen konvexen Potentials $\tilde{h}(x)$, dessen Gradient $\nabla \tilde{h}(x)$ sowohl ein Proximal-Operator als auch affin-äquivariant ist.

Theoretische Grundlagen:

1. Skalierungsäquivarianz: Ein Lemma zeigt, dass wenn eine Funktion $\Psi$ homogen vom Grad 2 ist ($\Psi(ax) = a^2\Psi(x)$), ihr Gradient homogen vom Grad 1 ist ($\nabla \Psi(ax) = a\nabla \Psi(x)$), also skalierungsäquivariant.
2. Verschiebungsäquivarianz: Um Verschiebungsäquivarianz zu erreichen, wird das Signal in seine Mittelwertkomponente (Projektion auf den Vektor $\mathbf{1}$) und seine zentrierte Komponente (Projektion auf das orthogonale Komplement) zerlegt.

Architektonischer Aufbau:

• Basis: Ein Input Convex Neural Network (ICNN), dessen Gewichte nicht-negativ und Aktivierungsfunktionen monoton steigend sind (z. B. SortPool), um Konvexität zu garantieren.
• Skalierung: Um Homogenität vom Grad 2 zu erreichen, werden alle Bias-Terme im ICNN entfernt und die Ausgabe des Netzes quadriert: $h_\theta(x) = \Psi_\theta(x)^2$.
• Verschiebung: Um die volle Affin-Äquivarianz zu erreichen, wird das Potential wie folgt definiert:
  $$\tilde{h}(x) = h((I - P)x) + \frac{1}{2}\|Px\|_2^2$$
  Dabei ist $P$ der Projektionsoperator auf den Mittelwert ($\frac{1}{n}\mathbf{1}\mathbf{1}^\top$). Der Term $h((I-P)x)$ behandelt die zentrierten Daten (skalierungsäquivariant), während der quadratische Term $\frac{1}{2}\|Px\|_2^2$ die Mittelwertkomponente korrekt behandelt.
• Operator: Der gelernte Proximal-Operator ist dann einfach der Gradient: $\hat{x} = \nabla \tilde{h}_\theta(x)$.

Training:
Die Netzwerke werden mit der Proximal-Matching-Loss-Funktion trainiert. Diese Loss-Funktion stellt sicher, dass der gelernte Operator im Grenzwert den MAP-Schätzer (Maximum A Posteriori) für die zugrunde liegende Datenverteilung approximiert.

3. Wichtige Beiträge

1. Theoretische Charakterisierung: Die Autoren leiten notwendige und hinreichende Bedingungen her, unter denen der Gradient einer konvexen Funktion affin-äquivariant ist (Homogenität vom Grad 2 des Potentials und spezifische Zerlegung der Eingabe).
2. AE-LPN Architektur: Entwicklung einer neuronalen Netzwerk-Architektur, die diese Bedingungen durch strukturelle Einschränkungen (keine Bias-Terme, Quadrierung der Ausgabe, Mittelwert-Subtraktion) erzwingt.
3. Garantieerhalt: Im Gegensatz zu Normalisierungstricks (z. B. Z-Score-Normalisierung) bleibt bei AE-LPNs die mathematische Garantie erhalten, dass die gelernte Abbildung ein exakter Proximal-Operator ist. Dies ist entscheidend für die Konvergenz von Algorithmen wie Plug-and-Play.
4. Robustheit: Demonstration, dass diese Struktur die Generalisierungsfähigkeit auf Verteilungen außerhalb des Trainingsdatensatzes (Out-of-Distribution) signifikant verbessert.

4. Ergebnisse

Die Methode wurde an synthetischen Daten und realen Bilddaten (BSDS500) evaluiert:

• Synthetische Beispiele (Split-Normal-Verteilung):

  • Nur die AE-LPNs konnten den wahren Proximal-Operator und das log-prior korrekt lernen, während sie gleichzeitig skalierungsäquivariant blieben.
  • Herkömmliche LPNs lernten zwar den Operator, waren aber nicht äquivariant.
  • LPNs mit dem „Normalisierungstrick" waren äquivariant, aber der resultierende Operator war kein exakter Proximal-Operator mehr (keine Konvergenzgarantie).

• Bildentrauschung über verschiedene Rauschniveaus:

  • Modelle wurden mit einem Rauschlevel von $\sigma=0.1$ trainiert und an verschiedenen Test-Rauschlevels evaluiert.
  • AE-LPNs zeigten eine deutlich höhere Robustheit gegenüber Änderungen des Rauschlevels im Vergleich zu nicht-äquivarianten LPNs.
  • Skalierungsäquivarianz erwies sich als wichtiger Faktor für die Robustheit als reine Verschiebungsäquivarianz.

• Affine Transformationen:

  • Bei der Anwendung von Transformationen $g(x) = \alpha x + (1-\alpha)$ (Simulation von Helligkeitsänderungen) zeigten AE-LPNs eine nahezu perfekte Äquivarianz: Der Fehler zwischen $f(g(x))$ und $g(f(x))$ war vernachlässigbar (hoher PSNR).
  • Dies bestätigt, dass AE-LPNs die affine Struktur der Eingabe korrekt verarbeiten, während andere Modelle dies nicht leisten.

5. Bedeutung und Ausblick

Die Arbeit löst ein fundamentales Problem beim Einsatz von Deep Learning für inverse Probleme: Sie ermöglicht es, die Flexibilität datengetriebener Modelle mit der mathematischen Stabilität und Struktur von klassischen Regularisierern zu verbinden.

• Praktischer Nutzen: AE-LPNs sind besonders nützlich in Szenarien, in denen sich die Rauschcharakteristika oder die Skalierung der Daten zwischen Trainings- und Einsatzumgebung ändern (Out-of-Distribution-Generalisierung).
• Zukünftige Anwendungen: Die Autoren sehen Potenzial für Anwendungen in der CT-Rekonstruktion, wo Variationen in Fensterung und Intensitätsskalierung häufig sind, sowie für eine verbesserte Sample-Komplexität beim Training.

Zusammenfassend stellen AE-LPNs einen wichtigen Schritt dar, um die theoretische Fundierung von „Learned Proximal Networks" zu stärken und sie für robustere, physikalisch konsistente Anwendungen in der Bildverarbeitung nutzbar zu machen.