OKRidge: Scalable Optimal k-Sparse Ridge Regression

Abstract

We consider an important problem in scientific discovery, namely identifying sparse governing equations for nonlinear dynamical systems. This involves solving sparse ridge regression problems to provable optimality in order to determine which terms drive the underlying dynamics. We propose a fast algorithm, OKRidge, for sparse ridge regression, using a novel lower bound calculation involving, first, a saddle point formulation, and from there, either solving (i) a linear system or (ii) using an ADMM-based approach, where the proximal operators can be efficiently evaluated by solving another linear system and an isotonic regression problem. We also propose a method to warm-start our solver, which leverages a beam search. Experimentally, our methods attain provable optimality with run times that are orders of magnitude faster than those of the existing MIP formulations solved by the commercial solver Gurobi.

1. Introduction

We are interested in identifying sparse and interpretable governing differential equations arising from nonlinear dynamical systems. These are scientific machine learning problems whose solution involves sparse linear regression. Specifically, these problems require the exact solution of sparse regression problems, with the most basic being sparse ridge regression:

$$\underset{\mathit{\beta }}{\text{min}}\parallel \mathit{y}-\mathit{X}\mathit{\beta }{\parallel }_2^2+{\lambda }_2\parallel \mathit{\beta }{\parallel }_2^2\text{subject}\text{to}\parallel \mathit{\beta }{\parallel }_0\le k,$$ (1)

where $k$ specifies the number of nonzero coefficients for the model. This formulation is general, but in the case of nonlinear dynamical systems, the outcome $y$ is a derivative (usually time or spatial) of each dimension $x$. Here, we assume that the practitioner has included the true variables, along with many other possibilities, and is looking to determine which terms (which transformations of the variables) are real and which are not. This problem is NP-hard [49], and is more challenging in the presence of highly correlated features. Selection of correct features is vital in this context, as many solutions may give good results on training data, but will quickly deviate from the true dynamics when extrapolating past the observed data due to the chaotic nature of complex dynamical systems.

Both heuristic and optimal algorithms have been proposed to solve these problems. Heuristic methods include greedy sequential adding of features [25, 16, 54, 22] or ensemble [31] methods. These methods are fast, but often get stuck in local minima, and there is no way to assess solution quality due to the lack of a lower bound on performance. Optimal methods provide an alternative, but are slow since they must prove optimality. MIOSR [9], a mixed-integer programming (MIP) approach, has been able to certify optimality of solutions given enough time. Slow solvers cause difficulty in performing cross-validation on ${\lambda }_2$ (the ${\ell }_2$ regularization coefficient) and $k$ (sparsity level).

We aim to solve sparse ridge regression to certifiable optimality, but in a fraction of the run time. We present a fast branch-and-bound (BnB) formulation, OKRidge. A crucial challenge is obtaining a tight and feasible lower bound for each node in the BnB tree. It is possible to calculate the lower bound via the SOS1 [9], big-M [10], or the perspective formulations (also known as the rotated second-order cone constraints) [32, 4, 59]; the mixed-integer problems can then be solved by a MIP solver. However, these formulations do not consider the special mathematical structure of the regression problem. To calculate a lower bound more efficiently, we first propose a new saddle point formulation for the relaxed sparse ridge regression problem. Based on the new saddle point formulation, we propose two novel methods to calculate the lower bound. The first method is extremely efficient and relies on solving only a linear system of equations. The second method is based on ADMM and can tighten the lower bound given by the first method. Together, these methods give us a tight lower bound, used to prune nodes and provide a small optimality gap. Additionally, we propose a method based on beam-search [58] to get a near-optimal solution quickly, which can be a starting point for both our algorithm and other MIP formulations. Unlike previous methods, our method uses a dynamic programming approach so that previous solutions in the BnB tree can be used while exploring the current node, giving a massive speedup. In summary, our contributions are:

1. We develop a highly efficient customized branch-and-bound framework for achieving optimality in $k$-sparse ridge regression, using a novel lower bound calculation and heuristic search.

2. To compute the lower bound, we introduce a new saddle point formulation, from which we derive two efficient methods (one based on solving a linear system and the other on ADMM).

3. Our warm-start method is based on beam-search and implemented in a dynamic programming fashion, avoiding redundant calculations. We prove that our warm-start method is an approximation algorithm with an exponential factor tighter than previous work.

On benchmarks, OKRidge certifies optimality orders of magnitude faster than the commercial solver Gurobi. For dynamical systems, our method outperforms the state-of-the-art certifiable method by finding superior solutions, particularly in high-dimensional feature spaces.

2. Preliminary: Dual Formulation via the Perspective Function

There is an extensive literature on this topic, and a longer review of related work is in Appendix A. If we ignore the constant term ${\mathit{y}}^T\mathit{y}$, we can rewrite the loss objective in Equation (1) as:

$${ℒ}_{\text{ridge}}\left(\mathit{\beta }\right)≔{\mathit{\beta }}^T{\mathit{X}}^T\mathit{X}\mathit{\beta }-2{\mathit{y}}^T\mathit{X}\mathit{\beta }+{\lambda }_2\sum _{j=1}^p{\beta }_j^2,$$ (2)

with $p$ as the number of features. We are interested in the following optimization problem:

