Fast Probabilistic Whitening Transformation for Ultra-High Dimensional Genetic Data

Abstract

Statistical methods often make assumptions about independence between the samples or features of a dataset. Yet correlation structure is ubiquitous in real data, so these assumptions are often not met in practice. Whitening transformations are widely applied to remove this correlation structure. Existing approaches to whitening are based on standard linear algebra, rather than a probabilistic model, and application to high dimensional datasets with $n$ samples and $p$ features is problematic as $p$ approaches or exceeds $n$. Moreover, the computational time becomes prohibitive since the naive transform is cubic in $p$. Here we propose a probabilistic model for data whitening and examine its properties based on first principles as $p$ increases. We demonstrate the statistical properties of the probabilistic model and derive a remarkably efficient algorithm that is linear instead of cubic time in the number of features. We examine the out-of-sample performance of the probabilistic whitening model on simulated data, and real genotype data. In an application to impute z-statistics from unobserved genetic variants from a genome-wide association study of schizophrenia, the probabilistic whitening transformation, had the lowest mean square error while being up to an order of magnitude faster than other methods. Using this approach, we also identify tandem repeats that explain genetic regulatory signals for disease-relevant genes. Analyses are implemented in our novel open source R packages decorrelate and imputez.

1. Introduction

Correlation structure is ubiquitous in real world datasets. Yet statistical models often make assumptions about independence between a dataset’s rows or columns that are not satisfied practice. Whitening transformations are a widely used preprocessing step to remove correlation structure in the data before performing other analyses [1]. In general, for a data matrix $Y_{n\times p}$ with $n$ samples and $p$ features, the procedure computes the sample covariance matrix between features, ${\text{Σ}}_{p\times p}$, and transforms the observed data by post-multiplying by the non-unique whitening matrix ${\text{Σ}}^{-\frac{1}{2}}$ defined using standard linear algebra [1]. Transforming rows instead follows a complementary procedure. Covariance and their corresponding whitening matrices, along with correlation and their corresponding ‘decorrelating’ matrices, play a central role in many standard analysis methods such as multivariate regression, canonical correlation analysis, linear and quadratic discriminate analysis, Gaussian mixture models, Maholanobis distance, multivariate outlier detection, and independent component analysis [2–4], linear mixed models [5, 6], as well as deep learning applications [7–9].

For low dimensional data where $n\gg p$ seen in classical statistical settings, whitening transforms the data to have an identity sample covariance matrix. This procedure assumes that the true covariance matrix is known, or is well estimated by the sample covariance matrix. Yet the use of the sample covariance matrix for this transformation can be problematic for two distinct reasons. Computationally, evaluating the covariance matrix and creating the whitening matrix requires $\mathcal{O}\left(p^2\right)$ memory and $\mathcal{O}\left(p^3\right)$ time, which becomes infeasible for large $p$ [1]. Statistically, the transformation must be adapted for the high dimensional case since the sample covariance matrix is no longer full rank or invertible. In modern high dimensional datasets, these computational and statistical limitations raise issues for the application of whitening transforms either as a preprocessing step, or a central part of widely used statistical methods.

Beyond general applications in statistical analysis, the field of statistical genetics makes extensive use of whitening transformations in order to account for the correlation between nearby genetic variants. This correlation structure is termed ‘linkage disequlibirum’ or ‘LD’ and is a major complication in downstream analysis of regression summary statistics (i.e. coefficient estimates or z-statistics) from genome-wide association studies [10]. It is common practice to estimate the correlation matrices from an external ‘reference panel’ of samples not included in the study of interest [10]. These ‘out-of-sample’ estimates are then used either as a preprocessing step to transform the summary statistics to be approximately uncorrelated, or internally in a statistical model. This approach has become standard in the field for phenotype prediction [11, 12], statistical fine-mapping [13–15], imputing z-statistics [16], detecting anomalous z-statistics [17], multiple testing correction [18], identifying genome annotations enriched for risk variants [19], and simulating knockoff datasets with similar correlation structure [20]. While the literature on whitening transformations has focused on estimating and applying the transform on the same data [1], an understanding of the performance of these out-of-sample applications has been lacking [21].

Direct estimation of covariance and precision matrices is a huge field of research [22–24]. Many approaches apply regularization to ensure the estimated covariance or precision matrix has specific properties. Widely used methods apply regularization so that the estimated covariance matrix is positive definite by shrinking the sample eigen-values [25, 26], or ensure sparsity of the covariance matrix and its inverse by shrinking regression coefficients corresponding to partial correlation coefficients [27]. Others regularize the condition number [28]. There have been many methods formulated as a convex combination of two matrices. Some methods use a convex combination between a non-sparse and sparse components [29, 30]. Stein-type estimators use a convex combination between the sample covariance and a target matrix [25, 26, 31]. A probabilistic approach modeling the observed data as a multivariate Gaussian with an inverse Wishart prior on the covariance gives a posterior mean that is a convex combination between the sample covariance matrix and the prior scale matrix [32–36].

