Translation Invariance of Neural Operators for the FitzHugh-Nagumo Model

Título: Invarianza de Translación de Operadores Neuronales para el Modelo de FitzHugh–Nagumo

1. Planteamiento del Problema

El modelo de FitzHugh–Nagumo (FHN) es una herramienta fundamental en la electrofisiología computacional y la neurociencia para describir la dinámica de células excitables y la propagación del potencial de acción. Matemáticamente, el sistema combina una ecuación en derivadas parciales (EDP) parabólica de reacción-difusión con una ecuación diferencial ordinaria (EDO) rígida (stiff).

El desafío central abordado en este trabajo es la capacidad de los Operadores Neuronales (NOs) para aprender la solución de este sistema, específicamente su invarianza de translación. El modelo FHN posee una propiedad física clave: si se mantiene constante un conjunto de parámetros, la solución ante una corriente aplicada es invariante respecto a la traslación temporal de dicha corriente (la solución se desplaza en el tiempo pero mantiene su forma).

El problema se divide en dos aspectos:

1. Generalización fuera de distribución (OOD): Evaluar si los NOs pueden generalizar dinámicas donde la corriente aplicada se traslada tanto en el espacio como en el tiempo, aunque el modelo solo haya sido entrenado con corrientes en posiciones espaciales variables pero en un tiempo fijo.
2. Comparativa de Arquitecturas: Determinar qué arquitectura de NO (entre siete variantes) logra un equilibrio óptimo entre precisión, costo computacional y capacidad de capturar esta invarianza física.

2. Metodología

Los autores proponen una estrategia de entrenamiento novedosa diseñada para evaluar explícitamente la invarianza de translación y reducir el costo computacional de la generación de datos:

• Estrategia de Datos:

  • Entrenamiento: El conjunto de datos se genera variando la intensidad de la corriente aplicada y su posición espacial, pero manteniendo fijo el tiempo de aplicación.
  • Prueba (Test): El conjunto de prueba introduce un escenario más desafiante (OOD) donde la corriente se traslada tanto en el espacio como en el tiempo.
  • Ventaja: Al no necesitar muestrear la corriente en diferentes momentos para el entrenamiento, se reduce drásticamente el costo de generación de datos.

• Arquitecturas Evaluadas:
  Se compararon siete arquitecturas de Operadores Neuronales:

  1. CNOs (Convolutional Neural Operators): Basados en redes convolucionales que preservan la equivalencia continua.
  2. DONs (Deep Operator Networks): Redes con ramas "branch" (entrada) y "trunk" (coordenadas), incluyendo variantes con codificadores CNN (DONs-CNN) y descomposición POD (POD-DONs).
  3. FNOs (Fourier Neural Operators): Utilizan la transformada de Fourier para parametrizar kernels integrales. Incluye variantes: TFNOs (Tensorizadas de Tucker) y LocalNOs (kernels localizados).

• Métricas de Evaluación:

  • Precisión: Error relativo $L_2$ y $L_1$ en conjuntos de entrenamiento y prueba.
  • Eficiencia: Tiempo de entrenamiento, número de parámetros entrenables y tiempo de inferencia.
  • Costo Compuesto: Se definieron funciones de costo ($C$ y $C_{log}$) que ponderan el error, el número de parámetros y el tiempo de entrenamiento para una comparación justa entre modelos de diferentes escalas.

3. Contribuciones Clave

• Evaluación de Invarianza de Translación: Es uno de los primeros estudios que utiliza una estrategia de entrenamiento específica (tiempo fijo, espacio variable) para probar explícitamente la capacidad de los NOs de aprender la invarianza temporal inherente a los modelos de electrofisiología.
• Benchmark Exhaustivo: Se realiza una comparación rigurosa de siete arquitecturas de NOs en un modelo de EDP/EDO rígido, algo que rara vez se hace de manera conjunta.
• Análisis de Compromisos (Trade-offs): Se identifica claramente la compensación entre la precisión en datos de entrenamiento (donde FNOs destaca) y la robustez en datos de prueba con traslaciones (donde CNOs destaca).
• Identificación de Limitaciones Físicas: Se destaca la dificultad de los modelos para predecir correctamente el "efecto umbral" (si una corriente es suficiente para disparar un potencial de acción), lo cual genera errores significativos si se predice incorrectamente.

4. Resultados Principales

• CNOs (Convolutional Neural Operators):
  • Rendimiento: Son los únicos modelos que logran generalizar eficazmente a las dinámicas trasladas en el conjunto de prueba, capturando la invarianza de translación con un error mediano $L_2$ de ~0.09.
  • Costo: Requieren el mayor tiempo de entrenamiento (~10 horas) y tienen un alto costo computacional, aunque su precisión en entrenamiento es comparable a otros.
• FNOs (Fourier Neural Operators):
  • Rendimiento: Logran el menor error de entrenamiento ($L_2 \approx 0.008$), superando a todos los demás. Sin embargo, fallan en generalizar a las dinámicas trasladas en el conjunto de prueba, mostrando errores altos y alta sensibilidad a la inicialización.
  • Costo: Tienen un número masivo de parámetros (151.1M) y tiempos de inferencia lentos, a pesar de un tiempo de entrenamiento moderado (4 horas).
• DONs y Variantes (DONs, DONs-CNN, POD-DONs):
  • Rendimiento: Muestran una eficiencia computacional superior (entrenamiento e inferencia más rápidos, pocos parámetros). Sin embargo, no generalizan bien a las dinámicas trasladas, fallando en replicar la invarianza de translación.
  • Ventaja: Son flexibles para geometrías no cuadradas sin modificaciones estructurales.
• TFNOs y LocalNOs:
  • Comparten las limitaciones de los FNOs en la generalización de dinámicas trasladas, aunque TFNOs reduce la memoria y LocalNOs mejora la velocidad de inferencia gracias a menos parámetros.

5. Significado y Conclusión

Este trabajo proporciona un marco de referencia crítico para la aplicación de Operadores Neuronales en problemas de física compleja y rígida.

• Implicaciones Prácticas: Si el objetivo es la máxima precisión en datos dentro de la distribución de entrenamiento, los FNOs son la opción preferida. Sin embargo, para aplicaciones del mundo real donde las condiciones de contorno o la temporalidad varían (escenarios OOD con traslaciones), los CNOs son la única opción robusta, a pesar de su mayor costo de entrenamiento.
• Limitaciones Detectadas: Se identifica que la dificultad principal no es solo la aproximación de la función, sino la predicción correcta de fenómenos físicos discretos como el umbral de disparo de un potencial de acción. Un error en la predicción de si se alcanza o no el umbral genera outliers masivos que degradan el rendimiento global.
• Futuro: La investigación futura debe centrarse en mejorar la capacidad de los NOs para identificar correctamente los umbrales de activación y en desarrollar estrategias que combinen la eficiencia de los DONs con la robustez de los CNOs.

En resumen, el estudio demuestra que la elección de la arquitectura de NO no es universal; depende estrictamente de si la prioridad es la precisión en datos conocidos (FNO) o la generalización robusta ante traslaciones espaciotemporales (CNO).