Original: A M ESSAGE P ASSING P ERSPECTIVE ON L EARNING D YNAMICS OF C ONTRASTIVE L EARNING
Revised: A M ESSAGE P ASSING P ERSPECTIVE ON L EARNING D YNAMICS OF C ONTRASTIVE L EARNING

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In recent years, contrastive learning achieves impressive results on self-supervised visual representation learning, but there still lacks a rigorous understanding of its learning dynamics.
Revised: In recent years, contrastive learning achieves impressive results on self-supervised visual representation learning, but there still lacks a rigorous understanding of its learning dynamics.

Original: In this paper, we show that if we cast a contrastive objective equivalently into the feature space, then its learning dynamics admits an interpretable form.
Revised: In this paper, we show that if we cast a contrastive objective equivalently into the feature space, then its learning dynamics admits an interpretable form.

Original: Specifically, we show that its gradient descent corresponds to a specific message passing scheme on the corresponding augmentation graph.
Revised: Specifically, we show that its gradient descent corresponds to a specific message passing scheme on the corresponding augmentation graph.

Original: Based on this perspective, we theoretically characterize how contrastive learning gradually learns discriminative features with the alignment update and the uniformity update.
Revised: Based on this perspective, we theoretically characterize how contrastive learning gradually learns discriminative features with the alignment update and the uniformity update.

Original: Meanwhile, this perspective also establishes an intriguing connection between contrastive learning and Message Passing Graph Neural Networks (MP-GNNs).
Revised: Meanwhile, this perspective also establishes an intriguing connection between contrastive learning and Message Passing Graph Neural Networks (MP-GNNs).

Original: This connection not only provides a unified understanding of many techniques independently developed in each community, but also enables us to borrow techniques from MP-GNNs to design new contrastive learning variants, such as graph attention, graph rewiring, jumpy knowledge techniques, etc.
Revised: This connection not only provides a unified understanding of many techniques independently developed in each community, but also enables us to borrow techniques from MP-GNNs to design new contrastive learning variants, such as graph attention, graph rewiring, jumpy knowledge techniques, etc.

Original: We believe that our message passing perspective not only provides a new theoretical understanding of contrastive learning dynamics, but also bridges the two seemingly independent areas together, which could inspire more interleaving studies to benefit from each other. 
Revised: We believe that our message passing perspective not only provides a new theoretical understanding of contrastive learning dynamics, but also bridges the two seemingly independent areas together, which could inspire more interleaving studies to benefit from each other. The code is available at https://github.com/PKU-ML/ Message-Passing-Contrastive-Learning .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Contrastive Learning (CL) has become arguably the most effective approach to learning visual representations from unlabeled data (Chen et al., 2020b;
Revised: Contrastive Learning (CL) has become arguably the most effective approach to learning visual representations from unlabeled data (Chen et al., 2020b;

Original: He et al., 2020; Chen et al., 2020c;d; 2021;
Revised: He et al., 2020; Chen et al., 2020c;d; 2021;

Original: Caron et al., 2021).
Revised: Caron et al., 2021).

Original: However, till now, we actually know little about how CL gradually learns meaningful features from unlabeled data.
Revised: However, till now, we actually know little about how CL gradually learns meaningful features from unlabeled data.

Original: Recently, there has been a burst of interest in the theory of CL in order to explain its effectiveness.
Revised: Recently, there has been a burst of interest in the theory of CL in order to explain its effectiveness.

Original: However, despite the remarkable progress that has been made, existing theories of CL are established for either an arbitrary function f in the function class
Revised: However, despite the remarkable progress that has been made, existing theories of CL are established for either an arbitrary function f in the function class

Original: F (Saunshi et al., 2019; Wang et al., 2022) or the optimal f ∗ with minimal contrastive loss (Wang &
Revised: F (Saunshi et al., 2019; Wang et al., 2022) or the optimal f ∗ with minimal contrastive loss (Wang &

Original: Isola, 2020; HaoChen et al., 2021; Wang et al., 2022).
Revised: Isola, 2020; HaoChen et al., 2021; Wang et al., 2022).

Original: Instead, a theoretical characterization of the learning dynamics is largely overlooked, which is the focus of this work.
Revised: Instead, a theoretical characterization of the learning dynamics is largely overlooked, which is the focus of this work.

Original: Perhaps surprisingly,from a feature space perspective, we find out that the optimization dynamics of contrastive learning corresponds to a specific message passing scheme among different samples.
Revised: Perhaps surprisingly, we find out that the optimization dynamics of contrastive learning corresponds to a specific message passing scheme among different samples.

Original: Specifically, based on a reformulation of the alignment and uniformity losses of the contrastive loss, we show that the derived alignment and uniformity updates actually correspond to message passing on two different graphs: the alignment update on the augmentation graph defined by data augmentations, and the uniformity update on the affinity graph defined by feature similarities.
Revised: Specifically, based on a reformulation of the alignment and uniformity losses of the contrastive loss into the feature space, we show that the derived alignment and uniformity updates actually correspond to message passing on two different graphs: the alignment update on the augmentation graph defined by data augmentations, and the uniformity update on the affinity graph defined by feature similarities.

Original: Therefore, the combined contrastive update is a competition between two message passing rules.
Revised: Therefore, the combined contrastive update is a competition between two message passing rules.

Original: Based on this perspective, we further show that the equilibrium of contrastive learning can be achieved when the two message rules are balanced, i.e., when the learned distribution P θ matches the ground-truth data distribution P d , which provides a clear picture for understanding the dynamics of contrastive learning.
Revised: Based on this perspective, we further show that the equilibrium of contrastive learning can be achieved when the two message rules are balanced, i.e., when the learned distribution P θ matches the ground-truth data distribution P d , which provides a clear picture for understanding the dynamics of contrastive learning.

Original: Meanwhile, as message passing is a general paradigm in many scenarios, the message passing perspective of contrastive learning above also allows us to establish some intriguing connections to these seemingly different areas.
Revised: Meanwhile, as message passing is a general paradigm in many scenarios, the message passing perspective of contrastive learning above also allows us to establish some intriguing connections to these seemingly different areas.

Original: One particular example is in graph representation learning.
Revised: One particular example is in graph representation learning.

Original: Message Passing Graph Neural Networks (MP-GNNs) are the prevailing designs in modern Graph
Revised: Message Passing Graph Neural Networks (MP-GNNs) are the prevailing designs in modern Graph

Original: Neural Networks (GNNs), including numerous variants like GCN (Kipf & Welling, 2017), GAT (Veliˇckovi´c et al., 2018), and even the Transformers (Vaswani et al., 2017).
Revised: Neural Networks (GNNs), including numerous variants like GCN (Kipf & Welling, 2017), GAT (Veliˇckovi´c et al., 2018), and even the Transformers (Vaswani et al., 2017).

Original: There is a vast literature studying its diffusion dynamics and representation power (Li et al., 2018; Oono & Suzuki, 2020;
Revised: There is a vast literature studying its diffusion dynamics and representation power (Li et al., 2018; Oono & Suzuki, 2020;

Original: Wang et al., 2021; Dong et al., 2021; Xu et al., 2019; 2018).
Revised: Wang et al., 2021; Dong et al., 2021; Xu et al., 2019; 2018).

Original: Therefore, establishing a connection between contrastive learning (CL) and MP-GNNs will hopefully bring new theoretical and empirical insights for understanding and designing contrastive learning methods.
Revised: Therefore, establishing a connection between contrastive learning (CL) and MP-GNNs will hopefully bring new theoretical and empirical insights for understanding and designing contrastive learning methods.

Original: In this work, we illustrate this benefit from three aspects: 1) we establish formal connections between the basic message passing mechanisms in two domains; 2) based on this connection, we discover some close analogies among the representative techniques independently developed in each domain; and 3) borrowing techniques from MP-GNNs, we design two new contrastive learning variants, and demonstrate their effectiveness on benchmark datasets.
Revised: In this work, we illustrate this benefit from three aspects: 1) we establish formal connections between the basic message passing mechanisms in two domains; 2) based on this connection, we discover some close analogies among the representative techniques independently developed in each domain; and 3) borrowing techniques from MP-GNNs, we design two new contrastive learning variants, and demonstrate their effectiveness on benchmark datasets.

Original: We summarize our contributions as follows:
Revised: We summarize our contributions as follows

Original:  • Learning Dynamics.
Revised: : • Learning Dynamics.

Original: We reformulate of the contrastive learning into the feature space and develop a new decomposition of the alignment and uniformity loss.
Revised: We reformulate of the contrastive learning into the feature space and develop a new decomposition of the alignment and uniformity loss.

Original: Based on this framework, we show that the alignment and uniformity updates correspond to two different message passing schemes, and characterize the equilibrium states under the combined update.
Revised: Based on this framework, we show that the alignment and uniformity updates correspond to two different message passing schemes, and characterize the equilibrium states under the combined update.

Original: This message perspective provides a new understanding of contrastive learning dynamics.
Revised: This message perspective provides a new understanding of contrastive learning dynamics.

Original: • Connecting CL and MP-GNNs.
Revised: • Connecting CL and MP-GNNs.

Original: Through the message passing perspective of contrastive learning (CL), we show that we can establish an intriguing connection between CL and MP-GNNs.
Revised: Through the message passing perspective of contrastive learning (CL), we show that we can establish an intriguing connection between CL and MP-GNNs.

Original: We not only formally establish the equivalence between alignment update and graph convolution, uniformity update and self-attention, but also point out the inherent analogies between important techniques independently developed in each domain.
Revised: We not only formally establish the equivalence between alignment update and graph convolution, uniformity update and self-attention, but also point out the inherent analogies between important techniques independently developed in each domain.

Original: • New Designs Inspired by MP-GNNs.
Revised: • New Designs Inspired by MP-GNNs.

Original: We also demonstrate the empirical benefits of this connection by designing two new contrastive learning variants borrowing techniques from MP-GNNs: one is to avoid the feature collapse of alignment update by multi-stage aggregation, and one is to adaptively align different positive samples with by incorporating the attention mechanism.
Revised: We also demonstrate the empirical benefits of this connection by designing two new contrastive learning variants borrowing techniques from MP-GNNs: one is to avoid the feature collapse of alignment update by multi-stage aggregation, and one is to adaptively align different positive samples with by incorporating the attention mechanism.

Original: Empirically, we show that both techniques leads to clear benefits on benchmark datasets.
Revised: Empirically, we show that both techniques leads to clear benefits on benchmark datasets.

Original: In turn, their empirical successes also help verify the validness of our established connection between CL and MP-GNNs
Revised: In turn, their empirical successes also help verify the validness of our established connection between CL and MP-GNNs.

Original:  A M ESSAGE P ASSING P ERSPECTIVE ON C ONTRASTIVE L EARNING
Revised:  A M ESSAGE P ASSING P ERSPECTIVE ON C ONTRASTIVE L EARNING

Original: In this section, we develop a message passing perspective for understanding the dynamics of contrastive learning.
Revised: In this section, we develop a message passing perspective for understanding the dynamics of contrastive learning.

