Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Axial Neural Networks for Dimension-Free Foundation Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Axial Neural Networks (set-based and graph-based), permutation-equivariant over tensor axes, enabling dimension-agnostic architectures, and converts PDE foundation models (CViT, MPP) into dimension-agnostic variants.
Problem
Training foundation models on physics data (PDE solutions) with varying spatial dimensionalities is hard because existing architectures fix a maximum dimension, use separate encoders per dimension, or are tied to 2D patchify operations, and prior dimension-agnostic methods are computationally expens
Approach
Treat tensor axes as set or graph elements and share parameters across axes via axis-permutation equivariant layers (set-based SXNN via permutation aggregation; graph-based GXNN via lifting over axis pairs plus GNN-style subsequent layers), so one set of weights handles any tensor rank.
Key results
- On the GP kernel toy task, XCNN (graph-based) shows superior test accuracy across 2D-5D inputs, while CNN-3D performs poorly even on 3D data.
- Set-based SXCNN with same depth/width as GXCNN has 6x fewer parameters and 80% lower wall-clock time but degrades in high dimensions.
- Fine-tuning X-MPP pretrained on 1D, 2D, and 3D data outperforms MPP trained only on 2D on 2D PDEs; X-MPP accelerates fine-tuning on unseen 1D and 3D PDEs with little data (except BG 1D).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use graph-based axial convolutions instead of a dimension-specific 3D CNN | GP kernel prediction · GP kernel prediction, 2D–5D | test accuracy | Helped |
| Use a dimension-specific 3D CNN with zero-padding or flattening for other dimensions | GP kernel prediction · GP kernel prediction, higher dimensions including 3D | test accuracy | Hurt |
| Use graph-based axial convolutions rather than set-based axial convolutions (ablation) | XCNN · GP kernel prediction, high dimensions | accuracy | Helped |
| Increase the depth and width of set-based axial CNN (ablation) | SXCNN · GP kernel prediction, high dimensions | accuracy | Helped |
| Train with both 1D and 2D PDEs rather than only 2D PDEs (ablation) | MPP · PDE foundation model pretraining | NRMSE | Hurt |
| Pretrain on 1D, 2D, and 3D PDEs rather than only 2D PDEs | X-MPP · 2D PDE fine-tuning | NRMSE | Helped |
| Use dimension-agnostic X-MPP rather than dimension-specific MPP for unseen-dimension fine-tuning | MPP · Unseen 1D and 3D PDE fine-tuning, except BG 1D | NRMSE | Helped |
| Use dimension-agnostic X-MPP instead of inflated MPP layers | MPP · CFD 3D fine-tuning | NRMSE | Helped |
| Replace InstanceNorm with LayerNorm (ablation) | MPP and X-MPP · Multidimensional PDE training | NRMSE | No clear effect |
Surprising findings
- Naively pretraining MPP jointly on 1D and 2D data significantly degrades performance, indicating failure to learn a unified representation across dimensions.
- The set-based axial convolution behaves like a convolution with a symmetric kernel, limiting expressivity and motivating the graph-based variant.
Limitations
- Axial cross-attention between tensors of different dimensions is not yet fully explored.
- Graph-based XNN is computationally expensive, especially for high-dimensional data.
- No error bars reported because pretraining is computationally expensive; the dimension-agnostic PDE solver is not strictly axis-permutation equivariant due to positional encoding and patch de-embedding.
Experimental setup
Datasets: PDEBench PDEArena Synthetic Gaussian process kernel dataset (2D-5D)
Benchmarks: PDEBench PDEArena GP kernel classification toy dataset
Metrics: NRMSE accuracy parameter count wall-clock time
Baselines: MPP CViT CNN-3D SXCNN
Builds on: Deep Sets Graph Neural Networks Multiple Physics Pretraining CViT ViT
Closest papers by meaning
- Graph Neural PDE Solvers with Conservation and Similarity-Equivariance (ICML 2024 · similarity 0.86)
- HyPINO: Multi-Physics Neural Operators via HyperPINNs and the Method of Manufactured Solutions (NeurIPS 2025 · similarity 0.85)
- Riemann Tensor Neural Networks: Learning Conservative Systems with Physics-Constrained Networks (ICML 2025 · similarity 0.85)
- PINNacle: A Comprehensive Benchmark of Physics-Informed Neural Networks for Solving PDEs (NeurIPS 2024 · similarity 0.84)
- Number Theoretic Accelerated Learning of Physics-Informed Neural Networks (AAAI 2025 · similarity 0.84)
- Global convergence of adjoint-optimized neural PDEs (NeurIPS 2026 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.