Yet whitening transformations used in a preprocessing step or internally in a statistical model do not require direct estimates of the covariance matrix. Instead, covariance matrices are merely nuisance parameters: essential to the model, but not of direct interest. In data whitening, only the transformed data, rather the covariance matrix itself, is of interest. Here we seek to account for the covariance, rather than estimate it directly, and we propose an implicit covariance approach. This important distinction motivates the probabilistic framework, statistical theory and computationally efficient algorithms we use here to perform whitening transformations without directly estimating the covariance or whitening matrices.

Bayesian or, more broadly, probabilistic formulations of existing techniques can lead to deeper understanding of model behavior, generalization of the model, regularization to incorporate prior knowledge, or incorporation into a hierarchical model [3, 4]. Widely used Gaussian mixture models and extensions follow from a probabilistic formulation of the k-means problem, and probabilistic formulations of principal components analysis have long been influential in statistical machine learning [37, 38]. Yet we are not aware of such a rigorous probabilistic formulation for the widely used whitening transformation.

Here we propose a probabilistic formulation of the whitening transformation under Gaussian noise, and incorporate regularization in the form of an inverse-Wishart prior distribution on the covariance matrix. This produces a well-known estimator with a single hyperparameter that can be efficiently estimated using an empirical Bayes approach. Using first principles of the probabilistic model, we give a direct interpretation of the hyperparameter and examine the statistical behavior of this whitening transformation. We demonstrate that the probabilistic whitening transformation does not require direct estimation of the covariance or whitening matrices, and derive a remarkably efficient approach to apply the transformation to ultra-high dimensional data.

We examine the empirical performance of the probabilistic whitening transformation on simulated data and genetic data from the 1000 Genomes Project [40]. We apply the whitening transformation to a genome-wide association study of schizophrenia [41], and impute z-statistics of unobserved variants using an external reference panel. Finally, we use this approach to impute genetic regulatory signals for tandem repeats effecting gene expression in the human brain.

2. Results

2.1. Introduction to whitening transformations

Consider observed data $Y_{n\times p}$ with $n$ samples and $p$ features drawn from a distribution with mean vector $\overrightarrow{\mu }$, identity covariance between rows, and covariance between features given by ${\text{Σ}}_{p\times p}$. Let

$$Y\sim \left[\overrightarrow{\mu },I_n,\text{Σ}\right]$$ (1)

indicate these properties without specifying particular distributional assumptions. Transforming the data by the non-unique whitening matrix ${\text{Σ}}^{-\frac{1}{2}}$ produces mean and covariances according to

$$Y{\text{Σ}}^{-\frac{1}{2}}\sim \left[\overrightarrow{\mu }{\text{Σ}}^{-\frac{1}{2}},I_n,I_p\right].$$ (2)

where $\text{Σ}$ is known. This follows from the definition of covariance and the standard property that ${\text{Σ}}^{-\frac{1}{2}}{\text{ΣΣ}}^{-\frac{1}{2}}=I$. While it is common to assume a multivariate Gaussian distribution for the observed data, this transformation applies to the first two moments of any multivariate distribution. Applying this transformation scales and rotates the data to produce an identity covariance matrix. The whitening matrix is not uniquely defined and different transformation matrices result in different rotations all with identity covariance [1].

The sample covariance is $\widehat{\text{Σ}}=n^{-1}{Y}^{T}Y$, following mean centering of the columns. Let the singular value decomposition of the scaled observed data be

$$Y/\sqrt{n}=U_{n\times p}{D}_{p\times p}{V}_{p\times p}^T,$$ (3)

with singular values $\overrightarrow{d}$, so that the eigen decomposition of the sample covariance is $\widehat{\text{Σ}}=V{D}^2{V}^T$ and the whitening matrix ${\widehat{\text{Σ}}}^{-\frac{1}{2}}=V{D}^{-1}{V}^T$ transforms the data to have identity covariance (SI 5.1).

This definition of the whitening matrix minimizes the total squared distance between original and transformed data, and is termed the ‘zero-phase components analysis’ or ZCA whitening [1]. This three step transform, one for each matrix multiplication, can be interpreted as rotating, scaling, then unrotating the observed data to give an identity sample covariance (Figure 1). This whitening transform is the most intuitive since it returns the data to its original axes, and it is the foundation for the transformations we describe below.

Figure 1: Intuition for data whitening transformation.

A) Original data, B) Data rotated along principal components, C) Data rotated and scaled, D) Data rotated, scaled and rotated back to original axes. Green arrows indicate principal axes and lengths indicate eigen-values.

In practice, the true value of $\text{Σ}$ is not known and it must be estimated from the observed data $Y$. An exact whitening transformation that produces an identity covariance matrix is only possible when $n>p$ so that that sample covariance matrix, $\widehat{\text{Σ}}$, is invertible. Even in this case, the transform can be unstable when the trailing singular-values are very small, since evaluating $D^{-1}=\text{diag}(1/\overrightarrow{d})$ becomes very unstable for small values of $d_i$ and are sensitive to precision limitations in numerical linear algebra. In high dimensional data with $n<p$, the trailing singular values are exactly zero, so some form of regularization is needed. An exact whitening transformation no longer exists in these cases, we want to learn a transformation that that transformed data has a covariance close to identify by some metric (SI 5.2).

