Precise asymptotics of reweighted least-squares algorithms for linear diagonal networks

Summary

This paper studies the trajectory of a certain iterative scheme in the high-dimensional limit. In each step, there're two update steps on a fresh batch: 1) regularized least squares; 2) element-wise nonlinear transformation. This type of iteration is motivated and encompasses several interesting iterations such as IRLS, AGOP, etc. The result is a precise characterization of the empirical distribution of the iterates at any finite time in the proportional limit. This allows one to compute many interesting summary statistics.

Strengths

Though the iteration studied in this paper may look exotic at the first glance, it is well motivated (see the first two pages). The proof uses, unsurprisingly, CGMT since the problem is given in a form that's ready for it. That said, working out the details is not a trivial task. I went through the whole argument in the appendix and got the impression that everything is treated in a fairly rigorous way. In particular, all technicalities associated with CGMT such as justifying the exchange of min / max and proving uniform convergence have been properly addressed. In doing so, the authors have invoked various tricks / auxiliary results commonly used in this literature.

Weaknesses

As the authors have already discussed, one obvious drawback of CGMT is that it needs sample splitting (which may make the iteration under consideration incompatible with its motivating special cases). I actually think that CGMT has the potential to handle repeated batches, or even full batch, though the resulting state evolution is likely to be much less interpretable. So I don't really object to the sample splitting assumption. Also, CGMT is capable of offering non-asymptotic guarantees. I wonder if the authors are interested in doing so (no pressure).

Questions

Minor comments: 1. In the equation between line 66/67, the fresh dataset $\boldsymbol{y}^{(t)}, \boldsymbol{X}^{(t)}$ hasn't been defined (correct me otherwise). 1. Line 49, typo: is --> are 1. Line 164, typo: $t = 0, 1, \dots, T$ 1. Equation between line 165/166, why not add superscript $(t)$ to $\boldsymbol{\epsilon}$? After all, the noise is also fresh and is not shared by other batches. 1. Regarding boundedness of $\boldsymbol{\theta}^*$, as mentioned, it's only used in Lemma 3. Over there, I wonder if it's possible to remove this by truncating $\boldsymbol{\theta}^*$ entry-wise at a large constant $K$, doing the proof (in particular the argument in line 682-692), then sending $K\to\infty$. (To clarify, this is a minor assumption and removing it requires technical work that may not worth it. So I don't mind if the authors leave it as is.) 1. Line 246-249, the authors meant to discuss a second difference from prior works. What's done in this paper is discussed but what's done in prior work is not. So the comparison is unclear. 1. There are two notation sections that are almost identical. I suggest remove one of them. 1. Line 462, what is a "proper function"? Do you mean a "proper convex function"? 1. Line 466, "restate": please put a reference for CGMT. 1. Equation between line 577/578, $V$ should be $V_t$. 1. Equation between line 579/580, $v_i$ should be $v_i^{(t)}$, $V$ should be $V_t$. 1. Line 580, typo: "satisfies".

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

NA

Summary

The paper considers a general algorithm of the form shown in Lines 66-67, as it includes several interesting methods as special cases. The main result is Theorem 1, which proves asymptotic convergence in probability of any function $g\in PL(2)$ to the corresponding expectation. One important special case of this result is the convergence of the test errors of the iterates. Synthetic experiments show that the errors predicted by Theorem 1 align well with the actual performance of the algorithms.

Strengths

- The paper is well-written, easy to follow, and well-motivated. - One interesting contribution is that the paper understands a few different lines of research and is able to unify them using a general algorithm, where the reweighting function can be chosen to implement a specific type of method. - Theorem 1 is a powerful result as it holds for any function in the specific class $PL(2)$ or for any continuous and bounded function. - The proofs do not seem to be trivial to me. While I am not familiar with the proof techniques, the paper illustrates the technical points in the paragraph of Line 240, which makes the high-level proof idea clear and provides cues for interested readers to go deeper into the technical details.

Weaknesses

