SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $s_{\lambda} = \sum_{i=1}^{R} \frac{\lambda_i}{\lambda_i + \lambda}$ ,  $\lambda_1, \ldots, \lambda_R$  are the eigenvalues of the Gram matrix  $G$  of matrix  $A$ , and  $\lambda \geq 0$  is a regularization parameter. For comparison, our runtime for constant failure probability  $\delta$  is
$$
O \left(R \sum_ {i = 1} ^ {N} \mathrm {n n z} \left(U _ {i}\right) + \frac {R ^ {3}}{\varepsilon} \log \left(\prod_ {i = 1} ^ {N} I _ {i}\right) \log R\right).
$$
Woodruff and Zandieh's method provides a significant advantage for large column count  $R$  or high regularization parameter  $\lambda$ . As a result, it is well-suited to the problem of regularized low-rank approximation when the column count  $R$  is given by the number of data points in a dataset. On the
other hand, the algorithm has poor dependence on the matrix count  $N$  and error parameter  $\varepsilon$ . For tensor decomposition,  $R$  is typically no larger than a few hundred, while high accuracy ( $\epsilon \approx 10^{-3}$ ) is required for certain tensors to achieve a fit competitive with non-randomized methods (see section 4.3, Figures 4 and 5). When  $\lambda$  is small, we have  $s_{\lambda} \approx R$ . Here, Woodruff and Zandieh's runtime has an  $O(R^3)$  dependence similar to ours. When  $R \leq \log^4 R \log N$ , our sampler has faster construction time as well.
Finally, we note that our sampling data structure can be constructed using highly cache-efficient, parallelizable symmetric rank- $R$  updates (BLAS3 operation dSYRK). As a result, the quadratic dependence on  $R$  in our algorithm can be mitigated by dense linear algebra accelerators, such as GPUs or TPUs.
A.4 Proof of Theorem 3.1
Theorem 3.1 appeared in a modified form as Lemma 10 in the work by Malik [18]. This original version used the definition  $\tilde{G}_{>k} = \Phi * \bigotimes_{a=k+1}^{N} G_k$  in place of  $G_{>k}$  defined in Equation (4), where  $\Phi$  was a sketched approximation of  $G^+$ . Woodruff and Zandieh [29] exhibit a version of the theorem with similar modifications. We prove the version stated in our work below.
Proof of Theorem 3.1. We rely on the assumption that the Khatri-Rao product  $A$  is a nonzero matrix (but it may be rank-deficient). We begin by simplifying the expression for the leverage score of a row of  $A$  corresponding to multi-index  $(i_1, \dots, i_N)$ . Beginning with Equation (2), we derive
$$
\begin{array}{l} \ell_ {i _ {1}, \dots , i _ {N}} = A [ (i _ {1}, \dots , i _ {N}),: ] G ^ {+} A [ (i _ {1}, \dots , i _ {N}),: ] ^ {\top} \\ = \left\langle A \left[ \left(i _ {1}, \dots , i _ {N}\right), : \right] ^ {\top} A \left[ \left(i _ {1}, \dots , i _ {N}\right), : \right], G ^ {+} \right\rangle \\ = \left\langle \left(\bigoplus_ {a = 1} ^ {N} U _ {a} [ i _ {a},: ]\right) ^ {\top} \left(\bigoplus_ {a = 1} ^ {N} U _ {a} [ i _ {a},: ]\right), G ^ {+} \right\rangle \\ = \langle \bigoplus_ {a = 1} ^ {N} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ], G ^ {+} \rangle \\ = \langle \bigotimes_ {a = 1} ^ {k - 1} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ], U _ {k} [ i _ {k},: ] ^ {\top} U _ {k} [ i _ {k},: ] * \bigotimes_ {a = k + 1} ^ {N} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ], G ^ {+} \rangle \\ = \left\langle \bigotimes_ {a = 1} ^ {k - 1} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ], U _ {k} [ i _ {k},: ] ^ {\top} U _ {k} [ i _ {k},: ], G ^ {+} * \bigotimes_ {a = k + 1} ^ {N} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ] \right\rangle . \tag {13} \\ \end{array}
$$
We proceed to the main proof of the theorem. To compute  $p(\hat{s}_k = s_k \mid \hat{s}_{< k} = s_{< k})$ , we marginalize over random variables  $\hat{s}_{k+1} \ldots \hat{s}_N$ . Recalling the definition of  $h_{< k}$  from Equation (3), we have
$$
\begin{array}{l} p \left(\hat {s} _ {k} = s _ {k} \mid \hat {s} _ {<   k} = s _ {<   k}\right) \propto \sum_ {i _ {k + 1}, \dots , i _ {N}} p \left(\left(\hat {s} _ {<   k} = s _ {<   k}\right) \wedge \left(\hat {s} _ {k} = s _ {k}\right) \wedge \bigwedge_ {u = k + 1} ^ {N} \left(\hat {s} _ {u} = i _ {u}\right)\right) \tag {14} \\ \propto \sum_ {i _ {k + 1}, \dots , i _ {N}} \ell_ {s _ {1}, \dots , s _ {k}, i _ {k + 1}, \dots , i _ {N}}. \\ \end{array}
$$
The first line above follows by marginalizing over  $\hat{s}_{k + 1},\dots,\hat{s}_N$ . The second line follows because the joint random variable  $(\hat{s}_1,\dots,\hat{s}_N)$  follows the distribution of statistical leverage scores on the rows of
A. We now plug in Equation (13) to get
$$
\begin{array}{l} \sum_ {i _ {k + 1}, \dots , i _ {N}} \ell_ {s _ {1}, \dots , s _ {k}, i _ {k + 1}, \dots , i _ {N}} \\ = \sum_ {i _ {k + 1}, \dots , i _ {N}} \langle \bigotimes_ {a = 1} ^ {k - 1} U _ {a} [ s _ {a},: ] ^ {\top} U _ {a} [ s _ {a},: ], U _ {k} [ s _ {k},: ] ^ {\top} U _ {k} [ s _ {k},: ], G ^ {+} * \bigotimes_ {a = k + 1} ^ {N} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ] \rangle \\ = \sum_ {i _ {k + 1}, \dots , i _ {N}} \left\langle h _ {<   k} h _ {<   k} ^ {\top}, U _ {k} [ s _ {k},: ] ^ {\top} U _ {k} [ s _ {k},: ], G ^ {+} * \bigotimes_ {a = k + 1} ^ {N} U _ {a} [ i _ {a},: ] ^ {\top} U _ {a} [ i _ {a},: ] \right\rangle \\ = \left\langle h _ {<   k} h _ {<   k} ^ {\top}, U _ {k} \left[ s _ {k},: \right] ^ {\top} U _ {k} \left[ s _ {k},: \right], G ^ {+} * \bigotimes_ {a = k + 1} ^ {N} \sum_ {i _ {a} = 1} ^ {I _ {a}} U _ {a} \left[ i _ {a},: \right] ^ {\top} U _ {a} \left[ i _ {a},: \right] \right\rangle \\ = \left\langle h _ {<   k} h _ {<   k} ^ {\top}, U _ {k} \left[ s _ {k},: \right] ^ {\top} U _ {k} \left[ s _ {k},: \right], G ^ {+} * \bigotimes_ {a = k + 1} ^ {N} G _ {a} \right\rangle \\ = \left\langle h _ {<   k} h _ {<   k} ^ {\top}, U _ {k} [ s _ {k},: ] ^ {\top} U _ {k} [ s _ {k},: ], G _ {> k} \right\rangle . \tag {15} \\ \end{array}
$$
We now compute the normalization constant  $C$  for the distribution by summing the last line of Equation (15) over all possible values for  $\hat{s}_k$ :
$$
\begin{array}{l} C = \sum_ {s _ {k} = 1} ^ {I _ {k}} \left\langle h _ {<   k} h _ {<   k} ^ {\top}, U _ {k} \left[ s _ {k},: \right] ^ {\top} U _ {k} \left[ s _ {k},.: \right], G _ {> k} \right\rangle \\ = \left\langle h _ {<   k} h _ {<   k} ^ {\top}, \sum_ {s _ {k} = 1} ^ {I _ {k}} U _ {k} [ s _ {k},: ] ^ {\top} U _ {k} [ s _ {k},: ], G _ {> k} \right\rangle \tag {16} \\ = \langle h _ {<   k} h _ {<   k} ^ {\top}, G _ {k}, G _ {> k} \rangle . \\ \end{array}
$$
For  $k = 1$ , we have  $h_{<k} = [1, \dots, 1]^{\top}$ , so  $C = \langle G_k, G_{>k} \rangle$ . Then  $C$  is the sum of all leverage scores, which is known to be the rank of  $A$  [30]. Since  $A$  was assumed nonzero,  $C \neq 0$ . For  $k > 1$ , assume that the conditioning event  $\hat{s}_{<k} = s_{<k}$  occurs with nonzero probability. This is a reasonable assumption, since our sampling algorithm will never select prior values  $\hat{s}_1, \dots, \hat{s}_{k-1}$  that have 0 probability of occurrence. Let  $\tilde{C}$  be the normalization constant for the conditional distribution on  $\hat{s}_{k-1}$ . Then we have
$$
\begin{array}{l} 0 <   p \left(\hat {s} _ {k - 1} = s _ {k - 1} \mid \hat {s} _ {<   k - 1} = s _ {<   k - 1}\right) \\ = \tilde {C} ^ {- 1} \left\langle h _ {<   k - 1} h _ {<   k - 1} ^ {\top}, U _ {k - 1} [ s _ {k - 1},: ] ^ {\top} U _ {k - 1} [ s _ {k - 1},: ], G _ {> k - 1} \right\rangle \\ = \tilde {C} ^ {- 1} \left\langle h _ {<   k} h _ {<   k} ^ {\top}, G _ {> k - 1} \right\rangle (17) \\ = \tilde {C} ^ {- 1} \left\langle h _ {<   k} h _ {<   k} ^ {\top}, G _ {k} * G _ {> k} \right\rangle (17) \\ = \tilde {C} ^ {- 1} \left\langle h _ {<   k} h _ {<   k} ^ {\top}, G _ {k}, G _ {> k} \right\rangle \\ = \tilde {C} ^ {- 1} C \\ \end{array}
$$
Since  $\tilde{C} > 0$ , we must have  $C > 0$ .
A.5 Proof of Lemma 3.2
We detail the construction procedure, sampling procedure, and correctness of our proposed data structure. Recall that  $T_{I,F}$  denotes the collection of nodes in a full, complete binary tree with  $\lceil I / F \rceil$  leaves. Each leaf  $v \in T_{I,F}$  holds a segment  $S(v) = \{S_0(v)..S_1(v)\} \subseteq \{1..I\}$ , with  $|S(v)| \leq F$  and  $S(u) \cap S(v) = \emptyset$  for distinct leaves  $u, v$ . For each internal node  $v$ ,  $S(v) = S(L(v)) \cup S(R(v))$ , where  $L(v)$  and  $R(v)$  denote the left and right children of node  $v$ . The root node  $r$  satisfies  $S(r) = \{1..I\}$ .
**Construction:** Algorithm 3 gives the procedure to build the data structure. We initialize a segment tree  $T_{I,F}$  and compute  $G^v$  for all leaf nodes  $v \in T_{I,F}$  as a sum of outer products of rows from  $U$  (lines 1-3). Starting at the level above the leaves, we then compute  $G^v$  for each internal node as
the sum of  $G^{L(v)}$  and  $G^{R(v)}$ , the partial Gram matrices of its two children. Runtime  $O(IR^2)$  is required to compute  $I$  outer products across all iterations of the loop on line 3. Our segment tree has  $\lceil I / F \rceil - 1$  internal nodes, and the addition in line 6 contributes runtime  $O(R^2)$  for each internal node. This adds complexity  $O(R^2(\lceil I / F \rceil - 1)) \leq O(IR^2)$ , for total construction time  $O(IR^2)$ .
To analyze the space complexity, observe that we store a matrix  $G^v \in \mathbb{R}^{R \times R}$  at all  $2\lceil I / F \rceil - 1$  nodes of the segment tree, for asymptotic space usage  $O(\lceil I / F \rceil R^2)$ . We can cut the space usage in half by only storing  $G^v$  when  $v$  is either the root or a left child in our tree, since the sampling procedure in Algorithm 4 only accesses the partial Gram matrix stored by left children. We can cut the space usage in half again by only storing the upper triangle of each symmetric matrix  $G^v$ . Finally, in the special case that  $I < F$ , the segment tree has depth 1 and the initial binary search can be eliminated entirely. As a result, the data structure has  $O(1)$  space overhead, since we can avoid storing any partial Gram matrices  $G^v$ . This proves the complexity claims in point 1 of Lemma 3.2.
Algorithm 3 BuildSampler  $(U\in \mathbb{R}^{I\times R},F,Y)$
1: Build tree  $T_{I,F}$  with depth  $d = \lceil \log \lceil I / F\rceil \rceil$   
2: for  $v\in \mathrm{leaves}(T_{I,F})$  do  
3:  $G^{v}:= \sum_{i\in S(v)}U[i,:]^{\top}U[i,:]$   
4: for  $u = d - 2\dots 0$  do  
5: for  $v\in \mathrm{level}(T_{I,F},u)$  do  
6:  $G^{v}:= G^{L(v)} + G^{R(v)}$
Sampling: Algorithm 4 gives the procedure to draw a sample from our proposed data structure. It is easy to verify that the normalization constant  $C$  for  $q_{h,U,Y}$  is  $\langle hh^{\top}, G^{\mathrm{root}(T_{I,F}), Y} \rangle$ , since  $G^{\mathrm{root}(T_{I,F})} = U^{\top}U$ . Lines 8 and 9 initialize a pair of templated procedures  $\tilde{m}$  and  $\tilde{q}$ , each of which accepts a node from the segment tree. The former is used to compute the branching threshold at each internal node, and the latter returns the probability vector  $q_{h,U,Y}[S_0(v): S_1(v)]$  for the segment  $\{S_0(v)...S_1(v)\}$  maintained by a leaf node. To see this last fact, observe for  $i \in [I]$  that
$$
\begin{array}{l} \tilde {q} (v) [ i - S _ {0} (v) ] \\ = C ^ {- 1} U [ i,: ] \cdot (h h ^ {\top} \circledast Y) \cdot U [ i,: ] ^ {\top} \tag {18} \\ = C ^ {- 1} \left\langle h h ^ {\top}, U [ i,: ] ^ {\top} U [ i,: ], Y \right\rangle \\ = q _ {h, U, Y} [ i ]. \\ \end{array}
$$
The loop on line 12 performs the binary search using the two templated procedures. Line 18 uses the procedure  $\tilde{q}$  to scan through at most  $F$  bin endpoints after the binary search finishes early.
The depth of segment tree  $T_{I,F}$  is  $\log \lceil I / F \rceil$ . As a result, the runtime of the sampling procedure is dominated by  $\log \lceil I / F \rceil$  evaluations of  $\tilde{m}$  and a single evaluation of  $\tilde{q}$  during the binary search. Each execution of procedure  $\tilde{m}$  requires time  $O(R^2)$ , relying on the partial Gram matrices  $G^{v}$  computed during the construction phase. When  $Y$  is a general p.s.d. matrix, the runtime of  $\tilde{q}$  is  $O(FR^2)$ . This complexity is dominated by the matrix multiplication  $W \cdot (hh^\top \circledast Y)$  on line 5. In this case, the runtime of the "RowSampler" procedure to draw one sample is  $O(R^2\log \lceil I / F\rceil +FR^2)$ , satisfying the complexity claims in point 2 of the lemma.
Now suppose  $Y$  is a rank-1 matrix with  $Y = uu^{\top}$  for some vector  $u$ . We have  $hh^{\top} \circledast Y = (h \circledast u)(h \circledast u)^{\top}$ . This gives
$$
\tilde {q} _ {p} (h, C, v) = \operatorname {d i a g} \left(W \cdot \left(h h ^ {\top} * u u ^ {\top}\right) \cdot W\right) = \left(W \cdot \left(h * u\right)\right) ^ {2}
$$
where the square is elementwise. The runtime of the procedure  $\tilde{q}$  is now dominated by a matrix-vector multiplication that costs time  $O(FR)$ . In this case, we have per-sample complexity  $\mathcal{O}(R^2\log \lceil I / F\rceil + FR)$ , matching the complexity claim in point 3 of the lemma.
Correctness: Recall that the inversion sampling procedure partitions the interval  $[0,1]$  into  $I$  bins, the  $i$ -th bin having width  $q_{h,U,Y}[i]$ . The goal of our procedure is to find the bin that contains the uniform random draw  $d$ . Since procedure  $\tilde{m}$  correctly returns the branching threshold (up to the offset "low") given by Equation (7), the loop on line 12 correctly implements a binary search on the list of bin endpoints specified by the vector  $q_{h,U,Y}$ . At the end of the loop,  $c$  is a leaf node that
Algorithm 4 Row Sampling Procedure
Require: Matrices  $U, Y$  saved from construction, partial Gram matrices  $\{G^v \mid v \in T_{I,F}\}$ .  
1: procedure  $m_p(h, C, v)$   
2: return  $C^{-1} \langle hh^\top, G^v, Y \rangle$   
3: procedure  $q_p(h, C, v)$   
4:  $W := U[S(v), :]$   
5: return  $C^{-1} \operatorname{diag}(W \cdot (hh^\top \otimes Y) \cdot W^\top)$   
6: procedure RowSample(h)  
7:  $C := \langle hh^\top, G^{\mathrm{root}(T_{I,F})}, Y \rangle$   
8:  $\tilde{m}(\cdot) := m_p(h, C, \cdot)$   
9:  $\tilde{q}(\cdot) := q_p(h, C, \cdot)$   
10:  $c := \operatorname{root}(T_{I,F})$ , low = 0.0, high = 1.0  
11: Sample  $d \sim \mathrm{Uniform}(0.0, 1.0)$   
12: while  $c \notin \operatorname{leaves}(T_{I,F})$  do  
13: cutoff := low +  $\tilde{m}(L(c))$   
14: if cutoff ≥ d then  
15:  $c := L(c)$ , high := cutoff  
16: else  
17:  $c := R(c)$ , low := cutoff  
18: return  $S_0(v) + \arg \min_{i \geq 0} \left( \log + \sum_{j=1}^{i} \tilde{q}(c) [j] < d \right)$
maintains a collection  $S(c)$  of bins, one of which contains the random draw  $d$ . Since the procedure  $\tilde{q}$  correctly returns probabilities  $q_{h,U,Y}[i]$  for  $i \in S(c)$  for leaf node  $c$ , (see Equation (18)), line 18 finds the bin that contains the random draw  $d$ . The correctness of the procedure follows from the correctness of inversion sampling [22].
A.6 Cohesive Proof of Theorem 1.1
In this proof, we fully explain Algorithms 1 and 2 in the context of the sampling procedure outlined in section 3.2. We verify the complexity claims first and then prove correctness.
**Construction and Update:** For each matrix  $U_{j}$ , Algorithm 1 builds an efficient row sampling data structure  $Z_{j}$  as specified by Lemma 3.2. We let the p.s.d. matrix  $Y$  that parameterizes each sampler be a matrix of ones, and we set  $F = R$ . From Lemma 3.2, the time to construct sampler  $Z_{j}$  is  $O(I_jR^2)$ . The space used by sampler  $Z_{j}$  is  $O(\lceil I_j / F\rceil R^2) = O(I_jR)$ , since  $F = R$ . In case  $I_{j} < R$ , we use the special case described in Appendix A.5 to get a space overhead  $O(1)$ , avoiding a term  $O(R^{2})$  in the space complexity.
Summing the time and space complexities over all  $j$  proves part 1 of the theorem. To update the data structure if matrix  $U_{j}$  changes, we only need to rebuild sampler  $Z_{j}$  for a cost of  $O(I_jR^2)$ . The construction phase also computes and stores the Gram matrix  $G_{j}$  for each matrix  $U_{j}$ . We defer the update procedure in case a single entry of matrix  $U_{j}$  changes to Appendix A.7.
Sampling: For all indices  $k$  (except possibly  $j$ ), lines 1-5 from Algorithm 2 compute  $G_{>k}$  and its eigendecomposition. Only a single pass over the Gram matrices  $G_{k}$  is needed, so these steps cost  $O(R^{3})$  for each index  $k$ . Line 5 builds an efficient row sampler  $E_{k}$  for the matrix of scaled eigenvectors  $\sqrt{\Lambda_k} \cdot V_k$ . For sampler  $k$ , we set  $Y = G_{k}$  with cutoff parameter  $F = 1$ . From Lemma 3.2, the construction cost is  $O(R^3)$  for each index  $k$ , and the space required by each sampler is  $O(R^3)$ . Summing these quantities over all  $k \neq j$  gives asymptotic runtime  $O(NR^3)$  for lines 2-5.
The loop spanning lines 6-12 draws  $J$  row indices from the Khatri-Rao product  $U_{\neq j}$ . For each sample, we maintain a "history vector"  $h$  to write the variables  $h_{<k}$  from Equation (3). For each index  $k \neq j$ , we draw random variable  $\hat{u}_k$  using the row sampler  $E_k$ . This random draw indexes a scaled eigenvector of  $G_{>k}$ . We then use the history vector  $h$  multiplied by the eigenvector to sample a row index  $\hat{t}_k$  using data structure  $Z_k$ . The history vector  $h$  is updated, and we proceed to draw the next index  $\hat{t}_k$ .
As written, lines 2-5 also incur scratch space usage  $O(NR^3)$ . The scratch space can be reduced to  $O(R^3)$  by exchanging the order of loops on line 6 and line 8 and allocating  $J$  separate history vectors  $h$ , once for each draw. Under this reordering, we perform all  $J$  draws for each variable  $\hat{u}_k$  and  $\hat{t}_k$  before moving to  $\hat{u}_{k+1}$  and  $\hat{t}_{k+1}$ . In this case, only a single data structure  $E_k$  is required at each iteration of the outer loop, and we can avoid building all the structures in advance on line 5. We keep the algorithm in the form written for simplicity, but we implemented the memory-saving approach in our code.
From Lemma 3.2, lines 9 and 10 cost  $O(R^2 \log R)$  and  $O\left(R^2 \log \lceil I_k / R \rceil\right)$ , respectively. Line 11 costs  $O(R)$  and contributes a lower-order term. Summing over all  $k \neq j$ , the runtime to draw a single sample is
$$
O \left(\sum_ {k \neq j} \left(R ^ {2} \log \lceil I _ {k} / R \rceil + R ^ {2} \log R\right)\right) = O \left(\sum_ {k \neq j} R ^ {2} \log \max  \left(I _ {k}, R\right)\right).
$$
Adding the runtime for all  $J$  samples to the runtime of the loop spanning lines 2-6 gives runtime  $O\left(NR^3 + J\sum_{k \neq j} R^2 \log \max (I_k, R)\right)$ , and the complexity claims have been proven.
Correctness: We show correctness for the case where  $j = -1$  and we sample from the Khatri-Rao product of all matrices  $U_{k}$ , since the proof for any other value of  $j$  requires a simple reindexing of matrices. To show that our sampler is correct, it is enough to prove the condition that for  $1 \leq k \leq N$ ,
$$
p \left(\hat {t} _ {k} = t _ {k} \mid h _ {<   k}\right) = q _ {h _ {<   k}, U _ {k}, G _ {> k}} \left[ t _ {k} \right], \tag {19}
$$
since, by Theorem 3.1,  $p(\hat{s}_k = s_k \mid \hat{s}_{<k} = s_{<k}) = q_{h_{<k}, U_k, G_{>k}}[s_k]$ . This would imply that the joint random variable  $(\hat{t}_1, \dots, \hat{t}_N)$  has the same probability distribution as  $(\hat{s}_1, \dots, \hat{s}_N)$ , which by definition follows the leverage score distribution on  $U_1 \odot \dots \odot U_N$ . To prove the condition in Equation (19), we apply Equations (9) and (11) derived earlier:
$$
\begin{array}{l} p \left(\hat {t} _ {k} = t _ {k} \mid h _ {<   k}\right) \\ = \sum_ {u _ {k} = 1} ^ {R} p \left(\hat {t} _ {k} = t _ {k} \mid \hat {u} _ {k} = u _ {k}, h _ {<   k}\right) p \left(\hat {u} _ {k} = u _ {k} \mid h _ {<   k}\right) \quad \text {(B a y e s ＂ R u l e)} \\ = \sum_ {u _ {k} = 1} ^ {R} w [ u _ {k} ] \frac {W [ t _ {k} , u _ {k} ]}{\| W [ : , u _ {k} ] \| _ {1}} \quad \text {(E q u a t i o n s (9) a n d (1 1) , i n r e v e r s e)} \\ = q _ {h _ {<   k}, U _ {k}, G _ {> k}} \left[ t _ {k} \right]. \tag {20} \\ \end{array}
$$
A.7 Efficient Single-Element Updates
Applications such as CP decomposition typically change all entries of a single matrix  $U_{j}$  between iterations, incurring an update cost  $O(I_jR^2)$  for our data structure from Theorem 1.1. In case only a single element of  $U_{j}$  changes, our data structure can be updated in time  $O(R\log I_j)$ .
Proof. Algorithm 5 gives the procedure when the update  $U_{j}[r,c] \coloneqq \hat{u}$  is performed. The matrices  $G^{v}$  refer to the partial Gram matrices maintained by each node  $v$  of the segment trees in our data structure, and the matrix  $\tilde{U}_j$  refers to the matrix  $U_{j}$  before the update operation.
Let  $T_{I_j,R}$  be the segment tree corresponding to matrix  $U_j$  in the data structure, and let  $v \in T_{I_j,R}$  be the leaf whose segment contains  $r$ . Lines 3-5 of the algorithm update the row and column indexed by  $c$  in the partial Gram matrix held by the leaf node.
The only other nodes requiring an update are ancestors of  $v$ , each holding a partial Gram matrix that is the sum of its two children. Starting from the direct parent  $A(v)$ , the loop on line 6 performs these ancestor updates. The addition on line 8 only requires time  $O(R)$ , since only row and column  $c$  change between the old value of  $G^v$  and its updated version. Thus, the runtime of this procedure is  $O(R\log [I_j / R])$  from multiplying the cost to update a single node by the depth of the tree.
Algorithm 5 UpdateSampler  $(j,r,c,\hat{u})$
1: Let  $u = \tilde{U}_j[r,c]$   
2: Locate  $v$  such that  $r \in S(v)$   
3: Update  $G^v[c,:] += (\hat{u} - u)\tilde{U}_j[r,:]$   
4: Update  $G^v[;,c] += (\hat{u} - u)\tilde{U}_j[r,:]^{\top}$   
5: Update  $G^v[c,c] += (\hat{u} - u)^2$   
6: while  $v \neq \text{root}(T_{I_j,R})$  do  
7:  $v_{\text{prev}} := v, v := A(v)$   
8: Update  $G^v := G^{\text{sibling}(v_{\text{prev}})} + G^{v_{\text{prev}}}$
A.8 Extension to Sparse Input Matrices
Our data structure is designed to sample from Khatri-Rao products  $U_{1} \odot \ldots \odot U_{N}$  where the input matrices  $U_{1}, \ldots, U_{N}$  are dense, a typical situation in tensor decomposition. Slight modifications to the construction procedure permit our data structure to handle sparse matrices efficiently as well. The following corollary states the result as a modification to Theorem 1.1.
Corollary A.1 (Sparse Input Modification). When input matrices  $U_{1}, \ldots, U_{N}$  are sparse, point 1 of Theorem 1.1 can be modified so that the proposed data structure has  $O\left(R\sum_{j=1}^{N} nnz(U_{j})\right)$  construction time and  $O\left(\sum_{j=1}^{N} nnz(U_{j})\right)$  storage space. The sampling time and scratch space usage in point 2 of Theorem 1.1 does not change. The single-element update time in point 1 is likewise unchanged.
Proof. We will modify the data structure in Lemma 3.2. The changes to its construction and storage costs will propagate to our Khatri-Rao product sampler, which maintains one of these data structures for each input matrix.
Let us restrict ourselves to the case  $F = R, Y = [1]$  in relation to the data structure in Lemma 3.2. These choices for  $F$  and  $Y$  are used in the construction phase given by Algorithm 1. The proof in Appendix A.5 constrains each leaf  $v$  of a segment tree  $T_{I,F}$  to hold a contiguous segment  $S(v) \subseteq [I]$  of cardinality at most  $F$ . Instead, choose each segment  $S(v) = \{S_0(v),..S_1(v)\}$  so that  $U[S_0(v):S_1(v),\cdot ]$  has at most  $R^2$  nonzeros, and the leaf count of the tree is at most  $\lceil \mathrm{nnz}(U) / R^2\rceil +1$  for input matrix  $U \in \mathbb{R}^{I\times R}$ . Assuming the nonzeros of  $U$  are sorted in row-major order, we can construct such a partition of  $[I]$  into segments in time  $O(\mathrm{nnz}(U))$  by iterating in order through the nonzero rows and adding each of them to a "current" segment. We shift to a new segment when the current segment cannot hold any more nonzeros.
This completes the modification to the data structure in Lemma 3.2, and we now analyze its updated time / space complexity.