# Axial Neural Networks for Dimension-Free Foundation Models

arXiv: https://arxiv.org/abs/2510.13665 · HTML: https://mlsanity.com/paper/2510.13665/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 42 references that are in the map.
- First posted on arXiv 15 Oct 2025.

## Builds on (most cited first, 26 of 26)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset](https://mlsanity.com/paper/1705.07750/) (2017)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [General E(2)-Equivariant Steerable CNNs](https://mlsanity.com/paper/1911.08251/) (2019)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [Generalizing Convolutional Neural Networks for Equivariance to Lie Groups on Arbitrary Continuous Data](https://mlsanity.com/paper/2002.12880/) (2020)
- [Omnivore: A Single Model for Many Visual Modalities](https://mlsanity.com/paper/2201.08377/) (2022)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [mPLUG-2: A Modularized Multi-modal Foundation Model Across Text, Image and Video](https://mlsanity.com/paper/2302.00402/) (2023)
- [DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2403.03542/) (2024)
- [A Dynamical Model of Neural Scaling Laws](https://mlsanity.com/paper/2402.01092/) (2024)
- [GenTron: Diffusion Transformers for Image and Video Generation](https://mlsanity.com/paper/2312.04557/) (2023)
- [CViT: Continuous Vision Transformer for Operator Learning](https://mlsanity.com/paper/2405.13998/) (2024)
- [Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data](https://mlsanity.com/paper/2411.06646/) (2024)
- [FMint: Bridging Human Designed and Data Pretrained Models for Differential Equation Foundation Model](https://mlsanity.com/paper/2404.14688/) (2024)
- [Any-dimensional equivariant neural networks](https://mlsanity.com/paper/2306.06327/) (2023)
- [Jointly Trained Image and Video Generation using Residual Vectors](https://mlsanity.com/paper/1912.07991/) (2019)
- [Dimension Agnostic Neural Processes](https://mlsanity.com/paper/2502.20661/) (2025)

## Abstract (arXiv)

The advent of foundation models in AI has significantly advanced general-purpose learning, enabling remarkable capabilities in zero-shot inference and in-context learning. However, training such models on physics data, including solutions to partial differential equations (PDEs), poses a unique challenge due to varying dimensionalities across different systems. Traditional approaches either fix a maximum dimension or employ separate encoders for different dimensionalities, resulting in inefficiencies. To address this, we propose a dimension-agnostic neural network architecture, the Axial Neural Network (XNN), inspired by parameter-sharing structures such as Deep Sets and Graph Neural Networks. XNN generalizes across varying tensor dimensions while maintaining computational efficiency. We convert existing PDE foundation models into axial neural networks and evaluate their performance across three training scenarios: training from scratch, pretraining on multiple PDEs, and fine-tuning on a single PDE. Our experiments show that XNNs perform competitively with original models and exhibit superior generalization to unseen dimensions, highlighting the importance of multidimensional pretraining for foundation models.