🤖 machine learning

Translation Invariance of Neural Operators for the FitzHugh-Nagumo Model

This paper benchmarks seven Neural Operator architectures for solving the FitzHugh-Nagumo model, revealing that while Convolutional Neural Operators excel at handling translated spatio-temporal dynamics, Deep Operator Networks offer superior efficiency, and Fourier Neural Operators achieve the lowest training error despite poor generalization to out-of-distribution scenarios.

Original authors: Luca Pellegrini

Published 2026-03-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Luca Pellegrini

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to predict how a heart cell will react when you poke it with electricity. This isn't just a simple "poke and watch" game; the cell is a complex, fast-moving system that behaves differently depending on where you poke it and when you poke it.

This paper is a report card on a new generation of AI tools called Neural Operators. Think of these not as standard AI that memorizes pictures, but as AI that learns the rules of physics itself. The researchers wanted to see which of these AI tools could best predict the behavior of a specific heart cell model (the FitzHugh-Nagumo model) when the electrical "poke" moves around in space and time.

Here is the breakdown of their experiment and findings, using some everyday analogies:

The Big Challenge: The "Moving Target" Problem

The researchers set up a tricky test.

Training (Learning): They taught the AI by showing it electrical pokes at different spots on the cell, but always at the exact same time.
Testing (The Exam): They then asked the AI to predict what happens if the poke happens at a different time and a different spot.

In the real world, heart cells have a special property called Translation Invariance. This is a fancy way of saying: "If I poke the cell here at 5 seconds, it reacts the same way as if I poke it there at 10 seconds, just shifted in time." The AI's job was to figure out this rule without being explicitly told it exists.

The Contenders: A Race of AI Architects

The paper tested seven different types of Neural Operators. Imagine them as seven different types of chefs trying to learn a complex recipe:

FNOs (Fourier Neural Operators): The Speedy Math Whiz.
- Strength: They are incredibly fast at learning the recipe when the ingredients are in the kitchen (training set). They made the fewest mistakes during practice.
- Weakness: When you moved the ingredients to a different counter (the test set with time shifts), they got confused. They are like a chef who memorized the recipe perfectly but can't adapt if you change the order of steps. They are also very expensive to run (slow inference).
CNOs (Convolutional Neural Operators): The Patient Architect.
- Strength: They are the only ones who truly understood the "moving target" rule. When the poke moved in time and space, they still got the prediction right. They are the most reliable for real-world scenarios.
- Weakness: They are slow and expensive to train. It took them 10 hours to learn the basics, compared to minutes for others. They are like a master builder who takes a long time to lay the foundation but builds a house that never leaks.
DONs (Deep Operator Networks): The Fast & Frugal.
- Strength: They are lightning fast to train and very cheap to run.
- Weakness: They failed the test. They memorized the specific examples they saw but couldn't generalize. When the poke moved, they gave up. They are like a student who memorized the answers to the practice test but fails the real exam because the questions were slightly rearranged.
The Others (TFNOs, LocalNOs, POD-DONs): The Specialists.
- Some were good at saving memory, others at handling local details, but none of them could match the CNOs in handling the "moving" aspect of the problem.

The Verdict: The Trade-Off

The paper highlights a classic trade-off in AI: Speed vs. Robustness.

If you want the fastest, cheapest model for a controlled environment where nothing changes, DONs are great.
If you want the most accurate model for a controlled environment, FNOs are the winners.
But, if you are dealing with the real world, where things move, shift, and change unpredictably (like a beating heart), CNOs are the only ones that didn't break.

The "Gotcha" Moment

The researchers also found a common problem for all these AIs: The Threshold Effect.
Imagine a light switch. If you push it 10%, nothing happens. If you push it 11%, the light turns on. The AI struggled to figure out that tiny difference. If the electrical poke wasn't quite strong enough to trigger a heart beat, the AI sometimes guessed it would, or guessed it wouldn't, creating huge errors. This is the next big hurdle they need to solve.

In a Nutshell

This study is a benchmark for the future of medical AI. It tells us that while some AI models are faster and cheaper, they might fail when the situation changes slightly. To build reliable AI for medical simulations (like predicting heart arrhythmias), we might need to accept slower, more expensive models (like CNOs) because they are the only ones that truly understand the physics of movement and time.

1. Problem Statement

The paper addresses the challenge of applying Neural Operators (NOs) to learn the solution operators of stiff Partial Differential Equations (PDEs), specifically the FitzHugh–Nagumo (FHN) model. The FHN model is a standard representation of excitable cells in computational electrophysiology, characterized by a coupled system of a parabolic reaction-diffusion PDE (transmembrane potential) and a stiff ordinary differential equation (gating variable).

The core problem investigated is translation invariance. Physically, if a stimulus (applied current) is shifted in time or space within the FHN model, the resulting solution should simply be a translated version of the original solution. The authors aim to determine if current NO architectures can intrinsically learn this physical property. They specifically test the models' ability to generalize to out-of-distribution (OOD) scenarios where the stimulus is translated in both time and space, a scenario not explicitly present in the training data.

