SMPE: A Framework for Multi-Dimensional Permutation Equivariance

In $1$ , a one-dimensional invariant network was constructed based on the proposed variable node attention pooling mechanism, which effectively addressed graph classification problems. Due to its permutation invariance, it is applicable to graphs of different sizes.

The work in $2$ introduces the Substructure Assembling Network (SAN), a graph neural network model that effectively learns interpretable and discriminative graph representations for classification. It overcomes challenges by hierarchically assembling substructures using the Substructure Assembling Unit (SAU) and enhancing performance with Soft Sequence with Context Attention (SSCA).

The goal of these two works is to model one-dimensional invariant mappings, which is relevant to our work but not the primary focus of our investigation. We have incorporated these two works into the 'Related Work' section and provided additional discussions.

$1$ Meltzer P, Mallea M D G, Bentley P J. Pinet: A permutation invariant graph neural network for graph classification $J$ . arXiv preprint arXiv:1905.03046, 2019.

$2$ Yang Y, Guan Z, Zhao W, et al. Graph substructure assembling network with soft sequence and context attention $J$ . IEEE Transactions on Knowledge and Data Engineering, 2022, 35(5): 4894-4907.

The hyperparameters used in our experiments are listed in Appendix E.4 of the revised paper. For example

• Pooling functions: Without loss of generality, all the pooling operations are set to be the mean function.

• Training strategy: We perform batch iterations for $8\times 10^{4}$ times with a batch size of 512, and reduce the learning rate by half after $4\times 10^{4}$ iterations. This strategy for three tasks is the same. For the design task of MIMO and cell-free MIMO system transmission, we randomly assign SNR values to each training and testing sample.

• Activation & normalization: To accelerate convergence and enhance the stability of the training process, we applied the ReLU activation and LN to $\mathtt{ETN2D}$ and $\mathtt{ETN3D}$.

• Monte-Carlo simulation: Each experimental result is the average of three valid experiments.

Many thanks for your careful reviews and constructive comments.

The running times of all the methods have been compared in Appendix F.2. Specifically, the following table provides a comparison of the runtime for each network across different tasks. The runtime refers to the average time (ms) taken for a batch during GPU inference, with a batch size of 1000. In the first row, '1', '2' and '3' denote experiments in Sections 4.1, 4.2, and 4.3, respectively. ${\mathcal{A}}$, ${\mathcal{B}}$, and ${\mathcal{C}}$ represent the scenario configurations for each task during inference. It can be observed that the running time of $\mathtt{SMPEN-DS}$ is comparable to that of $\mathtt{ETN}$, and the running time of $\mathtt{SMPEN2D-TF}$ is similar to that of $\mathtt{CITransNet2D}$.

The source code and associated datasets will be provided soon, and we are currently making efforts to organize the code and add comments to improve reproducibility.

First, we would like to illustrate the primary sources of complexity in the framework. The computational complexities and parameter quantities are shown as follows

The total number of input features is $\Pi\_{n=1}^NK\_n$. This implies that under the condition of a fixed input dimensions $N$, the computational complexity orders of $\mathtt{SMPEL-DS}$ and $\mathtt{SMPEL-TF}$ are lower than the second order of $\Pi\_{n=1}^NK\_n$, and the parameter quantity is also independent of $\Pi\_{n=1}^NK\_n$. Therefore, the main factor that limits the scalability of the framework is the exponential expansion of network size with the growth of dimension $N$. The number $2^N$ in this table arises from the algorithm's requirement to obtain pooled information for every subset in the power set of ${\mathbb N}\backslash\{\emptyset\}$ and incorporate it into the computation.

The exploration of efficient alternatives that do not significantly compromise model capacity is an interesting topic. We are currently conducting further research about this problem. From a practical implementation perspective, the following method can achieve a trade-off between performance and complexity: Train the network for multiple times, with each training randomly selecting $H$ subsets from the $2^N-1$ subsets to include in the computation. Then, the best performing networks are retained. This approach allows for the identification of the most important $H$ combinations, effectively reducing the scale from $2^N-1$ to $H$.

