🤖 machine learning

PIMSM: Physics-Informed Multi-Scale Mamba for Stable Neural Representations under Distribution Shift

The paper proposes Physics-Informed Multi-Scale Mamba (PIMSM), a state-space architecture that aligns model discretization with physical timescales to overcome temporal kernel mismatch, thereby achieving robust and stable neural representations for scientific time series under severe distribution shifts.

Original authors: Sangyoon Bae, Shinjae Yoo, Jiook Cha

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sangyoon Bae, Shinjae Yoo, Jiook Cha

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

The Big Problem: The "One-Size-Fits-All" Trap

Imagine you are trying to learn a new language. Most AI models today are like students who memorize a specific textbook perfectly. They are great at answering questions from that exact book (the "training data"). But the moment you ask them a question from a slightly different book, or in a different accent, or with fewer words, they get confused.

In the scientific world, this happens with time-series data (like brain scans or weather reports).

Brain scans (fMRI): The brain doesn't just have one speed. It has fast electrical spikes and slow, rolling waves of activity.
Weather: The atmosphere has fast wind gusts and slow seasonal shifts.

Current AI models often treat these complex, multi-speed signals as a single, messy pattern to be memorized. They don't respect the fact that nature operates on different "time scales" simultaneously. When the conditions change (e.g., a different MRI machine, a different weather station, or a shorter recording), these models fail because they learned the wrong "rhythm."

The Solution: PIMSM (The "Musical Conductor")

The authors propose a new model called PIMSM (Physics-Informed Multi-Scale Mamba).

Think of a standard AI model as a drummer who only knows one beat: Boom-bap, Boom-bap. No matter if the song is a fast jazz tune or a slow ballad, the drummer tries to force that one beat onto it. It works okay for the practice song, but fails when the tempo changes.

PIMSM is like a Musical Conductor who listens to the orchestra first.

Listening (The Spectrum): Before the music starts, PIMSM looks at the "sheet music" (the frequency spectrum) of the signal. It identifies the "knees" or turning points where the music shifts from fast notes to slow notes.
Assigning Roles: Instead of using one beat, PIMSM assigns different "drummers" (or memory settings) to different parts of the song.
- One drummer handles the fast, high-pitched notes (fast updates).
- Another handles the slow, deep bass notes (long memory).
Anchoring to Reality: Crucially, PIMSM ties these drummers to real-world time units. If the data is measured in "seconds" or "hours," the model's memory is calibrated in seconds or hours. It doesn't just learn abstract numbers; it learns physical time.

How It Works (The "Physics-Informed" Part)

The paper argues that nature follows specific physical rules.

The "Knee" Frequencies: Just as a road might have a sharp curve (a "knee") where you must slow down, the data has specific frequencies where the behavior changes. PIMSM finds these curves automatically.
The Map: It uses these curves to build a map. It tells the AI: "For this part of the signal, remember the last 5 seconds. For that part, remember the last 5 hours."
No Guessing: Unlike other models that guess how to remember things, PIMSM is "physics-informed." It uses the actual shape of the data's energy to decide how to remember it.

The Experiments: Testing in the "Wild"

The authors tested this idea in two very different worlds to see if it was robust (sturdy) when things changed.

1. The Brain (fMRI Data)

The Challenge: They took brain scans and artificially chopped off the beginning of the recording (simulating a situation where you only have a few seconds of data). They also tried to train the model on "resting" brains and test it on "active" brains.
The Result: Standard models got confused and their internal "understanding" of the brain drifted apart. PIMSM stayed stable. It was like a navigator who could still find the destination even if you only showed them a tiny fragment of the map. It kept its "memory" aligned with the brain's actual physical rhythms.

2. The Weather (Weather-5K)

The Challenge: They trained the model on weather data from some cities and asked it to predict weather for completely new cities it had never seen before (out-of-distribution).
The Result: PIMSM was the most accurate at predicting temperature, wind, and pressure for these new locations. It didn't need to be retrained for the new cities; its understanding of "fast wind" and "slow pressure changes" was universal enough to work anywhere.

