Understanding Augmentation-based Self-Supervised Representation Learning via RKHS Approximation and Regression

Thank you so much for your constructive comments! We have uploaded a new version of our paper, and would like to answer your questions below.

Comments

1. What we mean by “disentangle the effects of the model and the augmentation”: We mean that our bounds (a) are nonparametric, (b) do not depend on any model complexity or inductive bias, and hence (c) allow the user to choose any model of any size and architecture for the encoder $\hat{\Phi}$, which is what’s actually learned during upstream (see our remark in Section 3). Comparing our bounds to prior work:
   • HaoChen et al. (2021) and Saunshi et al. (2022): Their bounds contain model complexity (e.g. Rademacher complexity) and thus depend on the model.
   • Cabannes et al. (2023): Their bounds are nonparametric, similar to ours. However, they assumed the function class to be a subset of the RKHS of a predefined kernel $K$, which is determined by the model architecture. Thus, their results are still model-dependent.

A major novelty of our work compared to prior work is removing the effect of the model and only focusing on the augmentation. We believe that this is important, because it helps “clearly lay out the underlying structure produced by the augmentation distribution”, as the reviewer also pointed out.

2. What does $f^{\star}$ lie in a particular RKHS mean, and is this assumption reasonable?
   • First, we’d like to note that: Assuming the target function to belong to some (usually compact) function class is necessary for proving any nonparametric learning guarantee. Such assumptions are known as priors in statistical learning; the seminal works (DeVore et al., 2004; Bauer et al., 2007) have discussed in detail why priors are inevitable for statistical guarantees.
   • An RKHS (or more generally, a separable Hilbert space, e.g. Fischer & Steinwart (2020)) is a common choice of such a prior in most existing learning theory work. A concern is that, regardless of the definition of this Hilbert space, a small perturbation can make $f^*$ no longer belong to the space, as the reviewer pointed out. However, this has not been deemed a serious issue in learning theory community, since if the perturbation is small enough (e.g. on a subset of measure zero), then it won’t have a large impact on the statistical guarantees.
   • Interpreting our Assumption 1: Our induced RKHS is determined by the augmentation, which is selected based on human prior knowledge of $f^*$. Thus, so long as the human prior knowledge is correct, “$f^*$ belonging to the induced RKHS” can be derived from the soft invariance condition. For instance, if the augmentation is image translation and we know that image semantics are invariant w.r.t. translation, then soft invariance implies that $f^*$ belongs to the RKHS induced by translation.

Of course, a good model architecture also can contain human prior knowledge, like the design of convolutional networks. Cabannes et al. (2023) had a discussion on this in their Appendix D.3, where they considered linearization of one-block convolutional networks. However, in general, it is not quite clear how the inductive bias of a deeper, more complicated practical model (that cannot be linearized) can be formulated as a kernel.

DeVore et al. Mathematical methods for supervised learning. 2004.
Bauer et al. On Regularization Algorithms in Learning Theory. 2007.
Fischer & Steinwart. Sobolev Norm Learning Rates for Regularized Least-Squares Algorithms. 2020.

3. Regarding the claim in Saunshi et al. (2022):
   • We’d like to first clarify that we interpret “non-vacuous” as “meaningful”, rather than “close to empirical observations”. We do believe that our bounds, which do not contain any model complexity, are meaningful. We are definitely not claiming anything strong such as “augmentation is all you need”, and we emphasize in Appendix A that combining inductive biases of the model class can lead to tighter bounds. What we do claim is that, if the augmentation is good enough (which is a strong assumption), then we can achieve good generalization without the help of model inductive bias. This is why we disagree with Saunshi et al.’s strong claim that any guarantees independent of the inductive bias are vacuous.
   • Regarding the disjoint augmentation example in Saunshi et al. (2022): Thank you for this great question! This example showcases a poor choice of augmentation, for which $K_X$ defined in our Eqn. (3) is zero everywhere, so surely our bounds won’t work. Their observation was that even for this very bad augmentation, we can still obtain generalization guarantees if there is a good model (feature map). To see why: Suppose there is a good feature map $\Phi$. Then, the augmentation $x \mapsto a$ is disjoint and thus is very bad. However, the augmentation $x \mapsto \Phi(a)$ is not disjoint and in fact could be quite good. And if we view $x \mapsto \Phi(a)$ as the augmentation, then our model-free bounds work again.

