Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper shows sketching methods to speed up iterative Hessian sketch (IHS) and Hessian regression. With decently accurate subspaces learned from QR decomposition, IHS and Hessian regression achieve faster convergence than data oblivious methods.
main_review: **Pros:**

- This paper shows provable speed-ups on the second-order methods introduced in previous works, and further demonstrates with experiments.
- The previous works are properly surveyed and discussed in the context of the proposed methods.


**Cons, and questions to authors:**

Despite the above contributions, the writing is ambiguous and does not properly draw relationships among different parts, especially among Section 2, 3 and 4. I would increase my score if the authors could properly address the doubts. 

- How does COUNT-SKETCH introduced in Section 3.1 relate to the methods proposed in Section 4 and 5? It seems both Section 4 and 5 rely on the QR method in Section 3.2; the results in 3.1 are more of standard leverage score sampling results and are irrelevant to the proposed approaches.
- How can the matrix $T$ be quickly chosen in Algorithm 2 Line 7, and what is the relationship between $T$ and $S_i$, and $T$ and $A$? It is not clear from the writing why such a $T$ is readily available. 
- Cost of each iteration: at the bottom of Page 7, the authors claim that “no additional computational cost is incurred in generating $S$ other than solving the iteration step using COUNT-SKETCH”, and Guassian matrices and sparse JL transforms "will be considerably slower". However, the iteration step itself requires running Algorithm 2 or 3, each of which includes QR on the original $A$ matrix (Algorithm 2 Line 8, Algorithm 3 Line 2), which are also expensive. I would suggest the authors compare per-iteration complexities of the approaches shown in Figure 1. Besides, it is not clear how Algorithm 2 Line 10 is implemented.
- Section 3.2 at the bottom of Page 5 says “We found empirically that not squaring this loss function works better than squaring it”. Why does (not-)squaring matter if $S$ is solved by just minimizing the loss $\mathcal{L}(S, A_i)$ itself?


**And some minor issues:**

- The notations at the bottom of Page 5 and in Lemma 3.3 at the beginning of Page 6 are inconsistent: Page 5 says $SA_i = Q_i R_i^{-1}$, whereas Page 6 says $SA = QR$. Also, it would be better if the authors could expand the $\mathcal{L}(S, A_i)$ term at the bottom of Page 5 to explicitly show how S goes into the loss. 
- What are the $A_i$, $b_i$s in Section 6.1 below Equation 7, are they just $A$ and $b$ (the features and labels of the corresponding datasets)?
- It may be nice to show these numerics: the convergence of first-order methods in FIgure 1-4 as a baseline, and how the spectrum of $A$ looks like. 
summary_of_the_review: As mentioned above, despite the contributions, the writing (if not the method itself) is often confusing. I would suggest the authors address these ambiguities in both the author response and the revised version.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors extend a line of work focused on sketching the Hessian for convex problems to help accelerate second order optimization methods. In particular, they present an algorithm for learning weights of a sketching matrix (with one non-zero entry per columns chosen uniformly) by using gradient descent; the Hessians used for training are treated as draws from a distribution of Hessians. They also discuss how leverage scores can be used to improve convergence rates by ensuring that "heavy" rows are sampled with probability 1. 

The authors show empirically that the learned sketches improve convergence rates and reduce the number of iterations for several problems. They also provide theoretical results stating a reduction in rows required in the sketch when heavy leverage scores rows are known and error/time complexity bounds for the Hessian sketch/regression problems.
main_review: ## Strengths:
The paper is well written and gives a nice background of existing literature with strong motivation. The technical detail is of good quality and the numerical experiments provided show that the method performs well on selected data sets when comparing number of iterations required. The idea is simple and intuitively appealing. The problems addressed, specifically least squares, are ubiquitous and finding efficient ways to solve high dimensional problems is of perennial interest. In addition, error/complexity bounds are provided when used in Iterative Hessian Sketch and Fast Regression.


