Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Summary:
Previous work has shown that sketching the Hessian can improve the running time of second-order optimization methods. However, these prior sketching techniques were data-oblivious. This paper leverages learned sketches to improve the convergence rate of second-order optimization methods.

Strengths: 
+ The learned sketch generally converges better than the other baseline.

Weaknesses:
+ The technical novelty of the paper seems incremental. The learned sketch algorithm from (Liu et al., 2020) serves as a replacement for the data-oblivious sketch in these second-order optimization methods.

Questions:
1) In experiments, is it possible to plot the convergence rate of the non-sketched, full-sized Hessian? I would like to compare the best convergence rate against the other baselines.
2) What is the total running time of the learned sketch against the other baselines?
3) The paper asks whether we should apply the same learned sketch each iteration or train the sketch continuously? However, this question was not thoroughly investigated in the paper. In the experiments, the sketching matrix is learned in each iteration. Is there evidence that supports this choice?

References:
1) Simin Liu, Tianrui Liu, Ali Vakilian, Yulin Wan, and David P. Woodruff. On learned sketches for randomized numerical linear algebra. arXiv:2007.09890 [cs.LG], 2020. URL https://arxiv.org/ abs/2007.09890.

review_title: Official Blind Review #4

final_score: rating: 6: Marginally above acceptance threshold
review_content: SUMMARY:

Sketching is a popular technique in numerical linear algebra for achieving various desirable properties (e.g., lower complexity, one pass methods). The present paper considers a particular kind of sketch for which the sketch matrix is learned from data. It shows how such learned sketches can be used in two types of problems: Hessian sketching (Sec. 3) and Hessian regression (Sec. 4). The authors give both algorithms and provide theoretical guarantees. They also apply these techniques to a number of both synthetic and real datasets in the experiments. For the most part, the experiments indicate that the proposed methods give a consistent, but not necessarily very large, improvement.

I think the idea of using learned sketches is interesting, and it seems like it has not been applied to the problems considered in this paper before. I think the paper strikes a good balance between algorithms, theory and experiments. I like the paper, but there are a few issues that must be addressed before it can be published. Most importantly, there seems to be errors in both Lemma 3.1 and Theorem 3.2. These should hopefully be easy to fix though. 

In the current state, I vote to reject the paper. But if my concerns, especially the issues with Lemma 3.1 and Theorem 3.2, can be addressed, I'll be happy to increase the rating. 


ADVANTAGES:

- Learned sketching is a fairly new and interesting idea.
- It seems like learned sketching has not been used for Hessian sketching and regression before.
- Good balance of algorithms, theory and experiments.


CONCERNS/QUESTIONS:

- In the list that describes performance improvements under "Our Contributions", it is not clear what the quantities $x^*$ and $X^*$ represent. In Sections 5 and 6, these quantities are used to represent the optimal solution to the unsketched problems, but here they seem to represent the solution for the sketched problem using one of the random sketches. If it's the latter, are $x^*$ and $X^*$ the best performing solutions produced by the competing methods? It would be helpful if you clarified this. 

- In Algorithm 1, how is $\alpha$ chosen? Is $\alpha$ updated adaptively, or is it just a fixed small number? Is there any rule of thumb for how to do this? 

- Lemma 3.1: Is seems like the bound on $\hat{Z}_2$ stated in the lemma is wrong. Based on the upper and lower bounds on $Z_2(S)$ towards the end of Section A, it seems like this bound should read
$$
\frac{Z_2(S)}{(1+\eta)^2} - 3 \eta \leq \hat{Z}_2 \leq \frac{Z_2(S)}{(1-\eta)^2} + 3 \eta.
$$

- Lemma 3.1: Since you use that $\max(\eta,\eta^2) = \eta$ when applying the result of Vershynin (2012) in the proof, you should add the condition that $\eta \leq 1$ in the lemma statement.

- Proof of Lemma 3.1: In Section A, 2nd sentence, you say "Since $T$ is a subspace embedding of the column space of $A$...". I think you should add that this is true with probability at least 0.99, to clarify where the 0.99 probability of success in the lemma statement comes from. 