$$\underset{\mathit{\beta }}{\text{min\hspace{0.17em}}}{ℒ}_{\text{ridge}}\left(\mathit{\beta }\right)\hspace{1em}\text{s.t.}\hspace{1em}\left(1-z_j\right){\beta }_j=0,\sum _{j=1}^p{z}_j\le k,z_j\in \left\{\mathrm{0,1}\right\},$$ (3)

where $k$ is the number of nonzero coefficients. With the sparsity constraint, the problem is NP-hard. The constraint $\left(1-z_j\right){\beta }_j$ in Problem (3) can be reformulated with the SOS1, big-M, or the perspective formulation (with quadratic cone constraints), which then can be solved by a MIP solver. Since commercial solvers do not exploit the special structure of the problem, we develop a customized branch-and-bound framework.

For any function $f\left(a\right)$, the perspective function is $g\left(a,b\right)≔bf\left(\frac{a}{b}\right)$ for the domain $b>0$ [32, 34, 26] and $g(a,b)=0$ otherwise. Applying to $f\left(a\right)=a^2$, we obtain another function $g(a,b)=\frac{a^2}{b}$. As shown by [4], replacing the loss term ${\beta }_j^2$ and constraint $\left(1-z_j\right){\beta }_j=0$ with the perspective formula ${\beta }_j^2/z_j$ in Problem (3) would not change the optimal solution. By the Fenchel conjugate [4], $g(\cdot ,\cdot )$ can be rewritten as $g(a,b)={\text{max}}_{c}ac-\frac{c^2}{4}b$. If we define a new perspective loss as:

$${ℒ}_{\text{ridge}}^{\text{Fenchel}}(\mathit{\beta },\mathit{z},\mathit{c})≔{\mathit{\beta }}^T{\mathit{X}}^T\mathit{X}\mathit{\beta }-2{\mathit{y}}^T\mathit{X}\mathit{\beta }+{\lambda }_2\sum _{j=1}^p\left({\beta }_j{c}_j-\frac{c_j^2}{4}{z}_j\right),$$ (4)

then we can reformulate Problem (3) as:

$$\underset{\mathit{\beta },\mathit{z}}{\text{min}}\underset{\mathit{c}}{\text{max}}{ℒ}_{\text{ridge}}^{\text{Fenchel}}\left(\mathit{\beta },\mathit{z},\mathit{c}\right)\hspace{1em}\text{s.t.}\hspace{1em}\sum _{j=1}^p{z}_j\le k,z_j\in \left\{\mathrm{0,1}\right\}.$$ (5)

If we relax the binary constraint {0, 1} to the interval [0, 1] and swap max and min (no duality gap, as pointed by [4]), we obtain the dual formulation for the convex relaxation of Problem (5):

$$\underset{\mathit{c}}{\text{max}}\underset{\mathit{\beta },\mathit{z}}{\text{min}}{ℒ}_{\text{ridge}}^{\text{Fenchel}}\left(\mathit{\beta },\mathit{z},\mathit{c}\right)\hspace{1em}\text{s.t.}\hspace{1em}\sum _{j=1}^p{z}_j\le k,z_j\in \left[\mathrm{0,1}\right].$$ (6)

While [4] uses the perspective formulation for safe feature screening, we use it to calculate a lower bound for Problem (3). However, directly solving the maxmin problem is computationally challenging. In Section 3.1, we propose two methods to achieve this in an efficient way.

3. Methodology

We propose a custom BnB framework to solve Problem (3). We have 3 steps to process each node in the BnB tree. First, we calculate a lower bound of the node, using two algorithms proposed in the next subsection. If the lower bound exceeds or equals the current best solution, we have proven that it does not lead to any optimal solution, so we prune the node. Otherwise, we go to Step 2, where we perform beam-search to find a near-optimal solution. In Step 3, we use the solution from Step 2 and propose a branching strategy to create new nodes in the BnB tree. We continue until reaching the optimality gap tolerance. Below, we elaborate on each step. In Appendix E, we provide visual illustrations of BnB and beam search as well as the complete pseudocodes of our algorithms.

3.1. Lower Bound Calculation

Tight Saddle Point Formulation

We first rewrite Equation (2) with a new hyperparameter $\lambda$:

$${ℒ}_{\text{ridge}-\lambda }\left(\mathit{\beta },\mathit{z}\right)≔{\mathit{\beta }}^T{\mathit{Q}}_{\lambda }\mathit{\beta }-2{\mathit{y}}^T\mathit{X}\mathit{\beta }+\left({\lambda }_2+\lambda \right)\sum _{j=1}^p{\beta }_j^2,$$ (7)

where ${\mathit{Q}}_{\lambda }≔{\mathit{X}}^T\mathit{X}-\lambda \mathit{I}$. We restrict $\lambda \in \left[0,{\lambda }_{\text{min}}\left({\mathit{X}}^T\mathit{X}\right)\right]$, where ${\lambda }_{\text{min}}(\cdot )$ denotes the minimum eigenvalue of a matrix. We see that ${\mathit{Q}}_{\lambda }$ is positive semidefinite, so the first term remains convex. This trick is related to the optimal perspective formulation [62, 28, 37], but we set the diagonal matrix $\text{diag}\left(\mathit{d}\right)$ in [28] to be $\lambda \mathit{I}$. We call this trick the eigen-perspective formulation. The optimal perspective formulation requires solving semidefinite programming (SDP) problems, which have been shown not scalable to high dimensions [28], and MI-SDP is not supported by Gurobi.

