Bayes beats Cross Validation: Fast and Accurate
Ridge Regression via Expectation Maximization

Ridge regression [25] is one of the most widely used statistical learning algorithms. Given training data ${\bf X}\in\mathbb{R}^{n\times p}$ and ${\bf y}\in\mathbb{R}^{n}$, ridge regression finds the linear regression coefficients $\hat{{\bm{\beta}}}_{\lambda}$ that minimize the $\ell_{2}$-regularized sum of squared errors, i.e.,

In practice, using ridge regression additionally involves estimating the value for the tuning parameter $\lambda$ that minimizes the expected squared error $\mathbb{E}({\bf x}^{\rm T}\!\hat{{\bm{\beta}}}_{\lambda}-y)^{2}$ for new data ${\bf x}$ and $y$ sampled from the same distribution as the training data. This problem is usually approached via the leave-one-out cross-validation (LOOCV) estimator, which can be computed efficiently by exploiting a closed-form solution for the leave-one-out test errors for a given $\lambda$. The wide and long-lasting use of the LOOCV approach suggests that it solves the ridge regression problem more or less optimally, both in terms of its statistical performance, as well as its computational complexity.

However, in this work, we show that LOOCV is outperformed by a simple expectation maximization (EM) approach based on a Bayesian formulation of ridge regression. While the two procedures are not equivalent, in the sense that they generally do not produce identical parameter values, the EM estimates tend to be of equal quality or, particularly in sparse regimes, superior to the LOOCV estimates (see Figure 1). Specifically, the LOOCV risk estimates can suffer from potential multiple and bad local minima when using iterative optimization, or misspecified candidates when performing grid search. In contrast, we show that the EM algorithm finds a unique optimal solution for large enough $n$ (outside pathological cases) without requiring any hard to specify hyper-parameters, which is a consequence of a more general bound on $n$ (Thm. 3.1) that we establish to guarantee the unimodality of the posterior distribution of Bayesian ridge regression—a result with potentially wider applications. In addition, the EM procedure is asymptotically faster than the LOOCV procedure by a factor of $l$ where $l$ is the number of candidate values for $\lambda$ to be evaluated (in the regime $p,q\in O(\sqrt{n})$ where $p$, $q$, and $n$ are the number of covariates, target variables, and data points, respectively). In practice, even in the usual case of $q=1$ and $l=O(1)$, the EM algorithm tends to outperform LOOCV computationally by an order of magnitude as we demonstrate on a test suite of datasets from the UCI machine learning repository and the UCR time series classification archive.

While the EM procedure discussed in this paper is based on a recently published procedure for learning sparse linear regression models [44], the adaption of this procedure to ridge regression has not been previously discussed in the literature. Furthermore, a direct adoption would lead to a main loop complexity of $O(p^{3})$ that is uncompetitive with LOOCV. Therefore, in addition to evaluating the empirical accuracy and efficiency of the EM algorithm for ridge regression, the main technical contributions of this work are to show how certain key quantities can be efficiently computed from either a singular value decomposition of the design matrix, when $p\geq n$, or an eigenvalue decomposition of the Gram matrix ${\bf X}^{\rm T}\!{\bf X}$, when $n>p$. This results in an E-step of the algorithm in time $O(r)$ where $r=\min(n,p)$, and an M-step found in closed form and solved in time $O(1)$, yielding an ultra-fast main loop for the EM algorithm.

These computational advantages result in an algorithm that is computationally superior to efficient, optimized implementations of the fast LOOCV algorithm. Our particular implementation of LOOCV actually outperforms the implementation in scikit-learn by approximately a factor of two by utilizing a similar preprocessing to the EM approach. This enables an $O(nr)$ evaluation for a single $\lambda$ (which is still slower than the $O(r)$ evaluation for our new EM algorithm; see Table 1 for an overview of asymptotic complexities of LOOCV and EM), and may be of interest to readers by itself. Our implementation of both algorithms, along with all experiment code, are publicly available in the standard package ecosystems of the R and Python platforms, as well as on GitHub¹¹1https://github.com/marioboley/fastridge.git.

In the remainder of this paper, we first briefly survey the literature of ridge regression with an emphasis on the use of cross validation (Sec. 2). Based on the Bayesian interpretation of ridge regression, we then introduce the EM algorithm and discuss its convergence (Sec. 3). Finally, we develop fast implementations of both the EM algorithm and LOOCV (Sec. 4) and compare them empirically (Sec. 5).