2.2. A probabilistic model for whitening transformations

Here we propose a whitening transform based on an probabilistic model of the observed data and demonstrate superior statistical performance compared to standard data whitening transformations [1], even in the low-dimension case where both are applicable. In proposing a probabilistic model, we consider the observed data to have a matrix normal distribution¹ with unknown mean, identity covariance between rows, and unknown covariance between columns. We then place an inverse Wishart (IW) prior on the covariance matrix between columns. The normal assumption is of course natural for unbounded continuous data, and the IW is natural because it is the conjugate prior so that the posterior distribution of the covariance is also IW [42]. Formally,

$$Y~\mathcal{M}\mathcal{N}\left(\overrightarrow{\mu },I_n,\text{Σ}\right)$$ (4)

$$\text{Σ}\mid T~\mathcal{I}\mathcal{W}((\delta -p-1)T,\delta )$$ (5)

where $T$ is the target matrix of the prior and $\delta >p-1$. Based on this model, the posterior expectation of the covariance is

$$E[\text{Σ}\mid Y]=(1-\lambda ){\widehat{\text{Σ}}}_{MLE}+\lambda T$$ (6)

$$\lambda =(\delta -p-1)/(\delta +n-p-1)$$ (7)

where ${\widehat{\text{Σ}}}_{MLE}=n^{-1}{Y}^{T}Y$ is the sample covariance matrix and $\lambda \in [0,1]$ by definition [32–35] The posterior expectation is therefore a convex combination between the sample covariance, ${\widehat{\text{Σ}}}_{MLE}$, and the prior target matrix, $T$.

Setting the target matrix to be symmetric positive definite (SPD) ensures that the posterior expected covariance is also SPD [33]. Restricting $T$ to be a minimally informative prior target corresponds to setting all prior covariances to zero. We focus here on the important case where the IW prior is centered around a scaled identity matrix, and we demonstrate that the estimate has exceptionally efficient computational scaling and produces estimates that are tractable to study analytically. Following previous work [25, 26, 33–35, 43] we set the target matrix to be $T=\nu I$ where $\nu =\mathrm{tr}\left({\widehat{\text{Σ}}}_{MLE}\right)/p$ is the average variance across all $p$ features. This gives the important special case we consider here with

$${\widehat{\text{Σ}}}_{\lambda }=(1-\lambda ){\widehat{\text{Σ}}}_{MLE}+\lambda \nu I$$ (8)

where ${\widehat{\text{Σ}}}_{\lambda }$ is defined as the posterior expected covariance from the probabilistic model.

In the important case of out-of-sample whitening where the whitening transform is learned from a training set before being applied to the observed data, the standard whitening transform motived by linear algebra assumes the correlation structure of the training and observed data are the same. This is a strong assumption that is rarely satisfied in practice. The finite size of the training set means that correlation structure will differ from the observed data, so the standard whitening transform will overfit the training data. Instead, the probabilistic model substantially relaxes this central assumption, and assumes only that the training and observed data are drawn from the same distribution.

In practice, the shrinkage can be applied to the correlation matrix while retaining the diagonal variance terms (SI 5.3).

2.3. Estimating the shrinkage parameter

Covariance estimators that are convex combinations of the sample covariance and target matrices have received much attention. The estimator for $\lambda$ depends on assumptions about the data and the specified loss function. Ledoit and Wolf [25] and Schäfer and Strimmer [26] propose minimizing the expected squared Frobenius loss between the true and estimated covariance matrices, then use a consistent estimate of $\lambda$ computed from the observed data. Further work extended this approach to improve estimation of $\lambda$ in the high dimensional setting [31, 43, 44]. Others use a probabilistic model [32–35], cross-validation [45, 46] or the out-of-sample likelihood [47, 48]. Yet these methods are designed for direct estimation of covariance or precision matrices, rather than data whitening.

Following the probabilistic formulation, we take a Bayesian approach to estimating $\lambda$, and we refer to this as the Gaussian Inverse Wishart empirical Bayes (GIW-EB) model. We integrate out the covariance matrix and maximize the likelihood of the observed data as a function of $\lambda$. Due to the analytic properties of the Gaussian and IW distributions, the covariance matrix can be treated as a nuisance parameter and marginalized out according to

$$p(Y\mid \delta ,\nu )=\int p(Y\mid \text{Σ},\delta ,\nu )p(\text{Σ}\mid \delta ,\nu )d\text{Σ},$$ (9)

