Denoising Low-Rank Data Under Distribution Shift: Double Descent and Data Augmentation

We thank the reviewer for the in-depth comments about our work and for agreeing with us that removing the Gaussian-like assumptions is one of our major contributions. If the reviewer feels that our series of responses adequately addresses their concerns, we hope that they will consider revising their score.

We would first like to clarify a misconception about the relative size of the signal and noise.

From assumption 1.2, we have that for $X_{trn} \in \mathbb{R}^{d \times N}$, we have that $\|\|X_{trn}\|\|_F^2= O(N)$. This implies that each data point has norm norm $O(1)$.

From assumption 2, we have that for $A_{trn} \in \mathbb{R}^{d \times N}$, $\mathbb{E}[A_{ij}^2] = \eta^2/d$. However, this means that for the whole matrix $A_{trn}$ we have that

$\mathbb{E}\left[\| \|A\_{trn}\|\|\_F^2\right] = \sum\_{i=1}^d \sum\_{j=1}^N \mathbb{E}\left[A\_{ij}^2\right] = d N \eta^2/d = \eta^2 N.$

Further, from assumption 2, we assume that $\eta^2 = \Theta(1)$.

Hence we have that the norm of the data matrix and the norm of the noise matrix are on the same scale. Hence we do not have vanishing noise.

In fact, based on our assumptions we have that $\mathbb{E}\left[ \|\|A_{trn}\|\|_F^2\right] = \Theta(N)$. However, we only have that $\|\|X\_{trn}\|\|\_F^2 = O(N)$. In particular, we allow for sublinear growth. Hence are applicable to situations where we have diminishing signal.

We picked the scaling so that data points and noise vectors have constant norm as $d, N \to \infty$ rather than having both norms go to $\infty$ as is the case in many prior works. Hence, due to this, the errors go to 0. However, if we undo the scaling, so multiply everything by $d$, we would have the asymptotic error be non-zero and would depend on $c$. This is not due to vanishing noise. It is due to the scaling we use for the problem, which is the same for both the signal and the noise.

We thank the reviewer for their positive comments. We agree with the reviewer that we offer a wide variety of theory results and that the difference in generalization between the under-parameterized and the over-parameterized regimes is fascinating. If the reviewer feels that our series of responses adequately addresses their concerns, we hope that they will consider revising their score.

We thank the reviewer for their interest in Corollary 5. We have updated the paper to mention Corollary 5 in the main paper.

We now address the reviewer's concerns.

Explanation of the terms in Theorem 1.

We thank the reviewer for pointing out this issue. We have added the following clarification. Our generalization error is given by

$\mathbb{E}\_{A\_{trn}, A\_{tst}}\left[\frac{\|Y\_{tst} - W\_{opt}(X\_{tst} + A\_{tst})\|^2}{N\_{tst}}\right]$

We can expand this into

$\frac{1}{N_{tst}}\left( \|Y_{tst} - W_{opt} X_{tst}\|^2 + \|W_{opt}A_{tst}\|^2 - 2Tr((Y_{tst}-W_{opt}X_{tst})^TW_{opt}A_{tst})\right).$

Due to the assumptions that entries of $A_{tst}$ have mean 0, when we compute the expectation, the cross term vanishes. Then, we are left with two terms. The first is $\frac{1}{N_{tst}} \left( \|Y_{tst} - W_{opt} X_{tst}\|^2\right)$, which we think of as the bias and the second is $\frac{1}{N_{tst}}\|W_{opt}A_{tst}\|^2$, which we think of as the variance.

