Learning Generative Embeddings using an Optimal Subsampling Policy for Tensor Sketching

Tucker decomposition has been considered as a multilinear generalization of SVD, and HOSVD is a constrained Tucker decomposition that ensures the orthogonality of factor matrices and all-orthogonality of the core tensor De Lathauwer et al. (2000). For the decomposition, a way to unfold tensor to matrix is needed to be defined. For order 3 tensor $A\in R^{n_1\times n_2\times n_3}$, "matrix unfolding" $A_{\left(1\right)}\in R^{n_1\times n_2{n}_3}$ contains the element $a_{i_1,i_2,i_3}$ at the position with row number $i_1$ and column number equal to $i_2\times n_3+n_2$. Other two matrix unfoldings $A_{\left(2\right)}$ and $A_{\left(3\right)}$ can be defined in the same manner. Then, we define “$k$-mode vectors,” defined as the $n_k$-dimensional vectors obtained from $A$ by varying the index $i_k$ and keeping the other indices fixed. With this definition, the 1-mode product of a tensor $A\in R^{n_1\times n_2\times n_3}$ by a matrix $U\in R^{m_1\times n_1}$, denoted by $A{\times }_{1}U$ is an $(m_1\times n_2\times n_3)$-tensor of which the entries are given by ${\sum }_{i_1}{a}_{i_1,i_2,i_3}{u}_{j_1,i_1}$. Other two mode product $A{\times }_{2}U$ and $A{\times }_{2}U$ can be also defined accordingly. Finally, order 3 SVD can be stated as following Theorem 2.1.

The algorithm for calculating core tensor $S$ and scaling matrices $U^{\left(k\right)}$ is given in the algorithm 1.

Thompson sampling is an algorithm for online decision problems where actions are taken sequentially in a manner that must balance between exploiting what is known to maximize immediate performance and investing to accumulate new information that may improve future performance Russo et al. (2018); Agrawal and Goyal (2012); Chapelle and Li (2011); Thompson (1933).

Suppose the agent has imperfect knowledge about the reward of each action of available action set $X$. Since the stochastic environment is assumed, the reward of each action is a random variable following a certain distribution denoted by $q_{\theta }$. Then, the imperfect knowledge about the reward of each action can be characterized by uncertain parameters, $\theta$. The distribution for the parameters, $p$, is introduced with the hyperparameters. To effectively balance exploitation and exploration, TS samples the parameters $\theta$ from $p$. Then, the expected reward of each action can be predicted by sampled $\theta$. The agent applies action selecting the one that has the highest expected reward. After applying action $x_t$, the agent observes an outcome $y_t$ and corresponding reward $r\left(y_t\right)$ from the stochastic environment. Using the observed $y_t,r\left(y_t\right)$, TS updates the hyperparameters using the Bayesian inference. By utilizing conjugated $p$ and $q_{\theta }$, the update can be done with a closed form solution.

Ahmadi-Asl et al. (2021) categorized randomized algorithms for the Tucker decomposition and HOSVD into four groups: Random projection Che and Wei (2019); Minster et al. (2020); Zhou et al. (2014); Sun et al. (2021); Wolf (2019); Batselier et al. (2018); Che et al. (2021), Sampling Caiafa and Cichocki (2010); Sun et al. (2009); Oseledets et al. (2008); Saibaba (2016); Tsourakakis (2010); Song et al. (2019); Perros et al. (2015); Traoré et al. (2019), Random least-squares Malik and Becker (2018), and Count sketch Wang et al. (2015); Gittens et al. (2020); Malik and Becker (2018, 2020). Numerical experiments showed that sampling group outperform random projection group in computation time with approximately the same results Ahmadi-Asl et al. (2021). The random projection based algorithms are not suitable for huge tensors stored outside-of-cores since they need to pass the original data tensor multiple times. Randomized pass-efficient algorithms have been introduced to overcome this issue by minimizing the number of passes Halko et al. (2011); Woodruff and others (2014); Boutsidis et al. (2016); Upadhyay (2018), but none utilize the sampling approach for the solution. SketchyCoreTucker is a unique algorithm that combines random projection and sampling to utilize the advantage of both approaches and overcome limitations rooted in using them separately.

