OKRidge: Scalable Optimal k-Sparse Ridge Regression

Summary

This paper proposes a novel branch-and-bound technique called OKRidge to solve sparse Ridge regression problems to learn differential equations. The paper makes 3 contributions: 1) A saddle point formulation of the optimization problem at hand which allows the use of a branch-and-bound algorithm leveraging the special structure (sparsity) of the problem. 2) This efficient branch-and-bound algorithm constructs tight lower-bounds thanks to the saddle-point formulation. 3) A beam search that leverages previous searches for past nodes to speed up the search process of the current node.

Strengths

The paper is well-written and easy to follow. Apart from some details (see question section below) that might be unclear to the non-expert readers, the context of the paper is clear and the literature review is comprehensive. Besides, several extensive experiments have been made. Given the number of hyperparameters in the method ($\lambda$, $\lambda_2$) and the hidden parameters that might impact the convergence speed of the proposed methodology (correlation $\rho$, setup size $n$ and $p$, $\textrm{SNR}$), the sensitivity analysis carried out in the experiments section is welcome.

Weaknesses

The paper does not have major weaknesses. Besides, as a suggestion, since your experiments assess the convergence speed as well as the quality of the sub-optimality gaps yielded by the solvers, you can use the Benchopt framework for more impactful figures: https://github.com/benchopt/benchopt . For instance, figure 1 could be summarized into a one-row figure.

Questions

- As a non-expert in the field of differential equation learning, it is unclear to me why $z_j \in \{0, 1\}$. Is it a generic formulation or does this integer constraint arise often in this field? A clarification in the intro would help the reader understand why it is necessary and why you compare yourselves to MIP solvers. - How does OKRidge scale with respect to the number of features? - "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" (l. 204): how can a Python routine (even very efficient) can beat an algorithm written in C/C++? Are you using a JIT compiler like Numba to compile critical parts of the code? If not, have you tried? It could significantly increase the runtime speed.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

I see one unaddressed limitation to the method: while OKRidge scales easily to hundreds of thousands of samples, it does not seem to scale as well on the number of features. For instance, as stated in l.132, deriving the lower bound requires inverting a $p\times p$ matrix, which does not scale well for $p \gg 10000$.

Summary

The paper proposes a lower bound calculation method and warm start method for sparse ridge regression via a branch-and-bound approach. Starting from a saddle point formulation, the lower bound on the original objective can be given by either a linear system or ADMM problem which can be efficiently solved. The warm-start is given by the beam search method. Experiments are provided to demonstrate the better efficiency of the proposed approach over baselines.

Strengths

1. propose a scalable method for sparse ridge regression. 2. demonstrate the feasibility of the proposed method in practice.

Weaknesses

1. The main contribution lies in efficient certifiable lower bound which however is just part of the whole algorithm. It would be better if the whole algorithm can be explicitly given in the main text. 2. There is no analysis that gives the complexity of the whole algorithm to achieve a acceptable optimality gap.

Questions

See above

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Limitations are addressed in the paper.

Summary

This paper proposes a method, namely OKRidge, to solve spare ridge regression problems and identify terms of differential equations of nonlinear dynamical systems. In particular, OKRidge is formulated as a branch-and-bound algorithm, and this paper provides a lower bound calculation.

Strengths

$\mathbf{1}.$ The proposed method is novel, mainly on the formulation of the optimization problem as prescribed by Eq (12), as well as the approach to solve the problem. $\mathbf{2}.$ The empirical analysis is thorough and results are promising compared with various solvers/algorithms, including commercial solver, e.g., GuRoBi. The experiments very well supports the proposal of the algorithm with adequate evidence. $\mathbf{3}.$ The proposed method is very complete, which covers a major branch-and-bound algorithm with novel lower bound approach and a heuristic to search in a tree.

Weaknesses

$\mathbf{1}.$ Despite the novel formulation of the optimization problem (as well as the lower bound derivation), it might be still combination of known approach to derive the package of approach, including BnB, dynamic programming and heuristics.

Questions

$\mathbf{1}.$ Please see $\mathbf{1}.$ in $\mathbf{Weakness}$ section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

This paper includes a limitation statement, and I do not have any concerns regarding a potential negative societal impact.

Summary