Original: We begin by reformulating the contrastive loss into the feature space with a new decomposition.
Revised: We begin by reformulating the contrastive loss into the feature space with a new decomposition.

Original: We then study the update rules derived from the alignment and uniformity losses, and explain their behaviors from a message passing perspective.
Revised: We then study the update rules derived from the alignment and uniformity losses, and explain their behaviors from a message passing perspective.

Original: When combined together, we also characterize how the two updates strike a balance at the equilibrium states.
Revised: When combined together, we also characterize how the two updates strike a balance at the equilibrium states.

Original: And we finish this section with a proof-of-idea to illustrate the effectiveness of our derive message passing rules.
Revised: And we finish this section with a proof-of-idea to illustrate the effectiveness of our derive message passing rules.

Original: 2. B ACKGROUND , R EFORMULATION , AND D ECOMPOSITION
Revised: 2. B ACKGROUND , R EFORMULATION , AND D ECOMPOSITION

Original: We begin our discussion by introducing the canonical formulation of contrastive learning methods in the parameter space, and present their equivalent formulation in the feature space.
Revised: We begin our discussion by introducing the canonical formulation of contrastive learning methods in the parameter space, and present their equivalent formulation in the feature space.

Original: Contrastive Learning (CL).
Revised: Contrastive Learning (CL).

Original: Given two positive samples ( x, x + ) generated by data augmentations, and an independently sampled negative sample x ′ , we can learn an encoder f θ : R d → R m with the wide adopted InfoNCE loss (Oord et al., 2018):
Revised: Given two positive samples ( x, x + ) generated by data augmentations, and an independently sampled negative sample x ′ , we can learn an encoder f θ : R d → R m with the wide adopted InfoNCE loss (Oord et al., 2018):

Original: [Equation]
Revised: [Equation]

Original: where the form term pulls positive samples ( x, x + ) together by encouraging their similarity, and the latter term pushes negative pairs ( x, x ′ ) apart.
Revised: where the form term pulls positive samples ( x, x + ) together by encouraging their similarity, and the latter term pushes negative pairs ( x, x ′ ) apart.

Original: In practice, we typically randomly draw M negative samples to approximate the second term.
Revised: In practice, we typically randomly draw M negative samples to approximate the second term.

Original: In contrastive learning, the encoder is parameterized by deep neural networks, making it hardly amenable for formal analysis.
Revised: In contrastive learning, the encoder is parameterized by deep neural networks, making it hardly amenable for formal analysis.

Original: Wen & Li (2021) resort to single-layer networks with strong assumptions on data distribution, but it is far from practice.
Revised: Wen & Li (2021) resort to single-layer networks with strong assumptions on data distribution, but it is far from practice.

Original: Instead, in this work, we focus on the dynamics of the learned data features f θ ( x ) in the feature space R m .
Revised: Instead, in this work, we focus on the dynamics of the learned data features f θ ( x ) in the feature space R m .

Original: As over-parameterized neural networks are very expressive and adaptive, we assume the feature matrix F θ to be unconstrained (updated freely ignoring dependences on parameters and sample complexity) as in HaoChen et al.
Revised: As over-parameterized neural networks are very expressive and adaptive, we assume the feature matrix F θ to be unconstrained (updated freely ignoring dependences on parameters and sample complexity) as in HaoChen et al.

Original: This assumption is also widely adopted in deep learning theory, such as, layer peeled models (Mixon et al., ; Zhu et al., 2021).
Revised: This assumption is also widely adopted in deep learning theory, such as, layer peeled models (Mixon et al., 2022; Zhu et al., 2021).

Original: The unconstrained feature assumption enables us to focus on the relationship between augmented samples.
Revised: The unconstrained feature assumption enables us to focus on the relationship between augmented samples.

Original: As pointed out by Wang et al.
Revised: As pointed out by Wang et al.

Original: (2022), the support overlap between augmented data is the key for contrastive learning to generalize to downstream tasks.
Revised: (2022), the support overlap between augmented data is the key for contrastive learning to generalize to downstream tasks.

Original: However, we still do not have a clear view of how data augmentations affect feature learning, which is the focus of this work.
Revised: However, we still do not have a clear view of how data augmentations affect feature learning, which is the focus of this work.

Original: For a formal exposition, we model data augmentations through the language of augmentation graph.
Revised: For a formal exposition, we model data augmentations through the language of augmentation graph.

Original: Following HaoChen et al.
Revised: Following HaoChen et al.

Original: (2021), we assume a finite sample space X (can be exponentially large).
Revised: (2021), we assume a finite sample space X (can be exponentially large).

Original: Given a set of natural data as ¯ X = { ¯ x | ¯ x ∈ R d } , we first draw a natural example ¯ x ∈ P d ( ¯ X ) , and then draw a pair of samples x, x + from a random augmentation of ¯ x with distribution A ( ·| ¯ x ) .
Revised: Given a set of natural data as ¯ X = { ¯ x | ¯ x ∈ R d } , we first draw a natural example ¯ x ∈ P d ( ¯ X ) , and then draw a pair of samples x, x + from a random augmentation of ¯ x with distribution A ( ·| ¯ x ) .

Original: We denote X as the collection of augmented “views”, i.e., X = ∪ ¯ x ∈ ¯ X supp ( A ( ·| ¯ x )) .
Revised: We denote X as the collection of augmented “views”, i.e., X = ∪ ¯ x ∈ ¯ X supp ( A ( ·| ¯ x )) .

Original: Regarding the N samples in X as N nodes, we can construct an augmentation graph G =
Revised: Regarding the N samples in X as N nodes, we can construct an augmentation graph G =

Original: ( X , A ) with an adjacency matrix
Revised: ( X , A ) with an adjacency matrix

Original: A representing the mutual connectivity.
Revised: A representing the mutual connectivity.

Original: The edge weight A xx ′ ≥ 0 between x and x ′ is defined as their joint probability A xx ′
Revised: The edge weight A xx ′ ≥ 0 between x and x ′ is defined as their joint probability A xx ′

Original: =
Revised: =

Original: P d ( x, x ′ ) =
Revised: P d ( x, x ′ ) =

Original: E ¯ x A ( x | ¯ x ) A ( x ′ | ¯ x ) .
Revised: E ¯ x A ( x | ¯ x ) A ( x ′ | ¯ x ) .

Original: The normalized adjacency matrix is ¯ A
Revised: The normalized adjacency matrix is ¯ A

Original: = D − 1 / 2 AD − 1 / 2 , where D =
Revised: = D − 1 / 2 AD − 1 / 2 , where D =

Original: deg( A ) denotes the diagonal degree matrix, i.e., D xx =
Revised: deg( A ) denotes the diagonal degree matrix, i.e., D xx =

Original: w x = (cid:80) x
Revised: w x = (cid:80) x

Original: A xx ′ .
Revised: A xx ′ .

Original: The symmetrically normalized Laplacian matrix L = I − ¯ A is known to be positive semi-definite (Chung, 1997).
Revised: The symmetrically normalized Laplacian matrix L = I − ¯ A is known to be positive semi-definite (Chung, 1997).

Original: As natural data and data augmentations are uniformly sampled in practice (Chen et al., 2020b), we assume uniform distribution w x = 1 /N for simplicity.
Revised: As natural data and data augmentations are uniformly sampled in practice (Chen et al., 2020b), we assume uniform distribution w x = 1 /N for simplicity.

Original: New Alignment and Uniformity Losses in Feature Space.
Revised: New Alignment and Uniformity Losses in Feature Space.

Original: Here, we reformulate the contrastive losses into the feature space, and decompose them into two parts, the alignment loss for positive pairs, and the uniformity loss for negative pairs ( L unif for InfoNCE loss, and L (sp)unif for spectral loss):
Revised: Here, we reformulate the contrastive losses into the feature space, and decompose them into two parts, the alignment loss for positive pairs, and the uniformity loss for negative pairs ( L unif for InfoNCE loss, and L (sp)unif for spectral loss):

Original: [Equation]
Revised: [Equation]

Original: where LSE( X ) =
Revised: where LSE( X ) =

Original: Tr( D log(deg( D exp( D − 1 / 2 XD − 1 / 2 )))) serve as a pseudo “norm”.
Revised: Tr( D log(deg( D exp( D − 1 / 2 XD − 1 / 2 )))) serve as a pseudo “norm”.

Original: The following proposition establishes the equivalence between original loss and our reformulated one.
Revised: The following proposition establishes the equivalence between original loss and our reformulated one.

Original: Proposition 1.
Revised: Proposition 1.

Original: When the x -th row of F is defined as F x = √ w x f θ ( x ) , we have:
Revised: When the x -th row of F is defined as F x = √ w x f θ ( x ) , we have:

Original: [Equation]
Revised: [Equation]

Original: Remark on Loss Decomposition.
Revised: Remark on Loss Decomposition.

Original: Notably, our decomposition of contrastive loss is different from the canonical decomposition proposed by Wang & Isola (2020).
Revised: Notably, our decomposition of contrastive loss is different from the canonical decomposition proposed by Wang & Isola (2020).

Original: Taking the InfoNCE loss as an example, they directly take the two terms of Eq.
Revised: Taking the InfoNCE loss as an example, they directly take the two terms of Eq.

Original: 1 as the alignment and uniformity losses, which, however, leads to two improper implications without feature normalization.
Revised: 1 as the alignment and uniformity losses, which, however, leads to two improper implications without feature normalization.

Original: First, minimizing their alignment loss L align ( θ ) = − E x,x + f θ ( x ) ⊤ f θ ( x + ) alone is an ill-posed problem, as it approaches −∞ when we simply scale the norm of F .
Revised: First, minimizing their alignment loss L align ( θ ) = − E x,x + f θ ( x ) ⊤ f θ ( x + ) alone is an ill-posed problem, as it approaches −∞ when we simply scale the norm of F .

Original: In this case, the alignment loss gets smaller while the distance between positive pairs gets larger.
Revised: In this case, the alignment loss gets smaller while the distance between positive pairs gets larger.

Original: Second, minimizing the uniformity loss alone will bridge the distance closer.
Revised: Second, minimizing the uniformity loss alone will bridge the distance closer.

Original: Instead, with our decomposition, the alignment loss (Eq. 2a) is positive definite and will bring samples together with smaller loss, and a larger uniformity loss (Eq. 2b) will expand their distances (Figure 1).
Revised: Instead, with our decomposition, the alignment loss (Eq. 2a) is positive definite and will bring samples together with smaller loss, and a larger uniformity loss (Eq. 2b) will expand their distances (Figure 1).

Original: Thus, our decomposition in Eq.
Revised: Thus, our decomposition in Eq.

Original: 2 seems more properly defined than theirs.
Revised: 2 seems more properly defined than theirs.

Original: Based on the reformulation in Eq. 2, we will study the dynamics of each objective in Section 2.2.1 and give a unified analysis of their equilibrium in Section 2.3, both from a message passing perspective.
Revised: Based on the reformulation in Eq. 2, we will study the dynamics of each objective in Section 2.2.1 and give a unified analysis of their equilibrium in Section 2.3, both from a message passing perspective.