Solving Problem (3) is equivalent to solving the following problem:

$$\underset{\mathit{\beta },\mathit{z}}{\text{min}}{ℒ}_{\text{ridge}-\lambda }\left(\mathit{\beta },\mathit{z}\right)\hspace{1em}\text{s.t.}\hspace{1em}\left(1-z_j\right){\beta }_j=0,\sum _{j=1}^p{z}_j\le k,z_j\in \left\{\mathrm{0,1}\right\}.$$ (8)

We get a continuous relaxation of Problem (3) if we relax {0, 1} to [0, 1].

We can now define a new loss analogous to the loss defined in Equation (4):

$${ℒ}_{\text{ridge}-\lambda }^{\text{Fenchel}}(\mathit{\beta },\mathit{z},\mathit{c})≔{\mathit{\beta }}^T{\mathit{Q}}_{\lambda }\mathit{\beta }-2{\mathit{y}}^T\mathit{X}\mathit{\beta }+\left({\lambda }_2+\lambda \right)\sum _{j=1}^p\left({\beta }_j{c}_j-\frac{c_j^2}{4}{z}_j\right).$$ (9)

Then, the dual formulation analogous to Problem (6) is:

$$\underset{\mathit{c}}{\text{max}}\underset{\mathit{\beta },\mathit{z}}{\text{min}}{ℒ}_{\text{ridge}-\lambda }^{\text{Fenchel}}\left(\mathit{\beta },\mathit{z},\mathit{c}\right)\hspace{1em}\text{s.t.}\hspace{1em}\sum _{j=1}^p{z}_j\le k,z_j\in \left[\mathrm{0,1}\right].$$ (10)

Solving Problem (10) provides us with a lower bound to Problem (8). More importantly, this lower bound becomes tighter as $\lambda$ increases. This novel formulation is the starting point for our work.

We next propose a reparametrization trick to simplify the optimization problem above. For the inner optimization problem in Problem (10), given any $\mathit{c}$, the optimality condition for $\mathit{\beta }$ is (take the gradient with respect to $\mathit{\beta }$ and set the gradient to 0):

$$\mathit{c}=\frac{2}{{\lambda }_2+\lambda }\left({\mathit{X}}^T\mathit{y}-{\mathit{Q}}_{\lambda }\mathit{\beta }\right).$$ (11)

Inspired by this optimality condition, we have the following theorem:

Theorem 3.1. If we reparameterize $\mathit{c}=\frac{2}{{\lambda }_2+\lambda }\left({\mathit{X}}^T\mathit{y}-{\mathit{Q}}_{\lambda }\gamma \right)$ with a new parameter $\gamma$, then Problem (10) is equivalent to the following saddle point optimization problem:

$$\underset{\mathit{\gamma }}{\text{max}}\underset{\mathit{z}}{\text{min}}{ℒ}_{\mathit{\text{ridge}}-\lambda }^{\mathit{\text{saddle}}}\left(\mathit{\gamma },\mathit{z}\right)\hspace{1em}\mathit{\text{s.t.}}\hspace{1em}\sum _{j=1}^p{z}_j\le k,z_j\in \left[\mathrm{0,1}\right],$$ (12)

where

$$\hspace{1em}{ℒ}_{\mathit{\text{ridge}}-\lambda }^{\mathit{\text{saddle}}}(\mathit{\gamma },\mathit{z})≔-{\mathit{\gamma }}^T{\mathit{Q}}_{\lambda }\mathit{\gamma }-\frac{1}{{\lambda }_2+\lambda }{\left({\mathit{X}}^T\mathit{y}-{\mathit{Q}}_{\lambda }\mathit{\gamma }\right)}^T\mathit{\text{diag}}\left(\mathit{z}\right)\left({\mathit{X}}^T\mathit{y}-{\mathit{Q}}_{\lambda }\mathit{\gamma }\right),$$ (13)

and $\mathit{\text{diag}}\left(\mathit{z}\right)$ is a diagonal matrix with $\mathit{z}$ on the diagonal.

To our knowledge, this is the first time this formulation is given. Solving the saddle point formulation to optimality in Problem (12) gives us a tight lower bound. However, this is computationally hard.

Our insight is that we can solve Problem (12) approximately while still obtaining a feasible lower bound. Let us define a new function $h\left(\gamma \right)$ as short-hand for the inner minimization in Problem (12):

$$h\left(\mathit{\gamma }\right)=\underset{\mathit{z}}{\text{min}}{ℒ}_{\text{ridge}-\lambda }^{\text{saddle}}\left(\mathit{\gamma },\mathit{z}\right)\hspace{1em}\text{s.t.}\hspace{1em}\sum _{j=1}^p{z}_j\le k,z_j\in \left[\mathrm{0,1}\right].$$ (14)

For any arbitrary $\gamma \in {\mathbb{R}}^p,h\left(\gamma \right)$ is a valid lower bound for Problem (3). We should choose $\gamma$ such that this lower bound $h\left(\gamma \right)$ is tight. Below, we provide two efficient methods to calculate such a $\gamma$.