Ridge regression [25] (also known as $\ell_{2}$-regularization) is a popular method for estimation and prediction in linear models. The ridge regression estimates are the solutions to the penalized least-squares problem given in (1). The solution to this optimization problem is given by:

When $\lambda\to 0$, the ridge estimates coincide with the minimum $\ell_{2}$ norm least squares solution [31, 22], which simplifies to the usual least squares estimator in cases where the design matrix ${\bf X}$ has full column rank (i.e. ${\bf X}^{\rm T}{\bf X}$ is invertible). Conversely, as $\lambda\to\infty$, the amount of shrinkage induced by the penalty increases, with the resulting ridge estimates becoming smaller for larger values of $\lambda$. Under fairly general assumptions [26], including misspecified models and random covariates of growing dimension, the ridge estimator is consistent and enjoys finite sample risk bounds for all fixed $\lambda\geq 0$, i.e., it converges almost surely to the prediction risk minimizer, and its squared deviation from this minimizer is bounded for finite $n$ with high probability. However, its performance can still vary greatly with the choice of $\lambda$; hence, there is a need to estimate the optimal value from the given training data.

Earlier approaches to this problem [e.g. 28, 14, 30, 29, 2, 34, 19, 48, 23, 7] rely on an explicit estimate of the (assumed homoskedastic) noise variance, following the original idea of Hoerl and Kennard [25]. However, estimating the noise variance can be problematic, especially when $p$ is not much smaller than $n$ [18, 20, 24, 46]. More recent approaches adopt model selection criteria to select the optimal $\lambda$ without requiring prior knowledge or estimation of the noise variance. These methods involve minimizing a selection criterion of choice, such as the Akaike information criterion (AIC) [1], Bayesian information criterion (BIC) [40], Mallow’s conceptual prediction (C_p) criterion [33], and, most commonly, cross validation (CV) [4, 49].

A particularly attractive variant of CV is leave-one-out cross validation (LOOCV), also referred to as the prediction error sum of squares (PRESS) statistic in the statistics literature [3]

where $\hat{y}_{i}=\tilde{\bf x}_{i}\hat{\bm{\beta}}^{-i}_{\lambda}$, and $\hat{\bm{\beta}}^{-i}_{\lambda}$ denotes the solution to (2) when the $i$-th data point $(\tilde{\bf x}_{i},y_{i})$ is omitted. LOOCV offers several advantages over alternatives such as 10-fold CV: it is deterministic, nearly unbiased [47], and there exists an efficient "shortcut" formula for the LOOCV ridge estimate [36]:

where ${\bf H}(\lambda)={\bf X}({\bf X}^{\rm T}{\bf X}+\lambda{\bf I}_{p})^{-1}{\bf X}^{\rm T}$ is the regularized “hat”, or projection, matrix and ${\bf e}={\bf y}-{\bf H}(\lambda){\bf y}$ are the residuals of the ridge fit using all $n$ data points. As it only requires the diagonal entries of the hat matrix, Eq. (4) allows for the computation of the PRESS statistic with the same time complexity $O(p^{3}+np^{2})$ as a single ridge regression fit.

Moreover, unless $p/n\to 1$, the LOOCV ridge regression risk as a function of $\lambda$ converges uniformly (almost surely) to the true risk function on $[0,\infty)$ and therefore optimizing it consistently estimates the optimal $\lambda$ [36, 22]. However, for finite $n$, the LOOCV risk can be multimodal and, even worse, there can exist local minima that are almost as bad as the worst $\lambda$ [42]. Therefore, iterative algorithms like gradient descent cannot be reliably used for the optimization, giving theoretical justification for the pre-dominant approach of optimizing over a finite grid of candidates $L=(\lambda_{1},\dots,\lambda_{l})$. Unfortunately, despite the true risk function being smooth and unimodal, a naïvely chosen finite grid cannot be guaranteed to contain any candidate with a risk value close to the optimum. While this might not pose a problem for small $n$ when the error in estimating the true risk via LOOCV is likely large, it can potentially become a dominating source of error for growing $n$ and $p$. Therefore, letting $l$ grow moderately with the sample size appears necessary, turning it into a relevant factor in the asymptotic time complexity.

