Hybrid Transformer–Mamba Architecture

• Hybrid Transformer–Mamba architecture integrates Transformer self-attention with state-space Mamba blocks to capture short- and long-range spatiotemporal dependencies.
• It employs a two-stage training regimen featuring physics-informed fine-tuning that drastically reduces PDE residuals while enhancing data fidelity.
• The model scales to arbitrary unstructured grids and query points, making it adaptable for diverse applications in physical field generation.

The hybrid Transformer–Mamba architecture is a class of deep neural network models that combine the expressivity of Transformer mechanisms (self-attention) with the state-space sequence modeling afforded by Mamba-type layers. This fusion is designed to efficiently capture both short- and long-range spatiotemporal dependencies on unstructured grids, particularly in the generation of physical fields. In the context of spatiotemporal field generation, the HMT (Hybrid Mamba–Transformer) backbone integrates a temporal state-space model ("Mamba block") for autoregressive propagation of global latent states with Galerkin-style Transformer blocks for spatial encoding/decoding. The architecture supports physics-informed fine-tuning via a residual-based correction module, directly reducing physical equation errors. Key innovations include permutation-invariant fusion of features, a point-query mechanism for local PDE residuals, a two-stage training regimen, and the introduction of the MSE-R metric to evaluate both data fidelity and physical realism (Du et al., 16 May 2025).

1. Architectural Composition and Data Flow

The HMT backbone processes unstructured spatial domains with temporal dynamics by integrating a Mamba state-space block and a Galerkin Transformer block:

• Inputs: Boundary/domain point set $X_{BD} = \{x_i\}$ with a binary identifier $Id_i \in \{0,1\}$, initial physical field $\phi(x_i, t_0)$, and arbitrary query positions $X_Q$.
• Spatial Encoding: Encoder $\mathcal{E}_1$ applies MLPs, k-NN local embedding, and Galerkin self-attention to produce global point features $G_0 \in \mathbb{R}^{N_{BD} \times N_g}$.
• Temporal Propagation: Mamba block $M$ pools $G_0$ to $z_0 = \mathrm{MaxPool}(G_0)$, then generates $z_i = \mathrm{Mamba}(z_0, ..., z_{i-1})$ autoregressively.
• Query Encoding and Fusion: Encoder $\mathcal{E}_2$ maps $X_Q$ to feature $H_Q$; decoder $D$ fuses $G_0$ and $z_i$ into per-point latents $H_i$, then applies Galerkin cross-attention to project to ${\phi}(X_Q, t_i)$.

Mathematically:

• $G_0 = \mathcal{E}_1(X_{BD}, Id, \phi(X_{BD},t_0))$
• $z_0 = \mathrm{MaxPool}(G_0)$
• $z_i = \mathrm{Mamba}(z_0, ..., z_{i-1})$
• $H_Q = \mathcal{E}_2(X_Q)$
• $H_i = \mathrm{Fuse}(G_0, z_i)$
• $$\hat{\phi}(X_Q, t_i) = \mathrm{FFN} \circ \mathrm{CrossAttn}(H_Q, H_i, G_0)$$

Coupling is achieved via residual connections across spatial (but not temporal) dimensions, with $G_0$ reused at every cross-attention block for stability.

2. Component Roles: Mamba and Galerkin-Transformer

• Mamba Block: Functions as a state-space sequence model for temporal feature propagation, integrating all previous latent vectors for smooth evolution. It autoregressively generates global state vectors $\{z_i\}_{i=1}^T$ based on pooled initial features. Its linear memory and compute profile allow scaling to long time horizons.
• Galerkin-Transformer:
  • Encoder: Employs global attention plus local geometric embeddings to model the spatial relationships among unstructured points.
  • Decoder: Utilizes cross-attention for mapping arbitrary queries to latent fields, ensuring permutation invariance and flexibility in output locations.

The synergy between Mamba (for memory-efficient temporal modeling) and Galerkin-Transformer (for permutation-invariant spatial attention) yields a backbone capable of handling physics contexts where point sets and query times are dynamic and potentially sparse.

3. Physics-Informed Fine-Tuning and Self-Supervision

After general pretraining, the HMT model often exhibits nontrivial physical equation residuals. Fine-tuning is performed via a physics-informed block:

• Residual Computation: For each query point, finite differences recompute spatial ($\partial_x \phi$) and temporal ($\partial_t \phi$) gradients; in Navier–Stokes problems, continuity ($R_1$) and momentum ($R_2$) residuals are explicitly constructed.
• Correction Module ($\mathcal{E}_3$): Residuals are encoded into correction vectors $\{\Delta z_i\}$, which are added to the temporal latents: $\check{z}_i = z_i + \Delta z_i$.
• Decoding and Loss: The refined latents $\check{z}_i$ are decoded via a separate FFN$_\text{FT}$; only $\mathcal{E}_3$ and FFN$_\text{FT}$ are trainable. The composite self-supervised loss balances field reconstruction against residual magnitude:

  $$L_2 = A_d \sum_{i=1}^T \| M_i \odot [\hat{\phi} - \tilde{\phi}] \|^2 + A_R \sum_{i=1}^T \| R(\tilde{\phi}) \|^2$$

where $M_i$ is a random masking matrix, $A_d$ and $A_R$ are task-tuned constants.

4. Training Regimen and Point-Query Gradient Evaluation

The training procedure involves:

• Stage 1: Data-driven pretraining ($$L_1 = \sum_{i,t} \| \hat{\phi}_i(t) - \phi_i^{GT}(t) \|^2$$) with full backbone optimization.
• Stage 2: Physics-informed fine-tuning (loss $L_2$), with backbone frozen and only correction layers trainable.

Gradient and residual computation leverage a point-query mechanism: for each $x_Q$, spatial gradients are estimated using neighbor offsets $\pm \Delta x$ and temporal gradients using time-difference queries. This facilitates efficient batch gradient computation on irregular meshes.

5. Quantitative Evaluation: MSE-R Metric and Performance

Accuracy and realism are jointly assessed via the MSE-R metric:

• $$\mathrm{MSE} = \frac{1}{N_QT} \sum_{i,t} \| \hat{\phi}_i(t) - \phi_i^{GT}(t) \|^2$$
• $$R = \frac{1}{N_QT} \sum_{i,t} \| \mathrm{Residuals}[\hat{\phi}_i(t)] \|$$

Empirically, physics-informed fine-tuning achieves up to two orders-of-magnitude reduction in PDE residuals. Representative dataset scores:

Model	Airfoil	Cylinder	Aneurysm	Acoustic	Simple-Car
GEO-FNO	0.6361	0.0844	0.02115	0.8573	0.1583
GINO	0.7100	0.0882	0.00091	0.9707	0.0898
TRANSOLVER	0.4841	0.0358	0.00019	0.4302	0.0620
HMT (no FT)	0.4432	0.0260	0.00015	0.4123	0.0652
HMT + FT	0.3917	0.0235	0.00008	0.4081	0.0648

HMT outperforms all baselines in four of five tasks before fine-tuning; FT further reduces MSE by 5–12%, and decreases residuals $R$ by up to 100× (e.g., $10^{-1} \rightarrow 10^{-3}$ for Airfoil). Sparse sampling experiments show up to 25% MSE reduction post-FT.

6. Model Scalability and Applicability

The architecture accommodates:

• Inputs on arbitrary unstructured grids, with boundary/interior point distinction.
• Arbitrary spatial or temporal queries, with permutation-invariant spatial decoding.
• Generalizable physics-informed corrections applicable to diverse PDE systems.

Datasets include Airfoil (2D), Cylinder (2D), Aneurysm (3D), Acoustic (2D grid), and Simple-Car (3D static), with backbone transformer/mamba depths and hidden sizes tunable per problem; finite-difference step sizes are dataset-specific.

7. Impact and Theoretical Implications

The Hybrid Transformer–Mamba architecture with physics-informed fine-tuning demonstrates:

• Effective suppression of physical inconsistency in data-driven spatiotemporal generative models.
• Scalable handling of long-term dynamics and arbitrary spatial configurations via SSM-backed temporal propagation and flexible attention-based spatial querying.
• A viable paradigm for augmenting neural field generators with explicit PDE residual loss, driving models toward physical law conformity.
• Quantitative evidence that hybridization is critical for both fidelity and efficiency, outperforming pure Transformer, pure Mamba, and contemporary geometric operator baselines in both data accuracy and physical realism, with lower computational overhead (Du et al., 16 May 2025).

This design is well-suited for advancing general-purpose generators of physical fields, especially in scientific domains requiring strong guarantees of physical law adherence and support for nonuniform, query-driven spatial sampling.