Resumen Técnico: Equivarianza de Normalización para Arquitecturas Base Arbitrarias

Enunciado del Problema

En tareas de predicción de imagen a imagen, particularmente en la eliminación ciega de ruido de imágenes, las redes neuronales a menudo sufren desplazamientos de distribución cuando el nivel de ruido en el momento de la prueba difiere del nivel de ruido de entrenamiento. Si bien la Equivarianza de Escala (EE) ($f(ay) = af(y)$) y la Equivarianza de Normalización (EN) estrictamente más fuerte ($f(ay + b\mathbf{1}) = af(y) + b\mathbf{1}$) son conocidas por mejorar la robustez ante tales desplazamientos, imponer estas propiedades en arquitecturas modernas de aprendizaje profundo presenta desafíos significativos.

Los métodos existentes imponen la EN restringiendo las capas internas de la red a familias específicas compatibles con la EN (por ejemplo, convoluciones con restricciones afines, no linealidades de ordenación y residuos afines). Este enfoque, aunque teóricamente sólido, tiene dos desventajas principales:

1. Incompatibilidad Arquitectónica: Descarta componentes estándar omnipresentes en las arquitecturas base modernas, como los mecanismos de atención (softmax) y la normalización por capas (LayerNorm), los cuales no conmutan con transformaciones afines.
2. Sobrecarga en Tiempo de Ejecución: Imponer la equivarianza a través de restricciones en las capas internas introduce costos computacionales significativos, con algunas arquitecturas base mostrando una ralentización de hasta 1.6× en comparación con los modelos estándar.

Metodología: El Envoltorio Equivariante de Normalización (WNE)

Los autores proponen un método que impone la EN a nivel de entrada-salida en lugar de a nivel de capa interna. Este enfoque se basa en una caracterización completa de la clase de funciones EN.

Caracterización Teórica

El artículo establece la Caracterización 1, demostrando que una función $f: \mathbb{R}^d \to \mathbb{R}^d$ es equivariante bajo normalización si y solo si admite una factorización normalizar-procesar-desnormalizar. Específicamente, para cualquier entrada $y$ con desviación estándar $\text{std}(y) > 0$:
$$f(y) = \text{std}(y) \cdot g(T_{NE}(y)) + \mu(y)\mathbf{1}$$
Donde:

• $\mu(y)$ y $\text{std}(y)$ son la media y la desviación estándar de la instancia.
• $T_{NE}(y) = \frac{y - \mu(y)\mathbf{1}}{\text{std}(y)}$ es el mapa de normalización que proyecta $y$ sobre la variedad normalizada $\mathcal{M}$ (media cero, desviación estándar unitaria).
• $g: \mathcal{M} \to \mathbb{R}^d$ es una función arbitraria que actúa sobre la entrada normalizada.

Esta caracterización implica que imponer la EN no requiere restringir la arquitectura interna; es una reparametrización equivalente de la clase de funciones.

El Envoltorio WNE

Basándose en esta teoría, los autores introducen WNE, un envoltorio libre de parámetros que puede aplicarse a cualquier arquitectura base existente $g_\theta$ (incluyendo CNNs y Transformers) sin modificar sus pesos o capas internas. El paso hacia adelante se define como:
$$f^{WNE}_\theta(y) = \text{std}(y) \cdot g_\theta(T_{NE}(y)) + \mu(y)\mathbf{1}$$

• Entrada: Imagen ruidosa $y$.
• Normalización: Calcular las estadísticas de instancia $\mu(y), \text{std}(y)$ y normalizar a $\tilde{y} = T_{NE}(y)$.
• Procesamiento: Pasar $\tilde{y}$ a través de la arquitectura base arbitraria $g_\theta$.
• Desnormalización: Reconstruir la salida escalando y desplazando: $\hat{x} = \text{std}(y)g_\theta(\tilde{y}) + \mu(y)\mathbf{1}$.

Esta construcción garantiza la EN exacta independientemente de los componentes internos de la arquitectura base (por ejemplo, atención, LayerNorm, BatchNorm). Para la estabilidad numérica, se añade un pequeño $\epsilon$ a la desviación estándar cuando está cerca de cero.

