Paper: Understanding Augmentation-based Self-Supervised Representation Learning via RKHS Approximation and Regression (ICLR-2024)

Reviewer comment (Novelty):
If I'm correct about the above relationship to Saunshi et al., I think this deserves more discussion, and I don't think you can say that you've refuted the claims of Saunshi et al. The paper could be improved by giving more context about the relationship of your results to those of Saunshi et al. For instance, is there any relationship between your augmentation complexity and Saunshi et al.'s approximately-disjoint augmentations? One might expect that approximately-disjoint augmentations would imply a high augmentation complexity; if so, that could be useful to discuss.

Author response:
2. **Claims in Saunshi et al. (2022)**: Thank you for this remark. We agree with your interpretation of their claims and have revised our discussion in Appendix A accordingly.

Reviewer comment (Evaluation):
Also, what do you mean by "disentangle the effects of the model and the augmentation"? It's not clear to me how the presented results disentangle these effects, although I do think they give a clear explanation of the role of the augmentations in particular.

Author response:
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.

Reviewer comment (Theory):
Can you give any intuition about what satisfying Assumption 1 actually means in practice? It seems unlikely to me that functions of practical interest could be *exactly* represented as expectations of some function $g$ on augmentations, and I also don't see any way to verify whether or not Assumption 1 holds.

Author response:
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.

Reviewer comment (Theory):
The proof of part (iv) of Proposition 1 also seems potentially incorrect. In particular I'm not sure how $\langle g^*, (I - \Gamma \Gamma^*) g^* \rangle_{P_A} \le \epsilon ||g^*||^2_{P_A}$ implies $\langle g_0, (I - \Gamma \Gamma^*) g_0 \rangle_{P_A} + ||g_1||^2_{P_A}  \le \epsilon (||g_0||^2_{P_A} + ||g_1||^2_{P_A})$; shouldn't $||g_1||^2_{P_A}$ be $\langle g_1, (I - \Gamma \Gamma^*) g_1 \rangle_{P_A}$ instead? It also seems like your proof here implies that $R(\Gamma^*) \subseteq \mathcal{F}_B(\Gamma; \epsilon)$ for all $\epsilon$. But that doesn't seem like it could be true; perhaps I'm misunderstanding something.

Author response:
1. **Part (iv) of Proposition 1**: We apologize for the typo in the original submission. The problem is fixed. Regarding the proof:
    - $\lVert g _1 \rVert _{P _A}^2$ is equal to $\langle g _1, (I - \Gamma \Gamma^*) g _1 \rangle _{P _A}$, because there is $\Gamma^* g _1 = 0$. We clarified this point in the updated pdf.
    - It should be $F_B(\Gamma; \epsilon) \subset R(\Gamma^*)$. The proof of (iv) only involves $f^* \in R(\Gamma^*)$, which is naturally implied by Assumption 1 (see the sentence after Assumption 1).

Reviewer comment (Theory):
In Theorem 2, don't the quantities $\overline{\lambda}_d^{-1}$ and $\gamma_G$ potentially depend on $N$? It seems like that could affect the asymptotic behavior of the bound in equation (20). Can you say anything about the convergence or boundedness of these quantities?

Author response:
2. **Asymptotic behavior of Theorem 2**: You are right that $\gamma _G$ and $\bar{\lambda} _d$ could potentially depend on $N$, and we do not know how to remove such independence at the moment. Intuitively, to remove such dependence, we need a quantitative measurement of the “smoothness” of the real top-d eigenfunctions, where better smoothness corresponds to fewer samples. Note that this dependence on the spectrum of the pretrained encoder also appeared in previous results (e.g. Theorem 4.5 of Saunshi et al. (2022), in which $\lambda _i$ depends on the pretrained encoder $\phi$).

