Quantitative Bounds for Sorting-Based Permutation-Invariant Embeddings

AI summary

Overview

Research area: Machine learning theory, specifically permutation-invariant representations for point sets and graphs; also touches on optimal transport (Wasserstein and sliced Wasserstein distances).

Technical level: Advanced. The paper is written in the language of linear algebra and group theory (singular values, full spark matrices, the symmetric group S_n, bi-Lipschitz constants) and is aimed at readers comfortable with mathematical proofs of injectivity and distortion bounds.

Scope (one sentence): The paper derives improved upper and lower bounds on how many one-dimensional projections are needed for sorting-based permutation-invariant embeddings to separate orbits, and provides the first quantitative bounds on how much these embeddings distort distances, as a function of the number of points n and the ambient dimension d.

What This Paper Is About

Many machine learning models on sets, point clouds, and graphs need a representation that does not change when the input rows are reordered. One popular recipe takes n points in ℝ^d, projects them onto D different one-dimensional directions, sorts each resulting list of n numbers, and concatenates the sorted lists. The question is how large D must be for this map to lose no information, and how badly it warps distances between inputs.

Prior work showed that for large enough D and "generic" projection matrices the map is injective and bi-Lipschitz, but left two gaps: nobody knew how small D could be, and nobody had explicit estimates of the bi-Lipschitz constants (the distortion). This paper closes much of both gaps.

Key Contributions

1. Improved injectivity bound for the sorted-projection embedding. The authors prove that a full spark projection matrix A ∈ ℝ^{d×D} yields an injective map β̄_A as soon as D ≥ n(d−1)+1, improving on the earlier quadratic-in-n requirement D ≥ rd((n−1)²+1).

2. A matching-order lower bound on the number of projections. They show that if ⌈D/(d−1)⌉ ≤ log₂(n)+1, then β̄_A fails to be injective for every choice of A, giving a logarithmic-in-n obstruction and hence an embedding-dimension lower bound of Ω(d·n log n).

3. First quantitative distortion bounds. They construct projection matrices whose bi-Lipschitz distortion grows quadratically in n and is completely independent of d, and prove that no choice of projection vectors can beat a distortion bound proportional to √n.

4. Extension to dimension-reduced variants. They show that the projected embeddings β̄_{A,L} and δ̄_{A,B} are injective at embedding dimension (2n−1)d, and that a sketching argument lets β̄_{A,L} at dimension proportional to nd (up to logarithmic terms) achieve distortion comparable to the full β̄_A.

Main Findings

• Injectivity upper bound (Theorem 4): For natural numbers n, D and d > 1, if A ∈ ℝ^{d×D} is full spark and D ≥ n(d−1)+1, then β̄_A is injective. The proof partitions columns by which difference vectors x_i − y_j they annihilate; full spark forces each such set to have size at most d−1, and the total count forces D ≤ n(d−1) unless the two inputs are permutations of each other.

• Injectivity lower bound (Theorem 5): If ⌈D/(d−1)⌉ ≤ log₂(n)+1, then for any A ∈ ℝ^{d×D} the map β̄_A is not injective. Equivalently, injectivity forces D = m(d−1) − r with integers m, r satisfying 0 ≤ r ≤ d−2 and m > log₂(n)+1. The construction partitions the D columns into k = ⌈D/(d−1)⌉ groups, picks a vector orthogonal to each group, and builds two non-permutation-equivalent point sets — one using sums over even index subsets, the other over odd index subsets — whose sorted projections coincide.

• Embedding-dimension table (Table 1): For β̄_A the best known upper bound on the embedding dimension M is n²(d−1)+n (from Theorem 4) and the best known lower bound is Ω(d·n log(n)) (from Theorem 5). For both δ̄_{A,B} and β̄_{A,L} the upper bound is (2n−1)d and the lower bound is nd, the latter attributed to prior work [JBM+23]. The paper notes that M ≥ nd is known to be necessary for any continuous, permutation-invariant injective function.

• Dimension-reduced embeddings (Theorems 9 and 10): Both δ̄_{A,B} and β̄_{A,L} are shown to be injective with embedding dimension (2n−1)d.