This paper proposes a novel algorithm to solve sparse ridge regression problems -- which are standard ridge regression problems with an additional l0 sparsity constraint (so this problem is NP Hard in general). The algorithm is a branch and bound based approach that uses a novel, tighter, lower bound computed by solving a saddle point reformulation of the problem using ADMM. The reformulation is done using a perspective transform approach along with a reparameterization trick. Additionally, beam search over the chosen non-zero coordinates is used as a heuristic to obtain better upper bounds while expanding the branch and bound tree. A provable guarantee on the optimality gap of the solution obtained from this algorithm is also provided. Experiments on synthetic data as well as data from dynamical systems like the Hopf bifurcation, the Lorenz system and a non-linear magnetohydrodynamical model demonstrate the efficacy of this algorithm compared to existing open source and commericial benchmarks.

Strengths

1. The paper is well very written and clear. 2. I greatly appreciate the detailed and helpful annotations in each step of the proofs as well as the algorithms. 3. The main result is obtained by combining a novel reformulation of the sparse ridge regression problem with ADMM to obtain a tighter lower bound that can be used in branch and bound. 4. Thorough experiments were conducted to study and analyze the performance of the described algorithm.

Weaknesses

1. Perhaps the authors could include experiments on more applications that arise from challenging physical systems beyond the three presented in the paper, since the beam search alone was sufficient for the synthetic problem. 2. The improvement over MIOSR in the case of the sparse differential equations like Lorenz, Hopf and MHD is not very significant (Figure 3), however it still matches the performance of using a proprietary solver.

Questions

1. How does this method compare to relaxing the sparsity constraint to a lasso regularization, and approximating a sparse solution from the solution of the lasso? 2. Have the authors conducted experiments using the simpler (but not as tight) lower bound obtained by solving the ridge regression linear system? If so, were there significant differences between this and the baseline? Further, are there any theoretical guarantees that can be shown for this simpler bound?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Yes, the authors addressed limitations of OKRidge in the case of smaller feature dimensions.

Summary

The authors propose OKRidge, a fast algorithm to solve k-sparse ridge regression problems with provable optimality. The proposed method is a customized branch-and-bound approach, leveraging the specific structure of the ridge regression problem in contrast to black-box solvers such as Gurobi. Namely, OKRidge first uses a new saddle point formulation specifically for efficiently computing lower bounds (either "roughly" by solving a rather simple linear system of more tightly with an ADMM scheme. Second, a beam search heuristic is proposed to find a near-optimal solution, with the global BnB structure in mind. The near-optimality of such solutions is provably guaranteed. Third, relatively straightforward branching and queuing strategies are also proposed to "complete the package", building a full custom solver. Last, the authors benchmark OKRidge against a wide variety of other approaches, not only limited to MIP formulations. Both performance of the solutions, and compute time are evaluated on synthetic benchmarks and on sparse identification of nonlinear dynamical systems.

Strengths

I think that the significance of this work lies in the careful and smart design and integration of multiple components, building into what seems to be a very efficient custom solver for a relatively broad class of problems. While each of this components, taken individually might not be ground-breaking, the method as a whole seems to build into a tool that has important significance. Although I am mostly an outsider and may very well miss some important relevant references, the authors are doing a very good job at placing their contribution in the context of a very abundant existing literature. The clarity of the paper is also to be praised. The appendices, although extremely voluminous (in part due to the hundreds of tables reported from the experiments), are completely self-contained. The proofs of the different theorems read well and are somewhat insightful. The Figures and algorithms summarized in Appendix E are also very helpful to bring all the pieces together. For the quality, although I am not familiar with large parts of the work, I was not able to find mistakes in the proofs and derivations. To be completely fair (maybe due to the very good quality of the presentation), some derivations can be a bit tedious, but seem relatively straightforward (which is not meant to downplay the significance of the contributions in any way).

Weaknesses

As stated in the previous section, it could be argued that none of the contributions, taken in isolation are "novel enough". But I don't believe that would be a fair argument. On the other hand, while the authors did a great job at presenting their method clearly and carefully, I find that it could have been helpful or interesting to include a little more comments on their results. For instance, it is great that provable guarantees are derived for the heuristic solution found by their beam-search (section 3.2.1) but it could have been nice to understand a bit more how tight those bounds are (or if they are insightful at all). More importantly, although the experimental results look good, do the authors consider this heuristic as an important limitation? Are there problems (or datasets) for which the sub-optimality of the beam-search is problematic? Also, it would be great to understand, in practice, where the compute time is being spent as it might be a good way to understand where the computational bottlenecks of OKRidge are and where there is potential room for further improvement.

Questions

I already included a few questions in the previous section. To clarify, I have a few more questions: 1) How do you estimate \lambda_min and \lambda_max, to determine \rho for ADMM? Is it critical to estimate it carefully? 2) It might be a silly question but do you have any insight on how to pick the optimality gap \epsilon in practice? Would it make sense to pick it adaptatively? 3) What about the scenario where you don't "know" k a priori but would like to find which level of sparsity offers a good trade-off? Can OKRidge be easily adapted to solve the problem for different levels and not have to rebuild the solution from scratch for every k? 4) In the end, while learning dynamical systems is an important problem of its own, OKRidge should be useful for other problems which can be cast as a sparse ridge regression. Why so much emphasis on dynamical systems?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Some limitations of OKRidge are clearly identified and pinpointed towards the end of the manuscript. Potential negative societal impacts were not discussed, but I don't foresee any.

