Unsupervised Adaptation of PDE Foundation Models

arXiv:2608.07053 · cs.AI · Submitted 2026-08-07 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: Paper Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Unsupervised Adaptation of PDE Foundation Models".

Jane: The paper was written by Ziye Song, Xin Yu, Ivor Tsang, Zhao Wei and Yueming Lyu from Nanyang Technological University and Adelaide University and Centre for Frontier AI Research, A*STAR.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Paper summary: TOM: Fresh paper on the table, and this one gave me goosebumps. The authors adapt PDE foundation models to brand-new equations without ever seeing a single ground-truth answer. That flips the usual scientific eye workflow on its head.

JANE: No interior labels at all? That's like a chef learning a new dish from the recipe and the heat settings, but never tasting a finished plate. How does the model even know it's doing well?

LALAM: The equations themselves become the teacher. The loss combines the PDE residual with boundary conditions, so physics supplies the supervision simulation data usually provides. And as far as this team can tell, nobody has pulled that off in the foundation-model setting before.

LU: It matters because simulation data is the scarce resource. The paper opens by reminding us that high-quality scientific simulations can eat months of supercomputer time, and every neural operator training sample inherits that cost.

MENG: Hold on — plenty of PDE foundation models already exist. Poseidon, PDE-Transformer, MPP. They all still demand dense labels when adapting to unseen equations.

JANE: So the pretraining is normal, but the adaptation step is what changes?

TOM: Exactly. Pretrain a neighborhood attention Transformer on six PDEBench families — compressible Navier-Stokes, shallow water, reaction-diffusion. Then freeze that backbone and tune only low-rank adapters against the physics objective.

LU: The adapter tweak is the sleeper hit. Standard LoRA learns unevenly across physical quantities, so they introduce NSLoRA, a Newton-Schulz orthogonalized variant. Each rank direction ends up pulling its own weight.

MENG: And the performance gap to full supervision is startling. On seven of the eight 2D benchmarks, they land within a factor of 2.5 of supervised LoRA finetuning. Zero interior labels involved.

JANE: They're also beating fully supervised neural operators?

TOM: On nine of the eleven downstream datasets, they outperform at least one neural operator baseline trained with complete ground truth. FNO, U-Net, CNextU-Net — supervised, and still losing.

LALAM: That's the headline for the field. Dense interior data is the wall between scientific eye and real deployment. If the governing equations plus boundary readings are enough, huge domains suddenly open up.

LU: The boundary term anchors the prediction at the rim, while the residual checks physics in the interior. Two weak signals, combined, carry the model nearly as far as the full answer key.

MENG: Nearly, but not all the way. The paper is careful about where it breaks, and the limitations section will walk us through those bruises.

TOM: Page one lays out that promise and the gap it fills. Let's start there.

Page 1 of the paper: TOM: We've got the promise and the problem. Page one builds the case for why the gap even exists.

JANE: Traditional numerical solvers — finite-difference, spectral — are accurate but painfully expensive. The paper ticks off that cost right in the opening paragraphs.

LALAM: Then neural operators arrived to learn mappings between function spaces. Fast at inference, sure, but their appetite for high-fidelity training data is brutal.

LU: Each training sample can itself require an expensive simulation. The paper frames that as the core bottleneck — you need data to make data.

MENG: PINNs were the first attempt to escape that trap. They embed the governing equations directly into the loss, so labels become optional. But they train per system, and the paper notes they're often suboptimal in practice.

JANE: Meta-learning extensions tried to share structure across related PDEs, yet their transfer decays when the equations genuinely differ. They generalize within families, then stumble across them.

TOM: That's where foundation models entered. Borrow the recipe from language and vision — pretrain broadly, adapt downstream. But here's the wrinkle the paper hammers home.

JANE: The adaptation stage still wants dense solution fields. So you're back to paying the simulation tax exactly when you want to escape it.

LU: Meanwhile, the governing equations and boundary conditions sit right there, known and cheap. Existing methods just weren't using them during adaptation.

MENG: This paper claims that first. First adaptation of a pretrained PDE foundation model driven strictly by the PDE residual and boundary conditions. No interior ground truth anywhere.

TOM: And there's an architectural angle hiding in the abstract. Downstream PDEs arrive at different spatial resolutions, so they adopt neighborhood attention — a local window over tokens instead of global attention.

JANE: Local windows fit PDEs intuitively.

TOM: Exactly. A disturbance spreads through neighbors, not through the whole domain at once. Global attention is checking every country when you only need the local weather.

LALAM: Plus, a fixed window size gives linear scaling in sequence length. That's a deployment win on top of an accuracy win.