so the marginal distribution is multivariate Student-t [42]. The hyperparameter $\delta$, and therefore $\lambda$, can be estimated directly from data using an empirical Bayes approach by maximizing $p(Y\mid \delta ,\nu )$ (SI 5.4). Given the special form of the target matrix, $\nu I$, we demonstrate that the data only enters the log-likelihood through the sample singular values. This has important implications for the probabilistic whitening transformation described below.

2.4. Properties of posterior expected covariance

The behavior of this probabilistic whitening transformation corresponding to the posterior expected covariance in Equation (8) can be understood in terms of the spectral decomposition of the observed data matrix, $Y$.

Based on the singular value decomposition of the observed data, the eigen decomposition of ${\widehat{\text{Σ}}}_{\lambda }$ is

$${\widehat{\text{Σ}}}_{\lambda }=(1-\lambda ){\widehat{\text{Σ}}}_{MLE}+\lambda \nu I$$ (10)

$$=V\mathrm{diag}\left[(1-\lambda ){\overrightarrow{d}}^2+\lambda \nu \right]V^T.$$ (11)

Therefore the eigen vectors of ${\widehat{\text{Σ}}}_{\lambda }$ are the same as for the sample covariance matrix, but the eigenvalues are shrunk so that the $j^{th}$ eigen-value is $(1-\lambda )d_j^2+\lambda \nu$. Since $\lambda \in [0,1]$, $\nu >0$, and $d_j^2\ge 0$, it follows that the eigen-values are always positive, so that ${\widehat{\text{Σ}}}_{\lambda }$ is aways positive definite and invertible, even when ${\widehat{\text{Σ}}}_{MLE}$ is not.

2.5. Statistical properties of probabilistic whitening transformation

The probabilistic whitening transformation can be understood in similar terms as the posterior expected covariance. Consider a fixed (i.e. non-random) latent data matrix $X_{n\times p}$ with uncorrelated columns, and let the observed data $Y_{n\times p}$ be generated by applying a rotation based on a symmetric positive definite covariance matrix ${\text{Σ}}_{p\times p}$ so that

$$Y~\left[X{\text{Σ}}^{\frac{1}{2}},I_n,\text{Σ}\right].$$ (12)

Given the observed data, $Y$, estimate the covariance matrix as ${\widehat{\text{Σ}}}_{\lambda }$ and apply the whitening transformation so that

$$Y{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}}~\left[X{\text{Σ}}^{\frac{1}{2}}{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}},I_n,{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}}\text{Σ}{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}}\right].$$ (13)

Since ${\widehat{\text{Σ}}}_{\lambda }$ converges to $\text{Σ}$ as $n\to \infty$ for a fixed $p$ [33], the whitening transform recovers the original signal according to

$$\underset{n\to \infty }{\lim }Y{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}}\stackrel{d}{\to }\left[X,I_n,I_p\right].$$ (14)

The whitening transformation can then be expressed in terms of singular vectors and modified singular values according to

$$\boxed{Y{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}}=U\mathrm{diag}\left(\sqrt{\frac{{\overrightarrow{d}}^2}{(1-\lambda ){\overrightarrow{d}}^2+\lambda \nu }}\right)V^T}$$ (15)

The central statement in Equation (15) of the whitening transformation in terms of the original singular vectors and modified singular values has both important statistical and computational implications we discuss below.

The whitening transformation retains the original singular vectors and acts only by modifying the sample singular values. From here we can examine the behavior of the transformation as a function of $\lambda$ (SI 5.5), and interpret the effect of $\lambda$ on the effective number of parameters used in covariance estimation (SI 5.6).

2.6. Scalar summaries of the correlation or covariance matrix

Common scalar summary statistics of covariance matrices can be computed efficiently based on the form of the spectral decomposition of ${\widehat{\text{Σ}}}_{\lambda }$ (Table 1). Standard values like trace, log determinant and condition number can be computed easily from the shrunken eigen-values. In addition the effective degrees of freedom [49], average correlation, average squared correlation [50], effective variance [51], effective number of features [52] are easily computed (SI 5.7).

Table 1: Computing scalar summary statistics of correlation matrices.

Terms are defined in the text except for letting the shrunken eigen-values be ${\tilde{d}}_i=d_i^2(1-\lambda )+\nu \lambda$, and the mean shrunken eigen-value be $a=\frac{1}{p}{\sum }_{i=1}^p{\tilde{d}}_i$. Let $\overrightarrow{\mathbf{1}}$ be a vector of length $p$ with all entries having value 1.

Statistic	Formula	Reference
Trace	${\sum }_{i=1}^p{\tilde{d}}_i$
Log det	${\sum }_{i=1}^p\log {\tilde{d}}_i$
Condition number	${\tilde{d}}_1/{\tilde{d}}_p$
Average squared correlation	${\sum }_{i=1}^p{\left({\tilde{d}}_i-a\right)}^2/p(p-1)a^2$	[50]
Effective degrees of freedom	$(1-\lambda )p(p-1)/2$	[49]
Average correlation	$\left({\overrightarrow{\mathbf{1}}}^T{\text{Σ}}_{\lambda }\overrightarrow{\mathbf{1}}-p\right)/(2p(p-1))$
Effective variance	$\exp \left[\frac{1}{p}{\sum }_{i=1}^p\log {\tilde{d}}_i\right]$	[51]
Effective number of features	${\overrightarrow{\mathbf{1}}}^T{\text{Σ}}_{\lambda }^{-1}\overrightarrow{\mathbf{1}}$	[52]

