Axial Neural Networks for Dimension-Free Foundation Models
Authors: Hyunsu Kim
Jonggeon Park
Joan Bruna
Hongseok Yang
Juho Lee
Venue: NeurIPS 2025 spotlight ()
Keywords: Parameter Sharing
Dimension Agnostic
Axis Permutation
Foundation Model
Message Passing
PDE
Architecture
Primary area: deep_learning

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
Hyunsu Kim', Jonggeon Park Joan Bruna Hongseok Yang', Juho Lee"
KKIIST, NNew York University, (kim, hyunsu, parkjonggeone hongseok, yang, juholee @kaistt aaasstaac.... ac, kr bruna@cimm nyu. edu
Abstract
The advent of foundation models in Al has significantly advanced ceneral-purposs learning, enabling remarkable capabilities zero-shot inference and in-context learning. However, training such models on physics data, including solutions partial differential equations (PDEs), poses a unique challenge due to varying di- mensionalities across different systems. Traditional approaches either fix max- imum dimension or employ separate encoders for different dimensionalities, re- sulting in inefficiencies. To address this, propose dimensionn-ggosste neural network architecture, the Axial Neural Network (XNN), inspired by parameter- sharing structures such Deep Sets and Graph Neural Networks. XNN general- izes across varying tensor dimensions while maintaining computational efficiency. We convert existing PDE foundation models into axial neural networks and eval- uate their performance across three training scenarios: training from scratch, pre- training on multiple PDEs, and fine-tuning on single PDE. Our experiments show that XNNs perform competitively with original models and exhibit superior generalization unseen dimensions, highlighting the importance of multidimen- sional pretraining for foundation models.
Introduction
The growing scale of deep learning models has led the emergence of general-purpose AI systems, often called foundation models. Trained over a large amount of unlabeled data with self supervi- sion, these models have shown impressive generalization performance, enabling effective zero-shot inference and in-context learning on a wide range of tasks. In practice, these models are further fine- tuned or post-trained for particular target tasks, achieving performance superior to models trained from scratch. A key requirement for developing highly-performing foundation models is the use of vast and diverse training data. In fact, the relationship between model's performance and the scale of data (and models) is known to follow version of power law, called scaling law [22, 18, 4].
This paper concerned with developing key techniques for building successful foundation mod- els for physics data, such as climate time-series data and solutions of partial differential equa- tions (PDEs). When training such model, we often have combine datasets from multiple sys- tems or differential equations that operate on different dimensionalities. For instance, the Burgers equation describing dissipative fluid flow is usually studied in one spatial dimension, whereas the Navier-Stokes Stokes equations are studied in two or three dimensions. Since the solutions of these equa- tions are typically represented as tensors whose elements are points spatial and temporal grids, different dimensionalities mean that those tensors storing solutions have different numbers of axes.
A straightforward way address such mixed-dimensional scenarios is to fix maximum dimen- sion and either pad lower-dimensional inputs with zeros or build separate encoders for different dimensions that share the same output space. However, both approaches inefficient for low-
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

dimensional data and incapable of handling inputs whose dimension exceeds the fixed maximum. Consequently, most of the prior works on PDEs have developed models tailored to a specific di- mension (typically 2D) [17, 28, 35, 21J. Extending these models to other dimensions is nontriv- ial. For instance, the commonly-use 16x16 patches for Transformer models [12], is inherently limited to 2D inputs. Although patchify operation can be designed for any specific dimension, a 2D patchifier is only applicable to 2D in- puts; processing ID or 3D data requires an entirely different set of parameters, as these models lack an intrinsic parameter-sharing mechanism across dimensions. While recent work has proposed dimension camension-agnostie methods based neural processes [24] and multilayer perceptrons [25], these methods still incur large computational costs: they either flatten high- nigh-dimensional data and pro- cess long sequences, leading to expensive attention layers, or pre-compute dimension-equivane weights via singular-value decomposition, which becomes infeasible in high dimensions.
patchify operation, which applies 2D convolution to extract
neural network by adopt- ing the core principle of surameter-sharing from setnuuttationquivvrrriiii architectures such as Deep Sets [38] and Graph Neural Networks [32, 15, 30J, closely related to De Finetti's ' theorem [11, 38, 3]. Our architecture achieves permutation equivariance over tensor axes. Concretely, we treat the axes of tensor as elements of a ser and introduce a permutation--quiiaarinn architecture which we refer to as the Axial Neural Network (XNN). Although such set-based XNN is simple and computation- find it inherently less expressive, and thus further propose an advanced version
We propose an efficient, dimension-agnostee and ciinensionn-eneralizaan
ally efficient, termed graph-basse based XNN, which captures relationships among axes by treating the axes of a tensor as vertices of graph. Finally, introduce dimension -anensiin- PDE foundational model trained and evaluated on PDEs of varying dimensionality within a single model. Crucially, and in contrast to traditional patchify operations, an XNN-based patchify operation leverages this parameter sharing to make it applicable to inputs of any dimension without modification.
To demonstrate the expressivity and benefits of multidimensional pretraining, we convert existing PDE foundation models [28, 35] to the variants based on our XNNs. We evaluate the resulting mod- els in three different settings: training single PDE from scratch, pretraining with multiple PDEs, and fine-tuning on a single PDE. We show that our variants perform competitively with their origi- nal counterparts. We also conduct experiments to demonstrate the imscen-dimensicn generalization ability of XNNs, which plays an important role in such a aamension-agnostie strategy. Our XNN architecture shows better performance in unseen dimension fine-tuning, which underscores the ne- cessity of multidimensional pretraining for foundation models. The implemented architectures are summarized in https: / /i ght.....t...g.ggggg...g...... com /i///iimmih// hammmmhyunnuuttnn
2 Backgrounds
2.1 Graph Neural Networks and Deep Sets
A Graph Neural Network (GNN) is deep neural network designed to process and make pre- dictions on data represented as a graph [30, 41]. GNNs are characterized by a message-passing mechanism, in which information is exchanged between nodes through their connections, called edges. Formally, given the feature vector æa of node a, its hidden representation h. is computed as
bl Db, eat))
ha = o Da.
bEngbr(a)
where Φ and sb are neural networks with parameters shared across all nodes, + denotes permutation-invarian aggregation operation, ngbr(a) is the set of neighbors of node a, and eab the edge feature between nodes a and The critical architectural feature of GNN that the parameters of ó and &i are shared across all nodes in the graph. This means the exact same functions are used to update each node's representation based on its local neighborhood.
This sarameter-sharing structure is the fundamental reason GNNs are coanuutaiio-equivariiii muting the input nodes simply changes the order of identical operations, leading to corresponding permutation in the output. This symmetry is therefore structurally embedded in the model design, allowing the GNN to generalize across graphs of different sizes and structures.
per-
Deep Set neural network for set-structured data [38] and can viewed as a special case of GNNs in which every node is equally connected to every other node. They can therefore be

---

expressed by the following simplification of Eq. 1:
K h = o Σ (xx))
where K the number of elements in the set. Similar to GNNs, the design relies on parameter sharing: the same function applied to every element æ; before permutation cermutation-ivaariatt aggre- gation (summation or maximum) is performed. This shared application of z ensures the model is permutation-invarian by construction.
2.2 Transpose and Axis- diss-eermtttionn Equivariance
The transpose of matrix flips the matrix over its diagonal, and the transpose of a tensor swaps two axis indices. Exchanging axes i,j E K] of rank-K tensor (Idd dd.....d, yields
To dy dd--d.--d---
where Tıj denotes the transpose between axes and An axis permutation is obtained by cumula- tive transposes, corresponding to reordering the axes. For a rank-4 tensor æ, for example,
(T dgddds I (1) dre2) daess de(s) -dgdzdadi1
(4)
where ı is the permutation associated with 11.
Equivariance is the property that function commutes with the action of symmetry group G. If and y are acted on by Px (x(gg and Py(g) for each group element E G, respectively, then Ø X y Geequivariant when
(px((g) æ) py(() 6(x), Vg G.
(5)
Let II be the group of all axis permutations of rank-K K tensor. A mapping is axis-permuuatioo equivariant if
(I(())) = [l((())). eIl II,
(6)
ie., permuting the input axes and then applying & produces the same result as applying e first and then permuting the output.
2.3 Cycle Notation for Axes Permutation
To express the reordering of axes throughout this paper, we require a precise notation for permuta- tions. A permutation is formally defined as bijection one-to-one mapping) from a set onto itself. In our context, the set consists of axis indices, such as {H, W, D3 for E iR HxWxD
First, review the standard cycle notation common in algebra. This notation group elements into disjoint cycles that show the path each element follows under the permutation. An element within cycle is mapped to the element immediately following it. The last element in cycle is mapped back to the first, completing the loop.
parentheses () to
For example, consider permutations on the set of three axis indices 21,2,3):
A permutation that maps + 3, 3 - 2, and 2 - 1 is written as the single cycle (1 (132 2).
A permutation that maps + 3, + 1, and leaves unchanged, 2 - 2, is written as 3). The element 2 is fixed point.
Adopting this cycle notation, define the transformation Tip the permutation that maps the original ordered set of axes d, x xdg...... ..x, to the new ordered arrangement di xd, x xxx xdi,- For example,
y (1)) y WXDSH y' = Tia(w), y' E DAWsH x @ HxWxD
& R HxWxD

---

3 Axial Neural Networks
We draw inspiration from permutationnequivaaie architectures such Deep Sets [38] and GNNs, which process set or graph data with a variable number of elements. A key advantage of these models is their ability to handle inputs of varying sizes by sharing the same parameters across all elements. We apply this core idea to the axes of a tensor, proposing neural network that can process input tensors of varying dimensions using a single set of parameters.
To this end, we introduce new type of neural network, the Axial Neural Network (XNN), which is equivariant to permutations of a tensor's 's axes. It achieves this by applying an identical transforma- tion with shared set of parameters to each axis, thereby treating them as interchangeable elements, similar the elements of set or the vertices of graph. We propose two variants: the set-based XNN and the graph-based XNN.
3.1 Set-Based Axial Neural Networks
Set-based Axial Neural Networks (SXNNs) are inspired by Deep Sets [38]. They treat the input as the set of all possible axis permutations a given tensor. Specifically,
Rank {
Rank 3 (1o (1, (x), 112 (x), [1s()), (L(x), (L()), (1:()), Rank 2 : (To( (x:); Rank K: (To():), IT (z), Ix)),
(8)
where [x) = æc A Deep- Set style aggregation as in Eq. 2 is then applied:
+ - (((())))) Hen
= o
with neural networks o, v. permutation-invariiat aggregation + (e.g. sum, mean, or max), and inverse permutation II-1
Theorem 3.1. Let & be an axis-permutation equivariant function (e.g,, another SXNN or point-
wise operation). The SXNN in Eq. 9 is axis periutermttttion equivariant for any rank-K tensor R xdx. (S B.1 for proof)
Linear. For instance, a simple SXNN may use & as the identity and da as a linear layer applied along the last axis, followed by max-pooling over the remaining axes for matching the output size. For rank-3 tensor x E H it formalized by
2
y
Σ u, (Pool1,2 (Linearg (LI())))), i=0
(10)
where Lineary applies a linear layer along the third axis (D). Pool1.2 pools over the other two axes (H and W), and they treat the remaining axis batch dimensions. Note that only three of six permutations are required, since pooling invariant to permutations across the pooled dimensions, i.e. (Rol(RR
Pool ( WxHxD), which omits the redundant permutations. Note also that pooling is equivalent to the approach commonly used in 3D inflation of 2D convolutional layers [6, 28). The transformation flow of the feature sizes in Eq. 10 is summarized as
No HxWxD Lineara,
Pool 22 hxuxd
H W x
HxWxD
Ih DxH Lincars
D x
Pool
dxhxu
11,
4 hxwxd,
I1 WXDxH
Lincars,
W x Dxh
Pool 2
w w xx
where the bold letters indicate the axes that the operation is applied. This construction generalizes to any rank- tensor with (K - -imdiemnsio pooling as
K-1
(Linear) ==
Σ m (Pool.... K- (Lineark K (I,())))) 1=0
(11)
Instead of downsampling like pooling, we may use upsampling operations (e.g., resize) to expand the tensor size if Linearg raises the output size.

---

Altaćhu:
Aluthal
Comean ComvDD
ConvaDoa ConviD 4
ha
hw
No
Ap
Artn(A.)
AltorA.)
(a) Lifting Layer
(b) Subsequent Layers
Figure 1: Axis Graphs.
Figure 2: Illustration of GXNN in 3D.
Convolution and Attention. For the convolutional layers, we assume the input tensor has an extra channel dimension C, i..., R IXXWXDXC The axial convolution is still applied along the spatial axes H, W. D. As in the linear case, we can construct the axial convolution using a ID convolutional layer and the axial self-attention layer applied over axis. In attention layers, the output sequence length matches the input, so it is unnecessary to apply a pooling or resizing to align output sizes.
K
K-1
(Conv)
Σ
II,' (Poolı..... K (I (Attn) Σ
nT,' (SelfAttn (IT (12)
Interestingly, the set-based axial attention is already used in a recent PDE foundation model [28] for reducing the computational complexity the Transformer. Splitting the operation across axes reduces the attention overhead. For instance, self-attention over IR HDW'D requires O((HWD)
complexity, whereas axial self-attention reduces this to O + 2 DD).
Non-linearity. & and vb in Eq. 9 can be arbitrary neural networks, including those with single non- linearity such ReLU or Sigmoid. Therefore, using any type of pointwise operation (including non-lincarities) does not violate the axis-permutation equivariance of XNN.
Expressivity. Although SXNNs provide strong expressibility for the dimension-agnosto architec- ture, their expressivity inefficient due to their symmetric structure; iee, they can universally ap- proximate any dimension-agnostic function but require relatively large width achieve this. For instance, consider a patch embedding for the Vision Transformer (ViT) [12] and we assume the patch embedding uses convolution layer with kernel size 2, stride size and both input and output channels are scalar-valued. conventional convolution, the operation on a 2 x pixel patch (2D) can be described as
Conv2D(æ))
aa
21
.12
arı + bzz + CT3
dz
(13)
F3
3:4
where denotes convolution, the first matrix is the kernel, and the second is 2 patch from the input image. On the other hand, the axial convolution in Eq. 12 with average pooling results in:
(on))
I1
T2
aæ1 + bxz
AvgPool
am + tazz thbrg
b
F3
0.33
b
2
ConvID(II
a
I1
43
ax
+ bxz
AvgPool am tarz + tbrs
b
F2
T4
0.22
bxz
a++b
2 0a1 + atb +
ath 2 33
bxra
D2
(14)
a+b
b
33
C4
which shows that the axial convolution behaves like convolution with a symmetric kernel. This symmetric kernel structure limits expressivity efficiency, increasing the number of output chan- nels is often necessary to mitigate this limitation. To address this issue, we also introduce a different type of XNN called the graph-bused XNN, which avoids the constraint entirely.
3.2 Graph-Based Axial Neural Networks
SXNN produces outputs that are inefficient in terms of expressivity compared to the standard neural networks, due to the simple aggregation H To overcome this limitation, can lift the input into an axes-permutation equivariant space, a strategy widely used the equivariant neural network literature [8, 9, 36, 13], and aggregate them in the intermediate layers in GNN.
Lifting Layer We imagine an undirected graph over the axes such as Fig. 1. We now apply the message-passing logic of GNNs as described in Eq. 1. Since the neighbors of d, are d2...., ... dg and
5

---

likewise for the other axes, we can informally express the update rule as
hdd δb
<((d di ① (dd dz)
(15)
where i denotes an arithmetic form of the permutation-invariant operator + and d1, d2,.... dx used conceptual feature representing each axis. The functions e and & be chosen based on the architecture or task. Although Eq. 15 follows parameter-sharing principles and permutation equivariance, practice, have to modify it to match the axes order after the aggregation @.
For example, in the case of convolutional layers, let æ E HB where is the number of channels. We can define é the identity function that returns the second argument (i.e. o(A, B) B) and ab H, W) as a Conv2D applied over axes H and W with the pooling layer at the end for matching the tensor sizes. We omit (H, H) as it is nontrivial, and the absence of it does not violate the axis permutation equivariance. When the indices of H, W, D) are {1, 1,2,30 2,3 Eq. 15 becomes
H = Ti3 13 Pool, 2onv2D2,a (1) - T T1 132 Pool, Conv2D T (1))
hw T23 T23 Pool, Conv2D T2a (2: T23 TeaTr23 Pool Conv2D (12)
(16)
h T33 Pool Conv2D (z) + T33 TT2 Pool, Conv2D T12(2))
where Taja denotes reordering axes 1, 3, to i, explained in 2.3. Of course, o and y need not be linear, and it can be multilayer perceptron. The transformation flow of the feature sizes of hr in Eq. 16 would be:
HxWxD
Tia
Dx W
Conv2D2.3,
Dxwxh
Pooli dx wxwwh
TisTi3
+ dx wx hH
TE3
W X D
Conv2Da.a
W xxxh
Pooh, x xxxxx
TisTish
Here the channel axis C' is omitted for simplicity. A CNN example is also illustrated in Fig. 2a. For ID and 2D cases, Eq. 15 reduces to generating the corresponding number of outputs.
(ID) hH ò (H, ( H)), (2D) h (H, & (H, H) @ ((H, W)), hw o (W, (W, W) @ (WW, H)).
(17)
The 1D case does not consider interaction with the other nodes and (HH H) is still nontrivial. Thus, instead of determining v H),, omit the self-edge term wb H) (W, W) but rather augment a ID tensor to a 2D tensor to fully utilize Conv2D. Possible augmentations include the outer product, repetition, and the diagonal matrix. We adopt repetition, which repeats the 1D tensor along a new axis to match the kernel size of Conv2D, and then averages it after lifting to recover the original ID tensor.
Generalization of the Lifting layers, Eq. 16, to rank-K tensor described as
hd, = o (Tocco)
K
(To)k- ) K> K>1, (18)
(x), +
iy
where [(a)())e)(d)) Tabed-
Theorem 3.2. Under some assumptions, the lifting layer of GXNN, Eq. 18, is axis-permutation equivariant for any rank-K K (except K = 1) tensor x E Rddoddx....xx xdx ($ B.2 for assumptions and proof)
Subsequent Layers. Unlike SXNN, which produces a single feature, we obtain three updated features hg, hw. and hp in GXNN. Due the lifting construction, these features are equivariant with respect to permutations of the input axes. In other words, the permutation of input axes results in the permutation of output features with rank Therefore, the subsequent layers will be a GNN whose input is graph with nodes hH, hw, and hp as the third graph in Fig. 1:
h/4 o hd.
vb (ha, hdd))
(19)
==1

