MTC: Multiresolution Tensor Completion from Partial and Coarse Observations

One of the common compression methods for tensors is CP decomposition (Hitchcock, 1927; Carroll and Chang, 1970), which approximates a tensor by multiple rank-one components. For example, let $\mathcal{X}\in\mathbb{R}^{I_{1}\times I_{2}\times I_{3}}$ be an arbitrary $3$-mode tensor of CP rank $R$, it can be expressed exactly by three factor matrices $\mathbf{U}\in\mathbb{R}^{I_{1}\times R},\mathbf{V}\in\mathbb{R}^{I_{2}\times R},\mathbf{W}\in\mathbb{R}^{I_{3}\times R}$ as

where $\circ$ is the vector outer product, and the double bracket notation $\llbracket\,\cdot,\cdot,\cdot\rrbracket$ is the Kruskal operator.

For the same tensor $\mathcal{X}$, a coarse observation is given by conducting a tensor mode product with an aggregation matrix (for example, on the first mode), $\mathbf{P}_{1}\in\mathbb{R}^{J_{1}\times I_{1}}$, where $J_{1}<I_{1}$,

where we use superscripts, e.g., $\mathcal{C}^{(1)}$, to denote the aggregated tensor. Other tensor background may be found in Appendix. We organize the symbols in Table 1.

To address the above objective, we consider to successively approximate the low-rank completion problem by coupled tensor decomposition. We first introduce two transformations and reformulate the objective into a coupled decomposition form:

• •

  Parameter replacement. We replace $\mathbf{Q}_{1}=\mathbf{P}_{1}\mathbf{U}$ and $\mathbf{Q}_{2}=\mathbf{P}_{2}\mathbf{V}$ as two new auxiliary factor matrices  and add extra constrains later since $\mathbf{Q}_{2}$ and $\mathbf{V}$ are related by a known $\mathbf{P}_{2}$ (but $\mathbf{P}_{1}$ is unknown, we does not consider the constraint for $\mathbf{Q}_{1}$ and $\mathbf{U}$).

• •

  Interim Tensor. For the first term in Equation (3.3), we construct an interim tensor $\tilde{\mathcal{X}}$ from an expectation-maximization (EM) (Acar et al., 2011a) approach,

  		$\displaystyle\mathcal{M}*(\mathcal{X}-\llbracket\mathbf{U,V,W}\rrbracket)$
  	$\displaystyle=$	$\displaystyle~{}\mathcal{M}*\mathcal{X}+(1-\mathcal{M})*\llbracket\mathbf{U,V,W}\rrbracket-\llbracket\mathbf{U,V,W}\rrbracket$
  (5)		$\displaystyle\approx$	$\displaystyle\underbrace{\left[\mathcal{M}*\mathcal{X}+(1-\mathcal{M})*\llbracket\mathbf{U}^{k},\mathbf{V}^{k},\mathbf{W}^{k}\rrbracket\right]}_{\mbox{the interim tensor $\tilde{\mathcal{X}}$ for next iteration}}-~{}\llbracket\mathbf{U},\mathbf{V},\mathbf{W}\rrbracket,$

  where $\mathbf{U}^{k},\mathbf{V}^{k},\mathbf{W}^{k}$ are the learned factors at the $k$-th iteration. The reconstruction $\mathcal{R}=\llbracket\mathbf{U}^{k},\mathbf{V}^{k},\mathbf{W}^{k}\rrbracket$ is expensive for large tensors. However, we do not need to compute it explicitly. In Section 5, we introduce a way to use the implicit form.

As a solution, we adopt a multiresolution strategy, as in Figure 1, which leverages a hierarchy of resolution to perform the coupled optimization effectively. The multiresolution factorization part (in Section 4) serves as the outer loop of our solution. At each resolution, we derive a joint normal equation for each factor and develop an efficient ALS-based solver (in Section 5) to handle multiple linear systems simultaneously.

We define

to be the index set of the date aspect at $d_{t+1}$-resolution. We sample at a regular interval as

yielding a $d_{t}$-resolution index set of size $\lceil\frac{I_{3(d_{t+1})}}{2}\rceil$.

For example, $\mathcal{C}^{(1)}_{(d_{t+1})}$ includes the date aspect, thus it will be subsampled using the index set $\mathcal{S}_{(d_{t},~{}date)}$.