LU: The remaining piece is NSLoRA. Standard LoRA's rank budget collapses in practice, leaving some physical quantities under-adapted. Orthogonalizing the low-rank factors with Newton-Schulz iterations rebalances that learning.

MENG: And the abstract's numbers are the hook — within 2.5× of supervised on most 2D benchmarks, ahead of at least one supervised baseline on nine of eleven datasets.

TOM: All without interior labels. Now page four shows the machinery that actually delivers those numbers.

Page 4 of the paper: JANE: The pitch is set. Page four turns it into equations and a blueprint.

TOM: The math is compact — a general spatiotemporal system takes the form partial-t of u plus N of u equals zero, with an initial condition and a boundary operator. The residual operator R then measures how badly any candidate field violates the equation.

JANE: And the task is next-step prediction. Feed eight frames of history, predict the ninth.

LU: The clever part is the shifted predictive design. The model outputs the entire shifted sequence, so every input frame doubles as supervision for the frame after it. One forward pass, dense temporal training signal.

MENG: But that design forces causality in attention. Each input frame already contains the exact value you're trying to predict for the next step. If attention could peek forward, the model would just copy and win trivially.

JANE: So temporal attention is masked to earlier frames only.

TOM: Right. Spatial attention stays bidirectional, since space has no such leakage. That asymmetry is dictated by the training objective itself.

LALAM: The backbone is a neighborhood attention Transformer. Each token attends within a window of size k along every axis — one temporal, the rest spatial. Locality is baked into the kernel.

LU: The encoder patches the input, sixteen by sixteen, and pushes each patch through strided convolutions into tokens. The decoder reverses the trip. Resolution in, resolution out.

MENG: What about datasets with completely different fields? Shear flow has velocity, pressure, and a tracer. Gray-Scott has two chemical concentrations. The channels don't line up at all.

TOM: The answer is an eighteen-slot shared representation. Three velocity slots, fifteen scalar slots. Each dataset writes into the slots matching its semantics, zero-fills the rest, and a binary mask tells the loss what's active.

JANE: Pressure parks in its slot, temperature in another. The empty slots just stay empty.

LU: Each sample also gets standardized per physical quantity before the encoder. Mean and variance computed per sample, so velocity and pressure live on comparable scales.

MENG: Pretraining is then plain supervised next-step prediction on PDEBench. Six subsets, simple MSE, no physics term in stage one.

LALAM: The philosophy is deliberate — first learn the vocabulary of physics broadly, then discipline the model to a specific equation in stage two. Pretraining supplies transferable representations, and the physics objective refines them.

JANE: Question is whether that financed architecture pays off on totally unseen equations. The tables on page seven put numbers on it.

TOM: Let's go read those.

Page 7 of the paper: TOM: Blueprint's in place. Page seven is the scoreboard — eleven downstream datasets, split into two suites.

JANE: Four come from The Well. Rayleigh-Bénard, Shear Flow, Gray-Scott, Active Matter — high-fidelity simulations of real physical processes.

LU: Seven more are synthesized from exact analytical solutions. Burgers, advection, Taylor-Green, wave, and advection-diffusion, spread across one, two, and three dimensions.

MENG: And the choice is principled. The unsupervised objective computes PDE residuals through finite differences, so any discretization error in the data corrupts the loss. Exact solutions give machine-precision residuals at every point.

TOM: They're candid that some popular benchmarks like PDEArena and PDEGym carry enough numerical error to break residual-based training. Data hygiene as a first-class concern.

JANE: The metric is VRMSE — variance-normalized root mean squared error. It divides the error by how much the field actually varies, so a score above one means you're worse than just predicting the time average.

LU: The baseline gauntlet includes FNO, TFNO, U-Net, and CNextU-Net, all trained from scratch with full labels. Plus PDE-Transformer and Poseidon, finetuned from released checkpoints.

MENG: The marquee number — their unsupervised framework beats PDE-Transformer under the identical physics objective by a geometric-mean factor of 9.9. Same loss, vastly better backbone.

TOM: And against the supervised upper bound, the same backbone finetuned with dense labels, the unsupervised version stays within 2.5× on seven of the eight 2D benchmarks.

JANE: So the supervised model saw every interior answer, and the unsupervised one saw only equations plus a boundary rim. Worst gap, a factor of 2.5.

LU: Against the neural operator baselines, it beats at least one of them on nine of the eleven datasets. Those baselines chewed on full ground truth the entire time.

LALAM: That compresses the entire argument into one table. The equations are data. Pretraining plus physics constraints carry a model most of the way to supervision quality.

MENG: One honesty flag — the UPAO gain is backbone-dependent. CNextU-Net improves on only fourteen of twenty-six transfer cells. TFNO improves on all of them.

TOM: So the pretrained representation is the amplifier. Strip it away and errors jump by more than an order of magnitude on the exact-solution suite.