2.7. Computational methods for data whitening

The computational time for applying the standard whiting transformation, even in the low dimensional setting, can quickly become intractable for moderate $p$. The time complexity of the standard algorithm is $\mathcal{O}\left(p^3+n{p}^2\right)$ (SI 5.8), and becomes very expensive for $p>100$ and can be intractable for $p>2000$ features. Yet the entire procedure can be substantially accelerated using properties linear algebra (Table 2). It is not necessary to explicitly compute the sample covariance and whitening matrices and then transform the observed data. Equation (15) rephrases the whitened data in terms of modified singular values, and the original left and right singular vectors. This gives a substantially faster algorithm. In the low dimensional setting of $n>p$, computing the SVD of the observed data is $\mathcal{O}\left(n{p}^2\right)$, estimating and applying the $\nu$ and $\lambda$ is negligable, and evaluating the matrix products is $\mathcal{O}\left(n{p}^2\right)$. The procedure is thus substantially reduced to $\mathcal{O}\left(n{p}^2\right)$ and uses $\mathcal{O}(np)$ memory.

Table 2:

Computational and memory complexity of whitening methods for $n$ samples and $p$ features using a rank $k$ transformation.

Method	Decomposition	Estimate $\lambda$	Transform	Total	Memory
Standard	$p^3+n{p}^2$	$\ge np$	$n{p}^2$	$p^3+n{p}^2$	$p^2$
GIW-EB $(n>p)$	$n{p}^2$	$n$	$n{p}^2$	$n{p}^2$	$np$
GIW-EB $(n<p)$	$n^{2}p$	$p$	$n^{2}p$	$n^{2}p$	$np$
GIW-EB, low rank	$knp$	$k$	$knp$	$knp$	$np$

2.7.2. Accelerating computation in the high dimensional setting

The standard formulation of the data whitening transformation does not apply to the high dimensional setting of $n<p$ since the sample covariance matrix is no longer invertible. Some regularization of the singular values is necessary. The probabilistic whitening transformation indeed shrinks the singular values towards $\nu$ and addresses the statistical challenge of the high dimensional setting. Yet this computational approach is still quadratic in $p$, and is infeasible for large $p$ typical in real datasets.

In general, the rank, $r$, of the sample covariance matrix is bounded above by $\min (n,p)$. In the high dimensional setting, at most only $r\le n$ singular values are nonzero. Evaluating the SVD is $\mathcal{O}\left(n^{2}p\right)$ in this case only produces the first $n$ singular vectors and values. So there remain $p-n$ uncomputed singular vectors with corresponding singular values of zero. Directly computing these is impractical and, in fact, unnecessary.

Again, the key insight comes from Equation (15). Since all $p$ singular values, including ones with value zero, are modified to be positive it seems that even left and right singular vectors with no variation (i.e. zero singular values) are required to compute the transformation. Yet here we demonstrate that in fact only non-zero singular values and their corresponding singular vectors are required. This insight further reduces the computational complexity of the entire transform.

The substantial improvement in computational complexity for the whitening transformation in the high dimensional settings comes from computing and using only singular vectors corresponding to nonzero singular values. The method is useful beyond this specific context, so here we derive a general theorem and discuss special cases and applications. The theorem is exact if all nonzero singular values are retained, and is approximate of only a partial set are retained. The derivation generalizes and simplifies the work of Lippert, et al. [6] (SI 5.9).

Theorem 1.

Let $C$ be a $p\times p$ symmetric positive semi-definite matrix, with $k\le p$ positive eigen values. Let the eigen decomposition of $C$ be $Vdiag(\overrightarrow{s})V^T$ where $s_i=0$ for $i>k$. Let the eigen-vectors be partitioned as $V=\left[V_1{V}_2\right]$, and the eigen-values be partitioned as $\overrightarrow{s}=\left[{\overrightarrow{s}}_1{\overrightarrow{s}}_2\right]$ to separate the nonzero and zero components. Let raising a matrix to an exponent $\alpha$ correspond to raising the eigen-values of the matrix to $\alpha$. Then

$$A{\left(C+\tau I_p\right)}^{\alpha }B=\left(A{V}_1\right)diag\left(s_1^{\alpha }\right)V_1^{T}B+{\tau }^{\alpha }\left(AB-\left(A{V}_1\right)V_1^{T}B\right)$$ (16)

where $V_1$ and ${\overrightarrow{s}}_1$ indicate the first $k$ eigen-vectors and values, respectively, and $A_{a\times p}$ and $B_{p\times b}$ are arbitrary matrices, $\tau$ is a positive scalar, and $I_p$ is the p-dimensional identity matrix.

