Translation Invariance of Neural Operators for the FitzHugh-Nagumo Model

Probleemstelling

De studie richt zich op de uitdaging om Neural Operators (NOs) toe te passen op de FitzHugh-Nagumo (FHN) model, een stelsel van partiële differentiaalvergelijkingen (PDE's) en stijve gewone differentiaalvergelijkingen (ODE's) dat wordt gebruikt in computationele elektrofysiologie om de dynamiek van opwekbare cellen (zoals hartspiercellen) te beschrijven.

De kern van het probleem ligt in twee aspecten:

1. Stijve dynamiek: Het FHN-model bevat snelle en trage tijdschalen, wat het leren van de oplossing voor neurale netwerken moeilijk maakt.
2. Translatie-invariantie: Een fundamentele fysische eigenschap van het FHN-model is dat de oplossing invariant blijft ten opzichte van tijdstranslatie van de aangelegde stroom (stimulus). Als de stimulus op een ander tijdstip wordt gegeven, verschuift de oplossing simpelweg in de tijd, maar blijft de vorm gelijk.
   • De uitdaging: Kunnen Neural Operators deze intrinsieke invariantie inherent leren, of moeten ze deze expliciet worden getraind met data die elke mogelijke tijdsverschuiving bevat? Het genereren van datasets met variaties in zowel ruimte als tijd is computatievermogen-intensief.

Methodologie

De auteurs hebben een nieuwe trainingsstrategie ontwikkeld om de translatie-invariantie te testen en de datasetgeneratie te optimaliseren.

1. Nieuwe Trainingsstrategie:

   • Trainingsset: De aangelegde stroom ($I_{app}$) varieert in ruimtelijke positie en intensiteit, maar de tijdspositie is vast. Dit reduceert de kosten voor het genereren van trainingsdata aanzienlijk.
   • Testset (Out-of-Distribution): De testset introduceert een veel uitdagender scenario waarbij de stimulus zowel in tijd als ruimte wordt verschoven.
   • Doel: Beoordelen welke architectuur deze translatie-invariantie kan generaliseren zonder expliciet getraind te zijn op tijdsverschuivingen.

2. Geëvalueerde Architecturen:
   Zeven verschillende Neural Operator-architecturen werden gebenchmarkt:

   • CNOs: Convolutional Neural Operators (ontworpen om structuurbehoudende continuïteit te garanderen).
   • DONs: Deep Operator Networks (met variaties: standaard, met CNN-encoder, en POD-DONs die Proper Orthogonal Decomposition gebruiken).
   • FNOs: Fourier Neural Operators (en variaties: TFNOs met Tucker-tensor decompositie en LocalNOs met gelokaliseerde kernen).

3. Evaluatiemetrics:
   De modellen werden beoordeeld op:

   • Nauwkeurigheid (L2 en L1 relatieve fouten) op trainings- en testsets.
   • Rekenkosten (trainingstijd, aantal trainbare parameters).
   • Inferentiesnelheid (latency).
   • Een gecombineerde kostenfunctie ($C$ en $C_{log}$) die fout, parameters en tijd weegt.

Belangrijkste Bijdragen

• Novel Training Strategy: Het introduceren van een trainingsparadigma dat gebruikmaakt van de fysieke translatie-invariantie om datasetkosten te verlagen, terwijl de generalisatie naar tijdsverschoven scenario's wordt getest.
• Uitgebreide Benchmark: Een van de eerste uitgebreide vergelijkingen van zeven verschillende NO-architecturen specifiek voor een stijf ionic model (FHN).
• Inzicht in Generalisatie: Het blootleggen van het fundamentele verschil in vermogen tussen architecturen om translatie-invariantie te leren, waarbij CNOs uitblinken in generalisatie en FNOs/DONs moeite hebben met out-of-distribution data.

Resultaten

De resultaten tonen duidelijke trade-offs tussen nauwkeurigheid, generalisatievermogen en rekenkosten:

1. Generalisatie (Translatie-invariantie):

   • CNOs (Convolutional Neural Operators): Dit is de enige architectuur die de translatie-invariantie effectief leert en generaliseert naar de testset (tijdsverschoven data). Ze behalen een mediane L2-fout van ongeveer 0,09 op de testset.
   • FNOs & Variaties: Bereiken de laagste trainingsfout (0,0081), maar presteren slecht op de testset. Ze genereren vaak grote fouten (soms >100) bij tijdsverschoven stimuli, wat aangeeft dat ze de invariantie niet hebben geleerd.
   • DONs & Variaties: Zijn zeer efficiënt in training, maar generaliseren niet goed naar de getranslateerde dynamiek. Ze kampen met hoge testfouten.

2. Rekenkosten en Efficiëntie:

   • FNOs: Hebben het grootste aantal parameters (151M) en de langste inferentietijd, ondanks een relatief snelle training (4 uur).
   • DONs: Zijn de meest efficiënte modellen met de kortste trainingstijd (15-30 min) en inferentietijd, maar ten koste van generalisatievermogen.
   • CNOs: Vereisen de meeste trainingskosten (10 uur) en hebben een hoger parameteraantal dan DONs, maar bieden de enige robuuste oplossing voor de geteste uitdaging.

3. Kostenfunctie Analyse:
   Wanneer rekening wordt gehouden met fout, parameters en tijd (via de $C_{log}$ metric), presteren FNOs goed op de trainingsset, maar stijgt hun "kost" drastisch op de testset door de slechte generalisatie. CNOs tonen een consistente, betrouwbare prestatie in beide scenario's.

Betekenis en Conclusie

De studie concludeert dat er een fundamenteel compromis bestaat bij het kiezen van een Neural Operator voor complexe ionic modellen:

• Als het doel is om snel te trainen en te infereren op data die sterk lijkt op de trainingsdata, zijn DONs de beste keuze.
• Als het doel is om fysisch correcte generalisatie te bereiken voor scenario's met translatie-invariantie (zoals tijdsverschoven stimuli), is CNO de enige betrouwbare optie, ondanks de hogere trainingskosten.
• FNOs, hoewel populair en nauwkeurig op trainingsdata, blijken kwetsbaar voor deze specifieke type generalisatieproblemen in het FHN-model.

De auteurs wijzen ook op een kritieke uitdaging voor alle NO's: het correct voorspellen van de drempelwaarde (threshold effect) voor het opwekken van een actiepotentiaal. Als het model niet precies kan bepalen of een stimulus sterk genoeg is, leidt dit tot grote uitbijters in de fouten.

Deze studie biedt een rigoureuze benchmark voor het veld van Scientific Machine Learning (SciML) en benadrukt dat architectiekeuze afhankelijk is van de specifieke vereisten voor generalisatie versus rekenkosten.