- Proof of Lemma 3.1: In Section A, you use the inverse of $R$ in multiple places. However, it seems like $R$ won't be invertible unless $A$ is of full rank. Can the proof be adapted for a case when $A$ is rank deficient? If not, you should add that $A$ is assumed to be full rank in the lemma statement.

- Proof of Lemma 3.1: This is related to the previous point. In the second to last equation on page 10, it seems like the equation
$$
\min_{x \in S^{d-1}} || S U x || = \min_{y \neq 0} \frac{|| S U W y ||}{|| W y ||}
$$
only will hold if W is full rank, which requires A to be full rank (since $AR^{-1} = UW$ and $U$ has $d$ columns and is full rank). Can the proof be adapted for a case when $A$ is rank deficient? If not, you should add that $A$ is assumed to be full rank in the lemma statement.

- Theorem 3.2: Given the error in Lemma 3.1, the bound in Theorem 3.2 needs to be updated accordingly. Also, in the proof in Section B, in the first inequality you use
$$
\frac{1}{\hat{Z}_1} \geq \frac{1}{\frac{1}{1+\eta} Z_1(S)}.
$$
There's a sign error here; it should read
$$
\frac{1}{\hat{Z}_1} \geq \frac{1}{\frac{1}{1-\eta} Z_1(S)}.
$$

- Proof of Theorem 3.2: In Section B, 2nd sentence, it would be helpful for the reader if you said "holds with probability 0.99" instead of "holds with high probability", and then also clarify that you do a union bound with the 0.99 probability from Lemma 3.1 to get a success probability of at least 0.98 in Theorem 3.2. This may be obvious to readers familiar with the area, but being clear with these things would make the paper accessible to a wider audience.

- Solving (4) deterministically would cost $O(nd^2)$. For Alg. 3 to be worthwhile, it therefore seems like the number $\min(\sigma_1/\sigma_1', \sigma_2/\sigma_2')$ in Theorem 4.1 must be smaller than $d$. Could you say something about why we expect this $\min$ to be small? Is it the case that the $\min$ may be large with some small probability?

- In the experiments, $m$ is chosen as $m = kd$ for some integer $k$. This is different from the $m = O(d^2)$ required for CountSketch to be a subspace embedding. Have you found that this choice $m = O(d)$ always works well in practice, or have you encountered datasets where such a choice has proven to be too small?

- For the experiments in Section 5.2, is there any reason why you only use learned sketches in the first round for the Gaussian and Swarm behavior datasets, rather than use them all rounds as for the other datasets and experiments?

- In Section 6, is there a reason why you choose $\eta = 1$ and $\eta = 0.2$ rather than using the rule for setting $\eta$ in Alg. 3?

- In Section C, you refer to "standard bounds for gradient descent". Can you please provide a reference for those that are unfamiliar with the literature?

- The paper ends abruptly with no conclusion.


MINOR CONCERNS/QUESTIONS:

- On page 2, in the 2nd paragraph, 3rd sentence, you say "If $A$ is tall-and-skinny, then $S$ is wide-and-fat...". Should it be "short-and-fat" rather than "wide-and-fat"?

- The usage of R is a bit confusing. It is used to mean the R matrix in a QR factorization in Alg. 2, the inverse of the R matrix from a QR factorization in Alg. 3, and the upper bound on the nuclear norm in Section 5.3. It would be less confusing if a the inverses in Alg. 3 and the nuclear norm bounds were called something other than R.

- In Section 5.2, for the Swarm behavior and Gisette datasets, you say that $B_i$ is of size $2430 \times 30$ and $5030 \times 30$, respectively. Should this be $2400 \times 30$ and $5000 \times 30$ since $n$ is 2400 and 5000 respectively for the two datasets?

#######################

Update:

The authors have addressed my questions adequately. In particular, my main concerns with the theoretical results and proofs have been fixed. I have updated the score from 3 to 6 for now.

review_title: Review of Learning-Augmented Sketches for Hessians

