🤖 machine learning

Physics-informed fine-tuning of foundation models for partial differential equations

This paper introduces a physics-informed fine-tuning framework that adapts pre-trained foundation models for partial differential equations by incorporating physical constraints directly into the objective, enabling accurate, data-efficient adaptation to new tasks and out-of-distribution scenarios without requiring extensive solution data.

Original authors: Vlad Medvedev, Leon Armbruster, Christopher Straub, Georg Kruse, Andreas Rosskopf

Published 2026-03-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vlad Medvedev, Leon Armbruster, Christopher Straub, Georg Kruse, Andreas Rosskopf

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 have a super-intelligent student named Poseidon. Poseidon has spent years reading thousands of textbooks on fluid dynamics, weather patterns, and how water flows through pipes. He is a "Foundation Model"—a master of general physics principles.

However, Poseidon has never seen a specific type of problem before: heat spreading through a metal plate (the Poisson equation). You need him to solve this new problem, but you only have one single example of the answer, or maybe even zero examples at all.

If you just ask Poseidon to guess based on his memory of water and wind, he might get it wrong because the rules are slightly different. This is the challenge the paper addresses: How do you teach a pre-trained expert a new, specific job when you don't have a textbook full of answers?

The Old Way: "Memorize the Answers"

Traditionally, to teach Poseidon this new job, you would need to show him thousands of examples of the metal plate heating up (Input) and the resulting temperature map (Output). He would memorize these pairs.

The Problem: Generating these thousands of examples takes a supercomputer weeks to calculate. It's expensive, slow, and often impossible to get enough data.
The Result: If you only give him a few examples, he starts "hallucinating"—he guesses patterns that look right statistically but break the laws of physics.

The New Way: "Physics-Informed Fine-Tuning"

The authors propose a smarter way. Instead of just showing Poseidon the answers, they give him a rulebook (the laws of physics) and say: "You don't need to memorize the answers. Just make sure your guesses follow the rules."

They introduce a framework called Physics-Informed Fine-Tuning. Here is how it works using three simple concepts:

1. The "Rulebook" (The PDE Residual)

Imagine Poseidon is drawing a map of the heat. The "Rulebook" is a mathematical equation that says, "If the heat is high here, it must flow to the cold spot there."

The Trick: Instead of checking if Poseidon's map matches a pre-drawn answer, the computer checks if his map breaks the rules.
The Analogy: It's like a teacher grading a student's essay not by comparing it to a model answer, but by checking if the student used correct grammar and logic. If the logic is sound, the answer is likely right, even if the teacher has never seen that specific topic before.

2. The "Hybrid" Strategy (The Best of Both Worlds)

The paper finds that the absolute best method is a Hybrid Approach.

The Setup: You give Poseidon a tiny bit of real data (maybe just 1 or 2 examples) AND the rulebook.
The Magic: While Poseidon looks at the 1 real example, the computer secretly generates thousands of fake scenarios (like random shapes of heat sources) and forces Poseidon to solve them using only the rulebook.
The Result: Poseidon learns the concept of heat flow so deeply that he can handle weird, extreme shapes he's never seen before. He becomes a true expert rather than a memorizer.

3. The "Zero-Shot" Miracle

The most exciting finding is that this method works even with zero labeled data.

The Scenario: You have no examples of the metal plate heating up.
The Solution: You just give Poseidon the rulebook. He generates his own "practice problems," solves them, and checks if they obey the physics.
The Outcome: He learns to solve the problem accurately without ever seeing a single "correct answer" from a human. It's like teaching someone to drive by giving them a physics textbook on friction and inertia, then letting them practice in a simulator where the car crashes if they break the laws of physics.

Why This Matters

Data Efficiency: We don't need to wait for expensive supercomputers to generate millions of training examples. We can learn from the laws of nature directly.
Safety & Reliability: Pure AI models sometimes make "statistically probable" but physically impossible predictions (like water flowing uphill). This method acts as a guardrail, ensuring the AI never breaks the laws of physics.
Generalization: When the AI encounters a weird, new situation (like a heat source shaped like a star instead of a circle), it doesn't panic. Because it understands the rules, it can figure out the answer.

The Bottom Line

This paper is like giving a genius student a new test. Instead of cramming with flashcards (data), we give them the fundamental laws of the universe (physics). The result? They can solve the test with almost no study time, and they get it right even when the questions are weird and unexpected.

This is a huge step forward for using AI in science, engineering, and climate modeling, where getting "perfect data" is often impossible, but the laws of physics are always true.

1. Problem Statement

While PDE Foundation Models (e.g., POSEIDON, WALRUS) have emerged as powerful universal surrogates pre-trained on diverse physical systems, adapting them to new downstream tasks remains challenging.

Data Scarcity: Adapting to unseen PDE classes often requires task-specific data, which is expensive or impossible to generate via traditional numerical solvers.
Distribution Shifts: Pre-trained models often fail when the distribution of initial/boundary conditions in the downstream task differs significantly from the pre-training data.
Limitations of Current Fine-Tuning: Standard data-driven fine-tuning requires large amounts of labeled input-output pairs. Purely physics-informed learning (PINNs) often suffers from unstable training and is rarely applied to fine-tuning existing foundation models.
The Gap: There is a lack of systematic study on how to leverage physical constraints (PDE residuals) to fine-tune data-based foundation models, particularly in data-scarce or zero-shot regimes.