Applying Theorem 1 to the whiting transform, and using the decomposition

$${\widehat{\text{Σ}}}_{\lambda }=V\text{diag}\left[(1-\lambda ){\overrightarrow{d}}^2+\lambda \nu \right]V^T,$$ (17)

restated here from Equation (11), it follows that

$$Y{\widehat{\text{Σ}}}_{\lambda }^{-\frac{1}{2}}=YV\text{diag}{\left[(1-\lambda ){\overrightarrow{d}}^2+\lambda \nu \right]}^{-\frac{1}{2}}{V}^T$$ (18)

$$=\left(Y{V}_1\right)\text{diag}{\left[(1-\lambda ){\overrightarrow{d}}_{1:k}^2+\lambda \nu \right]}^{-\frac{1}{2}}{V}_1^T+{\lambda }^{-\frac{1}{2}}{\nu }^{-\frac{1}{2}}\left(Y-\left(Y{V}_1\right)V_1^T\right)$$ (19)

where $V_1$ stores the first $k$ eigen vectors of ${\text{Σ}}_{\lambda }$ corresponding to the right singular vectors of $Y/\sqrt{n}$, and ${\overrightarrow{d}}_{1:k}^2$ are the first $k$ eigen values.

Using standard properties of matrix multiplication, the time complexity of evaluating these terms are defined as follows.

Corollary 1.1.

Since the time complexity of multiplying $a\times b$ and $b\times c$ matrices is $\mathcal{O}(abc)$, the time to evaluate the general formulation of Theorem 1 is $\mathcal{O}(abnp)$, once the SVD is already computed. In probabilistic data whitening in Equation (19), $B=I$ and $a=n$ so the time complexity reduces to $\mathcal{O}(knp)$.

In general, the whitening transformation can be applied in $\mathcal{O}(knp)$ time after the SVD is computed. When the first $n$ singular vectors and values are computed using the standard SVD, then $k=n$, so the entire procedure is $\mathcal{O}\left(n^{2}p\right)$, which is now linear in the number of features. This is a dramatic improvement from the naive method that is cubic in the number of features (Table 2).

2.7.3. Scaling to the ultra-high dimensional setting

While this algorithm is now linear in $p$, it is still quadratic in the number of samples, $n$. Yet empirical data can often be well approximated as a low rank matrix [53]. As $n$ and $p$ increase but the dimension of the underlying signal remains fixed, it is not necessary to compute all $\min (n,p)$ singular values and vectors. Instead a rank $k\ll \min (n,p)$ partial SVD can be computed in $\mathcal{O}(knp)$ time to provide a rank $k$ approximation of the observed data [54]. Moreover, the algorithm in Equation (19) can also be evaluated in $\mathcal{O}(knp)$ time given the partial SVD (Table 2). The complexity of this approach is now linear in both $n$ and $p$, and dependent on the rank $k$ approximation chosen by the analyst.

2.8. Beyond the whitening transformation

While this work focuses on the whitening transformation, applying Theorem 1 with other values for the general exponent α are also useful in practice. In the context of probabilistic whitening transformation with variables defined as in Equation (19), four common operations involving ${\widehat{\text{Σ}}}_{\lambda }^{\alpha }$ can be accelerated:

• $\alpha =1$ gives the posterior expected covariance matrix.

• $\alpha =-1$ gives the corresponding precision matrix.

• $\alpha =-\frac{1}{2}$ gives the whitening transformation.

• $\alpha =\frac{1}{2}$ gives the ‘coloring matrix’ so that post-multiplying by ${\widehat{\text{Σ}}}_{\lambda }^{\frac{1}{2}}$ produces a correlation of ${\widehat{\text{Σ}}}_{\lambda }$ between columns.

3. Results

3.1. Performance on simulated data

The formulation of the probabilistic whitening transformation provides two advances of practical interest in analyses of large-scale datasets. First, the implicit covariance approach reduces computational time for high-dimensional data from the standard cubic time method that creates and then inverts the covariance matrix. For example, in simulations with $n=5000$ and an increasing number of features the GIW-EB methods runs in < 2 min and the low rank GIW-EB with k=50 runs in < 1 min, when the standard method becomes intractable both in terms of time and memory (Figure 2A). Second, the probabilistic model motivates the empirical Bayes estimation of the shrinkage parameter which only uses the sample singular values of the data. Due to its probabilistic formulation, the whitening high-dimensional data using the GIW-EB method can produce lower out-of-sample error than other methods. For example, in a simulation Scenario 3 (described below), both full and low rank GIW-EB give estimated values of $\lambda$ that approach the optimal value while other methods give $\lambda$ values that are too small and have higher error (Figure 2B).

Figure 2: Performance of probabilistic whitening transformation.

