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

Thank you for the thorough review of our technical results and constructive feedback. We agree that obtaining non-asymptotic guarantees for this problem would be a very interesting technical challenge for future work. Here, we rely explicitly on a distributional characterization of the iterates, rather than obtaining a state evolution over deterministic scalars. This makes obtaining non-asymptotic guarantees more challenging (although likely possible) after the first 1-2 iterations.

We address your individual questions below. Thanks for catching some typos and the missing CGMT reference - we will fix all of these in our revision. Regarding a few individual questions:

• Boundedness of $\theta^*$: Thanks for this suggestion! Yes, we do believe this kind of truncation argument should allow us to weaken the boundedness requirement on $\theta^*$ to some weaker regularity condition on the distribution. However, we seem to run into some issues when trying to apply a truncation result recursively for each iteration (since we end up needing convergence of the empirical distribution of $(v, \theta^*)$ in a higher order Wasserstein metric). We will continue working on this and update the manuscript if we find a way around it!
• Line 246-249: Thanks for the catch. The prior work [CLOT21] assumes convergence of the initialization in $W_4$ and proves convergence of the estimator in $W_3$, which would not work for a recursive application of the result. We will add this in the revision.

References:

[CLOT21] Chang et al., “Provable benefits of overparameterization in model compression: From double descent to pruning neural networks,” AAAI 2021.

Thank you for the detailed review and constructive feedback. We address your individual questions and concerns below.

(1) “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.)”

We obtain all figures in the paper by performing the iteration in equation 4, which has a fresh batch at each step. We will be sure to specify this explicitly in each figure caption. In general, we do not expect the same exact behavior in the “full-batch” setting (indeed, the total amount of data observed is smaller by a factor of T), but a similar behavior holds when samples are randomly chosen at each iteration (with repetition allowed). Please see part (a) of our global response for more discussion and supplemental simulations for this point.

(2) “Could authors please comment on the technical hurdles if the independence assumption is not available?”

Please see part (a) of our global response for discussion of this point.

(3) “Prior analysis on IRLS does not have this assumption…”

We agree that prior non-asymptotic convergence results for IRLS do not have this assumption. However, we emphasize that the type of result we achieve is substantially different in flavor and relies on a very different set of technical tools. Unlike these works, we aim to provide an exact asymptotic characterization of the distribution of the iterates in the high-dimensional limit. While the tools required for this necessitate stronger assumptions on the data (Gaussian) and algorithm (online/sample-split setting), the resultant guarantees are much stronger and can apply to a wider range of problem settings and algorithms.

(4) “Is the error bound prediction accurate for squared test loss?...”

We have included simulations for the squared loss in Appendix D. Even though the squared loss is not PL(2), we do find that the error prediction from Theorem 1 seems to still be accurate in this case. So, even though our result doesn’t formally hold for squared loss, we still believe our theorem can still be useful for trajectory predictions in this case.

(5) “It seems a little bit weird that alternating minimization exhibits non-monotonic errors in Figure 1b…”

We were also surprised by this behavior! We find that, rather than depending on the loss used ($\ell_1$ or $\ell_2$), this seems to depend on the noise level, with the non-monotonicity appearing more prominently in the low-noise regime (Fig. 1b). We agree that this deserves more mention in the main paper, so we will be sure to mention this in Section 3.1.

(6) “Is it necessary to define pseudo-Lipschitz of order k?”

It is true that we only use the choice k=2 in our analysis – we will simplify the definition accordingly.

(7) “Lines 185-186: What is the precise relation between $\alpha$ and $p$?

Thanks for the question – this reweighting function is equivalent to the one used by the IRLS-p algorithm from [MF12] with the choice $p=2-4\alpha$. The analysis in [DDFG10] also makes an explicit connection between IRLS updates of this form and $\ell_p$-minimization for $0<p\leq1$.

References:

[MF12] Mohan and Fazel, "Iterative reweighted algorithms for matrix rank minimization," JMLR 2012.

[DDFG10] Daubechies et al., "Iteratively reweighted least squares minimization for sparse recovery," Communications on Pure and Applied Mathematics, 2010.

Dear authors, thank you for your reply to my comments.

I have taken a look at the rebuttal and also the comments from other reviewers.

I would like to keep my current, positive score, given the quality of the paper.

Thank you for the detailed review and constructive feedback. We address your individual questions and concerns below.

(1) “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.”

Thanks for the feedback - please see part (b) of our global response for the main take-home messages we hope to convey. In particular, we would like to emphasize the flexibility of our framework in characterizing the test error of several previously proposed algorithms which have been shown to perform well for sparse recovery tasks or to connect with feature learning in linear neural nets. Our results also provide theoretical support for the use of “weight sharing” to learn signals with additional group-sparse structure. We plan to emphasize these points more explicitly in the final discussion section.

(2) “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.”

Thanks for the suggestion - in the revision, we will include the appropriate references.

(3) “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?”

While the CGMT has been used extensively in recent literature to analyze estimation problems, a few new technical challenges arise when applying the result to an entire optimization trajectory, since each step relies on the previous iterate. The proof in [CPT23] relies on the fact that the “auxiliary optimization” obtained via the CGMT can be written in terms of a scalar functional of the previous iterate. However, the LDNN parameterization does not fall into the class of problems covered by this analysis technique. To deal with this, we instead obtain a distributional characterization of the iterates at each step. We do believe the proof technique for obtaining Wasserstein convergence guarantees could extend to studying other iterative algorithms in the sample-split setting. We state a few of these points in the last paragraph of the related work (Lines 141-147), but we will plan on adding a short description of the proof aspects we use which may be of more general interest.

(4) “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 setting that we emphasize here is actually the scenario where $nT < d$, so the total amount of observed data is smaller than the dimensionality. In such settings, effective algorithms must be able to take advantage of the signal sparsity to achieve good performance. While we focus on this more interesting regime, our technical results do not depend on this assumption and in fact hold for any constant value of $\kappa = d/n$.

(5) “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?”

Thanks for the feedback - please see part (b) of our global response about this point. We will be sure to add an additional remark about this in revision.

(6) “line 197: of the algorithm in (4) or in Equation (4).”

Thanks - we will fix this in the revision.

References:

[CPT23] Chandrasekher et al. “Sharp global convergence guarantees for iterative nonconvex optimization with random data,” The Annals of Statistics, 2023.

Thank you for the detailed response.

I have a question regarding the numerical results shown in the attached pdf. I thought that reusing batches is beneficial (in a somewhat different context of training say 2-layer neural nets with SGD) given a recent line of work [R1, R2, R3]. The message there is that to achieve the same error, the sample complexity is smaller (in order) if samples are reused (even twice per sample). So with the same sample size, the error of online / sample-splitting methods should be larger (cf.\ [R4, R5]). However, the experiment in the attached pdf suggests a seemingly opposite message since it compares sampling-splitting with reusing a single batch. This doesn't seem to be a fair comparison and indeed leads to an ``optimistic'' prediction. It seems more standard to compare the performance with and without sample-splitting on the same $nT$ samples.

[R1] Lee, J. D., Oko, K., Suzuki, T., & Wu, D. (2024). Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit. arXiv preprint arXiv:2406.01581.

[R2] Dandi, Y., Troiani, E., Arnaboldi, L., Pesce, L., Zdeborová, L., & Krzakala, F. (2024). The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220.

[R3] Arnaboldi, L., Dandi, Y., Krzakala, F., Pesce, L., & Stephan, L. (2024). Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405.15459.

[R4] Damian, A., Nichani, E., Ge, R., & Lee, J. D. (2024). Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36.

[R5] Arous, G. B., Gheissari, R., & Jagannath, A. (2021). Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22(106), 1-51.