Polynomial Width is Sufficient for Set Representation with High-dimensional Features

We thank Reviewer oY1k for the detailed summary of our results and acknowledging our theoretical significance. Regarding your questions, please see our response below:

1. Brief discussion on $\rho$.

Fine-grained analysis on $\rho$ remains a widely open question in this area. We agree with Reviewer oY1k that $\rho$ is highly non-smooth and hard to learn [1]. Nevertheless, we hypothesize that modern ML algorithms can approximate non-smooth inversion of the sum-pooling well. This argument is also supported by our experiments in Appendix K. We doubt that non-smooth points in a generic set are often located at the regime of no interest, meaning worst-case may not occur in practical scenarios. We also agree a rigorous study demystifying the $\rho$ mapping is very necessary and the suggested direction of bounding the smoothness of the inversion function is of high interest to explore in the future. Our results can be easily extended to equivariant functions, as a foundation to form multi-layer set functions. Under this scenario, $\phi$ may be simplified.

[1] Murphy et al., Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs

2. Presentation issues.

We have revised our introduction to clarify the bound of $L$ means the upper bound of the smallest $L$, and fixed all the typos pointed out.

3. Questions on Lemma E.1.

It is a good point to mention $\gamma_i$ can be chosen as prime numbers in Lemma E.1. But we would remark that the non-divisibility is the most crucial property to let Lemma E.1 hold, which allows for any other choices of $\gamma_i$ satisfying the condition in Lemma E.1.

We sincerely appreciate Reviewer jrmV's constructive advice to improve our paper's clarity. We have carefully revised the text and adopted a more informative notation system. Please check our updated version.

1. The exposition can be improved.

We have revised our high-level description of the proof sketch in Sec. 4 to enhance clarity. We also improved our notations in the following way: $w_i^{(1)}$ -> $e_i$ for canonical basis, $w_j^{(2)} \rightarrow \alpha_j$ for anchor construction, and $w_{i,j,k}^{(3)} \rightarrow \Gamma_{i,j,k}$ for linear coupling.

2. Question regarding linear coupling.

Sorry for the confusion. $a$ and $b$ here are both column vectors, and we stack them to form an $N$ by 2 matrices. The form of linear coupling may have appeared in other proof techniques, but to our best knowledge, it is the first time used to constrain the permutation orbits of matrices to ensure alignment.

We thank Reviewer a75R for acknowledging our mathematical rigor and high relevance with ML community. For your questions, please see our response below:

1. Confusion by Lemma 4.2 and its implication.

Thanks to Reviewer a75R's suggestion, we have revised the manuscript with an emphasis on "the same anchor is being applied over all the channels". The terminology "alignment" between $X$ and $X'$ refers to $X \sim X'$. Hence, our argument "once each coupled pairs are aligned, two data matrices are globally aligned" can be rephrased in a more rigorous way: "once the permutation orbits of each coupled pair intersect, the permutation orbits of the two data matrices also intersect", and can be translated into our mathematical statement: $[x_i, a] \sim [x'_i, a'] \Rightarrow X \sim X'$.

2. Why is Lemma 4.4 necessary?

Proof in Sec. 4.1.2 proceeds in the following way:

1. Sum-pooling $\sum_n \phi(x^{(n)}) = \sum_n \phi(x'^{(n)})$ implies column-wisely $Xw_i \sim X'w_i, \forall i \in [K]$ by injectivity of power mapping, which induces $(x_i - \gamma_k Xw^{(2)}_j) \sim (x'_i - \gamma_k X' w^{(2)}_j), \forall i,j,k$.
2. By Lemma 4.4, column-wisely $(x_i - \gamma_k Xw^{(2)}_j) \sim (x'_i - \gamma_k X' w^{(2)}_j), \forall i,j,k$ implies $[x_i, Xw^{(2)}_j] \sim [x'_j, X'w^{(2)}_j]$ for any $i,j$.
3. By union alignment (Lemma 4.2), $[x_i, Xw^{(2)}_j] \sim [x'_j, X'w^{(2)}_j]$ implies $X \sim X'$, which is essentially Eq. 5 $\Rightarrow$ Eq. 6.

Lemma 4.4 plays a critical role in our proof, as the construction therein mixes anchor channels with the original feature channels. And this construction guarantees that alignment between mixed channels can induce alignment between the tuples of anchor and feature channels.

3. Question regarding "The obtained results for $D' = 1$ can also be easily extended to $D' > 1$ as otherwise f can be written as $[f_1, \cdots, f_D']^\top$ and each $f_i$ has single output feature channel".

We apologize for the confusion caused. The original definition in Def. 2.2 states a general permutation-invariant function can admit vector outputs of dimension $D'$. Note that here $D'$ denotes the number of output channels (e.g., classification logits) instead of number of output set elements. We used the aforementioned argument to clarify that our investigation starting from Sec. 2.2 will only focus on scalar outputs without loss of generality. Our permutation-equivariant results are not as trivial as extending the dimension of the outputs. It requires some additional theoretical establishment Lemma G.1 to extend the proof of Theorem 3.1.

4. Practical significance and empirical verification.

In many practical applications of DeepSets-like architectures, even though the embedding dimension is not carefully chosen, it is observed that a reasonably small latent dimension is sufficient to let those networks perform well. Such empirical observations appear to contradict existing theoretical studies, in which exponential many neurons are often needed. In this work, we provide a more efficient construction of DeepSet embedding layers, which results in polynomially many hidden neurons for the first time, successfully explaining the empirical promise of DeepSets-like architectures. Designing experiments to directly testify our theoretical bound can be inherently challenging. We provide a preliminary experiment in Appendix K, in which we show that the minimal hidden dimension required to achieve a small training error scales with sequence length and feature dimension at a polynomial rate.

We genuinely thank Reviewer XKbR for acknowledging our theoretical soundness and providing a detailed summary of our technical insights. Regarding your questions, please see our response below:

1. Practical implications and relevance to the ML community.

When the input data is a set or a graph, machine learning models often require permutation invariance or equivariance. DeepSets or sum-pooling-based architectures are ubiquitous in these models. Representative examples encompass GNNs and PointNet. GNNs learn to pass information along graph topology via a neighborhood aggregation operation at each layer of computation, essentially corresponding to a set function [1]. PointNet [2] processes point clouds by representing points as an unordered set and follows a DeepSets-like architecture. However, the analysis of expressive power for most of these usages does not take the embedding dimension into consideration, as also acknowledged by Reviewer a75R. Our work provides a rigorous justification that moderately many neurons are sufficient to represent a high-dimensional set with DeepSets. This affirms the feasibility of the DeepSets architecture given high-dimensional features and explains why DeepSet-like operations can be effectively adopted in GNNs and PointNet.

[1] Xu et al., How Powerful are Graph Neural Networks?

[2] Qi et al., Pointnet: Deep learning on point sets for 3d classification and segmentation

2. Construction of $\rho$ remains black-box

Among most of works analyzing expressiveness of DeepSets architectures [1,2], identifying bijectivity (also known as orbit separation) of the sum-pooling layer is of most theoretical interest. We note that the resultant $\rho$ in our construction is not a "complete" black box. Our construction guarantees continuity of $\rho$, a crucial property which allows it to be universally approximated by a commonly-used neural network. We leave more fine-grained discussion of $\rho$ to our future work and a potential direction can be analyzing the learnability of $\rho$ depending on smoothness of the embedding layer ($\phi$).

[1] Zaheer et al., Deep Sets [2] Wagstaff et al., On the Limitations of Representing Functions on Sets

3. More intuitive interpretation of our proof.

We appreciate the reviewer's insightful interpretation of Lemma 4.3 and we have included it in our latest revision. We have also added some high-level proof ideas upon Reviewer jrmV's suggestion and would like to direct Reviewer XKbR to Fig. 2, an illustration of our proof, which Reviewer oY1k found helpful. Although the majority of our proofs are algebraic, a simple geometric interpretation could be that we identify a non-trivial subspace embedded in a higher-dimension ambient space such that for any two points in this subspace, the intersection of their orbits, generated by column-wise permutations, signifies the orbit intersection after a global row permutation.

4. Typos and confused notation.

Thanks for pointing out typos. We have fixed all the typos and added additional remarks for confusing notations in our updated manuscript. The notion of injectivity we adopted here also follows a standard definition, as "$X \not\sim X' \Rightarrow f(X) \neq f(X')$" is equivalent to "$f(X) = f(X') \Rightarrow X \sim X'$" by contrapositive. The term "prior arts" is usually used interchangeably with "prior results", referring to previous works showcasing state-of-the-art results, as exemplified in [1]. To avoid further confusion, we change it to "prior results".

[1] Tsakiris & Peng, Homomorphic sensing

5. Relaxation of the input space.

Most discussion on permutation-invariant models focuses on the invariance defined over the entire ambient space. Meanwhile, it is common in network expressiveness analysis to consider a general input (e.g., arbitrary open/close/compact sets). We appreciate reviewers' thoughtful suggestions on considering a constrained input space. If we know a part of elements in $X$'s order is irrelevant while other elements' order matter, one can apply the idea of DeepSets to the first subset of elements and concatenate the output with the representations of the remaining elements, and finally adopt standard fully connected networks to process the concatenation. This will still give universal approximation. Our theory is still applied to the DeepSet part. Some other relaxation can also be explored in our follow-up works. For instance, $X$ can be assumed from a low-dimensional manifold and further obey permutation invariance/equivariance. However, such discussion is beyond the scope of this work. Nevertheless, our results considering the most general inputs can already cover a lot of practical scenarios such as GNNs and PointNets.

6. Solution to hierarchical set.

Extending our results to hierarchical sets is beyond the scope of this paper. A straightforward solution based on our current results could be hierarchically apply the sum-pooling $\Phi$ and append a final output layer $\rho$. For instance, consider a set of depth 2: $X \in \mathbb{R}^{M \times N \times D}$, one can provably realize a function in terms of this special invariant structure by:

where $\rho = f \circ \Phi^{-1} \circ \Phi^{-1}$. If LP embedding layer is adopted, the minimal needed latent dimension for the innermost sum-pooling is upper bounded by $O(N^5D^2)$ and for the outer sum-pooling is upper bounded by $O(M^5N^{10}D^4)$.