An Augmentation-Aware Theory for Self-Supervised Contrastive Learning

We express our sincere gratitude to Reviewer Do5D for appreciating the novelty of our theory. We address your concerns below.

---

Weaknesses.

W1,W2,W3,W4,W7,W8. A. We will revise the typos as suggested.

W5. Mistakes in proof. (1) the decomposition at the beginning of Proof of Theorem 2.7 should have $p_{K-1}(\ldots)r_{K-1}(\ldots)$; (2) the last equation of Eq(31) should be $\sum_{m=1}^k$.

A. (1) The typo $r_K$ in the second line should be $r_{K-1}$, so that $p_{K-1}r_{K-1}$ appears in the second term of the decomposition. This typo does not affect the overall correctness of the proof.

(2) No, the summation should be $\sum_{m=k+1}^K$. According to the definition of $r_k$, there are $k$ different-class negatives, so the summation $\sum_{m=1}^K$ here breaks into two groups: the negatives $\bar{x}\_m$ having different classes from the anchor $\bar{x}$ ($\sum_{m=1}^k$) and the negatives $\bar{x}\_m$ sharing the same classes as the anchor $\bar{x}$ ($\sum\_{m=k+1}^K$). By fixing the abuse of notation $k$, Eq(31) becomes

$$\log(1+\sum_{m=1}^k\exp(-\mathbb{E}\_{\bar{x}\_m\sim\rho_{i_m}}\mathbb{E}\_{a',a_m}f(a(\bar{x}))^\top[f(a'(\bar{x}))-f(a_m(\bar{x}\_m))])+\sum_{m=k+1}^K\exp(-\mathbb{E}\_{\bar{x}_m\sim\rho_c}\mathbb{E}\_{a',a_m}f(a(\bar{x}))^\top[f(a'(\bar{x}))-f(a_m(\bar{x}\_m))])),$$

which simplifies the expectations and separates the two groups.

W6. Eq(26) not rigorous, making Eq(25) and Eq(29) not strictly equivalent. Necessity to modify Theorem 2.6.

A. (1) Eq(26) is rigorous according to the commutative law of addition. It holds for any $\\{i_1,\ldots,i_K\\}$ being a rearrangement of $[K]$. Eq(26) ensures that e.g. $\mathbb{E}\_{\bar{x}_1\sim\rho_1}\mathbb{E}\_{\bar{x}_2\sim\rho_2}\mathcal{L}=\mathbb{E}\_{\bar{x}_1\sim\rho_2}\mathbb{E}\_{\bar{x}_2\sim\rho_1}\mathcal{L}$, and therfore Eq(27) can all be represented by the same $r_k$, making Eq(25) and Eq(29) strictly equivalent.

(2) Theorem 2.6 is correct. Typos will be fixed according to W5(2).

---

Questions.

Q1. Understanding of augmentation overlapping.

A1. We respectively disagree from a mathematical viewpoint. In the original paper (Wang et al., 2021), augmentation overlap is formulated as $\mathcal{T}$-connectivity (Definition 4.3), where $x_i$ and $x_j$ are $\mathcal{T}$-connected if $\mathrm{supp}(p(x_i^+|x_i))\cap\mathrm{supp}(p(x_j^+|x_j))\neq\emptyset$, or equivalently, $\exists t_i,t_j \in \mathcal{T}$ s.t. $t_i(x_i)=t_j(x_j)$, meaning that $x_i$ and $x_j$ sharing exactly the same augmented views $t_i(x_i)=t_j(x_j)$, which is unrealistic.

---

Q2. Relationship between $P_X$ and $P_{\bar{X}}$.

A2. $\mathrm{P}_X$ is the distribution of $a(\bar{x})$, where $a\sim\mathrm{P}_A$ and $\bar{x}\sim\mathrm{P}\_{\bar{X}}$.

---

Q3. Which results does Lemma 2.8 represent?