2. Methodology

The authors propose a Physics-Informed Fine-Tuning (PIFT) framework that adapts pre-trained foundation models by incorporating physical laws directly into the optimization objective.

Base Model: The study utilizes POSEIDON, a Transformer-based foundation model pre-trained on fluid dynamics equations.
Downstream Task: The primary evaluation task is the 2D steady-state Poisson equation ( $-\Delta u = f$ ), which is an elliptic PDE fundamentally different from the hyperbolic/transport dynamics POSEIDON was pre-trained on.
Fine-Tuning Strategies:
1. Data-Driven: Standard supervised fine-tuning using labeled source-solution pairs.
2. Physics-Informed: The model is fine-tuned by minimizing the PDE residual ( $\|N_f(u)\|^2$ ) and boundary condition errors. Crucially, this can be done without labeled solution data (zero-shot regarding ground truth), relying solely on the governing equations.
3. Hybrid Strategy: A combination of data-driven loss and physics-informed loss.
  - Key Innovation: The hybrid approach employs on-the-fly augmentation. During training, the model generates unlabeled source terms (e.g., sinusoidal waves, random polygons) and evaluates them using only the physics loss. This expands the training distribution beyond the limited labeled dataset.
Implementation Details:
- Architecture: The POSEIDON backbone ( $\hat{\theta}$ ) and embeddings are transferred from pre-training, while input/output channel weights are re-initialized.
- Boundary Handling: Zero-padding is used at domain boundaries to enforce Dirichlet conditions ( $u=0$ ) and maintain spatial dimensions for finite-difference derivative calculations, avoiding soft penalty terms.
- Optimization: Gradients are backpropagated through the network to update weights based on the combined loss.

3. Key Contributions

First Systematic PIFT Framework: Introduces the first framework for physics-informed fine-tuning of data-based PDE foundation models, enabling adaptation to unseen PDE families.
Zero-Shot Adaptation: Demonstrates that foundation models can be adapted to new physical regimes using only physical constraints (PDE residuals) without any labeled solution data.
Hybrid Superiority: Proposes a hybrid fine-tuning strategy that combines limited labeled data with on-the-fly physics supervision. This yields superior generalization to out-of-distribution (OOD) scenarios compared to purely data-driven methods.
Efficiency: Shows that PIFT achieves competitive accuracy with significantly less training time and data than training operators (like FNOs) from scratch.

4. Experimental Results

The study compares POSEIDON (fine-tuned) against FNOs (trained from scratch) across interpolation (same distribution) and extrapolation (OOD) tasks.

Data Efficiency (Interpolation):
- In low-data regimes ( $M < 16$ ), fine-tuned foundation models significantly outperform FNOs.
- Hybrid POSEIDON can achieve high accuracy with just 1 labeled sample, whereas data-driven fine-tuning requires at least 32 samples to reach similar performance.
- Physics-informed FNOs outperform data-driven FNOs in low-to-mid data regimes due to regularization provided by physical constraints.
Generalization (Extrapolation):
- When tested on extreme, unseen source shapes (e.g., checkerboards, spirals), data-driven models fail significantly.
- Physics-Informed and Hybrid models exhibit robust stability. The hybrid model, trained on diverse unlabeled shapes via on-the-fly augmentation, achieves the lowest error in OOD scenarios, often outperforming its own interpolation error.
Physical Consistency:
- While data-driven models minimize pixel-wise error, they often learn "image-to-image" mappings that violate the underlying PDE.
- Physics-informed models produce solutions that strictly satisfy the governing equations (verified via Laplacian residuals), resulting in smoother, physically plausible predictions.
Computational Cost:
- Fine-tuning POSEIDON takes ~1 hour, significantly faster than training FNOs from scratch.
- Inference times are comparable (~1.6ms vs 1.7ms).

5. Significance and Impact

Paradigm Shift: This work establishes Physics-Informed Fine-Tuning as a scalable, data-efficient paradigm for scientific machine learning. It bridges the gap between the representational power of foundation models and the rigor of physical laws.
Solving the Data Bottleneck: It offers a solution for scientific domains where acquiring ground-truth data is prohibitively expensive or impossible (e.g., extreme weather, nuclear fusion, or rare material properties). Models can "self-supervise" using known physics.
Trustworthiness: By enforcing physical constraints, the resulting surrogates are more interpretable and reliable, reducing the risk of physically impossible predictions common in purely data-driven AI.
Future Directions: The authors suggest extending this to time-dependent PDEs, coupled multiphysics systems, and exploring physics-informed pre-training to further enhance foundation models.

In conclusion, the paper demonstrates that combining pre-trained foundation models with physics-informed objectives creates a robust, data-efficient pathway for solving complex PDEs, particularly in scenarios where data is scarce or the problem distribution shifts.