Title: Axial Neural Networks for Dimension-Free Foundation Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Parameter Sharing

Dimension Agnostic

Axis Permutation

Foundation Model

Message Passing

PDE

Architecture
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces axial neural networks (XNNs), which represeng a generalization of deep sets and are essentially networks equivariant with respect to permutations of the axes of tensors. This concept is interesting and allows the creation of more flexible graph networks which can combine different modalities with different number of dimensions.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- the exposition and core motivation are well and clearly written\n- simple method with a limited scope - make existing architeyctures for physics foundation models more flexible, in particular w.r.t. dimensionalit\n- the authors demonstrate this effectively using the PDEArena dataset by improving upon 2 existing foundation models\n\n**Weaknesses**\n- some work that is clearly related such as factorized convolutions/attention are not referenced\n- description of experimental method and baselines are somewhat lacking, it is not immediately clear how the efficacy of this method is tested.\n- more baselines would be helpful in quantifying the performance of the metthod, i.e. did this method improve upon a reasonable architecture with reasonable results?\n\nQUESTIONS:\n- for (4) the connection to de Finnetti's theorem is not entirely clear. Is P(\\theta) aprior on the parameters? and how does this enter the picture. A better explanation here would help\n- this method is clearly linked to some commutative properties, which are also satisfied by physical systems. This has been exploited by factorized methods such as factorized attention or convolution. Some acknowledgement/reference to these works would be appropriate\n- can the approach of generalizing from less to more dimensions help address the bottleneck posed by the availability of training data in SciML? An experiment in this regard would go a long way to convince readers about the usefulness of the method.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Axial Neural Networks (XNNs), a dimension-agnostic architecture that treats tensor axes as exchangeable elements to achieve axis-permutation equivariance. The authors develop two variants: set-based XNNs using Deep Sets aggregation and graph-based XNNs with message-passing between axes. They demonstrate the approach by adapting existing PDE foundation models (CViT, MPP) and show competitive performance on 2D tasks with superior generalization to unseen 1D/3D dimensions.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n- The paper is very clearly written and well-structured.\n\n- It addresses a real and important challenge in PDE foundation models—dimension mismatch. Existing solutions such as padding and separate encoders are indeed suboptimal, which makes the proposed approach valuable for the scientific machine learning community.\n\n- The connection to De Finetti's theorem provides a compelling justification for why axis-permutation equivariance enables generalization across dimensions. The mathematical framework is rigorous, with formal theorems and proofs that clearly establish the equivariance properties.\n\n**Weakness:**\n\n- I do not see any major weaknesses in the paper.\n\n- The experiments are conducted with relatively small networks (7M–12M parameters), which somewhat limits the strength of the claims regarding foundation model applicability. It would be helpful if the authors could demonstrate that XNNs can also benefit large-scale models in experiments.\n\nQUESTIONS:\n1. I am not fully convinced by the assertion that the \"patchify operation...is inherently limited to 2D inputs.\" (page 2, line 41-42). In practice, patchification can often be extended to 1D or 3D by using the appropriate convolution kernels (e.g., 1D or 3D convolutions). This point may somewhat overstate the limitations of existing approaches, and it would be helpful for the authors to clarify why such generalizations are considered inadequate in this context.\n\n2. I’m curious why CViT performs significantly worse on the SWE dataset in this paper. According to the original CViT paper, its performance on SWE (from PDEArena) is comparable to MPP. I noticed that this paper uses the SWE dataset from PDEBench—could this dataset difference be the reason for the discrepancy? Some clarification on this point would be appreciated.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces Axial Neural Networks, a dimension-agnostic architecture that ensures axis-permutation equivariance. When applied to PDE solvers, XNNs demonstrate strong generalization across 1D, 2D, and 3D tasks, outperforming some of the established baseline models in transfer learning and fine-tuning tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- The paper presents a compelling deep learning framework that enforces equivariance to axis permutations. The authors carefully design each component of their architecture, resulting in expressive set-based and graph-based equivariant models.\n    \n- They provide formal proofs to establish the equivariance properties of their proposed networks.\n\n- A well-designed toy example is introduced to demonstrate the expressiveness and dimension-agnostic capability of the model.\n\n- The authors conduct notable transfer learning experiments, particularly from 2D to 1D and 3D settings, highlighting the model’s ability to generalize effectively across unseen dimensions.\n\nWeaknesses\n\n- Although the authors claim to evaluate their architecture on state-of-the-art PDE foundation models, this is misleading. In practice, models like Poseidon and DPOT significantly outperform MPP on downstream tasks (see comparisons in [1]).\n\n- With regard to the evaluation of the foundation model: most downstream tasks (except CFD-3D) are relatively trivial. On CFD-3D, both MPP and X-MPP show very large errors. It is therefore difficult to assess the value of incorporating equivariance. The tested problems are either too easy or too hard. The authors are encouraged to evaluate their model on a wider set of downstream tasks such as those in [1], where foundation models achieve much stronger performance.\n\n- In Table 2, it is unclear whether some downstream task distributions were also used during pretraining. If so, this undermines the evaluation: in real-world settings, it is unrealistic to encounter downstream data from the exact same distribution as pretraining. The true value of a foundation model lies in its ability to generalize to unseen tasks.\n\n- The authors should much better clarify the experimental setup and better isolate seen vs. unseen tasks.\n\n- The paper does not discuss how the model’s error scales with increasing model size or the size of the pretraining dataset. These scaling properties are crucial for understanding the behavior and practical usefulness of foundation models.\n\n-----\n\n[1] Herde, M., Raonic, B., Rohner, T., Käppeli, R., Molinaro, R., de Bézenac, E., \\& Mishra, S. (2024). Poseidon: Efficient foundation models for pdes. Advances in Neural Information Processing Systems, 37, 72525-72624.\n\n[2] Hao, Z., Su, C., Liu, S., Berner, J., Ying, C., Su, H., ... \\& Zhu, J. (2024). Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. arXiv preprint arXiv:2403.03542.\n\nQUESTIONS:\n- Does the axis-permutation equivariance introduce any inductive bias that could limit performance on inherently asymmetric data?\n\n- How does the expressivity of GXNN compare to conventional ViT or CNN models when trained on high-dimensional data?\n\n- What effect would data augmentation techniques like rotations and transpositions during training have, compared to explicitly enforcing equivariance? Could such augmentations implicitly encourage the model to learn equivariant behavior?\n\nLIMITATIONS:\nThe authors did mention the limitations of their approach.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Axial Neural Networks (XNNs), a class of models that process tensors of arbitrary rank by treating each axis as a node in a graph and applying message passing to capture cross-axis dependencies. Two variants are introduced: SXNNs (using Deep Sets) and GXNNs (using GNN-style updates with convolutions or attention). The approach is designed to be equivariant to axis permutations and agnostic to axis semantics. XNNs are used to replace dimension-specific components in PDE surrogates (CViT and MPP), resulting in new variants (X-CViT, X-MPP) that show good performance in low-data and cross-dimensional generalization settings.\n\nSTRENGTHS AND WEAKNESSES:\nOriginality: Modeling tensor axes as graph nodes is a clean and general abstraction. The formulation unifies and extends Deep Sets and GNNs under a common interface. The idea is intuitive but nontrivial, and the paper develops it clearly. Related work is generally cited appropriately. The paper may benefit from more diverse set baselines (e.g., FNOs) in the experiments section.\n\nQuality: Theoretical claims (e.g., permutation equivariance) are sound and well formalized. Experimental results support the claim that XNNs improve parameter sharing and cross-dimensional generalization. However, several claims such as axis permutation robustness and generalization to new boundary conditions are not empirically validated. The ablation between SXNN and GXNN is performed only on a synthetic task. No error bars or standard deviations are reported, making it difficult to assess robustness or reproducibility.\n\nClarity: The paper is well structured and mostly easy to follow. Section 3.2 provides an easy to follow explanation of the GXNN update rule. Notation is consistent, but sometimes dense. The handling of domain-specific features (e.g., BCs, geometry, time) is not entirely clear.\n\nQUESTIONS:\n- Have you tested whether the model is empirically equivariant to axis permutations? For example, if you permute the input axes (e.g., swap spatial dimensions), does the output permute in the same way? This would verify that the symmetry proven in theory holds in practice.\n- Can the model generalize when the meaning of axes changes?  For example, if trained on data where axis 0 is time and axis 1 is space, can it still perform well when these roles are reversed at test time? This would test the claim that XNNs are agnostic to axis semantics, not just axis positions.\n- How are geometry and boundary conditions encoded? Is the model tied to Cartesian grids, or can it handle irregular domains?\n- Why not compare against FNOs? These are natural baseline for cross-resolution and few-shot generalization.\n\nLIMITATIONS:\nEquivariance to axis permutations:\nTheorems 3.1 and 3.2 prove that SXNNs and GXNNs are permutation-equivariant by construction, but there are no empirical tests. It’s unclear if the trained models preserve this symmetry under finite data and optimization noise. \n\nRobustness to axis semantics:\nClaimed in the abstract and introduction is that the model is said to be agnostic to the meaning of each axis.No experiments back this up. There’s no test where axes are reordered or their semantics are changed. These trivial unit tests can be included in the appendix. \n\nGeneralization to new geometries or BCs:\nAll tasks are defined on structured Cartesian grids with standard boundary conditions. It’s unclear if the model works for irregular domains or non-grid geometries.The tensor structure and axis-based convolutions don’t easily extend to curved or unstructured domains.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-01"}