The ablation experiments have been included in Appendix F.1 of the revised paper. Specifically, to validate the superiority of $g\_{\rm serial}$ (represented by $\mathtt{ParaN2D-TF}$) over $g\_{\rm para}$ (represented by $\mathtt{SerialN2D-TF}$), as well as the superiority of $g\_{\rm SMPE}$ over $g\_{\rm serial}$, we compared the performance of these three methods on the task in Section 4.2. In most scenarios, as observed from the table below, $g\_{\rm SMPE}$ outperforms $g\_{\rm serial}$, which, in turn, outperforms $g\_{\rm para}$.

The hyperparameters used in our experiments are listed in Appendix E.4 of the revised paper. These hyperparameters of the baseline models have been fine-tuned, and we reported the best results for comparison. For example:

• Pooling functions: Without loss of generality, all the pooling operations are set to be the mean function.

• Activation & normalization: To accelerate convergence and enhance the stability of the training process, we applied the ReLU activation and LN to $\mathtt{ETN2D}$ and $\mathtt{ETN3D}$.

• Parameter initialization: With the exception of networks $\mathtt{SMPEN2D-TF}$ and $\mathtt{SMPEN3D-TF}$ in the task of Section 4.3, which were initialized using the default method, we employed Kaiming initialization for the weight parameters of the linear layers in networks used for other tasks.

• Monte-Carlo simulation: Each experimental result is the average of three valid experiments.

In this work, we propose a plug-and-play paradigm for the construction of arbitrary high-dimensional equivariant networks with well-designed one-dimensional equivariant networks. Specifically, we first provide a definition of multi-dimensional equivariance based on algebraic definitions. Based on this, we further propose two methods to combine one-dimensional equivariant functions, which have been mathematically proven to satisfy the equivariant property in multiple dimensions. Moreover, we employed feature reuse and introduced global information at each layer to enhance performance while maintaining the multi-dimensional equivariance. Additionally, it has been proven that the constructed equivariant layer can degenerate into the multi-dimensional equivariant linear layer characterized by a complex mathematical expression. The substitutability of one-dimensional equivariant functions in the proposed framework implies its enhanced generality, effectively transforming the design problem of multi-dimensional equivariant networks into the selection of one-dimensional equivariant networks.

Thank you for your thorough reviews and valuable feedback.

The high-order permutation equivariance considers the high-order features of a single type of object. Assuming there are $K$ objects, their $N$th-order features among each other are denoted as ${{{\bf{\mathsf{X}}}}}\in{\mathbb R}^{K\times K\times \cdots\times K\times D\_X}$, where $K$ is repeated $N$ times in the dimension, so it can also be written as ${\bf{\mathsf{X}}}\in{\mathbb R}^{K^N\times D\_X}$ $1$ . The function $f$ exhibits $N$th-order permutation equivariance if it satisfies $\begin{aligned} f(\sigma\circ{\bf{\mathsf{X}}})=\sigma\circ f({\bf{\mathsf{X}}}),\ (\sigma\circ{\bf{\mathsf{X}}})\_{[k\_1,...,k\_N,:]} = {\bf{\mathsf{X}}}\_{[\sigma^{-1}\circ k\_1,...,\sigma^{-1}\circ k\_N,:]},\ \sigma\in {\mathbb S}\_K. \end{aligned}$ It can be observed that the same permutation $\sigma$ is employed across all dimensions. High-order permutation equivariance is an extension of one-dimensional permutation equivariance $1$

$2$ . The relevant literature on high-order equivariance also includes references $3$ , $4$ , $5$ , and so on. The multi-dimensional equivariance considers features among multiple types of objects, i.e., ${{{\bf{\mathsf{X}}}}}\in{\mathbb R}^{K\_1\times K\_2\times \cdots\times K\_N\times D\_X}$. A mapping $f$ is said to be $N$-dimensional equivariant if

