Translation Invariance of Neural Operators for the FitzHugh-Nagumo Model

1. Problématique et Contexte

L'étude se concentre sur l'application des Opérateurs Neuronaux (NO) à la résolution des équations aux dérivées partielles (EDP) régissant le modèle de FitzHugh-Nagumo (FHN). Ce modèle est fondamental en électrophysiologie computationnelle pour décrire la propagation des potentiels d'action dans les cellules excitables (neurones, cardiomyocytes).

Le défi principal réside dans la nature « raide » (stiff) du système FHN, qui couple une EDP parabolique (réaction-diffusion) à une équation différentielle ordinaire (EDO) rigide. De plus, le modèle possède une propriété physique clé : l'invariance par translation. Pour des paramètres fixes, un stimulus appliqué décalé dans le temps et l'espace produit une solution identique, simplement translatée.

L'objectif de l'article est d'évaluer la capacité de sept architectures d'Opérateurs Neuronaux à :

1. Apprendre les dynamiques complexes et raides du modèle FHN.
2. Capturer intrinsèquement l'invariance par translation, même lorsque les données de test présentent des stimuli décalés par rapport aux données d'entraînement (scénario out-of-distribution).

2. Méthodologie

Stratégie d'Entraînement Innovante

Les auteurs proposent une stratégie de génération de données novatrice exploitant l'invariance par translation :

• Jeu de données d'entraînement : Le courant appliqué ($I_{app}$) varie en intensité et en position spatiale, mais reste fixe dans le temps. Cela réduit considérablement le coût computationnel de la génération de données, car il n'est pas nécessaire d'échantillonner le stimulus à différents moments.
• Jeu de données de test : Introduit un scénario plus difficile où le courant est translaté à la fois dans le temps et l'espace. Cela permet de tester la capacité du modèle à généraliser l'invariance temporelle qu'il n'a pas vue explicitement lors de l'entraînement.

Architectures Évaluées

Sept architectures d'Opérateurs Neuronaux ont été comparées :

1. CNO (Convolutional Neural Operators) : Basés sur les réseaux de neurones convolutifs, conçus pour préserver l'équivalence continue et la structure.
2. DON (Deep Operator Networks) : Utilisent des réseaux « branch » (entrée) et « trunk » (coordonnées spatio-temporelles).
   • Variante : DON-CNN (avec encodeur CNN).
   • Variante : POD-DON (combinaison avec la Décomposition Orthogonale Propre pour la réduction de dimension).
3. FNO (Fourier Neural Operators) : Utilisent la transformée de Fourier pour paramétrer les noyaux intégraux et capturer les dépendances globales.
   • Variante : TFNO (FNO tensorisé de Tucker pour réduire la mémoire).
   • Variante : LocalNO (Utilise des noyaux intégraux localisés pour les discontinuités).

Métriques d'Évaluation

Les modèles ont été jugés sur :

• La précision (erreurs relatives $L_1$ et $L_2$) sur les ensembles d'entraînement et de test.
• L'efficacité computationnelle (temps d'entraînement, nombre de paramètres).
• La vitesse d'inférence.
• Des fonctions de coût combinant erreur, nombre de paramètres et temps d'entraînement ($C$ et $C_{log}$).

3. Résultats Clés

Les résultats, présentés dans le tableau 1 et les figures de l'article, révèlent des compromis significatifs :

• CNO (Convolutional Neural Operators) :

  • Performance : C'est la seule architecture capable de généraliser efficacement aux dynamiques translatées (test out-of-distribution), avec une erreur médiane $L_2$ de 0,09.
  • Coût : Nécessite un temps d'entraînement élevé (environ 10 heures) et un nombre de paramètres important (8,2 M), ce qui se traduit par un coût computationnel global plus élevé.

• FNO (Fourier Neural Operators) et variantes (TFNO, LocalNO) :

  • Performance : Obtiennent la meilleure précision sur l'ensemble d'entraînement (erreur $L_2$ de 0,0081). Cependant, ils échouent à capturer l'invariance par translation sur le jeu de test, produisant des erreurs très élevées (parfois > 100) pour de nombreux exemples translatés.
  • Coût : Temps d'entraînement rapide (4h pour FNO) mais temps d'inférence très lent, rendant l'architecture coûteuse pour la prédiction en temps réel.

• DON (Deep Operator Networks) et variantes :

  • Performance : Très efficaces en termes de temps d'entraînement et de nombre de paramètres (ex: DON-CNN avec 0,5 M de paramètres). Cependant, ils ne généralisent pas bien aux scénarios translatés, échouant à reproduire l'invariance du modèle.
  • Avantage : Ils offrent la meilleure efficacité globale pour l'entraînement et l'inférence, et sont flexibles pour des géométries non carrées.

• Problème de Seuil (Threshold Effect) :
  L'analyse a mis en évidence une difficulté commune à tous les NO : la prédiction correcte du déclenchement du potentiel d'action. Une erreur mineure dans l'estimation de l'intensité du stimulus (en dessous ou au-dessus du seuil) peut entraîner des erreurs massives (outliers), dégradant considérablement les métriques de précision.

4. Contributions Majeures

1. Benchmark Complet : Fournit la première comparaison rigoureuse de sept architectures d'Opérateurs Neuronaux appliquées spécifiquement au modèle FHN, incluant des variantes récentes comme TFNO et LocalNO.
2. Stratégie de Données Économe : Démonstration qu'il est possible d'évaluer l'invariance par translation en entraînant uniquement sur des données fixes dans le temps, réduisant ainsi drastiquement le coût de génération des données d'entraînement.
3. Analyse des Compromis (Trade-offs) : Identification claire que la haute précision sur les données d'entraînement (FNO) ne garantit pas la robustesse face aux translations (CNO), et que l'efficacité computationnelle (DON) ne compense pas le manque de généralisation physique.

5. Signification et Conclusion

Cette étude souligne que, bien que les Opérateurs Neuronaux soient puissants pour l'apprentissage d'opérateurs, leur capacité à capturer des propriétés physiques fondamentales comme l'invariance par translation dépend fortement de l'architecture choisie.

• Pour des applications nécessitant une robustesse face à des conditions variables (décalages temporels/spatiaux), les CNOs sont supérieurs, malgré leur coût d'entraînement.
• Pour des tâches où la vitesse d'entraînement et d'inférence est critique et où les conditions sont stables, les DONs sont préférables.
• Les FNOs, bien que très précis sur les données d'entraînement, semblent moins adaptés aux scénarios de translation sans mécanismes d'invariance explicites ou améliorations architecturales.

L'article conclut que le défi futur réside dans l'amélioration de la capacité des NOs à identifier correctement les seuils d'activation (effet de seuil), un aspect critique pour la fiabilité des simulations en électrophysiologie. Ce travail établit une référence importante pour le choix d'architectures dans le domaine du Scientific Machine Learning (SciML) appliqué aux modèles ioniqes complexes.