Dear Reviewer dUJ6,

Thank you for your time in reviewing this paper! We are still finalizing our response to you, and will submit it within the next 10 hours. Your review has an exceptionally high quality, and we want to make sure that our response could be equally good. We are definitely looking forward to having a discussion with you.

Best regards,
Authors

Thank you to the authors for the detailed response -- I will have a closer look. Please let me share a quick (and minor) response.

First of all, we don't have an individual response Reviewer to dUJ6. Is this okay?

Re: Chapter 11 of the book by Paulsen and Raghupathi, Section 11.2 in this chapter treats the RKHS defined by an integral operator (the range space construction), and so Proposition 11.1 is an instance of this. Also there are some superfluous notation and symbols (e.g., do we really need to display the Radon-Nikodym derivative, given that we have a concrete density function and don't use measure theoretic results?).

Re: soft-invariance, the following comment is not that important, but I just want to clarify what I meant (and the reason I found the term soft-invariance a bit weird). I do not expect a response to this comment.

My understanding of (1) is that this inequality controls the conditional variance of $g$ (averaged over x), and my comment is based on this interpretation. if a function is (exactly) invariant under some transformation, then we should expect the transformed function lies in the same set (i.e,, if $\tilde{g}$ is the transformed version of $g$, then we expect $\tilde{g} = g$ if $g$ is invariant). If the transformation is taking conditional expectation, then $g$ cannot be invariant since the conditional expectation maps $g$, which is a function of $a$, to a function of $x$. If $\mathcal{A}=\mathcal{X}$, I think calling the assumption soft-invariant makes sense, in that $g$ is invariant with respect to the Markov kernel (the conditional expectation); that is $g$ is soft-invariant as a function of $x$, if the distance between $\tilde{g}$ and $g$ is small, where $\tilde{g}(x) = \int g(x') P(d x'|x)$ with $P(d x'|x)$ the augmentation distribution.

Thank you for your review. We have uploaded a new version of our paper and we invite the reviewer to check it out. We would like to answer to your questions as follows:

Clarity: Thank you very much for this feedback! We agree that our paper is on the denser side and needs to be improved in terms of clarity. We have updated the paper based on the comments from the reviewers, including notation clarification, added illustrative examples, and rephrasing. Please let us know if there are further changes that you would like to see; we do hope our paper provides sufficient background to be self-inclusive, and would really appreciate your feedback.

Significance

1. Motivation and necessity of Assumption 1: First, we’d like to note that assuming the target function to belong to some function class is a necessity (rather than a convenience) for proving any nonparametric learning guarantee, and an RKHS is a common choice in lots of prior work. Second, in our Proposition 1 we show a duality between the eigenfunctions on $A$ and the eigenfunctions on $X$, and this essentially says that working with $A$ is equivalent to working with $X$ up to an invertible transformation. Please also refer to our discussion on Assumption 1 in Appendix A for more details.

2. Clarification on the arbitrary encoder: We apologize for the confusion. What we meant is that $\hat{\Phi}$ over $A$ can be an arbitrary function, while $\hat{\Psi}$ over $X$ is always defined as the average encoder of $\hat{\Phi}$ in this work. As we mentioned in Section 3, in practice during upstream people train $\hat{\Phi}$ instead of $\hat{\Psi}$, so they can use any model of any size and any architecture. We made this point clearer in Section 3 of the updated pdf.

3. Interpreting the soft invariance: The soft invariance is almost equivalent to Assumption 1.1 in Johnson et al. (2023), except that it adds a $\lVert g^* \rVert _{P _A}^2$ on the right so that it is homogeneous; we have clarified this in the updated pdf. Essentially, the soft invariance says that if $A, A’$ are augmented from the same original sample $X$, then the predictor would give them similar predictions. Similar conditions have appeared in lots of prior work, including Haochen et al. 2021; Saunshi et al. 2022; Johnson et al. 2023.

   • Could the reviewer please clarify which equation they are referring to by “$E[g(A) | x]$ and $g(a)$ are functions on different sets”? We couldn’t find in which equation we are comparing functions of different spaces.

4. Dimensionality dependency: We’d like to reemphasize that this dependency is a common issue and not a peculiar problem of our paper, as discussed in Section 5.1 and Section 6. Specifically, note that the more well-known model complexity based generalization bounds also cannot evade the curse of dimensionality.

Comments

We would also like to thank the reviewer for the comments, and we have updated the paper accordingly. Two further clarifications:

• "Strong" augmentation: the strength of an augmentation is determined by the joint distribution over $A$ and $X$, and can be quantitatively measured by the augmentation complexity introduced in this work. The reviewer pointed out that if $A$ and $X$ are independent, then the augmentation complexity is 1. This is correct, but please note that in practice any meaningful $A$ should depend on $X$ since $A$ is “augmented” from $X$, and hence the two should never be independent.
• Regarding Chapter 11 of the book by Paulsen and Raghupathi, could you please specify which results in Chapter 11 we should focus on?

Questions

1. How to implement Eqn. (9) if the exact values of $B$ and $\epsilon$ are unknown: thank you for the good question! In practice, one only needs to have an estimated upper bound $M$ of $\frac{ B }{ \sqrt{1-\epsilon} }$, and our bound of $\lVert \hat{f} - f^* \rVert$ is $O(M)$, so it will not make too much difference if $M$ has the same order as $\frac{ B }{ \sqrt{1-\epsilon} }$. This is the common practice for least-squares regression, including kernel ridge regression, where one doesn’t know the theoretical optimal regularization coefficient but uses an estimated value instead.

2. Why is $G^{-1/2}$ required: $G^{-1/2}$ is used for normalization. We added this clarification in the updated pdf.

3. “As shown in Wang et al.”: This is not an existence statement. We meant that Wang et al. showed this equation; we have edited this sentence in the updated pdf.

We hope that this response has addressed all your questions and concerns. The reviewer writes a very detailed, thoughtful and constructive review, which we really appreciate. We welcome the reviewer to have further discussions with us.

Thank you so much for your detailed review! We have uploaded a new version of our paper which we invite you to check out.

General remarks

1. Notations: We added a concrete example in Section 2, which clarifies how $P _{AX}$ can be chosen. Regarding your questions:
   • $P _A$, $p(a|x)$ and $p(x|a)$: In Section 2.1, paragraph 2: “Data augmentation induces a joint distribution $P _{AX}$, with marginal distributions $P _A$ and $P _X$”. Thus, $P _A$ exists from marginalization, similarly for $p(a|x)$ and $p(x|a)$.
   • $A$ and $X$: All capital letters (in normal font) are random variables; we have added this clarification.
   • $p$: In the updated pdf, we changed $p(x)$ to $P _X (x)$, and $p(a)$ to $P _A (a)$. The remaining $p(\cdot)$ should be clear from context.
   • $A, A’, F _B(\Gamma; \epsilon)$: $A, A’$ are two augmentations from the same original sample $X$. $F _B(\Gamma; \epsilon)$ comes from the soft invariance similar to prior work.
2. Definition of $\hat{\Psi}$: We would like to clarify that what we mean by “arbitrary” in Section 3 is that $\hat{\Phi}$ can be an arbitrary function; And in this work $\hat{\Psi}$ is always the average encoder of $\hat{\Phi}$ defined by Eqn. (10). As we mentioned before Eqn. (10), in practice during upstream people train $\hat{\Phi}$ instead of $\hat{\Psi}$, so they can use an arbitrary encoder of any size and architecture.
3. Suggesting new augmentations: Our method can compare between augmentations, but new augmentations depend on human prior knowledge of the downstream task.
4. Whether Assumption 1 can be verified: While it is not practically easy to find the exact value of $\epsilon$ to verify Assumption 1, we noted after Eqn. (8) that the exact value of $\epsilon$ does not affect the fact that the top-d eigenfunctions is optimal.
5. Is $G^{-1}$ well defined: $\hat{\Phi}$ is explicitly assumed to be full-rank (see the paragraph above Definition 3), so $G^{-1}$ is well-defined. We have added clarifications in the updated pdf.
6. What is $\lVert \cdot \rVert _n$: This notation is not used in the main body, but only in our proof. We have moved this to the appendix.
7. Choose d with a large gap between $\lambda _d$ and $\lambda _{d+1}$: We would like to clarify that a large gap between $\lambda _d$ and $\lambda _{d+1}$ is not necessary for our results. In fact, this is one of our major improvements over Haochen et al. (2021) and Saunshi et al. (2022). Their bounds depend on $(\lambda _d - \lambda _{d+1})^{-1}$. However, in reality $\lambda _d$ decays quite fast, so this gap is usually much smaller than $\lambda _d$. e.g. see Figure 1 of https://openreview.net/pdf?id=ByxY8CNtvr. Meanwhile, our bound does not depend on this gap but only $\lambda _d^{-1}$. Also, we are not suggesting practitioners to “guess for $d$” using our bounds. After all, our results are just upper bounds, not the real test error.
8. Regarding $\hat{\phi} = \bar{\phi}$: In Section 4, we study the special case $\hat{\phi} = \bar{\phi}$, the Monte-Carlo approximation of the optimal encoder. This is not a restriction, but rather an interesting special case because many popular self-supervised learning methods amount to this, as pointed out in Appendix C.
9. Combining augmentation and model: Yes, combining the augmentation complexity based bounds in this work, and previous model complexity based bounds, can lead to tighter bounds, as we discussed in Appendix A.
10. Lemma 9: $Q D _{\lambda} Q^{\top}$ is a $d \times d$ matrix where $Q, D$ are both bounded entrywise, hence the product $Q D _{\lambda} Q^{\top}$ is also bounded.

Minor comments

1. How to find the “sweet spot” in augmentation strength: We consider this as mainly an empirical observation determined by validation, demonstrated in Figure 3(c).
2. Reference for RKHS: Added in Section 2.2.
3. Are $K_X, K_A$ square integrable: Yes: $K _X$ is assumed to be square integrable after Eqn. (4), and $K _A$ is square integrable due to Proposition 1.
4. Fourth moment condition in Lemma 5: Lemma 5 is a general result for any $\hat{\phi}$, so it needs the fourth moment condition. Lemma 6 and Theorem 2 considers the special case $\hat{\phi} _i = \bar{\phi} _i$, and the fourth moment condition holds for this special case. We have added this clarification.
5. Exponential lower bounds for $\kappa$: There is no exponential lower bound. $\kappa$ can be polynomial if the augmentation is very strong; for example, in supervised pretraining where each sample is “augmented” to its upstream class.

We have also adopted your other suggestions, including updating (iv) and the proof of proposition 1, referring to examples 1-3 around Definition 2, revising the sentence after Definition 1, recalling $\lVert f^* \rVert$ in Section 3, etc.

Thank you again such a detailed, thoughtful and constructive review. We hope that we have addressed your questions and concerns, and welcome any further discussions.

Thank you for your discussion! We are glad that our previous response has answered most of your questions. We would like to respond to your new comments as follows:

1. Suggesting new augmentations: This is a great question. Let us provide a rough argument: For example, our analysis could explain why patchifying and discretizing an image into tokens (such as in ViT) is a good idea compared to pixel-level augmentation. For MAE, if we directly mask an image at the pixel level, then this augmentation is almost disjoint (as pointed out by Saunshi et al. 2022), and thus is quite bad. With patchifying and discretizing, we can make the masking augmentation less disjoint and thereby lower its augmentation complexity, leading to better generalization. Similarly, this argument could explain why discretizing techniques such as VQ-VAE or VQ-GAN can be helpful. Fully justifying this argument would require an extensive empirical study, which we leave to a future paper.

2. What if $f^{\star}$ does not belong to $H$, but is very close to $H$? We have a discussion on this question in Appendix A. Here we’d like to provide a more detailed discussion. The short answer to your question is: Assuming $f^*$ belonging to some function class is inevitable, but it is indeed possible to relax $f^* \in H$ such that $f^*$ belongs to some larger class, under some extra assumptions.

   • First, we have to assume that $f^*$ belongs to some function class, which usually requires $\lVert f^* \rVert _{\infty}$ to be bounded. Only assuming $f^*$ is close to $H$ is not sufficient for proving any bound. For example, for $f^* \in H$, suppose we add a tiny perturbation $f^* \rightarrow \tilde{f}^*$, such that $\tilde{f}^*$ is unbounded on a subset $X _0$ whose measure is very small but positive. Then, even though $\lVert f^* - \tilde{f}^* \rVert _{L^2(P _X)}$ can be very small, we still cannot prove any bound unless no training sample belongs to $X _0$, whose asymptotic probability is 0 when there are infinitely many samples. Basically, a single training sample from $X _0$ would have an uncontrollable effect on the predictor.
   • Second, relaxing “$f^* \in H$” is possible: We can assume that $f^*$ belongs to some larger function class. One such relaxation was done by Fischer & Steinwart (2020), who considered the “power spaces” $H^p$. The smaller $p$ is, the larger $H^p$ will be. And for $p < 1$, $H^p$ is not necessarily an RKHS, but still a separable Hilbert space. The reviewer could intuitively think of $H^p$ as the RKHS of $K^p$, assuming that $H^p$ is an RKHS. Then, Fischer & Steinwart proved tight statistic bounds under the assumptions that $f^* \in H^p$, the eigenvalues of $H$ decay fast enough, and $\lVert f^* \rVert _{\infty}$ is bounded (this is still inevitable).

From a theoretical perspective, $f^* \in H$ is indeed relaxable, but we feel that such relaxation would provide little extra help for people to understand the role of augmentation. That’s why in this work, we just assume that $f^* \in H$.

Fischer & Steinwart. Sobolev Norm Learning Rates for Regularized Least-Squares Algorithms. 2020.

3. Why we are not suggesting to “guess d”: In this work, we show the following trade-off controlled by $d$: A smaller $d$ leads to larger approximation error because the encoder has lower representational capacity, but a larger $d$ leads to larger estimation error because the encoder contains more noise, and thus gives rise to a higher downstream sample complexity. We do not claim that we can estimate the best $d$ that minimizes real approximation error + real estimation error, because what we have proved are upper bounds, and the $d$ that minimizes approximation error bound + estimation error bound need not to be the optimal $d$, even if the bounds are tight.
   • In statistical learning, “an upper bound is tight” is usually in a minimax sense, which means that one can prove a lower bound by constructing an example of a very bad task, and the upper bound is tight if it coincides with this lower bound on this worst-case example. But real tasks are rarely worst-case, so the upper bound can rarely match the real test error on a real task. This is why we cannot claim that we can estimate the best $d$ just with our upper bounds, even if they are extremely tight.

We are very happy to have further discussions with you, thank you again!

My comments and remarks have been well addressed. Thank you to the authors.

The notations have been adapted and clarified. My questions have been mostly answered. I still have three (whose responses do not need to be included in the paper)

1. Concerning "Suggesting new augmentations": Thank you for your response. I would like to continue shortly on this subject. Once one is able to compare between different augmentations one can choose the best. This is why I was wondering if your bounds give rise to suggestions for augmentation, too.

2. I find particularly interesting one of the concerns of reviewer dUj6 on the condition about where $f^*$ lives. I have read the authors' response to the question by reviewer dUj6 and it made me wonder if the case $f^* \notin H$ just adds an additional simple term to the error analysis (such as $err(f^*,…) = err_H(f,…) + \vert\vert f-f^*\vert\vert_V$ where $V$ is a space with less regularity (making it more likely to have $f^* \in V$ than having $f^* \in H$) and $f$ is a (well-chosen) element in $H$.

3. Why would the authors not suggest using their analysis to “guess d”? Does it relate to tightness of the error bounds?

Thank you for your review. We have uploaded a new version of our paper and we invite the reviewer to check it out. We would like to answer to your questions as follows:

Comments

1. Improving the clarity of Section 2.1: Thank you for the suggestion! Please note that we have largely revised Section 2.1, including adding an example of BERT that explains the notations, clarifying Eqn. (1), revising Proposition 1, and a new remark after Eqn. (8).

2. Comparing to Cabannes et al. (2023): As we mentioned in Section 6, the main difference is that Cabannes et al. (2023) assumed the function class to be a subset of $H$, which is the RKHS of a pre-selected kernel $K$ (which they said is determined by the model architecture). On the other hand, the main goal of this work is to decouple the effect of model and augmentation, in order to derive model-free learning guarantees. As a result, our paper does not have a “pre-selected kernel”. The only element that is pre-selected in our work is the augmentation method, which gives rise to the kernels. This helps isolate the effect of augmentation, allowing us to investigate its precise role.

3. The connection between our isometry property and RIP: Thank you for this great question. Our isometry property can be considered as a counterpart for RIP in infinite-dimensional functional spaces: RIP says that in finite-dimensional vector spaces, for a matrix $M$ and a vector $v$, $Mv$ preserves most variance of $v$, and therefore the classical method PCA finds the principal components that keep the most variance. Similarly, our isometry property says that $\Gamma^* \Gamma f^*$ preserves most variance of $f^*$ (the fraction of variance that is lost is less than $\epsilon$). Regarding whether $f^*$ can be encoded by the top eigenfunctions depends on how fast the eigenvalues of $\Gamma^* \Gamma$ decay: If it is fast enough then this is true, otherwise this is false. We have largely revised Section 2.1. Please let us know if the new version is clearer.

Questions

1. Correspondence of Eqn. (1) in Haochen et al.: We apologize for the confusion. We meant to say that the Laplacian over the augmentation graph was introduced in Haochen et al. (2021), and the same soft invariance condition was introduced in Johnson et al. (2023) in their Assumption 1.1. We have clarified this in Section 2.1 in the updated pdf.

2. Bound on the RKHS norm in the estimator: The norm bound can be considered as a capacity constraint and is necessary for proving any guarantees; for reference, please see Chapter 13 of Wainwright (2019). In general, there are two typical ways of estimating the regression function in nonparametric least-squares estimation: either restricting the optimization problem to some appropriately chosen subset (usually needs to be bounded), or adding a regularizer as in kernel ridge regression. In our particular case, using a regularizer is the same as adding a constraint on the RKHS norm, as shown in Lemma 35 of Cabannes et al. (2023).

3. Unbounded kernels: This is a good question. For our work, $\Gamma \Gamma^*$ is Hilbert-Schmidt because $\lVert K \rVert_{\infty}$ is bounded (by $\kappa^2$). To the best of our knowledge, a bounded $\lVert K \rVert_{\infty}$ is required by most statistical guarantees of kernel methods, from classical ones to more recent results. For examples of unbounded kernels and kernels not associated with an RKHS, please see Chapter 4 of Christmann and Steinwart (2008), as well as other papers from Ingo Steinwart. However, the learning guarantees provable for unbounded kernels are very limited, and they are far beyond the scope of this work.

4. Is $u _i \approx 0$ for $\lambda _i < 1 - \epsilon$?

   • Not necessarily, and here is a simple example: Suppose $\lambda _1 = 1$ and $\lambda _2 = 1 - \epsilon - \alpha$, for some very small $\alpha > 0$. And all other eigenvalues are 0. Then, Eqn. (8) is equivalent to $\frac{\alpha}{1 - \epsilon - \alpha} u _2^2 \le \epsilon u _1^2$, which holds as long as $\frac{u _2^2}{u _1^2} \le \frac{\epsilon (1 - \epsilon - \alpha)}{\alpha}$. Thus, $u _2$ can be very large, even much larger than $u _1$ if $\alpha$ is very small.
   • Moreover, $f^*$ is not necessarily low-dimensional. For example, if $\lambda _1 = 1$ and $\lambda _2 = \lambda _3 = \cdots = \lambda _{1000} = 1 - \epsilon - \alpha$, then with the same argument, $f^*$ can have 1000 dimensions.

Thank you also for the detailed writing suggestions in your minor comments. We have updated the draft accordingly.

We hope that this response has addressed all your questions and concerns. The reviewer raises some great questions which we really appreciate, and we welcome the reviewer to have further discussions with us.