JANE: Page ten's conclusion owns the remaining limits. Let's hear them.

Page 10 of the paper: JANE: Scoreboard's read, and the questions are sharpening. Page ten closes the core story and then gets honest about its bruises.

TOM: First, a tiny ablation that says a lot — how the PDE residual gets normalized. System residuals have wildly different scales, so the paper divides each sub-equation residual by its own measured magnitude.

LU: Skipping that normalization is a disaster. On Burgers, VRMSE jumps from 0.0034 with normalization to 0.1303 without. The gradient signal gets swallowed by the loudest equation.

MENG: Their best scheme also rescales the boundary weight after normalization. Table 6 shows it winning on most datasets, restoring the balance between the rim signal and the interior residual.

JANE: It's a tug-of-war between the boundary anchor and the physics residual. The weight tuning decides which side of the rope wins.

TOM: The conclusion then repeats the central claim — unsupervised adaptation works, but its gain is backbone-dependent. TFNO improves under UPAO, CNextU-Net barely moves, and their pretrained backbone sees the largest, most uniform gains.

LALAM: That's a mature result, honestly. The physics objective amplifies a good representation. It can't conjure one from random weights.

MENG: The limitations list is refreshingly specific. You need the governing equations of the target system, and you need boundary observations throughout the trajectory. Some field campaigns won't offer either.

LU: Also, the residual term weakens when simulation data carries discretization error. That's the reason the evaluation leans on exact analytical solutions — a necessity born from the residual's fragility.

JANE: And the scope stops at single-step prediction. Autoregressive rollouts, where the model eats its own outputs, are never tested. Long-horizon stability is wide open.

TOM: A model can nail one step and still wander into nonsense after fifty. That's a genuine caveat for real deployment.

LALAM: Still, the framing lands. Interior labels are the scarce resource in computational science. The equations are nearly free, and this paper shows how far the free signal can take you.

MENG: The bibliography that follows tells the lineage story — which previous systems supplied the pieces this paper assembles.

JANE: Let's trace that family tree next.

Page 13 of the paper: TOM: We've dissected the results and the limits. Page thirteen is pure references, but it reads like a family tree.

JANE: Raissi's PINN work anchors the physics-informed branch. Li's FNO and Lu's DeepONet anchor the neural operator branch. This paper welds those branches together.

LU: The dataset papers are here too. PDEBench supplies the pretraining corpus, and The Well provides four of the eleven downstream test sets. The evaluation would be impossible without that infrastructure.

MENG: The closest cousin is MORPH, listed right there. It also uses LoRA to adapt PDE foundation models, but it strictly demands dense ground truth. This paper takes that recipe and deletes the labels.

LALAM: The bibliography reads as a chain of missing ingredients. MORPH gave the LoRA pattern, The Well gave diverse simulation data, PDEBench gave the pretraining corpus.

TOM: The Taylor-Green entry is a beautiful detail. That's a 1937 fluid dynamics test case, and the appendix leans on its closed-form solution to validate modern learning systems.

JANE: Two traditions separated by almost a century, meeting in one table.

LU: Wang's PINN failure-mode papers appear too — the ones about training dynamics and causality. The temporal masking from page four was born from lessons those papers taught.

MENG: PI-MFM shows up as the other physics-informed foundation model. But it focuses on one-dimensional time-dependent PDEs, so its scope and this paper's scope barely overlap.

LALAM: And Shazeer's GLU paper explains the SwiGLU blocks in the feed-forward layers. Every transformer detail here is traceable to someone else's hard-won lesson.

LU: LeMON also appears — one network across nineteen PDEs via meta-learning. Another path toward cross-equation generalization, but still supervised at the end.

TOM: The references even point toward the appendix's practical tricks, including the channel mapping that lets one backbone eat any dataset.

JANE: Those weeds are on page sixteen. Let's get our hands dirty.

Page 16 of the paper: JANE: Family tree mapped. Page sixteen is pure engineering — channel slots and Newton-Schulz arithmetic.

TOM: The channel vocabulary gets fully spelled out. Pressure, temperature, buoyancy, density, height, passive tracer, gravitational potential — fifteen scalar names plus three velocity directions.

LU: Every dataset maps its fields into matching slots, and the active-channel mask decides what the loss sees. It's a fixed alphabet for physics.

MENG: Then NSLoRA gets its formula. Newton-Schulz iterates a polynomial — X goes to aX plus bX X-transpose X plus c of (X X-transpose) squared X — with coefficients 3.4445, minus 4.7750, 2.0315.

JANE: Five iterations of that?

TOM: Five iterations push singular values into a band around one. Not exact orthogonality, but enough to kill the near-zero directions that cause rank collapse.