---

Lifting
Hxw
H HXDXW
Conv2D
Hxwxd Bxdxx
Dxhxa a
DxwXA
Hxdxa
Wxaxd
Pool
Axuxx Axdxw
dxwxx
m mxdxk
w.X.X
H
dx
Subsequent
dxAxx
8 W xd
Dx H
xxDxx xx
Conv ID
Conv
Conv
Conv
ConvID
Conv
6xuxd
bxuxd
axxxx
Jxhxx
Pool
dx
wxdxh
Vxdxx
Zx
La
Axwxd xd
XN
XN
(a) SXCNN
(b) GXCNN
Figure 3: An example of SXNN and GXNN for CNN.
Therefore, the subsequent layers should also be dimension-agnostic message-passing architectures, with the lifted features hR, hw, and n as graph nodes. Eq. 19 also needs to be modified match the indices order in the aggregation.
For example, in the case of self-attention layers, can define á residual path, i.c. ((A, B(A)) = A + ((A), and vp as a self-attention applied over the last axis of the second ten- sor, i.e. (1A, B) = SclfAttng (B). Then, as also described in Fig. 2b, the GNN in Eq. 19 becomes
h' hn Ti3Ti3 SelfAttng (h.) T22 SelfAttnn (hw) T13 T3 SelfAttng(hb), by hw T SelfAttmg (h T23 22 SelfAttny(hw)) T2a SelfAttng
hDD = hD + Taa T13 SelfAttny (hx) + Taa T SelfAttn:(hw) T33 (efFAttny(hp)) (20)
equivalent to
13 SelfAttny(ha) T23 SelfAttng (hw) + T33 (llfAttns(hD)).
(21)
h/g hu + (mm hlw = hw (m)),,. h1D hD Tm))m,
where Ts are used for aligning the axis order to match the axes in the aggregation. Similarly, the Subsequent layers in Eq. 20 for rank-K tensor be written
h/2, o ha, Tha.TIOC (())) Fobc
(22)
Theorem 3.3. Under some assumptions, the subsequent layers of GXNN in Eq. 22 are axis- permutation equivariant for any rank-K K tensor ha, E proof)
Iddxxdxx... xidx..xx..xdi (5 B.3 for assumptions and
Pooling Layer. In typical CNNs such ResNet [19], before computing the output logit values using the linear head, global average pooling or global max pooling is applied over the height and width of the features to aggregate spatial information. Likewise, in GXNN, tensors through the lifting layer and subsequent layers, and we need to aggregate these features to merge information across axis permutations. Here is example and its generalized form:
obtain K feature
(rank 3) h' T (h')) @ Te (hw) @ Tyg (h/)), (rank K) h' T h's
(23)
The difference between SXNN and GXNN in simple CNN architecture can be seen in Fig. 3. SXCNN naturally satisfies the permutation equivariant structure by repeatedly stacking the same layer. On the other hand, GXCNN requires lifting layer at the beginning of the network and pooling layer the end. In the middle, subsequent layers can stacked repeatedly. Those XCNNs with added nonlinearity and normalization layers used in $ 5.1.
7

---