Original: 2. A M ESSAGE P ASSING P ERSPECTIVE ON A LIGNMENT AND U NIFORMITY 2.2. A LIGNMENT U PDATE IS M ESSAGE P ASSING ON A UGMENTATION G RAPH
Revised: 2. A M ESSAGE P ASSING P ERSPECTIVE ON A LIGNMENT AND U NIFORMITY 2.2. A LIGNMENT U PDATE IS M ESSAGE P ASSING ON A UGMENTATION G RAPH

Original: In fact, the reformulated alignment loss L align ( F ) =
Revised: In fact, the reformulated alignment loss L align ( F ) =

Original: Tr( F ⊤ LF ) is widely known as the Laplacian regularization that has wide applications in various graph-related scenarios (Ng et al., 2001;
Revised: Tr( F ⊤ LF ) is widely known as the Laplacian regularization that has wide applications in various graph-related scenarios (Ng et al., 2001;

Original: Ando &
Revised: Ando &

Original: Zhang, 2006).
Revised: Zhang, 2006).

Original: Its gradient descent with step size α > 0 gives the following update rule:
Revised: Its gradient descent with step size α > 0 gives the following update rule:

Original: [Equation]
Revised: [Equation]

Original: One familiar with graphs can immediately recognize that this update rule is a matrix-form message passing scheme of features F ( t ) with a propagation matrix P = (1 − 2 α ) I + 2 α ¯ A .
Revised: One familiar with graphs can immediately recognize that this update rule is a matrix-form message passing scheme of features F ( t ) with a propagation matrix P = (1 − 2 α ) I + 2 α ¯ A .

Original: For those less familiar with the graph literature, let us take a look at the local update of the feature F x of each x :
Revised: For those less familiar with the graph literature, let us take a look at the local update of the feature F x of each x :

Original: [Equation]
Revised: [Equation]

Original: We can see that the alignment loss updates each feature F x by a weighted sum that aggregates features from its neighborhood N x = { x ′ ∈ X| ¯ A xx ′ > 0 } .
Revised: We can see that the alignment loss updates each feature F x by a weighted sum that aggregates features from its neighborhood N x = { x ′ ∈ X| ¯ A xx ′ > 0 } .

Original: In the context of contrastive learning, this reveals that the alignment loss is to propagate features among positive pairs, and in each step, more alike positive pairs x, x + (with higher edge weights A xx + ) will exchange more information.
Revised: In the context of contrastive learning, this reveals that the alignment loss is to propagate features among positive pairs, and in each step, more alike positive pairs x, x + (with higher edge weights A xx + ) will exchange more information.

Original: Next, let us turn our focus to how this message passing scheme affects downstream classification.
Revised: Next, let us turn our focus to how this message passing scheme affects downstream classification.

Original: As for contrastive learning, because positive samples come from the same class, this alignment loss will help cluster same-class samples together.
Revised: As for contrastive learning, because positive samples come from the same class, this alignment loss will help cluster same-class samples together.

Original: However, it does not mean that any augmentations can attain a minimal intra-class variance.
Revised: However, it does not mean that any augmentations can attain a minimal intra-class variance.

Original: Below, we formally characterize the effect of data augmentations on the feature clustering of each class k .
Revised: Below, we formally characterize the effect of data augmentations on the feature clustering of each class k .

Original: We measure the degree of clustering by the distance between all intra-class features to a one-dimensional subspace M .
Revised: We measure the degree of clustering by the distance between all intra-class features to a one-dimensional subspace M .

Original: Proposition 2.
Revised: Proposition 2.

Original: Denote the distance between a feature matrix X and a subspace Ω by d Ω ( X ) :=
Revised: Denote the distance between a feature matrix X and a subspace Ω by d Ω ( X ) :=

Original: inf {∥ X − Y ∥
Revised: inf {∥ X − Y ∥

Original: F | Y ∈ Ω } .
Revised: F | Y ∈ Ω } .

Original: Assume that data augmentations are label-preserving.
Revised: Assume that data augmentations are label-preserving.

Original: Then, for samples in each class k , denoting their adjacency matrix as ¯ A k and their features at t -th step as F ( t ) k , the message passing in Eq.
Revised: Then, for samples in each class k , denoting their adjacency matrix as ¯ A k and their features at t -th step as F ( t ) k , the message passing in Eq.

Original: 4 leads to a smaller or equal distance to an one-dimensional subspace M
Revised: 4 leads to a smaller or equal distance to an one-dimensional subspace M

Original: [Equation]
Revised: [Equation]

Original: [Equation] ) where λ G k ∈ [0 , 2] denotes the intra-class algebraic connectivity of G (Chung, 1997), and M = { e 1 ⊗ y | y ∈ R c } , and e 1 is the eigenvector corresponding to the largest eigenvalue of ¯ A k .
Revised: where λ G k ∈ [0 , 2] denotes the intra-class algebraic connectivity of G (Chung, 1997), and M = { e 1 ⊗ y | y ∈ R c } , and e 1 is the eigenvector corresponding to the largest eigenvalue of ¯ A k .

Original: An important message of Proposition 2 is that the alignment update does not necessarily bring better clustering, which depends on the algebraic connectivity λ G k : 1) only a non-zero λ G k guarantees the strict decrease of intra-class distance, and 2) a larger λ G k indicates an even faster decrease.
Revised: An important message of Proposition 2 is that the alignment update does not necessarily bring better clustering, which depends on the algebraic connectivity λ G k : 1) only a non-zero λ G k guarantees the strict decrease of intra-class distance, and 2) a larger λ G k indicates an even faster decrease.

Original: According to spectral graph theory, λ G k is the second-smallest eigenvalue of the Laplacian L , and its magnitude reflects how connected the graph is.
Revised: According to spectral graph theory, λ G k is the second-smallest eigenvalue of the Laplacian L , and its magnitude reflects how connected the graph is.

Original: In particular, λ G k > 0 holds if and only if the intra-class graph is connected (Chung, 1997).
Revised: In particular, λ G k > 0 holds if and only if the intra-class graph is connected (Chung, 1997).

Original: Indeed, when the class is not connected, the alignment update can never bridge the disjoint components together.
Revised: Indeed, when the class is not connected, the alignment update can never bridge the disjoint components together.

Original: While in contrastive learning, the design of data augmentations determines G k as well as λ G k .
Revised: While in contrastive learning, the design of data augmentations determines G k as well as λ G k .

Original: Therefore, our message passing analysis quantitatively characterize how augmentations affect feature clustering in the contrastive learning dynamics.
Revised: Therefore, our message passing analysis quantitatively characterize how augmentations affect feature clustering in the contrastive learning dynamics.

Original: 2.2. U NIFORMITY U PDATE IS M ESSAGE P ASSING ON THE A FFINITY G RAPH
Revised: 2.2. U NIFORMITY U PDATE IS M ESSAGE P ASSING ON THE A FFINITY G RAPH

Original: Contrary to the alignment loss, the uniformity loss instead penalizes the similarity between every sample pair x, x ′ ∈ X to encourage feature uniformity.
Revised: Contrary to the alignment loss, the uniformity loss instead penalizes the similarity between every sample pair x, x ′ ∈ X to encourage feature uniformity.

Original: Perhaps surprisingly, from a feature space perspective, its gradient descent rule, namely the uniformity update, also corresponds to a message passing scheme with the affinity graph G ′ with a “fake” adjacency matrix A ′ : 
Revised: Perhaps surprisingly, from a feature space perspective, its gradient descent rule, namely the uniformity update, also corresponds to a message passing scheme with the affinity graph G ′ with a “fake” adjacency matrix A ′ : [Equation]

Original: [Equation]
Revised: [Equation]

Original: where ¯ A ′ sym
Revised: where ¯ A ′ sym

Original: = D ′ exp − 1 A ′ exp + A ′ exp D ′
Revised: = D ′ exp − 1 A ′ exp + A ′ exp D ′

Original: exp − 1 , and A ′ exp =
Revised: exp − 1 , and A ′ exp =

Original: exp( D − 1 / 2 F ( t ) F ( t ) ⊤ D − 1 / 2 ) , D ′ exp =
Revised: exp( D − 1 / 2 F ( t ) F ( t ) ⊤ D − 1 / 2 ) , D ′ exp =

Original: deg( A ′ exp ) .
Revised: deg( A ′ exp ) .

Original: When we further apply stop gradient to the target branch (as adopted in BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2021)), we can further simplify it to (global update with stop gradient) F ( t +1) =
Revised: When we further apply stop gradient to the target branch (as adopted in BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2021)), we can further simplify it to (global update with stop gradient) F ( t +1) =

Original: [(1+ α )
Revised: [(1+ α )

Original: I − α ¯ A ′ ]
Revised: I − α ¯ A ′ ]

Original: F ( t ) , where ¯ A ′ =
Revised: F ( t ) , where ¯ A ′ =

Original: D ′ exp − 1 A ′ exp .
Revised: D ′ exp − 1 A ′ exp .

Original: We can also interpret the uniformity update as a specific message passing scheme on the affinity graph G ′ =
Revised: We can also interpret the uniformity update as a specific message passing scheme on the affinity graph G ′ =

Original: ( X , A ′ ) with node set X and adjacency matrix A ′ .
Revised: ( X , A ′ ) with node set X and adjacency matrix A ′ .

Original: In comparison to the augmentation graph G that defines the edge weights via data augmentations, the affinity graph is instead constructed through the estimated feature similarity between ( x, x ′ ) .
Revised: In comparison to the augmentation graph G that defines the edge weights via data augmentations, the affinity graph is instead constructed through the estimated feature similarity between ( x, x ′ ) .

Original: The contrastive learning with two graphs resembles the wake-sleep algorithm in classical unsupervised learning (Hinton et al., 1995).
Revised: The contrastive learning with two graphs resembles the wake-sleep algorithm in classical unsupervised learning (Hinton et al., 1995).

Original: Further, we notice that the affinity matrix ¯ A ′ is induced by the RBF kernel k RBF ( u, v ) =
Revised: Further, we notice that the affinity matrix ¯ A ′ is induced by the RBF kernel k RBF ( u, v ) =

Original: exp( −∥ u − v ∥ 2 / 2) when the features are normalized, i.e., ∥ f ( x ) ∥ = 1 (often adopted in practice (Chen et al., 2020b)).
Revised: exp( −∥ u − v ∥ 2 / 2) when the features are normalized, i.e., ∥ f ( x ) ∥ = 1 (often adopted in practice (Chen et al., 2020b)).

Original: Generally speaking, the RBF kernel is more preferable than the linear kernel as it is more non-Euclidean and non-parametric.
Revised: Generally speaking, the RBF kernel is more preferable than the linear kernel as it is more non-Euclidean and non-parametric.

Original: Besides, because the message passing in the alignment update utilizes ¯ A , which is also non-negative and normalized, the InfoNCE uniformity update with ¯ A ′ nce can be more stable and require less hyperparameter tuning to cooperate with the alignment update.
Revised: Besides, because the message passing in the alignment update utilizes ¯ A , which is also non-negative and normalized, the InfoNCE uniformity update with ¯ A ′ nce can be more stable and require less hyperparameter tuning to cooperate with the alignment update.