LU: The magnitudes come from rescaling by the Frobenius norms of the original matrices. That replaces standard LoRA's fixed alpha-over-r with an adaptive per-module amplitude.

JANE: The warm-up protocol is pragmatic. They optimize with standard LoRA first, then flip to the NSLoRA forward pass once validation improvement stalls. Same matrices, same optimizer state, only the forward path changes.

LU: That's how the ablation stays fair. Standard LoRA and NSLoRA share the same rank budget, same warm-up, same training loop. The comparison isolates the orthogonalization.

MENG: The spectral analysis justifies the surgery. Standard LoRA's stable rank sits between 1.31 and 2.20 despite a nominal budget of sixteen. The rank budget is silently collapsing.

JANE: NSLoRA lifts stable rank by 5.2 percent on average, up to 15.6 percent on Rayleigh-Bénard.

TOM: Modest numbers, but the per-channel tables show the gains concentrate on the weak quantities — pressure, tracer, the second velocity component. Exactly the rebalancing they promised.

LALAM: The philosophy is lovely. Don't buy more capacity, use what you already paid for. Sixteen directions are plenty if all sixteen actually carry signal.

LU: And they measure the cost of the fix. Newton-Schulz runs 1.38× faster end to end than SVD-based orthogonalization, because it's built from plain matrix multiplications.

TOM: With the adapters sorted, the remaining risk is the data. The appendix's dataset construction is just as careful.

JANE: Page nineteen builds exact-solution benchmarks with closed-form fields. Let's look at those.

Page 19 of the paper: TOM: Adapters sorted. Page nineteen handles the data side — finishing the pretraining menu and starting the exact-solution banquet.

JANE: Alongside the compressible Navier-Stokes runs, there's shallow water with radial dam breaks, FitzHugh-Nagumo reaction-diffusion, and incompressible Navier-Stokes with random forcing.

LU: That's a varied diet. Different equations, different spatial dims, different physical scales. The backbone is forced to learn what physics holds in common.

MENG: Then the evaluation data appears, built on a clever trick. Every dataset has a closed-form analytical solution, so the stored fields satisfy the PDE exactly at every grid point.

JANE: The residual can then be computed to machine precision. No solver noise, no discretization error, no ambiguity about whether the data actually obeys the equation.

TOM: First up is 2D Taylor-Green. The velocity is a product of cosines and sines decaying with e to the minus two-nu-t, and the pressure decays at twice that rate.

LU: They sweep viscosity uniformly between 0.01 and 0.1 per trajectory. Two hundred trajectories, each a different physical regime, all evaluated on a 256-by-256 grid.

MENG: Why such obsessive care? Because UPAO's residual is only as trustworthy as the data precision. If the stored fields violate the equation by even a little, the loss punishes the model for the solver's mistakes.

LALAM: Most learning benchmarks never check whether their data solves the PDE. Here, the data defines the truth and the residual simultaneously, so the data must be exact. That's rigorous experimental hygiene.

MENG: The same pattern carries through the wave, advection-diffusion, and Burgers datasets. Random modes, random phases, parameter sweeps, all closed-form.

JANE: The grids are periodic too, which makes Fourier differentiation exact — no boundary corrections needed for these particular cases.

TOM: These exact-solution cases stress the method where it should shine. But The Well datasets on page twenty-two are the hostile stress test.

LU: Active matter and Rayleigh-Bénard will hurt more. Let's see whether the method takes the pain.

Page 22 of the paper: JANE: The clean benchmarks look good. Page twenty-two plunges into the messy, high-fidelity world.

TOM: It wraps the clean cases first — 1D advection with pure translating sine modes, and three dee advection on a 64-cubed grid. Then The Well arrives.

LU: Active Matter comes first. Rod-like particles swimming in a Stokes fluid, governed by a coupled Smoluchowski-Stokes system. The paper uses the concentration field plus two velocity components.

MENG: The parameter sweep is wide — five dipole strengths and nine alignment strengths. That's a demanding spread of active behaviors.

JANE: Next is Rayleigh-Bénard. Boussinesq convection, Rayleigh number from a million to ten billion, Prandtl number from 0.1 to 10.0. Pure chaos territory at the high end.

TOM: And the domain is 512 by 128, strongly elongated. Remember the appendix's aspect-ratio routing? This dataset gets the stretched neighborhood attention kernel.

LU: Velocity, pressure, and buoyancy all have to be predicted. Four active channels, each with its own physical character, and the trajectories run for hundreds of snapshots.

MENG: Seven of the eleven downstream tasks are equations completely unseen during pretraining. This is the true generalization test — new equations, not merely new parameters.