• Distortion lower bound (Theorem 18): For any choice of projection vectors, the distortion of β̄_A can never be better than a bound proportional to √n.

• Distortion upper bounds via construction (Theorems 14 and 15): Two probabilistic constructions of the projection matrix A — plus an explicit construction for d = 2 — make β̄_A achieve bi-Lipschitz distortion scaling like n², independent of d. These require D on the order of n²d and n⁴d respectively.

• Distortion after sketching (Theorem 20): Using a sketching argument, β̄_{A,L} with embedding dimension proportional to nd, up to logarithmic terms, can achieve bi-Lipschitz distortion similar to that of β̄_A.

• Sharp case and phase retrieval link for n = 2: For n = 2 the upper and lower bounds coincide, giving the sharp threshold D ≥ 2d−1 for orbit separation when A is full spark. Via a connection to real phase retrieval established in prior work, this recovers the classical result that 2d−1 measurements are necessary and sufficient for sign retrieval in ℝ^d provided the measurement vectors form a full spark frame.

• Numerical experiments: For small parameters d > 1 and n > 2, experiments based on a criterion from [BHS25, Proposition 3.8 on p. 14] show the new bounds are typically suboptimal: there exist D < n(d−1)+1 and A ∈ ℝ^{d×D} for which β_A still separates orbits. For n = 2 the results are optimal. The paper's phase diagram (Figure 1, drawn for d = 2) shows a widening gap between the non-separating region and the guaranteed separating region as n grows.

• Connection to sliced Wasserstein distance: β_A with columns sampled from the unit sphere is exactly a finite-dimensional Monte Carlo approximation of the sliced 2-Wasserstein distance between empirical measures. Corollary 16 states that for D ≳ dn² log(n√d + log n), the sampled sliced distance gives, with high probability, a bi-Lipschitz approximation of W₂ on empirical measures of support size n, with distortion of order Õ(n²). Theorem 18 is the corresponding converse. The paper notes that O(n²) distortion for d = 2 was previously obtained in [CCO17], that bounds for higher dimensions were substantially weaker [Wei23], and that for measures with infinite support bi-Lipschitz equivalence between Wasserstein and sliced Wasserstein is impossible [BG21], though Hölder-type bounds exist [Bon13].

• Prior results the paper builds on: Theorem 1 of [BHS25] gives that the upper Lipschitz constant of β̄_A equals σ₁(A) (the largest singular value), and that for D = n!(d−1)+1 with full spark the lower Lipschitz constant is at least min over d-column subsets of σ_d(A(I)). Theorem 2 of [RD23] gives a lower Lipschitz bound for D ≥ rd((n−1)²+1). Theorem 3 ([DG24] and [BTW24]) gives injectivity of δ̄_{A,B} for D ≥ 2nd+1 for Lebesgue almost every (A, B), implying bi-Lipschitzness whenever injective, with total embedding dimension M = Dn = 2n²d + n.

Methodology in Plain English

The authors work purely with the geometry of sorted projections. A key simplification is to view a permutation-invariant function as a function on orbits — equivalence classes of point sets related by reordering — and to measure the distance between two orbits as the smallest Frobenius-norm difference over all reorderings. Injectivity and the bi-Lipschitz property are then statements about this quotient space.

For the injectivity results, they argue by contradiction. If two non-equivalent point sets produce identical sorted projections, they look at every pair consisting of one row from each set and ask which projection directions are orthogonal to the difference of that pair. If the projection matrix has full spark (every d columns are linearly independent), each such set of directions is small. Counting how many directions are "used up" across all rows forces a relationship between D, n and d — which yields both the sufficient condition D ≥ n(d−1)+1 and the necessary condition involving log₂(n). The lower-bound construction is explicit and combinatorial: it groups columns, finds orthogonal vectors, and builds point sets indexed by even versus odd subsets so that sorting hides the difference.

For the distortion results, they construct projection matrices probabilistically and then bound the worst-case ratio between how far apart two orbits are and how far apart their embeddings land. One construction is entirely explicit for the two-dimensional case. A separate argument shows that no matrix can do better than √n, by exhibiting inputs whose distances are squeezed. For the dimension-reduced versions, they borrow a "sketching" technique — projecting the already-sorted representation down with a random linear map — and show the distortion survives.