Reviewer comment (Novelty):
I noticed that parts of your analysis resemble some of the statements of Johnson et al. (2023) at a high level; in particular, your soft invariance assumption is similar to their "approximate view invariance" assumption, and the decomposition in Theorem 1 is somewhat similar to their bound in Proposition 4.2. The details are different (their analysis focuses on approximating functions $\mathcal{A} \to \mathbb{R}$ rather than $\mathcal{X} \to \mathbb{R}$ and doesn't include the approximation error of the encoder) but it might still be worth including a brief discussion of the similarities.

Author response:
3. **Comparing to Johnson et al. (2023)**:
    - The soft constraint is equivalent to Assumption 1.1 in Johnson et al. (2023). We have clarified this point in the new Section 2.
    - Our Theorem 1 is in fact far distinct from Proposition 4.2 of Johnson et al. (2022), for the following reasons:
        1. Proposition 4.2 of Johnson et al. assumed that there is a pretrained upstream encoder $r _d^*$, fixing which they bound the excess risk of the linear probe. So essentially, their bound is the same as a bound of a linear model, where the inputs are $r _d^*(x)$ and the labels are the downstream labels. And note that their bound is only for the estimation error. They do not consider the approximation error because $r _d^*$ is fixed.
        2. Regarding our Theorem 1: First, it bounds both approximation error and estimation error. The approximation error is bounded using $\tau$, and Theorem 2 is bounding this $\tau$ when the encoder is the optimal encoder. Second, regarding the estimation error, while we also consider the linear probe, we use a completely different set of proving tools (local Gaussian and localized Rademacher complexity). Third, we are bounding the risk itself instead of the excess risk (that is we don’t assume that there is a good linear probe on top of the pretrained encoder), and our result also includes Gaussian label noise.

We did not feel the need to add this comparison in the remark of Theorem 1 given how different the results are, but we are happy to include it in our revision if the reviewer finds it necessary.

Reviewer comment (Writing):
It took me a while to locate the definition of $\hat{f}$ when reading Theorem 1; perhaps you could add a reference back to equation (9) where it is defined.

Author response:
3. Reminding $\hat{f}$ in Theorem 1: Thanks for this great suggestion! Added.

Reviewer comment (Writing):
A major issue is the lack of clarity. The paper is currently dense, not accessible to a general audience (as someone who is not well-versed in the area, I find the paper really challenging to read), and thus requires a major revision. Because of this, I have to admit that I have not fully understood the contribution of the paper.

Author response:
**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.

Reviewer comment (Theory):
It is not clear how reasonable the assumption on the target function is. It seems that the assumption is conveniently introduced so that we can analyse the problem using a particular RKHS. Is there any motivating example for using $\mathbb{E}_A[g^*(A)|x]$ as a target? It should be more straightforward to work with (functions on) the augmentation space $\mathcal{A}$ rather than $\mathcal{X}$, as in Johnson et al., 2023.

Author response:
**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.

Reviewer comment (Theory):
The encoder has to be of the form (10), as otherwise $\hat{\Psi}$ are not included in the RKHS. Therefore, the claim that the encoder $\hat{\Psi}$ can be arbitrary is not accurate to me.

Author response:
**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.

Reviewer comment (Theory):
The soft-invariance assumption is not straightforward to interpret. How does it measure "invariance"? Since $\mathbb{E}[g(A)|x]$ and $g(a)$ are functions on different sets, the invariance is not well-defined.

Author response:
**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.

Reviewer comment (Writing):
The problem is not well presented. It would have been more helpful if the authors had provided a high-level problem description and define symbols accordingly (e.g., as in the introduction of Johnson et al., 2023). For example, how representation learning is performed is not really mentioned using symbols in the paper (e.g., $\hat{\Phi}$ is not mentioned up to page 4).

Author response:
**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.

Reviewer comment (Writing):
Some symbols are loosely defined or mentioned before formally defined; e.g., $\hat{\mathcal{H}}_d$ is mentioned before the formal definition of $\hat{H}_{\Gamma}$.

Author response:
**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.

Reviewer comment (Theory):
For Proposition 1, the authors might want to check Chapter 11 of "An Introduction to the Theory of Reproducing Kernel Hilbert Spaces" by Paulsen and Raghupathi.

Author response:
Dear Reviewer kPdq,

Thank you for your response. "The range of a Hilbert-Schmidt operator can form an RKHS" is a well-known result. For example, a similar theorem was stated in Theorem 4.9 of Schölkopf & Smola (2002). The purpose of Proposition 1 is to introduce to our readers the necessary background of RKHS that leads to the key isometry property.

We have added a citation to the book you mentioned after Proposition 1.

Best regards,  
Authors

Reviewer comment (Writing):
In Section 5, the term "strong" augmentation is obscure. One may consider an augmentation distribution that is independent of $x$, in which case the augmentation complexity is 1. Is the dependency between $A$ and $X$ considered a measure of "strong"?

Author response:
**"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.

Reviewer comment (Theory):
Equation (9) requires a suitable regularisation that depends on $B$ and $\varepsilon$. In some sense, the analysis addresses a well-specified case. What happens if this is violated?

Author response:
**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.

Reviewer comment (Writing):
Regarding the statement "As shown in Wang et al. (2022b), there is...": Is this an existence statement?

Author response:
**“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.

Reviewer comment (Writing):
The key proof element is Assumption 1. I wished Section 2.1 was written differently so as to improve its clarity, emphasizing the assumptions that really matter based on my understanding, that is 1) $\Gamma \Gamma^*$ is Hilbert-Schmidt 2) eq. (8) + $f^* \in R(\Gamma^*)$ rather than the more obscure (1) or Ass.1. The writing of Proposition 1 can be definitely improved.

