SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {E} \left(\frac {\langle X t , X g \rangle}{\| X g \| _ {2} \| X t \| _ {2} ^ {2}}\right) ^ {2} = \operatorname {E} \left(\frac {t ^ {T} V \Sigma^ {2} V ^ {T} g}{\| U g ^ {\prime} \| _ {2} \| X t \| _ {2} ^ {2}}\right) ^ {2} = \operatorname {E} \left(\frac {t ^ {T} V \Sigma g ^ {\prime}}{\| g ^ {\prime} \| _ {2} \| X t \| _ {2} ^ {2}}\right) ^ {2} = \frac {1}{d _ {1}} \frac {\| t ^ {T} V \Sigma \| _ {2} ^ {2}}{\| X t \| _ {2} ^ {2}} = \frac {1}{d _ {1}}.
$$
So this way we recreate the ideal bound from Theorem B.1. Note that  $\frac{\|X\|_{-2}^2}{\|X\|_F^2} \leq 1 / d_1$ . Of course it comes with the negative side of having to compute the SVD of  $X$ . But since this is just a theoretical algorithm, it's still interesting and shows how we would ideally update  $T_i$ . See Figure 6 for the effect of this change experimentally.
It's an interesting problem how it might inspire a better CCE algorithm. Somehow we'd have to get information about the the SVD of  $X$  into our sparse super-space approximations.
We now show how to extend the vector case to general matrices.
Lemma B.3. Let  $X \in \mathbb{R}^{n \times d_1}$  be a matrix with singular values  $\sigma_1 \geq \dots \geq \sigma_{d_1} \geq 0$ . Define  $\rho = \sigma_{d_1}^2 / \sum_i \sigma_i^2$ , then for any  $T \in \mathbb{R}^{n \times d_2}$ ,
$$
\operatorname {E} _ {G \sim N (0, 1) ^ {d _ {1} \times k}} \left[ \inf  _ {M \in \mathbb {R} ^ {k \times d _ {2}}} \| X (G M - T) \| _ {F} ^ {2} \right] \leq (1 - \rho) ^ {k} \| X T \| _ {F} ^ {2}.
$$
Proof. The case  $k = 1, d_2 = 1$  is proven above in Theorem B.2.
Case  $k = 1$ : We first consider the case where  $k = 1$ , but  $d_{2}$  can be any positive integer (at most  $k$ ). Let  $T = [t_{1}|t_{2}|\ldots |t_{d_{2}}]$  be the columns of  $T$  and  $M = [m_1|m_2|\ldots |m_{d_2}]$  be the columns of  $M$ . Then the  $i$ th column of  $X(GM - T)$  is  $X(Gm_{i} - t_{i})$ , and since the squared Frobenius norm of a
Figure 6: SVD aligned noise converges faster. In the discussion we mention that picking the random noise in  $H_{i}$  as  $g = V\Sigma^{-1}g^{\prime}$ , where  $g^{\prime}\sim N(0,1)^{d_{1}\times (k - d_{2})}$ , can improve the convergence rate from  $(1 - \rho)^{ik}$  to  $(1 - 1 / d)^{ik}$ , which is always better. In this graph we experimentally compare this approach (labeled "smart noise") against the IID gaussian noise (labeled "noise"), and find that the smart noise indeed converges faster – at least once we get close to zero noise. The graph is over 40 repetitions where  $X$  is a random rank-10 matrix plus some low magnitude noise.
We also investigate how much we lose in the theorem by only considering  $M$  on the form  $[I|M']$ , rather than a general  $M$  that could take advantage of last rounds  $T_{i}$ . The plots labeled "half noise" and "half smart noise" are limited in this way, while the two others are not. We observe that the effect of this is much larger in the "non-smart" case, which indicates that the optimal noise distribution we found might accidentally be tailored to our analysis.
matrix is simply the sum of the squared column 12 norms, we have
$$
\begin{array}{l} \operatorname {E} [ \| X (G M - T) \| _ {F} ^ {2} ] = \operatorname {E} \left[ \sum_ {i = 1} ^ {d _ {2}} \| X (G m _ {i} - t _ {i}) \| _ {2} ^ {2} \right] \\ = \sum_ {i = 1} ^ {d _ {2}} \operatorname {E} \left[ \| X (G m _ {i} - t _ {i}) \| _ {2} ^ {2} \right] \\ \leq \sum_ {i = 1} ^ {d _ {2}} (1 - \rho) \mathrm {E} [ \| X t _ {i} \| _ {2} ^ {2} ] \tag {10} \\ = (1 - \rho) E \left[ \sum_ {i = 1} ^ {d _ {2}} \| X t _ {i} \| _ {2} ^ {2} \right] \\ = (1 - \rho) \operatorname {E} [ \| X T \| _ {F} ^ {2} ]. \\ \end{array}
$$
where in (10) we applied the single vector case.
Case  $k > 1$ : This time, let  $g_1, g_2, \ldots, g_k$  be the columns of  $G$  and let  $m_1^T, m_2^T, \ldots, m_k^T$  be the rows of  $M$ .
We prove the lemma by induction over  $k$ . We already proved the base-case  $k = 1$ , so all we need is the induction step. We use the expansion of the matrix product  $GM$  as a sum of outer products
Figure 7: Expectation,  $\operatorname{E}\left[\frac{x}{px + (1 - p)y}\right]$  when  $x, y$  are IID with Exponential (blue) or Chi Square distribution (Orange). In both cases the expectation is  $\geq 1$  when  $p \leq 1/2$ , just as Theorem B.4 predicts.
$$
G M = \sum_ {i = 1} ^ {k} g _ {i} m _ {i} ^ {T}:
$$
$$
\begin{array}{l} \operatorname {E} [ \| X (G M - T) \| _ {F} ^ {2} ] = \operatorname {E} \left[ \left\| X \left(\sum_ {i = 1} ^ {k} g _ {i} m _ {i} ^ {T} - T\right) \right\| _ {F} ^ {2} \right] \\ = \operatorname {E} \left[ \left\| X \left(g _ {1} m _ {1} ^ {T} + \left(\sum_ {i = 2} ^ {k} g _ {i} m _ {i} ^ {T} - T\right)\right) \right\| _ {F} ^ {2} \right] \\ \leq (1 - \rho) \mathrm {E} \left[ X \left(\left\| \sum_ {i = 2} ^ {k} g _ {i} m _ {i} ^ {T} - T\right) \right\| _ {F} ^ {2} \right] \tag {11} \\ \leq \left(1 - \rho\right) ^ {k} \operatorname {E} \left[ \| X T \| _ {F} ^ {2} \right]. \\ \end{array}
$$
where (11) used the  $k = 1$  case shown above, and (12) used the inductive hypothesis. This completes the proof for general  $k$  and  $d_{2}$  that we needed for the full theorem.
B.1 Technical lemmas
It remains to show an interesting lemma used for proving the vector case in Theorem B.2.
Lemma B.4. Let  $a_1, \ldots, a_n \geq 0$  be IID random variables and assume some values  $p_i \geq 0$  st.  $\sum_{i} p_i = 1$  and  $p_n \leq 1 / n$ . Then
$$
E \left[ \frac {a _ {n}}{\sum_ {i} p _ {i} a _ {i}} \right] \geq 1.
$$
This completes the original proof with  $p_i = \frac{\sigma_i^2}{\sum_j \sigma_j^2}$  and  $a_i = g_i^2$ .
Proof. Since the  $a_i$  are IID, it doesn't matter if we permute them. In particular, if  $\pi$  is a random permutation of  $\{1, \dots, n-1\}$ ,
$$
\begin{array}{l} E \left[ \frac {a _ {n}}{\sum_ {i} p _ {i} a _ {i}} \right] = E _ {a} \left[ E _ {\pi} \left[ \frac {a _ {n}}{p _ {n} a _ {n} + \sum_ {i} p _ {i} a _ {\pi_ {i}}} \right] \right] (12) \\ \geq E _ {a} \left[ \frac {a _ {n}}{E _ {\pi} \left[ p _ {n} a _ {n} + \sum_ {i <   n} p _ {i} a _ {\pi_ {i}} \right]} \right] (13) \\ = E _ {a} \left[ \frac {a _ {n}}{p _ {n} a _ {n} + \sum_ {i <   n} p _ {i} \left(\frac {1}{n - 1} \sum_ {j <   n} a _ {j}\right)} \right] (14) \\ = E _ {a} \left[ \frac {a _ {n}}{p _ {n} a _ {n} + (1 - p _ {n}) \sum_ {i <   n} \frac {a _ {i}}{n - 1}} \right], (15) \\ \end{array}
$$
where Equation (13) uses Jensen's inequality on the convex function  $1 / x$ .
Now define  $a = \sum_{i=1}^{n} a_i$ . By permuting  $a_n$  with the other variables, we get:
$$
\begin{array}{l} E _ {a} \left[ \frac {a _ {n}}{p _ {n} a _ {n} + (1 - p _ {n}) \sum_ {i <   n} \frac {a _ {i}}{n - 1}} \right] = E _ {a} \left[ \frac {a _ {n}}{p _ {n} a _ {n} + \frac {1 - p _ {n}}{n - 1} (a - a _ {n})} \right] (16) \\ = E _ {a} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \frac {a _ {i}}{p _ {n} a _ {i} + \frac {1 - p _ {n}}{n - 1} (a - a _ {i})} \right] (17) \\ = E _ {a} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \frac {a _ {i} / a}{\frac {1 - p _ {n}}{n - 1} - \left(\frac {1 - p _ {n}}{n - 1} - p _ {n}\right) a _ {i} / a} \right] (18) \\ = E _ {a} \left[ \frac {1}{n} \sum_ {i = 1} ^ {n} \phi \left(a _ {i} / a\right) \right], (19) \\ \end{array}
$$
$$
\phi (q _ {i}) = \frac {q _ {i}}{\frac {1 - p _ {n}}{n - 1} - (\frac {1 - p _ {n}}{n - 1} - p _ {n}) q _ {i}}
$$
is convex whenever  $\frac{1 - p_n}{n - 1} / \left( \frac{1 - p_n}{n - 1} - p_n \right) = \frac{1 - p}{1 - np} > 1$ , which is true when  $0 \leq p_n < 1/n$ . That means we can use Jensen's again:
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \phi \left(a _ {i} / a\right) \geq \phi \left(\frac {1}{n} \sum_ {i} \frac {a _ {i}}{a}\right) = \phi \left(\frac {1}{n}\right) = 1,
$$
which is what we wanted to show.
C Correspondence between theory and practice
Theorem 3.1 quite tightly matches the empirical behavior of Algorithm 1 (Dense CCE for Least Squares). In Figure 8, we show the convergence when solving the least squares problem  $\mathrm{argmin}_T||XT - Y||_F^2$  using the method.
Figure 8: Convergence of the least squares problem. We show and compare the convergence of the least square methods using different methods. We found that the convergence of Theorem 3.1 is close to that of Algorithm 1.
As discussed in Appendix B, we can often strengthen the result from using  $\rho$  (smallest singular value / the squared Frobenius norm) to just using  $1 / d_{1}$ . This is because the smallest singular value issue only becomes relevant when  $Y$  is chosen to exactly align with the worst possible subspace of  $X$ . In practice this is rarely the case. In the particular case of Figure 8, we chose  $X$  and  $Y$  as iid standard normals.
D Hashing
If  $h: [n] \to [m]$  and  $s: [n] \to \{-1, 1\}$  are random functions, a Count Sketch is a matrix  $H \in \{0, -1, 1\}^{m \times n}$  where  $H_{i,j} = s(i)$  if  $h(i) = j$  and 0 otherwise. Charikar et al. [2002] showed that if  $m$  is large enough, the matrix  $H$  is a dimensionality reduction in the sense that the norm  $\| x \|_2$  of any vector in  $\mathbb{R}^n$  is approximately preserved,  $\| Hx \|_2 \approx \| x \|_2$ .<sup>6</sup>
This gives a simple theoretical way to think about the algorithms above: The learned matrix  $T' = H^T T$  is simply a lower dimensional approximation to the real table that we wanted to learn. While the theoretical result requires the random "sign function"  $s$  for the approximation to be unbiased, in practice this appears to not be necessary when directly learning  $T'$ . Maybe because the vectors can simply be slightly shifted to debias the result.
There are many strong theoretical results on the properties of Count Sketches. For example, Woodruff [2014] showed that they are so called "subspace embeddings" which means the dimensionality reduction is "robust" and doesn't have blind spots that SGD may accidentally walk into. However, the most practical result is that one only needs  $h$  to be a "universal hash function" ala Carter and Wegman [1977], which can be as simple and fast as the "multiply shift" hash function by Dietzfelbinger et al. [1997].
If Count Sketch shows that hashing each  $i \in [n]$  to a single row in  $[m]$ , we may wonder why methods like Hash Embeddings use multiple hash functions (or DHE uses more than a thousand.) The answer can be seen in the theoretical analysis of the "Johnson Lindenstrauss" transformation and in particular the "Sparse Johnson Lindenstrauss" as analyzed by Cohen et al. [2018]. The analysis shows that if the data being hashed is not very uniform, it is indeed better to use more than one hash function (more than 1 non-zero per column in  $H$ .) The exact amount depends on characteristics in the data distribution, but one can always get away with a sparsity of  $\epsilon$  when looking for a  $1 + \epsilon$  dimensionality reduction. Hence we speculate that DHE could in general replace the 1024 hash functions with something more like Hash Embeddings with an MLP on top. Another interesting part of the Cohen et al. [2018] analysis is that one should ideally split  $[m]$  in segments, and have one hash function into each segment. This matches the implementations we based our work on below.
This also implies that inner products are approximately preserved by the dimensionality reduction.
E How to store the hash functions
We note that unlike the random hash functions used in Step 1, the index pointer functions obtained from clustering takes space linear in the amount of training data or at least in the ID universe size. At first this may seem like a major downside of our method, and while it isn't different from the index tables needed after Product Quantization, it definitely is something extra not needed by purely sketching based methods.
We give three reasons why storing this table is not an issue in practice:
1. The index pointer functions can be stored on the CPU rather than the GPU, since they are used as the first step of the model before the training/inference data has been moved from the CPU to the GPU. Furthermore the index lookup is typically done faster on CPUs, since it doesn't involve any dense matrix operations.
2. The index pointers can replace the original IDs. Unless we are working in a purely streaming setting, the training data has to be stored somewhere. If IDs are 64 bit integers, replacing them with four 16-bit index pointers is net neutral.
3. Some hashing and pruning can be used as a prepossessing step, reducing the universe size of the IDs and thus the size of the index table needed.
F Different strategies for CCE
We include other graphs about CCE in Figure 9. They are all on the Kaggle dataset and were run three times. These graphs helped us develop insights on CCE and choose the correct versions for Figure 4a and Figure 4b.
G AUC Graphs
We also evaluate the models using AUC, which is another commonly used metric for gauging the effectiveness of a recommendation model. For example, it was used in [Kang et al., 2021]. AUC provides the probability of getting a correct prediction when evaluating a test sample from a balanced dataset. Therefore, a better model is implied by a larger AUC. In this section, we plot the graphs again using AUC; see Figure 10 and Figure 11.
H Table Collapse
Table collapsing was a problem we encountered for the circular clustering method as described in Appendix A. We describe the problem and the metric we used to detect it here, since we think they may be of interest to the community.
Suppose we are doing  $k$ -means clustering on a table of 3 partitions in order to obtain 3 index pointer functions  $h_j^c$ . These functions can be thought as a table, where the  $(i,j)$ -entry is given by  $h_j^c(i)$ .
There are multiple failure modes we have to be aware of. The first one is column-wise collapse:
1	0	0
1	1	2
1	0	3
:	:	:
1	3	1
In this table the first column has collapsed to just one cluster. Because of the way  $k$ -means clustering works, this exact case of complete collapse isn't actually possible, but we might get arbitrarily low entropy as measured by  $H_{1}$ , which we define as follows: For each column  $j$ , its column entropy is defined to be the entropy of the probability distribution  $p_{j} \colon h_{j}^{c}([n]) \to [0,1]$  defined by
$$
p _ {j} (x) = \frac {\# \{i : h _ {j} ^ {c} (i) = x \}}{n}.
$$
(a) Kaggle dataset, CCE, best of 10 epoch: We ran different versions of CCE for 10 epochs. Here ct means the number of clustering done, and cf refers to the number of batches between clusterings. Since each epoch has around 300,000 batches, we essentially clustered once every epoch. The performance increases with more clustering. Another observation is that as  $m$  increases, a few lines were merged due to early stopping. We found that CCE ct6 cf300000 performs the best, which becomes the CCE model in Figure 4a.
(b) Kaggle dataset, CCE, Strategy 1: We ran different versions of CCE for 1 epoch under the constraint that all clusterings must finish before half of an epoch. It turns out that there is a balance between the number of clusterings and the 'quality' of the clustering, represented by the number of batches between clusterings. We found that CCE ct2 cf75000 performs the best, which becomes the CCE model in Figure 4b.
(c) Kaggle dataset, CCE, Strategy 2: Strategy 2 here tries to have all the clustering finish by  $2/3$  of an epoch. These runs did not perform well. It turned out that we need to let the model have time to converge after all the clusterings.
Figure 9: Strategies for CCE that gave us insight.
(d) Kaggle dataset, CCE, Strategy 3: This strategy perfectly summarizes all the previous findings. Increasing the number of clusterings in general gives better performance; Letting the model 'rest' after clustering increases the performance; Increasing the interval between clusterings give better result.
Then we define  $H_{1}$  to be the minimum entropy of the (here 3) column-entropies.
The second failure mode is pairwise collapse:
1	0	1
2	2	3
1	0	3
3	1	0
2	2	1
(a) Kaggle dataset, Best of 10 epoch, AUC.
(b) Kaggle dataset, 1 epoch, AUC.
Figure 10: The AUC version of Figure 4.
(c) Terabyte dataset, 1 epoch, AUC.
In this case the second column is just a permutation of the first column. This means the expanded set of possible vectors is much smaller than we would expect. We can measure pairwise collapse by computing the entropy of the histogram of pairs, where the entropy of the column pair  $(j_{1}, j_{2})$  is defined by the column entropy of  $h_{j_{1}}^{c}(\cdot) + \max(h_{j_{1}}^{c}) h_{j_{2}}^{c}(\cdot)$ . Then we define  $H_{2}$  to be the minimum of such pair-entropies for all  $\binom{3}{2}$  pairs of columns.
Pairwise entropy can be trivially generalized to triple-wise and so on. If we have  $c$  columns we may compute each of  $H_{1},\ldots ,H_{c}$ . In practice  $H_{1}$  and  $H_{2}$  may contain all the information we need.
H.1 What entropies are expected?
The maximum value for  $H_{1}$  is  $\log k$ , in the case of a uniform distribution over clusters. The maximum value for  $H_{2}$  is  $\log \binom{k}{2} \approx 2\log k$ . (Note  $\log n$  is also an upper bound, where  $n$  is the number of points in the dataset / rows in the table.)
With the CE method we expect all the entropies to be near their maximum. However, for the Circular Clustering method this is not the case! That would mean we haven't been able to extract any useful cluster information from the data.
Instead we expect entropies close to what one gets from performing Product Quantization (PQ) on a complete dataset. In short:
(a) Kaggle dataset, CCE, best of 10 epoch, AUC.
(b) Kaggle dataset, CCE, Strategy 1, AUC.
(c) Kaggle dataset, CCE, Strategy 2, AUC.
Figure 11: The AUC version of Figure 9.
(d) Kaggle dataset, CCE, Strategy 3, AUC.
1. Too high entropy: We are just doing CE more slowly.
2. Too low entropy: We have a table collapse.
3. Golden midpoint: Whatever entropy normal PQ gets.