Fast Lower Bound Calculation

First, we provide a fast way to choose $\gamma$. The choice of $\gamma$ is motivated by the following theorem:

Theorem 3.2. The function $h\left(\gamma \right)$ defined in Equation (14) is lower bounded by

$$h\left(\gamma \right)\ge -{\gamma }^T{\mathit{Q}}_{\lambda }\mathit{\gamma }-\frac{1}{{\lambda }_2+\lambda }{∥{\mathit{X}}^T\mathit{y}-{\mathit{Q}}_{\lambda }\mathit{\gamma }∥}_2^2.$$ (15)

Furthermore, the right-hand size of Equation (15) is maximized if $\gamma =\hat{\gamma }={\text{argmin}}_{\mathit{\alpha }}{ℒ}_{\text{ridge}}\left(\mathit{\alpha }\right)$, where in this case, $h\left(\gamma \right)$ evaluated at $\hat{\gamma }$ becomes

$$h\left(\hat{\gamma }\right)={ℒ}_{\text{ridge}}\left(\hat{\gamma }\right)+\left({\lambda }_2+\lambda \right){\mathit{\text{SumBottom}}}_{p-k}\left(\left\{{\hat{\gamma }}_j^2\right\}\right),$$ (16)

where ${\mathit{\text{SumBottom}}}_{p-k}(\cdot )$ denotes the summation of the smallest $p-k$ terms of a given set.

Here we provide an intuitive explanation of why $h\left(\hat{\gamma }\right)$ is a valid lower bound. Note that the ridge regression loss is strongly convex. Assuming that the strongly convex parameter is $\mu$ (see Appendix B), by the strong convexity property, we have that for any $\gamma \in {\mathbb{R}}^p$,

$${ℒ}_{\text{ridge}}\left(\gamma \right)\ge {ℒ}_{\text{ridge}}\left(\hat{\gamma }\right)+\nabla {ℒ}_{\text{ridge}}\left(\hat{\gamma }{)}^T\right(\gamma -\hat{\gamma })+\frac{\mu }{2}\parallel \gamma -\hat{\gamma }{\parallel }_2^2.$$ (17)

Because $\hat{\gamma }$ minimizes ${ℒ}_{\text{ridge}}(\cdot )$, we have $\nabla {ℒ}_{\text{ridge}}\left(\hat{\gamma }\right)=0$. For the $k$-sparse vector $\gamma$ with $\parallel \gamma {\parallel }_0\le k$, the minimum for the right-hand side of Inequality (17) can be achieved if ${\gamma }_j={\hat{\gamma }}_j$ for the top $k$ terms of ${\hat{\gamma }}_l^2$ ‘s. This ensures the bound applies for all $k$-sparse $\gamma$. Thus, the $k$-sparse ridge regression loss is lower-bounded by

$${ℒ}_{\text{ridge}}\left(\gamma \right)\ge {ℒ}_{\text{ridge}}\left(\hat{\gamma }\right)+\frac{\mu }{2}{\mathit{\text{SumBottom}}}_{p-k}\left(\left\{{\hat{\gamma }}_j^2\right\}\right)$$

for $\gamma \in {\mathbb{R}}^p$ with $\parallel \gamma {\parallel }_0\le k$. For ridge regression, the strong convexity $\mu$ parameter can be chosen from $\left[\mathrm{0,2}\left({\lambda }_2+{\lambda }_{\text{min}}\left({\mathit{X}}^T\mathit{X}\right)\right)\right]$. If we let $\mu =2\left({\lambda }_2+\lambda \right)$, we obtain $h\left(\hat{\gamma }\right)$ in Theorem 3.2.

The lower bound $h\left(\hat{\gamma }\right)$ can be calculated extremely efficiently by solving the ridge regression problem (solving the linear system $\left({\mathit{X}}^T\mathit{X}+{\lambda }_2\mathit{I}\right)\mathit{\gamma }={\mathit{X}}^T\mathit{y}$ for $\gamma$) and adding the extra $p-k$ terms. However, this bound is not the tightest we can achieve. In the next subsection, we discuss how to apply ADMM to maximize $h\left(\gamma \right)$ further based on Equation (14).

Tight Lower Bound via ADMM

Let us define $p≔{\mathit{X}}^T\mathit{y}-{\mathit{Q}}_{\lambda }\gamma$. Starting from Problem (12), if we minimize $z$ in the inner optimization under the constraints ${\sum }_{j=1}^p{z}_j\le k$ and $z_j\in \left[\mathrm{0,1}\right]$ for $\forall j$, we have $z_j=1$ for the top $k$ terms of $p_j^2$ and $z_j=0$ otherwise. Then, Problem (12) can be reformulated as follows:

$$-\underset{\gamma }{\text{min}}\left(F\left(\gamma \right)+G\left(\mathit{p}\right)\right)\hspace{1em}\text{s.t.}\hspace{1em}{\mathit{Q}}_{\lambda }\gamma +\mathit{p}={\mathit{X}}^T\mathit{y},$$ (18)