Original: Thus, our message passing view not only provides a new understanding of the uniformity update, but also helps explain the popularity of the InfoNCE loss.
Revised: Thus, our message passing view not only provides a new understanding of the uniformity update, but also helps explain the popularity of the InfoNCE loss.

Original: [Figure]
Revised: [Figure]

Original: T HE E QUILIBRIUM B ETWEEN
Revised: [Figure]

Original: THE 
Revised: THE A LIGNMENT AND U NIFORMITY U PDATE

Original: A LIGNMENT AND U NIFORMITY U PDATE

Original: Combining the alignment update and the uniformity update discussed above, we arrive at our final update rule for contrastive learning with the (regularized) InfoNCE loss (with stop gradient):
Revised: Combining the alignment update and the uniformity update discussed above, we arrive at our final update rule for contrastive learning with the (regularized) InfoNCE loss (with stop gradient):

Original: [Equation]
Revised: [Equation]

Original: which is a special message passing scheme using the difference between the augmentation (data) and affinity (estimated) adjacency matrices A diff =
Revised: which is a special message passing scheme using the difference between the augmentation (data) and affinity (estimated) adjacency matrices A diff =

Original: ¯
Revised: ¯

Original: A − ¯ A ′ .
Revised: A − ¯ A ′ .

Original: Accordingly, we can reformulate the contrastive loss more intuitively as:
Revised: Accordingly, we can reformulate the contrastive loss more intuitively as:

Original: [Equation]
Revised: [Equation]

Original: where sg( · ) means the stop gradient operator.
Revised: where sg( · ) means the stop gradient operator.

Original: Equilibrium Distribution.
Revised: Equilibrium Distribution.

Original: The reformulation above (Eq. 10) also reveals the stationary point of the learning dynamics, that is, ¯ A diff =
Revised: The reformulation above (Eq. 10) also reveals the stationary point of the learning dynamics, that is, ¯ A diff =

Original: ¯ A − A ′ = 0 , which implies the model equilibrium as follows.
Revised: ¯ A − A ′ = 0 , which implies the model equilibrium as follows.

Original: Proposition 3.
Revised: Proposition 3.

Original: The contrastive learning dynamics with the InfoNCE loss saturates when the real and the fake adjacency matrices agree, i.e., ¯ A
Revised: The contrastive learning dynamics with the InfoNCE loss saturates when the real and the fake adjacency matrices agree, i.e., ¯ A

Original: = ¯ A ′ nce , which means that element-wisely, we have:
Revised: = ¯ A ′ nce , which means that element-wisely, we have:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: That is, the ground-truth data conditional distribution P d ( x + | x ) equals to the Boltzmann distribution
Revised: That is, the ground-truth data conditional distribution P d ( x + | x ) equals to the Boltzmann distribution

Original: P θ ( x + | x ) estimated from the encoder features.
Revised: P θ ( x + | x ) estimated from the encoder features.

Original: In this way, our analysis reveals that contrastive learning (with InfoNCE loss) implicitly learns a probabilistic model P θ ( x + | x ) (Eq. 11), and its goal is to approximate the conditional data distribution
Revised: In this way, our analysis reveals that contrastive learning (with InfoNCE loss) implicitly learns a probabilistic model P θ ( x + | x ) (Eq. 11), and its goal is to approximate the conditional data distribution

Original: P d ( x + | x ) defined via the augmentation graph G .
Revised: P d ( x + | x ) defined via the augmentation graph G .

Original: This explains the effectiveness of contrastive learning on out-of-distribution detection tasks (Winkens et al., 2020) that rely on good density estimation.
Revised: This explains the effectiveness of contrastive learning on out-of-distribution detection tasks (Winkens et al., 2020) that rely on good density estimation.

Original: In practical implementations of InfoNCE, e.g., SimCLR, a temperature scalar τ is often introduced to rescale the cosine similarities.
Revised: In practical implementations of InfoNCE, e.g., SimCLR, a temperature scalar τ is often introduced to rescale the cosine similarities.

Original: With a flexible choice of τ , the model distribution P θ could better approximate sharp (like one-hot) data distribution P d .
Revised: With a flexible choice of τ , the model distribution P θ could better approximate sharp (like one-hot) data distribution P d .

Original: A detailed comparison to related work is included in Appendix C.
Revised: A detailed comparison to related work is included in Appendix A.

Original: 2. A P ROOF - OF - IDEA E XPERIMENT
Revised: 2. A P ROOF - OF - IDEA E XPERIMENT

Original: At last, we verify the derived three message passing rules (alignment, uniformity, contrastive) through a synthetic experiment using on a synthetic augmentation graph to mimic contrastive learning in the feature space.
Revised: At last, we verify the derived three message passing rules (alignment, uniformity, contrastive) through a synthetic experiment using on a synthetic augmentation graph to mimic contrastive learning in the feature space.

Original: From Figure 1, we can see that the results generally align well with our theory.
Revised: From Figure 1, we can see that the results generally align well with our theory.

Original: First, the alignment update (Section 2.2.1) indeed bridges intra-class samples closer, but there is also a risk of feature collapse.
Revised: First, the alignment update (Section 2.2.1) indeed bridges intra-class samples closer, but there is also a risk of feature collapse.

Original: Second, the uniformity update (Section 2.2.2) indeed keeps all samples uniformly distributed.
Revised: Second, the uniformity update (Section 2.2.2) indeed keeps all samples uniformly distributed.

Original: Third, by combining the alignment and uniformity updates, the contrastive update (Section 2.3) can successfully cluster intra-class samples together while keeping inter-class samples apart, even after long iterations.
Revised: Third, by combining the alignment and uniformity updates, the contrastive update (Section 2.3) can successfully cluster intra-class samples together while keeping inter-class samples apart, even after long iterations.

Original: In this way, we show that like the actual learning behaviors of contrastive learning, the derived contrastive message passing rules can indeed learn stable and discriminative features from unlabeled data.
Revised: In this way, we show that like the actual learning behaviors of contrastive learning, the derived contrastive message passing rules can indeed learn stable and discriminative features from unlabeled data.

Original:  C ONTRASTIVE L EARNING (CL) AND M ESSAGE P ASSING G RAPH N EURAL N ETWORKS (MP-GNN S ): C ONNECTIONS , A NALOGIES , AND N EW D ESIGNS
Revised:  C ONTRASTIVE L EARNING (CL) AND M ESSAGE P ASSING G RAPH N EURAL N ETWORKS (MP-GNN S ): C ONNECTIONS , A NALOGIES , AND N EW D ESIGNS

Original: In the previous section, we have shown the contrastive learning dynamics corresponds nicely to a message passing scheme on two different graphs, ¯ A and ¯ A ′ .
Revised: In the previous section, we have shown the contrastive learning dynamics corresponds nicely to a message passing scheme on two different graphs, ¯ A and ¯ A ′ .

Original: In a wider context, message passing is a general technique with broad applications in graph theory, Bayesian inference, graph embedding, etc.
Revised: In a wider context, message passing is a general technique with broad applications in graph theory, Bayesian inference, graph embedding, etc.

Original: As a result, our interpretation of contrastive learning from a message passing perspective enables us to establish some interesting connections between contrastive learning and these different areas.
Revised: As a result, our interpretation of contrastive learning from a message passing perspective enables us to establish some interesting connections between contrastive learning and these different areas.

Original: A noticeable application of message passing is graph representation learning.
Revised: A noticeable application of message passing is graph representation learning.

Original: In recent years, Message
Revised: In recent years, Message

Original: Passing Graph Neural Networks (MP-GNNs) become the de facto paradigm for GNNs.
Revised: Passing Graph Neural Networks (MP-GNNs) become the de facto paradigm for GNNs.

Original: Various techniques have been explored in terms of different kinds of graph structures and learning tasks.
Revised: Various techniques have been explored in terms of different kinds of graph structures and learning tasks.

Original: By drawing a connection between contrastive learning and MP-GNNs, we could leverage these techniques to design better contrastive learning algorithms, as is the goal of this section.
Revised: By drawing a connection between contrastive learning and MP-GNNs, we could leverage these techniques to design better contrastive learning algorithms, as is the goal of this section.

Original: To achieve this, we first establish the basic connections between the basic paradigms in Section 3.1.
Revised: To achieve this, we first establish the basic connections between the basic paradigms in Section 3.1.

Original: Built upon these connections, in Section 3.2, we point out the common ideas behind many representative techniques independently proposed by each community in different names.
Revised: Built upon these connections, in Section 3.2, we point out the common ideas behind many representative techniques independently proposed by each community in different names.

Original: Lastly, we show that we could further borrow some new techniques from the MP-GNN literature for designing new contrastive learning methods in Section 3.3.
Revised: Lastly, we show that we could further borrow some new techniques from the MP-GNN literature for designing new contrastive learning methods in Section 3.3.

Original: 3. C ONNECTIONS BETWEEN THE T WO P ARADIGMS
Revised: 3. C ONNECTIONS BETWEEN THE T WO P ARADIGMS

Original: Regarding the alignment update and the uniformity update in contrastive learning, we show that they do have a close connection to message passing schemes in graph neural networks.
Revised: Regarding the alignment update and the uniformity update in contrastive learning, we show that they do have a close connection to message passing schemes in graph neural networks.

Original: Graph Convolution and Alignment Update.
Revised: Graph Convolution and Alignment Update.

Original: A graph data with N nodes typically consist of two inputs, the input feature X ∈ R N × d , and the (normalized) adjacency matrix ¯ A ∈ R n × n .
Revised: A graph data with N nodes typically consist of two inputs, the input feature X ∈ R N × d , and the (normalized) adjacency matrix ¯ A ∈ R n × n .

Original: For a feature matrix H ( t ) ∈ R n × m at the t -th layer, the canonical graph convolution scheme in GCN (Kipf & Welling, 2017) gives H ( t +1) =
Revised: For a feature matrix H ( t ) ∈ R n × m at the t -th layer, the canonical graph convolution scheme in GCN (Kipf & Welling, 2017) gives H ( t +1) =

Original: ¯ AH ( t ) .
Revised: ¯ AH ( t ) .

Original: Later, DGC (Wang et al., 2021) generalizes it to
Revised: Later, DGC (Wang et al., 2021) generalizes it to

Original: H ( t +1) =
Revised: H ( t +1) =

Original: (1 − ∆ t ) I + ∆ t ¯ AH ( t ) with a flexible step size ∆ t for finite difference, which reduces the canonical form when ∆ t = 1 .
Revised: (1 − ∆ t ) I + ∆ t ¯ AH ( t ) with a flexible step size ∆ t for finite difference, which reduces the canonical form when ∆ t = 1 .

Original: Comparing it to the alignment update rule in Eq.
Revised: Comparing it to the alignment update rule in Eq.

Original: 4, we can easily notice their equivalence when we take the augmentation graph ¯
Revised: 4, we can easily notice their equivalence when we take the augmentation graph ¯

Original: A
Revised: A

Original: as the input graph for a GCN (or
Revised: as the input graph for a GCN (or

Original: DGC), and adopt a specific step size ∆ t = 2 α .
Revised: DGC), and adopt a specific step size ∆ t = 2 α .