In order to enhance the practicality and complexity of the task, we incorporated real-world datasets and conducted experiments on the transmission design task discussed in Section 4.3. We replaced the original results in Table 3 with the new results and move the original Table 3 to the appendix as supplementary results.

The new dataset is provided by QuaDRiGa $1$ . QuaDRiGa, short for QUAsi Deterministic RadIo channel GenerAtor, is used for generating realistic radio channel impulse responses for system-level simulations of mobile radio networks. These simulations are used to determine the performance of new digital-radio technologies in order to provide an objective indicator for the standardization process in bodies like 3rd Generation Partnership Project (3GPP). Furthermore, the channel generator is continuously updated with the evolution of wireless communication standards $2$ . For more details, please refer to the official website of QuaDRiGa: https://quadriga-channel-model.de.

The dataset configuration we used is as follows

where the scenario '3GPP\_38.901\_UMa\_NLOS ' refers to a specific configuration defined by the 3GPP in technical specification 38.901

$3$ . This scenario is designed to model wireless communication in urban macro environments where non-line-of-sight (NLOS) conditions are prevalent. In this scenario, QuaDRiGa simulates a multi-path channel model that includes the effects of buildings, structures, and other obstacles typically found in urban macro environments. The model considers the propagation of radio waves in NLOS conditions, accounting for reflections, diffractions, and scattering from surrounding objects.

The dataset from the generator closely approximates real-world scenarios, ensuring the practical value of experimental results.

$1$ Jaeckel S, Raschkowski L, Börner K, et al. QuaDRiGa: A 3-D multi-cell channel model with time evolution for enabling virtual field trials $J$ . IEEE transactions on antennas and propagation, 2014, 62(6): 3242-3256.

$2$ Wang C X, Bian J, Sun J, et al. A survey of 5G channel measurements and models $J$ . IEEE Communications Surveys & Tutorials, 2018, 20(4): 3142-3168.

$3$ 3GPP TR 38.901 v16.1.0, "Study on channel model for frequencies from 0.5 to 100 GHz," Tech. Rep., 2019.

The computation methods of these two distance metrics are as follows

About graph neural networks: Graph neural networks are designed to solve graph-related problems, which often consider modeling of nodes and edges

$1$ . Due to the unordered nature of nodes, graph neural networks possess permutation equivariance. It is worth noting that most graph neural networks can be viewed as one-dimensional equivariant networks, as evidenced by references $2$ and $3$ . As our work primarily focuses on the construction of multi-dimensional equivariant networks, our approach is more general than graph neural networks that target one-dimensional equivariant network design.

About pooling methods: The pooling function is utilized to aggregate multiple features, and this characteristic can be employed to construct permutation-invariant networks. In Sections 3.2 and 3.3 of the paper, pooling functions are respectively introduced to acquire global information and provide invariance. Most pooling functions can be applied within the proposed framework, as the multi-dimensional equivariance maintained by our framework does not depend on the selection of pooling function. Therefore, designing pooling functions is beyond the scope of this paper.

$1$ Wu Z, Pan S, Chen F, et al. A comprehensive survey on graph neural networks $J$ . IEEE transactions on neural networks and learning systems, 2020, 32(1): 4-24.

$2$ Vignac C, Loukas A, Frossard P. Building powerful and equivariant graph neural networks with structural message-passing $J$ . Advances in neural information processing systems, 2020, 33: 14143-14155.

$3$ Velickovic P, Cucurull G, Casanova A, et al. Graph attention networks $J$ . stat, 2017, 1050(20): 10-48550.

There are various ways to apply one-dimensional equivariant functions to multiple dimensions, but not all of them can maintain the equivariant property in multiple dimensions. For example, one can flatten the equivariant dimensions of the multi-dimensional data and apply one-dimensional equivariant functions. Alternatively, one can fix one equivariant dimension and flatten the remaining dimensions along with the feature dimension to create a new feature dimension. However, these approaches do not satisfy the multi-dimensional equivariance.