Many sampling-based algorithms compute each factor matrix by sampling the "columns" of the corresponding unfolding matrix or equivalently sampling the fibers of the original data tensor. However, our P-SCT samples the "rows" of the unfolding matrix or equivalently sampling the tensor slices. By doing so, P-SCT assures all sampled slices are intersected with each other and progressively updates the approximation of the core from the intersection of all sampled slices.

The main differentiation of our P-SCT is sample efficiency earned by active learning the informativeness distribution utilizing the idea of Thompson sampling. Many sampling-based algorithms use different sampling protocols to minimize the number of samples and computation time, but none use active learning for progressiveness. Sampling based on leverage scores are proposed in Saibaba (2016); Perros et al. (2015); Cheng et al. (2016). These algorithms are not sample efficient due to prior scanning entire original data tensor to compute the leverage scores. CUR decomposition can be considered a sampling technique with a difference that the sampling procedure is performed heuristically instead of randomly Caiafa and Cichocki (2010); Oseledets (2011); Saibaba (2016). It is known that CUR approximations are less accurate than the SVD, and the quality of decomposition quite depends on selection of fibers. The optimal selection has been considered an NP-hard problem, and thus heuristic algorithms have been used for computing suboptimal solutions. Our progressive sketching approach can be applied to CUR decomposition to overcome this issue too. Recently, Song et al. (2019) extended matrix CUR decomposition to tensor CURT decomposition and proposed an randomized algorithm to compute a low-rank CURT approximation to a tensor Song et al. (2019). But this method does not use reinforcement learning to acquire sample efficiency.

Random least-squares and count-sketch groups are a randomized algorithms for ALS-based computation of the best mulitilinear rank approximation such as Higher-Order Orthogonal Iteration (HOOI) De Lathauwer et al. (2000). A randomized least-squares HOOI algorithm is proposed in Malik and Becker (2018), but numerical experiment in Ahmadi-Asl et al. (2021) showed that sampling group have better computation efficiency. Count-sketch technique has been implemented in several algorithms, mostly for CANDECOMP/PARAFAC (CP) decomposition Wang et al. (2015); Gittens et al. (2020); Malik and Becker (2020).

There are other notable researches that share the same objective of approximating tensor decomposition with the highest accuracy and minimal computation. Oh et al. (2018) proposed P-TUCKER, a scalable Tucker decomposition that resolves the computational and memory bottleneck of ALS by row-wise updating factor matrices with parallelization Oh et al. (2018). FREDE is a graph embedding based on matrix sketching that applies Frequent Directions — a deterministic matrix sketching in the row-updates model — on a matrix-factorization interpretation of a neural embedding, VERSE Tsitsulin et al. (2021). Others introduced a randomized CP decomposition algorithms Wang et al. (2015); Gittens et al. (2020); Malik and Becker (2020). The CP decomposition factorizes a tensor into a sum of component rank-one tensors Kolda and Bader (2009), and it is a special case of Tucker decomposition with hyper-diagonal core tensor. Another type of tensor decomposition is Tensor-Train (TT) decomposition which requires iterative matrix SVD computation, and thus SketchyCoreSVD or any randomized matrix SVD algorithm can be applied. Tucker, CP and TT decompositions produces different factors of which best for different applications. This study focuses on Tucker decomposition, and still SketchyCoreTucker and P-SCT have unique ability to learn/provide informativeness distributions of each mode to the following computations, generally the main operation of interest.