- The first weakness is the assumption that each batch of data is independent. Personally I feel that this is a very strict assumption as I can not imagine a situation where it is satisfied. For example, almost all optimization algorithms (say SGD) would also reuse the previously seen data. Perhaps it is difficult to remove this assumption, but I think more discussions about it could be beneficial. Specifically: - Could the authors please provide experiments similar to Figure 1 and allow the algorithm to reuse the data? (it is not very clear to me whether Figure 1 takes fresh batches or not.) - Could authors please comment on the technical hurdles if the independence assumption is not available? - Prior analysis on IRLS does not have this assumption, see for example: Theorem 2 of "On the Convergence of IRLS and Its Variants in Outlier-Robust Estimation" and in particular Figure 2a shows the error prediction; Theorem 1 of "Globally-convergent Iteratively Reweighted Least Squares for Robust Regression Problems"; Theorem 3.2 of "Iteratively Reweighted Least Squares for Basis Pursuit with Global Linear Convergence Rate". Would their analysis be useful in any way to remove the independence assumption? - The second weakness is that the experimental evaluations in their current form are not very deep. For example: - Is the error bound prediction accurate for $\ell_2$ squared test loss? The paper commented that $\ell_2$ squared loss does not belong to $PL(2)$, but having this experiment is important as $\ell_2$ squared losses are more commonly seen in my opinion and this experiment can verify whether the assumption of $PL(2)$ is too strong. - It seems a little bit weird that alternating minimization exhibits non-monotonic errors in Figure 1b (purple). Is that because these are test errors, or because the $\ell_1$ loss is a bit different from the training loss (objective function)? How about $\ell_2$ losses or the prediction (residual) errors on test (training) data? All these can be easily verified by experiments, and having them would give a better understanding of the behavior of alternating minimization, and it would also make the analysis more complete and convincing in my humble opinion.

Questions

- Is it necessary to define pseudo-Lipschitz of order $k$? It seems that for most of the cases, $k$ is simply chosen to be $2$. - Lines 185-186: What is the precise relation between $\alpha$ and $p$? It is not very clear why we need two different symbols.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See above.

Summary

In this paper, the authors propose theoretical analysis of the high-dimensional dynamics of the reweighted least-squares methods in the context of "linear diagonal networks". The general algorithm is given in Equation (4) and includes alternating minimization (AM), reparameterized IRLS, and linear recursive feature machines (lin-RFM) as special cases with different choices of reweighting function, see Table 1. In the high-dimensional regime where the observation dimension $d$ and sample size $n$ are both large and comparable, the authors provide, in Theorem 1, distributional characterization of the algorithm behavior for a fixed number of iteration $T$ as $n,d \to \infty$ at the same pace. The result is then extended the multi-variate setting in Theorem 2. Simulations results are provided to support the proposed theoretical analyses.

Strengths

The paper is in general in good shape. It is well written and easy to follow. The proposed analysis is based on Convex Gaussian Min-Max Theorem and improves previous efforts.

Weaknesses

It seems that the problem under study is of interest, but it would be great if the author would better motivate the study and setting, e.g., by providing some take-home messages. I have some detailed comments below

Questions

* line 77 contribution: I suggest that the authors could consider adding pointers to precise theoretical/simulation results when talking about contributions. For example, referring to Theorem 1 for the first contribution. * line 91, perhaps add a paragraph on the related technique tool of Convex Gaussian Min-Max Theorem beyond [7], what is the technical challenge here, and whether the obtained technical results are of more general interest? * line 172-175: the assumption of $nT > d$ is an artifact of the proof, or very intrinsic to the obtained theoretical results? I somehow feel this is a bit unrealistic, or at least, a regime that is not of great interest to study. Could the authors comment on this? Also, this should be stated explicitly in the theorem. * the authors mentioned a few time the keyword "feature learning", but after reading the paper I did not find any in-depth discussion on this. Perhaps having an additional remark to make the connection? minor: * line 197: of the algorithm in (4) or in Equation (4).

Rating

5

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

I do not see any potential negative social impact of this work.