Factor $\mathbf{W}_{(d_{t})}\in\mathbb{R}^{I_{3(d_{t})}\times R}$ corresponds to the date aspect at $d_{t}$. The initialization of $d_{t+1}$ factor $\tilde{\mathbf{W}}_{(d_{t+1})}\in\mathbb{R}^{I_{3(d_{t+1})}\times R}$ is given by

if $I_{3(d_{t+1})}$ is an even number, then

As shown in Figure 2, from $\mathcal{C}^{(1)}_{(d_{t+1})}$ to $\mathcal{C}^{(1)}_{(d_{t})}$, only the “dark” slabs are selected for the $d_{t}$-resolution. Tensor slabs of $\mathcal{C}^{(1)}_{(d_{t})}$ along the date aspect are controlled by $\mathbf{W}_{(d_{t})}$, where one row corresponds to one slab. After solving the $d_{t}$ problem, the solution factor $\mathbf{W}_{(d_{t})}$ provides an estimate for those “dark” slabs. To also estimate the “light” slabs, we assume that it could be approximated by the neighboring “dark” slabs. The actual interpolation follows Equation (16), (17), (18), where we average neighboring rows.

We define

to be the index set of ICD-10 aspect at $d_{t+1}$-resolution.

We order the count vector and find the indices corresponding to the largest half (composing the dense part of the tensor, while the unselected half is the sparse part), denoted as $\mathcal{S}_{(d_{t},~{}ICD-10)}$, as the index set at $d_{t}$, so that

The tensors and aggregation matrices that include ICD-10 aspect, e.g., $\mathcal{C}^{(1)}_{(d_{t+1})}$, are subsampled along this mode.

$\mathbf{V}_{(d_{t})}\in\mathbb{R}^{{I_{2(d_{t})}}\times R}$ is the corresponding solution factor for ICD-10 at $d_{t}$. We get the $d_{t+1}$ initialization $\tilde{\mathbf{V}}_{(d_{t+1})}\in\mathbb{R}^{{I_{2(d_{t+1})}}\times R}$ as follows. The rows of $\tilde{\mathbf{V}}_{(d_{t+1})}$ are direct copies from the rows of $\mathbf{V}_{(d_{t})}$ if the corresponding indices are selected in $\mathcal{S}_{(d_{t},~{}ICD-10)}$, otherwise, we fill the rows with random entries (in the paper, we use i.i.d. samples from $[-1,1]$, and users might customize their configuration based on the applications).

As for the categorical mode (in Figure 3), the rationality of our strategy comes from two underlying heuristics: (i) the dense half of the indices is likely to account for a large volume of the tensor, such that after the optimization at $d_{t}$, we already recover a dominant part at $d_{t+1}$; (ii) the selected indices (slabs) preserve the dense half of the tensor, while from randomized algorithm perspective, a slab may be regarded as a training sample, composited by the other factor matrices. Preserving those dense slabs would be beneficial for the estimation of other factors.

To solve a low-rank CP approximation of a given tensor, ALS is one of the most popular algorithms (Sidiropoulos et al., 2017). It has been applied for both decomposition (Larsen and Kolda, 2020; Sidiropoulos et al., 2017) and completion (Song et al., 2019; Karlsson et al., 2016). However, our approach differs by considering a reconstructed (interim) tensor at each step in a coupled decomposition setting. Below is the standard ALS algorithm.

Given a tensor $\mathcal{X}\in\mathbb{R}^{I_{1}\times I_{2}\times I_{3}}$, to find a rank-$R$ CP decomposition with factors $\mathbf{U}\in\mathbb{R}^{I_{1}\times R}$, $\mathbf{V}\in\mathbb{R}^{I_{2}\times R}$, $\mathbf{W}\in\mathbb{R}^{I_{3}\times R}$, the following loss function is minimized,

ALS algorithm optimally solves for one of the factor matrices while keeping the others fixed and then cycles through all of the factor matrices. A sub-iteration of ALS for $\mathbf{U}$ is executed by solving

where $\mathbf{V}\odot\mathbf{W}\in\mathbb{R}^{I_{2}I_{3}\times R}$ is the Khatri-Rao product and $\mathbf{X}_{1}\in\mathbb{R}^{I_{1}\times I_{2}I_{3}}$ is the unfolded tensor along the first mode.

In our case, at each resolution $d_{t}$, we have a coupled tensor decomposition problem as defined in equation ((2)). Similar to standard ALS, if we solve for one factor matrix while keeping all the other unknowns fixed, the problem reduces to a quadratic problem which can be solved optimally by coupling normal equations with respect to the factor at each granularity.

To reduce clutter, we elaborate on the full resolution form of Equation (6). Specifically, we solve a least squares problem with the following equations.