## Weaknesses
The paper claims to provide a framework for learned sketching that applies to a large number of problems in convex optimization but still focuses on least squares problems. It seems that the primary contribution of this paper is to apply learned sketches for least squares problems as introduced in works by Liu et al (_On learned sketches for randomized numerical linear algebra [2020]_ and _Learning the positions in countSketch [2020]_) to iterative hessian sketch and fast regression. Although theoretical bounds are provided, the contribution is more limited than initially stated. 



## Questions and comments:
1. Is it common that the Hessian can be decomposed as $A^T A$ with $A \in \mathbb{R}^{n \times d}$ and $n \gg d$ for other convex problems? This is common for LS but is it often observed elsewhere that makes it more useful for non LS problems.
2. For learning the locations of the non-zero entries as discussed in 3.1, how is the oracle trained to predict these heavy rows? Are the number of occurrences of heavy rows tabulated and the the $k$ most common rows to have leverage scores over some threshold selected? More detail on this front would be helpful. 
3. Although the sketches can be learned offline and finished "within 5 minutes", it is unclear how the sketch training time compares with the time to solve the problem. A comparison of training time and time to convergence would be helpful to evaluate its merit.  



## Minor comments
- $p \in \mathbb{N}^n$ rather than being a real vector
- The notation $|(A_i, b_i)|$ isn't entirely clear. It clear indicates the sixe 
summary_of_the_review: Since the paper primarily focuses on the application of a learned sketch to an iteratively solved LS problem (both components well established elsewhere) the contribution seems marginal. I believe the paper is marginally below the acceptance threshold as is. The paper can be improved by addressing questions and comments above, in particular, evidence that the method is applicable to more general convex problems and a more comprehensive comparison for total time to solve compared to naive sketching.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper considers how learned CountSketches can be used in a variety of optimization tasks. The authors propose a method for predicting rows with high leverage score based on the training data. These rows are then sampled deterministically alongside either a standard CountSketch, or the learned CountSketch introduced by Indyk et al. (2019). Some novel theory is presented for this new sketch. 
main_review: # Strengths

S1. Learning of sketch matrices is an interesting twist to the more standard random sketches out there.

S2. The theory for the learned heavy row sketch is a nice attempt at providing justification for why the proposed learning method might improve results. 

# Weaknesses

W1. It is unclear how useful these learned sketches actually are. In applications, it seems like the learned sketches could easily break. To put it in machine learning terms, the heavy row sketch uses row indices as features to predict whether or not rows are heavy. This is somewhat akin to classifying images on a social networking site based on the ID of the user who uploaded them: if the user with that ID uploaded mostly cat pictures in the training data, then classify the newly uploaded picture as a cat picture. This means that if the rows of a design matrix in the testing data are permuted, then the learned heavy row sketch isn’t useful anymore. This could for example happen in the electricity dataset if two residents switch homes. It therefore seems like, in practice, you would need to retrain the model frequently to ensure that the leverage score estimates remain accurate.

W2. The datasets in the experiments are very small scale, and can be solved very easily using deterministic methods. It would have been more interesting to see the performance on larger scale datasets where these kind of techniques would potentially be more useful. 

W3. The improvements that the learned sketches yield don't seem that substantial in some of the experiments. For example, in Figures 1 and 4-6, it seems like they just save a few iterations compared to doing standard CountSketch.

W4. Some parts of the paper are unclear; see questions below.

# Questions

Q1. How is line 10 in Alg. 2 computed?

Q2. Below Theorem 4.2, you discuss how “a better subspace embedding can lead to a faster convergence.” But it’s not clear from the discussion why this is the case. In particular, it’s not clear why the discussion around improved accuracy guarantees would lead to improved convergence rates. What am I missing?

Q3. Out of curiosity, how do you implement the various CountSketch matrices? Do you form sparse matrices in Python with the relevant nonzeros and multiply, or do you do something more refined to make the computation faster?

Q4. The datasets are poorly introduced. What does $b$ represent in the Electric and GHG datasets? How is the data in $A$ in the GHG dataset organized---is time along the rows with different columns corresponding to different measurement sites? 