Original: In other words, the gradient descent of F using the alignment loss (an optimization step) is equivalent to the graph convolution of F based on the augmentation graph ¯ A (an architectural module).
Revised: In other words, the gradient descent of F using the alignment loss (an optimization step) is equivalent to the graph convolution of F based on the augmentation graph ¯ A (an architectural module).

Original: Their differences in names and forms are only superficial, since inherently, they share the same message passing mechanism.
Revised: Their differences in names and forms are only superficial, since inherently, they share the same message passing mechanism.

Original: Oversmoothing and Feature Collapse.
Revised: Oversmoothing and Feature Collapse.

Original: The equivalence above is the key for us to unlock many intriguing connections between the two seemingly different areas.
Revised: The equivalence above is the key for us to unlock many intriguing connections between the two seemingly different areas.

Original: A particular example is the oversmoothing problem, one of the centric topics in MP-GNNs with tons of discussions, to name a few, Li et al.
Revised: A particular example is the oversmoothing problem, one of the centric topics in MP-GNNs with tons of discussions, to name a few, Li et al.

Original: (2018), Oono & Suzuki (2020), Shi et al.
Revised: (2018), Oono & Suzuki (2020), Shi et al.

Original: Oversmoothing refers to the phenomenon that when we apply the graph convolution above for many times, the node features will become indistinguishable from each other and lose discriminative ability.
Revised: Oversmoothing refers to the phenomenon that when we apply the graph convolution above for many times, the node features will become indistinguishable from each other and lose discriminative ability.

Original: Meanwhile, we have also observed the same phenomenon for the alignment update in Figure 1.
Revised: Meanwhile, we have also observed the same phenomenon for the alignment update in Figure 1.

Original: Indeed, it is also well known in contrastive learning that with the alignment loss alone, the features will collapse to a single point, known as complete feature collapse .
Revised: Indeed, it is also well known in contrastive learning that with the alignment loss alone, the features will collapse to a single point, known as complete feature collapse .

Original: Hua et al.
Revised: Hua et al.

Original: (2021) show that when equipped with BN at the output, the features will not fully collapse but still collapse to a low dimensional subspace, named dimensional collapse .
Revised: (2021) show that when equipped with BN at the output, the features will not fully collapse but still collapse to a low dimensional subspace, named dimensional collapse .

Original: In fact, both two feature collapse scenarios with the augmentation graph can be characterized using existing oversmoothing analysis of a general graph (Oono & Suzuki, 2020).
Revised: In fact, both two feature collapse scenarios with the augmentation graph can be characterized using existing oversmoothing analysis of a general graph (Oono & Suzuki, 2020).

Original: Therefore, this connection also opens new paths for the theoretical analysis of contrastive learning.
Revised: Therefore, this connection also opens new paths for the theoretical analysis of contrastive learning.

Original: Self-Attention and Uniformity Update.
Revised: Self-Attention and Uniformity Update.

Original: Besides the connection in the alignment update, we also notice an interesting connection between the uniformity update rule and the self-attention mechanism.
Revised: Besides the connection in the alignment update, we also notice an interesting connection between the uniformity update rule and the self-attention mechanism.

Original:  Self-attention is the key component of Transformer (Vaswani et al., 2017), which recently obtains promising progresses in various scenarios like images, video, text, graph, molecules (Lin et al., 2021).
Revised: [Table] [Table] Self-attention is the key component of Transformer (Vaswani et al., 2017), which recently obtains promising progresses in various scenarios like images, video, text, graph, molecules (Lin et al., 2021) [Table]

Original: For an input feature H ( t ) ∈ R N × m at the t -th step, the self-attention module gives update:
Revised: For an input feature H ( t ) ∈ R N × m at the t -th step, the self-attention module gives update:

Original: [Equation]
Revised: [Equation]

Original: H

Original: [Equation]
Revised: [Equation]

Original: t ) , where¯ A
Revised: ) , where

Original: [Equation]
Revised: [Equation]

Original: When comparing it to the uniformity update in Eq. 8, we can see that the two rules are equivalent when α = − 1 .
Revised: When comparing it to the uniformity update in Eq. 8, we can see that the two rules are equivalent when α = − 1 .

Original: Because the step size is negative, self-attention is actually a gradient ascent step that will maximize the feature uniformity loss (Eq. 2b).
Revised: Because the step size is negative, self-attention is actually a gradient ascent step that will maximize the feature uniformity loss (Eq. 2b).

Original: Consequently, stacking self-attention will reduce the uniformity between features and lead to more collapsed representations, which explains the feature collapse behavior observed in Transformers (Shi et al., 2022).
Revised: Consequently, stacking self-attention will reduce the uniformity between features and lead to more collapsed representations, which explains the feature collapse behavior observed in Transformers (Shi et al., 2022).

Original: Different from previous work (Dong et al., 2021), our analysis above stems from the connection between contrastive learning and MP-GNNs, and it provides a simple explanation of this phenomenon from an optimization perspective.
Revised: Different from previous work (Dong et al., 2021), our analysis above stems from the connection between contrastive learning and MP-GNNs, and it provides a simple explanation of this phenomenon from an optimization perspective.

Original: 3. A NALOGIES IN E XISTING T ECHNIQUES
Revised: 3. A NALOGIES IN E XISTING T ECHNIQUES

Original: In the discussion above, we establish connections between the two paradigms.
Revised: In the discussion above, we establish connections between the two paradigms.

Original: Besides, we know that many important techniques and variants have been built upon the basic paradigms in each community.
Revised: Besides, we know that many important techniques and variants have been built upon the basic paradigms in each community.

Original: Here, through a unified perspective of the two domains, we spot some interesting connections between techniques developed independently in each domain, while the inherent connections are never revealed.
Revised: Here, through a unified perspective of the two domains, we spot some interesting connections between techniques developed independently in each domain, while the inherent connections are never revealed.

Original: Below, we give two examples to illustrate benefits from these analogies.
Revised: Below, we give two examples to illustrate benefits from these analogies.

Original: NodeNorm / LayerNorm and ℓ 2 Normalization.
Revised: NodeNorm / LayerNorm and ℓ 2 Normalization.

Original: In contrastive learning, since SimCLR (Chen et al., 2020b), a common practice is to apply ℓ 2 normalization before calculating the InfoNCE loss, i.e., f ( x ) / ∥
Revised: In contrastive learning, since SimCLR (Chen et al., 2020b), a common practice is to apply ℓ 2 normalization before calculating the InfoNCE loss, i.e., f ( x ) / ∥

Original: f ( x ) ∥ .
Revised: f ( x ) ∥ .

Original: As a result, the alignment loss becomes equivalent to the cosine similarity between the features of positive pairs.
Revised: As a result, the alignment loss becomes equivalent to the cosine similarity between the features of positive pairs.

Original: Yet, little was known about the real benefits of ℓ 2 normalization.
Revised: Yet, little was known about the real benefits of ℓ 2 normalization.

Original: As another side of the coin, this node-wise feature normalization technique, in the name of NodeNorm (Zhou et al., 2021), has already been shown as an effective approach to mitigate the performance degradation of deep GCNs by controlling the feature variance.
Revised: As another side of the coin, this node-wise feature normalization technique, in the name of NodeNorm (Zhou et al., 2021), has already been shown as an effective approach to mitigate the performance degradation of deep GCNs by controlling the feature variance.

Original: Zhou et al.
Revised: Zhou et al.

Original: (2021) also show that LayerNorm (Ba et al., 2016) with both centering and normalization can also obtain similar performance.
Revised: (2021) also show that LayerNorm (Ba et al., 2016) with both centering and normalization can also obtain similar performance.

Original: These discussions in MP-GNNs also help us understand the effectiveness of ℓ 2 normalization.
Revised: These discussions in MP-GNNs also help us understand the effectiveness of ℓ 2 normalization.

Original: PairNorm / BatchNorm and Feature Centering.
Revised: PairNorm / BatchNorm and Feature Centering.

Original: Another common technique in contrastive learning is feature centering, i.e., f ( x ) − µ , where µ is the mean of all sample (node) features.
Revised: Another common technique in contrastive learning is feature centering, i.e., f ( x ) − µ , where µ is the mean of all sample (node) features.

Original: In DINO (Caron et al., 2021), it is adopted to prevent feature collapse.
Revised: In DINO (Caron et al., 2021), it is adopted to prevent feature collapse.

Original: Similarly, Hua et al.
Revised: Similarly, Hua et al.

Original: (2021) also show that BatchNorm (Ioffe & Szegedy, 2015) (with feature centering as a part) can alleviate feature collapse.
Revised: (2021) also show that BatchNorm (Ioffe & Szegedy, 2015) (with feature centering as a part) can alleviate feature collapse.

Original: On the other side, PairNorm (Zhao & Akoglu, 2020) combines feature centering and feature normalization (mentioned above) to alleviate oversmoothing.
Revised: On the other side, PairNorm (Zhao & Akoglu, 2020) combines feature centering and feature normalization (mentioned above) to alleviate oversmoothing.

Original: In particular, they show that PairNorm can preserve the total feature distance of the input (initial features).
Revised: In particular, they show that PairNorm can preserve the total feature distance of the input (initial features).

Original: This, in turn, could be adopted as an explanation for how feature centering works in contrastive learning.
Revised: This, in turn, could be adopted as an explanation for how feature centering works in contrastive learning.

Original: 3. I NSPIRED N EW D ESIGNS
Revised: 3. I NSPIRED N EW D ESIGNS

Original: Besides the analogies drawn above, we find that there are also many other valuable techniques from the MP-GNN literature that have not found their (exact) counterparts in contrastive learning.
Revised: Besides the analogies drawn above, we find that there are also many other valuable techniques from the MP-GNN literature that have not found their (exact) counterparts in contrastive learning.

Original: In this section, we explore two new connections from different aspects: multi-stage aggregation for avoiding feature collapse, and graph attention for better feature alignment. 
Revised: In this section, we explore two new connections from different aspects: multi-stage aggregation for avoiding feature collapse, and graph attention for better feature alignment. [Figure]

Revised:      / L QHD U  $ FF X U D F\

Revised: [Figure]

Original: 3.3. N EW C ONNECTION I: M ULTI - STAGE G RAPH A GGREGATION
Revised: 3.3. N EW C ONNECTION I: M ULTI - STAGE G RAPH A GGREGATION

Original: As discussed above, the alignment loss alone will lead to feature collapse in contrastive learning.
Revised: As discussed above, the alignment loss alone will lead to feature collapse in contrastive learning.

Original: Although existing works like BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2021) can alleviate feature collapse via asymmetric architectural designs, we are still not fully clear how they work.
Revised: Although existing works like BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2021) can alleviate feature collapse via asymmetric architectural designs, we are still not fully clear how they work.

Original: Instead, the connection between feature collapse and oversmoothing (Section 3.1) inspires us to borrow techniques from the MP-GNNs to alleviate the feature collapse issue of contrastive learning.
Revised: Instead, the connection between feature collapse and oversmoothing (Section 3.1) inspires us to borrow techniques from the MP-GNNs to alleviate the feature collapse issue of contrastive learning.