TOM: The paper reports one honest blemish. On Gray-Scott, TFNO wins the transfer contest, because its global Fourier modes capture the low-wavenumber patterns better than a local 5-by-5 window.

JANE: Local attention has a blind spot for globally coherent structures, and they print that result anyway. Science working as advertised.

LU: Active matter, convection, shear instabilities — all adapted without a single interior ground-truth field. Just equations and boundary readings.

LALAM: That's the whole arc of the paper in one sentence. Labels are scarce, equations are abundant, and the pretrained backbone knows how to ride the physics signal.

MENG: The concluding segment is next — what this unlocks, and what stays locked.

Conclusion: LALAM: We've walked the whole paper. Now we step back and ask what it means.

TOM: The headline is simple. A pretrained PDE foundation model can adapt to an unseen equation using only the governing equation and boundary data. No interior labels anywhere.

JANE: And the gap to supervision is just a factor of 2.5 on most 2D benchmarks. Meanwhile, fully supervised neural operators end up behind on nine of eleven datasets.

LU: The reusable piece is NSLoRA. A cheap Newton-Schulz orthogonalization that stops rank collapse and rebalances learning across physical quantities. That trick will travel beyond PDEs.

MENG: The pretraining stage still costs serious compute, sure. But it's paid once, and adaptation becomes a label-free, low-rank tweak.

LALAM: That's the foundation-model promise fully realized for physics. Amortize the expensive learning up front, then personalize for free.

JANE: The deeper lesson is that physics itself is a supervision signal. Generations of scientists encoded the world's behavior into these equations, and this framework finally asks the equations to do the teaching.

TOM: The open wounds — autoregressive stability, discretized real-world data, equations that aren't known — become the field's to heal.

MENG: The most exciting part for me is the reframing. The answer key was never the only teacher.

LU: Think about climate modeling, where simulations cost weeks of supercomputing. If equations plus boundary readings can fine-tune a foundation model, regional studies get dramatically cheaper.

JANE: And engineering design — airfoils, reactors, pipelines — where sensors give boundary data but interior measurements are impossible.

LALAM: That's exactly the deployment wall this paper starts to knock down.

TOM: We'll be watching for the follow-ups. Autoregressive rollouts, noisy data, unknown equations — those are the next papers.

JANE: Thanks for sticking with us. The next paper on the pile is already calling.

Ziye Song, Zhao Wei, Xin Yu, Ivor Tsang, Yueming Lyu

Nanyang Technological University · Adelaide University · Centre for Frontier AI Research, A*STAR

cs.AI

Submitted: 2026-08-07

Updated: 2026-08-10

License: http://creativecommons.org/licenses/by-nc-sa/4.0/

Importance score: 77/100

Key concepts

PDE residual: A measure of how much a predicted solution violates the governing partial differential equation. It is computed using finite differences and serves as a training signal, replacing the need for labeled simulation data.
Neighborhood attention: A transformer attention mechanism where each token only attends to nearby tokens within a local window, rather than the entire sequence. This matches the local nature of physical interactions and scales linearly with sequence length.
NSLoRA: A variant of low-rank adaptation that orthogonalizes the low-rank factors using Newton-Schulz iterations. This prevents uneven learning across physical quantities, ensuring each rank direction contributes effectively during fine-tuning.

Terminology

Adaptation — 49× in this paper · explained in 1 episode
Transformer — 37× in this paper · explained in 2 episodes
Foundation Model — 31× in this paper · explained in 3 episodes
Benchmark — 29× in this paper · explained in 2 episodes
Foundation Models — 20× in this paper · explained in 7 episodes
PDE residual — 20× in this paper · explained in 2 episodes
Shear Flow — 19× in this paper
Active Matter — 14× in this paper
Ablation — 11× in this paper · explained in 1 episode
Accuracy — 7× in this paper · explained in 1 episode
Unsupervised PDE-Based Adaptation Objective — 7× in this paper
Transferability — 6× in this paper · explained in 4 episodes
Ablation study — 5× in this paper · explained in 1 episode
Alignment — 5× in this paper · explained in 1 episode
DeepONet — 5× in this paper · explained in 2 episodes
Zero-Shot Transfer — 5× in this paper · explained in 4 episodes
Core — 4× in this paper · explained in 2 episodes
Fidelity — 4× in this paper · explained in 2 episodes
Fine-Tuning — 4× in this paper · explained in 5 episodes
Generalization — 4× in this paper · explained in 13 episodes
AdamW Optimizer — 3× in this paper · explained in 1 episode
Low-Rank Adaptation (LoRA) — 3× in this paper · explained in 9 episodes
Data Augmentation — 2× in this paper · explained in 7 episodes
Information Leakage — 2× in this paper · explained in 2 episodes
Meta-learning — 2× in this paper · explained in 2 episodes
Parameter-Efficient Fine-Tuning — 2× in this paper · explained in 2 episodes
Physics-Informed Neural Networks (PINNs) — 2× in this paper · explained in 6 episodes
Redundancy — 2× in this paper · explained in 2 episodes
Ablation Studies — 1× in this paper · explained in 1 episode
Activation Function — 1× in this paper · explained in 1 episode
Bulk Viscosity — 1× in this paper · explained in 2 episodes
Continual Learning — 1× in this paper · explained in 2 episodes
Downstream Tasks — 1× in this paper · explained in 2 episodes
Hypernetwork — 1× in this paper · explained in 3 episodes
In International Conference on Learning Representations — 1× in this paper
Inverse Problem — 1× in this paper · explained in 2 episodes
Physical Constraints — 1× in this paper · explained in 7 episodes
Transfer Learning — 1× in this paper · explained in 10 episodes