Summary

In this manuscript the authors consider the ridge regression with a hard sparsity constraints. Since the problem is posed in the fixed design setting, the paper is on the optimization theory, while the main motivation for its study authors find in the applications related to data-driven methods for dynamical systems. Using the work of Atamturk and Gomez 2020 on max-min relaxation of the original problem, novel lower bounds are developed (one in closed form and the other using ADMM to improve the tightness). These results coupled with Beam-Search of Wiseman and Rush 2016 are used within Bound-and Branch framework to propose an algorithm with a theoretical bound on the optimality gap. Experiments are run on synthetic data to compare with SOTA methods for the studied problem, as well as on datasets from Brunton, Proctor and Kutz 2016 on toy problems in dynamical systems to showcase the application in the sparse identification of (continuous-time deterministic) equations of motion from data under the (strong) assumptions that the finite-dimensional basis of a function space is know and that dataset contains states and their time derivatives.

Strengths

The paper is very well written, apart of issues stated below. All the important concepts are properly introduced in the main body and extensively elaborated upon in the appendix. Main contribution well presented and the novelty of the new ideas w.r.t. existing work is made very clear. Up to my knowledge, relevant references are cited. The theoretical results are correct, except minor issues explained below. Numerical experiments are very detailed.

Weaknesses

In the following I give the main weaknesses of this paper: __(A) Title vs content:__ The paper is content wise purely work on optimization. Namely, the link to ML and to learning of dynamical systems is only through the problem which is an empirical risk minimization for linear model with $\ell_2$ regularization and $\ell_0$ constraints. Hence, there is no really any learning of dynamical systems as the title suggests. Therefore, since I find that the paper is in its current form of interest to ML community, I propose the authors to remove from the title “for Learning Dynamical Systems”. If not, then much more than elaborated discussion needs to be done to justify the title. In particular: 1. Assumptions under which the system is learned from data providing explaining when they are met in practice. My opinion is that relying just on SINDy approach is very limited since this aspect is quite vague. 2. Comparison with other SOTA methods to learn dynamical systems from data is needed (meaning those solving the true problem and not an empirical one) with the discussion on benefits and limitations of the current proposal. Also some insights in how statistical learning theory of the proposed method might be developed. 3. At least one experiment on the performance of the method on the real dataset of challenging dynamical system (unknown non-polynomial nonlinearities in ODE). __(B) Error in Thm. 3.1.:__ When $Q_\lambda$ is singular, i.e. when $\lambda = \lambda_{\min}(X^\top X)$, not every $c$ is in the range of $X^\top$, and, hence, reparametrization doesn’t hold, yielding that (12) is not equivalent to (10). Consequently, when $\sigma_{\min}(X)=0$, e.g. when number of features exceeds number of samples, Thm. 3.1. on which all else is based doesn’t hold. This aspect needs to be checked and clarified. __(C) Notion of optimality:__ In the paper, starting from the title and abstract, the word _optimal_ is used. I believe that many of such uses are not entirely clear. To me it seams that sometimes it means just feasible k-sparse solution. Although upper bound and lower bound guarantees are shown, in the paper I do not find theoretical guarantee for the convergence of the proposed method to the optimal solution stated. Can the authors, please, elaborate on this? To conclude, the current evaluation score I have given reflect the aforementioned issues. Depending on the authors reply, I am willing to reconsider it.

Questions

Please elaborate on the claim in the line 57/58. What is the meaning of outperforms? Seams that on the validation set it performs similarly as MIOSR.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Apart of currently discussed limitations, if the title is kept as is, the authors need more discussion as noted above.

Summary