Original: In the MP-GNN literature, a principled solution to oversmoothing is the jump knowledge technique proposed in JK-Net (Xu et al., 2018).
Revised: In the MP-GNN literature, a principled solution to oversmoothing is the jump knowledge technique proposed in JK-Net (Xu et al., 2018).

Original: Instead of adopting features only from the last step H ( T ) (which might already oversmooth), JK-Net aggregates features from multiple stages, e.g., H (1) , . . ., H ( T ) , [Figure]
Revised: Instead of adopting features only from the last step H ( T ) (which might already oversmooth), JK-Net aggregates features from multiple stages, e.g., H (1) , . . .

Original:      / L QHD U  $ FF X U D F\

Original: [Figure] as the final feature for classification.
Revised: , H ( T ) , as the final feature for classification.

Original: In this way, the knowledge from previous non-over-smoothing steps could directly jump into the final output, which helps improve its discriminative ability.
Revised: In this way, the knowledge from previous non-over-smoothing steps could directly jump into the final output, which helps improve its discriminative ability.

Original: This idea of multi-stage aggregation is also widely adopted in many other MP-GNNs, such as APPNP (Klicpera et al., 2019), GCNII (Chen et al., 2020a), SIGN (Rossi et al., 2020), etc.
Revised: This idea of multi-stage aggregation is also widely adopted in many other MP-GNNs, such as APPNP (Klicpera et al., 2019), GCNII (Chen et al., 2020a), SIGN (Rossi et al., 2020), etc.

Original: In this work, we transfer this idea to the contrastive learning scenario, that is, to alleviate feature collapse by combining features from multiple previous epochs.
Revised: In this work, we transfer this idea to the contrastive learning scenario, that is, to alleviate feature collapse by combining features from multiple previous epochs.

Original: Specifically, we create a memory bank M to store the latest s -epoch features from the same natural sample ¯ x .
Revised: Specifically, we create a memory bank M to store the latest s -epoch features from the same natural sample ¯ x .