• •

  For $\mathbf{U}$, equations are derived using $\tilde{\mathbf{X}}_{1}$ and $\mathbf{C}^{(2)}_{1}$, which are the unfoldings of $\tilde{\mathcal{X}}$ and $\mathcal{C}^{(2)}$ along the first mode. Thus, we have the following two equations,

  (24)		$\displaystyle(\mathbf{V}\odot\mathbf{W})\mathbf{U}^{\top}$	$\displaystyle=\tilde{\mathbf{X}}_{1}^{\top},$
  (25)		$\displaystyle\lambda_{2}(\mathbf{Q}_{2}\odot\mathbf{W})\mathbf{U}^{\top}$	$\displaystyle=\lambda_{2}\mathbf{C}_{1}^{(2)\top}.$

• •

  Similarly, for $\mathbf{V}$, equations are derived using $\tilde{\mathbf{X}}_{2}$ and $\mathbf{C}^{(1)}_{2}$.

• •

  For $\mathbf{W}$, equations are derived using $\tilde{\mathbf{X}}_{3}$, $\mathbf{C}^{(1)}_{3}$ and $\mathbf{C}^{(2)}_{3}$.

• •

  For $\mathbf{Q}_{1}$, equations are derived using $\mathbf{C}^{(1)}_{1}$.

• •

  For $\mathbf{Q}_{2}$, equations are derived using $\mathbf{C}^{(2)}_{2}$ only. However, $\mathbf{Q}_{2}$ and $\mathbf{V}$ satisfy the constraint $\mathbf{Q}_{2}=\mathbf{P_{2}V}$.

Considering the constraint issue, e.g., between $\mathbf{V}$ and $\mathbf{Q}_{2}$, we use the equations to update one of them (the one corresponds to fine granular aspect) and later update the other one. For our example, $\mathbf{V}$ is updated by the linear equations, while $\mathbf{Q}_{2}$ is updated by the constraint $\mathbf{Q}_{2}=\mathbf{P}_{2}\mathbf{V}$ and new $\mathbf{V}$, instead of using linear equations²²2Note that, since these two factors are coupled, we could have transformed the $\mathbf{Q}_{2}$ equations to supplement $\mathbf{V}$ equations by pseudo inverse (i.e. inverse of $\mathbf{P}_{2}^{T}\mathbf{P}_{2}$), such that $\mathbf{V}$ would associate with three linear equations. However, the non-uniqueness of $\mathbf{P}_{2}^{T}\mathbf{P}_{2}$ makes the optimization unstable.. The optimization flow follows the ALS routine, where during each iteration, we sequentially update $\mathbf{W,Q_{1},U}$, $\mathbf{V}$ by the associated linear equations, and then manually update a new $\mathbf{Q_{2}}$.

$\mathbf{W,Q_{1},U}$, $\mathbf{V}$ are solved by using joint normal equations. For each of them, we consider stacking multiple linear systems and formulating joint normal equations with respect to each factor for the objective defined in equation (6). Using $\mathbf{U}$ as the example, the normal equations are given as

where the components are weighted accordingly. The formulation could be further simplified as

This formulation is amenable to large scale problems as the right hand sides can exploit sparsity and the use of the matricized tensor times Khatri-Rao product (MTTKRP) kernel with each tensor. To avoid dense MTTKRP with $\tilde{\mathcal{X}}$, the implicit form given in Equation (7) can be used to efficiently calculate the MTTKRP by

where $\mathcal{R}=\llbracket\mathbf{U}^{k},\mathbf{V}^{k},\mathbf{W}^{k}\rrbracket$. The above simplification requires sparse MTTKRP and a sparse reconstruction tensor $\mathcal{M}*\mathcal{R}$, for which efficient programming abstractions are available (Singh et al., 2021).

To solve the joint normal equations, we employ a two-stage approach.

• •

  Stage 1: We start from the initialization given by the low resolution and apply the weighted Jacobi algorithm for $N$ (e.g., 5) iterations, which smooths out the error efficiently.

• •

  Stage 2: Starting from the $(N+1)-th$ iteration, we apply Cholesky decomposition to the left-hand side (e.g., Equation (5.2.2)) and decompose the joint normal equations into solving two triangular linear systems, which could generate more stable solutions. Finally, the triangular systems are solved exactly.

In practice, we find that the combination of Stage 1 and Stage 2 works better than directly applying Stage 2 in light of both convergence speed and final converged results, especially when the number of observed entries is small. We perform ablation studies for this part.

