SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. Let  $\sigma$  be as in the statement of the proposition, and let  $\mu$  be a signed Borel measure that is finite and is supported on a compact set  $K \subseteq \mathbb{R}^d$ . Furthermore, assume that  $\int_{\mathbb{R}^d} \sigma(\boldsymbol{a} \cdot \boldsymbol{x} + b) d\mu(\boldsymbol{x}) = 0$  for all  $\boldsymbol{a}$  and  $b$ . We need to prove that  $\mu = 0$ .
Note that by the linearity of  $\mu$ , we have that  $\int f d\mu = 0$  for every  $f$  in the space spanned by functions of the form  $\sigma(\boldsymbol{a} \cdot \boldsymbol{x} + b)$ . Additionally, since functions in this space can approximate any continuous function on  $K$  uniformly ([32]), Propositions 3.3 and 3.8), and  $\mu$  is compactly supported, we have that  $\int f d\mu = 0$  for every continuous function. Finally, by the Riesz representation theorem ([36], Theorem 6.19) we know that a signed (and more generally complex) measure on  $K$  is defined uniquely by the integrals of continuous functions and thus  $\mu = 0$  as required.
Proof. The proof follows the proof of Theorem 3.3 with some modifications. Set  $m = 2n(d + 1) + 1$ . A measure in  $\mathcal{M}_{\leq n}(\mathbb{R}^d)$  is determined (albeit not uniquely) by a matrix  $\pmb{X} = (\pmb{x}_1, \dots, \pmb{x}_n) \in \mathbb{R}^{d \times n}$  that represents  $n$  points in  $\mathbb{R}^d$ , and a weight vector  $\pmb{w} \in \mathbb{R}^n$ . Let  $\mathbb{M}$  denote the space of pairs of measure parameters
$$
\mathbb {M} = \left\{\left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}\right) \in \mathbb {R} ^ {n} \times \mathbb {R} ^ {n} \times \mathbb {R} ^ {d \times n} \times \mathbb {R} ^ {d \times n} \right\},
$$
$$
F \left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}; \boldsymbol {y}, \sigma\right) = \sum_ {i = 1} ^ {n} w _ {i} \exp \left(- \sigma^ {- 2} \| \boldsymbol {x} _ {i} - \boldsymbol {y} \| ^ {2}\right) - \sum_ {i = 1} ^ {n} w _ {i} ^ {\prime} \exp \left(- \sigma^ {- 2} \| \boldsymbol {x} _ {i} ^ {\prime} - \boldsymbol {y} \| ^ {2}\right). \tag {29}
$$
We prove the proposition by showing that for Lebesgue almost every  $\pmb{y}_1, \dots, \pmb{y}_m$  and  $\sigma_1, \dots, \sigma_m$ ,
$$
\begin{array}{l} \left\{\left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}\right) \in \mathbb {M} \mid \sum_ {i = 1} ^ {n} w _ {i} \delta_ {\boldsymbol {x} _ {i}} = \sum_ {i = 1} ^ {n} w _ {i} ^ {\prime} \delta_ {\boldsymbol {x} _ {i} ^ {\prime}} \right\} \\ \stackrel {(*)} {=} \left\{\left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}\right) \in \mathbb {M} \mid F \left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}; \boldsymbol {y}, \sigma\right) = 0, \forall \boldsymbol {y} \in \mathbb {R} ^ {d}, \sigma \in \mathbb {R} _ {+} \right\} \\ \stackrel {(*)} {=} \left\{\left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}\right) \in \mathbb {M} \mid F \left(\boldsymbol {w}, \boldsymbol {w} ^ {\prime}, \boldsymbol {X}, \boldsymbol {X} ^ {\prime}; \boldsymbol {y} _ {i}, \sigma_ {i}\right) = 0, \forall i = 1, \dots , m \right\} \\ \end{array}
$$
As in the proof of Theorem 3.3, equality  $(^{**})$  follows from the finite witness theorem. The equality  $(^*)$  follows on the one hand from the fact that whenever the measures defined by  $(\pmb {w},\pmb {X})$  and  $(\pmb{w}^{\prime},\pmb{X}^{\prime})$  are the same, necessarily all integrals of functions against these measures are the same, and so  $F(\pmb {w},\pmb{w}^{\prime},\pmb {X},\pmb{X}^{\prime};\pmb {y},\sigma) = 0$  for all  $\pmb {y}\in \mathbb{R}^d,\sigma \in \mathbb{R}_+$ .
On the other hand, if the measure  $\mu$  defined by  $(\boldsymbol{w}, \boldsymbol{X})$  and the measure  $\mu'$  defined by  $(\boldsymbol{w}', \boldsymbol{X}')$ , are not the same, then it is sufficient to show that for some choice of  $\boldsymbol{y} \in \mathbb{R}^d$ ,  $\sigma \in \mathbb{R}_+$  we have  $F(\boldsymbol{w}, \boldsymbol{w}', \boldsymbol{X}, \boldsymbol{X}'; \boldsymbol{y}, \sigma) \neq 0$ . To see this is indeed the case, choose some  $\boldsymbol{y}_0 \in \mathbb{R}^d$  which the non-zero matrix  $\mu - \mu'$  assigns a non-zero weight  $w_0$ . Then
$$
\lim_{\sigma \to 0,\sigma >0}F(\boldsymbol {w},\boldsymbol{w}^{\prime},\boldsymbol {X},\boldsymbol{X}^{\prime};\boldsymbol{y}_{0},\sigma) = w_{0}\neq 0.
$$
and so for small enough  $\sigma$  we have that  $F(\boldsymbol{w}, \boldsymbol{w}', \boldsymbol{X}, \boldsymbol{X}'; \boldsymbol{y}_0, \sigma)$  which concludes the proof.
D.2 Proofs for Section 5
Proposition 5.1. Let  $n \geq 2$ ,  $d, m \in \mathbb{N}$ , and let  $f: \mathbb{R}^d \to \mathbb{R}^m$  be differentiable at some  $\pmb{x}_0 \in \mathbb{R}^d$ . Then the induced moment function  $\hat{f}: S_{\leq n}(\mathbb{R}^d) \to \mathbb{R}^m$  defined in (1) is not bi-Lipschitz.
Proof. We choose some arbitrary  $\pmb{d} \in \mathbb{R}^d$  with unit norm, and focus on multiset of two elements in  $\mathbb{R}^d$  of the form  $S_{\epsilon} = \{\{\pmb{x}_0 + \epsilon \pmb{d}, \pmb{x}_0 - \epsilon \pmb{d}\} \}$ . We note that the Wasserstein distance  $W_2(S_\epsilon, S_0)$  is  $\sqrt{2}\epsilon$ . Thus, it is sufficient to show that
$$
\lim  _ {\epsilon \rightarrow 0} \frac {\left\| \hat {f} \left(S _ {\epsilon}\right) - \hat {f} \left(S _ {0}\right)\right\|}{\left| \epsilon \right|} = 0, \tag {30}
$$
for this implies that there is no positive  $c$  for which (8) holds. Indeed, denoting the differential of  $f$  at  $x_0$  by  $J$ , we have
$$
\begin{array}{l} \frac {\| \hat {f} (S _ {\epsilon}) - \hat {f} (S _ {0}) \|}{| \epsilon |} = \frac {\| f (\boldsymbol {x} _ {0} + \epsilon \boldsymbol {d}) + f (\boldsymbol {x} _ {0} - \epsilon \boldsymbol {d}) - 2 f (\boldsymbol {x} _ {0}) \|}{| \epsilon |} \\ = \frac {\| f (\boldsymbol {x} _ {0} + \epsilon \boldsymbol {d}) - f (\boldsymbol {x} _ {0}) - \epsilon J \boldsymbol {d} + f (\boldsymbol {x} _ {0} - \epsilon \boldsymbol {d}) - f (\boldsymbol {x} _ {0}) + \epsilon J \boldsymbol {d} \|}{| \epsilon |} \\ \leq \frac {\| f (\boldsymbol {x} _ {0} + \epsilon \boldsymbol {d}) - f (\boldsymbol {x} _ {0}) - \epsilon J \boldsymbol {d} \|}{| \epsilon |} + \frac {\| f (\boldsymbol {x} _ {0} - \epsilon \boldsymbol {d}) - f (\boldsymbol {x} _ {0}) - (- \epsilon J \boldsymbol {d}) \|}{| \epsilon |} \stackrel {{\epsilon \rightarrow 0}} {{\rightarrow}} 0, \\ \end{array}
$$
where the convergence to zero of the last expression follows from the definition of differentiability at a point.
D.3 Proofs for Section 6
We now prove Corollary 6.1:
Proof. By Proposition 1.3 in [9], it is sufficient to show that there exists  $\mathbf{A} \in \mathbb{R}^{m \times d}$ ,  $\mathbf{b} \in \mathbb{R}^d$  such that the permutation invariant function
$$
\left(\boldsymbol {x} _ {1}, \dots , \boldsymbol {x} _ {n}\right) \mapsto \sum_ {j = 1} ^ {n} \sigma \left(\boldsymbol {A} \boldsymbol {x} _ {j} + \boldsymbol {b}\right) \tag {31}
$$
is orbit separating. This means that any two elements in  $K^n$  that are not related by a permutation will be separated by the function in (31). Any pair of elements not related by a permutation correspond to two distinct multisets in  $S_{\leq n}(\mathbb{R}^d)$  with exactly  $n$  points. By Theorem 3.3, for almost every choice of  $A, b$  the function in (31) will be injective on  $S_{\leq n}(\mathbb{R}^d)$ , and thus for such choice this function is indeed invariant and orbit separating.
We now prove Corollary 6.2:
Proof. Our first step is to show that  $\mathcal{P}_{\leq n}(K)$  is compact with respect to the Wasserstein metric. Since  $\mathcal{P}_{\leq n}(K)$  is the image of the compact set
$$
Q := \left\{\left(\boldsymbol {w}, \boldsymbol {X}\right) \in \mathbb {R} ^ {n} \times K ^ {n} \mid \sum_ {i = 1} ^ {n} w _ {i} = 1 \text {a n d} w _ {j} \geq 0, , j = 1, \dots , n \right\}
$$
under the function
$$
\left(\boldsymbol {w}, \boldsymbol {X}\right) \mapsto \sum_ {i = 1} ^ {n} w _ {i} \delta_ {\boldsymbol {x} _ {i}},
$$
it is sufficient to show that this function is continuous, as the image of a compact set under a continuous map is compact. Thus, given a sequence of  $(\boldsymbol{w}^{(k)},\boldsymbol{X}^{(k)})$  which converges to some  $(\boldsymbol{w},\boldsymbol{X})$ , we need to show that in Wasserstein space, the measures  $\mu_{k}\coloneqq \sum_{i = 1}^{n}w_{i}^{(k)}\delta_{\boldsymbol{x}_{i}^{(k)}}$  converge to the measure  $\mu \coloneqq \sum_{i = 1}^{n}w_{i}\delta_{\boldsymbol{x}_{i}}$ . Since the Wasserstein distance metrizes the weak topology on measures ([40], Theorem 6.9), it is sufficient to see that the integral of every continuous  $s:K\to \mathbb{R}$  against the sequence of measures converge to the limit measure. Indeed:
$$
\int_ {K} s (\boldsymbol {x}) d \mu^ {(k)} (\boldsymbol {x}) = \sum_ {i = 1} ^ {n} w _ {i} ^ {(k)} s (\boldsymbol {x} _ {i} ^ {(k)}) \rightarrow \sum_ {i = 1} ^ {n} w _ {i} s (\boldsymbol {x} _ {i}) = \int_ {K} s (\boldsymbol {x}) d \mu (\boldsymbol {x})
$$
Thus we have shown that  $\mathcal{P}_{\leq n}(K)$  is compact.
Next, by Theorem 3.3, for almost every choice of  $A, b$  the function
$$
q (\mu) := \int_ {\boldsymbol {x} \in K} \sigma (\boldsymbol {A x} + \boldsymbol {b}) d \mu (\boldsymbol {x})
$$
is injective on  $\mathcal{M}_{\leq n}(\mathbb{R}^d)$ , and so in particular on the compact subset  $\mathcal{P}_{\leq n}(K)$ . As  $q$  is also continuous, and a continuous injective function defined on a compact set is a homeomorphism, it follows that  $q^{-1}: q(\mathcal{P}_{\leq n}(K)) \to \mathcal{P}_{\leq n}(K)$  is continuous. We then have that  $f = (f \circ q^{-1}) \circ q$ . We can then use Tietze's extension theorem to extend  $f \circ q^{-1}$  from its compact domain to a continuous function  $F$  defined on all of  $\mathbb{R}^m$ , and we then obtain  $f = F \circ q$  as required.
Proof. The message passing iterations discussed in the theorem define new node features  $\pmb{h}_v^{(t)}$  from the previous node features  $\pmb{h}_v^{(t-1)}$  via
$$
\boldsymbol {h} _ {v} ^ {(t)} = \sum_ {u \in \mathcal {N} (v)} \sigma \left(\boldsymbol {A} ^ {(t)} \left(\eta^ {(t)} \boldsymbol {h} _ {v} ^ {(t - 1)} + \boldsymbol {h} _ {u} ^ {(t - 1)}\right) + \boldsymbol {b} ^ {(t)}\right).
$$
This can be rewritten as
$$
\boldsymbol {m} _ {v, u} ^ {(t)} = \eta^ {(t)} \boldsymbol {h} _ {v} ^ {(t - 1)} + \boldsymbol {h} _ {u} ^ {(t - 1)}
$$
$$
\boldsymbol {h} _ {v} ^ {(t)} = f ^ {(t)} \left(\{\{\boldsymbol {m} _ {v, u} ^ {(t)}, | u \in \mathcal {N} (v) \} \}\right) = \sum_ {u \in \mathcal {N} (v)} \sigma \left(\boldsymbol {A} ^ {(t)} \boldsymbol {m} _ {v, u} ^ {(t)} + \boldsymbol {b} ^ {(t)}\right), t = 1 \dots , T
$$
The final 'readout' step creates a global graph features from the node features of the last iteration  $T$  via
$$
\boldsymbol {h} _ {G} = f ^ {(T + 1)} \left(\{\{\boldsymbol {h} _ {v} ^ {(T)} \mid v \in V \}\}\right) = \sum_ {v \in V} \sigma \left(\boldsymbol {A} ^ {(T + 1)} \boldsymbol {h} _ {v} ^ {(T)} + \boldsymbol {b} ^ {(T + 1)}\right)
$$
It is known [44] that every pair of graphs  $G, \hat{G} \in \mathcal{G}_{\leq n}(\Sigma)$  which are not be separated by  $T$  iterations of the WL test, will not be separated by our message passing procedure, regardless of the choice of the parameters
$$
\boldsymbol {\theta} = \left(\boldsymbol {A} ^ {(1)}, \dots , \boldsymbol {A} ^ {(T + 1)}, \boldsymbol {b} ^ {(1)}, \dots , \boldsymbol {b} ^ {(T + 1)}, \eta^ {(1)} \dots , \eta^ {(T)}\right).
$$
We need to prove the opposite direction.
Since the collection of all graph-pairs from  $\mathcal{G}_{\leq n}(\Sigma)$  is countable, it is sufficient to show that any fixed pair of graphs  $G, \hat{G} \in \mathcal{G}_{\leq n}(\Sigma)$  that can be separated by  $T$  iterations of the WL test, can be separated by almost every choice of parameters.
Let us fix such a pair  $G, \hat{G} \in \mathcal{G}_{\leq n}(\Sigma)$  which can be separated by  $T$  iterations of WL. Note that the final features  $h_G = h_G(\pmb{\theta})$  and  $h_{\hat{G}} = h_{\hat{G}}(\pmb{\theta})$  obtained from the message passing procedures are an analytic function of the parameters  $\pmb{\theta}$ , and therefore to show separation for almost every  $\pmb{\theta}$  it is sufficient to show existence of  $\pmb{\theta}$  for which  $h_G(\pmb{\theta}) \neq h_{\hat{G}}(\pmb{\theta})$ .
To show that a single such  $\pmb{\theta}$  exists, we choose the parameters of the functions recursively in the order  $(\eta^{(1)},A^{(1)},b^{(1)}),(\eta^{(2)},A^{(2)},b^{(2)}),\ldots$  in which they are applied in the message passing procedure. When choosing the parameters  $(\eta^{(t)},A^{(t)},b^{(t)})$  corresponding to the  $t$ th step, it is sufficient to show that, if at the previous  $(t - 1)$  timestamp, for two different nodes  $v,w$  we had
$$
\boldsymbol {h} _ {v} ^ {(t - 1)} \neq \boldsymbol {h} _ {w} ^ {(t - 1)} \text {o r} \{\{\boldsymbol {h} _ {u} ^ {(t - 1)}, u \in \mathcal {N} (v) \} \} \neq \{\{\boldsymbol {h} _ {u} ^ {(t - 1)}, u \in \mathcal {N} (w) \} \} \tag {32}
$$
then  $\pmb{h}_v^{(t)}$  and  $\pmb{h}_w^{(t)}$  will not be the same. Here  $v$  and  $w$  are nodes in either  $G$  or  $\hat{G}$ .
Our first goal is to choose  $\eta^{(t)}$  so that, for all given nodes  $v,w$  in  $G$  or  $\hat{G}$  such that (32) is satisfied, we have
$$
\{\left\{\boldsymbol {m} _ {v, u} ^ {(t)}, \mid u \in \mathcal {N} (v) \right\} \} \neq \{\left\{\boldsymbol {m} _ {w, x} ^ {(t)}, \mid x \in \mathcal {N} (w) \right\}. \tag {33}
$$
To choose  $\eta^{(t)}$  in this way, we note that since all previous parameters were already determined, and since we are only interested in a single pair of graphs, there is only a finite number of features
$$
\Sigma^ {(t)} = \left\{\boldsymbol {h} _ {v} ^ {(t - 1)} | v \text {i s a n o d e i n} G \text {o r} \hat {G} \right\}
$$
which we are interested in. Since for fixed vectors  $\pmb{x}_1, \pmb{x}_2, \pmb{y}_1, \pmb{y}_2$  of the same dimension with  $\pmb{x}_1 \neq \pmb{x}_2$ , there can be at most a single  $\eta$  satisfying the equation
$$
\eta \boldsymbol {x} _ {1} + \boldsymbol {y} _ {1} = \eta \boldsymbol {x} _ {2} + \boldsymbol {y} _ {2}. \tag {34}
$$
we conclude that all but a finite number of  $\eta$  satisfy
$$
\begin{array}{l} \eta \boldsymbol {h} _ {v} ^ {(t - 1)} + \boldsymbol {h} _ {u} ^ {(t - 1)} \neq \eta \boldsymbol {h} _ {w} ^ {(t - 1)} + \boldsymbol {h} _ {x} ^ {(t - 1)}, \\ \forall \boldsymbol {h} _ {v} ^ {(t - 1)}, \boldsymbol {h} _ {u} ^ {(t - 1)}, \boldsymbol {h} _ {w} ^ {(t - 1)}, \boldsymbol {h} _ {x} ^ {(t - 1)} \in \Sigma^ {(t - 1)} \text {s u c h t h a t} \boldsymbol {h} _ {v} ^ {(t - 1)} \neq \boldsymbol {h} _ {w} ^ {(t - 1)} \\ \end{array}
$$
We choose  $\eta^{(t)}$  which satisfies the inequality above. This implies that if the left-hand side in (32) is indeed an inequality  $h_v^{(t-1)} \neq h_w^{(t-1)}$ , then  $m_{v,u}^{(t)} \neq m_{w,x}^{(t)}$  for all  $(u,x) \in \mathcal{N}(v) \times \mathcal{N}(w)$ , which in turn implies the inequality of multisets obtained from  $w$  and  $v$  ((33). On the other hand, if the left-hand side in (32) is an equality  $h_v^{(t-1)} = h_w^{(t-1)}$  and the multisets in the right-hand side of (32) are distinct, then clearly the multisets in (33) are distinct too. Thus, we have proved that we can choose  $\eta^{(t)}$  so that (32) implies (33).
It remains to show that we can choose the parameters of  $f^{(t)}$  so that, if (33) holds for some nodes  $v, w$  of  $G$  or  $\hat{G}$ , then the features  $\pmb{h}_v^{(t)}$  and  $\pmb{h}_w^{(t)}$  obtained from applying  $f^{(t)}$  to the multisets in (33) will be distinct. Indeed, since we are only requiring  $f^{(t)}$  to be injective on a finite collection of finite sets, we have that all these multisets are contained in  $S_{\leq n}(\Sigma)$  for some finite  $\Sigma$ , and therefore by Theorem 3.3 there is a choice of parameters  $(\pmb{A}^{(t)}, \pmb{b}^{(t)})$  which is injective on  $S_{\leq n}(\Sigma)$ , even when  $m = 1$ . Thus, we obtained a recursive procedure for choosing the parameters  $\pmb{\theta}$  such that  $G$  and  $\hat{G}$  will be separated, which concludes our proof.
E Numerical Experiments: Additional Details
E.1 Empirical injectivity and Bi-Lipschitzness
To generate the results shown in Figure 2 (and Figure 4 below), we ran multiple independent test instances in which we generated two random matrices  $X_{1}, X_{2} \in \mathbb{R}^{d \times n}$ , representing two sets of  $n$  vectors in  $\mathbb{R}^d$ . With exact details appearing below,  $X_{1}, X_{2}$  were generated such that: (1) Each entry of  $X_{1}$  and  $X_{2}$  has expectation zero and a standard deviation (STD) of 1; (2)  $X_{1}$  and  $X_{2}$  differ in exactly  $n_{\Delta}$  randomly chosen columns, with the parameter  $n_{\Delta}$  chosen uniformly at random from  $\{1, \ldots, n\}$ ; (3) each entry of the  $n_{\Delta}$  nonzero columns of  $\Delta X = X_{2} - X_{1}$  has expectation zero and  $\mathrm{STD} = \rho$ , with the parameter  $\rho$  itself drawn uniformly (once per test instance) from  $[\rho_{\min}, \rho_{\max}]$ . We used  $\rho_{\min} = 0.02$ ,  $\rho_{\max} = 1$ . In other words, the relative difference  $\rho$  between non-identical columns of  $X_{1}$  and  $X_{2}$  is chosen at each instance randomly between  $2\%$  and  $100\%$ . The motivation for this construction was to test various types and magnitudes of differences between multisets.
We then randomly generated  $\bar{A} \in \mathbb{R}^{\bar{m} \times d}$  and  $\bar{b} \in \mathbb{R}^{\bar{m} \times 1}$ , from which we took subblocks to be used as parameters for the function  $f(x; A, b)$  of Equation (3): for various values of  $m \in [\bar{m}]$ , we took  $A_{m}, b_{m}$  to be the top  $m$  rows of  $\bar{A}$  and  $\bar{b}$  respectively, and used them to construct the embedding  $\hat{f}(X; A_{m}, b_{m}, \sigma): \mathbb{R}^{d \times n} \to \mathbb{R}$ :
$$
\hat {f} \left(X; \boldsymbol {A} _ {m}, \boldsymbol {b} _ {m}, \sigma\right) = \sum_ {i = 1} ^ {n} \sigma \left(\boldsymbol {A} _ {m} x _ {i} + \boldsymbol {b} _ {m}\right),
$$
with  $x_{i}$  denoting the columns of  $X$ . The entries of  $\bar{A}$  and  $\bar{b}$  were drawn from Gaussian distributions chosen such that for each row  $a_{k}$  of  $\bar{A}$  and corresponding entry  $b_{k}$  of  $b$ , and each column  $x_{i}$  of  $X_{1}$  or  $X_{2}$ , the input to the activation  $\sigma$ ,  $a_{k} \cdot x_{i} + b_{k}$ , has expectation zero and  $\mathrm{STD} = 1$ ; specifically,
$$
\mathbb {E} \left[ a _ {k} \cdot x _ {i} \right] = \mathbb {E} \left[ b _ {k} \right] = 0 \quad \text {a n d} \quad \mathrm {S T D} \left[ a _ {k} \cdot x _ {i} \right] = \mathrm {S T D} \left[ b _ {k} \right] = \frac {1}{\sqrt {2}}.
$$
For various activation functions  $\sigma$ , we calculated the ratio
$$
r \left(X _ {1}, X _ {2}\right) = \frac {\| \hat {f} \left(X _ {1} ; \boldsymbol {A} _ {m} , \boldsymbol {b} _ {m} , \sigma\right) - \hat {f} \left(X _ {2} ; \boldsymbol {A} _ {m} , \boldsymbol {b} _ {m} , \sigma\right) \| _ {2}}{W _ {2} \left(X _ {1} , X _ {2}\right)},
$$
with  $W_{2}(\cdot ,\cdot)$  being the 2-Wasserstein distance. We used a Sinkhorn approximation of  $W_{2}(\cdot ,\cdot)$ , calculated by the GeomLoss Python library [11]. Finally, for each activation  $\sigma$  and embedding dimension  $m$ , we took  $c$  and  $C$  to be the minimum and maximum of  $r(X_1,X_2)$  respectively over all test instances, and recorded the ratio  $c / C$ . In each experimental setting, we ran between 500,000 and 2 million independent instances, depending on the values of  $d$  and  $n$ . The results appear in Figure 4. It can be seen that a similar behaviour is exhibited across different values of  $d$  and  $n$ . In particular, all PwL activations have  $c / C = 0$  at low  $m$  and all analytic activations have positive  $c / C$  in all settings tested.
Probabilistic distributions of data At each instance, as mentioned above, we first chose  $n_{\Delta} \sim$  Uniform  $\left[\{1, \ldots, n\}\right]$  and  $\rho \sim$  Uniform  $[\rho_{\min}, \rho_{\max}]$ . We then randomly choose  $n_{\Delta}$  columns labelled by  $J$ , at which  $X_1, X_2$  should differ. Let  $I = [n] \setminus J$ . Denote by  $X[:, \Lambda]$  the subblock of the matrix  $X$  with columns indexed by  $\Lambda$ . The entries of  $X_1[:, I] = X_2[:, I]$  were drawn i.i.d. from Normal  $(0, 1)$ . For  $X_1[:, J]$  and  $X_2[:, J]$ , we generated two random matrices  $U, V \in \mathbb{R}^{d \times n_{\Delta}}$ , each of whose entries are i.i.d. Gaussian with expectation 0,  $\mathrm{STD} = \sqrt{1 - \frac{1}{12} \left( \rho_{\max}^2 + \rho_{\max} \rho_{\min} + \rho_{\min}^2 \right)}$  and  $\mathrm{STD} = 1$  for  $U, V$  respectively. We then set
$$
X _ {1} [:, J ] = U - \frac {1}{2} \rho V, X _ {2} [:, J ] = U + \frac {1}{2} \rho V.
$$
Lastly, we generated the entries of  $\bar{A}$  and  $\bar{b}$  from Normal  $\left(0,\frac{1}{\sqrt{2d}}\right)$  and Normal  $\left(0,\frac{1}{\sqrt{2}}\right)$  respectively. We now show that: (i) each entry of  $X_{1}(:,J]$  and  $X_{2}(:,J]$  has expectation zero and  $\mathrm{STD} = 1$ , (ii) given  $\rho$ , each entry of  $\Delta X(:,J]$  has expectation 0 and  $\mathrm{STD} = \rho$ , and (iii) All inputs to the activation  $\sigma$  have expectation zero and  $\mathrm{STD} = 1$ . Let  $x_{1},x_{2}$  be two corresponding entries of  $X_{1}(:,J],X_{2}(:,J]$ , and let  $u,v$  be their corresponding entries of  $U,V$ . Then
$$
x _ {1} = u - \frac {1}{2} \rho v, \quad x _ {2} = u + \frac {1}{2} \rho v.
$$
Figure 4: Empirical injectivit and bi-Lipschitzness
The empirical ratio  $c / C$  of Equation (8) as a function of the embedding dimension  $m$ . The results for different sizes of vector-sets  $n$  and ambient dimension  $d$  are shown. For low  $m$ , the piecewise-linear ReLU and HardTanH have  $c / C$  exactly zero. See Appendix E.1 for a full description of the experimental setting.
We now calculate the expectation and variance of  $x_{1}, x_{2}$ . Since  $U, V$  and  $\rho$  are independent, we have that
$$
\mathbb {E} \left[ x _ {1} \right] = \mathbb {E} \left[ u - \frac {1}{2} \rho v \right] = \mathbb {E} \left[ u \right] - \frac {1}{2} \mathbb {E} \left[ \rho \right] \mathbb {E} \left[ v \right] = 0 - \frac {1}{2} \frac {\rho_ {\min } + \rho_ {\max }}{2} \cdot 0 = 0,
$$
and by a similar reasoning,  $\mathbb{E}\left[x_2\right] = 0$ . The variance of  $x_{1}$  is given by:
$$
\begin{array}{l} \operatorname {V a r} \left[ x _ {1} \right] = \operatorname {V a r} \left[ u - \frac {1}{2} \rho v \right] = \operatorname {V a r} [ u ] + \frac {1}{4} \left(\operatorname {V a r} [ \rho ] + \mathbb {E} [ \rho ] ^ {2}\right) \left(\operatorname {V a r} [ v ] + \mathbb {E} [ v ] ^ {2}\right) - \mathbb {E} [ \rho ] ^ {2} \mathbb {E} [ v ] ^ {2} \\ = \operatorname {V a r} [ u ] + \frac {1}{4} \left(\frac {\left(\rho_ {\max } - \rho_ {\min }\right) ^ {2}}{1 2} + \frac {\left(\rho_ {\max } + \rho_ {\min }\right) ^ {2}}{4}\right) (1 + 0) - \mathbb {E} [ \rho ] ^ {2} \cdot 0 \\ = \operatorname {V a r} [ u ] + \frac {1}{1 2} \left(\rho_ {\max } ^ {2} + \rho_ {\max } \rho_ {\min } + \rho_ {\min } ^ {2}\right) \\ = \left(1 - \frac {1}{1 2} \left(\rho_ {\max } ^ {2} + \rho_ {\max } \rho_ {\min } + \rho_ {\min } ^ {2}\right)\right) + \frac {1}{1 2} \left(\rho_ {\max } ^ {2} + \rho_ {\max } \rho_ {\min } + \rho_ {\min } ^ {2}\right) = 1. \\ \end{array}
$$
By a similar argument,  $\operatorname{Var}[x_2] = 1$ . Thus, (i) holds. Let  $\Delta x = x_{2} - x_{1} = \rho v$ . Then
$$
\mathbb {E} \left[ \Delta x \right] = \mathbb {E} \left[ x _ {2} \right] - \mathbb {E} \left[ x _ {1} \right] = 0 - 0 = 0.
$$
Moreover, the conditional variance of  $\Delta x$  given  $\rho$  is
$$
\operatorname {V a r} [ \Delta x | \rho ] = \operatorname {V a r} [ \rho v | \rho ] = \rho^ {2} \cdot \operatorname {V a r} [ v | \rho ] = \rho^ {2} \cdot \operatorname {V a r} [ v ] = \rho^ {2},
$$
and thus (ii) holds. Finally, we show that (iii) holds. We already have established that each entry of  $X_{1}$  and of  $X_{2}$  has zero mean and  $\mathrm{STD} = 1$ . Let  $a_{k}, b_{k}$  be a row or  $\bar{A}$  and its corresponding entry of  $\bar{b}$ . Let  $x_{i}$  be an arbitrary column of  $X_{1}$  or  $X_{2}$ . Then