P-SCT is for any order $K$ tensor field, so let us assume order 3 tensor $A\in R^{N_1\times N_2\times N_3}$ for explanation purposes, which is huge in size, and its spectrum is decaying fast enough for low rank approximation. Its matricizations along the three dimensions are denoted by $A_{\left(1\right)}\in R^{N_1\times \left(N_2{N}_3\right)}$, $A_{\left(2\right)}\in R^{N_2\times \left(N_3{N}_1\right)}$, $A_{\left(3\right)}\in R^{N_3\times \left(N_1{N}_2\right)}$. Let us denote row space of $A_{\left(k\right)}$ as $Ro{w}_k=\{A_{\left(k\right)}^{(1,:)},\dots ,A_{\left(k\right)}^{(N_k,:)}\}$. Then, $Ro{w}_k$ where $k=1,2,3$ are bandit arms of TS. As an action, we choose a $Ro{w}_k$ and sample from it without replacement. Note that a row vector of a matrix unfolding is a slice matrix in original geometric dimension of order 3 tensor. Thus, P-SCT samples order $K-1$ slice tensors to sketch an order $K$ tensor field.

The amount of information contained in the samples ($\equiv$ order $K-1$ slice tensors) is the reward of the action. For Tucker decomposition, the amount of information can be measured as accumulated variation in data fibers along with $K-1$ dimensions. For instance, let us assume there is a sample $A_{\left(1\right)}^{({\delta }_1,:)}$ which is equivalent to the slice matrix $A^{({\delta }_1,:,:)}$. Both rows and columns (i.e., data fibers) of $A^{({\delta }_1,:,:)}$ are used in Tucker decomposition. We define the Sum of Absolute Differences ($SAD$) to capture the accumulated variation in these data fibers. $SAD$ of $A_{\left(1\right)}^{({\delta }_1,:)}\equiv A^{({\delta }_1,:,:)}$ is

where the absolute differences of data in data fibers of $A^{({\delta }_1,:,:)}$ are accumulated and normalized by the size of slice matrix. The normalization is required because what we want to measure is the amount of information per usage of the memory for the sample efficient Tucker decomposition. $SAD$ of samples from other $Ro{w}_k$ can be calculated in the same manner.

Dirichlet distribution is used to balance exploitation and exploration of $Ro{w}_k$. We update the concentration parameters of Dirichlet distribution by the entropy of normalized $SAD$ of samples from $Ro{w}_k$. By doing so, P-SCT samples more from $Ro{w}_k$ of which information is uniformly distributed rather than concentrated in narrow region. It effectively reduces the mutual information between samples.

At first, we assume concentration parameters of the Dirichlet distribution as unity to make the distribution equivalent to the uniform distribution $\alpha =\{{\alpha }_1,{\alpha }_2,{\alpha }_3\}=\{1,1,1\}$. Then, sampling ratios $p=\{p_1,p_2,p_3\}$ are drawn from the Dirich $\left(\alpha \right)$. By multipling the batch-size per round $N_{batch}$ to $p$, we get the number of samples $n=\{n_1,n_2,n_3\}$ that we newly sample from each $Ro{w}_k$.

We utilize $SAD$ once again here for the sampling. Rather than randomly sample, we assign weights by $SAD$. As a result, P-SCT greedily search samples that contain highest amount of information. Combined two search schemes — 1) TS with Dirichlet distribution that balances $Ro{w}_k$ selection minimizing the mutual information between samples and 2) $SAD$-weighted sampling that maximizes the use of information of samples — is the core of P-SCT algorithm that learns optimal subsampling policy for streaming tensor sketching.

At the first round, we know nothing about the input tensor. Thus, weights are set as unity that can be denoted by $W=\{w_1,w_2,w_3\},w_k=\{1,...,1\},n\left(w_k\right)=N_k.$ where $w_k$ is weights for $Ro{w}_k$ and $n\left(w_k\right)$ denotes the length of weights. Let us define sampled indices $\Delta =\{{\Delta }_1,{\Delta }_2,{\Delta }_3\}$ and unsampled indices $\Omega =\{{\Omega }_1,{\Omega }_2,{\Omega }_3\}$ and an allowed total number of samples $N_{allow}$. Then, the main steps of P-SCT are the following. In the algorithm, $\lfloor \cdot \rfloor$ represents floor function.