3.3 Example: Dimennionn-Agootic PDE Solver
One of the important applications is to solve PDEs. Solving PDEs with Al for reducing the cost of numerical PDE solvers, which often requires supercomputers, is rapidly rising field in modern machine learning [1, 40, 31, 27]. Each PDE has different spatial dimensionalityy and its solutions are represented tensor whose elements points in spatial and temporal grids. Although the neural operator [23] has a crucial benefit in solving PDEs, ViT is still commonly used in PDE foundation models due to its strong generalization. The model for Multiple Physics Pretraining (MPP) [28] is one such model that serves powerful performance in multiple 2D PDE training.
We provide an example of a dimension- sists of patch embedding. multiple attention layers, and patch de-embedding. The patch embedding and patch de-embedding are CNNs. Thus, in the axial implementation, we use the patch embedding as the lifting layer and the rest as the subsequent layers. The attention layers are the same as the SelfAttn example described in Eq. 20, and the patch de-embedding is convolution variant of it. The details and illustrations of the example can be referred to in $ A.
PDE solver by merging GXNN and MPP. MPP con-
4 Related Work
Several studies have explored dimension-agnostic architectures. Levin and Díaz [25] proposed any- dimensional equivariant neural networks that leverage representation stability from algebraic topol- ogy, enabling models trained on fixed input dimensions to generalize to arbitrary sizes. Similarly, Lee al. [24] introduced dimension-agnoster neural processes, which incorporate dimension ag- gregator block to unify inputs of varying dimensions into a shared latent space. These approaches of- fer strong potential for constructing flexible operators that scale beyond conventional grid-dependent solvers. From a practical standpoint, multi-modal training, particularly joint training on images and videos, has also been widely explored [10, 37, 14, 26, 7]. However. these methods typically treat images as individual video frames and either introduce temporal attention layers or use separate embedders for video, without adopting a truly cameeniion-agnstee approach.
5 Experiments
5.1
Toy Dataset: Gaussian Process Kernel Prediction
To demonstrate the dimension training and inference pabilities of XNN, we construct a toy dataset of synthetic Gaussian process (GP) data. The task is binary classification between two kernels: radial basis function (RBF) and periodic, given cialiiddieesionall data (2D, 3D, 4D, and 5D) randomly sampled from the synthetic GP. We build two model to train it: a CNN with 3D convolutions and an XCNN with graph-based axial convolutions composed of 2D convolutions. Although our architecture naturally supports multi-dimensional data, the 3D CNN is restricted to only 3D spatial inputs. Therefore, for 2D data, we pad with zeros to construct 3D-like inputs, and for 4D and 5D data, we flatten the last axes to obtain 3D shapes (e.g., SD: fGx16x16x 16 16xC + 16x16xz 4096xC, where is the channel dimension). Further experimental details are provided in $ D. As shown in Fig. 4, XCNN built with 2D convolutions exhibits superior test accuracy across all dimensions. Notably, despite using 3D convolutions, CNN-3D performs poorly on nigher-dimensional data, including 3D, highlighting the robustness of XNN's s -agnsstii design.
3D
90-
XCNN
Accuracy (56)
01.00
70
A8.32
60-
50
Dimensions
Figure 4: Accuracy on GP kernel prediction.
Comparison Between SXNN and GXNN. SXNN exhibits a trade-ofi off between computational effi- ciency and expressivity inefficiency as described in Eq. 13. To examine this and verify the necessity of GXNN, we build a set-based axial CNN (SXCNN) composed of Conv 1Ds and a graph-based axial CNN (GXCNN) from the previous experiment. We then compare their number of parameters (#Params), wall clock time (W. Clock) of forward computation, and inference performance (Acc.) using the GP toy dataset. Table I shows the comparisons, where SXCNN-L. is an enlarged version of SXCNN for fair comparison with GXCNN. Note that SXCNN with the same depth and width as
8

---

Table I: SXCNN VS. GXCNN.
Table 2: Test NRMSE of PDE solvers on 2D PDEs.
SXCNN SXCNN-L
GXCNN
Model
Pretrain
FineTune
DR
NS
SWE
CRD MO.1
CFD
Depth
5
CVIT
X
X
0.1078
0.1876'
Width
256
128
X-CVIT
x
X
0.0322
0.1148
0.1948'
4Phrums
150K
791K
899K
W. Clock
80ms
9tmms
101ms
X
0.0015
2D Avc.
95,94
92.60
95.45
X-MPP
X
X
0.0118
0.00112
0.0118
D.0K33
Acc.
MPP
20
0,0447
0.0087
0.404
0.0499
5D Acc.
Acc.
62.09
62.64
55.18
70.18
70.86
X-MPP
MPP
ID,20,33
1D,2D
X
0,2183
0.0430
0.0086
0.0265
0.1881"
0.0428
0.0517
0.2199
20
0,0022
0
X-MPP ID,20,3D
0.0058
0.0011
0.0209
0.0294
GXCNN has 6x fewer parameters and 80% reduced wall clock time, but its performance degrades in high dimensions. After increasing the depth and width, it performs fairly well in high dimensions, but GXCNN still exhibits superior performance, indicating the necessity of GXNN.
5.2 PDE Solver Foundation Models
In this experiment, we evaluate the effectiveness of XNN architecture for multi-dimensional training. A compelling use case would be multi-PDE solution training, which involves different PDBs with varying dimensionalities. We train PDE solver foundation models on the solutions of time-homogeneous PDEs drawn from a variety of physical systems. The data are sourced from widely used PDE solution benchmark datasets: PDEBench [33] and PDEArena [16]. The list of PDEs and their details are provided in $ E.
We implemented our architecture in two ofate-of-the-arl PDE foundation model baselines: CVIT [35] and MPP [28]. Both are based on ViT [12], incorporating patch embedding and multiple attention layers with some revisions optimized for PDE learning. Note that these baselines designed only for 2D data, as the 2D convolutional layer for patch embedding is ceieension-dependee attention layers are not. We build their XNN variants, termed X-CVIT and X-MPP, which include axial linears, axial convolutions, and axial attentions. The specifications are detailed in $ C.
though the
Throughout the experiments, we followed the basic training procedures of baselines such as CViT and MPP, which take a few timesteps as input and predict the next timestep of a PDE solution. We set CVIT, X-CVIT to take
2 timesteps as input, and MPP, X-MPP to take 4, in contrast to the original MPP paper, which used s = 16. We also replace InstanceNorm [34] of MPP by LayerNorm [2] due training instability observed on 1D and 3D data. The evaluation metric is the Normalized Root Mean Squared Error (NRMSE), defined in $ D.
The handling of domain-speent specific features such as boundary conditions, geometry, and time is deter- mined by the baseline methods we modified for XNN; i.e., X-MPP follows the handling method of MPP. Additionally, the baseline methods target PDE solutions defined only on regular grids. Ac- cording to MPP, the boundary conditions, geometry, and time are not separately input to the model. Instead, MPP is pretrained tions, allowing it to learn general patterns of PDE solutions. Since we eventually finetune the model on specific PDE with known boundary conditions and geometry, it is not necessary to encode them separately as inputs. The only thing MPP handles separately is the periodicity of the boundary condition. Depending on the periodicity of the boundary condition, MPP determines whether to use sequential position bias or periodic position bias in the attention layer.
multiple PDE solutions with varying boundary conditions and equa-
To isolate the effect of our architectural modifications, we use the smallest versions of CVTT and MPP as backbones: CViT-S and MPP-Ti, with 12M and 7M parameters, respectively. Unlike Me- Cabe et al. [28], we exclude incompressible fluid dynamics from training due to its large data size relative to model capacity. The model size is smaller than the pretraining dataset (~40MM, so pre- training offers limited benefit over training from scratch, but it is sufficient to highlight the advantage of multidimensional training.
Expressivity of XNN. We evaluated the architectural expressivity of XNN by measuring how well performs compared to dimension-spe speecifc models in three different settings: single PDE training, multiple PDE training, and single PDE finetuning. Since the baselines are designed for 2D PDEs, we evaluate performance only on 2D PDEs: Diffusion--Raationn (DR), Incompressible Navier-Stokes (NS), Shallow Water Equation (SWE), and Compressible Fluid Dynamics (CFD) with Mach num- bers M = 0.1 and 1. 1.0. We measure NRMSE for each PDE. Notably, for 1D-2D joint training in MPP, we convert 1D PDE solutions to 2D by padding with zero values.

---

BG 1D
ADV 1D
CFD 1D
DS 1D
CFD 3D
910-2
x10-2
1.5 R10-2
0.60
NRMSE
0.55
0.5
0.50
MPP
X-MPP
500
1000
1500
500
1000
200
400
600
0
200
500
1000
#Daaa
#Data
*Data
eData
#Data
Figure 5: Test NRMSE of finetuning on unseen dimensions (ID and 3D).
The results are shown in Table 2. The notation (f) in the table denotes evaluation on PDEArena, meaning the rest are on PDEBench. (*) indicates results trained with both CFD M0.1 and CFD M1.0, meaning more challenging setup compared to training respectively. The underscores denote the best result compared to the competitors. Since we report based on the baseline codebase, the empty slots denote PDEs that the codebase does not support. In every scenario, the axial variants exhibit competitive results compared to the non-axial baselines. In particular, due to the benefit of multidimensional pretraining, the finetuned X-MPP (trained with ID, 2D, and 3D data) outperforms MPP trained only with 2D. Note that pretraining MPP with both 1D and 2D) leads to significant degradation, it fails to learn a unified representation space across dimensions.
this experiment, we demonstrate the dimension- generalization (few-shot learning) ability of XNN on data from unseen dimensions. To do so, we pretrain both MPP and X-MPP on 2D PDE data and compare their finetuning performance on un- seen ID and 3D PDEs, demonstrating the benefit of multi-dimensional pretraining in XNN-based foundation models. We use three ID PDEs and one 3D PDE: Diffusion Sorption (DS), Burgers' equation (BG), ID Compressible Fluid Dynamics (CFD ID, and 3D Compressible Fluid Dynamics (CFD 3D). We evaluate using the same metric, NRMSE, and compare how finetuning performance improves the size of the given dataset increases.
Unseen Dimension Generalization.
For 3D finetuning of MPP, we use the inflation technique described in McCabe et al. [28]. repeats P P kernel of the 2D convolution layers P times and divides by P. The weights of the linear projection for the additional variable the 3D PDE are initialized with the average of the trained weights corresponding to the existing variables. For 1D finetuning of MPP, we augment the input tensor from 1D to 2D only in the patch embedding at the beginning and the patch de-embedding at the end, that the attention layers in the middle operate on purely ID PDEs. In contrast, model naturally extends to 3D and reduces to ID. For X-MPP, we use shared linear projection across different PDEs for unified representation learning. This contrasts with MPP, where different línear projection learns each PDE during fine-tuning.
As shown in Fig. 5, except for BG ID, X-MPP (orange) accelerates the finetuning performance with only a small amount of data. In particular for 3D, the redundancy of the inflated layers in MPP (blue) limits 3D generalization, whereas the dimension-agnootie architecture is free from this issue, thereby efficiently utilizing the representation trained from 2D PDEs.
Conclusion
In this work, we introduce XNNs, family of efficient and expressive architectures designed to be agnostic to input dimensionality. Motivated by the limitations of prior PDE models constrained to fixed dimensions, XNNs leverage axis permutation equivariance over tensor axes naturally generalize across ID, 2D, and 3D domains, which can be applied to develop a dimension-agnostic PDE solver foundation model.
In the PDE solving problem, our empirical results highlight the benefits of multidimensional pre- training and the superior finetuning ability of XNNs compared to the dimension-specie models such as MPP. These findings underscore the importance of designing foundation models capable of operating across varying spatial dimensions, which is critical step toward scalable and adaptable scientific machine learning systems.
Limitations and Future Work. An axial version of the cross attention between two tensors of different dimensions is not yet fully explored. Addressing this issue as a natural extension of this work is a valuable future direction that would further widen the usability of XNNs. Additionally, improving the computational efficiency of GXNN, especially for high-dimensional data, also promising direction towards the practical deployment of XNNs in large-scale scientific simulations.
10

---

Acknowledgement
We thank Heejun Lee for his thoughtful technical support, Jules Berman for his insightful con- cerns, and Kyunghyun Cho for facilitating our wonderful collaborative research. This work was partly supported by Institute
Information & communications Technology Planning & Evalua- Artificial Intelli-
tion(IITP) grant funded by the Korea (MST) (No. Noo.S--001--1199000 gence Graduate School Program(KAlST)) No.RS-2024 2024-00509279, Global Al Frontier Lab; No.RS- 2022-1120011,, Meta-learning Applicable to Real- world Problems) and Artificial intelligence in- dustrial convergence cluster development project funded by the Ministry of Science and ICT(MSIT, Korea) & Gwangju Metropolitan City. H. Yang supported by the National Research Foundation of Korea(NRF) grant funded by the Korean (MST) (No. RS-2023-00029660),
References
[1] Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Jo- hannes Brandstetter. Universal physies transformers: framework for efficiently scaling neu- ral operators. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, De- cember 10 15, 2024, 2024.
[2] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR, abs/ abs/1607.064500 2016. URL
org/ wttpoorg//abboooooo/as ab641667.640000 06450.
[3] Benjamin Bloem Reddy and Yee Whye Teh. Probabilistic symmetries and invariant neural net- works. J. Mach. Learn. Res., 21:90:1-90:61, 2020. URL https:// /tmpp::///mr...////////// org/ papers/ v21 /21/19- 322 html.
[4] Blake Bordelon, Alexander B. Atanasov, and Cengiz Pehlevan. A dynamical model of neu- ral scaling laws. In Forty-first International Conference Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
[5] James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //github. com/ jax jax-ml/ ml/jax.
[6] João Carreira and Andrew Zisserman. Quo vadis, action recognition? A new model and the kinetics dataset. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pages 4724-4727 4733. IEEE Computer Society, 2017.
[7] Shoufa Chen, Mengmeng Xu, Jiawei Ren, Yuren Cong, Sen He, Yanping Xie, Animesh Sinha, Ping Luo, Tao Xiang, and Juan MaanMl Pérez-Ría. Gentron: Diffusion transformers for image and video generation. IEEE/CVF Conference on Computer Vision and Pattern Recogni- tion. CVPR 2024, Seattle, WA, USA, June 16-22, 2024, pages 6441-6451, IEEE, 2024. doi: ALLL09/CVVRR2233.2224..00 2024. 60616.
URL https: / doii org/ 10. 1109/ 100.001111/////2233333
[8] Taco Cohen and Max Welling. Group equivariant convolutional networks. In Proceedings of The 33rd International Conference on Machine Learning (ICML 2016), 2016.
[9] Taco S. Cohen and Max Welling. Steerable enns. In International Conference Learning Representations (ICLR), 2017.
[10] Yatin Dandi, Aniket Das. Soumye Singhal, Vinay P. Namboodiri, and Piyush Rai. Jointly trained image and video generation using residual vectors. In IEEE Winter Conference on Applications of Computer Vision, WACV 2020, Snowmass Village, CO, USA, March 1-5, 2020, pages 3017-3031. IEEE, 2020. doi: 10. (.L110//NACVA55552220000033333 2020. 202.5/22322222.99333033
URL https:
// //doi. org 10.
[11] de Finetti and B. Funzione caratteristica di un fenomeno aleatorio. Attidella R. Academia Nazionale dei Lincei, Serie, 6.(4):25129) 1931.