Summary

The paper addresses a fundamental limitation of PDE foundation models: "Pretrained partial differential equation (PDE) foundation models can generalize across different equations, but adapting them to unseen PDE systems typically requires dense solution data, which is often expensive or unavailable. The authors note that existing neural operators performance critically depends on large-scale, high-fidelity datasets, where each training sample may itself require expensive numerical simulation. While physics-informed neural networks (PINNs) incorporate governing equations directly into training, PINNs are typically optimized for individual PDE systems and can be difficult to train, often leading to suboptimal accuracy in practice. Recent PDE foundation models can be adapted to new equations with improved accuracy. However, their adaptation stage remains largely data-driven, typically requiring dense ground-truth solution fields for fine-tuning."

The authors propose "the first unsupervised PDE-based finetuning paradigm that adapts a pretrained PDE foundation model to unseen equations relying exclusively on the governing PDEs and boundary conditions, thereby eliminating the requirement for any interior ground-truth solution fields. The key insight is that while labeled solution data is scarce, the governing equations and boundary conditions are readily available, yet existing approaches do not leverage this information during adaptation."

The framework operates in two stages: "(i) a pretraining stage, which establishes transferable representation across heterogeneous equations, and (ii) an adaptation stage, which finetunes the model on a target equation under the unsupervised PDE-based adaptation objective. The backbone is a neighborhood attention (NA) Transformer, chosen because window-restricted attention inherently aligns with the localized nature of PDE dynamics and scales linearly with the sequence length given a fixed window size. The model is pretrained on six PDEBench subsets covering compressible Navier-Stokes in 1D, 2D, and 3D, shallow water, diffusion-reaction, and incompressible Navier-Stokes using supervised next-step prediction with loss Lpre = (1/(Tinomega))‖Ût+1 − Ut+1‖2. The architecture features a shifted predictive design that imposes dense supervision across every temporal subsequence from 1 to Tin within a single forward, and the unified channel representation maps heterogeneous physical fields into a shared C=18-dimensional representation via deterministic slot assignments with binary masks. Temporal causality is enforced in attention to preclude such information leakage where the predicted output ûτ+1 would directly exploit the value of uτ+1 under an unmasked temporal attention, degenerating the learning objective into a trivial identity mapping."

The UPAO consists of two terms. The boundary condition term supervises the prediction against the boundary reference field on a peripheral band of width 1 grid points along the spatial boundary: LBC = (1/Sb) Σ x∈Sb ‖û t+1(x) − u* t+1(x)‖2. The PDE residual term evaluates the discrete residual R(û) on every grid point of Sp through finite difference schemes, where Diffusive and gradient operators use central differences, while convective nonlinearities such as those in Navier–Stokes or Burgers employ a conservative upwind discretization following [9]. A crucial component is equation-wise normalization: LPDE = sqrt((1/K) Σ k (1/(s2 k Sp)) Σ x∈Sp ‖R k(û)(x)‖2), where s2 k is the residual magnitude evaluated on the input history window Ut. The total loss is L = λBC · LBC + LPDE, with λBC = 104 as default. The pretrained backbone remains frozen while LoRA modules are optimized.

The paper identifies an optimization bottleneck in standard LoRA: the effective rank frequently collapses because no constraint prevents redundancy among the columns of B and rows of A. The resulting rank deficiency restricts the adaptation of less dominant physical quantities. They introduce NSLoRA, which applies Newton-Schulz (NS) iteration to each low-rank matrix independently and rescale them by the original Frobenius norms: ∆W = ‖A‖ F · ‖B‖ F · NS(B) NS(A). The NS iteration X ← aX + b(XXT)X + c(XXT)2X uses coefficients (a,b,c) = (3.4445, −4.7750, 2.0315) with five iterations. Since ‖A‖ F and ‖B‖ F are computed from the learnable matrices at every forward pass, the rescaling ‖A‖ F · ‖B‖ F replaces the fixed α/r with an adaptive per-module amplitude. NSLoRA follows a standard LoRA warm-up phase, switching to the orthogonalized forward path once the validation VRMSE fails to improve for five consecutive epochs, while retaining the warm-up matrices A, B and the AdamW optimizer state across the switch. NS orthogonalization is 1.38× faster end-to-end than SVD orthogonalization (158.19 ms vs 217.75 ms per forward pass).