In our work, we propose a plug-and-play paradigm for the construction of arbitrary high-dimensional equivariant networks with well-designed one-dimensional equivariant networks. Unlike the examples mentioned, these proposed combination methods have been mathematically proven to satisfy the equivariant property in multiple dimensions. Moreover, we employed feature reuse and introduced global information at each layer to enhance performance while maintaining the multi-dimensional equivariance. Additionally, it has been proven that the constructed equivariant layers can degenerate into the multi-dimensional equivariant linear layers described in $1$ , which is characterized by a complex mathematical expression. In summary, while our proposed method has a straightforward expression, its foundation and underlying analysis are not straightforward. The concise expression reduces the cost of further exploitation and analysis, reflecting the principle we advocate.

$1$ Hartford J, Graham D, Leyton-Brown K, et al. Deep models of interactions across sets $C$ . International Conference on Machine Learning. PMLR, 2018: 1909-1918.

The involved mappings are relatively complex non-linear mappings. A linear mapping between two vector spaces needs to satisfy additivity, which means $f(X+Y)=f(X)+f(Y)$. Clearly, the pseudo-inverse and the optimal solution for downlink transmission design do not satisfy this property. Specifically, for the transmission design task in Sections 4.2 and 4.3, finding the optimal results requires solving Problem (26) and Problem (28), respectively. These two problems are typical NP-hard problems and are difficult to solve $1$ . The typical solving algorithms usually involve iterative operations on high-dimensional matrices and require multiple initializations to select the optimal results $2$ , which is very time consuming and can not guarantee optimality.

$1$ Zhao X, Lu S, Shi Q, et al. Rethinking WMMSE: Can its complexity scale linearly with the number of BS antennas? $J$ . IEEE Transactions on Signal Processing, 2023, 71: 433-446.

$2$ Christensen S S, Agarwal R, De Carvalho E, et al. Weighted sum-rate maximization using weighted MMSE for MIMO-BC beamforming design $J$ . IEEE Transactions on Wireless Communications, 2008, 7(12): 4792-4799.

We sincerely thank you for your careful reviews and constructive comments.

The proposed framework can be easily extended to higher dimensions, and we illustrate the framework of network construction by using a task of 4D mapping as an example.

Consider the precoding design for cell-free multi-user multiple-input multiple-output (MU-MIMO) systems There is a CPU controlling $M$ distributed APs to support $K$ UEs, where each AP is equipped with $N\_{T}$ antennas and each AP is equipped with $N\_{R}$ antennas. The channel between the $m$-th AP and the $k$-th UE is denoted as ${{\bf H}}\_{k,m}\in{\mathbb{C}}^{N\_{T}\times N\_{R}}$, and the corresponding precoding matrix is ${\bf W}\_{k,m}\in{\mathbb{C}}^{N\_{T}\times N\_{R}}$. By stacking the channel and precoding matrices belonging to all APs and UEs, we have ${{\bf{\mathsf{H}}}}\in{\mathbb{C}}^{M\times K\times N\_{T}\times N\_{R}}$ and ${{\bf{\mathsf{W}}}}\in{\mathbb{C}}^{M\times K\times N\_{T}\times N\_{R}}$. Similar to the task in Section 4.3, we consider the optimization problem of maximizing the sum rate under transmit power constraints. It can be proved that the mapping from ${{\bf{\mathsf{H}}}}$ to certain optimal ${\bf{\mathsf{W}}}^{\star}=f({{\bf{\mathsf{H}}}})$ is 4D permutation equivariant.