Q5. Also, does |$(A,b)$_train|=400 and |$(A,b)$_test|=100  mean that you have a total of 500 pairs of design matrices and corresponding vectors $b$, and that you use 400 for training and the rest for testing? If yes, are all the reported results the average results over all pairs $(A,b)$ in the testing data? 

Q6. It’s not clear why some methods are not given in some plots. Why is “learned(combined)” not plotted in Figure 1 when it appears in Figure 2? Why are “learned(value-only)” and “learned(combined)” not included in Figure 3? Why are “learned(value-only)”, “learned(combined)” and CountSketch not included in Figure 4?

Q7. What is the difference between the three plots in Figure 1? Please explain in the caption. Also, the y-axis labels are cut off and not showing properly. 

Q8. In Section 6.2, why are you using Newton’s method to solve a least squares problem? Is there any benefit to doing that rather than just doing a single sketch-and-solve on the least squares problem itself?

Q9. In the last paragraph before the conclusion, you say that you “only run one iteration” for each subroutine. What does that mean? It looks like you do multiple iterations in Figure 4?

Q10. One of the questions you mention in the introduction is “(2) should we apply the same learned sketch in each iteration, or learn it in the next iteration by training on a data set involving previously learned sketches from prior iterations?” You then say that you will answer this question, but I didn’t see this discussed in the paper. You briefly mention it in the paragraph “Training” in Section 6.2, but I didn’t see it properly addressed anywhere. Did I miss something? 

# Other minor things

- In footnote on page 3, “inut” should be “input”.

- The “Conclusion” heading is too big for a paragraph heading. You should put it as a proper section heading.

- You should specify the unit used in the “runtime” label on the x-axis of Figure 6. I’m assuming it’s seconds?

summary_of_the_review: The core idea of learning sketch matrices is interesting. The theory that provides a potential justification for why learned heavy row sketching might improve on the oblivious CountSketch is welcome. However, the “learned oracle” for the heavy row sketch seems very easy to break: Simply permuting rows would mean that it has to be relearned. Overall, it’s not clear how useful the learned sketch would be in practice for that reason. The experiments are all done on very small problems that don’t need sketching. Parts of the paper are also unclear as outlined above. For these reasons, I think it’s below the acceptance threshold.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper studies matrix sketching algorithms where the sketching matrix is `learned', or generated adaptively. Theoretically, this is done by favoring more the rows of A with high leverage scores, while practically this is done by associating that set with rows of large norms. It evaluates this algorithm on both L_1 and L_2 regression, and observes improved test errors.
main_review: My background is more in theoretical matrix algorithms, so to me, the idea of adaptively generating sketches based on leverage scores is somehow more fundamental than sketching. The authors do do a good job addressing this related literature, but their results are mostly direct consequences of matrix concentration bounds.

The experimental results are interesting: I'm not aware of such evaluations of test errors in previous works. They also point to significant gains of these methods over oblivious sketching based methods. On the other hand, these studies mostly take place on small to moderate sized data where the performance gains from sketching is unclear. So I feel they serve mostly as proof of concept of the utility of such sketching, and perhaps also point out that different formulations/objectives are necessary for giving better test errors.
summary_of_the_review: I feel this paper is well thought out, and the experimental results are highly interesting. However, I also feel that the studies taken here (adaptively select rows) can be taken one step deeper: in particular, it would be very interesting to see theoretical analyses of why such learned sketching give improved results (over oblivious sketches).
review_title: null

Meta-review: This paper proposes a new contribution in the recent literature on learning distributions of sketches. While all reviewers have recognized the overall good quality of the presentation, two factors seem to weight heavily on a negative decision: clarifications on the contribution's scope (presented as a tool for general Hessians in the introduction, but ultimately only applied to least-square errors of linear predictors, to recover an explicit factorization of the Hessian matrix) and links with existing literature; weakness of experiments whose small scale does not justify using sketches in the first place. Since this is a "learning" approach, I am particularly sensitive to the latter point, and therefore am inclined to reject, but I encourage the authors to address these two issues with the current draft.
Decision: Reject