Why This Matters (The Takeaway)

The paper claims that for scientific AI to be truly useful, it can't just be a "pattern matcher." It needs to respect the physics of time.

Old Way: "I will memorize the pattern of this specific dataset." (Fails when the dataset changes).
PIMSM Way: "I will learn the physical rhythm of this signal, so I can adapt to new speeds, new locations, and new conditions."

By anchoring the AI's memory to the actual physical timescales of the data (like seconds for a heartbeat or hours for a storm), the model becomes much more stable. It doesn't just fit the data; it understands the structure of the data, allowing it to survive when the real world throws it a curveball.

Summary Analogy

If a standard AI is a parrot that repeats what it heard in the classroom, PIMSM is a musician who understands the theory of rhythm. If the parrot hears a song played slightly faster, it gets lost. The musician, however, understands that "fast" and "slow" are relative concepts tied to the beat, so they can play the song correctly no matter who is conducting or what instrument is used.

Technical Summary: PIMSM – Physics-Informed Multi-Scale Mamba for Stable Neural Representations under Distribution Shift

1. Problem Statement

Scientific foundation models aim to learn reusable representations across varying datasets, acquisition protocols, and deployment domains. However, current sequence backbones often treat scientific temporal structures as unconstrained patterns to be fitted, ignoring the physical reality that natural dynamical systems (e.g., neural and atmospheric time series) are organized by interacting processes across multiple physical timescales.

The authors identify a specific failure mode termed temporal kernel mismatch. This occurs when a model fits in-distribution dynamics using an effective memory policy that is not anchored to the signal's physical timescales. Consequently, when deployment conditions shift (e.g., changes in temporal resolution, sampling rate, brain state, or geographic domain), the model's learned memory policy fails to preserve task-relevant dynamics, leading to representation drift and degraded transfer performance. The core challenge is to design an architecture that preserves multiscale structure to ensure representations remain reusable under shift, rather than merely optimizing for in-distribution correlation fitting.

2. Methodology: Physics-Informed Multi-Scale Mamba (PIMSM)

PIMSM is a state-space architecture designed to align temporal scales with physical signal properties. It replaces fully unconstrained temporal parameterization with a spectrum-guided hierarchy.

2.1 Core Mechanism: Spectrum-to-Discretization Mapping

The method leverages the observation that neural and atmospheric signals often exhibit piecewise power-law spectral structures.

Spectral Analysis (SpectralHyperNet): The model estimates the Power Spectral Density (PSD) of the input time series and fits a piecewise power-law form: $P(f) \propto f^{-\beta_k}$ for frequency regimes separated by "knee frequencies" ( $f_k$ ). These knees represent spectral transition points between different temporal regimes.
Scale-Specific Parameterization: The estimated knee frequencies define $K$ frequency bands. For each band, an energy-weighted representative frequency ( $f_{c,k}$ ) is computed.
Acquisition-Unit Anchoring: These representative frequencies are mapped to scale-specific discretization parameters ( $\Delta_k$ $Δ_{k}$ ) for the State Space Model (SSM) heads. Crucially, $\Delta_k$ $Δ_{k}$ is expressed in acquisition time units (e.g., TRs for fMRI, hours for weather).
- Higher frequency bands map to larger $\Delta$ (faster updates).
- Lower frequency bands map to smaller $\Delta$ (longer memory).
- A hierarchical constraint ( $\Delta_1 \ge \Delta_2 \ge \dots$ ) enforces temporal ordering.
Regularization of Compensation: To prevent the continuous-time state matrix $A$ from compensating for the fixed $\Delta$ (which would render the physical anchoring ineffective), PIMSM applies a scale regularizer on $A$ , keeping its decay scale near a physically motivated initialization.

2.2 Architecture

PIMSM integrates these parameters into a Mamba-2 backbone:

Multi-Scale Heads: The model employs multiple SSM heads grouped into $K$ scale groups (default $K=3$ ). Heads within the same group share the discretization parameter $\Delta_k$ but maintain independent state transitions.
Cross-Scale Attention: A lightweight cross-scale attention module fuses the scale-specific states, modeling interactions between different temporal regimes.
Training Objective: The total loss includes the task loss, spectral fit loss (to ensure accurate knee estimation), seam continuity loss, and regularization terms for $\Delta$ (anchoring to acquisition units) and $A$ (preventing scale drift).

3. Key Contributions

The paper outlines three primary contributions:

Temporal-Scale Alignment as a Requirement: The authors formulate "temporal kernel mismatch" as a diagnostic for why scientific time-series representations fail under distribution shifts. They argue that preserving physical timescales is a prerequisite for reusable representations.
Physics-Informed Multi-Scale Parameterization: PIMSM is proposed as a structured architecture that derives scale-specific kernels from spectrum-estimated knee frequencies, expresses discretization steps in acquisition units, and regularizes the state matrix to ensure these scales remain interpretable and transferable.
Multi-Axis Robustness Evaluation: The model is evaluated across diverse structured shifts in both neural (fMRI) and meteorological (Weather-5K) domains, demonstrating improved robustness without modality-specific adaptation or shift-specific supervision.

4. Experimental Results

The authors evaluated PIMSM against parameter-matched baselines (Mamba2, Transformer, 1D-CNN) across four axes of shift:

Axis (i) Temporal-Context Truncation (HCP Motor fMRI):
- Setup: Models trained on full motor blocks but evaluated on severely truncated inputs (first 2 or 3 TRs).
- Result: PIMSM achieved the smallest accuracy gap between full-block and truncated inputs and the highest representation similarity (CKA and dCor) between the two conditions. Ablations showed that spectrum-derived initialization was critical; learnable or random $\Delta$ values resulted in lower stability.
Axis (ii) Low-Resource Generalization (HCP Motor fMRI):
- Setup: Training with 1% to 100% of labeled data.
- Result: PIMSM significantly outperformed baselines in the 1% data regime, indicating superior sample efficiency when supervised signals are extremely limited.
Axis (iii) Brain-State Shift (Resting-State $\to$ Task):
- Setup: Models trained on resting-state fMRI and evaluated on task-evoked fMRI without adaptation.
- Result: PIMSM achieved the highest accuracy and substantially lower variance across seeds compared to baselines, suggesting reduced sensitivity to the shift in acquisition state.
Axis (iv) Cross-Domain Shift (Weather-5K):
- Setup: Spatial out-of-distribution (OOD) forecasting on held-out weather stations.
- Result: PIMSM attained the lowest Mean Absolute Error (MAE) across all 24 variable-horizon endpoints. It reduced MAE by an average of 15.2% compared to the best non-PIMSM baseline, with gains increasing at longer forecasting horizons (up to 17.9% at 168h).

5. Significance and Claims

The paper claims that PIMSM demonstrates that temporal-scale alignment serves as a practical inductive bias for scientific foundation models. By anchoring the model's memory policy to the signal's physical organization (via spectrum-derived knees and acquisition units), the architecture preserves the geometry of representations under severe distribution shifts.

The authors emphasize that these gains are not merely accuracy improvements but evidence that the model can maintain reusable dynamics across different contexts (e.g., from resting to task states, or from one geographic region to another). The work suggests that for scientific signals, preserving structure is as important as fitting correlations, particularly when deployment conditions (temporal resolution, labels, or domain) change.

The study acknowledges limitations, noting that it assumes a small number of spectrum-derived knees suffices to summarize temporal organization and that the HCP experiments test input truncation rather than scanner-level TR changes. However, the results support the hypothesis that physics-informed temporal parameterization can serve as a reusable design principle for models operating under real-world constraints like short recordings, scarce labels, and spatial OOD scenarios.