For this task, we concatenate the real parts and imaginary parts of ${{\bf{\mathsf{H}}}}$, and the shapes of the input and the target output are ${\widetilde {{\bf{\mathsf{H}}}}}\in{\mathbb R}^{M\times K\times N\_{T}\times N\_{R}\times 2}$ and ${\widetilde {{\bf{\mathsf{W}}}}}\in{\mathbb R}^{M\times K\times N\_{T}\times N\_{R}\times 2}$, respectively. For this, the architecture of layer ${\rm SMPEL}:{\mathbb R}^{M\times K\times N\_{T}\times N\_{R} \times D}\to {\mathbb R}^{M\times K\times N\_{T}\times N\_{R}\times D}$ can be represented as follows: $\begin{aligned} {\rm SMPEL}({\bf{\mathsf{X}}}) = {\rm FC}\_1(\sigma({\rm LN}({\bf{\mathsf{M}}}))),\\ {\bf{\mathsf{M}}} = {\rm FC}\_2([{\bf{\mathsf{O}}}\_1, {\bf{\mathsf{O}}}\_2, {\bf{\mathsf{O}}}\_3, {\bf{\mathsf{O}}}\_4, {\bar {\bf{\mathsf{X}}}}\_{{\mathbb P}\_1}, ..., {\bar {\bf{\mathsf{X}}}}\_{{\mathbb P}\_{15}}]),\\ {\bf{\mathsf{O}}}\_0 = {\bf{\mathsf{X}}},\ {\bf{\mathsf{O}}}\_{n} = {\rm PE1D}\_{n}({\bf{\mathsf{O}}}\_{n-1}),\ n\in\{1, 2, 3, 4\}, \end{aligned}$ where ${\rm FC}\_1:{\mathbb R}^{D}\to{\mathbb R}^{D}$ and ${\rm FC}\_2:{\mathbb R}^{(4+15)D}\to{\mathbb R}^{D}$ represent the linear layer applied at the last dimension. ${\rm PE1D}\_1:{\mathbb R}^{M\times D}\to {\mathbb R}^{M\times D}$, ${\rm PE1D}\_2:{\mathbb R}^{K\times D}\to {\mathbb R}^{K\times D}$, ${\rm PE1D}\_3:{\mathbb R}^{N\_{T}\times D}\to {\mathbb R}^{N\_{T}\times D}$, and ${\rm PE1D}\_4:{\mathbb R}^{N\_{R}\times D}\to {\mathbb R}^{N\_{R}\times D}$ denotes the 1D equivariant layer operating at the 1,2,3, and 4-th dimensions, respectively.

Then, the four-dimensional equivariant neural network used to model $f$ can be constructed as follows $\begin{aligned} {\widetilde {{\bf{\mathsf{W}}}}} = \sigma\_{\rm out}({\rm FC}\_{\rm out}({\rm SMPEN}({\rm FC}\_{\rm in}({\widetilde {{\bf{\mathsf{H}}}}})))), \end{aligned}$ where ${\rm FC}\_{\rm in}:{\mathbb R}^{2}\to {\mathbb R}^{D}$ and ${\rm FC}\_{\rm out}:{\mathbb R}^{D}\to {\mathbb R}^{2}$ are used to is used to change the feature dimensions. $\sigma\_{\rm out}(\cdot)$ is designed to ensure that ${\widetilde {{\bf{\mathsf{W}}}}}$ satisfies the power constraints. At this point, the construction of the four-dimensional permutation-equivariant network from ${\widetilde {{\bf{\mathsf{H}}}}}$ to ${\widetilde {{\bf{\mathsf{W}}}}}$ is completed. Following this approach, the framework can be easily extended to higher dimensions.

This is an interesting topic that is currently being considered. Although we are unable to provide rigorous mathematical conclusions at this moment, we can analyze the following points from an experimental perspective and offer practical recommendations.