Author response:
**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).

Reviewer comment (Theory):
I think the comparison with the source condition of Cabannes et al. could be more developed since the authors' assumption ($f^* \in R(\Gamma^*)$) looks quite similar. I looked for the formula corresponding to (1) in HaoChen et al. but could not find it.

Author response:
**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.

Reviewer comment (Theory):
The final predictor of Definition 1 used a bound constraint on the norm. I am not sure whether it is common in some literature since I am used rather to regularized versions. These formulations require better justification.

Author response:
**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).

Reviewer comment (Evaluation):
I also have doubts with the comparison with the RIP.

Author response:
**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.

Reviewer comment (Theory):
Can you quote some articles that used a bound constraint on the final predictor as you do? In general imposing constraints simplifies proofs, is this reason of your choice?

Author response:
**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).

Reviewer comment (Theory):
Can you give examples of augmentations that do not lead to $\int K_X(x,x) dp(x) <\infty$? $p(a|x)=\delta_x$ for instance? Similarly, give some examples where $\Gamma \Gamma^*$ is clearly Hilbert-Schmidt (Gaussian blur I guess)? I ask this because it is not a given that $K_X(x,x)\in\mathbb{R}$ in my opinion. There are kernel integral operators that are not related to RKHSs.

Author response:
**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.

Reviewer comment (Theory):
Is it true that (8) implies that $u_i \approx 0$ for $\lambda_i < 1-\epsilon$? It seems to me to be the case. Eq.(8) can then be interpreted as requiring that $f^*$ is mostly encoded only on the first eigenvalues of $\Gamma \Gamma^*$, in other words that it is 'simple/low dimensional', and that has little to do with RIP I would say. Please discuss further this aspect, and if possible suggest a reformulation of Sec 2.1.

Author response:
**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.

Reviewer comment (Writing):
I find, that the assumptions and notations should be pointed out more clearly (see questions).

Author response:
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.

Reviewer comment (Theory):
The assumptions and notations should be pointed out more clearly. For instance, it seems it is assumed, but I did not see it stated explicitly, that \(P_\mathcal{A}\) has a density, and that \(p(a|x)\) and \(p(x|a)\) exist.

Author response:
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)$.

Reviewer comment (Writing):
What is A in \(E(g(A)|x)\)? Similarly, X is not defined as a random variable (it is clear from the context but should be defined).

Author response:
1. **Notations**: We added a concrete example in Section 2, which clarifies how $P _{AX}$ can be chosen. Regarding your questions:
    - $A$ and $X$: All capital letters (in normal font) are random variables; we have added this clarification.

Reviewer comment (Writing):
The notational difference between \(p\) and \(p(\cdot,\cdot)\) is small.

