A new initialisation to Control Gradients in Sinusoidal Neural network

Resumen Técnico: Una Nueva Inicialización para Controlar los Gradientes en Redes Neuronales Sinusoidales

Enunciado del Problema
Las Representaciones Neuronales Implícitas (INR), particularmente aquellas que utilizan funciones de activación sinusoidales (SIREN), son efectivas para aproximar señales de alta frecuencia en aplicaciones como la reconstrucción de imágenes, la codificación de señales y la resolución de ecuaciones diferenciales parciales (EDP). Sin embargo, las arquitecturas SIREN estándar enfrentan un compromiso crítico: si bien mitigan el sesgo espectral de las Perceptrones Multicapa (MLP) estándar al aprender altas frecuencias, aumentar la profundidad de la red a menudo conduce a la aparición de artefactos de alta frecuencia espurios. Este fenómeno degrada la generalización y provoca la explosión de los gradientes. Por el contrario, los esquemas de inicialización estándar (por ejemplo, Xavier, Kaiming) o la inicialización original de SIREN no proporcionan garantías teóricas para la estabilidad de los gradientes en redes periódicas profundas, lo que a menudo resulta en gradientes que se desvanecen o explotan, obstaculizando la dinámica de entrenamiento.

Metodología
Los autores proponen un nuevo esquema de inicialización, derivado analíticamente, para redes SIREN diseñado para operar en el "borde del caos" (EOC). Este régimen asegura que las varianzas de las pre-activaciones permanezcan estables y que los gradientes ni se desvanezcan ni exploten a medida que aumenta la profundidad de la red.

La metodología se basa en un análisis teórico en el límite de ancho infinito ($N \to \infty$) y profundidad infinita ($L \to \infty$):

1. Distribución de Pre-activación: Los autores derivan una expresión en forma cerrada para el punto fijo de la varianza de la pre-activación ($\sigma_a^2$) utilizando la función $W$ de Lambert. Demuestran que la varianza converge a un punto fijo determinado por el parámetro de escalado de los pesos $c_w$ y la varianza del sesgo $c_b$.
2. Estabilidad del Gradiente: Mediante el análisis del Jacobiano del mapeo de capa a capa, los autores derivan la condición para la estabilidad del gradiente. Definen una varianza de gradiente reescalada $\sigma_g$ y hacen cumplir $\sigma_g = 1$ para prevenir la explosión o el desvanecimiento del gradiente. Esta restricción establece una relación específica entre $c_w$ y $c_b$.
3. Control Espectral: Los autores identifican dos soluciones específicas a lo largo de la curva de estabilidad:
   • $\sigma_a = 1$: Corresponde a la intuición original de SIREN (las pre-activaciones siguen una distribución $N(0,1)$), lo que los autores muestran que conduce a un ensanchamiento espectral y aliasing en redes profundas.
   • $\sigma_a = 0$ (Propuesto): Establece la varianza de la pre-activación en cero. En este régimen, las capas profundas se comportan casi linealmente ($\sin(z) \approx z$), suprimiendo efectivamente el ensanchamiento del espectro de Fourier y previniendo la aparición de frecuencias más allá del límite de Nyquist.
4. Análisis del Núcleo Tangente Neuronal (NTK): El artículo investiga el impacto de la inicialización en el NTK. Demuestran que la inicialización propuesta ($\sigma_a=0$) resulta en valores propios del NTK que escalan linealmente con la profundidad, mientras que los gradientes de entrada permanecen normalizados. Esto contrasta con la inicialización original de SIREN, donde los valores propios del NTK y los gradientes explotan exponencialmente con la profundidad.

Contribuciones Clave

• Inicialización en Forma Cerrada: La derivación de un esquema de inicialización explícito para SIREN que garantiza distribuciones de gradiente invariantes entre capas y un espectro de Fourier independiente de la profundidad.
• Justificación Teórica para $\sigma_a = 0$: La identificación de que establecer la varianza de la pre-activación en cero es óptimo para controlar el espectro de frecuencias, prevenir el aliasing y asegurar una dinámica de entrenamiento estable en arquitecturas profundas.
• Dinámica del NTK: Una demostración de que controlar la propagación del gradiente en la inicialización moldea directamente los valores propios del NTK, determinando así la velocidad de aprendizaje de los diferentes componentes de frecuencia y evitando el sesgo espectral inherente a otras configuraciones de redes profundas.

Resultados
La inicialización propuesta ($\sigma_a = 0$) fue evaluada frente al esquema original de SIREN, la variante $\sigma_a = 1$ y otras arquitecturas de última generación (por ejemplo, FINER, WIRE, Tanh con características de Fourier, ReLU con codificación posicional) en diversas tareas:

• Ajuste de Funciones: En la aproximación de funciones multiescala en 1D, 2D y 3D, el método propuesto logró consistentemente un error de generalización más bajo en comparación con las líneas base, particularmente a medida que aumentaba la profundidad de la red.
• Reconstrucción de Imágenes y Video: En tareas de ajuste de imágenes y reconstrucción de video (utilizando el conjunto de datos atmosférico ERA-5), el método propuesto produjo características más nítidas y significativamente menos artefactos de alta frecuencia en comparación con SIREN original y otros métodos.
• Redes Neuronales Informadas por Física (PINNs): El método fue probado en la resolución de la ecuación de Burgers 1D, las ecuaciones de Navier-Stokes estacionarias 2D y la ecuación de calor 2D en geometrías complejas. La inicialización propuesta arrojó soluciones estables, mientras que la inicialización original de SIREN produjo resultados ruidosos e inestables en las tareas de Navier-Stokes y ecuación de calor debido a componentes de alta frecuencia no controlados.
• Desharollado: En experimentos de eliminación de ruido en imágenes, la inicialización propuesta actuó como un regularizador implícito efectivo, reconstruyendo señales limpias mientras ignoraba el ruido de alta frecuencia, a diferencia de otros métodos que tendían a sobreajustar el ruido.

Significado
El artículo afirma que este trabajo proporciona una comprensión teórica más profunda de la inicialización de INR, avanzando más allá del ajuste empírico hacia un enfoque basado en la teoría de sistemas dinámicos. Al estabilizar la propagación del gradiente y controlar el espectro de frecuencias en la inicialización, el método propuesto permite que las redes sinusoidales profundas aprendan detalles de alta frecuencia sin introducir artefactos espurios. Esto permite una dinámica de entrenamiento estable con tasas de aprendizaje que escalan linealmente con la profundidad, haciendo que las arquitecturas SIREN sean más robustas y aplicables a problemas complejos y de alta dimensión, como la computación científica y la reconstrucción de señales multiescala. Los autores señalan que, aunque la motivación surge de las INR, la inicialización en forma cerrada es aplicable a cualquier arquitectura profunda con activaciones periódicas donde se desee la estabilidad del gradiente y el control de la frecuencia.