1. The impact of the number of objects $K\_1,...,K\_N$ on performance:
We consider the performance of a trained network from two perspectives: the testing performance in the training configuration and in other configurations. During testing in the training configuration, the number of of objects has a negligible impact on the network performance. However, when tested in other configurations, if the number of a certain type of objects in the training configuration is small, the change of this number will have a significant impact on the network performance. For example, in Table 3, the performance of SMPEN3D-TF trained under configuration ${\mathcal{B}}$ shows poor performance when tested under configuration ${\mathcal{A}}$. Conversely, if the number of such objects is large in the training configuration, minor variations in number have a negligible impact on the performance. For instance, in Table 2, SMPEN2D-TF trained under configuration ${\mathcal{B}}$ exhibits excellent performance when tested under configuration ${\mathcal{C}}$. Intuitively, when the number difference is small, the modeled equivariant mappings are similar, while the distinctions in the modeled mappings become more significant as the number difference grows large.

2. The impact of the relationship between the number of object types $N$ and the network dimensionality on performance:
When the dimension of the equivariant network is equal to the number of object types in the input, the input matches the network. With adequate training, the network can achieve good performance under the training configuration. When the dimension of the equivariant network is smaller than the number of object types in the input, multiple types of objects are reshaped into a single dimension for input. Due to equivariance, the network loses its ability to distinguish the merged types, resulting in a decrease in performance. For example, in Table 3, SMPEN3D-TF trained under configuration ${\mathcal{A}}$ performs better than SMPEN2D-TF when tested under configuration ${\mathcal{A}}$. However, merging multiple types of objects into a single type will increase the number of the new type of objects. According to the analysis of the first point, this operation may lead to better performance of the network when tested under other configurations. For instance, in Table 3, SMPEN2D-TF trained under configuration ${\mathcal{B}}$ performs better than SMPEN3D-TF when tested under configuration ${\mathcal{A}}$.

In conclusion, we provide the following recommendations for network selection.

• If the number of each type of objectis is relatively large, choose an equivariant network with the same dimension as the number of types. For example, for the task in Table 2, we choose SMPEN2D-TF instead of SetTrans1D.

• If the number of each type of objectis is relatively small, to achieve better testing performance under the training configuration, it is still advisable to choose an equivariant network with the same dimension as $N$. To achieve better testing performance under other configurations, selecting a network with lower equivariant dimension can be considered. For instance, for the task in Table 3, SMPEN3D-TF performs better under the training configuration, while SMPEN2D-TF performs better under other configurations.

Finally, many thanks for your careful reviews and constructive comments.

The weighted sum of the processed features and the global information is performed through equations (5) in the paper. The processing of the linear layer ${\rm FC}\_2$ can be rewritten as $\begin{aligned} {\rm FC}\_2([{\bf{\mathsf{O}}}\_1, ..., {\bf{\mathsf{O}}}\_N, {\bar{\bf{\mathsf{X}}}}\_{{\mathbb P}\_1}, ..., {\bar{\bf{\mathsf{X}}}}\_{{\mathbb P}\_{2^N-1}}])=\sum\_{n=1}^{N}{\bf{\mathsf{O}}}\_n\times\_{(N+1)}{\bf W}^{\rm PE}\_n + \sum\_{i=1}^{2^N-1}{\bar{\bf{\mathsf{X}}}}\_{{\mathbb P}\_i}\times\_{(N+1)}{\bf W}^{\rm GI}\_i +\_{(N+1)} {\bf b}, \end{aligned}$

Where ${\bf W}^{\rm PE}\_n\in{\mathbb R}^{D\times D}$ and ${\bf W}^{\rm GI}\_i\in{\mathbb R}^{D\times D}$ are the matrix weights representing $w^{\rm PE}\_n$ and $w^{\rm GI}\_{{\mathbb P}\_i}$, and ${\bf b}\in{\mathbb R}^{D}$ is the bias of this layer. The operation '$\times\_{(N+1)}$' denotes that the multiplication acts on the $(N+1)$-th dimension of the tensor (i.e., the last dimension) $1$ . Similarly, the bias is added along the last dimension, which does not affect the equivariance. By training this linear layer ${\rm FC}\_2$, the weights $w^{\rm PE}\_n$ and $w^{\rm GI}\_{{\mathbb P}\_i}$ can be learned.