A) Run times are shown for standard (green), GIW-EB (red) and low rank GIW-EB with $k=50$ (dark red) whitening transformations. Results shown for $n=5000$ samples and an increasing number of features. Points indicate observed wall time and lines show loess smooth. B) The empirical out-of-sample root mean squared error for a simulated dataset across the range of $\lambda$ values is shown by the black curve. The $\lambda$ value giving the lowest error is show by the purple point. Values of $\lambda$ estimated by other methods are shown by vertical lines.

Since an exact whitening transformation produces an identity covariance matrix when applied to the data it was estimated from, the accuracy of an approximate whitening transformation can be evaluated by estimating it from a subset of the data and computing the empirical covariance after applying the transform to the rest of the dataset. This out-of-sample accuracy of the whitening transformation is captured by the root mean squared error of the difference between the empirical covariance of the transformed data and the identity matrix.

We consider simulations with $n=500$ and the number of features $p$ increasing from 20 to 2000, with correlation between features in each of 4 scenarios:

1. constant correlation of $\rho =0.6$

2. auto-correlation with $\rho =0.95$

3. 4 blocks of constant correlation of $\rho =0.6$

4. 4 blocks of auto-correlation with $\rho =0.95$

In Scenario 3, most methods show equivalent performance for small values of $p$, error increases as $p$ approaches $n$, and then error decreases and converges to the optimal root mean squared error (Figure 3A). The ‘Oracle’ method, shown in black here, transforms the data using the true covariance, so it gives a lower bound on the error. Using a whitening transformation without shrinkage by setting $\lambda =0$ performs poorly even for small $p$ and error only increases with $p$. Using the pseudoinverse improves performance when $p>n$, but does not compete with shrinkage methods. Examining different approaches to estimate the shrinkage parameter, the GIW-EB method gives the smallest error across all values of $p$ and the low rank GIW-EB with k = 50 outperforms other methods for moderate $p$. Both full and low rank GIW-EB give errors that decrease with $p$ and avoids the increase in error at $p\approx n$ seen with other shrinkage methods.

Figure 3: Out-of-sample performance of whitening transformation.

Simulated data for $n=500$ sample are drawn from a multivariate normal distribution with known covariance matrix and 11 whitening transformations are then applied to 500 additional samples from the same distribution. Of these, 9 methods either estimate the shrinkage parameter $\lambda$ from the data or use a fixed value. Setting $\lambda =0$ corresponds to no shrinkage, and pseudoinverse does not perform any shrinkage. The oracle method performs a whitening transformation using the true correlation matrix from which the data are simulated, and so gives the best out-of-sample performance. Performance is evaluated by comparing the correlation matrix of the transformed data to the identity matrix, and computing the root mean squared error (rMSE). A) rMSE of the whitening transformation for an increasing number of features B) Zoom-in on y-axis from (A). Points indicate observed values and lines show loess smooth.

Overall, the GIW-EB method consistently gives the lowest out-of-sample root mean square error across a range of $p$ values in the 4 simulation scenarios (Figure S1-4). The low rank GIW-EB method is competitive in Scenarios 1 and 3, were the true covariance is low rank.

3.2. Performance on genetic data from 1000 Genomes Project

Since many analyses of summary statistics from genome-wide association studies apply a whitening transformation estimated from an external reference panel, we evaluated the out-of-sample performance of these whitening transformations on genetic data from the 1000 Genomes Project [40]. Individuals of European, Asian and African ancestry were analyzed separately, and the whitening transformation was performed on independent LD blocks including varying number of genetic variants defined in each population [55]. The whitening transform was trained on half of the samples and evaluated on the other half. Despite the genetic data not being multivariate normal, shrinkage methods performed well across most windows in European individuals (Figure 4A). The GIW-EB matches or is competitive with other methods that estimate $\lambda$ from the data. The low rank GIW-EB method with $k=50$ is most competitive for windows with < 1,500 variants. The GIW-EB methods and methods with fixed $\lambda$ values use the implicit covariance algorithm, and these give the best computational performance with a wall time of ∼5 min compared to 3–30 hours for the other methods (Figure 4B). Error profiles and computational time was similar for individuals of African and Asian ancestry (Figure S5).

Figure 4: Out-of-sample whitening performance on genetic data.

A) Root mean squared error (rMSE) compared to baseline for European individuals stratified by the number of genetic variants in the LD window. Values indicate the values averaged across windows, and bars indicate the 95% confidence interval. When normalized rMSE values exceed the axes, the value is shown in white text. B) Wall time for each method shown on a log₁₀ scale. Methods using GIW-EB or fixed $\lambda$ values use the implicit covariance algorithm.

3.3. Application to imputing GWAS summary statistics

Genome-wide associations studies (GWAS) release summary statistics from their data corresponding to z-statistics from the set of genetic variants analyzed in the study. Since the number of genetic variants analyzed in the original study can vary widely, there has been interest in imputing z-statistics of additional variants observed in an external reference panel. Since the correlation between z-statistics of two genetic variants matches the linkage disequilibrium between the two variants, Pasaniuc, et al. [16] proposed a method to impute z-statistics from unobserved variants using the set of observed z-statistics and the correlation structure between variants in a genomic region. This is substantially faster and cheaper than performing ‘genotype imputation’ on the full dataset [56].