As a further disadvantage, LOOCV (or CV in general) is sensitive to sparse covariates, as illustrated in Figure 1 where the performance of LOOCV, relative to the proposed EM algorithm, degrades as the noise variance $\sigma^{2}$ grows. In the sparse covariate setting, a situation common in genomics, the information about each coefficient is concentrated in only a few observations. As LOOCV drops an observation to estimate future prediction error, the variance of the CV score can be very large when the predictor matrix is very sparse, as the estimates depend on only a small number of the remaining observations. In the most extreme case, known as the multiple means problem [27], ${\bf X}={\bf I}_{n}$, and all the information about each coefficient is concentrated in a single observation. In this setting, the LOOCV score reduces to $\sum y_{i}^{2}$, and provides no information about how to select $\lambda$. In contrast, the proposed EM approach explicitly ties together the coefficients via the probabilistic Bayesian interpretation of $\lambda$ as the inverse-variance of the unknown coefficient vector. This “borrowing of strength” means that the procedure provides a sensible estimate of $\lambda$ even in the case of multiple means (see Appendix A).

The ridge estimator (2) has a well-known Bayesian interpretation; specifically, if we assume that the coefficients are a priori normally distributed with mean zero and common variance $\tau^{2}\sigma^{2}$ we obtain a Bayesian version of the usual ridge regression procedure, i.e.,

where $\pi(\cdot)$ is an appropriate prior distribution assigned to the variance hyperparameter $\tau^{2}$. For a given $\tau>0$ and $\sigma>0$, the conditional posterior distribution of $\bm{\beta}$ is also normal [32]

where the posterior mode (and mean) $\hat{{\bm{\beta}}}_{\tau}$ is equivalent to the ridge estimate with penalty $\lambda=1/\tau^{2}$ (we rely on the variable name in the notation $\hat{{\bm{\beta}}}_{x}$ to indicate whether it refers to (6) or (2)).

To obtain efficient implementations of the E-Step of the EM algorithm as well as of the LOOCV shortcut formula, one can exploit the fact that the ridge solution is preserved under orthogonal transformations. Specifically, let $r=\min(n,p)$ and $m=\max(n,p)$ and let ${\bf U}{\bm{\Sigma}}{\bf V}^{\rm T}={\bf X}$ be a compact singular value decomposition (SVD) of ${\bf X}$. That is, ${\bf U}\in\mathbb{R}^{n\times r}$ and ${\bf V}\in\mathbb{R}^{p\times r}$ are semi-orthonormal column matrices, i.e., ${\bf U}^{\rm T}{\bf U}={\bf I}_{n}$ and ${\bf V}^{\rm T}{\bf V}={\bf I}_{p}$, and ${\bm{\Sigma}}=\mathrm{diag}(s_{1},\dots,s_{r})\in\mathbb{R}^{r\times r}$ is a diagonal matrix that contains the non-zero singular values ${\bf s}=(s_{1},\dots,s_{r})$ of ${\bf X}$. With this decomposition, and an additional $O(nr)$ pre-processing step to compute ${\bf c}={\bm{\Sigma}}{\bf U}^{\rm T}{\bf y}$, we can compute the ridge solution ${\bm{\alpha}}_{\tau}\in\mathbb{R}^{r}$ for a given $\tau$ with respect to the rotated inputs ${\bf X}{\bf V}$ in time $O(r)$ via

where ${\bf a}\odot{\bf b}$ denotes the element-wise Hadamard product of vectors ${\bf a}$ and ${\bf b}$. The compact SVD itself can be obtained in time $O(mr^{2})$ via an eigendecomposition of either ${\bf X}^{\rm T}{\bf X}={\bf V}{\bm{\Sigma}}^{2}{\bf V}^{T}$ in case $n\geq p$ or ${\bf X}{\bf X}^{\rm T}={\bf U}{\bm{\Sigma}}^{2}{\bf U}^{\rm T}$ in case $n<p$ followed by the computation of the missing ${\bf U}={\bf X}{\bf V}{\bm{\Sigma}}^{-1}$ or ${\bf V}={\bf X}^{\rm T}{\bf U}{\bm{\Sigma}}^{-1}$.

