Smooth and probabilistic PARAFAC model with auxiliary covariates

In the study of multivariate longitudinal data in economics, researchers have combined tensor decomposition with auto-cross-covariance estimation and autoregressive models (Fan et al., 2008; Lam et al., 2011; Fan et al., 2011; Bai and Wang, 2016; Wang et al., 2019, 2021). These approaches do not work well with highly sparse data or do not scale well with the feature dimensions, which are important for working with medical data. Functional PCA (Besse and Ramsay, 1986; Yao et al., 2005) is often used for modeling sparse longitudinal data in the matrix-form. It utilizes the smoothness of time trajectories to handle sparsity in the longitudinal observations, and estimates the eigenvectors and factor scores under this smoothness assumption. Yokota et al. (2016) and Imaizumi and Hayashi (2017) introduced smoothness to tensor decomposition models, and estimated the model parameters by iteratively solving penalized regression problems. The methods above don’t consider the auxiliary covariates $\bm{Z}$.

It has been previously discussed that including $\bm{Z}$ could potentially improve our estimation. Li et al. (2016) proposed SupSFPC (supervised sparse and functional principal component) and observed that the auxiliary covariates improve the signal estimation quality in the matrix setting for modeling multivariate longitudinal observations. Lock and Li (2018) proposed SupCP which performs supervised multiway factorization model with complete observation and employs a probabilistic tensor model (Tipping and Bishop, 1999; Mnih and Salakhutdinov, 2007; Hinrich and Mørup, 2019). Although an extension to sparse tensor is straightforward, SupCP does not model the smoothness and can be much more affected by severe missingness along the time dimension.

SPACO can be viewed as an extension of functional PCA and SupSFPC to the setting of three-way tensor decomposition (Acar and Yener, 2008; Sidiropoulos et al., 2017) using the parallel factor (PARAFAC) model (Harshman and Lundy, 1994; Carroll et al., 1980). It uses a probabilistic model and jointly models the smooth longitudinal data with potentially high-dimensional non-temporal covariates $\bm{Z}$. We refer to the SPACO model as SPACO- when no auxiliary covariate $\bm{Z}$ is available. SPACO- itself is an attractive alternative to existing tensor decomposition implementations with probabilistic modeling, smoothness regularization, and automatic parameter tuning. In our simulations, we compare SPACO with SPACO to demonstrate the gain from utilizing $\bm{Z}$.

Let $\bm{X}\in\mathbb{R}^{I\times T\times J}$ be a tensor for some sparse multivariate longitudinal observations, where $I$ is the number of subjects, $J$ is the number of features, and $T$ is the number of total unique time points. For any matrix $A$, we let $A_{i:}/A_{:i}$ denote its $i^{th}$ row/column, and often write $A_{:i}$ as $A_{i}$ for the $i^{th}$ column for convenience. Let $\bm{X}_{I}\coloneqq\begin{pmatrix}\bm{X}_{:,:,1}&\cdots&\bm{X}_{:,:,J}\end{pmatrix}\in\mathbb{R}^{I\times(TJ)}$, $\bm{X}_{T}\coloneqq\begin{pmatrix}\bm{X}_{:,:,1}^{\top}&\cdots&\bm{X}_{:,:,J}^{\top}\end{pmatrix}\in\mathbb{R}^{T\times(IJ)}$, $\bm{X}_{J}\coloneqq\begin{pmatrix}\bm{X}^{\top}_{:,1,:}&\cdots&\bm{X}_{:,T,:}^{\top}\end{pmatrix}\in\mathbb{R}^{J\times(IT)}$ be the matrix unfolding of $\bm{X}$ in the subject/feature/time dimension respectively. We also define:

We assume $\bm{X}$ to be a noisy realization of an underlying signal array $\mathbb{F}=\sum_{k=1}^{K}\bm{U}_{k}\circledcirc\bm{\Phi}_{k}\circledcirc\bm{V}_{k}$. We stack $\bm{U}_{k}/\bm{\Phi}_{k}/\bm{V}_{k}$ as the columns of $\bm{U}/\bm{\Phi}/\bm{V}$, denote the rows of $\bm{U}/\bm{\Phi}/\bm{V}$ by $\bm{u}_{i}/\bm{\phi}_{t}/\bm{v}_{j}$, and their entries by $u_{ik}/\phi_{tk}/v_{jk}$. We let $x_{itj}$ denote the $(i,t,j)$-entry of $\bm{X}$. Then,