where $F\left(\mathit{\gamma }\right)≔{\mathit{\gamma }}^T{\mathit{Q}}_{\lambda }\mathit{\gamma }$ and $G\left(\mathit{p}\right)≔\frac{1}{{\lambda }_2+\lambda }{\mathit{\text{SumTop}}}_k\left(\left\{p_j^2\right\}\right)$. The solution to this problem is a dense vector that can be used to provide a lower bound on the original $k$-sparse problem. This problem can be solved by the alternating direction method of multipliers (ADMM) [17]. Here, we apply the iterative algorithm with the scaled dual variable $q$ [33]:

$${\mathit{\gamma }}^{t+1}=\underset{\gamma }{\text{argmin}}F\left(\gamma \right)+\frac{\rho }{2}{∥{\mathit{Q}}_{\lambda }\gamma +{\mathit{p}}^t-{\mathit{X}}^T\mathit{y}+{\mathit{q}}^t∥}_2^2$$ (19)

$${\mathit{\theta }}^{t+1}=2\alpha {\mathit{Q}}_{\lambda }{\gamma }^{t+1}-(1-2\alpha )\left({\mathit{p}}^t-{\mathit{X}}^T\mathit{y}\right)$$ (20)

$${\mathit{p}}^{t+1}=\underset{\mathit{p}}{\text{argmin}}G\left(\mathit{p}\right)+\frac{\rho }{2}{∥{\mathit{\theta }}^{t+1}+\mathit{p}-{\mathit{X}}^T\mathit{y}+{\mathit{q}}^t∥}_2^2$$ (21)

$${\mathit{q}}^{t+1}={\mathit{q}}^t+{\mathit{\theta }}^{t+1}+{\mathit{p}}^{t+1}-{\mathit{X}}^T\mathit{y},$$ (22)

where $\alpha$ is the relaxation factor, and $\rho$ is the step size.

It is known that ADMM suffers from slow convergence when the step size is not properly chosen. According to [33], to ensure the optimized linear convergence rate bound factor, we can pick $\alpha =1$ and $\rho =\frac{2}{\sqrt{{\lambda }_{\text{max}}\left(Q_{\lambda }\right){\lambda }_{\text{min}>0}\left(Q_{\lambda }\right)}}$¹, where ${\lambda }_{\text{max}}(\cdot )$ denotes the largest eigenvalue of a matrix, and ${\lambda }_{\text{min}>0}(\cdot )$ denotes the smallest positive eigenvalue of a matrix.

Having settled the choices for the relaxation factor $\alpha$ and the step size $\rho$, we are left with the task of solving Equation (19) and Equation (21) (also known as evaluating the proximal operators [52]). Interestingly, Equation (19) can be evaluated by solving a linear system while Equation (21) can be evaluated by recasting the problem as an isotonic regression problem.

Theorem 3.3. Let $F\left(\gamma \right)={\mathit{\gamma }}^T{\mathit{Q}}_{\lambda }\mathit{\gamma }$ and $G\left(\mathit{p}\right)=\frac{1}{{\lambda }_2+\lambda }{\mathit{\text{SumTop}}}_k\left(\left\{p_j^2\right\}\right)$. Then the solution for the problem ${\mathit{\gamma }}^{t+1}={\text{argmin}}_{\mathit{\gamma }}F\left(\mathit{\gamma }\right)+\frac{\rho }{2}{∥{\mathit{Q}}_{\lambda }\mathit{\gamma }+{\mathit{p}}^t-{\mathit{X}}^T\mathit{y}+{\mathit{q}}^t∥}_2^2$ is

$${\mathit{\gamma }}^{t+1}={\left(\frac{2}{\rho }\mathit{I}+{\mathit{Q}}_{\lambda }\right)}^{-1}\left({\mathit{X}}^T\mathit{y}-{\mathit{p}}^t-{\mathit{q}}^t\right).$$ (23)

Furthermore, let $\mathit{a}={\mathit{X}}^T\mathit{y}-{\mathit{\theta }}^{t+1}-{\mathit{q}}^t$ and $𝒥$ be the indices of the top $k$ terms of $\left\{\left|a_j\right|\right\}$. The solution for the problem ${\mathit{p}}^{t+1}={\text{argmin}}_{\mathit{p}}G\left(\mathit{p}\right)+\frac{\rho }{2}{∥{\mathit{\theta }}^{t+1}+\mathit{p}-{\mathit{X}}^T\mathit{y}+{\mathit{q}}^t∥}_2^2$ is $p_j^{t+1}=\mathit{\text{sign}}\left(a_j\right)\cdot {\hat{v}}_j$, where

