Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Axial Neural Networks for Dimension-Free Foundation Models

Hyunsu Kim, Jonggeon Park, Joan Bruna, Hongseok Yang, Juho Lee

TL;DR

This work addresses the challenge of training foundation models on physics data across different spatial dimensions by introducing Axial Neural Networks (XNNs), which are permutation-equivariant with respect to tensor axes. It presents two variants, SXNN (set-based) and GXNN (graph-based), to enable dimension-free processing and efficient parameter sharing across dimensions. Through experiments on toy GP tasks and PDE solver benchmarks, the authors show that XNNs match or exceed the performance of dimension-specific baselines and, importantly, generalize better to unseen dimensions when trained multidimensionally. The results argue for multidimensional pretraining as a key component in scalable physics-informed foundation models, paving the way for flexible PDE solvers that operate across 1D–3D domains without rearchitecting for each dimension.

Abstract

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.

Axial Neural Networks for Dimension-Free Foundation Models

TL;DR

This work addresses the challenge of training foundation models on physics data across different spatial dimensions by introducing Axial Neural Networks (XNNs), which are permutation-equivariant with respect to tensor axes. It presents two variants, SXNN (set-based) and GXNN (graph-based), to enable dimension-free processing and efficient parameter sharing across dimensions. Through experiments on toy GP tasks and PDE solver benchmarks, the authors show that XNNs match or exceed the performance of dimension-specific baselines and, importantly, generalize better to unseen dimensions when trained multidimensionally. The results argue for multidimensional pretraining as a key component in scalable physics-informed foundation models, paving the way for flexible PDE solvers that operate across 1D–3D domains without rearchitecting for each dimension.

Abstract

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.
Paper Structure (36 sections, 3 theorems, 74 equations, 8 figures, 7 tables)

This paper contains 36 sections, 3 theorems, 74 equations, 8 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Backgrounds
  3. Graph Neural Networks and Deep Sets
  4. Transpose and Axis-Permutation Equivariance
  5. Cycle Notation for Axes Permutation
  6. Axial Neural Networks
  7. Set-Based Axial Neural Networks
  8. Graph-Based Axial Neural Networks
  9. Lifting Layer.
  10. Example: Dimension-Agnostic PDE Solver
  11. Related Work
  12. Experiments
  13. Toy Dataset: Gaussian Process Kernel Prediction
  14. PDE Solver Foundation Models
  15. Conclusion
  16. ...and 21 more sections

Key Result

Theorem 3.1

Let $\phi$ be an axis-permutation equivariant function (e.g., another SXNN or pointwise operation). The SXNN in eq:sxnn is axis-permutation equivariant for any rank-$K$ tensor $\boldsymbol{x} \in \mathbb{R}^{d_1 \!\times\! d_2 \!\times\! \ldots \!\times\! d_K}$. (app:sec:sxnn-proof for proof)

Figures (8)

  • Figure 1: Axis Graphs.
  • Figure 1: SXCNN vs. GXCNN.
  • Figure 2: Illustration of GXNN in 3D.
  • Figure 3: An example of SXNN and GXNN for CNN.
  • Figure 4: Accuracy on GP kernel prediction.
  • ...and 3 more figures

Theorems & Definitions (6)

  • Theorem 3.1
  • Theorem 3.2
  • Theorem 3.3
  • proof
  • proof
  • proof