where $\Lambda_{f}=\mathrm{diag}\{s_{1}^{2},\ldots,s_{K}^{2}\}$ is a $K\times K$ diagonal covariance matrix. Even though $\bm{X}$ is often of high rank, we consider the scenario where the rank $K$ of $\mathbb{F}$ is small.

Without covariates, we set the prior mean parameter ${\bm{\eta}}_{i}={\bm{0}}$. If we are interested in explaining the heterogeneity in ${\bm{\eta}}_{i}$ across subjects with auxiliary covariates $\bm{Z}\in\mathbb{R}^{I\times q}$, then we may model ${\bm{\eta}}_{i}$ as a function of ${\bf z}_{i}\coloneqq\bm{Z}_{i,:}$. Here, we consider a linear model $\eta_{ik}={\bf z}_{i}^{\top}\bm{\beta}_{k},\quad\forall k=1,\ldots,K$. To avoid confusion, we will always call $\bm{X}$ the “features”, and $\bm{Z}$ the “covariates” or “variables”.

We refer to $\bm{U}$ as the subject scores, which characterize differences across subjects and are latent variables. We refer to $\bm{V}$ as the feature loadings, which reveal the composition of the factors using the original features and could assist the downstream interpretation. Finally, $\bm{\Phi}$ is referred to as the time trajectories, which can be interpreted as function values sampled from some underlying smooth functions $\phi_{k}(t)$ at a set of discrete-time points, e.g., $\bm{\Phi}_{k}=(\phi_{k}(t_{1}),\ldots,\phi_{k}(t_{T}))$.

Recalling that $\bm{X}_{I}\in\mathbb{R}^{I\times(TJ)}$ is the unfolding of $\bm{X}$ in the subject direction, we write $\vec{i}$ for the indices of observed values in the $i^{\text{th}}$ row of $\bm{X}_{I}$, and $\bm{X}_{I,\vec{i}}$ for the vector of these observed values. Each such observed value $x_{itj}$ has noise variance $\sigma_{j}^{2}$, and we write $\Lambda_{\vec{i}}$ to represent the diagonal covariance matrix with diagonal values $\sigma_{j}^{2}$ being the corresponding variances for $\varepsilon_{ijt}$ at indices in $\vec{i}$. Similarly, we define $\{\vec{t},\;\bm{X}_{T,\vec{t}},\;\Lambda_{\vec{t}}\}$ for the unfolding $\bm{X}_{T}\in\mathbb{R}^{T\times(IJ)}$, and $\{\vec{j},\;\bm{X}_{J,\vec{j}},\;\Lambda_{\vec{j}}\}$ for the observed indices, the associated observed vector and diagonal covariance matrix for the $j^{th}$ row in $\bm{X}_{J}\in\mathbb{R}^{J\times(IT)}$. We set $\Theta=\{\bm{V},\bm{\Phi},{\bm{\beta}},\left(\sigma^{2}_{j},j=1,\ldots,J\right),\left(s^{2}_{k},k=1,\ldots,K\right)\}$ to denote all model parameters. Set $\mathbb{H}=(\bm{V}\odot\bm{\Phi})$ and ${\bf f}_{i}=\bm{X}_{I,\vec{i}}-\mathbb{H}_{\vec{i}}{\bm{\eta}}_{i}$. If $\bm{U}$ is observed, the complete data log-likelihood is

Set $\tilde{\Sigma}_{i}=\Lambda_{\vec{i}}+\mathbb{H}_{\vec{i}}\Lambda_{f}\mathbb{H}_{\vec{i}}^{\top}$. The marginalized log likelihood integrating out the randomness in $\bm{U}$ enjoys a closed form (Lock and Li, 2018):

Set $\Sigma_{i}=(\mathbb{H}_{\vec{i}}^{\top}\Lambda_{\vec{i}}^{-1}\mathbb{H}_{\vec{i}}+\Lambda_{f}^{-1})^{-1}$. We can also equivalent express the marginal likelihood as below.

We use the form in eq. (4) to derive the updating formulas and criteria for rank selection since it does not involve the inverse of a large non-diagonal matrix.