In summary, after an $O(mr^{2})$ pre-processing step, we can obtain rotated ridge solutions for an individual candidate $\tau$ in time $O(r)$. Moreover, for the optimal $\tau^{*}$, we can find the ridge solution $\hat{{\bm{\beta}}}_{\tau^{*}}={\bf V}\hat{{\bm{\alpha}}}_{\tau^{*}}$ with respect to the original input matrix via an $O(pr)$ post-processing step. Below we show how the key statistics that have to be computed per candidate $\tau$ (and $\sigma$) can be computed efficiently based on $\hat{{\bm{\alpha}}}_{\tau}$, the pre-computed ${\bf c}$, and SVD. For the EM algorithm, these are the posterior squared norm and sum of squared errors, and for the LOOCV algorithm, this is the PRESS statistic. While the main focus of this work is the EM algorithm, the fast computation of the PRESS shortcut formula appears to be not widely known (e.g., the current implementation in both scikit-learn and glmnet do not use it) and may therefore be of independent interest.

In this section, we compare the predictive performance and computational cost of LOOCV against the proposed EM method. We present numerical results on both synthetic and real-world datasets. To implement the LOOCV estimator, we use a predefined grid, $L=(\lambda_{1},\ldots,\lambda_{l})$. We use the two most common methods for this task: (i) fixed grid - arbitrarily selecting a very small value as $\lambda_{\rm min}$, a large value as $\lambda_{\rm max}$, and construct a sequence of $l$ values from $\lambda_{\rm max}$ to $\lambda_{\rm min}$ on log scale; (ii) data-driven grid - find the smallest value of $\lambda_{\rm max}$ that sets all the regression coefficient vector to zero ²²2For ridge regression, $\lambda_{\rm max}=\infty$. Following the glmnet package, the sequence of $\lambda$ is derived for $\alpha=0.001$. The penalty function used by glmnet is $\lambda[(1-\alpha)\|{\bm{\beta}}\|^{2}_{2}+\alpha\|{\bm{\beta}}\|_{1}]$, where $\alpha=0$ corresponds to ridge regression. (i.e. $\hat{\bm{\beta}}=0$), multiply this value by a ratio such that $\lambda_{\rm min}=\kappa\lambda_{\rm max}$ and create a sequence from $\lambda_{\rm max}$ to $\lambda_{\rm min}$ on log scale. The latter method is implemented in the glmnet package in combination with an adaptive $\kappa$ coefficient

which we replicate here as input to our fast LOOCV algorithm (Appendix, Table 4) to efficiently recover the glmnet LOOCV ridge estimate. ³³3glmnet LOOCV is computed directly by model refitting via coordinate-wise descent which can be slow.

We consider a fixed grid of ${\bm{\lambda}}=(10^{-10},\ldots,10^{10})$ and the grid based on the glmnet heuristic; in both cases, we use a sequence of length 100. The latter is a data-driven grid, so we will have a different penalty grid for each simulated or real data set. Our EM algorithm does not require a predefined penalty grid, but it needs a convergence threshold which we set to be $\epsilon=10^{-8}$. All experiments in this section are performed in Python and the R statistical platform. Datasets and code for the experimental results is publicly available. As is standard in penalized regression, and without any loss of generality, we standardized the data before model fitting. This means that the predictors are standardized to have zero mean, standard deviation of one, and the target has a mean of zero, i.e., the intercept estimate is simply $\hat{\beta}_{0}=(1/n)\sum y_{i}$.

The introduced EM algorithm is a robust and computationally fast alternative to LOOCV for ridge regression. The unimodality of the posterior guarantees a robust behavior for finite $n$ under mild conditions relative to LOOCV grid search, and the SVD preprocessing enables an overall faster computation with an ultra-fast $O(k\min(n,p))$ main loop. Combining this with a procedure such as orthogonal least-squares to provide a highly efficient forward selection procedure is a promising avenue for future research. As the Q-function is an expected negative log-posterior, it offers a score on which the usefulness of predictors themselves may be assessed, i.e., a model selection criteria, resulting in a potentially very accurate and fast procedure for sparse model learning. An important open problem is the theoretical analysis of the expected number of EM iterations $k$ that is required for convergence. The empirical evidence suggests that $k$ converges to a constant, and is thus negligible in the asymptotic time complexity. This is in alignment with the convergence of the posterior to a multivariate normal distribution. However, such intuitive and empirical arguments cannot replace a rigorous worst-case analysis.

We thank the anonymous reviewers for their valuable feedback that led to a substantially improved theoretical analysis. This work was supported by the Australian Research Council (DP210100045).