$$\begin{gathered} \stackrel{\mathit{ˆ}}{\mathit{v}}=\underset{\mathit{v}}{\text{argmin}}\sum _{j=1}^p{w}_j{\left(v_j-b_j\right)}^2\hspace{1em}\mathit{\text{s.t.}}\hspace{1em}{v}_i\ge v_l\mathit{\text{if}}\left|a_i\right|\ge \left|a_l\right| \\ {w}_j=\left\{\begin{array}{ll}1& \mathit{\text{if}}j\notin 𝒥\\ 1+\frac{2}{\rho \left({\lambda }_2+\lambda \right)}& \mathit{\text{otherwise}}\end{array},b_j=\frac{\left|a_j\right|}{w_j}.\right. \end{gathered}$$ (24)

Problem (24) is an isotonic regression problem and can be efficiently solved in linear time [12, 21].

3.2. Beam-Search as a Heuristic

After finishing the lower bound calculation in Section 3.1, we next explain how to quickly reduce the upper bound in the BnB tree. We discuss how to add features, keep good solutions, and use dynamic programming to improve efficiency. Lastly, we give a theoretical guarantee on the quality of our solution.

Starting from the vector 0, we add one coordinate at a time into our support until we reach a solution with support size $k$. At each iteration, we pick the coordinate that results in the largest decrease in the ridge regression loss while keeping coefficients in the existing support fixed:

$$j^{*}\in \underset{j}{\text{argmin}}\underset{\alpha }{\text{min}}{ℒ}_{\text{ridge}}\left(\mathit{\beta }+\alpha {\mathit{e}}_j\right)\hspace{1em}⟺\hspace{1em}{j}^{*}\in \underset{j}{\text{argmax}}\frac{{\left({\nabla }_j{ℒ}_{\text{ridge}}\left(\mathit{\beta }\right)\right)}^2}{{∥{\mathit{X}}_{:j}∥}_2^2+{\lambda }_2},$$ (25)

where ${\mathit{X}}_{:j}$ denotes the $j$-th column of $\mathit{X}$, and the right-hand side uses an analytical solution for the line-search for $\alpha$. This is similar to the sparse-simplex algorithm [6]. However, after adding a feature, we adjust the coefficients restricted on the new support by minimizing the ridge regression loss.

The above idea does not handle highly correlated features well. Once a feature is added, it cannot be removed [61]. To alleviate this problem, we use beam-search [58, 43], keeping the best $B$ solutions at each stage of support expansion:

$$j^{*}\in \text{arg}{\text{BottomB}}_j\left(\underset{\alpha }{\text{min}}{ℒ}_{\text{ridge}}\left(\mathit{\beta }+\alpha {\mathit{e}}_j\right)\right),$$ (26)

where $j^{*}\in$ ${\text{arg}\text{BottomB}}_j$ means $j^{*}$ belongs to the set of solutions whose loss is one of the B smallest losses. Afterwards, we finetune the solution on the newly expanded support and choose the best B solutions for the next stage of support expansion. A visual illustration of beam search can be found in Figure 6 in Appendix E which also contains the detailed algorithm.

Although many methods have been proposed for sparse ridge regression, none of them have been designed with the BnB tree structure in mind. Our approach is to take advantage of the search history of past nodes to speed up the search process for a current node. To achieve this, we follow a dynamic programming approach by saving the solutions of already explored support sets. Therefore, whenever we need to adjust coefficients on the new support during beam search, we can simply retrieve the coefficients from the history if a support has been explored in the past. Essentially, we trade memory space for computational efficiency.

3.2.1. Provable Guarantee

Lastly, using similar methods to [29], we quantify the gap between our found heuristic solution $\hat{\beta }$ and the optimal solution ${\beta }^{*}$ in Theorem 3.4. Compared with Theorem 5 in [29], we improve the factor in the exponent from $\frac{m_{2k}}{M_{2k}}$ to $\frac{m_{2k}}{M_1}$ (since $M_1\le M_{2k}$, where $M_1$ and $M_{2k}$ are defined in [29]).

Theorem 3.4. Let us define a $k$-sparse vector pair domain to be ${\Omega }_k≔\left\{(\mathit{x},\mathit{y})\in {\mathbb{R}}^p\times {\mathbb{R}}^p:\parallel \mathit{x}{\parallel }_0\le \right.\left.k,\parallel \mathit{y}{\parallel }_0\le k,\parallel \mathit{x}-\mathit{y}{\parallel }_0\le k\right\}$. Any $M_1$ satisfying $f\left(\mathit{y}\right)\le f\left(\mathit{x}\right)+\nabla f\left(\mathit{x}{)}^T\right(\mathit{y}-\mathit{x})+\frac{M_1}{2}\parallel \mathit{y}-\mathit{x}{\parallel }_2^2$ for all $(\mathit{x},\mathit{y})\in {\Omega }_1$ is called a restricted smooth parameter with support size 1, and any $m_{2k}$ satisfying $f\left(\mathit{y}\right)\ge f\left(\mathit{x}\right)+\nabla f\left(\mathit{x}{)}^T\right(\mathit{y}-\mathit{x})+\frac{m_{2k}}{2}\parallel \mathit{y}-\mathit{x}{\parallel }_2^2$ for all $(\mathit{x},\mathit{y})\in {\Omega }_{2k}$ is called a restricted strongly convex parameter with support size $2k$. If $\stackrel{\mathit{ˆ}}{\mathit{\beta }}$ is our heuristic solution by the beam-search method, and ${\mathit{\beta }}^{*}$ is the optimal solution, then:

$${ℒ}_{\mathit{\text{ridge}}}\left({\mathit{\beta }}^{*}\right)\le {ℒ}_{\mathit{\text{ridge}}}\left(\stackrel{\mathit{ˆ}}{\mathit{\beta }}\right)\le \left(1-e^{-m_{2k}/M_1}\right){ℒ}_{\mathit{\text{ridge}}}\left({\mathit{\beta }}^{*}\right).$$ (27)

3.3. Branching and Queuing

Branching:

The most common branching techniques include most-infeasible branching and strong branching [2, 1, 15, 7]. However, these two techniques require having fractional values for the binary variables $z_j$’s, which we do not compute in our framework. Instead, we propose a new branching strategy based on our heuristic solution $\stackrel{\mathit{ˆ}}{\mathit{\beta }}$: we branch on the coordinate whose coefficient, if set to 0, would result in the largest increase in the ridge regression loss ${ℒ}_{\text{ridge}}$ (See Appendix E for details):

$$j^{*}=\underset{j}{\text{argmax}}{ℒ}_{\text{ridge}}\left(\stackrel{\mathit{ˆ}}{\mathit{\beta }}-{\hat{\beta }}_j{\mathit{e}}_j\right)$$ (28)

The intuition is that the coordinate with the largest increase in ${ℒ}_{\text{ridge}}$ potentially plays a significant role, so we want to fix such a coordinate as early as possible in the BnB tree.

Queuing:

Besides the branching strategy, we need a queue to pick a node to explore among newly created nodes. Here, we use a breadth-first approach, evaluating nodes in the order they are created.

4. Experiments

We test the effectiveness of our OKRidge on synthetic benchmarks and sparse identification of nonlinear dynamical systems (SINDy)[19]. Our main focus is: assessing how well our proposed lower bound calculation speeds up certification (Section 4.1), and evaluating solution quality of OKRidge on challenging applications (Section 4.2). Additional extensive experiments are in Appendix G and H. Our algorithms are written in Python. Any improvements we see over commercial MIP solvers, which are coded in C/C++, are solely due to our specialized algorithms.

4.1. Assessing How Well Our Proposed Lower Bound Calculation Speeds Up Certification

Here, we demonstrate the speed of OKRidge for certifying optimality compared to existing MIPs solved by Gurobi [35]. We set a 1-hour time limit and an optimality gap of relative tolerance 10⁻⁴. We use a value of 0.001 for ${\lambda }_2$. Our 4 baselines include MIPs with SOS1, big-M $(M=50$ to prevent cutting off optimal solutions), perspective [4], and eigen-perspective formulations ($\lambda ={\lambda }_{\text{min}}\left({\mathit{X}}^T\mathit{X}\right)$) [28]. In the main text, we use plots to present the results. In Appendix G, we present the results in tables. Additionally, in Appendix G, we conduct perturbation studies on ${\lambda }_2\left({\lambda }_2=0.1\right.$ and $\left.{\lambda }_2=10\right)$ and $M(M=20$ and $M=5)$. Finally, still in Appendix G, we also compare OKRidge with other MIPs including the MOSEK solver [3], SubsetSelectionCIO [11], and L0BNB [39].

Similar to the data generation process in [11, 48], we first sample $x_i\in {\mathbb{R}}^p$ from a Gaussian distribution $𝒩(0,\Sigma )$ with mean 0 and covariance matrix $\Sigma$, where ${\Sigma }_{ij}={\rho }^{|i-j|}$. Variable $\rho$ controls the feature correlation. Then, we create the coefficient vector ${\mathit{\beta }}^{*}$ with $k$ nonzero entries, where ${\mathit{\beta }}_j^{*}=1$ if $j$ mod $(p/k)=0.$ Next, we construct the prediction $y_i=x_i^T{\mathit{\beta }}^{*}+{ϵ}_i$, where ${\mathit{ϵ}}_i\stackrel{i.i.d.}{~}𝒩\left(0,\frac{{∥X{\mathit{\beta }}^{*}∥}_2^2}{\mathrm{S}\mathrm{N}\mathrm{R}}\right)$. SNR stands for signal-to-noise ratio (SNR), and we choose SNR to be 5 in all our experiments.

In the first setting, we fix the number of samples with $n=100000$ and vary the number of features $p\in \left\{\mathrm{100,500,1000,3000,5000}\right\}$ and correlation levels $\rho \in \left\{\mathrm{0.1,0.5,0.9}\right\}$ (See Appendix G for $\rho =0.3$ and $\rho =0.7)$. We warm-started the MIP solvers by our beam-search solutions. The results can be seen in Figure 1. From both figures, we see that OKRidge outperforms all existing MIPs solved by Gurobi, usually by orders of magnitude.

Figure 1:

Comparison of running time (top row) and optimality gap (bottom row) between our method and baselines, varying the number of features, for three correlation levels $\rho =\mathrm{0.1,0.5,0.9}(n=100000,k=10)$. Time is on the log scale. Our method is generally orders of magnitude faster than other approaches. Our method achieves the smallest optimality gap, especially when the feature correlation $\rho$ is high.

In the second setting, we fix the number of features to $p=3000$ and vary the number of samples $n\in \left\{\mathrm{3000,4000,5000,6000,7000}\right\}$ and the correlation levels $\rho \in \left\{\mathrm{0.1,0.5,0.9}\right\}$ (see Appendix G for $\rho =0.3$ and $=0.7$). As in the first setting, we also warm-started the MIP solvers by our beam-search solutions. The results are in Figure 2 When $n$ is close to $p$ or the correlation is high $(\rho =0.9)$, no methods can finish within the 1-hour time limit, but OKRidge prunes the search space well and achieves the smallest optimality gap. When $n$ becomes larger in the case of $\rho =0.1$ and $\rho =0.5$, OKRidge runs orders of magnitude faster than all baselines.

Figure 2:

Comparison of running time (top row) and optimality gap (bottom row) between our method and baselines, varying sample sizes, for three correlation levels $\rho =\mathrm{0.1,0.5,0.9}(p=3000,k=10)$. Time is on the log scale. When $\rho =0.1$ and $\rho =0.5$, OKRidge is generally orders of magnitude faster than other approaches. In the case $\rho =0.9$, we achieve the smallest optimality gap as shown in the bottom row.

4.2. Evaluating Solution Quality of OKRidge on Challenging Applications

On previous synthetic benchmarks, many heuristics (including our beam search method) can find the optimal solution without branch-and-bound. In this subsection, we work on more challenging scenarios (sparse identification of differential equations). We replicate the experiments in [9] using three dynamical systems from the PySINDy library [27, 42]: Lorenz System, Hopf Bifurcation, and magnetohydrodynamical (MHD) model [24]. The Lorenz System is a 3-D system with the nonlinear differential equations:

$$dx/dt=-\sigma x+\sigma y,dy/dt=\rho x-y-xz,dz/dt=xy-\beta z$$

where we use standard parameters $\sigma =10,\beta =8/3,\rho =28$. The true sparsities for each dimension are (2, 3, 2). The Hopf Bifurcation is a 2-D system with nonlinear differential equations:

$$dx/dt=\mu x+\omega y-A{x}^3-Ax{y}^2,dy/dt=-\omega x+\mu y-A{x}^{2}y-A{y}^3$$

where we use the standard parameters $\mu =-0.05,\omega =1,A=1$. The true sparsities for each dimension are (4, 4). Finally, the MHD is a 6-D system with the nonlinear differential equations:

$$\begin{gathered} d{V}_1/dt=4V_2{V}_3-4B_2{B}_3,d{V}_2/dt=-7V_1{V}_3+7B_1{B}_2,d{V}_3/dt=3V_1{V}_2-3B_1{B}_2 \\ d{B}_1/dt=2B_3{V}_2-2V_3{B}_2,d{B}_2/dt=5V_3{B}_1-5B_3{V}_1,d{B}_3/dt=9V_1{B}_2-9B_1{V}_2. \end{gathered}$$

The true sparsities for each dimension are (2, 2, 2, 2, 2, 2).

We use all monomial features (candidate functions) up to 5th order interactions. This results in 56 functions for the Lorentz System, 21 for Hopf Bifurcation, and 462 for MHD. Due to the high-order interaction terms, features are highly correlated, resulting in poor performance of heuristic methods.

4.2.1. Baselines and Experimental Setup

In addition to MIOSR (which relies on the SOS1 formulation), we also compare with three common baselines in the SINDy literature: STLSQ [54], SSR [16], and E-STLSQ [31]. The baseline SR3 [25] is not included since the previous literature [9] shows it performs poorly. We compare OKRidge with other baselines using the SINDy library [27, 42]. We follow the experimental setups in [9] for model selection, hyperparameter choices, and evaluation metrics (please see Appendix F for details). In Appendix H, we provide additional experiments on Gurobi with different MIP formulations and comparing with more heuristic baselines.

4.2.2. Results

Figure 3 displays the results. OKRidge (red curves) outperforms all baselines, including MIOSR (blue curves), across evaluation metrics. On the Lorenz System, all methods recover the true feature support when the training trajectory is long enough. When the training trajectory length is short, i.e., the left part of each subplot, (or equivalently, when the number of samples is small), OKRidge performs uniformly better than all other baselines. On the Hopf Bifurcation, all heuristic methods fail to recover the true support, resulting in poor performance. On the final MHD, OKRidge maintains the top performance and outperforms MIOSR on the true positivity rate. This demonstrates the effectiveness of OKRidge, which incurs lower runtimes and yields better metric scores under high-dimensional settings. The highest runtimes are incurred for the MHD (with 462 candidate functions/features), which are shown in Figure 4.

Figure 3:

Results on discovering sparse differential equations. On various metrics, OKRidge outperforms all other methods, including MIOSR which uses a commercial (proprietary) MIP solver.

Figure 4:

Running time comparison between OKRidge and MIOSR on the MHD system with 462 candidate functions. OKRidge is significantly faster than the previous state of the art.

Limitations of OKRidge

When the feature dimension is low (under 100s), Gurobi can solve the problem to optimality faster than OKRidge. This is observed on the synthetic benchmarks $(p=100)$ and also on the Hopf Bifurcation $(p=21)$. Since Gurobi is a commercial proprietary solver, we cannot inspect the details of its sophisticated implementation. Gurobi may resort to an enumeration/brute-force approach, which could be faster than spending time to calculate lower bounds in the BnB tree. This being said, OKRidge is still competitive with Gurobi in the low-dimensional setting, and OKRidge scales favorably in high-dimensional settings.

5. Conclusion

We presented a method for optimal sparse ridge regression that leverages a novel tight lower bound on the objective. We showed that the method is both faster and more accurate than existing approaches for learning differential equations ‒ a key problem in scientific discovery. This tool (unlike its main competitor) does not require proprietary software with expensive licenses and can have a significant impact on various regression applications.