A new initialisation to Control Gradients in Sinusoidal Neural network

Résumé technique : Une nouvelle initialisation pour contrôler les gradients dans les réseaux de neurones sinusoïdaux

Énoncé du problème
Les représentations neuronales implicites (INR), en particulier celles utilisant des fonctions d'activation sinusoïdales (SIREN), sont efficaces pour approximer des signaux haute fréquence dans des applications telles que la reconstruction d'images, le codage de signaux et la résolution d'équations aux dérivées partielles (EDP). Cependant, les architectures SIREN standard font face à un compromis critique : bien qu'elles atténuent le biais spectral des réseaux de neurones à plusieurs couches (MLP) standards en apprenant les hautes fréquences, l'augmentation de la profondeur du réseau conduit souvent à l'émergence d'artefacts haute fréquence spurieux. Ce phénomène dégrade la généralisation et provoque une explosion des gradients. À l'inverse, les schémas d'initialisation standards (par exemple, Xavier, Kaiming) ou l'initialisation SIREN originale ne fournissent pas de garanties théoriques pour la stabilité des gradients dans les réseaux périodiques profonds, entraînant souvent des gradients qui s'évanouissent ou explosent, ce qui entrave la dynamique d'apprentissage.

Méthodologie
Les auteurs proposent un nouveau schéma d'initialisation analytique pour les réseaux SIREN, conçu pour fonctionner à la « limite du chaos » (EOC). Ce régime garantit que les variances des pré-activations restent stables et que les gradients ne s'évanouissent ni n'explosent à mesure que la profondeur du réseau augmente.

La méthodologie repose sur une analyse théorique dans la limite d'une largeur infinie ($N \to \infty$) et d'une profondeur infinie ($L \to \infty$) :

1. Distribution des pré-activations : Les auteurs dérivent une expression en forme fermée du point fixe de la variance des pré-activations ($\sigma_a^2$) en utilisant la fonction de Lambert $W$. Ils montrent que la variance converge vers un point fixe déterminé par le paramètre de mise à l'échelle des poids $c_w$ et la variance du biais $c_b$.
2. Stabilité des gradients : En analysant la jacobienne de la mapping couche par couche, les auteurs dérivent la condition de stabilité des gradients. Ils définissent une variance de gradient redimensionnée $\sigma_g$ et imposent $\sigma_g = 1$ pour prévenir l'explosion ou l'évanouissement des gradients. Cette contrainte établit une relation spécifique entre $c_w$ et $c_b$.
3. Contrôle spectral : Les auteurs identifient deux solutions spécifiques le long de la courbe de stabilité :
   • $\sigma_a = 1$ : Correspond à l'intuition originale de SIREN (les pré-activations suivent une loi $N(0,1)$), ce que les auteurs montrent conduire à un élargissement spectral et à un aliasing dans les réseaux profonds.
   • $\sigma_a = 0$ (Proposé) : Fixe la variance des pré-activations à zéro. Dans ce régime, les couches profondes se comportent presque linéairement ($\sin(z) \approx z$), supprimant efficacement l'élargissement du spectre de Fourier et empêchant l'émergence de fréquences au-delà de la limite de Nyquist.
4. Analyse du noyau tangentiel neuronal (NTK) : L'article examine l'impact de l'initialisation sur le NTK. Ils démontrent que l'initialisation proposée ($\sigma_a=0$) entraîne des valeurs propres du NTK qui évoluent linéairement avec la profondeur, tandis que les gradients d'entrée restent normalisés. Cela contraste avec l'initialisation SIREN originale, où les valeurs propres du NTK et les gradients explosent exponentiellement avec la profondeur.

Contributions clés

• Initialisation en forme fermée : La dérivation d'un schéma d'initialisation explicite pour SIREN garantissant des distributions de gradients invariantes à travers les couches et un spectre de Fourier indépendant de la profondeur.
• Justification théorique de $\sigma_a = 0$ : L'identification du fait que la mise à zéro de la variance des pré-activations est optimale pour contrôler le spectre de fréquence, prévenir l'aliasing et assurer une dynamique d'apprentissage stable dans les architectures profondes.
• Dynamique du NTK : Une démonstration montrant que le contrôle de la propagation des gradients à l'initialisation façonne directement les valeurs propres du NTK, déterminant ainsi la vitesse d'apprentissage des différentes composantes fréquentielles et évitant le biais spectral inhérent à d'autres configurations de réseaux profonds.

Résultats
L'initialisation proposée ($\sigma_a = 0$) a été comparée au schéma SIREN original, à la variante $\sigma_a = 1$ et à d'autres architectures de l'état de l'art (par exemple, FINER, WIRE, Tanh avec caractéristiques de Fourier, ReLU avec encodage de position) sur diverses tâches :

• Ajustement de fonctions : Dans l'approximation de fonctions multi-échelles en 1D, 2D et 3D, la méthode proposée a systématiquement atteint une erreur de généralisation plus faible que les méthodes de référence, en particulier à mesure que la profondeur du réseau augmentait.
• Reconstruction d'images et de vidéos : Dans les tâches d'ajustement d'images et de reconstruction vidéo (utilisant le jeu de données atmosphérique ERA-5), la méthode proposée a produit des caractéristiques plus nettes et significativement moins d'artefacts haute fréquence que le SIREN original et d'autres méthodes.
• Réseaux de neurones informés par la physique (PINN) : La méthode a été testée sur la résolution de l'équation de Burgers 1D, des équations de Navier-Stokes stationnaires 2D et de l'équation de la chaleur 2D dans des géométries complexes. L'initialisation proposée a produit des solutions stables, tandis que l'initialisation SIREN originale a généré des résultats bruyants et instables pour les tâches d'équations de Navier-Stokes et de la chaleur en raison de composantes haute fréquence incontrôlées.
• Débruitage : Dans les expériences de débruitage d'images, l'initialisation proposée a agi comme un régularisateur implicite efficace, reconstruisant des signaux propres tout en ignorant le bruit haute fréquence, contrairement à d'autres méthodes qui avaient tendance à surajuster le bruit.

Signification
L'article affirme que ce travail fournit une compréhension théorique plus profonde de l'initialisation des INR, allant au-delà du réglage empirique pour adopter une approche fondée sur la théorie des systèmes dynamiques. En stabilisant la propagation des gradients et en contrôlant le spectre de fréquence à l'initialisation, la méthode proposée permet aux réseaux sinusoïdaux profonds d'apprendre des détails haute fréquence sans introduire d'artefacts spurieux. Cela permet une dynamique d'apprentissage stable avec des taux d'apprentissage qui évoluent linéairement avec la profondeur, rendant les architectures SIREN plus robustes et applicables à des problèmes complexes et de haute dimension tels que l'informatique scientifique et la reconstruction de signaux multi-échelles. Les auteurs notent que, bien que la motivation provienne des INR, l'initialisation en forme fermée est applicable à toute architecture profonde avec des activations périodiques où la stabilité des gradients et le contrôle de la fréquence sont souhaités.