Model parameters in eq. (3) or eq. (4) are not identifiable due to (1) parameters rescaling from $(\bm{\Phi}_{k},\bm{V}_{k},\bm{\beta}_{k},s_{k}^{2})$ to $(c_{1}\bm{\Phi}_{k},c_{2}\bm{V}_{k},c_{3}\bm{\beta}_{k},c_{3}^{2}s_{k}^{2})$ for any $c_{1}c_{2}c_{3}=1$, and (2) reordering of different component $k$ for $k=1,\ldots,K$. More discussions of the model identifiability can be found in Lock and Li (2018). Hence, adopting similar rules from Lock and Li (2018), we require

• (C.1)

  $\|\bm{V}_{k}\|_{2}^{2}=1$, $\|\bm{\Phi}_{k}\|_{2}^{2}=T$.

• (C.2)

  The latent components are in decreasing order based on their overall variances $\Lambda_{f,kk}+\|\bm{Z}\bm{\beta}_{k}\|_{2}^{2}/I$, and the first non-zero entries in $\bm{V}_{k}$ and $\bm{\Phi}_{k}$ to be positive, e.g., $v_{k1}>0$ and $\phi_{k1}>0$ if they are non-zero.

To deal with the sparse sampling along the time dimension and take into consideration that features are often smooth over time in practice, we assume that the time component $\bm{\Phi}_{k}$ is sampled from a slowly varying trajectory function $\phi_{k}(t)$, and encourage smoothness of $\phi_{k}(t)$ by directly penalizing the function values via a penalty term $\sum_{k}\lambda_{1k}\bm{\Phi}_{k}^{\top}\Omega\bm{\Phi}_{k}$. This paper considers a Laplacian smoothing (Sorkine et al., 2004) with a weighted adjacency matrix $\Gamma$. Let $\mathcal{T}(t)$ represent the associated time for $\bm{X}_{.,t.,}$. We define $\Omega$ and $\Gamma$ as

Practitioners may choose other forms for $\Omega$. If practitioners want $\phi_{k}(t)$ to have slowly varying derivatives, they can also use a penalty matrix that penalizes changes in gradients over time.

Further, when the number of covariates $q$ in $\bm{Z}$ is moderately large, we may wish to impose sparsity in the $\bm{\beta}$ parameter. We encourage such sparsity by including a lasso penalty (Tibshirani, 2011) in the model. In summary, our goal is then to find parameters maximizing the expected penalized log-likelihood, or minimizing the penalized expected deviance loss, under norm constraints:

Only the identifiability constraint (C.1) has entered the objective. We can always guarantee (C.2) by changing the signs in $\bm{V}$, $\bm{\Phi}$, $\bm{\beta}$ and reordering the components afterward without changing the achieved objective value.

Eq. (2.2) describes a non-convex problem. We will find locally optimal solutions via an alternating update procedure: (1) fixing other parameters and updating $\bm{\beta}$ via lasso regressions; (2) fixing $\bm{\beta}$ and updating other model parameters using the EM algorithm. We give details of our iterative estimation procedure in Section 3.

One initialization approach is to form a Tucker decomposition $[\bm{U}_{\perp},\bm{\Phi}_{\perp},\bm{V}_{\perp};\bm{G}]$ of $\bm{X}$ using HOSVD/MLSVD (De Lathauwer et al., 2000) where $\bm{G}\in{\mathbb{R}}^{K_{1}\times K_{2}\times K_{3}}$ is the core tensor and $\bm{U}_{\perp}\in{\mathbb{R}}^{I\times K_{1}}$, $\bm{\Phi}_{\perp}\in{\mathbb{R}}^{T\times K_{2}}$, $\bm{V}_{\perp}\in{\mathbb{R}}^{J\times K_{3}}$ are unitary matrices multiplied with the core tensors along the subject, time and feature directions respectively ($K_{1}/K_{2}/K_{3}$ is the smallest between $K$ and $I/T/J$), and then perform PARAFAC decomposition on the small core tensor $\bm{G}$ (Bro and Andersson, 1998; Phan et al., 2013). We initialize SPACO with Algorithm 2, which combines the above approach with functional PCA (Yao et al., 2005) to work with sparse longitudinal data. Algorithm 2 consists of the following steps: (1) perform SVD on $\bm{X}_{J}$ to get $\bm{V}_{\perp}$; (2) project $\bm{X}_{J}$ onto each column of $\bm{V}_{\perp}$ and perform functional PCA to estimate $\bm{\Phi}_{\perp}$; (3) run a ridge-penalized regression of rows of $\bm{X}_{I}$ on $\bm{V}_{\perp}\otimes\bm{\Phi}_{\perp}$, and estimate $\bm{U}_{\perp}$ and $\bm{G}$ from the regression coefficients