The framework is evaluated on eleven downstream datasets: four from The Well (Rayleigh–Bénard, Shear Flow, Gray–Scott, Active Matter) and seven analytically derived exact-solution datasets (Burgers 1D/2D, Advection 1D/3D, Taylor–Green 2D, Wave 2D, Advection–Diffusion 2D), with seven of these eleven tasks involve unseen PDEs during pretraining. The metric is VRMSE (variance-normalized root mean squared error).

Key results from Table 1:

our framework reduces VRMSE by a geometric-mean factor of 9.9 compared to PDE-Transformer with standard LoRA under identical UPAO
Compared to the supervised LoRA upper bound on the same backbone, our unsupervised framework remains within 2.5× on seven of the eight 2D benchmarks
Against the four competitive neural operator baselines, our framework outperforms at least one on nine of the eleven datasets
Representative VRMSE values: Ours achieves 0.0060 on TG 2D, 0.0014 on Wave 2D, 0.0097 on AdvDiff 2D, 0.0023 on Burgers 2D, 0.1812 on Rayleigh–Bénard, 0.1112 on Shear Flow, 0.1099 on Gray–Scott, 0.1215 on Active Matter, 0.0112 on 1D Advection, 0.0022 on 3D Advection, 0.0979 on 1D Burgers.

Table 2 evaluates cross-dataset transferability. "Our zero-shot transfer from the PDEBench-pretrained backbone is stronger than every CNextU-Net source checkpoint on Active Matter and Gray–Scott and weaker on the remaining five targets, and the subsequent UPAO finetuning reduces the target VRMSE by 1.4x to 7.0x on the four benchmarks from The Well. Importantly, The magnitude of the UPAO gain is nevertheless backbone-dependent rather than exclusive to our architecture. On CNextU-Net, UPAO improves over zero-shot transfer on 14 of the 26 comparable off-diagonal cells and degrades it on the remaining 12. On TFNO, it improves all 26 cells."

Effect of pretraining (Table 3): "Optimizing all parameters from random initialization under UPAO degrades VRMSE on all eight 2D benchmarks. The degradation reaches more than an order of magnitude on the four exact-solution datasets and remains a factor of 1.2 to 5.3 on the four datasets in The Well."

Effect of NSLoRA vs standard LoRA (Table 4): NSLoRA reduces VRMSE on all eight 2D benchmarks by an average of approximately 4%. The per-channel breakdown (Table 14) reveals that these improvements are concentrated on less dominant physical quantities such as Burgers Vx, Shear Flow tracer, and TG pressure.

Spectral analysis (Table 5): "Despite a nominal rank budget of r = 16, the stable rank of standard LoRA lies between 1.31 and 2.20 on the eight 2D datasets, confirming the rank deficiency that NSLoRA addresses. NSLoRA consistently lifts the stable rank, participation ratio, and entropy rank, raising the stable rank by 5.2% on average and by up to 15.6% on Rayleigh–Bénard."

Normalization strategy (Table 6): The RESCALED NORM configuration, which applies per-sub-equation normalization and sets λBC = 104 / min k s2 k, achieves superior performance on most of the eight datasets compared to NORM and NO NORM.

LoRA vs full finetuning (Table 13): LoRA matches or exceeds full finetuning on half of them, supporting the parameter-efficient design choice. Boundary weight sensitivity (Table 12) confirms λBC = 104 as a strong default.

The paper concludes: "We propose an unsupervised PDE-based finetuning framework that adapts a pretrained PDE foundation model to unseen equations strictly through the governing PDEs and boundary conditions, eliminating the requirement for interior ground-truth fields. We design a NA Transformer that accommodates varying spatial resolutions. We introduce NSLoRA to rebalance learning across heterogeneous physical quantities."

Stated limitations: "The framework requires that the governing equations of the target PDE are known and that boundary observations are available throughout the trajectory, restricting applicability when such priors cannot be obtained. In addition, the PDE residual term provides weaker supervision when the simulation data carries discretization error. Evaluation is restricted to single-step prediction, leaving autoregressive stability untested."

