Poseidon-B — Joint fine-tune (Cylinder + Kelvin–Helmholtz)

Paper: Measuring and Mitigating Forgetting when Extending a PDE Foundation Model to Simulated Flow Regimes (Representations for the Physical Sciences Workshop @ NeurIPS 2026)

A fine-tuned checkpoint of Poseidon-B (scOT, ~158M parameters) extended to two flow regimes that are absent from Poseidon's pretraining distribution, added simultaneously in a single training run:

1. Cylinder — wall-bounded, incompressible flow past a circular cylinder (vortex shedding), Re ≈ 62–180.
2. KHC — a compressible Navier–Stokes Kelvin–Helmholtz double shear layer.

The goal of this checkpoint is to add these new regimes without catastrophic forgetting of Poseidon's original families. It is trained jointly on the two new regimes together with a small replay buffer of two pretraining families, so that after adaptation the original families remain near (or below) their original error while the two new regimes are fit well.

Research model. Like the base model, this is intended for research use only.

Model description

• Architecture: ScOT (Swin-V2-style multiscale operator transformer, swinv2 backbone), identical to Poseidon-B — the new regimes were added without changing the architecture or the input/output channels.
• Input/output: a 4-channel physical state (ρ, u, v, p) on a 128 × 128 grid.
• Conditioning: lead-time conditioned. The time input is a normalized lead time t = k / (T − 1), where k is the frame step and T the trajectory length on the 21-frame cadence the model was trained on (see Usage).
• Solid-body encoding: the cylinder is written into the density channel — ρ = 1 in fluid, ρ = 0 in solid — so the geometry enters through an existing channel and doubles as a geometry mask. No extra input channels are added.
• Parameters: ~158M (full fine-tune; all weights trainable).

Key config values (see config.json): image_size=128, num_channels=4, num_out_channels=4, embed_dim=96, depths=[8,8,8,8], num_heads=[3,6,12,24], window_size=16, use_conditioning=true.

Training

Full fine-tuning from Poseidon-B with a joint data mixture (new regimes at weight 1.0, replay families at weight 0.25):

Hyperparameters:

• Optimizer: AdamW, lr 2e-4, weight decay 1e-5
• Schedule: OneCycle, 10% warm-up
• Loss: L1, gradient clipping at 1.0
• Batch size: 192, 31 optimizer steps/epoch × 50 epochs = 1,550 updates (fixed budget)
• Mixed precision (AMP), seed 0
• All trajectories resampled to a common length of T = 21 frames; frame step k ∈ [1, 12]
• Checkpoint selection: lowest validation error on 30 held-out cylinder trajectories

Per-family normalization is used: each family is normalized by its own channel statistics (the exact stats used are shipped in norm_stats_used.json). The binary geometry channel is left un-normalized.

Evaluation