Imputing z-statistics from unobserved variants can be stated in terms of out-of-sample whitening (SI 5.10), and we evaluated the performance of different whitening methods on imputation accuracy. We used summary statistics for 5.6M variants from schizophrenia GWAS of 58K cases and 77K controls of European ancestry [41]. Z-statistics for 15M variants were imputed using a reference panel of 407 unrelated individuals of European ancestry from the 1000 Genomes Project [40] using a sliding window of 1 Mb and a flanking region of 250 kb across all autosomes. Accuracy was evaluated by comparing the imputed and observed z-statistic for 460K randomly selected variants with minor allele frequency ≥ 5%. GIW-EB gave high imputation accuracy across a broad range of z-statistic values (Figure 5A), and give smallest genome-wide root mean squared error (Figure 5B). Imputation error increases for small minor allele frequencies (MAF), and GIW-EB gives the smallest squared error across the range of MAF values (Figure 5C). Given 407 samples and a mean of ∼ 3100 variants in each genome window, the low rank GIW-EB with $k=50$ was not a good approximation in this case. The GIW-EB methods and methods with fixed $\lambda$ values use the implicit covariance algorithm, and these give the best computational performance and require < 1 hr for genome-wide analysis for this reference panel compared to 4–126 hours for other methods (Figure 5D). Overall, GIW-EB method is the most accurate and fastest approach to impute GWAS summary statistics.

Figure 5: Imputation of GWAS summary statistics.

A) Imputed z-statistics plotted against the observed z-statistics for 461,707 SNPs. B) Root mean squared prediction error for each method. Red vertical dashed line indicates the value for GIW-EB. C) Square prediction error as a function of minor allele frequency. Smooth curves are fit by a generalized additive model that is the default in ggplot2::geom_smooth(). GIW-EB (k=50) was excluded from this plot due to high imputation error. D) Compute time required for each method. When values exceed the axes, values are shown in text. Methods using GIW-EB or fixed $\lambda$ values use the implicit covariance algorithm.

3.4. Application to impute eQTL signals for tandem repeats

Studies of genetic risk for disease and genetic regulation of gene expression have focused on common single nucleotide polymorphisms (SNPs) because they can be measured cheaply at scale. In recent work, Ziaei Jam, et al. [57] created a population reference panel including donors from 1000 Genomes Project of 1.7M tandem repeats that are not captured by existing panels. Here, we integrate this resource with summary statistics from a study of genetic regulation of gene expression in the human brain using almost 4K bulk RNA-seq samples [58]. As a proof of principle, we applied this approach to 3 disease-relevant genes with complex genome structure and we identified tandem repeats that explain the genetic regulatory signal (Figure 6). C9orf72 is a key risk gene for amyotrophic lateral sclerosis, with expansion of a intronic hexanucleotide repeat conferring disease risk and effecting gene expression[59]. TSPAN14 and TYK2 have been implicated in risk for Alzheimer’s disease [60], and here we identify tandem repeats with a stronger association with gene expression than SNPs in the region.

Figure 6: Imputation of eQTL summary statistics in human brain implicates tandem repeats..

Results are shown for A) C9orf72, B) TSPAN14, and C) TYK2. Black points indicate observed variants included in the original eQTL analysis, red points indicates imputed variants, and red triangles indicate imputed tandem repeats.

4. Discussion

The whitening transformation is a central, if often unstated, component of many widely used statistical methods. While many standard statistical methods have been reformulated with a probabilistic model, or extended to handle regularization and modern high dimensional datasets, the whitening transform has received less attention. The standard whitening transformation is limited by computational time, memory usage, and gives poor out-of-sample performance on modestly sized datasets. Regularizing the covariance estimate improves statistical performance, but is still very computationally demanding for high dimensional datasets.

We proposed a regularized probabilistic whitening transformation and examined its statistical properties based on the singular value decomposition of the data matrix. The particular form of the probabilistic formulation enables an implicit covariance approach that avoids direct estimation of the covariance matrix and reduces the computational time from cubic to linear in the number of features. This can reduce computational time by a factor of ≥ 100 in practical applications.

In simulations and applications high-dimensional genotype data from the 1000 Genomes Project, we demonstrated that our GIW-EB method matches or exceeds other methods in out-of-sample whitening benchmarks. In an application to impute z-statistics from a schizophrenia GWAS, GIW-EB gave the best imputation performance. As a proof of principle, we demonstrate that z-statistic imputation powered by our GIW-EB whitening transform can identify tandem repeats explaining the genetic regulatory signals for disease-relevant genes for amyotrophic lateral sclerosis and Alzheimer’s disease.

As the scale of genomic datasets continues to increase, the implicit covariance approach of the probabilistic model will enables application of whitening transforms on high dimensional data where existing methods are intractable. Future applications include single cell gene expression and larger genetic references panels.