In a noiseless model with $\delta=0$ and complete temporal observations, one may replace the functional PCA step of Algorithm 2 with standard PCA. Then $[\bm{U},\bm{\Phi},\bm{V}]$ becomes a PARAFAC decomposition of $\frac{1}{1+\delta}\bm{X}$.

Proof of Lemma 4.1 is given in Appendix B.2.

In our synthetic experiments in Section 5, we observe that the inclusion of $Z$ can result in a better estimate of the subject scores when the true subject scores strongly depend on $Z$. A natural following-up question is if we can ascertain the importance of such covariates when they exist. Here, we consider the construction of approximated p-values from conditional independence/marginal independence tests between $\bm{Z}_{j}$ and $\bm{U}_{k}$:

Both questions are of practical interest.

Testing for marginal independent between $Y$ and any covariate $Z_{j}$ is straightforward. Let $t(\bm{Z}_{j},{\bf y})$ be a test statistic. Let $T\coloneqq t(\bm{Z}_{j},{\bf y})$ and $T^{*}\coloneqq t(\bm{Z}_{j}^{*},{\bf y})$ , where $\bm{Z}_{j}^{*}$ is a permutation of $Z_{j}$. Then, under the null hypothesis $H_{0}^{marginal}$, $T$ are exchangeable with independent copies of $T^{*}$. Hence, $P(T>t^{*}_{1-\alpha}|{\bf y})\leq\alpha$ for any $\alpha\in(0,1)$ where $t^{*}_{1-\alpha}$ is the $(1-\alpha)$-percentile of the empirical distribution formed by copies $T^{*}$ and $\infty$.

Suppose that we have access to the conditional distribution of $Z_{j}|Z_{j^{c}}$. Let $t(\bm{Z}_{j},\bm{Z}_{j^{c}},{\bf y})$ be a test statistic. Let $T\coloneqq t(\bm{Z}_{j},\bm{Z}_{j^{c}},{\bf y})$ and $T^{*}\coloneqq t(\bm{Z}_{j}^{*},\bm{Z}_{j^{c}},{\bf y})$, where $\bm{Z}_{j}^{*}$ is an independent copy generated from the conditional distribution $Z_{j}|Z_{j^{c}}$ and $\bm{Z}_{j^{c}}$. Then, under the null hypothesis $H_{0}^{partial}$, $T$ and $T^{*}$ have the same law conditional on $\bm{Z}_{j^{c}}$ and ${\bf y}$. So $P(T>t^{*}_{1-\alpha}|\bm{Z}_{j^{c}},{\bf y})\leq\alpha$, where $t^{*}_{1-\alpha}$ is the $(1-\alpha)$-percentile of the conditional distribution of $T^{*}$ (Candès et al., 2018).

Proof of Lemma 4.2 is given in Appendix B.3. As a result, when the term $\Delta_{i}(\delta)=0$, we can apply the traditional randomization tests with response $\tilde{y}_{i}(\delta)$ and features $\tilde{\bf z}_{i}(\delta)$ for subject $i$. Proposition 4.3 gives our construction of the observed and randomized test statistics .

In practice, $\bm{\beta}_{\ell}$ for the other factors $\ell\neq k$ are also estimated from the data. When $\delta\neq 0$, we could have $\Delta_{i}(\delta)\neq 0$, which renders the randomization test invalid. Thus, we typically set $\delta=0$. For convenience, we drop the $\delta$ argument when $\delta=0$, e.g., $T_{partial}=T_{partial}(\delta)$. Algorithm 3 summarizes our proposal.

The conditional randomization test requires generating $Z_{j}^{*}$ from the conditional distribution $Z_{j}|Z_{j^{c}}$. We estimate the conditional distribution of $Z_{j}^{*}$ via proper exponential family distribution . More details on the generations of $\bm{Z}_{j}^{*}$ are provided in Appendix A.4.