Reported as per-channel NRMSE (RMSE divided by the channel's value range), with a no-change persistence baseline for reference. All numbers below are 3-seed means.

New regimes

Both are well below their persistence baselines.

Retention on Poseidon's pretraining families (factor relative to the unchanged base model; < 1.0 = improved, > 1.0 = degraded):

All four tracked families land between 0.63× and 1.16× their original error, i.e. the two new regimes were added while the original families were largely retained or improved — including families that were never placed in the replay buffer.

Rollout stability (cylinder). Autoregressive rollout over held-out trajectories stays stable: median NRMSE grows only from ~0.007 to ~0.017 over 20 steps, with field correlation above 0.9 at every step (median ≈ 0.985 at step 20).

Intended use & limitations

• Intended use: research on operator learning, PDE surrogates, transfer/retention, and fast approximate rollouts for the two supported regimes (cylinder wake and compressible KH shear layers) and Poseidon's original families.
• Resolution: all fields are on a 128 × 128 grid; claims are restricted to wake-scale dynamics (fine scales are smoothed).
• Extrapolation: the cylinder holdout sits adjacent to the training band; accuracy degrades smoothly with distance in Reynolds number and is only validated up to ~33% beyond the training maximum, and on 2D flow (below the onset of 3D wake instability).
• Rollout horizon: validated within ~4 shedding periods. The frame-resampled model has a minimum usable lead of 1/20 of the trajectory window; shorter leads degrade sharply.
• Not a general CFD solver. It is a regime-specific surrogate, not a replacement for a PDE solver outside the tested distributions.

Usage

The checkpoint is saved in Hugging Face save_pretrained format (config.json + pytorch_model.bin) and loads with the scOT model class. Two conventions matter for getting the reported accuracy:

1. Normalize per channel with the target family's statistics (shipped in norm_stats_used.json), and de-normalize the output with the same statistics.
2. Frame cadence. The model was trained on trajectories resampled to T = 21 frames. The time input is the normalized lead t = k / (T − 1) = k / 20, where k is the frame step on that 21-frame cadence. A raw trajectory at a different frame count must first be resampled to 21 frames (e.g. by evenly spaced index selection) for the lead time to be meaningful. k = 2 (t = 0.1) is the canonical lead this checkpoint was selected on.

import numpy as np
import torch
from scOT.model import ScOT  # from the Poseidon / scOT codebase

model = ScOT.from_pretrained("rwmasood/poseidon-b-joint-cyl-kh").eval().cuda()

# Cylinder channel statistics used at training time (from norm_stats_used.json).
# Channels are (rho, u, v, p). For the cylinder regime rho encodes the body: 1=fluid, 0=solid.
mean = torch.tensor([0.992431640625, 0.9874249696731567,
                     -6.629727431572974e-05, 0.9546216726303101]).view(1, 4, 1, 1).cuda()
std  = torch.tensor([0.08666647970676422, 0.1872677057981491,
                     0.15668894350528717, 0.11970974504947662]).view(1, 4, 1, 1).cuda()

# A raw trajectory of physical fields (rho, u, v, p), shape (n_frames, 4, 128, 128).
traj = np.load("cylinder_case.npy")

# Resample to the 21-frame cadence the model was trained on.
T = 21
idx = np.linspace(0, traj.shape[0] - 1, T).round().astype(int)
traj = torch.from_numpy(traj[idx]).float().cuda()

# Predict frame i + k from frame i, with lead time k / (T - 1).
i, k = 0, 2
x = (traj[i:i + 1] - mean) / std
t = torch.full((1,), k / (T - 1), device=x.device)   # 2 / 20 = 0.1

with torch.no_grad():
    pred_norm = model(pixel_values=x, time=t).output   # scOT returns a ScOTOutput with `.output`

pred = pred_norm * std + mean                          # de-normalize to physical units

For the KHC regime, swap in that family's statistics instead:

# KHC channel statistics (rho, u, v, p), from norm_stats_used.json
mean = torch.tensor([1.4987393617630005, 0.15815676748752594,
                     -0.00015812707715667784, 11.184250831604004]).view(1, 4, 1, 1)
std  = torch.tensor([0.32741427421569824, 0.6542320847511292,
                     0.4637199640274048, 3.0745689868927]).view(1, 4, 1, 1)

Autoregressive rollout: feed the model's de-normalized prediction back in as the next input.

Sanity check: on the 30 held-out cylinder trajectories at lead 2 (21-frame cadence), this checkpoint gives a mean per-channel NRMSE of ≈ 0.017, matching the reported cylinder number.

Files

• config.json — model configuration (ScOT / swinv2)
• pytorch_model.bin — model weights
• norm_stats_used.json — per-family normalization statistics (cylinder, KHC, CE, RP)

Environment notes

• transformers==4.29.2 (the version the base model config targets).
• The scOT package is required to instantiate the model class.
• On recent NVIDIA GPUs (Blackwell / sm_120), install a matching CUDA build of PyTorch.

Base model & attribution

Fine-tuned from camlab-ethz/Poseidon-B Cylinder and KHC training data were generated with the PyFR solver.

Citation

If you use this checkpoint, please cite:

@inproceedings{yourkey2026,
  title     = {Measuring and Mitigating Forgetting when Extending a PDE Foundation Model to Simulated Flow Regimes},
  author    = {Wasif Masood and Shaoib Noor},
  booktitle = {Representations for the Physical Sciences (RPS) Workshop at NeurIPS 2026, Paris, 2026.},
  year      = {2026},
  url       = {https://openreview.net/forum?id=XwMrQSsWF3}
}

License

Released under CC-BY-NC-4.0, inheriting the base model's license (research/non-commercial use).