final_score: rating: 4: Ok but not good enough - rejection
review_content: The paper proposed a learned variant of the well-known iterative Hessian sketch (IHS) method of Pilanci and Wainwright, for efficiently solving least-squares regression. The proposed method is essentially a learned variant of the count-sketch, where the positions of the non-zero entries are random while the value is learned. While getting a learned variant for IHS is an interesting direction, the current theoretical contribution of this paper is only incremental, and most importantly, the reviewer is unconvinced for the practicality of the current approach.

The learned sketching matrix S^t is computed at each iteration, which should introduce a significant computational overhead. Then the plots comparing the learned IHS with unlearned IHS is unfair since it is regarding the iteration count. The iteration count does not reflect such overhead, and the reviewer believe that the wall-clock time comparision would be more sensible.

The reviewer found that the authors are over-selling their numerical results. For example in Figures 1, 2, 6, 7, that the learned IHS has the same linear convergence rate as the unlearn IHS, and actually only slightly faster overall (only 1 iteration faster if we view the plots horizontally), but the authors' claims are like “We can observe that the classical sketches have approximately the same order of error and the learned sketches reduce the error by a 5/6 to 7/8 factor in all iterations”, “We observe that the classical sketches yield approximately the same order of error, while the learned sketches improve the error by at least 30% for the synthetic dataset and surprisingly by at least 95% for the Tunnel dataset.”, which the authors seem to purposefully report from vertical views to exaggerate the results. This is very misleading to the readers which are not familiar with the related works.

********* Update ********

The reviewer appreciates the efforts authors have made. However, the response does not fully address the issues. In the new experiments the authors used small sketch size to show clearer advantages of the learned IHS (Figure 2 & 4), but at that regime all of the methods including learned-IHS are converging only slowly and not practical compare to the slightly larger sketch size choices. The reviewer believes that such a comparison is not meaningful. In order to truly demonstrate the benefits of learned-IHS (which seems to be robust to small sketch sizes), the authors should choose for each algorithm the best sketch size and then compare them in run time, at least between learned-IHS and Count-sketch IHS.

Beside the flaws in numerical experiments, the reviewer found that the theoretical contribution of the current version is incremental. From the reviewer's point of view, the meaningful analysis for learned IHS should certainly be how the converge relates to  the statistics of the training data, to show the benefit of the learned sketch theoretically (i.e. to show how much better the learned sketch SA compare to the unlearned ones in terms of Z_1 and Z_2). The authors have avoided such type of analysis and the main results are "safe-guarded" by the concentration of the random sketch, which are easy to derive.

Overall, the idea of combining IHS with learned sparse sketch is interesting, but the reviewer believes that the current version needs significant amount of rework to be publishable in top conferences.



review_title: An interesting approach for accelerated sketching-based regression solver, but the practicality seems unconvincing, and the numerical results could be misleading

Meta-review: This paper considers convex optimization problems whose solutions involve the solution of linear systems defined in terms of the Hessian. It presents algorithms that reduce the runtime of standard iterative approaches to solving these problems by iteratively sketching the Hessian; the novelty lies in the fact that the authors use the idea of learned sketches which have been used prior for problems in data mining. In particular, the authors use the approach to learning sketches of Liu et al., 2020 to learn the entries in sparse sketching matrices for the Hessian, and propose using the Iterative Hessian Sketch algorithms of Pilanci and Wainright, 2016 to iteratively solve the concerned optimization problem. The advantages of learned sketches are that they may allow using smaller sketch sizes while making progress on the problem, as they are learned to work well on the distribution of Hessians from which the problem instance is drawn.

The consensus of the reviews is that the idea of using learned sketches for convex optimization seems to be novel, and this paper is an interesting attempt, but falls short of the level of contribution required for publication in ICLR. The main concern is that the theory provided for the use of the learned sketches is incremental: the analysis does not reflect the fact that the sketches are learned; instead, the algorithm builds in a safeguard by using both a random sketch and a learned sketch, and the analysis uses the properties of the random sketch to proceed. The empirical results are suggestive, but the convergence rates of the learned and random sketches do not vary much, so the benefits seem marginal for most of the problems considered (with the exception of a standard least squares problem, for which we know learned sketching performs well).

 The paper is recommended to be rejected, as the theory is weak, and the empirical results are borderline. 
Decision: Reject