Relating this to Theorem 1 (let us ignore the big-O and little-o terms), we have two terms for both the over and under-parameterized regimes. The first term (that depends on $L$ is the bias term. The second term (that is independent of the $L$) is the variance term.

The extra term in the overparameterized case appears due to a technical reason. In the under-parameterized case, certain products of matrices cancel nicely and either become 0 or the identity. Hence, we do not need to estimate them using random matrix theory. However, this is not the case for the overparameterized case. Unfortunately, we get a term for which the error between the true value and our estimate is $O(\|\Sigma\|^/N^2)$.

Proof Sketch

The proof is broken down into the following steps.

1. Obtain a formula for $W_{opt}$. In particular, we know that this $W_{opt} = X(X+A)^+$. However, this step expands $(X+A)^+$ and simplifies the expressions. (Lemma 1 for the overparameterized and Lemma 9 for the under-parameterized case)
2. Substitute our expression into our expressions for the bias and the variance. Then, we can compute these norms as the trace of various matrices. (Lemma 2, 10 for the bias, the variance is done in the proof of the main theorem)
3. This is the most challenging step - show that the variety of the products concentrate to diagonal matrices. This is quite difficult, and Lemma 3,4,5,6,7,11,12,13,14,15 are proving these concentration results.

Appendix organization.

We thank the reviewer for pointing this out. We have reorganized the appendix. We have also added more details to Appendix C.

Overfitting incorrect

We thank the reviewer for pointing this out. However, we believe the error is a typo (not a fundamental one). The result has been updated. In the theorem for the underparameterized case, we were missing a $c$. So the relative error is $c/(1-c)$, which approaches 0 as $c \to 0$. Also, we see in the highly overparameterized case where we have a relative error of $1/(c-1)$, which also approaches 0 as $c \to \infty$. This is the highly overparameterized case and the highly underparameterized case both benignly overfit in this case, and we see the more complex tempered overfitting in the proportional regime.

Questions

Theorem 2 bound

Yes, that is correct. We have fixed it.

c vs 1/c

We did it because we test on real data. Recall that $c = d/N$. However, for real data, $d$ is fixed (for CIFAR d is 3332=3072), and is something that we cannot change. So, we vary $c$ by varying $N$ and plotted our graphs against $1/c = N/d$. This reads naturally once one notices that $N$ was the varied quantity.

Figure 1a

This is not correct. Figure 1a, the solid line connects the theoretical values. If the reviewer looks closely (zooms in on the old pdf), the scatter points and the line do not exactly match. There are close, but it is not exact. However, to avoid this confusion. We have sampled more theory points and plotted a smoother version of the theory curves, making it more evident that the solid lines are always theory.

Figure 1c vs Figures 9 and 10

For Figures 9 and 10, we are doing this for distributions in our low-rank subspace. Hence, the curves exactly line up. Figure 1c is for an experiment without the exact low-rank assumptions needed by our results. Hence, there is a mismatch.

—-------------

We thank the reviewer again. With the above clarifications and the changes to the paper, the reviewer is willing to accept the paper.

Let us continue to build some more intuition.

Why do things concentrate

Let $A$ be a $p \times q$ matrix with IID entries that have mean 0 and variance 1 and bounded fourth moment. Then if we look at $\frac{1}{q} AA^T$, let $\lambda_1, \ldots, \lambda_k$ be the eigenvalues, and let $\mu_{p,q}$ be the uniform distribution on these eigenvalues.

Note here the randomness in $\mu_{p,q}$ is over which eigenvalue of $AA^T$ we are picking AND NOT over $A$

Then the Marchenko Pastur theorem tells us that $\mu_{p,q}$ converges to marchenko pastur weakly. Hence this tells us that for EVERY A, this converges, so we have pathwise convergence instead of in probability or expectation.

Then we saw above the expressions of the form $\sum_{i=1}^N \frac{1}{N} \frac{1}{\sigma_i^2}$ can be written as $\mathbb{E}\_{\mu\_{d,N}}\left[\frac{1}{\sigma^2}\right]$ where the expectation is taken not over $A$ but over this $\mu_{d,N}$ measure. Hence, everything concentrates.

Difference between $c>1$ and $c < 1$

Continuing with the notation from above. Suppose $p < q$, then we notice that $AA^T$ is $p \times p$. It can be shown that $A$ is full rank.

However, in our work, for one case we see $AA^T/q$ and in the other $A^TA/q$ (Recall $P$, $Q$ terms at the very beginning of the previous comment).

However, the Marchenko Pastur theorem doesn't apply to $A^TA/q$, it doesn't have the correct normalization. So to correct we need to multiply by $q/p$. For our results $q = d, p = N$. Hence, we are multiplying by $c$.

This is the second reason that the powers of $c$ are different

Understanding $L$

$L$ is the cooridnates of $X_{tst}$ in the basis given by $U$. Then suppose $X_{tst} = U_{tst} \Sigma_{tst} V_{tst}^T$. Then we see that $L = U^TU_{tst} \Sigma_{tst} V_{tst}^T$. Then, using the unitary invariance of the norm, we have that

$\|\|\beta\_U^T( \Sigma\_{trn}^2c + \eta\_{trn}^2 I)^{-1} L \|\|\_F^2 = \|\|\beta\_U^T( \Sigma_{trn}^2c + \eta\_{trn}^2 I)^{-1} U^TU_{tst} \Sigma_{tst} \|\|\_F^2$

Hence we see that the error depends on an alignment term between the training and test data $U^TU_{tst}$ and the singular values $\Sigma_{tst}$.

Why no $V$ terms

Since we are doing linear least squares problems, we see that everything depends on the Gram matrix of the data (this is also known as the kernel trick). Hence, we only get dependence on $U$ and not on $V$. In the classical regime, this would mean we depend on the eigenvectors and eigenvalus of the covariance matrix and not on anything else.

Different powers of $c$

This is due to two reasons

1. The lack of symmetry the expressions for $W_{opt}$ (See previous reply)
2. The need to multiply by $d/N$ to get convergence (see the section on concentration in this reply).

We hope that the above helps better understand the proof. If there are steps that are unclear or parts where the reviewer would like more detail, we are happy to clarify and provide more details.

We thank the reviewer for their detailed and in-depth review. We also agree with the reviewer that the non-asymptotic validation with real data is a strength of our paper. If the reviewer feels that our series of responses adequately addresses their concerns, we hope that they will consider revising their score.

Here are some short responses to the weaknesses before we delve deeper:

Weakness 1: As a punchline, using the spectrum of the noise allows us to work with possibly dependent data. However, as discussed in the paper, most results already make assumptions that encompass our assumptions on the noise. The fact that we only use assumptions about the noise is a strength of the work, not a weakness. Being able to handle such a broad setting comes with many complications, discussed in detail below.

Weakness 2: It is already hard to analyze the inverse of a sum of random matrices (X+A)^{-1}. This is, in fact, a fundamental issue with analyzing input noise settings. Our idea is to use matrix inverse perturbation results, which comes with several complications already. The problem is further complicated by having to deal with pseudo-inverses, not only complicated by that. We provide more details below.

We will now explain why deriving our test error results is not easy. The difficulty in the proof comes from the placement of the noise and when the expectation is taken. Both of these are crucial.

Placement of the noise

In our setup, we have noise on the input data instead of the output. However, for comparison, let us briefly consider the situation where we have noise on the output.

The case of output noise

When the noise is on the output, we have some training data $X$ and $y = \beta^TX + \epsilon$ where $\epsilon$ is noise. In this situation, if we solve the least squares problem, we have our estimate:

$\hat{\beta}^T = yX^+ = \beta^TXX^+ + \epsilon X^+$

Now we have two terms. The first only depends on $X$ and not the noise. Hence, this is straightforward to analyze. The second term is linear in the noise.

Continuing from here, suppose we care about $\|\|\beta - \hat{\beta}\|\|^2 = \|\beta\|^2 + \|\hat{\beta}\|^2 - 2\beta^T\hat{\beta}$ (as is the case in prior work). We see that $\|\beta\|$ is a constant.

Next $\|\hat{\beta}\| = \|\beta^TXX^+\|^2 + \|\epsilon X^+\|^2 + 0$ where the cross term is zero due to the independence of the noise. Here $XX^+$ is roughly constant, so it is easy to analyze, and if $\epsilon$ is Gaussian, then $\epsilon X^+$ is Gaussian, and the norm is easy to calculate.

For the cross term $\beta^T\beta = \beta^TXX^+\beta + \beta^T(X^+)^T\epsilon^T$. The second term is zero due to the independence of $\epsilon$, and the first is a quadratic form dependent on $XX^+$.

Comparison to our setup (noisy inputs)

Let us now compare this to our setup. Since we have noise on the input, our estimate is

$\hat{\beta} = \beta^T X (X+A)^+$

To start off, we already cannot decouple the noise and the data. Hence, this is one challenge that we had to overcome. In the output noise setup, since the two decouple, we only need to understand the spectrum of $X$. In our case, we need to understand the spectrum of $X+A$.

In the classical regime ($d$ fixed, $N$ to infinity), this would be easier since $(X+A)(X+A)^T$ would converge (with appropriate normalization) to the covariance matrix $XX^T$. However, there are two issues in our problem that complicate this:

1. First, we are not in the classical regime. In the proportional regime, as noted by Cheng and Montanari and many other prior works, $(X+A)(X+A)^T$ is not a consistent estimate for the covariance matrix.

2. Second, we actually do not have terms of the form $(X+A)(X+A)^T$ (as was the case in the output noise situation). We have terms of the form $X(X+A)^T$. While this is subtle, it makes a significant difference.

Points 1 and 2 combined imply that we must understand the eigenstructure of $X+A$. Specifically, point two asks us to study $X(X+A)^T$, which implies that we need to understand the alignment between the eigenvectors of $X$ and the eigenvectors of $X+A$. Studying the eigenvectors of a perturbed matrix is very difficult. Further, we don't need to understand just the leading eigenspaces, but all eigenvectors.

We believe that the fact we got around this hurdle with accessible techniques is, in fact, a strength of our work.

When the Expectation is Taken

Another fact to highlight is that we compute the expectation over training data noise after computing the test error for $W_{opt}$. Recall that $W_{opt}$ is trained to minimize the MSE error over a single noise instance. If we computed the test error for a linear function minimizing the expected MSE error (corresponding to $W^*$), the result is much simpler. This latter case is exactly the case of learning a regularized auto-encoder.

We kindly ask if the reviewer, due to the subjective nature of their concern and the discussion with us, would be willing to increase their score to 5. (Note we have limited the audience of this comment).

Apology if I am interrupting your conversion with reviewer RbRx. I think it is extremely inappropriate to directly ask the reviewer for a specific score. And trying to hide this request with a limited view permission goes against the spirit of peer review.

We thank the reviewer for their comment and engagement and apologize for the summary.

My main suggestion for improvement would be for the authors to more critically (and fairly) situate their work in the context of previous literature.

We thank the reviewer for their comment. Our main difficulty is that in the denoising setting, there seems to be very few works. (Sonthalia and Nadakuditi 2023, Pretorius et al. 2018, and Cui and Zdeborova 2023). The reviewer agrees that we generalize the first two. We have modified the paper to highlight the strengths of all three papers. Primarily, Cui and Zdeborova 2023.

As an aside, I would like to echo Reviewer jx5r regarding the cherry picking of comments in the summary.

We apologize. We had hoped to identify the main concerns for each reviewer. We have updated the summary.

I would like to clarify some points of my review and provide some concrete suggestions for improvement. Also, to clarify, I agree with the authors that simple, accessible proofs are desirable and that technical difficulty for the sake of technical difficulty is not a useful exercise. My apologies if my initial review signaled otherwise.

We are thankful for the concrete suggestions. We hope that we have addressed most of these in the updated manuscript.

We are also thankful that the reviewer agrees that simple proofs are desirable. Hence, we hope that the simplicity is not held against us.

You’ve provided a phenomenological study in which you propose a simple model which reflects some aspects of more realistic situations and analyze the simpler set-up to provide insight ... is that not an insight already evident from, e.g., [1]?

[1] does tell us that avoiding $c = 1$ is necessary. However, that is not the only insight from our paper. Specifically, our insights are novel when considering distribution shifts for denoising. This is where we believe many of our insights appear. As we mentioned, there are very few works that study denoising. As far as we can tell, only Sonthalia and Nadukiti 2023 consider distribution shifts, but in a limited manner.

With the importance of denoising, we believe that having a simple model for which we can theoretically justify (known) phenomena is essential. Specifically, we show

1. Double descent occurs even when the test data is from a different distribution.
2. We show that data augmentation hurts in the $c > 1$ regime. While this can be thought of as understood from prior work such as [1]. Due to the inherent dependence in the data when doing data augmentation. We believe having a simple model where such results are theoretically backed is important.
3. As far as we can tell, The results on tempered and benign overfitting for denoising are also new. These do provide the interesting qualitative insight that when we have a small amount of training data, we can improve the denoiser by increasing the denoiser (even with distribution shift) by increasing over parameterization.

Thank you to the authors for their response regarding the difficulty in the analysis ... Many of the estimates used in the proof of Theorem 1 seem to come from that paper.

We agree with the reviewer that we do not have new, deep, technical mathematical results. However, the analysis here is not straightforward. We highlight some of the difficulties.

1. Lemma 4, which considers the concentration of the inverse of matrices. This is not needed in Sonthalia and Nadakuditi 2023 as they don't have to invert matrices.
2. Lemma 5, for bounding the variance of each entry, we consider the square of the matrix. We show that we can estimate this. Hence, by symmetry, we can bound the variance of the sums of the terms and, then, the individual terms. This argument is new and is not present in Sonthalia and Nadakuditi 2023.
3. Lemma 7, we had to bound the variance of the product of random variables. Sonthalia and Nadakuditi 2023 did not need to do this.

Again, regarding technical insights and novelty, ... without mention of any differences (and perhaps strengths that the previous work may have and weaknesses of their own).

We tried to highlight the differences between prior work and ours — specifically, the difference in data setting. We have added a paragraph in the introduction concerning some of our method's limitations and highlighted the strength of some prior work.

There seem to be two directly related works cited by the authors: Pretorius et al. (2018) and Sonthalia and Nadakuditi (2023). I do agree that the results obtained here are improvements upon these.

We agree with the reviewer.

Cheng and Montanari (2022)

We want first to clarify that we are looking at the same paper. Is the reviewer referring to https://arxiv.org/pdf/2210.08571.pdf ?

Further, we are uncertain where we say Cheng and Montanari have input noise. Could the reviewer please point that out to us? We cite this paper in 4 places.

1. To motivate the problem of having poorly conditioned covariances. This was done to help motivate the study for low rank data.
2. As a reference to prior work that mentions the proportional regime.
3. When we talk about their data assumption. How they have assumptions similar to $x\_i = \Sigma^{1/2}z\_i$. Directly quoting from their paper

``We allow the feature vector to be high-dimensional, or even infinite-dimensional, in which case it belongs to a separable Hilbert space, and assume either $z\_i = \Sigma^{-1/2}x\_i$ to have i.i.d. entries, or to satisfy a certain convex concentration property.''

Further, their paper does have additive Marchenko Pastur noise and, unless we are mistaken, is not about kernel regression. Again, quoting directly from their paper.

``we revisit ridge regression (l2-penalized least squares) on i.i.d. data $(x\_i, y\_i), i \le n$, where $x\_i$ is a feature vector and $y\_i = \langle \beta, x\_i\rangle + \xi\_i \in \mathbb{R}$ is a response.''

Page 4 details the assumptions on $\xi\_i$, which are mean 0, variance $\tau^2$, and I.I.D. Theorems 1,2,3 on pages 9, 10, and 11 seem to be the main results. However, none of them are for kernel regression. Theorem 1 is for the ridge regression situation, and Theorems 2 and 3 are for the ridgeless case for the over and under-parameterized cases.

Due to these differences, we believe we and the reviewer are potentially referring to different papers.

Related to the previous point: It is likely that I have missed something crucial here ... Again, perhaps I’ve mis-interpreted this point.

Apologies, but we are not certain we follow.

The point here was to compare our data assumptions with prior work. Many prior works assume $x_i \sim \mathcal{N}(0,\Sigma)$ and IID, whereas our paper studies $x_i \in \mathcal{V}$ in some low dimensional subspace. We do not need the data to be centered in our paper. Hence, we are unsure about the reviewer's point.

To give one more example of potentially unfair comparisons, the authors provide a comparison to Cui and Zdeborova (2023),

We have updated the manuscript to mention that Cui and Zdeborova 2023 have a more general model.

One other big missing piece is comparisons and citations to results already in the random matrix theory literature.

We do cite the original works showing convergence of the spectrum. However, since, as the reviewer points out and we agree with, we are not proving new random matrix theory results, we do not have an in-depth literature review. However, if the reviewer feels there are papers that we should mention, please let us know.

We thank the reviewer again for their comments. We hope that with these discussions, we clear doubts in relation to our contribution and relation to prior work. We hope that the updated manuscript reflects this.

Thanks very much to the authors for pointing this out: you are entirely correct that the Cheng and Montanari (2022) paper does not study kernel ridge regression, my apologies for this (this was an unfortunate typo that I did not catch after erasing a sentence). I have updated the comment to reflect this. Thanks once more, and my apologies again for this typo. I also did not mean to imply that you mentioned their work as studying input noise.

I will try to make clearer what I was trying to say while clarifying the centering issue I brought up:

The point here was to compare our data assumptions with prior work. Many prior works assume and IID, whereas our paper studies in some low dimensional subspace. We do not need the data to be centered in our paper. Hence, we are unsure about the reviewer's point.

To clarify, what I am wondering is whether it is more accurate to compare your results to a setting in which data is not centered, but the covariance of the noise is well-behaved? I.e., a more closely related model to yours seems to me to be the situation in which the observed data is X_{trn} + G, where the entries of G are standard Gaussian. If I am correct about this, it is quite different from the $\mathcal{N}(0, \Sigma)$ setting you compare to, and my reference to Cheng and Montanari (2022) was to contrast these. I hope this is clearer, and apologies for the confusion.

We restored the comment. The limited visibility allowed the PCs, SACs, ACs, and all reviewers could see the post. We appreciate the feedback of the reviewer.

I would first thank the reviewers for their detailed responses. However, I feel cherry-picking the content of my review is not conducive to a productive discussion. Your summary of my review focused my opinion over the first two introductory sections of this paper, and barely mentioned my complaints about the presentation of your main results in Section 3. Therefore, I kindly request the authors to update their summary to better reflect on the reviewers' criticisms and expand on the specific steps taken to improve the paper.