The paper also notes the dataset selection rationale: "Since UPAO relies on finite difference approximations to compute the PDE residual term, it is highly sensitive to the numerical precision of the data. Certain simulation benchmarks, such as PDEArena and PDEGym, often introduce discretization errors that severely compromise the accuracy of computed PDE residuals. Therefore, we utilize datasets synthesized from exact analytical solutions, which guarantee absolute precision at discrete coordinates."

Improvements for AI systems

Unsupervised adaptation to unseen PDEs without labeled data. The AI system can fine-tune a pretrained PDE foundation model to new equations using only the governing PDEs and boundary conditions, eliminating the need for interior ground-truth solution fields. This reduces dependence on expensive numerical simulations and enables adaptation when dense solution data is unavailable or costly. Empirically, this unsupervised fine-tuning stays within 2.5× of the supervised LoRA upper bound on most 2D benchmarks and improves zero-shot transfer VRMSE by 1.4×–7.0× on The Well datasets.
Rank-stabilized parameter-efficient fine-tuning (NSLoRA). The AI system can fine-tune large models with low-rank adapters without suffering from rank collapse, which otherwise restricts adaptation of less dominant physical quantities (e.g., pressure, tracer fields). By applying Newton–Schulz orthogonalization and adaptive per-module amplitude rescaling, the improved system lifts the stable rank by 5.2% on average (up to 15.6%) and achieves an additional 4% VRMSE reduction over standard LoRA across all eight 2D benchmarks. The switch is 1.38× faster than SVD-based orthogonalization, making it practical for large-scale fine-tuning.
Spatial-resolution-agnostic sequence modeling with temporal causality. The AI system can process PDE fields at varying spatial resolutions using a neighborhood attention transformer with shifted predictive design and dense temporal supervision. The enforced temporal causality prevents information leakage, so the model cannot trivially copy the next-step input; instead, it learns physically consistent forward dynamics. The window-restricted attention scales linearly with sequence length, enabling efficient handling of long 2D/3D spatial sequences.
Physics-residual supervision robust to multi-scale equations. The AI system can balance supervision across heterogeneous physical quantities with different magnitudes through equation-wise normalization of PDE residuals. This prevents large-magnitude equations from dominating the loss and improves accuracy on most benchmarks. The residual term uses central differences for diffusive/gradient operators and conservative upwind discretization for convective nonlinearities, making the system applicable to Navier–Stokes, Burgers, advection, and wave equations.
Effective adaptation from a unified cross-equation pretrained backbone. The AI system can transfer representations across heterogeneous PDE families by pretraining on six diverse PDEBench subsets (compressible Navier–Stokes 1D/2D/3D, shallow water, diffusion-reaction, incompressible Navier–Stokes) with a shared channel representation. After adaptation, it outperforms at least one of four competitive neural operator baselines on nine of eleven downstream datasets, including unseen equations (e.g., Burgers, Taylor–Green, Wave, Advection–Diffusion) and simulation-based benchmarks (Rayleigh–Bénard, Shear Flow, Gray–Scott, Active Matter). This enables zero-shot transfer and fast fine-tuning to new equations with minimal data.
Data-quality-aware applicability. The system can avoid performance degradation due to discretization errors by leveraging exact analytical solution datasets for residual-based adaptation. When boundary observations are available throughout the trajectory, the boundary-condition term provides a stable anchor, with a default weight of 104 shown to be robust. The improved system can thus be deployed in settings where simulation data is numerically precise and governing equations are known.
Reduced computational and data cost for scientific machine learning. By replacing supervised dense ground-truth fine-tuning with PDE-based unsupervised losses and LoRA-style parameter-efficient updates, the AI system requires far fewer labeled samples and less memory/bandwidth for training. This makes it feasible to adapt large foundation models to new scientific domains on modest hardware, while still approaching the accuracy of fully supervised fine-tuning.

Abstract

Pretrained partial differential equation (PDE) foundation models can generalize across different equations, but adapting them to unseen PDE systems typically requires dense solution data, which is often expensive or unavailable. To address this limitation, we propose an unsupervised PDE-based finetuning framework that eliminates the need for ground-truth solutions. We first pretrain a neighborhood attention Transformer on diverse time-dependent PDEs spanning varying spatial scales, yielding transferable representations across heterogeneous equations. In the adaptation stage, we construct a physics-based objective using the PDE residual and boundary conditions, and finetune the model on unseen equations via low-rank adaptation (LoRA). To address the uneven learning across physical quantities in standard LoRA, we introduce NSLoRA, a Newton-Schulz orthogonalized variant that rebalances adaptation. Our method achieves performance comparable to supervised LoRA finetuning without requiring any ground-truth solutions, while consistently outperforming competitive neural operator baselines and recent PDE foundation models across heterogeneous PDE benchmarks spanning multiple spatial dimensions.