Original: At the t -th epoch, for x generated from ¯ x , we have stored old features of it as z ( t − 1) ¯ x 
Revised: At the t -th epoch, for x generated from ¯ x , we have stored old features of it as z ( t − 1) ¯ x , z ( [Equation] , z ( t −

Original: [Equation]
Revised: [Equation]

Original:  where r = min( s, t ) .
Revised: ) ¯ x where r = min( s, t ) .

Original: Then we replace the original positive feature f θ ( x + ) with the aggregation of the old features in the memory bank, where we simply adopt the sum aggregation for simplicity, i.e., z ¯ x = 1 r (cid:80) ri =1 z ( t − i )
Revised: Then we replace the original positive feature f θ ( x + ) with the aggregation of the old features in the memory bank, where we simply adopt the sum aggregation for simplicity, i.e., z ¯ x = 1 r (cid:80) ri =1 z ( t − i )

Original: ¯ x .
Revised: ¯ x .

Original: Next, we optimize the encoder network f θ with the following multi-stage alignment loss alone ,
Revised: Next, we optimize the encoder network f θ with the following multi-stage alignment loss alone ,

Original: [Equation]
Revised: [Equation]

Original: Afterwards, we will push the newest feature z = f θ ( x ) to the memory bank, and drop the oldest one if exceeded.
Revised: Afterwards, we will push the newest feature z = f θ ( x ) to the memory bank, and drop the oldest one if exceeded.

Original: In this way, we could align the new feature with multiple old features that are less collapsed, which, in turn, also prevents the collapse of the updated features.
Revised: In this way, we could align the new feature with multiple old features that are less collapsed, which, in turn, also prevents the collapse of the updated features.

Original: Result Analysis.
Revised: Result Analysis.

Original: We empirically evaluate this method on CIFAR-10 with ResNet-18 backbone (details in Appendix B.1). From Figure 2a, we can see that as we adopt a larger memory bank with a large aggregation step s , the linearprobing accuracycontinues to grow larger and larger. Although the overall accuracy is still relatively low, we can observe a clear benefit of multi-stage aggregation for preventing feature collapse, which previously has onlybeen possible via asymmetric architectural designs.
Revised: From Figure 2a, we can see that when directly applying the multi-stage alignment objective, a larger number of aggregation epochs could indeed improve linear accuracy and avoid full feature collapse. Although the overall accuracy is still relatively low, we can observe a clear benefit of multi-stage aggregation (better than random guess when s ≥ 2 ), which was only possible using asymmetric modules, e.g., SimSiam’s predictor (Chen & He, 2021).

Original: The success of applying multi-stage aggregation  to contrastive learning not only suggests a new approach for negative-free contrastive learning, but also helps verify our established connection between contrastive learning and MP-GNNs. In Section D, we also show that theproposed multi-stage aggeregation can also bring clear benefits  on benchmark datasets.
Revised: Inspired by this observation, we further combine multi-stage aggregation with SimSiam to further alleviate feature collapse. As shown in Table 1a, the multi-stage mechanism indeed brings cosistent and significant improvements over SimSiam on all three datasets, which clearly demonstrates the benefits of multi-stage aggregation on improving existing non-contrastive methods by further alleviating their feature collapse issues.

Revised: Experiment details can be found in Appendix C.1.

Original: 3.3. N EW C ONNECTION II: G RAPH A TTENTION
Revised: 3.3. N EW C ONNECTION II: G RAPH A TTENTION

Original: A noticeable defect of existing contrastive loss is that they assign an equal weight for different positive pairs.
Revised: A noticeable defect of existing contrastive loss is that they assign an equal weight for different positive pairs.

Original: As these positive samples are generated by input-agnostic random data augmentations, there could be a large difference in semantics between different positive samples.
Revised: As these positive samples are generated by input-agnostic random data augmentations, there could be a large difference in semantics between different positive samples.

Original: As marked in red (Figure 2b), some augmented views may even lose the center objects, and aligning with these bad cases will bridge samples from different classes together and distort the decision boundary.
Revised: As marked in red (Figure 2b), some augmented views may even lose the center objects, and aligning with these bad cases will bridge samples from different classes together and distort the decision boundary.

Original: Therefore, in the alignment update, it is not proper to assign equal weights as in vanilla alignment loss (Eq. 5).
Revised: Therefore, in the alignment update, it is not proper to assign equal weights as in vanilla alignment loss (Eq. 5).

Original: Instead, we should adaptively assign different weights to different pairs based on semantic similarities .
Revised: Instead, we should adaptively assign different weights to different pairs based on semantic similarities .

Original: Revisiting MP-GNNs, we find that there are well-established techniques to address this issue.
Revised: Revisiting MP-GNNs, we find that there are well-established techniques to address this issue.

Original: A well-known extension of graph convolution is graph attention proposed in GAT (Veliˇckovi´c et al., 2018), and Transformers can be seen as a special case of GAT applied to a fully connected graph.
Revised: A well-known extension of graph convolution is graph attention proposed in GAT (Veliˇckovi´c et al., 2018), and Transformers can be seen as a special case of GAT applied to a fully connected graph.

Original: GAT extends GCN by taking the semantic similarities between neighborhood features (corresponding to positive pairs in the augmentation graph) into consideration.
Revised: GAT extends GCN by taking the semantic similarities between neighborhood features (corresponding to positive pairs in the augmentation graph) into consideration.

Original: Specifically, for a node x with neighbors
Revised: Specifically, for a node x with neighbors

Original: [Table]

Original: N x , GAT aggregates neighbor features with the attention scores as
Revised: N x , GAT aggregates neighbor features with the attention scores as

Original: [Equation]
Revised: [Equation]

Original: , where α ( x, x+
Revised: α [Equation] x

Original: [Equation]
Revised: [Equation]

Original:  In GAT, they adopt additive attention to compute the so-called attention coefficients e ( x, x ′ ) .
Revised: x [Equation] In GAT, they adopt additive attention to compute the so-called attention coefficients e ( x, x ′ ) .

Original: Alternatively, we can adopt the dot-product attention as in Transformers, e.g., e ( x, x ′ )
Revised: Alternatively, we can adopt the dot-product attention as in Transformers, e.g., e ( x, x ′ )

Original: =
Revised: =

Original: H ( t ) x
Revised: H ( t ) x

Original: ⊤ H ( t ) x ′ .
Revised: ⊤ H ( t ) x ′ .

Original: Motivated by the design of graph attention, we revise the vanilla alignment loss by reweighting positive pairs with their estimated semantic similarities.
Revised: Motivated by the design of graph attention, we revise the vanilla alignment loss by reweighting positive pairs with their estimated semantic similarities.

Original: Following Transformer, we adopt the dotproduct form to estimate the attention coefficient as e ( x, x + ) =
Revised: Following Transformer, we adopt the dotproduct form to estimate the attention coefficient as e ( x, x + ) =

Original: f θ ( x ) ⊤ f θ ( x + ) .
Revised: f θ ( x ) ⊤ f θ ( x + ) .

Original: Indeed, Figure 2c shows that the contrastive encoder can provide decent estimation of semantic similarities between positive samples.
Revised: Indeed, Figure 2c shows that the contrastive encoder can provide decent estimation of semantic similarities between positive samples.

Original: Therefore, we adopt this attention coefficient to adaptively reweight the positive pairs, and obtain the following attention alignment loss (with the uniformity loss unchanged),
Revised: Therefore, we adopt this attention coefficient to adaptively reweight the positive pairs, and obtain the following attention alignment loss (with the uniformity loss unchanged),

Original: [Equation]
Revised: [Equation]

Original: Specifically, we assign the attention score α ( x, x + ) as α( x, x + ) = [Equation] x ′
Revised: Specifically, we assign the attention score α ( x, x + ) as α

Original: [Equation]
Revised: [Equation]

Original: β
Revised: x,

Original: [Equation]
Revised: [Equation]

Original: ) , where we simply adopt M negative samples in the mini-batch for computing the normalization, and β is a hyperparameter to control the degree of reweighting.
Revised: , where we simply adopt M negative samples in the mini-batch for computing the normalization, and β is a hyperparameter to control the degree of reweighting.

Original: We also detach the gradients from the weights as they mainly serve as loss coefficients.
Revised: We also detach the gradients from the weights as they mainly serve as loss coefficients.

Original: The corresponding attention alignment update is(
Revised: The corresponding attention alignment update is

Original: [Equation]
Revised: [Equation]

Original: with [Equation] Compared to the vanilla alignment update (Eq. 5), the attention alignment update adaptively aligns different positive samples according to the feature-level semantic similarity α ( x, x + ) .
Revised: Compared to the vanilla alignment update (Eq. 5), the attention alignment update adaptively aligns different positive samples according to the feature-level semantic similarity α ( x, x + ) .

Original: Thus, our proposed attention alignment can help correct the potential sampling bias of existing data augmentation strategies, and learn more semantic-consistent features.
Revised: Thus, our proposed attention alignment can help correct the potential sampling bias of existing data augmentation strategies, and learn more semantic-consistent features.

Original: Result Analysis.
Revised: Result Analysis.

Original: We follow the basic setup of SimCLR (Chen et al., 2020b) and consider two popular backbones: ResNet-18 and ResNet-50.
Revised: We follow the basic setup of SimCLR (Chen et al., 2020b) and consider two popular backbones: ResNet-18 and ResNet-50.

Original: We consider image classification tasks on CIFAR-10,
Revised: We consider image classification tasks on CIFAR-10,

Original: CIFAR-100 and ImageNet-100.
Revised: CIFAR-100 and ImageNet-100.

Original: We train 100 epochs on ImageNet-100 and 200 epochs on CIFARand CIFAR-100.
Revised: We train 100 epochs on ImageNet-100 and 200 epochs on CIFARand CIFAR-100.

Original: More details can be found in Appendix B.2.
Revised: More details can be found in Appendix C.2.

Original: As shown in Table 1, the proposed attention alignment leads to a consistent 1% increase in linear probing accuracy among all settings, which helps verify that our attentional reweighting of the positive samples indeed lead to consistently better performance compared to the vanilla equally averaged alignment loss.
Revised: As shown in Table 1b, the proposed attention alignment leads to a consistent 1% increase in linear probing accuracy among all settings, which helps verify that our attentional reweighting of the positive samples indeed lead to consistently better performance compared to the vanilla equally averaged alignment loss.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we proposed a new message passing perspective for understanding the learning dynamics of contrastive learning.
Revised: In this paper, we proposed a new message passing perspective for understanding the learning dynamics of contrastive learning.

Original: We showed that the alignment update corresponds to a message passing scheme on the augmentation graph, while the uniformity update is a reversed message passing on the affinity graph, and the two message passing rules strike a balance at the equilibrium where the two graphs agree.
Revised: We showed that the alignment update corresponds to a message passing scheme on the augmentation graph, while the uniformity update is a reversed message passing on the affinity graph, and the two message passing rules strike a balance at the equilibrium where the two graphs agree.

Original: Based on this perspective as a bridge, we further developed a deep connection between contrastive learning and MP-GNNs.
Revised: Based on this perspective as a bridge, we further developed a deep connection between contrastive learning and MP-GNNs.

Original: On the one hand, we established formal equivalences between their message passing rules, and pointed out the inherent connections between some existing techniques developed in each domain.
Revised: On the one hand, we established formal equivalences between their message passing rules, and pointed out the inherent connections between some existing techniques developed in each domain.

Original: On the other hand, borrowing techniques from the MP-GNN literature, we developed two new methods for contrastive learning, and verified their effectiveness on real-world data.
Revised: On the other hand, borrowing techniques from the MP-GNN literature, we developed two new methods for contrastive learning, and verified their effectiveness on real-world data.

Original: Overall, we believe that the discovered connection between contrastive learning and message passing could inspire more studies from both sides to benefit from each other, both theoretically and empirically.
Revised: Overall, we believe that the discovered connection between contrastive learning and message passing could inspire more studies from both sides to benefit from each other, both theoretically and empirically.

Revised: A CKNOWLEDGEMENT

Revised: Yisen Wang and Zhouchen Lin were supported by the National Key R&D Program of China (2022ZD0160304), the major key project of PCL, China (No. PCL2021A12), the NSF China (No. 62006153, 62276004), Open Research Projects of Zhejiang Lab (No. 2022RC0AB05), Huawei

Revised: Technologies Inc., Qualcomm, and Project 2020BD006 supported by PKU-Baidu Fund.

Original:  AP ROOF OF M AIN T HEORETICAL R ESULTS
Revised: Yifei Wang thanks Xiaojun Guo from PKU for valuable feedbacks on an early draft of this work. We thank anonymous reviewers for constructive comments A R ELATED W ORK

Revised: Because our work analyzes contrastive learning in terms of the update rules of the alignment and uniformity losses, it is closely related to Wang & Isola (2020) that firstly propose the alignmentand-uniformity (A+U) perspective of the contrastive learning objective.

Revised: However, there is a key difference between the two works.

Revised: In particular, Wang & Isola (2020) only show that if an encoder with perfect A+U exists , it would be the minimizer of the InfoNCE loss (see Theorem 1 (Wang & Isola, 2020)).

Revised: But what we are actually interested is the opposite : will the minimization of the

Revised: InfoNCE loss guide us to the desired A+U properties?

Revised: This is exactly the focus of our work by studying the learning dynamics of contrastive learning.

Revised: Besides, we also extend the analysis A+U on the pretraining task (Wang & Isola, 2020) to its theoretical effect on downstream generalization as discussed in Propositions 2.

Revised: Comparing to Wang et al.

Revised: , they only assume the existence of perfect alignment, while our analysis shows that the contrastive learning process helps improve feature clustering.

Revised: Another closely related work is HaoChen et al.

Revised: (2021) that develops the augmentation graph framework that we build upon.

Revised: Different from ours, they only characterize the optimal solution of the spectral loss L sp directly through eigen-decomposition (ED), which leads to following drawbacks.

Revised: First, because N can be exponentially large, ED ( O ( N 3 ) complexity) is not directly applicable in practice.

Revised: Instead, our analysis focuses on the scalable first-order optimization methods that are adopted in practice with good performance in limited training steps.

Revised: Second, their theory only applies to the spectral loss and fails to characterize the widely adopted InfoNCE loss, which does not admit a closed-form solution.

Revised: In comparison, our message passing analysis naturally incorporates the

Revised: InfoNCE loss, as shown in Section 2.2.2.

Revised: Besides, Wen & Li (2021) study the contrastive learning process, but only with single-layer networks under toy problems, which is hardly practical.

Revised: Tian et al.

Revised: (2021) study the dynamics of non-contrastive learning, but only focus on the predictor parameters.

Revised: Different from theirs, we adopt a feature space view and analyze the overall evolution of sample features for contrastive learning in practice.

Revised: Many works before like Wang & Liu (2021) also intuitively discuss the behaviors of contrastive learning gradients.

Revised: Instead, we provide a formal and systematic study of the learning dynamics of the samples features by characterizing its effect on the downstream performance.

Revised: In particular, we establish an intriguing connection between contrastive learning and message passing, which is the first time that the two distinctive techniques from different areas are formally bridged together.

Revised: B P ROOF OF M AIN T HEORETICAL R ESULTS

Original: Proposition 1.
Revised: Proposition 1.

Original: When the x -th row of F is defined as F x = √ w x f θ ( x [Equation] , we have:
Revised: When the x -th row of F is defined as F x = √ w x f θ ( x ) , we have:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We first prove the two equations Eq.(2a) and Eq.(2b), which states the relations between F and f ( x ) .[Equation]
Revised: We first prove the two equations Eq.(2a) and Eq.(2b), which states the relations between F and f ( x ) .

Original: For the alignment loss

Original: [Equation]
Revised: [Equation]

Original: we have

Original: [Equation]
Revised: [Equation]

Original: For the uniformity loss [Equation] we have
Revised: For the uniformity loss, we have

Original: [Equation]
Revised: [Equation]

Original: Therefore [Equation] we have
Revised: Therefore, we have

Original: [Equation]
Revised: [Equation]

Original: Proposition 2.
Revised: Proposition 2.

Original: Denote the distance between a feature matrix X and a subspace Ω by d Ω ( X ) :=
Revised: Denote the distance between a feature matrix X and a subspace Ω by d Ω ( X ) :=

Original: inf {∥ X − Y ∥
Revised: inf {∥ X − Y ∥

Original: F | Y ∈ Ω } .
Revised: F | Y ∈ Ω } .

Original: Assume that data augmentations are label-preserving.
Revised: Assume that data augmentations are label-preserving.

Original: Then, for samples in each class k , denoting their adjacency matrix as ¯ A k and their features at t -th step as F ( t ) k , the message passing in Eq.
Revised: Then, for samples in each class k , denoting their adjacency matrix as ¯ A k and their features at t -th step as F ( t ) k , the message passing in Eq.

Original: 4 leads to a smaller or equal distance to an one-dimensional subspace M
Revised: 4 leads to a smaller or equal distance to an one-dimensional subspace M

Original: [Equation]
Revised: [Equation]

Original: [Equation] where λ G k ∈ [0 , 2] denotes the intra-class algebraic connectivity of G (Chung, 1997), and M = { e 1 ⊗ y | y ∈ R c } , and e 1 is the eigenvector corresponding to the largest eigenvalue of ¯ A k .
Revised: 22) where λ G k ∈ [0 , 2] denotes the intra-class algebraic connectivity of G (Chung, 1997), and M = { e 1 ⊗ y | y ∈ R c } , and e 1 is the eigenvector corresponding to the largest eigenvalue of ¯ A k .

Original: Proof.
Revised: Proof.

Original: Under the condition that the data augmentations are label-preserving, there are no edges between inter-class nodes.
Revised: Under the condition that the data augmentations are label-preserving, there are no edges between inter-class nodes.

Original: Thus, the propagation of the alignment update only happens within each class.
Revised: Thus, the propagation of the alignment update only happens within each class.

Original: We denote the propagation matrix as Q k = (1 − 2 α )
Revised: We denote the propagation matrix as Q k = (1 − 2 α )

Original: I +

Revised: I +

Original: α ¯ A k .
Revised: α ¯ A k .

Original: Then we have F( t +1) k
Revised: Then we have F

Original: = Q k F ( t ) k .
Revised: [Equation]

Original: Since ¯ A k is a symmetrically normalized adjacent matrix, we can denote the spectrum of ¯ A k as 1 = λ 1 ≥ · · · ≥ λ n ≥ − 1 .
Revised: Since ¯ A k is a symmetrically normalized adjacent matrix, we can denote the spectrum of ¯ A k as 1 = λ 1 ≥ · · · ≥ λ n ≥ − 1 .

Original: Then the spectrum of Q k is q i = 1 − 2 α + 2 αλ
Revised: Then the spectrum of Q k is q i = 1 − 2 α + 2 αλ

Original: i , i = 1 , 2 , . . . , N , and specifically, q 1 = 1 .
Revised: i , i = 1 , 2 , . . . , N , and specifically, q 1 = 1 .

Original: We denote the corresponding orthonormal basis of Q k as [ e 1 , . . .
Revised: We denote the corresponding orthonormal basis of Q k as [ e 1 , . . .

Original: , e n ] since Q k is a symmetric matrix.
Revised: , e n ] since Q k is a symmetric matrix.

Original: Accordingly, we could decompose F ( t ) k according to the orthonormal basis as
Revised: Accordingly, we could decompose F ( t ) k according to the orthonormal basis as

Original: [Equation]
Revised: [Equation]

Original:  Therefore, the projection of F ( t ) k to the space M = { e 1 ⊗ y | y ∈ R c } gives the first component of the decomposition (Eq.23), i.e. , e 1 ⊗ h 1 , and the residual is the remaining components orthogonal to e 1 , i.e. , (cid:80) ni
Revised: [Equation] Therefore, the projection of F ( t ) k to the space M = { e 1 ⊗ y | y ∈ R c } gives the first component of the decomposition (Eq.23), i.e. , e 1 ⊗ h 1 , and the residual is the remaining components orthogonal to e 1 , i.e. , (cid:80) ni

Original: =2 e i ⊗ h i , with the following distance,
Revised: =2 e i ⊗ h i , with the following distance,

Original: [Equation]
Revised: [Equation]

Original: Similarly, we can also deduce the distance of the updated features.
Revised: Similarly, we can also deduce the distance of the updated features.

Original: Since Qe i
Revised: Since Qe i

Original: =
Revised: =

Original: q i e i , we have
Revised: q i e i , we have

Original: [Equation]
Revised: [Equation]

Original: Thus, its projection to M is also e 1 ⊗ h 1 , and we have
Revised: Thus, its projection to M is also e 1 ⊗ h 1 , and we have

Original: [Equation]
Revised: [Equation]

Original: Combining Eqs.24,26, we arrive at
Revised: Combining Eqs.24,26, we arrive at

Original: [Equation]
Revised: [Equation]

Original: Note that 1 − λ 2 is just the algebraic connectivity of G .
Revised: Note that 1 − λ 2 is just the algebraic connectivity of G .

Original: Therefore, with q 2 = 1 − 2 α + 2 λ 2 , we have
Revised: Therefore, with q 2 = 1 − 2 α + 2 λ 2 , we have

Original: [Equation]
Revised: [Equation]

Original: By taking square root on this inequality

Original: [Equation]
Revised: [Equation]

Original:  we arrive at the final result[Equation]
Revised: By taking square root on this inequality, we arrive at the final result

Original: [Equation]
Revised: [Equation]

Original: which completes the proof
Revised: which completes the proof.

Original: [Equation]

Original: Proposition 3.
Revised: Proposition 3.

Original: The contrastive learning dynamics with the InfoNCE loss saturates when the real and the fake adjacency matrices agree, i [Equation]
Revised: The contrastive learning dynamics with the InfoNCE loss saturates when the real and the fake adjacency matrices agree, i.e., ¯ A

Original: e., ¯ A = ¯ A ′ nce , which means that element-wisely, we have:
Revised: = ¯ A ′ nce , which means that element-wisely, we have:

Original: [Equation]
Revised: [Equation]

Original: That is, the ground-truth data conditional distribution P d ( x + | x ) equals to the Boltzmann distribution
Revised: That is, the ground-truth data conditional distribution P d ( x + | x ) equals to the Boltzmann distribution

Original: P θ ( x + | x ) estimated from the encoder features.
Revised: P θ ( x + | x ) estimated from the encoder features.

Original: Proof.
Revised: Proof.

Original: Notice that the contrastive update takes the form of
Revised: Notice that the contrastive update takes the form of

Original: [Equation]
Revised: [Equation]

Original: Therefore, when ¯ A =
Revised: Therefore, when ¯ A =

Original: ¯ A ′ nce , we have F ( t +1) =
Revised: ¯ A ′ nce , we have F ( t +1) =

Original: F ( t ) , which means that the contrastive learning dynamics with the InfoNCE loss saturates.
Revised: F ( t ) , which means that the contrastive learning dynamics with the InfoNCE loss saturates.

Original: At this time, we have
Revised: At this time, we have

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: which implies the result in the proposition.
Revised: which implies the result in the proposition.

Original: B E XPERIMENTAL D ETAILS B. D ETAILS OF M ULTI - STAGE G RAPH A GGREGATION
Revised: C E XPERIMENTAL D ETAILS C. D ETAILS OF M ULTI - STAGE G RAPH A GGREGATION

Original: We follow the default  settings of SimSiam (Chen &He, 2021), except the predictor.
Revised: For the experiments in Figure 2a, we follow the default designs and settings of SimSiam (Chen &

Original:  We use ResNet50 as the backbone andhave a two-layer projector following that. We pretrain the encoder forepochs and then evaluate linear accuracy on the pretrained encoder.
Revised: He, 2021) on CIFAR-10, except that we remove its predictor to focus on the effect of multi-stage aggregation. We adopt ResNet-50 as the backbone and pretrain for 100 epochs.

Original: During the training process, we use  memory banks to store the feature of latest k epochs.
Revised: During the training process, we use a memory bank to store the features of all samples from the last k epochs.

Original: When aggregating the stored features, we set theweight of each feature as 1 r .
Revised: To obtain the final aggregated feature, we average the features from the last k epochs. Afterwards, we update the memory bank with the feature of the positive sample in the current epoch, and discard old features if they are not from the last k epochs. For the experiments in Table 1a, we also follow SimSiam’s setting and consider image classification tasks on three benchmark datasets: CIFAR-10, CIFAR-100 and ImageNet-100. We train ResNetand ResNet-50 backbones for 200 epochs with the default hyperparameters of SimSiam (Chen &

Original: B. D ETAILS OF G RAPH A TTENTION E XPERIMENTS
Revised: He, 2021). Specifically, we use the SGD optimizer with 256 batch size and 5e-5 weight decay. To implement the multi-stage variant of SimSiam, after 50 warmup epochs, we replace the target feature in SimSiam (which is orginally the representation of a positive sample) with the multi-stage aggeragated features over the last 2 epochs (current epoch and the last epoch). C. D ETAILS OF G RAPH A TTENTION E XPERIMENTS

Original: We use SimCLR (Chen et al., 2020b) as our baseline and ResNet-18,ResNet-50 as the backbones.
Revised: We use SimCLR (Chen et al., 2020b) as our baseline and ResNet-18,ResNet-50 as the backbones.

Original: We follow the backbone with a projection MLP to map the features to the 128-dimensional projection space.
Revised: We follow the backbone with a projection MLP to map the features to the 128-dimensional projection space.

Original: We consider image classification tasks on CIFAR-10, CIFAR-100 and ImageNet-100.
Revised: We consider image classification tasks on CIFAR-10, CIFAR-100 and ImageNet-100.

Original: There are two stages in contrastive learning, i.e., self-supervised pretraining and linear evaluation.
Revised: There are two stages in contrastive learning, i.e., self-supervised pretraining and linear evaluation.

Original: We pretrain the encoder for 100 epochs on ImageNet-100 and for 200 epochs on CIFAR-10 and CIFAR100.
Revised: We pretrain the encoder for 100 epochs on ImageNet-100 and for 200 epochs on CIFAR-10 and CIFAR100.

Original: We use the LARS optimizer with cosine annealed learning rate schedule and 512 batch size.
Revised: We use the LARS optimizer with cosine annealed learning rate schedule and 512 batch size.

Original: After the pretraining process, we train a linear classifier following the frozen backbones using the
Revised: After the pretraining process, we train a linear classifier following the frozen backbones using the

Original: SGD optimizer.
Revised: SGD optimizer.

Original: To obtain a more accurate estimation of semantic similarity, we use the features before the projection layer and warmup the model for 30 epochs.
Revised: To obtain a more accurate estimation of semantic similarity, we use the features before the projection layer and warmup the model for 30 epochs.

Original: C D ISCUSSIONS ON R ELATED W ORK

Original: Our work is closely to Wang & Isola (2020) that propose the alignment-and-uniformity (A+U) perspective on contrastive learning.

Original: Nevertheless, they only show that if an encoder with perfect A+U exists , it would be the minimizer of the InfoNCE loss (see Theorem 1 (Wang & Isola, 2020)).

Original: But what we are actually interested is the opposite : will the minimization of the InfoNCE loss guide us to the desired A+U properties?

Original: This is exactly the focus of our work by studying the learning dynamics of contrastive learning.

Original: Besides, we also extend the analysis A+U on the pretraining task (Wang &

Original: Isola, 2020) to its theoretical effect on downstream generalization as discussed in Propositions 2.

Original: Comparing to Wang et al.

Original: , they only assume the existence of perfect alignment, while our analysis shows that the contrastive learning process helps improve feature clustering.

Original: Another closely related work is HaoChen et al.

Original: (2021) that develops the augmentation graph framework that we build upon.

Original: Different from ours, they only characterize the optimal solution of the spectral loss L sp directly through eigen-decomposition (ED), which leads to following drawbacks.

Original: First, because N can be exponentially large, ED ( O ( N 3 ) complexity) is not directly applicable in practice.

Original: Instead, our analysis focuses on the scalable first-order optimization methods that are adopted in practice with good performance in limited training steps.

Original: Second, their theory only applies to the spectral loss and fails to characterize the widely adopted InfoNCE loss, which does not admit a closed-form solution.

Original: In comparison, our message passing analysis naturally incorporates the

Original: InfoNCE loss, as shown in Section 2.2.2.

Original: Besides, Wen & Li (2021) study the contrastive learning process, but only with single-layer networks under toy problems, which is hardly practical.

Original: Tian et al.

Original: (2021) study the dynamics of non-contrastive learning, but only focus on the predictor parameters.

Original: Different from theirs, we adopt a feature space view and analyze the overall evolution of sample features for contrastive learning in practice.

Original: Many works before like Wang & Liu (2021) also intuitively discuss the behaviors of contrastive learning gradients.

Original: Instead, we provide a formal and systematic study of the learning dynamics of the samples features by characterizing its effect on the downstream performance.

Original: In particular, we establish an intriguing connection between contrastive learning and message passing, which is the first time that the two distinctive techniques from different areas are formally bridged together.

Original: At last, similar to our Proposition 3, Zimmermann et al. (2021) also show that contrastive learning can recover the data distribution from a nonlinear ICA perspective.

Original: The main difference is that they derive this result based on the learning objective, while ours focuses on how the learning dynamics attains this equilibrium.

Original: Besides, to obtain this distribution recovery property, Zimmermann et al.

Original: (2021) rely heavily on the data generation process assumption, while ours does not.

Original: D I NCORPORATING M ULTI - STAGE A GGREGATION TO E XISTING M ETHODS

Original: In Section 3.3.1, we have shown that multi-stage aggeragtion, when used alone, can alleviate feature collapse without asymmetric modules.

Original: Nevertheless, we find that this technique alone is insufficient to attain state-of-the-art results.

Original: In view of this property, in this section, we further explore combining multi-stage aggregation with existing techniques to alleviate feature collapse, such as the predictor module in SimSiam (Chen & He, 2021).

Original: Hopefully, combining the two effective techniques can better avoid feature collapse and learn diverse features.

Original: Setup.

Original: As in Section 3.3.2, we consider image classification tasks on three benchmark datasets:

Original: CIFAR-10, CIFAR-100 and ImageNet-100, and train a ResNet-18 backbone for 200 epochs with the default hyperparameters of SimSiam (Chen & He, 2021).

Original: Specifically, we use the SGD optimizer with 256 batch size and 0.00005 weight decay.

Original: To implement the multi-stage variant of SimSiam, denoted as SimSiam+Multi , we replace the canonical positive feature in SimSiam with a multi-stage aggeragated feature (Eq. 13) after 50 warmup epochs.

Original: Specifically, the aggregated positive feature is computed by averaging the positive features from the current epoch and the last epoch.

Original: [Table]

Original: As shown in Table 2, we find that SimSiam-MultiStage brings cosistent and significant improvements on all three benchmark datasets, which clearly demonstrates the benefits of multi-stage aggregation on improving existing non-contrastive methods by further alleviating their feature collapse issues.