The authors propose a new framework for computing the k-sparse ridge regression. Specifically, the framework considers the Lagrangian dual and relaxes the lower bound to an ADMM problem with an isotonic proximal operator, which has a linear time solution. Combined with beam search and branch-and-bound, the authors demonstrate the framework's effectiveness in the experiment. Further, they claim to prove a tighter approximation bound based on the beam-search result.

Strengths

The paper's main contribution is the relation between the k-sparse problem with an isotonic proximal operator, which can be computed in linear time ($O(n log k)$ for sorting the coefficients.) I believe this technique is new and may have a significant impact on other works with k-sparse constraints. Actually, the framework can be generalized to any strongly convex and Lipschitz C^1 optimization with k-sparse constraints. The proposed Lagrangian dual lower-bound and isolation of the binary variables only depends on the k-sparse constraints, not the ridge regression itself. In addition, the proofs are clearly written with annotations on every inequality, which helps a lot in verification.

Weaknesses

Despite the proofs being clearly annotated, there are still major bugs in the main theorems. However, I think these bugs may be corrected by weakening the theoretical results, and these errors may not impact the practical value of the work. Specifically, the $L_{\text{ridge-}\lambda}^\text{saddle}$ problem should be a lower bound of equation (12), not an equivalent problem. To prove equivalence, we need to prove that the set of optimality conditions are the same, and the error can be easily checked by taking stationary condition on c in eq (10), which leads to the missing slackness condition $2\beta_j=c_jz_j$ which doesn't appear in the stationary condition of the new formulation in eq (12) unless other assumptions are made. And I think the subtle error in the derivation is that the proof of Theorem 3.1 is not rigorous, i.e., not every step in the L569 derivation is an equivalence relation. The equation (40) may not be one-to-one, and the step of taking $\beta=\gamma$ with additional constraint eq (40), reparameterization of $c$, and moving $\gamma$ from the minimum operator to the maximum operator may not be rigorous. There may be a way to prove it more rigorously, and in the worst case, the derivation becomes a relaxation, which is still helpful. However, we will need to see the more formal derivation first. Second, there are also bugs in the authors' proof for the beam search in Theorem 3.4. In Section 3.2, the authors select the new support by the Gauss-Southwell rule of largest improvement. However, when combined with beam search, it is required to select the candidates with the lowest objective value, which is a conflict. As we see in Algorithm 1 and Algorithm 7 in the Appendix, the actual implementation generates the next step using the B largest improvements for each candidate and chooses the B best candidates with the best objective value, which differs from the proof in Theorem 3.4. On the other hand, the theorem and proofs on the beam search solution are independent to the number of candidates $B$ in the authors' proof, so there must be some error or hidden assumption (e.g., $B\geq p^k$) here. However, I still consider this draft as a good paper and recommend it for acceptance. Just that these bugs need to be addressed, or the Theorems need to be weakened. However, please address these main issues in the rebuttal.

Questions

1. The proposed method is an algorithm for general sparse ridge regression with experiments on the dynamic system. There is no other utilization of the property of the dynamic system here other than the data itself. So I don't think the title is appropriate. 2. Appendix L562: missing z. 3. Appendix L567: The page overflows, so part of the equations are hidden. 4. Appendix L588: The conclusion (Therefore $\gamma=$) is your assumption.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

4 excellent

Limitations

NA

My rating of the paper won't change as it was already very high but thanks a lot for the detailed answers.

For the tightness of the beam search solution bound, I had not even realized that $B$ was not even incorporated. It looks like a difficult problem. Maybe it would be worth mentioning it in the discussion as future work as it seems important and probably relevant to other bounds / problems too. Same comment could hold for the comment on the use of randomized linear algebra tools for speed-up. For my last comment on the emphasis on dynamical systems, I think the motivation was already clearly stated. I guess that my remark eluded more to the points that two other reviewers made, and I agree with them, that apart from potential "strategical reasons" (i.e. make sure that you reach the dynamical systems crowd), not much justifies to have it in the title as nothing in the algorithm or its analysis leverages anything that is specific to it. I agree that it should be removed. (And also, to be clear, I'm not going to leave the decision to you and it won't affect my evaluation / rating of the submission.)

Rising my score

I thank the authors for their detailed replies. After reading other reviews as well as the authors replies to them, I am confident that the the promised steps of the revision will be implemented appropriately and that the paper’s quality will be improved. Hence, I rise my score to accepted.

Comment by kX5s

Thanks for the clarification. I have adjusted my score.