---

[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 1 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[13] Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing volutional neural networks for equivariance to lie groups on arbitrary continuous data. In Proceedings of The 37th International Conference Machine Learning (ICML 2020), 2020.
[14] Rohit Girdhar, Mannat Singh, Nikhila Ravi, Laurens van der Maaten, Armand Joulin, and Ishan Mísra. Omnivore: A single model for many visual modalities. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18- 24, 2022, pages 1609111 IEEE, 2022. doi: 10. Pttp9/CVVPS228822222220 https: / ttt.:/oooo/o//oo...ooo//oo/ooo/ wwg/ooggg///oooo 111111/////0 11w111199///998988 2022. 01563.
URL
[15] Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In IEEE International Joint Conference Neural Networks, IJCNN 2005, Montreul, QC, Canada, July 31 August 4, 2005, pages 729-734. IEEE, 2005.
[16] Jayesh K Gupta and Johannes Brandstetter. Towards muttisppatioeemmooalssalll generalized pde modeling. arXiv preprint arXiv:2209.1560 15616, 2022.
[17] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. DPOT: auto-regressive denoising operator transformer for large-scale PDE pre-training. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
[18] Alexander Havrilla and Wenjing Liao. Understanding scaling laws with statistical and ap- proximation theory for transformer neural networks on intrinsically dimensinnional data. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Infor- mation Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770-778. IEEE Computer Society, 2016.
[20] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, An- dreas Steiner, and Mare van Zee. Flax: neural network library and ecosystem for JAX, 2024. URL http: : /ttp//gttttb.te/ee/// como google/ ffaa..
[21] Maximilian Herde, Bogdan Raonic, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Em- manuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes. In Advances Neural Information Processing Systems 38: Annual Conference on Neural In- formation Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[22] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URL https: /ttp:///aaivvvv
[23] Nikola B. Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhat- tacharya, Andrew M. Stuart, and Anima Anandkumar. Neural operator: Learning maps be- tween function spaces with applications to pdes. J. Mach. Learn. Res., 14.89:1-89:97 2023.
URL hwtph:ttt///t /mmrr.mrrr/.//// org/
tttp:/pers///oeeeeeeeaa. papers/v24 /warrss/22/211 1524 html.
[24] Hyungi Lee, Chacyun Jang, Dongbok Lee, and Juho Lee. Dimension agnostic neural pro- cesses, 2025. URL https: /ttp:://axxiivvo.e
2066..
[25] Eitan Levin and Mateo Díaz. Any-dimensional equivariant neural networks. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, International Conference on Artificial Intelligence and Statistics, 2-4 May 2024, Palau de Congressos, Valencia, Spain, volume
12

---

238 of Proceedings of Machine Learning Research, pages 2781..71. PMLR, 2024. URL https: https/://pccceeiiiee wtt.ss.....pppeee.... www.ppesss/veee v2/8//22////////// Levin24a html.
[26] Kunchang Li, Yali Wang, Yinan He, Yizhuo Li, Yi Wang, Limin Wang, and Yu Qiao. Uni- formerv2: Unlocking the potential of image vits for video understanding. In IEEE/CVF In- ternational Conference on Computer Vision, ICCV 2023, Puris, France, October 1-6, 2023, URL https:
pages 1632-1643. IEEE, 2023. doi: 10. AA110//CCV51100.2222.00000
//doi. , www/o0gg//.0.o 1011119990VVV/ ICCV51070. 2023. 00157.
[27] Yuxuan Liu, Jingmin Sun, Xinjie He, Griffin Pinney, Zecheng Zhang, and Hayden Schaeffer. Prose-fd: A multimodal pde foundation model for learning multiple operators for forecasting fluid dynamics, 2024. URL https::////. atttss://arae /ttp:///aavvv.../iaa org/ abs/ abs/2409 2409 09811.
[28] Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles D. Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Géraud Krawezik, François Lanusse, Mariel Pettee, Tîberiu Tesileanu, Kyunghyun Cho, and Shirley Ho. Multiple physics pretraining for spatiotemporal surrogate models. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[29] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf. Edward Yang, Zach De Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, sesh-performacce deep learning library, 2019. URL https:://////c //tp:///aaiiv..a.aaeeeaeeaa .rw. abs/1912. 01703.
[30] Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Mon- fardini. The graph neural network model. IEEE Trans. Neural Networks, 20(1):61-80, 2009. doi: 10.1 N0.1009/NN..2088.005555 2008 . 2005605.
URL
[31] Zezheng Song, Jiaxin Yuan, and Haizhao Yang. Fmint: Bridging human designed and data pretrained models for differential equation foundation model, 2024. URL htlthttps:/aaaee het/abb/24444 14 188.
[32] Alessandro Sperduti and Antonina Starita. Supervised neural networks for the classification of structures. IEEE Trans. Neural Networks, 8(3):714-735, 1997. doi: 10.1109/72.572108. URL
setps:://oos http:///oo.og/o////////o wtt.://../.......//...... 572100
[33] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Ale- siani, Dirk Pflüger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. In Advances in Neural Information Processing Systems 35: Annual Confer- Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
[34] Dmitry Ulyanov, Andrea Vedaldi, and Victor S. Lempitsky. Instance normalization: The miss- ing ingredient for fast stylization. CoRR, abs/1607.08022. 2016. URL http://arxix: org/abs/ / / 006000000 08022.
[35] Sifan Wang, Jacob H Seidman, Shyam Sankaran, Hanwen Wang, and George Pappas Paris. Cvit: Continuous vision transformer for op-erator learning. arXiv preprint arXiv:2405.13998 3, 2024.
[36] Maurice Weiler and Gabriele Cesa. General e(2) equivvaraa steerable cnns. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019). 2019.
[37] Haiyang Xu, Qinghao Ye, Ming Yan, Yaya Shi, Jiabo Ye, Yuanhong Xu, Chenliang Li, Bin Bi. Qi Qian, Wei Wang, Guohai Xu, Ji Zhang, Songfang Huang, Fei Huang, and Jingren Zhou. mplug-2: A modularized multi-modal foundation model across text, image and video. In International Conference Machine Learning, ICML. 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings s of Machine Learning Research, pages 28728-38748. PMLR, 2023.
13

---

[38] Manzil Zaheer, Satwik Kottur. Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. In Advances Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, 2017.
[39] Matthew D. Zeiler, Dilip Krishnan, Graham W. Taylor, and Rob Fergus. Deconvolutional net- works. In 2010 IEEE Computer Sociery Conference Computer Vision and Pattern Recogni- tion, pages 2528-2533 2010. doi: 10.1 10.110/CVPR.220005553999
[40] Hang Zhou, Yuezhou Ma, Haixu Wu, Haowen Wang, and Mingsheng Long. Unisolver: Pde- conditional transformers universal pde solvers, 2024. URL https: httpss://axxi.orrg
[41] Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. Graph neural networks: A review of methods and applications. A/ Open, 1:57-81, 2020. doi: 1..1111/..OPPEN.222..0..... org..l....11... 0016/116660 j aiopen.nn.... 2021 01. 001.
URL https:
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction reflects our main contribution: novel architec- ture for dimension- free training and inference.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper discuss the limitations in the last paragraph of the conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specificattin, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Ys]]
15

---

Justification:
The paper provides the assumptions and proofs in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
The paper provides experimental details the experiment section and the
appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data also be provided via
often one good way accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
16

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We share the source code with configuration files in the submission and the data is publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https: : //ee://nnnn /ttp//ips.. cc/ ssione Ponoo) for more details.
public/ gwwdgeiii//eege Catee/CoeSeaaai
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. cc/ pcc/publiii cilbiic//ee gu i des CodeeSumm ssi Poooooooooooo icyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper provides the experiemental details including hyperparameters in the appendix section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [No]
Justification: expensive.
Error bars are not reported because the pretraining is too computationally
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative еπτor rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The paper provide such information in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't t make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:/ /ttess/////uii./.eeee
Answer: [Yes]
Justification:
The paper only targets mathematical data, which does not involve any harmful
aspect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The paper only targets mathematical data, which does not involve any societal impacts.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
18

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper trains with only PDEs.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions to access the model or implementing safery filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The paper clarifies the source of baseline codes and data.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
19

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papersw thoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification:
The paper will release the code with the documentation after accept.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing and human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
No crowdsoureing and human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
20

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLM is used only for editing.
Guidelines:
The answer NA means that the method development in this rescarch does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: / ht////ew /neuriii i cc cc/Conferenccss/ 2025 2/LL/LLM) for what should or should not be deseribed.
21

---