Why This Matters

Impact on research: The paper converts a qualitative existence story ("injectivity and bi-Lipschitzness hold for large enough D and generic A") into quantitative statements about how large D must be and how much distortion any method must incur. It also supplies the first non-trivial lower bounds on both the number of projections and the achievable distortion, framing the remaining gap between upper and lower bounds as a concrete open problem. The explicit link to sliced Wasserstein distances means the distortion bounds transfer directly into statements about how well a Monte Carlo sliced distance approximates the true 2-Wasserstein distance.

Real-world applications cited or motivated in the paper:

• Learning on multisets and permutation-invariant point sets, where a model must produce the same output regardless of input ordering.
• Graph deep learning, where outputs should be invariant to permutations of graph nodes.
• Metric-based learning tasks such as nearest neighbor search and clustering, which benefit from knowing that nearby orbits map to nearby vectors and far orbits map to far vectors.
• Comparing empirical distributions through the 2-Wasserstein and sliced 2-Wasserstein distances, where the embedding acts as a fast, sorting-based surrogate for an otherwise cubic-time computation (the Hungarian method costs O(n³), while sorting costs O(n log n)).

Industry relevance: Practitioners who build set- or graph-based models must choose a representation dimension that balances expressiveness against compute and memory. This paper's dimension bounds (n²(d−1)+n versus the Ω(d·n log n) lower bound for the sorting embedding, and (2n−1)d versus nd for the projected variants) give direct guidance on where the practical sweet spot currently lies and how far the theory is from the known information-theoretic floor of nd.

Future Directions

• Closing the dimension gap for β̄_A. The upper bound n²(d−1)+n and the lower bound Ω(d·n log n) differ substantially for n > 2; the paper explicitly leaves open whether either bound is sharp in that regime.

• The d > 2 case of the logarithmic lower bound. The authors note that for d = 2 prior work shows the logarithmic lower bound is nearly attainable — injectivity holds for n ≤ 2^{cD/log D} for generic matrices with D ≥ D₀ — and that it remains unclear whether similar bounds hold when d > 2.

• Tightening distortion constants. The constructions achieve n² distortion independent of d, while Theorem 18 forces at least √n distortion for any matrix. Whether the true optimal distortion lies closer to √n or to n² is unresolved.

• Understanding why numerical experiments beat the guarantees. For n > 2 the experiments exhibit matrices with D < n(d−1)+1 that still separate orbits, indicating the sufficient condition in Theorem 4 is not necessary; characterizing exactly when fewer projections suffice is an open problem.

Target Audience

This paper is for machine learning theorists and mathematically inclined practitioners working on permutation-invariant and equivariant representations, graph neural networks, and expressive set functions. It will also interest researchers in optimal transport who study sliced Wasserstein approximations of the 2-Wasserstein distance, and applied mathematicians working on frames, spark conditions, and phase retrieval, given the sharp n = 2 threshold of 2d−1 projections. Readers without a background in linear algebra, matrix singular values, and Lipschitz analysis will find the proofs and the bound statements hard going.

Authors’ abstract

We study permutation-invariant embeddings of $d$-dimensional point sets, which are defined by sorting $D$ independent one-dimensional projections of the input. Such embeddings arise in graph deep learning where outputs should be invariant to permutations of graph nodes. Previous work showed that for large enough $D$ and projections in general position, this mapping is injective, and moreover satisfies a bi-Lipschitz condition. However, two gaps remain: firstly, the optimal size $D$ required for injectivity is not yet known, and secondly, no estimates of the bi-Lipschitz constants of the mapping are known. In this paper, we make substantial progress in addressing both of these gaps. Regarding the first gap, we improve upon the best known upper bounds for the embedding dimension $D$ necessary for injectivity, and also provide a lower bound on the minimal injectivity dimension. Regarding the second gap, we construct matrices of projection vectors, so that the bi-Lipschitz distortion of the mapping depends quadratically on the number of points $n$, and is completely independent of the dimension $d$. We also show that for any choice of projection vectors, the distortion of the mapping will never be better than a bound proportional to the square root of $n$. Finally, we show that similar guarantees can be provided even when linear projections are applied to the mapping to reduce its dimension.