Technical Summary: Discovering Physical Directions in Weight Space

Problem Statement

Neural operators have demonstrated significant success in creating scalable surrogate solvers for parametric partial differential equations (PDEs) through large-scale pretraining and in-context adaptation. However, a fundamental question remains regarding the nature of adaptation within a physical family: when a shared operator is fine-tuned to multiple regimes (e.g., varying viscosity or diffusivity) within the same physical family, do the resulting weight updates merely define isolated "regime experts," or do they reveal a reusable, structured physical geometry in weight space?

Existing adaptation strategies face limitations. Conditional operators require architectures and training corpora that cover the entire physical state space with high fidelity. Fine-tuning into multiple regime-specific experts creates fragmented checkpoints without a principled method to compose them for unseen regimes. Generic model-merging techniques (e.g., model soups, task arithmetic) typically treat downstream tasks as unordered, selecting merge coefficients based on validation loss rather than physically meaningful coordinates. Consequently, these methods often fail to generalize effectively to out-of-distribution (OOD) regimes, particularly when extrapolation is required.

Methodology: Calibration-Conditioned Merge (CCM)

The paper proposes Calibration-Conditioned Merge (CCM), a post-hoc coordinate readout method that treats same-anchor endpoint experts not as isolated checkpoints, but as finite-difference probes of a local physical direction in weight space.

1. Theoretical Framework

The method assumes that for a parametric PDE family, downstream regimes are ordered points along an underlying physical axis. Starting from a shared family anchor $\theta_0$, the authors fine-tune two endpoint experts: $\theta_{low}$ (low regime) and $\theta_{high}$ (high regime).

The weight updates relative to the anchor are decomposed into two components:

• Shared Adaptation ($\Delta_+$): The mean of the endpoint residuals, capturing solver corrections shared across the family.
  $$\Delta_+ = \frac{\Delta_{low} + \Delta_{high}}{2}$$
• Signed Physical Direction ($\Delta_-$): The signed contrast of the residuals, capturing the variation oriented along the physical axis.
  $$\Delta_- = \frac{\Delta_{high} - \Delta_{low}}{2}$$

This decomposition defines a one-dimensional coordinate line in weight space:
$$\theta(\alpha) = \theta_0 + \Delta_+ + \alpha\Delta_-$$
where $\alpha = -1$ and $\alpha = +1$ recover the low and high experts, respectively. Static averaging corresponds to $\alpha=0$.

2. The CCM Algorithm

CCM selects a target composition coordinate $\hat{\alpha}$ based on available test-time information, deploying a single merged checkpoint $\hat{\theta}_\lambda = \theta(\hat{\alpha})$ without further fine-tuning, routing, or ensembling. Three selection strategies are proposed:

• CCM-Coord: Uses known physical metadata (e.g., normalized diffusivity) to directly map the physical coordinate $s(\lambda)$ to $\hat{\alpha}$.
• CCM-Scale: Accounts for potential mismatches between the physical coordinate scale and the weight-space scale by learning a global scaling factor $\gamma$ on validation data ($\hat{\alpha} = \gamma s(\lambda)$).
• CCM-Prefix: When metadata is unreliable or missing, the method infers $\hat{\alpha}$ by minimizing the loss over a short observed rollout prefix (e.g., the first $K$ time steps) against a finite bank of candidate coordinates.

Key Contributions

1. Physical-Coordinate View of Fine-Tuning: The paper formulates same-anchor endpoint experts as finite-difference probes of a local solver trajectory. It demonstrates that endpoint-anchor residuals encode reusable physical structure (a signed direction) rather than arbitrary checkpoint drift.
2. Calibration-Conditioned Merge (CCM): A novel method for training-free transfer across PDE regimes. It composes a single checkpoint using physical metadata, calibrated scales, or short rollout prefixes, avoiding the need for target-regime training or dynamic routing.
3. Empirical Evidence of Calibratable Directions: Through extensive experiments, the authors validate that endpoint fine-tuning reveals a signed physical direction useful for both interpolation and extrapolation. This is evidenced by the strong correlation between the optimal merge coordinate and independent physical parameters, and the superior performance of CCM over static averaging in OOD regimes.

Experimental Results

The method was evaluated on three distinct PDE families: reaction–diffusion dynamics, viscosity-parameterized 2D Navier–Stokes equations, and radial dam-break dynamics.

• Decomposition Validity: Experiments on controlled axes showed that while static averaging ($\alpha=0$) is effective when shared adaptation dominates, it fails when target regimes require a directed shift along the physical axis. Introducing a non-zero $\alpha$ restores endpoint-specific physics.
• Coordinate Alignment: In reaction–diffusion and Navier–Stokes tasks, the diagnostic oracle coordinate (selected via rollout loss) aligned almost perfectly with the normalized physical coordinate ($r=1.00$ and $r=0.99$, respectively), confirming the physical interpretability of the weight-space direction.
• Extrapolation Performance: CCM achieved its most significant gains in extrapolative (OOD) regimes where static averaging attenuates endpoint-specific physics:
  • Reaction–Diffusion: Reduced OOD rollout error by 54.2% relative to the family anchor.
  • Navier–Stokes: Reduced OOD rollout error by 42.8%.
  • Radial Dam-Break: Reduced OOD rollout error by 13.8% (using prefix-based selection).
• Robustness: The gains persisted across different model scales (FNO small/base/large), backbone architectures (including DPOT-style), and random seeds. "Wrong-sign" controls (flipping the direction) significantly increased error, confirming that the improvement relies on the correct physical orientation, not arbitrary interpolation.

Significance and Claims

The paper claims that same-anchor neural PDE experts form a continuous family rather than isolated checkpoints. By decomposing endpoint residuals, the authors show that neural solvers encode continuous physical variation as interpretable weight-space directions.

The significance of this work lies in operationalizing the hypothesis that endpoint fine-tuning exposes a calibratable physical direction. This allows for:

1. Training-free regime transfer: Adapting to new physical parameters without retraining or complex routing mechanisms.
2. Diagnostic capability: Serving as a tool to understand how neural solvers represent physical parameters, distinguishing between shared solver corrections and regime-specific physical variations.
3. Extrapolation: Providing a principled method to extrapolate beyond the training distribution by traversing the signed physical direction, a task where naive averaging fails.

The authors remain modest regarding limitations, noting that CCM presupposes a common local geometry and compatible architectures. The method does not provide a universal OOD guarantee, and the coordinate assumption may break down if endpoint gaps are too large or if the checkpoint manifold exhibits high curvature. Furthermore, the prefix-based approach requires a brief observation horizon, precluding true zero-shot deployment in all scenarios.