1. Sample from $Ro{w}_k$

   • Draw sampling ratios $p=\{p_1,p_2,p_3\}$ from Dirich $\left(\alpha \right)$.

   • For $Ro{w}_k$, sample $n_k=\lfloor p_k{N}_{batch}\rfloor$ indices from unsampled indices ${\Omega }_k$ following weights $w_k$. Denote the indices of the sampled rows to be ${\Delta }_k$.

2. Compute $SAD$ and update TS parameters

   • Compute $SAD$ of $A_{\left(k\right)}^{({\Delta }_k,:)}$. Denote $SAD$ of $A_{\left(k\right)}^{({\delta }_k,:)},{\delta }_k\in {\Delta }_k$ as $SAD\left({\delta }_k\right)$.

   • Compute $SAD$ entropy $H_k=-{\sum }_{{\delta }_k\in {\Delta }_k}\frac{SAD\left({\delta }_k\right)}{{\sum }_{{\delta }_k}SAD\left({\delta }_k\right)}\log \left(\frac{SAD\left({\delta }_k\right)}{{\sum }_{{\delta }_k}SAD\left({\delta }_k\right)}\right).$

   • Update the concentration parameter ${\alpha }_k\leftarrow H_k$.

   • Update ${\Omega }_k\leftarrow {\Omega }_k\setminus {\Delta }_k$.

   • Linearly interpolate $SAD$ of unsampled rows $SAD\left({\Omega }_k\right)$ based on the index differences using $SAD\left({\Delta }_k\right)$.

   • Set $SAD\left({\Delta }_k\right)=0$ to avoid sampling duplicates, and update $w_k$ by normalized $SAD$.

   • Repeat Steps 1 and 2 until $n\left({\Delta }_1\right)+n\left({\Delta }_2\right)+n\left({\Delta }_3\right)>N_{allow}$.

3. Final SketchyCoreTucker

   • Compute rank-$r$ approximation using SketchyCoreTucker (see Appendix).

   • Instead of uniformly sampling rows of $A_{\left(1\right)}$,$A_{\left(2\right)}$, $A_{\left(3\right)}$, use ${\Delta }_1$, ${\Delta }_2$, ${\Delta }_3$ from Step 2.

   • For the columns of $A_{\left(1\right)}$,$A_{\left(2\right)}$, $A_{\left(3\right)}$, use corresponding fibers crossing the intersection $A^{({\Delta }_1,{\Delta }_2,{\Delta }_3)}$.

   • Dimension reduction parameters $k$ and $s$ can be decided as

     $$k=\lfloor r+\frac{1}{3}(n-r)\rfloor ,s=\lfloor r+\frac{2}{3}(n-r)\rfloor ,$$

     where $n=\left\{n\right({\Delta }_1),n({\Delta }_2),n({\Delta }_3\left)\right\}$.

Permutation agnostic P-SCT is achieved by considering the entropy of sample variance of latent embeddings based on J-L projected and $QR$ orthogonalized fibers of each tensor mode matrix unfolding. The entropy is used for the concentration/balancing parameter of Thompson sampling with a Dirichlet prior. P-SCT-Permute also substitutes $SAD$-based adoptive sampling with uniform sampling to avoid spatial interpolation. Steps 1 and 3 of the P-SCT algorithm are the same, and only Step 2 is different in P-SCT-Permute.

1. Sample from $Ro{w}_k$

2. Compute sample variance and update TS parameters

   • Denote ${\Theta }_1$,${\Theta }_2$, and ${\Theta }_3$ as the corresponding columns of $A_{\left(1\right)}$,$A_{\left(2\right)}$, $A_{\left(3\right)}$ that crossing the intersection $A^{({\Delta }_1,{\Delta }_2,{\Delta }_3)}$.

   • Compute latent embeddings by the map ${\Omega }_1\in R^{k_1\times n\left({\Theta }_1\right)}$ applied to $A_{\left(1\right)}^{(:,{\Theta }_1)}$ from the right and QR decomposition, $Q_1{R}_1=Y_1=A_{\left(1\right)}^{(:,{\Theta }_1)}{\Omega }_1^{\ast }$, and $Q_2$ and $Q_3$ accordingly.

   • Compute entropy of sample variance of $Q_k$

     $$H_k={\text{Entropy}}_j\left(\frac{1}{N_k-1}\sum _{i=1}^{N_k}{(Q_k^{(i,j)}-\frac{1}{N_k}\sum _{i=1}^{N_k}{Q}_k^{(i,j)})}^2\right)$$

   • Update the concentration parameter ${\alpha }_k\leftarrow H_k$.

   • Update ${\Omega }_k\leftarrow {\Omega }_k\setminus {\Delta }_k$.

   • Repeat Steps 1 and 2 until $n\left({\Delta }_1\right)+n\left({\Delta }_2\right)+n\left({\Delta }_3\right)>N_{allow}$.