$1$ Kolda, Tamara G., and Brett W. Bader. Tensor decompositions and applications. SIAM review 51(3), 455-500, 2009.

${\bar{\bf{\mathsf{X}}}}\_{{\mathbb P}}$ is obtained from pooling functions, which means it is aggregated from the information of each object. Its invariance to objects implies that it contains the global information of all object features. In each layer of the network, ${\bar{\bf{\mathsf{X}}}}\_{{\mathbb P}}$ is connected to the output after being weighted by learnable parameters $w^{\rm GI}$. This means that after passing through this layer, the global information of the input is still preserved in the output of this layer. Similarly, by stacking layers, the global information of the input is preserved in the network.

Due to the limitations of the main text, we have added Appendix E titled "Experimental Setup" in the revised version, which includes the specific mathematical expressions for the evaluation metric of each experiment and the selection of used pooling operations. For example, the evaluation metric ${\bar R}\_{\rm M}$ in Table 2 is defined as the average of $R\_{\rm M}$ across seven SNR values $\{0, 5, 10, 15, 20, 25, 30\}$, where the expression of $R\_{\rm M}$ is given by $\begin{aligned} R\_{\rm M}=\sum\_{k=1}^{K}\log\_2\left(1+\frac{|{\bf h}\_k^H{{\bf w}}\_k|^2}{\sigma^2+\sum\_{i\neq k}|{\bf h}^H\_k{\bf w}\_i|}\right), \end{aligned}$ where the definitions of symbols can be found in Section D.2.

Reference $2$ demonstrates that a general invariant function can be decomposed into $f({\bf x})=\rho\left({\rm sum}(\phi(x\_1),...,\phi(x\_K))\right)$, where $\phi$ and $\rho$ are the suitable transformations. More generally, reference $1$ reformulates this as $f({\bf x})=\rho\left({\rm pool}(\phi(x\_1),...,\phi(x\_K))\right)$. Furthermore, it was mentioned in $1$ that if the transformation $\phi$ appied at ${\bf x}$ is a stack of permutation equivariant layer, the model remains permutation invariant, and it is easy to prove that $f({\bf x})={\rm pool}({\rm PEL}({\bf x}))$ is a permutation invariant function, where ${\rm PEL}(\cdot)$ is the equivariant layer. Based on this, it can be demonstrated that a multi-dimensional equivariant function can be transformed into a multi-dimensional invariant function by pooling the output across multiple dimensions. Thus, we transform SMPEN into SMPIN in Section 3.3, expressed as ${\rm SMPIN}({\bf{\mathsf{X}}}) = {\rm Pool}\_{{\mathbb N}}\left({\rm SMPEN}({\bf{\mathsf{X}}})\right)$, where ${\rm Pool}\_{{\mathbb N}}$ represents the pooling function operating on dimensions in ${\mathbb N}$.

$1$ Lee J, Lee Y, Kim J, et al. Set transformer: A framework for attention-based permutation-invariant neural networks $C$ . International conference on machine learning. PMLR, 2019: 3744-3753.

$2$ Zaheer M, Kottur S, Ravanbakhsh S, et al. Deep sets $J$ . Advances in neural information processing systems, 2017, 30.

We sincerely appreciate your thorough reviews and insightful comments.

In this work, pooling functions are used in two instances: Firstly, in Section 3.2, a pooling function is utilized to extract global information from the input of each layer. Secondly, in Section 3.3, pooling functions are introduced across multiple dimensions to transform multi-dimensional equivariant functions into multi-dimensional invariant functions. Without loss of generality, both of the two types are set to be the mean function in the experiments, since the design of the pooling function is not the focus of this work.

We have supplemented this information in Section 4 and Appendex E.1.