Contribuciones Clave

1. Método: Un envoltorio libre de parámetros, listo para usar (WNE) que impone una equivarianza de normalización exacta entrada-salida en cualquier arquitectura base, incluidos los transformers, sin modificar las capas internas.
2. Teoría: Una caracterización completa que demuestra que la factorización normalizar-procesar-desnormalizar es tanto necesaria como suficiente para la EN, transformando el problema de una restricción arquitectónica interna a una parametrización entrada-salida.
3. Experimentos y Análisis:
   • Demostración de una robustez mejorada ante la discrepancia de niveles de ruido en la eliminación ciega de ruido tanto para arquitecturas base CNN (DnCNN) como Transformer (SwinIR, Restormer).
   • Análisis que muestra una sobrecarga de GPU negligible en comparación con las arquitecturas base.
   • Un análisis de coordenadas normalizadas que explica el mecanismo de robustez: bajo EN, la discrepancia de niveles de ruido desplaza principalmente la distribución de la "dificultad" (distancia en coordenadas normalizadas) en lugar de cambiar el problema de regresión en sí, permitiendo que el modelo opere en un régimen de interpolación a través de los niveles de ruido.

Resultados Experimentales

Los autores evaluaron WNE utilizando un diagnóstico de discrepancia de un solo nivel de ruido, entrenando modelos en un nivel de ruido fijo ($\sigma_{train} \in \{10, 25, 50\}$) y probando en un amplio rango de $\sigma_{test}$.

• Robustez: WNE supera significativamente a las arquitecturas base sin envolver bajo discrepancia de niveles de ruido. Para SwinIR, la base se degrada severamente cuando $\sigma_{test} \neq \sigma_{train}$, mientras que WNE mantiene un rendimiento estable. WNE coincide estrechamente con las tendencias de robustez de las arquitecturas base EN (NE-arch) que requieren extensas modificaciones internas.
• Rendimiento con Ruido Coincidente: En niveles de ruido coincidentes ($\sigma_{test} = \sigma_{train}$), WNE sigue siendo competitivo con la base sin envolver, mostrando ninguna degradación significativa en el PSNR pico.
• Eficiencia en Tiempo de Ejecución: En GPU, WNE añade una sobrecarga negligible (ratio $\approx 1.00\times$) en comparación con la base. Por el contrario, las arquitecturas base EN (NE-arch) incurren en una ralentización de hasta 1.6× debido al costo de imponer la equivarianza en las capas internas.
• Generalización: Los beneficios se extienden a:
  • Diferentes arquitecturas base (Restormer).
  • Imágenes en color (3 canales).
  • Ruido de sensor real (conjunto de datos SIDD).
  • Corrupciones no gaussianas (Uniforme, Laplace, Rayleigh, JPEG).
  • Reutilización iterativa en problemas inversos (desruidores entrenados con Noise2Noise utilizados en inpainting restringido), donde WNE permite una reconstrucción estable desde inicializaciones de alto ruido que hacen que los modelos sin envolver fallen.

Significado y Afirmaciones

El artículo afirma que imponer la equivarianza de normalización no es meramente una aproximación ni una restricción del espacio de funciones, sino una reparametrización equivalente que puede lograrse sin restricciones arquitectónicas.

El significado principal radica en desacoplar los beneficios de robustez de la EN de las limitaciones arquitectónicas de los métodos anteriores. Al mover la imposición de la equivarianza al límite entrada-salida mediante un envoltorio, los autores permiten el uso de componentes estándar y de alto rendimiento (como atención y LayerNorm) mientras retienen las garantías teóricas de la EN. Este enfoque ofrece una solución práctica para mejorar la robustez en la eliminación ciega de ruido y problemas inversos relacionados, sin incurrir en las penalizaciones computacionales asociadas con las restricciones de capas internas. Los autores señalan que esta perspectiva también proporciona una explicación geométrica clara de la robustez: estabiliza el problema de regresión en coordenadas normalizadas, asegurando que la superposición entre entrenamiento y prueba en "dificultad" se mantenga incluso bajo desplazamientos significativos de nivel de ruido.