XxHxWXC
RevIN
RevIN
T TxHxW x D
Field Embedding
Field Embedding
Patch Embedding
W x D
T' x WXC
Embodding
Einhediing
Embeddingg
T Txdxyxhxx
TxhxmxdxC
t-Attn.
t-Attn.
t-Attn.
t-Attn.
MLP
MLP
MLP
MLP
Attm
Atin
Attn
Attn
Attn
Attm
Attn
Attn
Attn
X-Attn.
y-Attn.
MLP
MLP
MLP
MLP
Txdx u X
Xxhx
Txhxwxde
Txh WXC
De-Embelding
De-Embedding
De-Embedding
Field Rocen.
Field Racon.
Feld Recon.
De-Embedding
TxDx
DxWXC
WxDXC
Field Recon.
Pooling
Xxxx
De-RevIN
De-RevIN
HxWxC
HxW DxC
(a) MPP
(b) Dimension- MPP
Figure 6: Comparison between MPP and its dimension-agnostir version.
A DimennioonAAnootic PDE Solver
In this section, present example of the dimension-agnosit PDE solver using the ViT-based model MPP as a backbone. MPP processes a PDE solution from timestep to (e.g., a E RTxHxWxC) and outputs the next timestep T (e.g., T1+1 #xxCC)) MPP consists of patch embedding, multiple self-attention layers, and patch de-embedding. In the self-attention layers. MPP separates the attention along the time axis (temporal self-attention, in depthwise manner. We also follow this separate temporal self-attention, but make the spatial self-attention dimension-agnostic. The comparison between MPP and its dimension-agnostic version illustrated in Fig. 6.
Patch Embedding. The patch embedding three-layer CNN, used as the lifting layer. The original patch embedding in MPP is as follows:
(mbBlock(#)) =
Embedder(z))
(24)
h Norm (emmmemeddder(()),
22

---

where Norm indicates InstanceNorm [34]. On the other hand, in the axial version, we use Embedder for in Eq. 18 with max aggregation as follows:
hd. max T(K-1)6) (AvgPool,.... K-2 (Einbedder--L. (To)K- (x00)))))
K-2
8æ1, JE[1,K]
ha: = Norm(ha,):
(25)
where T (a))))(()) denotes the transpose between the m-th and n2th axes. Through (I) the j-th and i-th axes are located at the (K - 1)-th and K-th positions, respec- tively. Then, Embedder-- -I, applies 2D CNN along the K 1 and K axes, and the remain- ing axes I,... K 2 are reduced via average pooling to match the output tensor size. Since the final Norm is pointwise operation, it preserves axis-permutation equivariance. We replace Instan- ceNorm with LayerNorm [2] to avoid training instability in multidimensional training.
Self-Attentions, MPP already employs set-based axial self-attention, which outputs single feature as follows:
K
h'= 1=1
TuKy (Teea K
1
h' h h'
(26)
In contrast, in the graph-based axial case, we obtain K features due to the lifting layer the patch embedding. Therefore, also set- bet-based structure but apply the same attention to each feature independently, which differs from the graph-based attention example in Eq. 21. The self-attention used is:
K Σ K
h1
Taoo (SelfALtn (T (To)c)))aa))) h44 has + h's Vi [],
(27)
Additional operations such as layer normalization and drop residual paths are included, but omitted here for simplicity. These operations follow MPP exactly.
Patch De-Embedding, In this block, we recover the input tensor's width and height using a CNN consisting of ConvTranssoos layers (also known deconvolutions) [39]. In the original imple- mentation, this CNN is defined as:
DeembBlock(4)) (h') GeL GeLUUXom((Com/I22((()))
Deembedder(ht (h) Caav/PDDDDemBBBocccceeeBBBBttt
(28)
h" Deembedder(h').
In the axial version, we define two Deembedder blocks: one for updating the node feature and the other for updating the neighhorhood feature (nhbr). We aggregate them using a operation. Then, Eq. 22 becomes:
h ((mude)
Deembedder (node)
K-1, K-1,k (h's,) K
d,
h ds
rrahbrr
max T(oCK) Deembedder (ubbrr JE(L,K
K (h',))
(29)
{ (nnduddo hh (ontbbb Vi [1,
h" d, max
where TA)KTT(K))) aligns the axis order of h'', to match that of h'y.
Unfortunately, this dimension-agnostic PDE solver is not disippermutaiion equivariant because the positional encoding (not mentioned in this section but present in the actual implementation) and the patch de-embedding layer break axis-permutation equivariance, even though the solver anyway works for every dimensionality.
B Proofs of Theorems
B.1 Theorem 3.1
Let d be an axis-permutation equivariant function (e.g., another SXNN or pointwise operation). The SXNN in Eq. 9 is axis seis-permuttttinn equivariant for any rank- K tensor pdiddx...dss
23

---

Proof. We prove the peiitteemutttion equivariance property defined in Eq. 6. For any axis permu- tation II' II. the permutation of the input Eq. before applying & becomes
Kt-1 n1 (+((((((((
(30)
Now, by substituting I1, I1;II' (equivalently, II, IIII
) we get
K1-1
K!-1
11," (I(I)
I (x)))())))
j=1
= II' K1-1
1 (+LI(())))
(31)
K!-1
II' ① (b
The last equality holds due the permutation-invariine nature of the operation θ
By the assumption that Φ is axis sxis-permutationn equivariant, it follows that
K!-1
K!-1
1 (ILL(II())))))
= o 11'
+
11,"
(v(IL())))
(32)
K1-1
= II'
n
(+(I()))))
i=1
which proves that Eq. 9 is axis sxis-permutation equivariant for all E II.
B.2 Theorem 3.2
Assumption All layers follow the same axis order. For example, if the lifting generates h and hw with shapes H x W x C' and W x H x., respectively, then every layer should produce h and hew with the same shapes.
Assumption 2. In the lifting layer, ø must be axis-permutation equivariant along axes 1, 1.2....,K'' K-1, and db along axes ........ K 2. For instance, if E R HxWXDDC for K = 3, then T12(0(x)) (T)2(zz) must hold, but (1=()())) = (T))(2)) is not required.
Under assumptions Assumption 1 and Assumption 2, the lifting layer of GXNN, Eq. 18, is axis- permutation equivariant for any rank-K (except K = 1) tensor æ E Rdiddx...xok
Proof. The lifting layer Lifting described in Eq. 18 can be equivalently written K
hd, = oi (x), θ T (O0001K-1)) (Tox)))))-)(((())
(33)
jØ
where (() T Tabed denotes the permutation from indices (1,2,3.4) to (a,b,c, as explained in $ 2.3. Applying axis permutation II(x)) the input gives
(((c((((()))) T T (iouc/ToooGGee (0000)K11)) v (Toux))(-=
(34)
Substituting (B))())(K) T(lckyll for =k) and II (9000)XX))))) T for i =(k),j -1), where IT' permutes only axes 1 through K 1 and IT" permutes only axes through K - we get
K
= (T)))
Fieeoooo II" ( (ITTTm)00x=-((1))
(())
(35)
24

---