The time and space complexity of MTC is asymptotically equal to that of applying standard CP-ALS on original tensor $\mathcal{X}$, which is $O(I_{1}I_{2}I_{3}R)$ and $O(I_{1}I_{2}I_{3})$, respectively.

Two COVID-related databases are considered in our evaluation:

• •

  The health tensor (HT) is a proprietary real-world dataset, constructed from IQVIA medical claim database, including a location by disease by time tensor, counting the COVID related disease cases in the US from Mar. to Aug. 2020 during the pandemic.

• •

  (ii) The Google COVID-19 Symptoms Search (GCSS)³³3https://pair-code.github.io/covid19_symptom_dataset/?country=GB is a public dataset, which shows aggregated trends in Google searches for disease-related health symptoms, signs, and conditions, in format location identifier by keyword by time.

For HT data, we use the most popular 1,200 zipcodes and 1,000 ICD-10 codes as the first two modes and use the first five months as the third mode (while the last month data is used for the prediction task), which creates $\mathcal{X}$. We aggregate the first mode by county and the second mode by CCS code separately to obtain two aggregated tensors, $\mathcal{C}^{(1)},\mathcal{C}^{(2)}$. We randomly sample 5% elements from the tensor as the partial observation, $\mathcal{X}*\mathcal{M}$. For the GCSS data, we use all 2,727 location identifiers and 422 search keywords as the first two modes and a complete 2020 year span as the third time mode. To create two coarse viewed tensors, we aggregate on the state level in the first mode and week level for the third mode, separately. Assume 5% of the elements are observed in GCSS. For both datasets, the first two modes are assumed categorical, while the last time mode is continuous.

In this experiment, we assume all the aggregation matrices are unknown. In the Appendix, we use synthetic data and GCSS to show that our model could achieve OracleCPD level accuracy if some aggregation $\mathbf{P}$ is given. Basic data statistics are listed in Table 2.

We include the following comparison models from different perspectives: (i) reference models: CPC-ALS (Karlsson et al., 2016): state of the art tensor completion model, only using partial observed data, OracleCPD (CP decomposition of the original complete tensor with ALS); (ii) related baselines (they are gradient based): Block gradient descent (BGD) (Xu and Yin, 2013), B-PREMA (Almutairi et al., 2021), CMTF-OPT (Acar et al., 2011b). They use both partial and coarse information; (iii) initialization models: MRTL (Park et al., 2020), TendiB (Almutairi et al., 2020), Higher-order SVD (HOSVD); (iv) our variants: $\texttt{MTC}_{multi-}$ removes multiresolution module; $\texttt{MTC}_{stage1-}$ removes stage 1 (Jacobi) in the solver; $\texttt{MTC}_{both-}$ removes both multiresolution part and stage 1. More details are in the Appendix.

Evaluation metrics include: (i) Percentage of Fitness (PoF) (Song et al., 2019), suppose the target tensor is $\mathcal{X}$ and the low-rank reconstruction is $\mathcal{R}=\llbracket\mathbf{U},\mathbf{V},\mathbf{W}\rrbracket$, the relative standard error (RSE) and PoF are defined by

(ii) CPU Time; (iii) Peak Memory.

We also consider (2) future tensor prediction, which belongs to a broader spatio-temporal prediction domain (Wang et al., 2020). We evaluate related to tensor-based baselines, while a thorough comparison with other models is beyond our scope. Specifically, we use Gaussian Process (GP) with radial basis function and white noise kernel along the date aspect to estimate a date factor matrix for the next month,

Then, the future disease tensor could be estimated by $\mathcal{R}_{future}=\llbracket\mathbf{U},\mathbf{V},\mathbf{W}_{future}\rrbracket$. We use the next-month real disease tensor to evaluate the prediction results, and again, we use PoF as the metric.

For HT dataset, we use CP rank $R=50$, $20$ iterations at each low resolution, and $N=5$ in solver stage 1, within which $5$ Jacobi rounds are performaned per iteration. For GCSS, we use $R=20$, $10$ iterations at each low resolution, $N=5$, and $10$ Jacobi rounds per iteration. By default, we set $200$ iterations at the finest resolution for both datasets, which ensures the convergence. The parameter $\lambda_{i}$ is set to be $\lambda=e^{-\frac{i}{20}}$ at the $i$-th iteration, which varies the focus from coarse information to fine-granular information gradually. We publish our dataset and codes in a repository⁴⁴4https://github.com/ycq091044/MTC.