2. Methodology

A. Novel Training Strategy

To evaluate translation invariance efficiently, the authors propose a data generation strategy that reduces computational costs:

Training Set: Generated with applied currents ( $I_{app}$ ) varying in spatial position and intensity, but with a fixed time of application.
Test/Validation Set: Introduces a challenging OOD scenario where the applied current is translated in both time and space.
Goal: To assess if the model learns the underlying invariance property rather than memorizing specific spatiotemporal coordinates.

B. Architectures Benchmarked

The study evaluates seven distinct Neural Operator architectures:

Convolutional Neural Operators (CNOs): Based on bandlimited functions and U-Net structures with skip connections, designed to preserve continuous equivalence.
Deep Operator Networks (DONs): Utilize a "branch" (processing input function) and "trunk" (processing coordinates) architecture.
- Variants: DONs-CNN (branch uses CNNs) and POD-DONs (trunk uses Proper Orthogonal Decomposition modes).
Fourier Neural Operators (FNOs): Parameterize integral kernels in the Fourier domain to capture global dependencies.
- Variants: TFNOs (Tucker Tensorized FNOs for dimensionality reduction) and LocalNOs (using localized kernels for high-frequency features).

C. Experimental Setup

Dataset: 2,000 training samples and 500 test samples generated using the Firedrake finite element library.
Hardware: Training performed on NVIDIA A16 GPUs; inference tested on an RTX 4070.
Metrics: Evaluation based on $L_2$ and $L_1$ relative errors, training time, number of trainable parameters, and inference speed.
Cost Function: A composite metric ( $C$ and $C_{log}$ ) was introduced to balance accuracy ( $E$ ), parameter count ( $P$ ), and training time ( $T$ ).

3. Key Results

A. Accuracy and Generalization

CNOs: Demonstrated the best generalization to the translated test set. They successfully captured the translation invariance property, achieving a median test $L_2$ error of 0.09. However, they required the highest training time (~10 hours).
FNOs: Achieved the lowest training error ($0.0081$), outperforming all others on in-distribution data. However, they failed to generalize well to translated dynamics, with many test cases showing extremely high errors (up to $100$).
DONs and Variants: Showed high efficiency but poor generalization. They struggled to replicate the translation invariance property, resulting in high test errors similar to FNOs.
TFNOs & LocalNOs: Followed the trend of their parent architectures, failing to reliably capture translated dynamics despite some specific successes.

B. Computational Efficiency

Training Speed: DONs were the fastest (15 mins for 0.5M parameters), followed by POD-DONs. FNOs were surprisingly fast (4 hours) despite having the largest parameter count (151.1M). CNOs were the slowest (10 hours).
Inference Speed: DONs and their variants were the most efficient for real-time evaluation. FNOs and TFNOs suffered from high inference latency. LocalNOs offered fast inference due to low parameter counts.

C. The "Threshold Effect" Challenge

A critical finding was that all architectures struggled to correctly identify whether a stimulus was sufficient to trigger an action potential (the threshold effect). Incorrectly predicting an action potential for sub-threshold stimuli (or vice versa) created significant outliers, degrading overall model performance. This is a difficult non-linear classification task embedded within the regression problem.

4. Key Contributions

Novel Benchmarking Strategy: Introduced a training/test split specifically designed to test translation invariance in stiff PDEs, significantly reducing the cost of dataset generation by avoiding time-shifted sampling during training.
Comprehensive Comparative Study: Provided the first rigorous side-by-side comparison of seven major NO architectures (CNO, DON, FNO, and their variants) on a stiff ionic model.
Identification of Trade-offs: Established a clear trade-off between training accuracy/efficiency and generalization to translated dynamics.
- FNOs are optimal for in-distribution accuracy but fail at translation.
- CNOs are robust to translation but computationally expensive to train.
- DONs are computationally efficient but lack generalization capabilities for this specific physical property.
Cost Metrics: Proposed and applied a weighted cost function ( $C_{log}$ ) to fairly compare architectures with vastly different scales of parameters and training times.

5. Significance

This work highlights a critical limitation in current Neural Operator research: high in-distribution accuracy does not guarantee physical consistency (invariance) in out-of-distribution scenarios.

For Computational Electrophysiology: The study suggests that while FNOs are powerful for standard forward problems, CNOs may be necessary for applications requiring robustness to stimulus location and timing shifts (e.g., cardiac defibrillation or neural stimulation planning), despite their higher training cost.
For Scientific ML: The paper underscores the difficulty of learning "stiff" dynamics and the specific challenge of the threshold effect in excitable media. It calls for future research into architectures that can better handle these non-linear threshold behaviors and physical invariances without prohibitive computational costs.

In summary, the paper provides a definitive benchmark showing that CNOs are currently the only architecture capable of reliably learning the translation invariance of the FHN model, while other popular architectures (FNO, DON) fail to generalize to translated dynamics despite superior training performance.