Let π, vT' be the element-wise permutation across axes induced by II. Since the aggregation is crrmutation-invariii and x (1) =(k) is equivalent + k (as x does not act on the K-th axis and k-th axis is transposed to K by T())))) and using Assumption (permutation equivariance of sb under I1"), we have
K +
= ((T
(
(36)
1tt
Applying Assumption again for under II' yields
I'' (T))KK())))))
(37)
which follows the definition of axes-permiine sxes-permutation equivariance as in Eq. 6.
B.3 Theorem 3.3
Assumption In the subsequent layers, both and db must be axis-permutation equivariant along axes 1,2,..., 1.
Under assumptions Assumption and Assumption 3, the subsequent layers of GXNN in Eq. 22 axis- ddisspermutation equivariant for any rank- K tensor ha, E [dixdxx..xxx
Proof. The subsequent layers Subsequent described in Eq. 22 can be equivalently written
h's, = ha, ((o()) Σ (lac)
(38)
Using Eq. 37, permutation of input leads to permutation of output of the lifting layer, which is the permutation of input of subsequent layers,
K
I''a, (()D))
Σ =[1==1
Toe
(39)
Because the index exchanges in T must adjust to the input permutation II', and both e and aa are II' l''equivariant by Assumption and the is permutation-invarian
I hds: T(O
(oic, IT'
(40)
B.4 Axes-Permutation Equivariance of Pooling Layer
Finally, the pooling layer aggregates the axis-wise features into a single feature, which remains equivariant to axis permutations II. As described in the example in Eq. 23, is equivalently written as
h" Σ Na,
(41)
By definition,
IT'T (1) (((())(() II (
T (K))
(42)
Using this identity and Eq. 40, permutation of the input yields
(T)
I''h',
Σ
IT'
Σ x'()==
T-1 Toi Ihh's
(
x(0==
(43)
K
Σ
IT
=1= h's
Σ
T- Tuunh's ITh".
(O(K)
700=1
which concludes that the axis-permutation of the input yields the axim-pernnnation of the output.
25

---

C X-CVIT and X-MPP
For X-MPP, we exactly followed the example of the dimension-agnostic PDE solver introduced in $ A.
For X-CViT, we mostly followed the structure of CViT but modified the dimension-dependen com- ponents, including patch embedding, spatial self-attention, coordinate query embedding, and de- coder cross-attention. For the remaining components, we refer the reader [35].
Patch Embedding. X-CVIT also uses patch embedding as a lifting layer. In the original CVIT, the patch embedding defined as
h = Cony2D(x))
(44)
In contrast, in X-CViT, the patch embedding is defined as
K
1
hd: K
Σ 541
(Tu-11000000) (MaxPooli.... (Conv2DR-1,K (T(O))))))--((())))))
(45)
where (()K))(X-1) is used to apply Conv2Dx-1.K along the i-th and j-th axes, and ((K-1))(K)()) reorders the axes back enable aligned feature aggregation through the summa- tion.
Spatial Self-Attentions. CVIT separates temporal and spatial attention and X-CVIT also follows the separated structure. In X-CViT, we use set-based axial self-attention for each feature. The original self-attention in CVTT is defined as
h' = SelfAttt(FFlateen(
(46)
which flattens the spatial axes of the input into a single sequence before applying self-attention (e.g.. Hx W + HW). In contrast, X-CVIT does not perform flattening but instead aggregates self-attention outputs along each axis:
h'', ha, +
T(Ko (Sef
(47)
where (()( is used to apply SelfAttn along the j-th axis, and T(K))) reorders the axes back for aggregation.
Pooling Axial Features. Right after computing the encoder, we still have K axial features lifted by the patch embedding. Thus, before applying the decoder, we aggregate them by averaging to obtain single feature.
Query Embedding, While MPP always predicts the entire spatial grid points at the next timestep. CViT selects where to predict by providing spatial coordinates, assuming the PDE solutions lie on a [0, 1|2 grid. To achieve this, CVTT calculates the distance between N query coor- dinates and all grid points in GK resulting in tensor of shape x G which is then encoded into a N x Q tensor, where Q is the feature size of the coordinate embedding. CVIT uses two linear layers to embed GK into R& for the distance tensor q E Nx (with K = 2 in 2D), treating N as a batch dimension:
Coordinate
q'
(48)
On the other hand, X-CVIT utilizes set-based axial linear layer with max aggregation as follows:
q' max T(K)G) (Normæ (Linearæ (GlobalMaxPooh....rr (Linear (R)) (Unflatte))
(E[1:K]
(49) where Unflatten reshapes KK - G G G, and T T() is used apply Lineark along the i-th axis. Note that GlobbalMxPoool...... converting x to R. This illustrates that XNN can be applied even when the number of input axes differs from the number of output axes.
reduces the size of axes through 1 to one,
Decoder Cross-Attention. the decoder cross-attention, we perform attention by treating the coordinate embedding g as the query and the features produced by the CVTT encoder as the key and
26

---

value. Since this operation involves a rank-1 1 vector and a rank- K tensor, we must carefully design the XNN architecture to handle such tiltiidiimesionall tross-attention. In CViT, the cross-attention is formulated as
h" = q' + Attn(q', Flatten(h'), Flatten(h')) (h')).
(50)
In contrast, X-CVIT uses SXNN with a Repeat function align dimensions and tensor sizes:
h" d' Σ GlobalMaxPooh,...te (Attn k-..... 1 (odd),T(CCH)) T(oCK) (h'); ((ICK((h))))
(51)
where Repeat 1-.. 144) expands and repeats the vector q' R along axes I,..., 1, result- ing in Repeat, K (q') RQXQX-xX.
D Experimental Details
D.1 GP Kernel Prediction
Synthetic Dataset Construction. To construct the synthetic dataset, samples were generated from zero-mean Gaussian processes using either the radial basis function kernel or the periodic kernel. For each kernel, data were generated in dimensions ranging from 5. In each dimension d (2,3, 3, 4, 5), a structured grid that equally divides the hypercube [-2, was created over the input space, with grid sizes tailored to 2 This results in output shapes as specified in Table 3, where the last channel of length was appended for convenience.
Table 3: Shape of tensors in the synthetic dataset
Dimensions
Shape
Dimensions
Shape
2D
(32, 32, 1)
4D
(8, 8, 8, 8, 1)
3D
(16, 16, 16, 1) 5D (4, 4, 4, 4, 4, 1)
Each grid point represents input the GP kernel, from which the full covariance matrix was computed for a batch of samples. Kernel parameters were randomly sampled, i.c., length within 0.61, scale within [0.1, 1], and period within [0.1, [0.1,0.5) 0.5) for the periodic kernel, introduce vari- ation. The resulting covariance matrices were corrected via eigenvalue clipping to ensure positive semidefiniteness. Each generated sample from the periodic kernel was labeled 0, while those from the RBF kernel were labeled 1. An equal number of samples from each kernel type was used to ensure balanced classes. All data were normalized to have zero mean and unit variance before being fed into the models. The dataset was then divided into training and validation subsets with an 80/20 split.
3D CNN. As 3D convolution operations are only compatible with 5D inputs, we forced the dimen- sions of the tensor to 5; for nower-dimensional samples, zero-padded along missing axes; for nigher-dimensional samples, we reshaped the tensor to be 5D by aggregating all the dimensions after the third dimension. We adopted conventional 3D convolutional neural network operating on tensors of shape (B, H, W, D, C), where is the batch size. The network consists of three stacked 3D convolutional layers, each with 32 filters and kernel sizes of 3 3 x 3, followed by LayerNorm and ReLU activations. After the convolutional layers, global max pooling is applied across all spa- tial dimensions, and a final dense layer maps the resulting features to single output for binary classification.
SXCNN. Set-based XCNN generalizes convolution to multidimensional inputs using directional convolutions along each axis. At its core is the SXConv module, which applies ID convolutions separately along each axis and merges the resulting features using an clement-wise maximum. Each SXConv operation is followed by LayerNorm and a ReLU activation. The network consists of five SXConv layers, with each layer operating at 64 hidden features, effectively doubling the base hidden dimensionality. The final feature maps are globally pooled across all spatial dimensions, and classification is performed using fully connected output layer.
GXCNN. Graph-based XCNN introduces a more complex interaction between spatial axes through the use of cross-axial convolutions. The architecture begins with lifting layer, which performs
27

---

2D convolutions across each pair of axes to produce set of intermediate representations. This followed by a series of XConv layers, where each layer simultaneously considers pairs of axes using separate node and neighhor convolutions. The feature maps from each interaction are merged using max operations after appropriate axis permutation. Each convolution operation is followed by LayerNorm and ReLU activation. The network consists of one XLift layer and four XConv layers, each with 32 hidden features. Axial max pooling was applied across each dimension, followed by global max pooling, and the pooled output passed to a dense layer for binary classification.
Hyperparameters. All models were trained using the Adam optimizer, with fixed learning rate of 0.001 and a batch size of 64. Training was conducted for 10 epochs for each model. The loss func-
tion used was binary cross-entropy computed from the sigmoid of the output logits. Training and evaluation were implemented in JAX 0.4.30 [5] and Flax 0.8.5 [20]. with PyTorch 2.7.0+cul 18 [29] used primarily for data preprocessing and batching. To ensure that performance comparisons were valid, the same hyperparameters and preprocessing procedures were applied across all models. The table below summarizes the hyperparameter settings used for training each of the three neural net- work models.
Table 4: Hyperparameters for each model architecture
Model
Architecture
Hidden Dim
Learning Rate
Batch Size
Epochs
CNN
3 Conv layers
le-3
64
10
SXCNN
5 SXConv layers
64
le-3
64
10
XCNN
1 XLift + 4 XConv layers
32
le-3
64
10
D.2 PDE Foundation Model
Hardware and Software. We implemented CViT and X-CVIT using JAX 0.4.30 [5] and FLAX 0.8.5 [20], while MPP and X-MPP were implemented using PyTorch 2.1.0+cu121 [29]. All exper- iments were conducted on NVIDIA GPUs: RTX 3090, RTX A6000, and RTX 5090. For the RTX 5090 machine, we used PyTorch 2.8.0+cu1 128 due to CUDA driver compatibility.
Training Loss. In CVIT training, the objective function is the l2 loss between the predicted next- timestep solution and the ground truth:
LeVi
B
| TEX
(52)
where B is the mini-batch, #x+1 is the ground-truth solution at the next timestep, f is the neural PDE solver being trained, t the timestep, and s the number of input timesteps.
However, different PDEs exhibit varying magnitudes in their state variables, which can lead to im- balanced training that overemphasizes PDEs with larger magnitudes. To address this, McCabe et al. [28] used the normalized mean squared error (NMSE), which scales each output to unit magnitude. We adopted the same loss function for X-MPP training:
| - ( 1
1 B
CMPP
(53)
| +
Since each PDE may produce tensors of different sizes, composing multiple PDEs in single mini- batch not feasible. Instead, we accumulate gradients across multiple mini-batches by summing them before performing parameter update.
Evaluation Metric. For evaluation, we use the normalized root mean squared error (NRMSE), which is the square root of the normalized mean squared error (NMSE), to compare with baseline methods. The metric is defined as
|.41 - f( o1))2
(54)
B
|
2 PEN
28

---

Table 5: Hyperparameters in the CVTT training.
DR
NS
SWE
Hyperpuram.
CVTT
X-CVI
CVT
X-CVI
CVTT
X-CVIT
(128,128)
(128,128)
(128,128)
(128.128)
(96,192)
(96,192)
Latent Dim, Embed Dim.
Depth, Attn Heads
512,384,5.6
512,384,66
312,384,56
512,384,5,6
512,384,56
512,384,56
Decoder Embeed Dim,
Decoder Alt Heads
Decoder Deyth
312,16,
512,16,1
312,16,1
512,16,1
512,1611
512,161
Input
ValTeea Splits
900050
9000050
65000010
65000010
5000/010
3600/010
Minibatch
16
16
16
52
32
Optim
Adam:
AdamW Decny 1P-6,
Adam
Adam|
Adam'W
Adam
LR Schedole
Wirm. Exp. Decay
Wirm. Exp.
Warm. Exp. Docay
Warm Exp.
Docny
Warm. Exp. Dansy
Warm Exp. Decay
IR end, peak
SE-4
IF-3
1F-6,
1E-6, 1F-3
1F-6,
18-3
decay, transit, warmup
0.9,
0.9, IE-5
Weight Docay
IE-5
IB-5
16-5
Grad. Clip
1.0
1.0
1.0
1.0
lieratioss
3E+5
3E-5
2B+5
2B+5
2B+5
2B+5
Table 6: Hyperparameters in the MPP from-scratch training and pretraining.
From Scraich
Pretraining
MPP
MPP
MPP
X- X-MPP
Hyperparam.
DR
SWE
CFD
DR
SWE
CFD
2D
1,2D
2D
1,2,3D
128
Input Timestops
16
16
16
16
16
4
Accumulation Steps
5
Epochs
120
120
120
120
120
120
120
120
120
110
Table 7: Hyperparameters in the MPP finetuning.
MPP
Hyperpeom
SWE2DD
DSID
CIDin BGTD
ARVID
TIEDD
SWE3T>
CHIDDD
1810
CHIID
IGG1
ATIIID CeIT
Inper Timesops
Accunalaioos
Kpoche
120
130
120
120
130
130
Hyperparameters. We report the hyperparameters used for training CViT, X-CViT, MPP, and X- MPP. The hyperparameters include the number of training epochs, train/val/test split, minibatch size, gradient accumulation steps, optimizer, weight decay, drop path probability, learning rate, learning rate scheduling, gradient clipping, and others.
For CVTT training, the hyperparameters are listed in Table 5. For MPP training, most hyperparame- ters shared across different settings, though some differ. The base values of the hyperparameters are:
Epochs: 120
Traie/Val/Test Splits: X %%10%% split on each dataset at the trajectory level, where X denotes subsample from 80% of the total dataset
Minibatch Size:
Accumulation Steps: 5
Optimizer: Adan
Weight Decay: IE-3
Drop Path Probability: 0.1
Learning Rate: DAdaptation
Learning Rate Scheduling: Cosine Decay
Gradient Norm Clipping: 1.0
These choices mostly follow the settings of McCabe al. [28]. The varying hyperparameters for
from-scratch training, pretraining, and finetuning described in Table 6 and Table 7.

---

E Partial Differential Equations
The PDE solution benchmarks are employed from PDEBench [33] and PDEArena [16]. Here are the specifications of the equations and their boundary conditions.
E.1 2D Shallow Water Equations.
The shallow water equations, derived from the general Navier-Stokes equations, provide reduced- order model for free-surface flows such as waves and dam breaks. the PDEBench benchmark, they are used to simulate 2D radial dam break scenario. The governing equations are given by:
O,h (hm)) +,,
1 4²h +
(hu) az
2 9-h2²
+ +,y(we)) -9hhh=b,
(55)
1
(hvo + Oy 1 + 2 gr16
+x((uv)) -grhOyb,
where h denotes the water depth, (u, v) are the horizontal velocity components, gr is the gravita- tional acceleration, and b(z, represents the bathymetry.
The initial condition corresponds to a circular bump the center of the domain Q = [-2.5, 2.5]²:
2.0,
if / <r,
1(1 = 0, x. v)
(56)
1.0, otherwise,
where the radius is randomly sampled from the uniform distribution U(0.3,0.7)
The simulation performed using the PyClaw finite volume solver. This PDE setup introduces realistic dynamics including shock propagation and wave reflections.
E.2 2D Compressible Fluid Dynamics.
The 2D compressible Navier-Stoker equations describe the dynamics of a compressible fluid, ac- counting for variations mass, momentum, and energy over time. These equations are fundamental for modeling gas flows where density changes are significant. The system consists of the conserva- tion laws for mass, momentum, and energy:
0,p + V (pv) = 0,
p00.v +v. Vv) =-vp (s+ !! 3
V(V V(V.v). v),
(57)
a (e+ (ove)) + V. (
1 501v13) v-
where p is the fluid density, v is the velocity vector, p is the pressure, e = p/([ - 1) is the internal energy (with r 5/3), and o' is the viscous stress tensor. The parameters 1) and & denote the shear and bulk viscosities, respectively.
To generate the data, the simulations employ second-order accurate HLLC finite volume solver for the inviscid terms, coupled with central differencing for the viscous contributions. This setup en- ables the benchmark to test model fidelity across wide range of physically realistic fluid dynamics problems.
E.3 2D Cerfusion-Reactiion Equation.
This equation models the interaction between two spatially distributed quantities (commonly an activator and an inhibitor) across a two-dimensional domain. It captures complex spatiotemporal behaviors such wave propagation and pattern formation, often seen in biological or chemical systems.
30

---

The system is described by:
deu = Du (Ozzut + uy + Ru (u,
(58)
0.v D (DxrV + Ogge) Ru(,,
where u(t, y) and vft, y) denote the concentrations of the activator and inhibitor, respec- tively. Du and Dr are their diffusion coefficients. The reaction dynamics are governed by the Nagumo model:
(u( =u u² k U1 Ro(u, v) v,
(59)
with à 5 10-3
In the benchmark setup, the simulation domain is x, E (-1,1) and t (0,5) The initial condi- tions generated using Gaussian noise, and Neumann boundary conditions (zero flux) are applied to ensure no flow across domain boundaries. Numerical solutions are computed using the finite volume method with fourth-order Runge- Kutta time integration.
E.4
2D Incompressible Navier-Stokes Equations (PDEArena).
The Navier-Stokes equations are a cornerstone of fluid dynamics, describing the motion of fluid substances under the influence of internal and external forces. In PDEArena, the two-dimensional incompressible Navier-Stokes equations are employed to investigate complex multi-scale flow phe- nomena. These equations govern the evolution of the velocity field v(t, IR2 in domain velocity-pressure
E R2, and are formulated in the velocity- pressure form as:
dv - (v v V)v Vp + +VV vvv=0, Öt
(60)
where p is the pressure, a is the kinematic viscosity (diffusion coefficient), and f represents external force, such buoyancy.
For simulation PDEArena, an additional scalar field introduced, representing a passive sealar (e.g.. particle concentration) that is advected by the velocity field and interacts with it through an external buoyancy force f = f)
The initial conditions include both the velocity and scalar fields. defined over a 128 x 128 grid with a resolution of Ar Ay = 0.25. The simulation time-step is At 1.5 seconds, and the domain is closed with Dirichlet boundary conditions v = 0 and Neumann conditions Üs/ /ör 0 for the scalar field.
The simulations are numerically solved using the #Flow framework and serve as a rich testbed for evaluating neural PDE surrogates in capturing advection-diffusion dynamics, vortex interactions, and response to varying force parameters.
E.5 Shallow Water Equations (PDEArena).
The shallow water equations are set of hyperbolic partial differential equations that describe the flow of a thin layer of incompressible fluid under the influence of gravity. They are derived from the incompressible Navier-Stokes equations by assuming that the horizontal length scales much larger than the vertical ones, leading to a vertically averaged flow model. In PDEArena, the shallow water equations are used model both local and global geophysical flow phenomena, such waves and large-scale atmospheric dynamics.
The equations govern the evolution of the fluid height htt,,) x) and the horizontal velocity field vt,, x) = (u, v) over a 2D domain x E IR,, and take the following form:
Oh at
+ (lv) O,
(61)
dv Øt
(v VJV + ghh = VV
where g is the gravitational acceleration, / the viscosity, and represents external forces such as wind stress or Coriolis effects.
31

---

The simulations are initialized with spatial fields of velocity and pressure, and are performed on a global grid with resolution 192 96, corresponding Ay = 3.75", and a temporal resolution of A = 48 hours. Periodic boundary conditions are applied in the longitudinal direction, while appropriate boundary conditions (e.g., reflective or free-slip) are used in the latitudinal direction.
spatial discretization of År 1.875%,
These simulations are generated using modified version of the SpeedyWeather. j1 framework. The setup allows for evaluating the performance of neural PDE surrogates on both velocity-pressure and vorticity-strean function formulations, capturing a wide range of scales and flow features rele- vant to climate and weather modeling.
E.6 ID Diffusion-Sorption Equation.
The diffusion-sorption equation models a diffusion process in porous media that is retarded by non-linear sorption mechanism. This type of process is relevant in real-world applications such as groundwater contaminant transport. In the PDEBench benchmark, transport of solute under the influence of sorption based on the Freundlich isotherm. The governing equation given by:
used simulate the ID
D R(u) R(u)
D Oenu((,
0,)(t,
(62)
where u(t, .z) denotes the concentration of the solute, D x 10 is the effective diffusion coefficient, and R(u) is the retardation factor that accounts for the sorption effect. The retardation factor is defined as:
1 1-6 Pohnatu
R(u)
(63)
with Ø = 0.29 the porosity, P. = 2880 the bulk density, hat = 3.5x100- the Freundlich coefficient, and n 0.874 the Freundlich exponent.
The initial condition sampled from a uniform distribution:
( 0, 1(0, for (0, 1).
(64)
The boundary conditions are given by:
0l = 1.0, uu4, = w(, 1),
(65)
where the second condition is Cauchy-type boundary involving a spatial derivative, introducing numerical challenges.
The simulation is performed using a finite volume method for spatial discretization and a fourth- order Runge- Kutta method for time integration. This PDE setup captures realistic nonlinear diffu- sion behaviors with singularities and complex boundary dynamics.
E.7 1D Burgers' Equation.
The Burgers' equation is fundamental nonlinear partial differential equation that models the in- terplay between convection and diffusion in fluid dynamics. It serves as a simplified prototype for the Navier-Stokes equations and is used study shock formation and dissipative processes. In the PDEBench benchmark, the 1D viscous Burgers' equation is used to simulate such nonlinear dynamics. The governing equation is given by:
0,((1, + Oe (EwR(..0)) L
(66)
where utt, is the velocity field and v is the diffusion coefficient, representing the kinematic vis- cosity.
The initial condition is constructed superposition of sinusoidal modes:
u(t 0,
A sin(ker + db:),
(67)
32

---

where the wave numbers k; 2TT La are randomly selected integers E [1, 7max A; ~ L(,, 1)
are amplitudes, o; ~ ((, 2 are phases, and L is the domain length. Additional operations such as applying the absolute value or a window function are applied with small probability to introduce further variability.
The domain is defined as x E (0, 1), and periodic boundary conditions are imposed:
w(t, 0) u(t, 1), (t, 0) = (t, 1).
(68)
The simulation is performed using second-order upwind finite difference scheme for the convective term and a central difference scheme for the diffusive term. This PDE setup is particularly suited for studying shock dynamics, nonlinear wave interactions, and the effect of viscosity on solution smoothness.
E.8
1D Advection Equation.
The advection equation is linear hyperbolic partial differential equation that models the transport of a conserved quantity without diffusion or reaction. It serves as a canonical example for study- ing wave propagation and translation phenomena in physics and engineering. In the PDEBench benchmark, the ID advection equation is used to simulate pure transport dynamics. The governing equation is given by:
0,t(t, x)) + ut, x) = 0,
(69)
where u(t, x)) represents the advected scalar quantity and the constant advection speed.
The initial condition defined as a superposition of sinusoidal waves:
u(t 0, Σ A sin(k; + di),
(70)
with k; 2xn./L, representing the wave numbers for randomly selected integers n; E [1, 7mmx]] amplitudes A; ~ L((0, 1), phases Φ; ~ L(00 2T). and L2 = 1 denoting the domain length. With small probability, transformations such as taking the absolute value or applying a window function are applied to the initial field to increase diversity.
The spatial domain is x E (0, 1), and periodic boundary conditions are employed:
mtt, 0) u(t, 1), a, 0) Otu(t, 1)-
(71)
The numerical solution is obtained using a second-order upwind finite difference scheme in both space and time. This PDE setup serves benchmark for evaluating models' ability to learn and reproduce translational dynamics with minimal distortion or dispersion.
E.9 ID Compressible Fluid Dynamics.
The ID compressible fluid dynamics (CFD) equations model the conservation of mass, momentum, and energy in a compressible fluid. Derived from the general compressible Navier-Stokes equations, they are used to simulate phenomena such as shock waves, rarefaction, and contact discontinuities. In the PDEBench benchmark, this setup includes various configurations such as inviscid flow, vis- cous flow, and shock-tube initial conditions. The governing equations are given by:
a.p Oz (pn) = 0,
(p(p)) + O,p (pu? + +p) 020, F. 0, [(F + p)u] 0, 0,()), (uo),
(72)
where P is the density, u is the velocity, p is the pressure, E + puu is the total energy with P r-1- and a the viscous stress term defined as a = np0,2 u for shear viscosity r). The ratio of specific heats is set to r = 5/3.
internal energy c =
The benchmark includes multiple initial condition types:
Random field: Initial p. u, and p fields generated as smooth random perturbations using a superposition of sine waves.
33

---

Shock tube: A Riemann problem where left and right states (p. p) are sampled from uniform distributions with sharp discontinuity a random position.
The spatial domain is E (0, 1) with two boundary condition types:
Periodic: Fields wrap around the domain,
Out-going: Ghost cells copy the nearest interior value to allow waves to exit.
The simulations are performed using second-order HLLC Riemann solver with MUSCL recon- struction inviscid cases, and central differencing for viscous terms. This PDE setup challeng- ing due to strong nonlinearity, shock formation, and sensitivity to initial and boundary conditions.
E.10
3D Compressible Fluid Dynamics.
The 3D compressible Navier-Stole equations govern the motion of gases where density, pressure, and velocity fields evolve in space and time. This system models conservation of mass, momen- tum, and energy in three dimensions, making essential for simulating realistic high-speed flows, turbulence, and shock dynamics. The equations are same as Eq. 57, but their variables are three dimensional, e.g. R
Numerical solutions are generated using a second-order HILLC scheme for the inviscid part and central differencing for viscosity. This 3D setting significantly increases the complexity of the benchmark, introducing challenges in both computational cost and physical realism for surrogate models.
34


Reviews:
[{"review_id": "ai6FP31ZRO", "reviewer": "Reviewer_XMsE", "summary": "This 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.", "questions": "- 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.", "limitations": "Equivariance 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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Originality: 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.", "was_revised": true, "final_justification": "I thank the authors for their discussion, and maintain my score. This is a solid paper.", "rebuttal": "Thank you for the positive comments. Below are our responses to the remaining concerns.\n\n---\n\n> The paper may benefit from more diverse set baselines (e.g., FNOs) in the experiments section.\n> Why not compare against FNOs? These are natural baseline for cross-resolution and few-shot generalization.\n\nAs you understood, our method was designed as an add-on to existing PDE models, so we did not include many baselines and instead focused on the comparison between a PDE model and its XNN version. Nonetheless, FNO can serve as a comparable example for cross-dimension functionality, so we compared them in terms of few-shot learning. As shown in the tables, X-MPP significantly outperforms FNO in the few-shot setting.\n\nTable 1. Test NRMSE in Burgers’ (BS) equation\n| #Data |     295    |     736    |    1471    |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.0485   |   0.0319   |   0.0214   |\n| X-MPP | **0.0337** | **0.0146** | **0.0038** |\n\nTable 2. Test NRMSE in Advection (ADV) equation\n| #Data |     197    |     491    |     981    |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.2516   |   0.2371   |   0.2323   |\n| X-MPP | **0.0795** | **0.0484** | **0.0202** |\n\nTable 3. Test NRMSE in compressed fluid dynamics (CFD) equation\n| #Data |     58     |     145    |     290    |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.0679   |   0.0522   |   0.0355   |\n| X-MPP | **0.0289** | **0.0125** | **0.0094** |\n\nTable 4. Test NRMSE in diffusion sorption (DS) equation\n| #Data |     12     |     30     |     60     |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.0350   |   0.0163   |   0.0078   |\n| X-MPP | **0.0098** | **0.0042** | **0.0025** |\n\n\n\n---\n\n> However, several claims such as axis permutation robustness and generalization to new boundary conditions are not empirically validated.\n\nWe did not include experiments on axis-permutation generalization because our focus is on number-of-axes generalization rather than axis-permutation generalization. In fact, in the PDE task, breaking axis-permutation equivariance using positional encoding performed better than enforcing equivariance. This is because retaining the semantics of each axis through positional encoding is beneficial.\n\n---\n\n> The ablation between SXNN and GXNN is performed only on a synthetic task.\n\nThe ablation between SXNN and GXNN on real data is a good suggestion. However, on real data, the best performance was achieved by mixing SXNN and GXNN across different layers, which is why we compared them on synthetic data.\n\n---\n\n> No error bars or standard deviations are reported, making it difficult to assess robustness or reproducibility.\n\nAs you pointed out, error bars or standard deviations are missing. Unfortunately, we were not able to run a sufficient number of experiments to measure the standard deviation because a single experiment (especially pretraining) takes too much time in our setting. We also mentioned this in the checklist. However, we partially observed that the results are consistent.\n\n---\n\n> The handling of domain-specific features (e.g., BCs, geometry, time) is not entirely clear.\n> How are geometry and boundary conditions encoded? Is the model tied to Cartesian grids, or can it handle irregular domains?\n\nThe handling of domain-specific features is determined by the baseline methods we modified for XNN; i.e., X-MPP follows the handling method of MPP. Additionally, the baseline methods target PDE solutions defined only on regular grids. According to MPP, the boundary conditions, geometry, and time are not separately input to the model. Instead, MPP is pretrained on multiple PDE solutions with varying boundary conditions and equations, allowing it to learn general patterns of PDE solutions. Since we eventually finetune the model on a specific PDE with known boundary conditions and geometry, it is not necessary to encode them separately as inputs. The only thing MPP handles separately is a periodicity of the boundary condition. Depending on the periodicity of the boundary condition, MPP determines to use between sequential position bias and periodic position bias in the attention layer. We will clarify this in the final version.\n\n---\n\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> Equivariance to axis permutations: Theorems 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\nAfter reviewing your comments, we realized that the GXNN used in the synthetic task was not implemented to be equivariant, although it was still dimension-agnostic. Therefore, we modified GXNN to be equivariant and re-evaluated its performance. As shown in Table 5, SXNN and the corrected GXNN satisfy equivariance, whereas CNN does not. In Table 6, the corrected GXNN still demonstrates strong performance compared to CNN on high-dimensional data.\n\nTable 5. Equivariance error\n|             | 3D CNN   | SXNN-L   | Corrected GXNN |\n|-------------|----------|----------|----------------|\n| Equiv. Err. | 0.453423 | 0.000000 |       0.000000 |\n\nTable 6. Test accuracy (%) in the synthetic task.\n|    | 3D CNN | GXNN  | Corrected GXNN |\n|----|--------|-------|----------------|\n| 2D |  97.78 | 95.45 |          95.83 |\n| 3D |  67.35 | 79.85 |          81.09 |\n| 4D |  65.93 | 70.86 |          72.94 |\n| 5D |  66.39 | 70.48 |          68.37 |\n\n---\n\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> Robustness to axis semantics: Claimed 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\nWe intended number-of-axes generalization rather than axis-permutation generalization from the equivariance. Moreover, axis-permutation generalization is not preferred in an application like a PDE solver. We investigated whether axis-permutation equivariance would help PDE generalization. To this end, we used an axis-permutation equivariant NN (XNN) with axis-permutation invariant positional encoding to guarantee equivariance of the PDE solver. However, it showed lower performance than XNN with *non-invariant positional encoding*, which breaks equivariance. Thus, although the XNN layers should be axis-permutation equivariant to achieve agnosticism to the number of axes, the final model *should not be equivariant* in order to preserve the semantics of each axis. We will elaborate on this perspective in the final version.\n\n---\n\n> Generalization to new geometries or BCs: All 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.\n\nGeneralization to new geometries or boundary conditions (BCs) strongly depends on the baseline PDE model used in XNN. Neural PDE solvers for irregular domains often introduce new encoders and objective functions [1,2]. Although irregular domains may break the axis-permutation equivariance of the PDE, XNN-based PDE solvers remain dimension-agnostic, and we are already breaking equivariance through non-invariant positional encodings. For example, MPP, one of our baseline PDE models, uses multiple position biases in the attention layers depending on the periodicity of the boundary conditions. The same approach is seamlessly applicable in X-MPP. On the other hand, non-grid approaches such as Neural Processes [3] take feature-coordinate pairs as input. Therefore, they do not require a dimension-agnostic architecture, as the coordinate information is already explicitly provided.\n\n[1] Khara et al., 2022, Neural PDE Solvers for Irregular Domains\n\n[2] Boussif et al., 2022, MAgNet: Mesh Agnostic Neural PDE Solver\n\n[3] Garnelo et al., 2018, Neural Processes"}, {"review_id": "JUNu5Efy8E", "reviewer": "Reviewer_Q3vK", "summary": "The 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.", "questions": "- 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?", "limitations": "The authors did mention the limitations of their approach.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns, so I increase the score to 4.", "rebuttal": "Thank you for your valuable feedback. Below are our responses to the remaining concerns.\n\n---\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\nThat is a valid point. I followed CViT’s statement claiming SOTA, which is why I referred to our evaluation as SOTA. However, we later figured out CViT was not the best model in our setting. We will correct the claim. As future work, it would be interesting to apply XNN to Poseidon, a SOTA PDE model.\n\n---\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\nWe included only 1D and 3D downstream tasks to demonstrate the unseen dimension generalization capability. For a comparable assessment, we evaluated our method on a new downstream task called 2D incompressible fluid dynamics (IFD), corresponding to seen dimensionality but unseen PDE data. As a result, the axial version of MPP, X-MPP, significantly outperformed MPP due to the embedded inductive bias in the architecture.\n\nTable 1. Finetuning performance on unseen PDE, IFD\n| #PretrainData |    11958   |    21912   |\n|:-----:|:----------:|:----------:|\n|  MPP  |   0.0350   |   0.0163   |\n| X-MPP | **0.0098** | **0.0042** |\n\n---\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> The authors should much better clarify the experimental setup and better isolate seen vs. unseen tasks.\n\nAs you said, Table 2 includes an unrealistic setting, such as including pretraining data in the downstream tasks. We would like to clarify that Table 2 is intended to show the expressivity of XNN, not its functionality as a foundation model. We tested the unseen dimension generalization capability of our model in Fig 4. In this case, the seen data are all 2D PDEs, and the unseen data are all 1D and 3D data (finetuned in the downstream task).\n\n---\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\nWe investigate the scaling properties with respect to the size of the pretraining dataset, in comparison to the Fourier Neural Operator (FNO), on 1D compressible fluid dynamics (CFD). FNO is trained only on CFD, whereas X-MPP is pretrained on multiple 2D PDEs. As shown in the table, X-MPP also scales well as the size of the pretraining dataset increases and outperforms FNO when provided with a sufficient amount of pretraining data.\n\nTable 2. Finetuning performance on 1D CFD.\n|           Model          |  X-MPP |  X-MPP |  X-MPP |   FNO  |\n|:------------------------:|:------:|:------:|:------:|:------:|\n| #PretrainData/#TotalData |  0.01  |   0.1  |    1   |    1   |\n|           NRMSE          | 0.3469 | 0.0198 | 0.0052 | 0.0232 |\n\n---\n\n> Does the axis-permutation equivariance introduce any inductive bias that could limit performance on inherently asymmetric data?\n\nActually the axis-permutation equivariance is used to design a dimension-agnostic architecture rather than to impose an inductive bias. Therefore, we do not need to guarantee equivariance for asymmetric data. Indeed, our PDE solver is not axis-permutation equivariant due to the non-invariant positional encodings, although each layer is equivariant. This is because the semantics of each axis are also important.\n\n---\n\n> How does the expressivity of GXNN compare to conventional ViT or CNN models when trained on high-dimensional data?\n\nTypically, equivariant models including XNNs sacrifice expressivity to guarantee equivariance. Hence, the expressivity of GXNN would also be lower than that of conventional ViT and CNN for the same number of model parameters. However, ViT or CNN requires an exponentially large number of parameters to process high-dimensional data, whereas GXNN can handle the same high-dimensional data with a relatively small number of parameters.\n\n---\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\nAlthough rotation is non-trivial in general PDEs, the transposition of spatial dimensions has the potential to produce a similar effect as equivariance. However, equivariant models are theoretically more robust because they directly constrain the loss surface through the restricted architecture [1].\n\n[1] Wang et al., 2022, Data Augmentation vs. Equivariant Networks: A Theory of Generalization on Dynamics Forecasting"}, {"review_id": "oBE3s2aS0V", "reviewer": "Reviewer_yeNs", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors adequately addressed the main concerns in the rebuttal. They clarified the scalability setup, provided justification for not using large models, and explained differences in CViT performance and patchification claims. While large-scale model validation remains missing, the core contributions are sound, and the method shows strong performance and generality. I recommend acceptance.", "rebuttal": "Thank you for the positive reviews. Below are our responses to the remaining concerns.\n\n---\n\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\nThat is a valid concern about this paper. We actually tried such a setting to demonstrate the benefits in large models, but our setup was not sufficient to test large-scale transformers due to the heavy memory requirements of processing PDE data. Instead, we leveraged a relatively large amount of data to demonstrate scalability. Below are the results on scalability with respect to the size of the pretraining dataset, in comparison to the Fourier Neural Operator (FNO), on 1D compressible fluid dynamics (CFD). FNO is trained only on CFD, whereas X-MPP is pretrained on multiple 2D PDEs. As shown in the table, X-MPP also scales well as the size of the pretraining dataset increases and outperforms FNO when provided with a sufficient amount of pretraining data.\n\nTable 1. Finetuning performance on 1D CFD.\n|           Model          |  X-MPP |  X-MPP |  X-MPP |   FNO  |\n|:------------------------:|:------:|:------:|:------:|:------:|\n| #PretrainData/#TotalData |  0.01  |   0.1  |    1   |    1   |\n|           NRMSE          | 0.3469 | 0.0198 | 0.0052 | 0.0232 |\n\n\n---\n\n> 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\nAs you pointed out, designing patchification is not restricted to specific dimensions. What we meant by “limited to 2D inputs” is that 2D patchification is only applicable to 2D inputs, not 1D or 3D, and requires modifying parameters to process 1D or 3D data. In contrast, XNN patchification is applicable to inputs of any dimension with the same parameters. We will clarify this in the final version.\n\n---\n\n> 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\nWe trained CViT on SWE from PDEArena, while MPP was trained on SWE from PDEBench. The main difference between the CViT paper and ours lies in the evaluation metric. PDE solutions typically involve multiple variables represented as a feature vector in the data. CViT used relative L2 error, which is the average of normalized MSE across variables, whereas we used NRMSE, which is the MSE normalized by the norm of the feature vector. Additionally, the CViT paper reported 5-timestep prediction error, while we reported 1-timestep prediction error."}, {"review_id": "93NZOdSOuK", "reviewer": "Reviewer_Rmj3", "summary": "The 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.", "questions": "- 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.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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?", "was_revised": true, "final_justification": "The paper is clearly written and the authors were able to address my initial concerns. I trust that with the suggested improvements and modifications the authors mentioned in the rebuttal, this manuscript should be published.", "rebuttal": "Thank you for the constructive comments. Below are our responses to the remaining concerns.\n\n---\n\n> some work that is clearly related such as factorized convolutions/attention are not referenced\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\nThey are good suggestions. There is a similarity between this work and factorized convolutions/attention in terms of low-rank approximation of a single operation. I will clearly state this in the final version.\n\n---\n\n> description of experimental method and baselines are somewhat lacking, it is not immediately clear how the efficacy of this method is tested.\n\nWe will add details of the experimental procedures in the final version. In a nutshell, we followed the basic training procedures of baselines such as CViT and MPP, which take a few timesteps as input and predict the next timestep of a PDE solution. These models are all based on Vision Transformers with some revisions optimized for PDE learning. We evaluated architectural expressivity in Table 2 in terms of from-scratch training, pretraining, and finetuning using only 2D PDE solutions. In Figure 4, we evaluated few-shot learning capability on data from unseen dimensions, that is, the models were pretrained on 2D PDEs only and finetuned with a few samples from 1D or 3D PDEs, demonstrating the benefit of multi-dimensional pretraining in XNN-based foundation models.\n\n---\n\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\n\nWe have conducted one more baseline FNO, a widely used architecture for PDE solutions. It also has dimension-agnostic functionality and is typically used for equation-specific learning. The Test NRMSE comparison with X-MPP on 1D few-shot learning task as below. As shown in the tables, X-MPP significantly outperforms FNO in the few-shot setting regardless of the size of the downstream dataset.\n\nTable 1. Test NRMSE in Burgers’ (BS) equation\n| #Data |     295    |     736    |    1471    |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.0485   |   0.0319   |   0.0214   |\n| X-MPP | **0.0337** | **0.0146** | **0.0038** |\n\nTable 2. Test NRMSE in Advection (ADV) equation\n| #Data |     197    |     491    |     981    |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.2516   |   0.2371   |   0.2323   |\n| X-MPP | **0.0795** | **0.0484** | **0.0202** |\n\nTable 3. Test NRMSE in compressed fluid dynamics (CFD) equation\n| #Data |     58     |     145    |     290    |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.0679   |   0.0522   |   0.0355   |\n| X-MPP | **0.0289** | **0.0125** | **0.0094** |\n\nTable 4. Test NRMSE in diffusion sorption (DS) equation\n| #Data |     12     |     30     |     60     |\n|:-----:|:----------:|:----------:|:----------:|\n|  FNO  |   0.0350   |   0.0163   |   0.0078   |\n| X-MPP | **0.0098** | **0.0042** | **0.0025** |\n\n---\n\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\nYes, $P(\\theta)$ is the prior distribution of the parameter $\\theta$. de Finetti’s theorem tells us that a permutation-invariant (exchangeable) neural network should take the form of a Deep Sets-style architecture, as shown in Eq. 2. Specifically, if $P(X_i|\\theta) = \\exp(\\langle \\psi(x), \\theta \\rangle - g(\\theta))$ and $P(\\theta) = \\exp(\\langle \\theta, \\alpha \\rangle - mg(\\theta) - \\psi(\\alpha, m))$ in terms of exponential families, then Eq. 4 can be written as $$P(X_1,\\ldots,X_n)=\\exp\\left(\\phi\\bigg(\\alpha+\\sum_{i=1}^n\\psi(X_i), m+n\\bigg)-\\phi(\\alpha,m)\\right),$$ which shares the same structure as Eq. 2. That is, an axis-permutation-invariant neural network should also share such an architecture by treating each element $X_i$ as an axis. This statement is also described in Eq. 2 of the Deep Sets paper [1]. The original description was insufficient for building intuition, so we will elaborate more in the final version.\n\n[1] Zaheer et al., 2018, Deep Sets\n\n---\n\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\nGeneralization from lower to higher dimensions indeed addresses the bottleneck, as the model trained on 1D and 2D data can perform well on 3D data after fine-tuning with a *small amount of 3D data*. We provided evidence for this in the last plot of Figure 4: with only 500 data points (trajectories), the model was able to generalize to 3D data."}]

Author rebuttal:

Meta-review: The paper addresses the challenge of building foundation models for physical dynamics when the input dimensionality may vary across different physical phenomena—for example, when training on multiple PDE simulations. Classical solutions include designing dedicated encoder-decoders for each case or fixing a maximum dimensionality and padding missing variables with zeros.
The paper introduces a novel dimension-agnostic framework, called Axial Neural Networks (XNNs), which can handle tensors of arbitrary rank by treating axes as set elements or graph nodes, and applying message passing to model cross-axis dependencies. The approach is inspired by previous work on equivariant neural networks. Two variants of the method are proposed and formally shown to be equivariant under axis permutations. This framework is instantiated on two PDE foundation model backbones (MPP and CViT) and evaluated on a range of PDE simulation problems.

The reviewers agree on the relevance of the problem for developing foundation models for physics. They find the approach original and non-trivial, with formal proofs supporting the claims and convincing experimental results across diverse settings. They consider that the authors’ responses and additional experiments satisfactorily address most of their questions.
I recommend acceptance.
Decision: Accept (spotlight)