Learning Affine-Equivariant Proximal Operators

1. Problem Statement

Inverse problems in signal processing and machine learning (e.g., image denoising, deblurring, super-resolution) often require regularization to solve ill-posed tasks. A powerful approach involves using Proximal Operators ($\text{prox}_R$) as sub-steps in iterative solvers.

• The Gap: While Learned Proximal Networks (LPNs) have been developed to learn data-driven regularizers that are guaranteed to be exact proximal operators (by defining them as gradients of convex potentials), they lack structural constraints.
• The Challenge: Many real-world data distributions exhibit symmetries, specifically shift (translation) and scale (intensity) equivariance. Standard LPNs do not inherently respect these symmetries.
• The Conflict: Existing methods to enforce equivariance in neural networks (e.g., normalization tricks or specific architectural designs like SortPool) often break the mathematical guarantee that the network represents a valid proximal operator (i.e., the gradient of a convex function). Specifically, a function can be affine-equivariant, but its gradient (the proximal operator) may not be.

Core Question: How can we design data-driven learned proximal operators that are provably exact proximal operators while simultaneously being affine-equivariant (invariant to shifts and scaling)?

2. Methodology: Affine-Equivariant Learned Proximal Networks (AE-LPNs)

The authors propose AE-LPNs, a framework that constructs proximal operators by enforcing specific conditions on the underlying convex potential function ($\psi$).

Theoretical Foundation

The proximal operator is defined as the gradient of a convex potential: $f(x) = \nabla \psi(x)$. To ensure $f$ is affine-equivariant ($f(ax+b) = af(x)+b$), the authors derive conditions for $\psi$:

1. Scale Equivariance: If $\psi$ is homogeneous of degree 2 (i.e., $\psi(ax) = a^2\psi(x)$), then its gradient $\nabla \psi$ is homogeneous of degree 1 (scale-equivariant).
2. Shift Equivariance: To handle shifts, the potential must decompose the input into its mean and its zero-mean component.

Architectural Construction

The AE-LPN is constructed through the following steps:

1. Base Convex Network: Start with an Input Convex Neural Network (ICNN). To ensure convexity, weights are constrained to be non-negative, and non-decreasing activation functions are used.
2. Scale Equivariance:
   • Remove all bias terms from the ICNN.
   • Use SortPool (a sorting-based nonlinearity) as the activation function, which is 1-homogeneous.
   • Square the output of this ICNN. If the base network output is $\Psi_\theta(x)$, the new potential is $h_\theta(x) = (\Psi_\theta(x))^2$. This ensures $h_\theta$ is homogeneous of degree 2.
3. Shift Equivariance:
   • Decompose the input $x$ into its mean component ($Px$) and zero-mean component ($(I-P)x$), where $P$ is the projection onto the all-ones vector.
   • Define the final potential $\tilde{h}(x)$ as:
     $$\tilde{h}(x) = h_\theta((I-P)x) + \frac{1}{2}\|Px\|_2^2$$
   • The first term handles the shape (scale equivariance), and the second term (quadratic in the mean) ensures the gradient correctly handles shifts.
4. Final Operator: The AE-LPN is the gradient of this constructed potential: $\hat{x} = \nabla \tilde{h}_\theta(x)$.

Training Strategy

The models are trained using Proximal Matching Loss ($L_{PM}$), which ensures the learned operator approximates the Maximum A Posteriori (MAP) estimator for the data distribution.

• Process: Clean signals $x$ are corrupted with Gaussian noise $y = x + z$. The network minimizes a risk function measuring the distance between the denoised output and the ground truth.
• Loss Function: A specific loss $d_\gamma$ is used that converges to the correct regularizer as $\gamma \to 0$.
• Schedule: Training begins with an $\ell_1$ loss pre-training phase, followed by optimization with the proximal matching loss where $\gamma$ is gradually decreased.

3. Key Contributions

1. Theoretical Proof: The paper proves that if a convex potential is constructed to be 2-homogeneous and properly decomposed regarding the mean, its gradient is provably an affine-equivariant proximal operator.
2. Novel Architecture: Introduction of AE-LPNs, which combine ICNNs, bias removal, SortPool activations, and a specific output squaring/decomposition strategy to satisfy both convexity and equivariance constraints simultaneously.
3. Resolution of the "Normalization Trick" Flaw: The authors demonstrate that standard normalization tricks (e.g., $z$-score normalization) break the convexity guarantee required for a network to be a true proximal operator. AE-LPNs solve this without relying on such tricks.
4. Robustness to Distribution Shift: The framework enables learned operators that generalize to noise levels and affine transformations unseen during training.

4. Experimental Results

The authors evaluated AE-LPNs on synthetic data and real-world image denoising (BSDS500 dataset).

• Synthetic Example (Split Normal Distribution):

  • Task: Learn the proximal operator of a split normal distribution.
  • Result: Only the AE-LPN successfully learned the true scale-equivariant proximal operator and the correct log-prior. The "normalization trick" produced a scale-equivariant denoiser but failed to be an exact proximal operator (violating the convex gradient property).

• Image Denoising (Out-of-Distribution Noise Levels):

  • Models were trained on noise level $\sigma=0.1$.
  • Test: Performance was evaluated across varying noise levels ($\sigma \in [0.05, 0.7]$).
  • Result: While standard LPNs performed best at the training noise level, AE-LPNs significantly outperformed them at higher and lower noise levels. This demonstrates superior robustness to distribution shifts in noise intensity. Scale equivariance was found to be the dominant factor in this robustness.

• Affine Transformation Invariance:

  • Test: The network was tested on images transformed by $g(x) = \alpha x + (1-\alpha)$ (simulating brightness changes).
  • Metric: PSNR between $f(g(x))$ and $g(f(x))$.
  • Result: AE-LPNs achieved near-perfect PSNR (negligible error), proving they strictly maintain the affine structure of the image. Other models (including the normalization trick variant) failed to maintain this structural consistency.

5. Significance and Impact

• Theoretical Rigor: This work bridges the gap between deep learning flexibility and mathematical guarantees in inverse problems. It ensures that learned denoisers are not just "black boxes" but valid proximal operators, which is crucial for the convergence of iterative algorithms like Plug-and-Play (PnP) ADMM.
• Practical Utility: By enforcing equivariance, AE-LPNs reduce the statistical complexity of learning and improve generalization. This is particularly valuable in medical imaging (e.g., CT reconstruction) and remote sensing, where data may vary in intensity scaling or offset due to acquisition conditions, and where training data covering all possible variations is unavailable.
• Future Directions: The authors suggest this framework could improve sample complexity in training and be applied to other inverse problems where windowing and intensity scaling are ubiquitous.