Author response:
1. **Notations**: We added a concrete example in Section 2, which clarifies how $P _{AX}$ can be chosen. Regarding your questions:
    - $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.

Reviewer comment (Theory):
Is \(\hat{\psi}\) always assumed to be of the form (10) or is (10) just an example?

Author response:
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.

Reviewer comment (Writing):
Can you explain the choice of the distribution \(P_{AX}\) a bit more? Particularly, it could be emphasized how the conditional distribution \(P(\cdot |x)\) has to be given. You can refer to the given examples.

Author response:
1. **Notations**: We added a concrete example in Section 2, which clarifies how $P_{AX}$ can be chosen.

Reviewer comment (Theory):
Can your approach be used to build or suggest an augmentation-based pretraining? Maybe under assumptions on prior information on the distribution \(P_X\)?

Author response:
3. **Suggesting new augmentations**: Our method can compare between augmentations, but new augmentations depend on human prior knowledge of the downstream task.

Reviewer comment (Theory):
Proposition 1, (iv): What does it mean that \(g^*\) satisfies Assumption 1? Assumption 1 is an Assumption on \(f^* \in \mathcal{L}^2(P_X)\), or does the statement in (iv) refer to (1)? The proof of statement (v) is done in the proof of statement (iv); I think those parts can be moved to the proof of statement (v).

Author response:
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.

Reviewer comment (Theory):
In (1), how are A and A' drawn from \(p(\cdot |x)\)? A motivation for (1) would be helpful. The space \(\mathcal{F}_B(\Gamma;\epsilon)\) is a central object, it would be good to motivate it more.

Author response:
1. **Notations**: We added a concrete example in Section 2, which clarifies how $P _{AX}$ can be chosen. Regarding your questions:
    -  $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.

Reviewer comment (Theory):
Definition 3: Is \(G^{-1}\) well-defined, even if \(d = \infty\)?

Author response:
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.

Reviewer comment (Experiments):
How is \(d\) chosen? Equation (20) would motivate to choose \(d\) such that there is a large gap between \(\lambda_d\) and \(\lambda_{d+1}\), but your result improves such a guess for \(d\) compared to HaoChen21 and Saunshi22.

Author response:
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.

Reviewer comment (Theory):
Is \(\bar{\phi}\) an eigenfunction of \(\Gamma^*\bar{\Gamma}\)? If so, how restrictive is the condition \(\hat{\phi} = \bar{\phi}\); is this just the choice of the encoder?

Author response:
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.

Reviewer comment (Experiments):
It is good to see a separation of model and augmentation in the error analysis. Do you think a combination, for instance, through augmentation based on the model, can be fruitful, too?

Author response:
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.

Reviewer comment (Theory):
Does Lemma 9 need a condition on boundedness of the matrix \(QDQ^T\) or absolute convergence of the sum \(\sum\limits_{i = 1}^\infty \lambda_i q_i^T q_i\)?

Author response:
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.

Reviewer comment (Experiments):
Did you compute the “sweet spot” mentioned in the introduction? Can it be computed?

Author response:
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).

Reviewer comment (Writing):
In Section 2.1 some reference for conditional expectation and RKHS could be added.

Author response:
2. **Reference for RKHS**: Added in Section 2.2.

Reviewer comment (Theory):
Proof of Proposition 1: The proof should also contain showing that \(\mathcal{H}_\Gamma\) is complete.

Author response:
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.

Reviewer comment (Theory):
Are \(K_X\) and \(K_A\) square integrable? It seems so from (4) and necessary from the proof of the statement that \(\Gamma^*\Gamma\) and \(\Gamma\Gamma^*\) are integral operators. Are conditions on \(P_{AX}\) needed that (4) holds?

Author response:
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.

Reviewer comment (Theory):
Lemma 5: Below Lemma 5, it is mentioned that the required 4th order bound holds. Why is the condition for such a bound included in the statement of the lemma?

Author response:
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.

Reviewer comment (Theory):
Are there exponential lower bounds on \(\kappa\) in Examples 2,3?

Author response:
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.