SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiachang Liu, Sam Rosen, Chudi Zhong, Cynthia Rudin
{jiachang.liu, sam.rosen, chudi.zhong}@duke.edu,
cynthiaa@cs.duke.edu
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.
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:
$$
\min  _ {\boldsymbol {\beta}} \| \boldsymbol {y} - \boldsymbol {X} \boldsymbol {\beta} \| _ {2} ^ {2} + \lambda_ {2} \| \boldsymbol {\beta} \| _ {2} ^ {2} \text {s u b j e c t t o} \| \boldsymbol {\beta} \| _ {0} \leq k, \tag {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  $\pmb{y}^T\pmb{y}$ , we can rewrite the loss objective in Equation (1) as:
$$
\mathcal {L} _ {\text {r i d g e}} (\boldsymbol {\beta}) := \boldsymbol {\beta} ^ {T} \boldsymbol {X} ^ {T} \boldsymbol {X} \boldsymbol {\beta} - 2 \boldsymbol {y} ^ {T} \boldsymbol {X} \boldsymbol {\beta} + \lambda_ {2} \sum_ {j = 1} ^ {p} \beta_ {j} ^ {2}, \tag {2}
$$
with  $p$  as the number of features. We are interested in the following optimization problem:
$$
\min  _ {\beta} \mathcal {L} _ {\text {r i d g e}} (\beta) \text {s . t .} (1 - z _ {j}) \beta_ {j} = 0, \sum_ {j = 1} ^ {p} z _ {j} \leq k, z _ {j} \in \{0, 1 \}, \tag {3}
$$
where  $k$  is the number of nonzero coefficients. With the sparsity constraint, the problem is NP-hard. The constraint  $(1 - z_{j})\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(a)$ , the perspective function is  $g(a, b) \coloneqq bf\left(\frac{a}{b}\right)$  for the domain  $b > 0$  [32, 34, 26] and  $g(a, b) = 0$  otherwise. Applying to  $f(a) = 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  $(1 - z_j)\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) = \max_c ac - \frac{c^2}{4} b$ . If we define a new perspective loss as:
$$
\mathcal {L} _ {\text {r i d g e}} ^ {\text {F e n c h e l}} (\boldsymbol {\beta}, \boldsymbol {z}, \boldsymbol {c}) := \boldsymbol {\beta} ^ {T} \boldsymbol {X} ^ {T} \boldsymbol {X} \boldsymbol {\beta} - 2 \boldsymbol {y} ^ {T} \boldsymbol {X} \boldsymbol {\beta} + \lambda_ {2} \sum_ {j = 1} ^ {p} \left(\beta_ {j} c _ {j} - \frac {c _ {j} ^ {2}}{4} z _ {j}\right), \tag {4}
$$
then we can reformulate Problem (3) as:
$$
\min  _ {\boldsymbol {\beta}, \boldsymbol {z}} \max  _ {\boldsymbol {c}} \mathcal {L} _ {\text {r i d g e}} ^ {\text {F e n c h e l}} (\boldsymbol {\beta}, \boldsymbol {z}, \boldsymbol {c}) \text {s . t .} \sum_ {j = 1} ^ {p} z _ {j} \leq k, \quad z _ {j} \in \{0, 1 \}. \tag {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):
$$
\max  _ {\boldsymbol {c}} \min  _ {\boldsymbol {\beta}, \boldsymbol {z}} \mathcal {L} _ {\text {r i d g e}} ^ {\text {F e n c h e l}} (\boldsymbol {\beta}, \boldsymbol {z}, \boldsymbol {c}) \text {s . t .} \sum_ {j = 1} ^ {p} z _ {j} \leq k, \quad z _ {j} \in [ 0, 1 ]. \tag {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.
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$ :
$$
\mathcal {L} _ {\text {r i d g e} - \lambda} (\boldsymbol {\beta}, \boldsymbol {z}) := \boldsymbol {\beta} ^ {T} \boldsymbol {Q} _ {\lambda} \boldsymbol {\beta} - 2 \boldsymbol {y} ^ {T} \boldsymbol {X} \boldsymbol {\beta} + (\lambda_ {2} + \lambda) \sum_ {j = 1} ^ {p} \beta_ {j} ^ {2}, \tag {7}
$$
where  $Q_{\lambda} \coloneqq \mathbf{X}^T \mathbf{X} - \lambda \mathbf{I}$ . We restrict  $\lambda \in [0, \lambda_{\min}(\mathbf{X}^T \mathbf{X})]$ , where  $\lambda_{\min}(\cdot)$  denotes the minimum eigenvalue of a matrix. We see that  $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  $\mathrm{diag}(d)$  in [28] to be  $\lambda 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:
$$
\min  _ {\beta , z} \mathcal {L} _ {\text {r i d g e} - \lambda} (\beta , z) \text {s . t .} (1 - z _ {j}) \beta_ {j} = 0, \sum_ {j = 1} ^ {p} z _ {j} \leq k, z _ {j} \in \{0, 1 \}. \tag {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):
$$
\mathcal {L} _ {\text {r i d g e} - \lambda} ^ {\text {F e n c h e l}} (\boldsymbol {\beta}, \boldsymbol {z}, \boldsymbol {c}) := \boldsymbol {\beta} ^ {T} \boldsymbol {Q} _ {\lambda} \boldsymbol {\beta} - 2 \boldsymbol {y} ^ {T} \boldsymbol {X} \boldsymbol {\beta} + (\lambda_ {2} + \lambda) \sum_ {j = 1} ^ {p} \left(\beta_ {j} c _ {j} - \frac {c _ {j} ^ {2}}{4} z _ {j}\right). \tag {9}
$$
Then, the dual formulation analogous to Problem (6) is:
$$
\max  _ {\boldsymbol {c}} \min  _ {\boldsymbol {\beta}, \boldsymbol {z}} \mathcal {L} _ {\text {r i d g e} - \lambda} ^ {\text {F e n c h e l}} (\boldsymbol {\beta}, \boldsymbol {z}, \boldsymbol {c}) \text {s . t .} \sum_ {j = 1} ^ {p} z _ {j} \leq k, z _ {j} \in [ 0, 1 ]. \tag {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  $c$ , the optimality condition for  $\beta$  is (take the gradient with respect to  $\beta$  and set the gradient to 0):
$$
\boldsymbol {c} = \frac {2}{\lambda_ {2} + \lambda} \left(\boldsymbol {X} ^ {T} \boldsymbol {y} - \boldsymbol {Q} _ {\lambda} \beta\right). \tag {11}
$$
Inspired by this optimality condition, we have the following theorem:
Theorem 3.1. If we reparameterize  $\mathbf{c} = \frac{2}{\lambda_2 + \lambda} (\mathbf{X}^T\mathbf{y} - \mathbf{Q}_\lambda \boldsymbol{\gamma})$  with a new parameter  $\boldsymbol{\gamma}$ , then Problem (10) is equivalent to the following saddle point optimization problem:
$$
\max  _ {\boldsymbol {\gamma}} \min  _ {\boldsymbol {z}} \mathcal {L} _ {r i d g e - \lambda} ^ {s a d d l e} (\boldsymbol {\gamma}, \boldsymbol {z}) \quad s. t. \quad \sum_ {j = 1} ^ {p} z _ {j} \leq k, \quad z _ {j} \in [ 0, 1 ], \tag {12}
$$
where  $\mathcal{L}_{ridge - \lambda}^{saddle}(\pmb {\gamma},\pmb {z}):= -\pmb{\gamma}^T\pmb {Q}_\lambda \pmb {\gamma} - \frac{1}{\lambda_2 + \lambda} (\pmb {X}^T\pmb {y} - \pmb {Q}_\lambda \pmb {\gamma})^T diag(\pmb {z})(\pmb {X}^T\pmb {y} - \pmb {Q}_\lambda \pmb {\gamma}),$  (13)
and  $\operatorname{diag}(z)$  is a diagonal matrix with  $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(\gamma)$  as short-hand for the inner minimization in Problem (12):
$$
h (\boldsymbol {\gamma}) = \min  _ {\boldsymbol {z}} \mathcal {L} _ {\mathrm {r i d g e} - \lambda} ^ {\text {s a d d l e}} (\boldsymbol {\gamma}, \boldsymbol {z}) \quad \text {s . t .} \quad \sum_ {j = 1} ^ {p} z _ {j} \leq k, \quad z _ {j} \in [ 0, 1 ]. \tag {14}
$$
For any arbitrary  $\gamma \in \mathbb{R}^p$ ,  $h(\gamma)$  is a valid lower bound for Problem (3). We should choose  $\gamma$  such that this lower bound  $h(\gamma)$  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(\gamma)$  defined in Equation (14) is lower bounded by
$$
h (\boldsymbol {\gamma}) \geq - \boldsymbol {\gamma} ^ {T} \boldsymbol {Q} _ {\lambda} \boldsymbol {\gamma} - \frac {1}{\lambda_ {2} + \lambda} \| \boldsymbol {X} ^ {T} \boldsymbol {y} - \boldsymbol {Q} _ {\lambda} \boldsymbol {\gamma} \| _ {2} ^ {2}. \tag {15}
$$
Furthermore, the right-hand size of Equation (15) is maximized if  $\gamma = \hat{\gamma} = \operatorname{argmin}_{\alpha} \mathcal{L}_{\mathrm{ridge}}(\alpha)$ , where in this case,  $h(\gamma)$  evaluated at  $\hat{\gamma}$  becomes
$$
h (\hat {\gamma}) = \mathcal {L} _ {\text {r i d g e}} (\hat {\gamma}) + (\lambda_ {2} + \lambda) \operatorname {S u m B o t t o m} _ {p - k} \left(\left\{\hat {\gamma} _ {j} ^ {2} \right\}\right), \tag {16}
$$
where  $\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(\hat{\gamma})$  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$ ,
$$
\mathcal {L} _ {\text {r i d g e}} (\boldsymbol {\gamma}) \geq \mathcal {L} _ {\text {r i d g e}} (\hat {\boldsymbol {\gamma}}) + \nabla \mathcal {L} _ {\text {r i d g e}} (\hat {\boldsymbol {\gamma}}) ^ {T} (\boldsymbol {\gamma} - \hat {\boldsymbol {\gamma}}) + \frac {\mu}{2} \| \boldsymbol {\gamma} - \hat {\boldsymbol {\gamma}} \| _ {2} ^ {2}. \tag {17}
$$
Because  $\hat{\gamma}$  minimizes  $\mathcal{L}_{\mathrm{ridge}}(\cdot)$ , we have  $\nabla \mathcal{L}_{\mathrm{ridge}}(\hat{\gamma}) = 0$ . For the  $k$ -sparse vector  $\gamma$  with  $\|\gamma\|_0 \leq 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
$$
\mathcal {L} _ {\mathrm {r i d g e}} (\boldsymbol {\gamma}) \geq \mathcal {L} _ {\mathrm {r i d g e}} (\hat {\boldsymbol {\gamma}}) + \frac {\mu}{2} \operatorname {S u m B o t t o m} _ {p - k} \left(\left\{\hat {\gamma} _ {j} ^ {2} \right\}\right)
$$
for  $\gamma \in \mathbb{R}^p$  with  $\| \gamma \|_0 \leq k$ . For ridge regression, the strong convexity  $\mu$  parameter can be chosen from  $[0, 2(\lambda_2 + \lambda_{\min}(\boldsymbol{X}^T\boldsymbol{X}))]$ . If we let  $\mu = 2(\lambda_2 + \lambda)$ , we obtain  $h(\hat{\gamma})$  in Theorem 3.2.
The lower bound  $h(\hat{\gamma})$  can be calculated extremely efficiently by solving the ridge regression problem (solving the linear system  $(\mathbf{X}^T\mathbf{X} + \lambda_2\mathbf{I})\boldsymbol{\gamma} = \mathbf{X}^T\mathbf{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(\gamma)$  further based on Equation (14).
Tight Lower Bound via ADMM
Let us define  $\pmb{p} \coloneqq \pmb{X}^T\pmb{y} - \pmb{Q}_{\lambda}\pmb{\gamma}$ . Starting from Problem (12), if we minimize  $\pmb{z}$  in the inner optimization under the constraints  $\sum_{j=1}^{p} z_j \leq k$  and  $z_j \in [0,1]$  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:
$$
- \min  _ {\boldsymbol {\gamma}} (F (\boldsymbol {\gamma}) + G (\boldsymbol {p})) \quad \text {s . t .} \quad \boldsymbol {Q} _ {\lambda} \boldsymbol {\gamma} + \boldsymbol {p} = \boldsymbol {X} ^ {T} \boldsymbol {y}, \tag {18}
$$
where  $F(\pmb{\gamma}) \coloneqq \pmb{\gamma}^T\pmb{Q}_{\lambda}\pmb{\gamma}$  and  $G(\pmb{p}) \coloneqq \frac{1}{\lambda_2 + \lambda} \text{SumTop}_k(\{p_j^2\})$ . 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  $\pmb{q}$  [33]:
$$
\boldsymbol {\gamma} ^ {t + 1} = \underset {\gamma} {\operatorname {a r g m i n}} F (\boldsymbol {\gamma}) + \frac {\rho}{2} \left\| \boldsymbol {Q} _ {\lambda} \boldsymbol {\gamma} + \boldsymbol {p} ^ {t} - \boldsymbol {X} ^ {T} \boldsymbol {y} + \boldsymbol {q} ^ {t} \right\| _ {2} ^ {2} \tag {19}
$$
$$
\boldsymbol {\theta} ^ {t + 1} = 2 \alpha \boldsymbol {Q} _ {\lambda} \boldsymbol {\gamma} ^ {t + 1} - (1 - 2 \alpha) (\boldsymbol {p} ^ {t} - \boldsymbol {X} ^ {T} \boldsymbol {y}) \tag {20}
$$
$$
\boldsymbol {p} ^ {t + 1} = \underset {\boldsymbol {p}} {\operatorname {a r g m i n}} G (\boldsymbol {p}) + \frac {\rho}{2} \| \boldsymbol {\theta} ^ {t + 1} + \boldsymbol {p} - \boldsymbol {X} ^ {T} \boldsymbol {y} + \boldsymbol {q} ^ {t} \| _ {2} ^ {2} \tag {21}
$$
$$
\boldsymbol {q} ^ {t + 1} = \boldsymbol {q} ^ {t} + \boldsymbol {\theta} ^ {t + 1} + \boldsymbol {p} ^ {t + 1} - \boldsymbol {X} ^ {T} \boldsymbol {y}, \tag {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_{\max}(Q_\lambda)\lambda_{\min > 0}(Q_\lambda)}}^1$ , where  $\lambda_{\max}(\cdot)$  denotes the largest eigenvalue of a matrix, and  $\lambda_{\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(\pmb{\gamma}) = \pmb{\gamma}^T\pmb{Q}_{\lambda}\pmb{\gamma}$  and  $G(\pmb{p}) = \frac{1}{\lambda_2 + \lambda} \text{SumTop}_k(\{p_j^2\})$ . Then the solution for the problem  $\pmb{\gamma}^{t+1} = \arg \min_{\pmb{\gamma}} F(\pmb{\gamma}) + \frac{\rho}{2} \| \pmb{Q}_{\lambda}\pmb{\gamma} + \pmb{p}^t - \pmb{X}^T\pmb{y} + \pmb{q}^t \|_2^2$  is
$$
\boldsymbol {\gamma} ^ {t + 1} = \left(\frac {2}{\rho} \boldsymbol {I} + \boldsymbol {Q} _ {\lambda}\right) ^ {- 1} \left(\boldsymbol {X} ^ {T} \boldsymbol {y} - \boldsymbol {p} ^ {t} - \boldsymbol {q} ^ {t}\right). \tag {23}
$$
Furthermore, let  $\mathbf{a} = \mathbf{X}^T\mathbf{y} - \pmb{\theta}^{t + 1} - \pmb{q}^t$  and  $\mathcal{I}$  be the indices of the top  $k$  terms of  $\{|a_j|\}$ . The solution for the problem  $\pmb{p}^{t + 1} = \mathrm{argmin}_{\pmb{p}}G(\pmb {p}) + \frac{\rho}{2}\| \pmb{\theta}^{t + 1} + \pmb {p} - \pmb{X}^T\pmb {y} + \pmb{q}^t\| _2^2$  is  $p_j^{t + 1} = \mathrm{sign}(a_j)\cdot \hat{v}_j$
$$
w h e r e \quad \hat {\boldsymbol {v}} = \underset {\boldsymbol {v}} {\operatorname {a r g m i n}} \sum_ {j = 1} ^ {p} w _ {j} \left(v _ {j} - b _ {j}\right) ^ {2} \text {s . t .} v _ {i} \geq v _ {l} \text {i f} | a _ {i} | \geq | a _ {l} | \tag {24}
$$
$$
w _ {j} = \left\{ \begin{array}{l l} 1 & \text {i f j \notin \mathcal {J}} \\ 1 + \frac {2}{\rho (\lambda_ {2} + \lambda)} & \text {o t h e r w i s e} \end{array} , \right. \qquad b _ {j} = \frac {| a _ {j} |}{w _ {j}}.
$$
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.
[33] also considers matrix preconditioning when computing the step size, but this is computationally expensive when the number of features is large, so we ignore matrix rescaling by letting  $\pmb{E}$  be the identity matrix in Section VI Subsection A of [33].
Starting from the vector  $\mathbf{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 \operatorname {a r g m i n} _ {j} \min  _ {\alpha} \mathcal {L} _ {\text {r i d g e}} (\boldsymbol {\beta} + \alpha \boldsymbol {e} _ {j}) \quad \Longleftrightarrow \quad j ^ {*} \in \operatorname {a r g m a x} _ {j} \frac {\left(\nabla_ {j} \mathcal {L} _ {\text {r i d g e}} (\boldsymbol {\beta})\right) ^ {2}}{\| \boldsymbol {X} _ {: j} \| _ {2} ^ {2} + \lambda_ {2}}, \tag {25}
$$
where  $X_{:j}$  denotes the  $j$ -th column of  $\mathbf{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 \arg \operatorname {B o t t o m B} _ {j} \left(\min  _ {\alpha} \mathcal {L} _ {\text {r i d g e}} \left(\beta + \alpha e _ {j}\right)\right), \tag {26}
$$