3. Final SketchyCoreTucker

Table 1 shows the complexity of R-SCT, P-SCT, P-SCT-Permute and one-shot algorithms (RP-HOSVD,R-STHOSVD,R-PET,R-ST, and R-HOID Ahmadi-Asl et al. (2021)) and an ALS-based algorithm (RP-HOOI Ibid.) for the order 3 tensor case. For simple comparison, we assume that $N_1=N_2=N_3=N$, $r_1=r_2=r_3=r$, $m_1=m_2=m_3=m$, $s_1=s_2=s_3=s$, and $k_1=k_2=k_3=k$. Also, it is assumed that the same $b=N_{batch}/3$ rows are sampled from each $Ro{w}_k$ every iteration for P-SCT and P-SCT-Permute. Thus, a product of total number of iteration $I$ and $b$ is $m$. Also, note that $N>m>r>s>k$.

A collection of time-varying, 2D slice stacks of Cardiac MRI is tested. $A\in R^{256\times 176\times 160}$ consists of 160 time snapshots 2D images, each of size 256 $\times$ 176. Based on the scree plots shown in Appendix, we choose $r=(20,20,5)$. We also choose $N_{allow}=300$. Table 2 and Figure 1 show that P-SCT and P-SCT-Permute produce good performance as other one-shot algorithms, although they use less than half of the data tensor. Figures 2 and 3 compares progressive performance to show the sample efficiency of P-SCT and P-SCT-Permute. P-SCT-Permute’s learning curve is slightly decaying faster, but its computation is slower than P-SCT.

NARR air temperature dataset is spatiotemporal series covering the North American continent Mesinger et al. (2006). Its grid resolution is 349 $\times$ 277, which is approximately 0.3 degrees (32km) resolution at the lowest latitude. It provides 3 hourly series from 1979/01/01 to the present. Air temperature is also collected from 29 pressure levels (i.e., 29 different altitudes). Therefore, it can be represented as a order 5 data tensor. Due to the limitation on the computational resource, only one altitude at 500 hpa pressure level is selected. Also, 1 month-long series on 2020 October is used for the experiment. Spatial coverage is also clipped to discard null values around boundaries. As a result, $A\in R^{(248,252,336)}$. Based on scree plots shown in Appendix, we choose $r=(50,10,20)$. $N_{allow}$ is assumed to be 300. Figures 5 and 6 compares progressive performance showing the sample efficiency of P-SCT and P-SCT-Permute. Interestingly, P-SCT-Permute performs better than P-SCT for the air temperature dataset. We can also see that the spatiotemporal structure of air temperature has been revealed by PS-SCT-Permute fastest.

The computation time of P-SCT and P-SCT-Permute is longer than R-SCT. This delay is inevitable since the progressive sketching approach adds more computations. The delay will be greater for higher order tensor or bigger sized data tensor. However, considering a combinatorial number of possible sampling ratio parameters, the amount of computation time we trade off for getting optimal sample efficient low rank approximation is sufficiently negligible. The optimally balanced performance of the progressive sketching approach is well addressed in Figure 1 that P-SCT maintains a good accuracy and computation time compared to fast-but-erroneous (RP-HOSVD, R-ST), and accurate-but-slow (RP-STHOSVD, R-HOID) one-shot algorithms. The computation time of PS-SCT-Permute should be improved in the future.