Quantitative Bounds for Sorting-Based Permutation-Invariant Embeddings

1. Problem Statement

The paper addresses the challenge of constructing permutation-invariant embeddings for sets of $n$ points in $d$-dimensional space ($X \in \mathbb{R}^{n \times d}$). In many machine learning applications (e.g., graph neural networks, point cloud processing), the input is a multiset where the order of elements is irrelevant. The goal is to map these sets to a vector space such that:

1. Permutation Invariance: The output remains unchanged if the input rows are permuted.
2. Orbit Separation (Injectivity): Distinct orbits (sets that are not permutations of each other) map to distinct vectors.
3. Bi-Lipschitz Continuity: The mapping preserves the geometry of the input space. Specifically, the distance between embedded vectors should be bounded both above and below by the minimum distance between the input sets under any permutation (the Frobenius orbit distance).

The specific focus is on the sorting-based embedding $\beta_A$, defined by projecting the input $X$ onto $D$ random directions $A \in \mathbb{R}^{d \times D}$, sorting the resulting 1D projections, and concatenating them:
$$\beta_A(X) = \begin{pmatrix} \downarrow(Xa_1) \\ \vdots \\ \downarrow(Xa_D) \end{pmatrix}$$
where $\downarrow(\cdot)$ denotes sorting in non-decreasing order.

Existing Gaps: Previous work established that $\beta_A$ is injective and bi-Lipschitz for sufficiently large $D$ and generic $A$. However, two critical quantitative gaps remained:

• Optimal Dimension: The exact minimum number of projections $D$ required for injectivity was unknown (previous bounds were super-exponential or quadratic in $n$).
• Distortion Bounds: No tight estimates existed for the bi-Lipschitz distortion (the ratio of upper to lower Lipschitz constants), particularly regarding its dependence on $n$ and $d$.

2. Methodology

The authors employ a combination of algebraic geometry, linear algebra, probability theory, and combinatorial arguments to derive rigorous bounds.

• Injectivity Analysis:

  • They utilize the concept of full spark matrices (matrices where every subset of $d$ columns is linearly independent).
  • They analyze the conditions under which $\beta_A(X) = \beta_A(Y)$ implies $X$ and $Y$ are permutations of each other. This involves counting the number of "bad" projections where the sorting order fails to distinguish between rows.
  • They apply the Finite Witness Theorem from real algebraic geometry to reduce the dimensionality of the domain required to prove injectivity for linear projections ($\delta_{A,B}$ and $\bar{\beta}_{A,L}$).

• Distortion Analysis:

  • Upper Bound: They introduce the concept of projective uniformity. A matrix $A$ is $(m, \delta)$-projectively uniform if, for any direction on the sphere, at least $D-m+1$ columns have a projection magnitude greater than $\delta$. This property ensures the lower Lipschitz constant is bounded away from zero.
  • Constructions: They construct specific matrices $A$ (deterministic for $d=2$, probabilistic for general $d$) that satisfy projective uniformity.
  • Lower Bound: They construct specific adversarial examples of input sets $X$ and $Y$ to prove that the distortion cannot be arbitrarily small, establishing a fundamental limit.

• Dimension Reduction: They use Oblivious Subspace Embeddings (OSE) (sketching) to show that the high-dimensional embedding $\beta_A$ can be compressed linearly while preserving the bi-Lipschitz properties.

3. Key Contributions and Results

A. Embedding Dimension Bounds (Injectivity)

The paper significantly tightens the bounds on the number of projections $D$ required for injectivity.

1. Upper Bound for $\beta_A$:

   • Result: If $A$ is full spark and $D \ge n(d-1) + 1$, then $\beta_A$ is injective.
   • Significance: This improves upon previous bounds that were quadratic in $n$ (e.g., $D \sim n^2$) or super-exponential. The embedding dimension $M = nD$ scales as $O(n^2 d)$.

2. Lower Bound for $\beta_A$:

   • Result: If $D < (d-1)\log_2(n) + 1$, then $\beta_A$ is not injective for any choice of $A$.
   • Significance: This establishes that the embedding dimension must be at least $\Omega(d \cdot n \log n)$. There remains a gap between the lower bound ($\sim n \log n$) and the upper bound ($\sim n^2$).

3. Improved Bounds for Linear Projections ($\delta_{A,B}$ and $\bar{\beta}_{A,L}$):

   • The authors show that for the linearly projected versions, injectivity is guaranteed for generic matrices when $D \ge (2n-1)d$.
   • This improves the previous requirement of $D \ge 2nd + 1$ and brings the embedding dimension closer to the theoretical lower bound of $nd$.

B. Bi-Lipschitz Distortion Bounds

The paper provides the first quantitative estimates for the distortion of these embeddings.

1. Upper Bound on Distortion (Positive Results):

   • Result: By constructing matrices with projective uniformity (e.g., Gaussian matrices or uniform sampling on the sphere), the authors prove that the distortion of $\beta_A$ scales as $O(n^2)$ (or $\tilde{O}(n^2)$ with logarithmic factors).
   • Key Insight: Crucially, this distortion bound is independent of the input dimension $d$.
   • Implication: For the sliced Wasserstein distance interpretation, this means the Monte-Carlo approximation using $D \sim n^2 d$ directions provides a bi-Lipschitz approximation of the true Wasserstein distance with distortion $O(n^2)$.

2. Lower Bound on Distortion (Negative Results):

   • Result: The authors prove that for any choice of projection vectors $A$, the distortion of $\beta_A$ is at least $\Omega(\sqrt{n})$.
   • Significance: This establishes a fundamental limit: one cannot achieve constant distortion (independent of $n$) with this specific sorting-based architecture. There is a gap between the achievable $O(n^2)$ and the lower bound $\Omega(\sqrt{n})$.

3. Dimension Reduction via Sketching:

   • Result: Using random linear projections (sketching), the authors show that the embedding dimension can be reduced to $M \sim nd$ (up to log factors) while preserving the $O(n^2)$ distortion of the original high-dimensional map.

4. Significance and Impact

• Theoretical Foundation for Graph Learning: The results provide rigorous justification for using sorting-based embeddings in Graph Neural Networks (GNNs) and set-based learning, guaranteeing that distinct graph structures are mapped to distinct representations with controlled geometric distortion.
• Wasserstein Distance Approximation: The work bridges the gap between sorting-based embeddings and the Sliced Wasserstein distance. It quantifies exactly how many projections are needed to approximate the Wasserstein distance with a guaranteed bi-Lipschitz constant, which is vital for metric learning tasks like clustering and nearest neighbor search.
• Optimality Gaps: The paper identifies precise gaps in current knowledge:
  • The gap between the necessary $D \sim n \log n$ and sufficient $D \sim n^2$ for injectivity.
  • The gap between the achievable $O(n^2)$ distortion and the $\Omega(\sqrt{n})$ lower bound.
• Practical Guidance: The results suggest that while sorting-based methods are theoretically sound, they require a number of projections scaling quadratically with the set size $n$ to guarantee good distortion properties, offering a trade-off between computational cost and geometric fidelity.

5. Conclusion

This paper makes substantial progress in the quantitative analysis of permutation-invariant embeddings. It refines the necessary and sufficient conditions for injectivity, moving from super-exponential to polynomial bounds, and provides the first concrete distortion estimates. While it proves that distortion must grow with $n$, it demonstrates that this growth can be controlled to $O(n^2)$ independently of the data dimension $d$, making these embeddings viable for high-dimensional set learning tasks.