Learning Affine-Equivariant Proximal Operators

Título: Aprendendo Operadores Proximais Equivariantes Afim (Learning Affine-Equivariant Proximal Operators)

Autores: Oriel Savir, Zhenghan Fang e Jeremias Sulam (Johns Hopkins University).

---

1. Problema e Motivação

Os operadores proximais são fundamentais em problemas inversos mal-postos (como remoção de ruído, super-resolução e compressão sensorial), atuando como sub-rotinas em solvers iterativos. Recentemente, as Redes de Aprendizado Proximal (LPNs) foram introduzidas para aprender operadores proximais exatos a partir de dados, garantindo que a rede aprendida seja o gradiente de um potencial convexo.

No entanto, as LPNs existentes carecem de garantias estruturais importantes. Em muitos cenários práticos, é crucial que os regularizadores (e seus operadores proximais correspondentes) respeitem propriedades de simetria dos dados, especificamente equivariância a deslocamentos (shift) e escala (scale).

• O Desafio: Métodos anteriores que tentam impor essa equivariância (como normalização de entrada ou arquiteturas específicas para denoisers genéricos) falham em manter a garantia teórica de que a rede representa um operador proximal exato.
  • A "truque de normalização" (subtrair a média e dividir pelo desvio padrão) quebra a propriedade de ser o gradiente de uma função convexa.
  • Arquiteturas que são equivariantes afim como funções não garantem que seus gradientes também o sejam.

O objetivo central do trabalho é projetar operadores proximais aprendidos que sejam exatos (garantidos matematicamente) e simultaneamente equivariantes a transformações afim escalares (deslocamento e escala).

---

2. Metodologia: AE-LPNs

Os autores propõem as Redes de Aprendizado Proximal Equivariantes Afim (AE-LPNs). A abordagem baseia-se na caracterização de que um operador proximal é o gradiente de uma função potencial convexa ($\psi$). Para garantir a equivariância afim no gradiente, eles derivam condições específicas sobre a função potencial.

Construção Teórica

1. Equivariância de Escala:

   • O gradiente de uma função é equivariante à escala se a função potencial for homogênea de grau 2.
   • Lema 1: Se $\Psi(ax) = a^2\Psi(x)$, então $\nabla\Psi(ax) = a\nabla\Psi(x)$.
   • Implementação: Constrói-se uma Rede Neural Convexa de Entrada (ICNN) sem termos de viés (bias) e com funções de ativação homogêneas (como SortPool). O resultado é então elevado ao quadrado ($h(x) = \Psi(x)^2$) para garantir a homogeneidade de grau 2.

2. Equivariância a Deslocamento (Shift):

   • Para lidar com deslocamentos aditivos ($x \to x + c\mathbf{1}$), a função potencial deve decompor o sinal em sua componente de média e sua componente ortogonal.
   • Seja $P$ o projetor no espaço gerado pelo vetor de uns ($\mathbf{1}$), e $(I-P)$ o projetor no complemento ortogonal.
   • Define-se a função potencial final como:
     $$\tilde{h}(x) = h((I - P)x) + \frac{1}{2}\|Px\|_2^2$$
   • O gradiente desta função satisfaz: $\nabla\tilde{h}(ax + c\mathbf{1}) = a\nabla\tilde{h}(x) + c\mathbf{1}$.

Arquitetura e Treinamento

• Arquitetura: Uma ICNN onde os pesos são não-negativos e as ativações são não-decrescentes (SortPool).
  • Remoção de todos os termos de viés (bias).
  • O sinal de entrada é centralizado (subtração da média) antes de entrar na ICNN.
  • A saída da ICNN é elevada ao quadrado e combinada com o termo quadrático da média.
• Função de Perda: Utiliza-se a Perda de Correspondência Proximal (Proximal Matching Loss - LPM).
  • O treinamento é supervisionado de forma não supervisionada: amostras limpas $x$ são corrompidas com ruído gaussiano $y = x + z$.
  • A rede é treinada para minimizar a distância entre a saída da rede e a amostra limpa, mas com uma perda específica que garante que a rede aprenda a densidade logarítmica dos dados (MAP), convergindo para o operador proximal correto quando o hiperparâmetro $\gamma \to 0$.

---

3. Resultados Experimentais

Os autores validaram o método em três cenários:

1. Aprendizado de Operadores Proximais em Distribuição Split-Normal (1D):

   • Em um exemplo construtivo unidimensional, apenas o AE-LPN conseguiu aprender o operador proximal verdadeiro e a função potencial convexa correta.
   • O método com "truque de normalização" produziu um denoiser equivariante, mas que não era um operador proximal exato (não era gradiente de um potencial convexo), violando a garantia teórica necessária para a convergência em métodos como Plug-and-Play.

2. Remoção de Ruído em Imagens (BSDS500) - Robustez a Níveis de Ruído:

   • Modelos foram treinados com ruído de nível $\sigma=0.1$ e testados em níveis variados.
   • Resultado: Embora as LPNs originais (mais flexíveis) tenham desempenho ligeiramente superior no nível de ruído de treinamento, os AE-LPNs demonstraram robustez superior em níveis de ruído fora da distribuição de treinamento (out-of-distribution).
   • A equivariância de escala mostrou-se mais crítica para essa robustez do que a de deslocamento.

3. Validação de Transformações Afim:

   • Testou-se a propriedade $f(g(x)) = g(f(x))$ onde $g(x) = \alpha x + (1-\alpha)$ (simulando mudanças de brilho).
   • Resultado: O AE-LPN foi o único modelo que manteve a estrutura afim da imagem com erro negligenciável (PSNR muito alto entre a saída transformada e a transformação da saída), provando que aprendeu um operador proximal verdadeiro e equivariante.

---

4. Contribuições Chave

1. Formulação Teórica: Derivação das condições necessárias e suficientes para que o gradiente de uma função convexa seja equivariante a transformações afim escalares (homogeneidade de grau 2 + decomposição de média).
2. Arquitetura AE-LPN: Proposta de uma arquitetura de rede neural que impõe essas restrições estruturais (remoção de bias, uso de SortPool, elevação ao quadrado) garantindo que a rede seja um operador proximal exato.
3. Garantia de Exatidão: Diferente de métodos heurísticos de normalização, o AE-LPN mantém a garantia matemática de ser um operador proximal, o que é essencial para a estabilidade e convergência em solvers de problemas inversos.
4. Robustez: Demonstração empírica de que a incorporação de simetrias afins melhora significativamente a generalização para distribuições de ruído e escalas não vistas durante o treinamento.

---

5. Significado e Impacto

Este trabalho preenche uma lacuna crítica entre o aprendizado profundo de regularizadores e a teoria de otimização convexa. Ao garantir que os operadores aprendidos respeitem simetrias físicas (escala e deslocamento) sem sacrificar a propriedade de serem operadores proximais exatos, o AE-LPN oferece:

• Melhor Generalização: Modelos mais robustos para cenários do mundo real onde as condições de aquisição (iluminação, escala, ruído) variam.
• Confiabilidade Teórica: Permite o uso seguro em frameworks iterativos complexos (como Plug-and-Play ou ADMM) onde a propriedade de ser um operador proximal é um pré-requisito para a convergência.
• Aplicações Futuras: O método tem potencial para ser aplicado em reconstrução de tomografia computadorizada (CT) e outras áreas onde variações de janela e escala de intensidade são comuns, além de reduzir a complexidade estatística necessária para o treinamento.