A3. Lemma 2.8 is part of Proposition 6 in Nozawa & Sato, 2021 and Lemma 4.3 in Arora et al., 2019. Compared with Proposition 6 in Nozawa & Sato, 2021, Lemma 2.8 does not have the $d(f)$ term and shows an equation because the left hand side is $\bar{\mathcal{R}}^{\mathrm{sup}}$ instead of $\bar{\mathcal{R}}^{\mathrm{un}}$ (without the need of using Jensen's inequality).

---

Q4. Advantages of decomposition.

A4. The decomposition makes separation between same-class and different-class negative terms possible. Specifically, it decomposes the contrastive risk according to the number of negative samples having different labels from the anchor sample. The decomposition is non-trivial, as it only works for contrastive losses that treat negative samples equally (Eq(26)). With this decomposition, we further decompose the risk into different-class negative terms and same-class negatives terms (Eq(31)). Then respectively investigating the two types of terms (Eq(33) and Eq(37)), we reach our main theorems.

---

Q5. Discussions about Theorem 2.4.

A5. Theorem 2.4 shows the finite-sample error bound as a complement to the other theorems based on the population distribution. It also shows the compatibility of our bound with generalization bounds in previous works.

---

Q6. Explain Lipschitz continuity.

A6. The Lipschitz continuity condition concerning $f$ with respect to the image $x$ for convolution neural networks has been widely made in previous works, see e.g. [1] On Lipschitz Bounds of General Convolutional Neural Networks. TIP, 2019. [2] Intriguing properties of neural networks. ICLR, 2014. Moreover, the consistency between our pixel-level (Figure 2) and representation-level (Figures 3 & 4) experiments also verifies the assumption.

---

Thanks for your insightful and constructive comments. According to our answers, we have clarified that our proofs have no factual mistakes (only typos to be fixed). We sincerely suggest the reviewer continue reviewing Section 3 and reconsider the overall recommendation score. We welcome any further questions during your follow up reviews.

We express our sincere gratitude to Reviewer 6JLt for appreciating the significance and reliability of our findings. We address your concerns below.

---

Q1. The error bound is not directly correlated with the final performance.

A1. We respectfully disagree. As shown in Section 4.2, we show that the optimal data augmentation parameter minimizing the distance sum also leads to the best downstream accuracy, indicating that lower downstream error bounds highly correlate with better downstream performance. Specifically, in Theorem 2.3 we show that the supervised error bound depends not only on the unsupervised loss but also on the maximum distance between augmentations and the minimum distance between samples of the same class. Typically, efforts focus on minimizing the unsupervised error to reduce the supervised error bound. However, the two distance terms, which depend on the augmentation parameter, are also crucial and must be minimized to improve classification accuracy. This is empirically verified by Figures 5 and 6, which show that classification accuracy is inversely proportional to the sum of these two distance terms.

---

Q2. The paper lacks empirical guidance, making it unhelpful for selecting appropriate augmentation parameters.

A2. On the one hand, our theoretical results in Section 2.2 show that the downstream classification risk depends not only on the unsupervised loss but also on the sum of two distance terms related to data augmentation, which provides empirical insight that the selection of data augmentation parameters is also vitally important. On the other hand, Section 3 demonstrates a theoretical trade-off between the two distance terms, indicating the existance of the optimal data augmentation parameters. Empirical varifications show that the optimal range of crop size should be large enough, with the upper bound being 1.0; the optimal probability of color jitter is around 0.6-0.8. (See Figures 5 & 6 and also the additional experiments on Imagenet100 in the response A2 to Reviewer Xxrg.) Furthermore, for future works, our theoretical results also guides a possible direction for designing new data augmentation methods to deliberately minimize some unsupervised surrogate of the two distances.

---

Q3. When selecting augmentation parameters, only suitable values lead to good performance. Could the authors explain why this occurs based on these theorems?

A3. In Theorems 2.1 and 2.3, we show that the supervised error bound depends not only on the unsupervised loss but also on the maximum distance between augmentations and the minimum distance between samples of the same class. Through model training, the unsupervised error can be reduced to a very small value. However, the impact of these two distance terms on classification performance is often overlooked. In fact, these distance terms are highly influenced by the augmentation parameter values. The augmentation parameter values that minimize the sum of these two distances may also yield the smallest classification error bound, as verified by Figures 5 and 6.

---

Thanks for your insightful and constructive comments. Hope our explanations and additional experiments can address your concerns.

We express our sincere gratitude to Reviewer cGdD for appreciating the soundness of our theoretical formulations and results. We address your concerns below.

---

Q1. Organization of Section 2.

A1. As suggested, we will make the discussions in Section 2.3 an independent section following Section 2 in our next version and also add a roadmap in the introductory section to enhance readability.

---

Q2. Lack of discussion regarding the contributions and comparisons with SOTA.

A2. The discussions of our contribution locate mainly in Section 2.3, where we compare our bound with previous theoretical works. Besides, we also conduct verification experiments in Section 4 to verify our new theoretical claims proposed in Sections 3.2 and 3.3. As the main contribution of this paper is to theoretically explain the role of data augmentation in self-supervised contrastive learning, the significance of the proposed error bound is not necessarily verified by ablations or experimental comparisons. Nonetheless, we could compare empirically the bound value with previous theoretical works. In detail, as shown in Figures 5 and 10, the sum of the two distance terms is around 1 (1.3 for TinyImagenet and 0.8 for CIFAR100) after 1000 epoch training, whereas we additionally measure the contrastive loss value to be around 42 (41.6 for TinyImagenet and 41.2 for CIFAR100). This indicates that compared with previous bounds, our bound relax the data generation assumption with the bound value increases less than 3% (1/42). See also A6.

---

Q3. Some experimental settings lack sufficient discussion and analysis in relation to the theoretical framework.

A3. The experimental section 4.1 is specifically in relation to theoretical analysis in Section 3 verifying the trade-off between the two distance terms, and the experimental section 4.2 is in relation to the error bound of Theorem 2.3, where the optimal augmentation parameter leads to the minimum distance sum and also the highest downstream accuracy.

To further enhance the relation of experiments to the theoretical framework, we here present an additional experiment about the formulations of data augmentation strategies. This experiment verifies the theoretical analysis in Section 3.3 about the combination of random crop and color distortion. Specifically, we present the pixel-level distances between data augmentations using only random crops and that using random crops plus color jitter on CIFAR100. As shown in the figure https://imgur.com/a/aj3Xp4i, on top of random crop, color distortion further reduces the minimum same-class different-image distance and enhances the maximum same-image different-augmentation distance.

---

Q4. Missing punctuation.

A4. We will fix the typo in our next version.

---

Q5. Do the authors analyze and discuss the impact of different augmentations and augmentation-aware error bounds in the two experimental settings?

A5. Yes, in both settings, we show that stronger augmentations, e.g. smaller range of crop size and higher probability of color jitter, lead to larger maximum same-image different-augmentation distance and smaller same-class different-image distance, which coincides with the theoretical analysis in Sections 3.2 and 3.3. From the perspective of our error bound, there exists a trade-off between the two distance terms w.r.t. the strength of data augmentation, indicating the existance of the optimal augmentation parameters. Related discussions can be found in Section 4.1.1 2nd paragraph, Section 4.1.2 2nd paragraph, and Section 4.2 2nd paragraph.

---

Q6. How does the proposed approach compare against state-of-the-art papers, and can the authors highlight its main contributions across these different settings?

A6. In fact, instead of proposing a new contrastive learning methods that competes against the state-of-the-art,, the aim of our error bound is to explain their working mechanisms. Our bound applies for any contrastive learnin methods as long as they use the InfoNCE loss function regardless of the backbones or other architectural techniques. It builds upon a more realistic data generation process described in Section 2.1 without relying on any further assumptions (Theorem 2.1). That is, compared with other theoretical papers deriving the error bound, our bound can easily adapt these previous bounds under our more realistic data generation, only with two additional trade-off distance terms added to these bounds. See also the third paragraph of Section 2.3 for more details.

---

Q7. The paper needs a detailed analysis of the experiments and discussions.

A7. See A2 and A3.

---

Thanks for your insightful and constructive comments. Hope our explanations and additional experiments can address your concerns.

We express our sincere gratitude to Reviewer Xxrg for appreciating our practical theoretical assumptions and novel theoretical results. We address your concerns below.

---

Q1. Empirical verification of Assumption 2.2.

A1. For each input image in the CIFAR100 dataset, we generate 100 different random augmentations, and plot the tsne visualizations of the random augmentated views, the mean of the augmented views, and the original input sample in the embedding space. We show the tsne plots of 4 randomly selected samples in https://imgur.com/a/NadpBSl. We see that the mean of augmented views lie near the embedding of the original input sample, which verifies Assumption 2.2. Besides, we also measure the $l$-2 distance between the mean of augmentations and the original sample in the embedding space. We show that on CIFAR100, the average distance is 0.083, which is small enough considering that the embeddings are normalized to have unit l-2 norm, indicating that the mean of augmented embeddings $\mathbb{E}_a f(a(\bar{x}))$ is close enough to input embedding $f(\bar{x})$.

---

Q2. The experiments are conducted only on small-scale datasets such as CIFAR-10/100. Since image resolution and dataset complexity could influence the theoretical results, an evaluation of larger, high-resolution datasets would provide a more comprehensive validation.

A2. Following your advice, we conduct additional verification experiments on Imagenet100, which contains 224*224 high resolution images from 100 classes. Due to time limitation, we run experiments with ResNet-18 with batchsize 64 for 50 epochs on 2 24G RTX 3090 GPUs. (The other settings follow those of TinyImagenet in our submitted paper.) The pixel-level verifications are shown in https://imgur.com/a/UcVPXwS, and the representation-level verification w.r.t. crop size is shown in https://imgur.com/a/jl00TVn, where the results conincides with that in Sections 4.1.1 and 4.1.2, i.e., as the augmentation strength increases, the max same-image different-augmentation distance increases and the min same-class different-image distance decreases. Besides, in the figure https://imgur.com/a/wPs9WOO, we verify that the optimal augmentation parameters with the smallest distance sum also leads to the highest downstream accuracy, which coincides with Section 4.2.

---

Q3. The study focuses exclusively on InfoNCE loss, which may limit the generalization of the theoretical findings.

A3. On the one hand, as one of the mainstream contrastive losses, InfoNCE-based contrastive learning still achieves state-of-the-art performances e.g. SimCLR, MoCo, etc. The extensive applications of InfoNCE loss partly demonstrate the generalization of our theory. On the other hand, although our theoretical results focus on the InfoNCE loss, we strongly believe that our analysis can potentially be generalized to other contrastive losses, e.g. BYOL, BarlowTwins, etc. Specifically, our analysis theoretically explains the effect of data augmentation to contrastive learning based only on the data generation procedure described in Section 2.1, with mild (Theorem 2.3) or even no further assumptions required (Theorem 2.1) and regardless of the network architectures. As long as a contrastive learning method selects positive/negative samples using random data augmentations, we can adapt the two trade-off distance terms to any error bound revealing the relationship between contrastive and downstream risks.

---

Q4. Page 1, top right. The statement “we usually use only two views in training instead of using multiple ones.” may be misleading for readers unfamiliar with SSL. While many traditional contrastive learning methods use two views, several recent approaches, such as MMCR and DINO, leverage multiple views.

A4. We will replace this statement with "we usually use only two views in training instead of using multiple ones in popular methods such as SimCLR and MoCo".

---

Thanks for your insightful and constructive comments. We sincerely hope you can reconsider the recommendation score if our explanations and additional experiments solves your concerns. We also welcome further questions and discussions.