A Message Passing Perspective on Learning Dynamics of Contrastive Learning
Authors: Yifei Wang
Qi Zhang
Tianqi Du
Jiansheng Yang
Zhouchen Lin
Yisen Wang
Venue: ICLR 2023 poster (2023)

Abstract: 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. 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. Specifically, we show that its gradient descent corresponds to a specific message passing scheme on the corresponding augmentation graph. Based on this perspective, we theoretically characterize how contrastive learning gradually learns discriminative features with the alignment update and the uniformity update. Meanwhile, this perspective also establishes an intriguing connection between contrastive learning and Message Passing Graph Neural Networks (MP-GNNs). 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. 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.

Published as a conference paper at ICLR 2023
A MESSAGE PASSING PERSPECTIVE ON LEARNING DY- NAMICS OF CONTRASTIVE LEARNING
Yifei Wang' Qi Zhang?* Tianqi Du Jiansheng Yang' Zhouchen Lin Yisen Wang2.3t
School of Mathematical Sciences, Peking University
National Key Lab of General Artificial Intelligence.
School of Intelligence Science and Technology, Peking University
3 Institute for Artificial Intelligence, Peking University
*Peng Cheng Laboratory
ABSTRACT
In recent years, contrastive learning achieves impressive results on self-supervised visual representation learning, but there still lacks rigorous understanding of its learning dynamics. In this paper, we show that if we cast a contrastive ob- jective equivalently into the feature space, then its learning dynamics admits an interpretable form. Specifically, we show that its gradient descent corresponds to specific message passing scheme on the corresponding augmentation graph. Based on this perspective, we theoretically characterize how contrastive learn- ing gradually learns discriminative features with the alignment update and the uniformity update. Meanwhile, this perspective also establishes an intriguing con- nection between contrastive learning and Message Passing Graph Neural Networks (MP-GNNs) This connection not only provides unified understanding of many techniques independently developed in each community, but also enables us to borrow techniques from MP-GNNs to design such as graph attention, graph rewiring, jumpy knowledge techniques, etc. We believe that our message passing perspective not only provides 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 hattps://gittaaa Message Passing ron www.rtttsiee-eeeeeeeeinn
contrastive learning variants,
INTRODUCTION
Contrastive Learning (CL.) has become arguably the most effective approach to learning visual representations from unlabeled data (Chen et al. 2020b He al., 2020. Chen et al. 2020c Wang 2021a Chen et al. 2020d: 2021 Caron et al. 2021). However, till now, we actually know little about how CL gradually learns meaningful features from unlabeled data. Recently, there has been a burst of interest in the theory of CL. However, despite the remarkable progress that has been made, existing theories of CL
et al.
established for either an arbitrary function f in the function class F (Saunshi et al. [019 Wang et al. 2022) or the optimal f" with minimal contrastive loss (Wang &
Isola 2020; HaoChen et al.) 2021 Wang al. 2022). Instead, theoretical characterization of the learning dynamics is largely overlooked, which is the focus of this work.
Perhaps surprisingly, we find out that the optimization dynamics of contrastive learning corresponds
to specific message passing scheme among different samples. Specifically, based on 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. Therefore, the combined contrastive update is a competition between two message passing rules. Based on this perspective, further show that the equilibrium of contrastive learning can be achieved when the two message rules are balanced, i.e., when the learned distribution Pe matches the ground-truth data distribution Pa. which provides a clear picture for understanding the dynamics of contrastive learning.
"Equal Contribution.
I Corresponding Author: Yisen Wang wyisen.wang
wang @ @ww.ceu.uu.............

---

Published as a conference paper at ICLR 2023
Meanwhile, message passing is 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. One particular example is in graph representation learning. Message Passing Graph Neural Networks (MP-GNNs) are the prevailing designs in modern Graph Neural Networks (GNNs), including numerous variants like GCN (Kipf & Welling| (017). GAT (Veličkovid
al. 2018), and even the Transformers (Vaswani et al.) 2017). There is a vast literature studying
its diffusion dynamics and representation power (Li ct al. 2018 Oono & Suzuki
2020 Wang et al. 2022) Therefore,
2021b Li et al. 2022 Dong al., 2021; Xu al., 2019: 2018 Chen al.
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. 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 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. We summarize our contributions as follows:
Learning Dynamics. We reformulate of the contrastive learning into the feature space and develop new decomposition of the alignment and uniformity loss. 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. This message perspective provides a new understanding of contrastive learning dynamics.
Connecting CL and MP-GNNs. Through the message passing perspective of contrastive learning (CL), we show that we can establish an intriguing connection between CL and MP-GNNs. 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.
New Designs Inspired by MP-GNNs. 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. Empirically, we show that both techniques leads to clear benefits on benchmark datasets. In turn, their empirical successes also help verify the validness of our established connection between CL and MP-GNNs.
2 A MESSAGE PASSING PERSPECTIVE ON CONTRASTIVE LEARNING
In this section, we develop a message passing perspective for understanding the dynamics of con- trastive learning. We begin by reformulating the contrastive loss into the feature space with new decomposition. We then study the update rules derived from the alignment and uniformity losses, and explain their behaviors from message passing perspective. When combined together, also characterize how the two updates strike a balance at the equilibrium states. And we finish this section with proof-of-idea to illustrate the effectiveness of our derive message passing rules.
2.1 BACKGROUND, REFORMULATION, AND DECOMPOSITION
We begin our discussion by introducing the canonical formulation contrastive learning methods in the parameter space, and present their equivalent formulation in the feature space.
Contrastive Learning (CL). Given two positive samples (z, x+) generated by data augmentations, and independently sampled negative sample x', we can learn encoder fo Rdd - R with the wide adopted InfoNCE loss (Oord et al. 2018);
(I) where the form term pulls positive samples (Ir, x) together by encouraging their similarity, and the latter term pushes negative pairs (ær, x') apart. In practice, we typically randomly draw M negative samples to approximate the second term.
Luce = Er,, [x) [a(x + Eg log Eg (exp(fa()) fo(2')),
In contrastive learning, the encoder is parameterized by deep neural networks, making it hardly amenable for formal analysis. Wen & Li (2021) resort to single-layer networks with strong assump- tions on data distribution, but it is far from practice. Instead, in this work, we focus on the dynamics
2

---

Published as a conference paper at ICLR 2023
of the learned data features [z( in the feature space Rm. As over- weer-parmetteiieed neural networks are very expressive and adaptive, we assume the feature matrix F6 to be unconstruined (updated freely ignoring dependences
parameters and sample complexity) as in HaoChen et al. (2021). ) This assumption is also widely adopted in deep learning theory, such as, layer peeled models (Mixon et al., 2022, Zhu et al. (02I). The unconstrained feature assumption enables us to focus on the relationship between augmented samples. As pointed out by Wang al. (2022). the support overlap between augmented data is the key for contrastive learning to generalize to downstream tasks, However, we still do not have clear view of how data augmentations affect feature learning, which is the focus of this work. For a formal exposition, we model data augmentations through the language of augmentation graph.
Following HaoChen et al. 2202 1) we assume a finite sample space X' (can be exponentially large). Given a set of natural data as X' {E I E Rd), we first draw natural example E E (a( and then draw a pair of samples x, xx from a random augmentation of z with distribution A( A(-
We denote X' as the collection of augmented "views", i.e., X' = Use & supp (A(- Regarding the N samples in X' as N nodes, we can construct an augmentation graph G = AA with an adjacency matrix A representing the mutual connectivity. The edge weight Azz' 2 0 between x
and x' is defined as their joint probability Azr' Pa(xr, x²) = Axr)(t (.). The normalized
adjaceney matrix is A D 1/1 AD-1// 112 where D deg(A) (A) denotes the diagonal degree matrix, Le., Daz = W/ = Eg. Azz' The symmetrically normalized Laplacian matrix L = I A is known to be positive semi-definite (Chung 1997). As natural data and data augmentations
uniformly
sampled in practice (Chen et al. 2020b). we assume uniform distribution Wa = 1//N for simplicity.
New Alignment and Uniformity Losses in Feature Space. 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 (Cunif for InfoNCE loss, and L for spectral loss):
Calign (F) = Tr LF) Ee, [f(x)) - f(2+)|1.
(2a)
Lunit (F) = LSE(FET)
= E2 log E.r exp( ((x') - E, |3
(2b)
where LSE(X) = Tr(D log( (deg( D Dxxp(D=122XD-110 12)))) serve as pseudo "norm"] The following proposition establishes the equivalence between original loss and our reformulated one.
Proposition 1. When the x-th row of F is defined as F₂ fo we have:
Lnce (0) = Lalign ( (Fo) + Lumit (Fo).
(3)
Remark Loss Decomposition. Notably, our decomposition of contrastive loss is different from the canonical decomposition proposed by Wang & Isola (2020). Taking the InfoNCE loss as an example, they directly take the two terms of Eq. [ as the alignment and uniformity losses, which, however, leads to two improper implications without feature normalization. First, minimizing their alignment loss Calign (0) = -Er, fo(x)) fe(x+) alone is an ill-posed problem, it approaches -oo when we simply scale the norm of F. In this case, the alignment loss gets smaller while the distance between positive pairs gets larger. Second, minimizing the uniformity loss alone will bridge the distance closer. 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 D. Thus, our decomposition in Eq. 2 seems more properly defined than theirs.
Based on the reformulation in Eq. [ we will study the dynamics of each objective in Section and give unified analysis of their equilibrium in Section 2.3 both from message passing perspective.
2.2
A MESSAGE PASSING PERSPECTIVE ON ALIGNMENT AND UNIFORMITY
2.2.1 1
ALIGNMENT UPDATE IS MESSAGE PASSING ON AUGMENTATION GRAPH
In fact, the reformulated alignment loss Lalign (F) = Tr(FT LF) is widely known as the Laplacian reqularization that has wide applications in various graph-related scenarios (Ng et al., 200T Ando &
Zhang 2006). Its gradient descent with step size a > 0 gives the following update rule:
(global update)
F(++1) = F(t) a V pio Latign (F(!)= = [(1 - -2a) + 2 F().
(4)
'ere log, exp are element-wise operations.
3

---

Published as a conference paper at ICLR 2023
One familiar with graphs can immediately recognize that this update rule is matrix-ffomm message
passing scheme of features F(t) with propagation matrix P P= (1 - 2nx) + 2a For those less familiar with the graph literature, let us take look at the local update of the feature F₂ of each x: (local update)
F((+1) = (1 - Fx) + 2a ExEE,, A x²
(5)
We can see that the alignment loss updates each feature Fz by weighted sum that aggregates features from its neighborhood N2 = {x' E '' > 0}. 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 Ars+) will exchange more information.
Next, let us turn our focus to how this message passing scheme affects downstream classification. As for contrastive learning, because positive samples come from the same class, this alignment loss will help cluster same-class samples together. However, it does not mean that any augmentations can attain a minimal intra-class variance. Below, we formally characterize the effect of data augmentations on the feature clustering of each class k. We measure the degree of clustering by the distance between all intra-class features to one- dimenniniio subspace M.
Proposition 2. Denote the distance between a feature matrix X and a subspace S2 by da(X)= == [I - YYLL || Y E 2.. Assume that data augmentations are label-preserving. Then, for samples in each class k; denoting their adjacency matrix as Ag and their features at I-th step as r(()) the message passing in Eq. Aleads to a smaller or equal distance to an one-dimensional subspace M
dm(F k S |1 - 22dgg=l d ()()))
(6)
where dg. [0.2] denotes the intra-clas, algebraic connectivity of 9 (Chung 1997). and M fer @ sly E } and e1 is the eigenvector corresponding to the largest eigenvalue of As-
An important message of Proposition 2 is that the alignment update does not necessarily bring better clustering, which depends on the algebraic connectivity >g. 1) only non-zero Xg, guarantees the strict decrease of intra-class distance, and 2) larger Àg. indicates an even faster decrease. According to spectral graph theory, Xg, is the second-smallest eigenvalue of the Laplacian L, and its magnitude reflects how connected the graph is. In particular, Aç. > O holds if and only if the intra-class graph is connected (Chung) 1997). Indeed, when the class is not connected, the alignment update can never bridge the disjoint components together. While in contrastive learning, the design of data augmentations determines G₁ as well as λg. Therefore, our message passing analysis quantitatively characterize how augmentations affect feature clustering in the contrastive learning dynamics.
2.2.2 UNIFORMITY UPDATE IS MESSAGE PASSING ON THE AFFINITY GRAPH
Contrary to the alignment loss, the uniformity loss instead penalizes the similarity between every sample pair x, x' E X' to encourage feature uniformity. Perhaps surprisingly, from a feature space perspective, its gradient descent rulc, namely the uniformity update, also corresponds to a message passing scheme with the affinity graph g' with a "fake" adjacency matrix A':
(global update)
p(++1)= = p(t)_ =(())_
aV F(t) Lunit (F(t)) [(1 2n) 20 FF())
where Abym = Despp -A'' + Alexp Dexp and Alexp
1012 (()) D- 1/2), D'ag
deg( A A/xo)) ). When we further apply stop gradient to the target branch (as adopted in BYOL Grill
let al. 2020) and SimSiam (Chen & He (021)). )) we can further simplify it to
where A' = exp We can also interpret the uniformity update as specific message passing scheme on the affinity graph G' (X, .A)) with node set X' and adjacency matrix A'. In comparison to the augmentation graph 9 that defines the edge weights via data augmentations, the affinity graph is instead constructed through the estimated feature similarity between (x, x'). The contrastive learning with two graphs resembles the wake-sleep algorithm in classical unsupervised learning (Hinton et al. 1995).
(global update with stop gradient) F(r-1) 1) = (
=''p' - -'
exp (8)
Further, we notice that the affinity matrix A' is induced by the RBF kernel KR (u, v) exp( U 2l2/22 when the features are normalized, i.e., [[[ = 1 (often adopted in practice (Chen et al. 2020b)). Generally speaking, the RBF kernel is more preferable than the linear kernel as it is more non- Euclideld and non- non-parametric. Besides, because the message passing in the alignment update utilizes A, which is also non- non-neeative and normalized, the InfoNCE uniformity update with A can be more stable and require less hyperparameter tuning to cooperate with the alignment update. 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.
4

---

Published as a conference paper at ICLR 2023
210
-05
-12
1.0 AR
63
531 18
1.0
Hi.rdomry papee bordal 1art
20 piiimdyy cotn
pnlrrtrdio (pdan 1850
00
222
Figure 1: Mimicking contrastive learning with our message passing rules on synthetic data. The data samples D are generated following the isotropic Gaussian distribution with means (-1, 0) and (1, 0) (two classes, blue and yellow), and variance 0.7. We then construct an augmentation graph 9 by drawing edged A, for any data pair satisfying ||x; - æzح12 2 I 0.4. For mimicking contrastive learning, we firstly initialize all features {fr] uniformly in [-1, 1]2, and then perform message passing of these features using G for 1,000 steps with step size 0.1. Best viewed in color.
2.3 THE EQUILIBRIUM BETWEEN THE ALIGNMENT AND UNIFORMITY UPDATE
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):
F(+1)= F()) a pt) Lace (F(t)) (()) a A()
which is a special message passing scheme using the difference between the augmentation (data) and affinity (estimated) adjacency matrices Aaii A A'. Accordingly, we can reformulate the contrastive loss more intuitively as:
((F) = (t(ss(Aaa)FF"))
TM(FFF") + Tr(sg(A))FF")) Callgn (F) + Lunit (F), (10)
where sg(-) means the stop gradient operator.
Equilibrium Distribution. The reformulation above (Eq. 10) also reveals the stationary point of the learning dynamics, that is, Aai = Á A' 0, which implies the model equilibrium as follows. Proposition 3. The contrastive learning dynamics with the InfoNCE loss saturates when the real and the fake adjacency matrices agree, i.e., A = Aloco which means that element-wisef, we have:
exp(fola) fo(x*))
V Vz, x+ X, Pa(x=(() Po(x*12) Exe exp(fe(x) fe(a=)
(11)
That is, the ground-truth data conditional distribution Pa(x equals to the Boltzmann distribution Po (x (x) estimated from the encoder features..
In this way, our analysis reveals that contrastive learning (with InfoNCE loss) implicitly learns probabilistic model Pa |x) (Eq, [1]. and its goal is to approximate the conditional data distribution Pa(x*(22) defined via the augmentation graph G. This explains the effectiveness of contrastive learning on out-of-distribution detection tasks (Winkens et al., (020)) that rely on good density estimation. In practical implementations of InfoNCE, e.g., SimCLR, temperature scalar 7 is often introduced rescale the cosine similarities. With a flexible choice of r, the model distribution Po could better approximate sharp (like one-hot) data distribution Pa. A detailed comparison to related work is included in Appendix
2.4 PRODF-OF-IDEA EXPERIMENT
At last, we verify the derived three message passing rules (alignment, uniformity, contrastive) through synthetic experiment using on synthetic augmentation graph to mimic contrastive learning in the feature space. From Figure [ we can see that the results generally align well with our theory. First, the alignment update (Section 2.2.1 indeed bridges intra-class samples closer, but there is also
5

---

Published as a conference paper at ICLR 2023
risk of feature collapse. Second, the uniformity update (Section 2.2.2) indeed keeps all samples uniformly distributed. 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. 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.
3 CONTRASTIVE LEARNING (CL) AND MESSAGE PASSING GRAPH NEURAL NETWORKS (MP-GNNs): CONNECTIONS, ANALOGIES, AND NEW DESIGNS
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'. In a wider context, message passing is a general technique with broad applications in graph theory, Bayesian inference, graph embedding, etc. As result, our interpretation of contrastive learning from message passing perspective enables us to establish some interesting connections between contrastive learning and these different areas.
A noticeable application of message passing is graph representation learning. In recent years, Message Passing Graph Neural Networks (MP-GNNs) become the de facto paradigm for GNNs. Various techniques have been explored in terms of different kinds of graph structures and learning tasks. By drawing a connection between contrastive learning and GNNs,N, we could leverage these techniques to design better contrastive learning algorithms, as is the goal of this section.
To achieve this, we first establish the basic connections between the basic paradigms in Section 3.1 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. Lastly, show that we could further borrow some new techniques from the MP-GNN literature for designing new contrastive learning methods in Section 3.3
3.1 CONNECTIONS BETWEEN THE TWO PARADIGMS
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.
Graph Convolution and Alignment Update. A graph data with N nodes typically consist of two inputs, the input feature X E RNxd xxd and the (normalized) adjacency matrix A E R=x= For feature matrix (tt) E Raxm at the t-th layer, the canonical graph convolution scheme in GCN (Kipf & Welling 2017) gives (1++1) = AH((), Later, DGC (Wang et al. 2021b) generalizes it to H = =(1- At)l + (t) with a flexible step size At for finite difference, which reduces the canonical form when At = Comparing it to the alignment update rule in Eq.44 we can easily notice their equivalence when we take the augmentation graph A as the input graph for a GCN (or DGC), and adopt a specific step size At 2a. 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). Their differences in names and forms are only superficial, since inherently, they share the same message passing mechanism.
Oversmoothing and Feature Collapse The equivalence above is the key for us to unlock many intriguing connections between the two seemingly different areas. A particular example is the oversmoothing problem, one of the centric topics in MP-GNNs with tons of discussions, to name few. Li et al. (2018), Oono & Suzuki (2020). Shi et al. (2022). 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. Meanwhile, we have also observed the same phenomenon for the alignment update in Figure [ Indeed, it is also well known in contrastive learning that with the alignment loss alone, the features will collapse point, known as complete feature collapse. Hua et al. (202I) 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. In fact, both two feature collapse scenarios with the augmentation graph can be characterized using existing oversmoothing analysis of general graph (Oono & Suzuki| 2020). Therefore, this connection also opens new paths for the theoretical analysis of contrastive learning.
single
Atl-Attention and Uniformity Update. Besides the connection in the alignment update, we also notice interesting connection between the uniformity update rule and the self-attention mechanism.

---

Published as a conference paper at ICLR 2023
Table 1: Comparison of linear probing accuracy (%) of contrastive learning methods and their variants inspired by MP-GNNs, evaluated on three datasets, CIFAR-IO (C10), CIFAR-100 (C100). and ImageNet- 100 (IN100), and two backbone networks: ResNet-18 (RN18) and ResNet-50 (RN50).
(a) SimSiam and its multi-stage variant (Eq.
13
(b) SimCLR and its attention variant (Eq.
15
Model
Pretraining
C10
C100
IN100
Model
Pretraining
C10
C100
IN100
SimSiam
83.8
56.3
68.8
SimCLR
84.5
56.1
62.3
RN18
RN18
+ Multi-stoe -Mulii-stgee
84.8
58,9
70,5
+ Attention
85.4
56,9
63.1
SimSiam 85.9
58.4
70.9
SimCLR 88.2
59.8
66.0
RN50
RN50
Multi-stage
87.0
59.8
72.3
+ Attention
89.4
60.7
66.7
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 al. 2021).
For an input feature (+) R NNxm at the t-th step, the self-attention module gives update:
H (++1) A' (+), where A' = D'xep A'xxp Alexp exp (1 H(T)).
(12)
When comparing it to the uniformity update in Eq. 8 we can see that the two rules are equivalent when a 1. Because the step size is negative, self-attention is actually a gradient ascent step that will maximize the feature uniformity loss (Eq. 2b). 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). Different from previous work (Dong al., 2,02T), 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.
3.2 ANALOGIES IN EXISTING TECHNIQUES
In the discussion above, we establish connections between the two paradigms. Besides, we know that many important techniques and variants have been built upon the basic paradigms in each community. Here, through 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. Below, give examples to illustrate benefits from these analogies.
NodeNorm LayerNorm and lg Normalization. In contrastive learning, since SimCLR (Chen et al. 2020b), a common practice is to apply & normalization before calculating the InfoNCE loss, i.e. f (=) (=)]]- As a result, the alignment loss becomes equivalent to the cosine similarity between the features of positive pairs. Yet, little was known about the real benefits of & normalization. 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 deep GCNs by controlling the feature variance. /Zhou et al. (2021) also show that Lay- enNorm (Ba et al. 2016) with both centering and normalization can also obtain similar performance. These discussions MP-GNNs also help us understand the effecttivnness of le normalization.
PairNorm/ / BatchNorm and Feature Centering. Another common technique in contrastive learning is feature centering, i.e., f(x) - μr, where is the mean of all sample (node) features. In DINO (Cron et al. 2021). ,, it is adopted to prevent feature collapse. Similarly, Hua al. (2021) also show that NarchNorm (lfff & Szegedy 2015) (with feature centering as part) can alleviate feature collapse. On the other side, PairNorm (Zhao & Akoglu 2020) combines feature centering and feature normalization (mentioned above) to alleviate oversmoothing. In particular, they show that PainNorm can preserve the total feature distance of the input (initial features). This, in turn, could be adopted as an explanation for how feature centering works in contrastive learning.
3.3 INSPIRED NEW DESIGNS
Besides the analogies drawn above, we find that there are also many other valuable techniques from the MP-GNN literature that have found their (exact) counterparts in contrastive learning. this section, we explore two new connections from different aspects: multi-stage aggregation for avoiding feature collapse, and graph attention for better feature alignment.
7

---

Published as a conference paper at ICLR 2023
50
Liner Accuraty
Ship
40
0.00
0.8L
-00
Truck
20
Bird
1,00
0.00
0.79
0:5
0.54
055
0.35
0.16
Rand
Aggmiation Epachs
1
2
3
5
6
Airplane
1.00
0.79
0.62
0.61
0.56
045
0.48
o.e
(a) mtlti-stage aggregation
(b) data augmentation examples
(c) estimated attention coefficients
Experiments on CIFAR-10: (a) Performance of multi-stage aggregation (linear accuracy (%)) and aggregation epochs s. (b) natural (Ist column) samples and their augmented views (2-4th columns) using SimCLR augmentations. (c) Estimated scores between the anchors (Ist column) and the augmented views (2-8th columns) using pretrained SimCLR model.
Figure 2:
3.3. NEW CONNECTION I: MULTI-STAGE GRAPH AGGREGATION
As discussed above, the alignment loss alone will lead to feature collapse in contrastive learning. Although existing works like BYOL (Grill al. 2020) and SimSiam (Chen & He 2021 )) can alleviate feature collapse via asymmetric architectural designs, we are still not fully clear how they work. Instead, the connection between feature collapse and oversmoothing (Section 3.1) inspires us to borrow techniques from the MP-GNNs alleviate the feature collapse issue of contrastive learning.
In the MP-GNN literature, a principled solution to oversmoothing is the jump knowledge technique proposed in JK-Net (Xu et al. 2018). Instead of adopting features only from the last step (T) (which might already oversmooth). JK-Net aggregates features from multiple stages, e.g., H1) ...., (I) ( as the final feature for classification. In this way, the knowledge from previous non-over-smoothinn steps could directly jump into the final output, which helps improve its discriminative ability. This idea of multi-stage aggregation is also widely adopted in many other MP-GNNs, such as APPNP (Klicpera al., 2019), GCNII (Chen et al. 2020a). SIGN (Rossi al. 2020), etc.
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. Specifically, we create a memory bank M to store the latest s-epoch features from the same natural sample ž. At the f-th epoch, for
22 (t-r) where = min(s, (). Then we replace the original positive feature fo( with the aggregation of the old features the
generated from I. we have stored old features of it as (f-1) 2 (5-2)
memory bank, where we simply adopt the sum aggregation for simplicity, i.e., - 2 (t-i) Next, we optimize the encoder network fo with the following multi multi-stage alignment loss alone,
Smuiti-allga (0) = Ex (fo(()) 20).
(13)
Afterwards, we will push the newest feature = = (o( to the memory bank, and drop the oldest one if exceeded. In this way, could align the new feature with multiple old features that less collapsed, which, in turn, also prevents the collapse of the updated features.
Result Analysis. From Figure 2a we 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 stage aggregation (better than random guess when s > 2), which was only possible using asymmetric modules, e.g., SimSiam's 's predictor (Chen & He 2021). Inspired by this observation, we further combine multi-stage aggregation with SimSiam to further alleviate feature collapse. As shown in Table Ta] 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. Experiment details can be found in Appendix C.I
3.3.2 NEW CONNECTION II: GRAPH ATTENTION
A noticeable defect of existing contrastive loss is that they assign an equal weight for different positive pairs, As these positive samples are generated by input-agnostic random data augmentations, there
8

---

Published as a conference paper at ICLR 2023
could be large difference in semantics between different positive samples. 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. Therefore, in the
alignment update, not proper to assign equal weights as in vanilla alignment loss (Eq [5. Instead, we should adaptively assign different weights to different pairs based on semantic similarities.
Revisiting MP-GNNs, we find that there are well-established techniques address this issue. A well-known extension of graph convolution is graph attention proposed in GAT (Veličković et al. 2018), and Transformers can be seen special case of GAT applied to a fully connected graph. GAT extends GCN by taking the semantic similarities between neighborhood features (corresponding to positive pairs in the augmentation graph) into consideration. Specifically, for node x with neighbors N2. GAT aggregates neighbor features with the attention scores as
(1+1) Ex+ EEN a(z, xt ! where a(z,x") Er'e, exp(e(z,x"))"
exp(e(ar.r ))
(14)
In GAT, they adopt additive attention to compute the so-called attention coefficients e( 2'), Altemn- tively, we can adopt the dot-product attention as in Transformers, e.g., e(z, (z) = = ! H!)
Motivated by the design of graph attention, we revise the vanilla alignment loss by reweighting positive pairs with their estimated semantic similarities. Following Transformer, we adopt the dot- product form to estimate the attention coefficient as e(z, x²) fo()) fo(x:)) Indeed, Figure 2c shows that the contrastive encoder can provide decent estimation of semantic similarities between positive samples. Therefore, we adopt this attention coefficient to adaptively reweight the positive pairs, and obtain the following atlention alignment loss (with the uniformity loss unchanged),
Cattn-aligh (0) T a(r, x (o) fo fo (x++)1² )11
(15)
2
Specifically, we assign the attention score a(I, x+) as a1x, x+) = fxp()-fo(z) fa(a=))
Σn' expl fela) fo(r")) where we simply adopt M negative samples in the mini-batch for computing the normalization, and B is hyperparameter to control the degree of reweighting. We also detach the gradients from the weights
as they mainly serve as loss coefficients. The corresponding attention aligument update is
N
(local update with attention) Fg - (1-2n)
F2 + 20
Σ Ax
(16)
xr
Compared to the vanilla ali alignment update (Eq. the attention alignment update adaptively aligns different positive samples according to the feature-level semantic similarity or(ar, x ). Thus, our pro- posed attention alignment can help correct the potential sampling bias of existing data augmentation strategies, and learn more simantic-consiseen features.
Result Analysis. We follow the basic setup of SimCLR (Chen et al. 2020b) and consider two popular backbones: ResNet-18 and ResNet-50. We consider image classification tasks on CIFAR-10, CIFAR-100 and ImageNet-100 100. We train 100 epochs ImageNet- 100 and 200 epochs on CIFAR-10 and -IFAR-100. More details can be found in Appendix C.2 As shown in Table [b the proposed
attention alignment leads to a consistent 1 5 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.
4 CONCLUSION
In this paper, we proposed message passing perspective for understanding the learning dynamics of contrastive learning. We showed that the alignment update corresponds to a message passing scheme the augmentation graph, while the uniformity update is reversed message passing on the affinity graph, and the two message passing rules strike a balance at the equilibrium where the two graphs agree. Based on this perspective as a bridge, we further developed deep connection between contrastive learning and MP-GNNs. 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. 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. 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.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
Yisen Wang and Zhouchen Lin were supported by the National Key R&D Program of China (2222D0060334)) the major key project of PCL., China (No. PCL2021A12), the NSF China (No. 62006153, 62276004), Open Research Projects of Zhejiang Lab (No. 2022RC0AB0S), Huawei Technologies Inc., Qualcomm, and Project 2020BD006 supported by PKU-Baidu Fund. Yifei Wang thanks Xiaojun Guo from PKU for valuable feedbacks on an early draft of this work. We thank anonymous reviewers for constructive comments.
REFERENCES
Rie Ando and Tong Zhang. Learning graph with laplacian regularization. In NeurIPS, 2006.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.064500 2016.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In ICCV, 2021.
Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. ICMI., 2020a.
Qi Chen, Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Optimization-induced graph implicit nonlinear diffusion. In ICML., 2022.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020b.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton. Big self-supervised models are strong semi-supervised learners. In NeurIPS, 2020c.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning, In CVPR, 2021.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.0422 2020d.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. arXiv preprint urXiv:2104.02057, 2021.
Fan RK Chung. Spectral graph theory. American Mathematical Society, 1997.
Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. ICML, 2021.
Jean-Bastien Grill, Florian Strub, Florent Altché, C. Tallec, Pierre H. Richemond, Elena Buchatskaya, C. Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, B. Piot, K. Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: new approach to self-supervised learning. In NeurlPS, 2020.
Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. In NeurlPS, 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020.
Geoffrey E Hinton, Peter Dayan, Brendan Frey, and Radford M Neal. The" wake-sleep" algorithm 1995.
for unsupervised neural networks. Science, 268(5214):1158-11
Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao. On feature decorrelation self-supervised learning. In ICCV, 2021.
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In /CML, 2015.
10

---

Published as a conference paper at ICLR 2023
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. ICLR, 2017.
Johannes Klicpera, Aleksandar Bojchevski, and Stephan Günnemann. Predict then propagate: Graph neural networks meet personalized pagerank. In ICLR, 2019.
Mingjie Li, Xiaojun Guo, Yifei Wang, Yisen Wang, and Zhouchen Lin. G en: Graph gaussian convolution networks with concentrated graph filters. In ICML, 2022.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. AAAI, 2018.
Tiayyag Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. arXiv preprint arXiv:2106.045:5 2021.
Dustin Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis, 20(2):1-13, 2022.
Andrew Ng. Michael Jordan, and Yair Weiss. On spectral clustering: Analysis and an algorithm. In NeurlPS, 2001.
Kenta Oono and Taiji Suzuki. Graph neural networks exponentially lose expressive power for node classification. In ICLR, 2020.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03348 2018.
Emanuele Rossi, Fabrizio Frasca, Ben Chamberlain, Davide Eynard, Michael Bronstein, and Federico Monti. SIGN: Scalable inception graph neural networks. arXiv preprint arXiv:2004.11198, 1/ 2020.
Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar, A theoretical analysis of contrastive unsupervised representation learning. In ICML, 2019.
Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen Lee, and James T Kwok. Revisiting over-smoothing in bert from the perspective of graph. 2022.
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding celf-supervised learning dynamics without contrastive pairs. ICML, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. In ICLR, 2018.
Feng Wang and Huaping Liu. Understanding the behaviour of contrastive loss. In CVPR, 2021.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere, 2020.
Yifei Wang, Zhengyang Geng, Feng Jiang, Chuming Li, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Residual relaxation for multi-view representation learning. In NeurIPS, 2021a.
Yifei Wang. Yisen Wang. Jiansheng Yang. and Zhouchen Lin. Dissecting the diffusion process in linear graph convolutional networks. In NewrIPS, 2021b.
Yifei Wang, Qi Zhang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Chaos is ladder: A new theoretical understanding contrastive learning via augmentation overlap. ICLR, 2022.
Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In ICML, 2021.
Jim Winkens, Rudy Bunel, Abhijit Guha Roy, Robert Stanforth, Vivek Natarajan, Joseph R Ledsam. Patricia MacWilliams, Pushmeet Kohli, Alan Karthikesalingam, Simon Kohl, et al. Contrastive training for improved out-o--distribution detection. arXiv preprint arXiv:2007.05566, 2020.
11

---

Published as a conference paper at ICLR 2023
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge networks. In ICML, 2018.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In ICLR, 2019.
Lingxiao Zhao and Leman Akoglu. PairNorm: Tackling oversmoothing in gnns. ICLR, 2020.
Kuangqi Zhou, Yanfei Dong, Kaixin Wang, Wee Sun Lee, Bryan Hooi, Huan Xu, and Jiashi Feng. Understanding and resolving performance degradation in graph convolutional networks. In CIKM, 2021.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features. In NeurIPS, 2021.
12

---

Published as a conference paper at ICLR 2023
A
RELATED WORK
Because our work analyzes contrastive learning in terms of the update rules the alignment and uniformity losses, it is closely related to Wang & Isola (2020) that firstly propose the alignment- and-uniformity (A+U) perspective of the contrastive learning objective. However, there is a key difference between the two works. In particular, Waag & 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)). But what we are actually interested is the opposite: will the minimization of the InfoNCE loss guide us to the desired A+U properties? This exactly the focus of our work by studying the learning dynamics of contrastive learning. Besides, we also extend the analysis A+U on the pretraining task (Wang& & Isola 2020) to its theoretical effect downstream generalization as discussed in Propositions 2 Comparing to Wang et al. (2022), they only assume the existence of perfect alignment, while our analysis shows that the contrastive learning process helps improve feature clustering.
Another closely related work is HaoChen et al. (2021) that develops the augmentation graph frame- work that we build upon. Different from ours, they only characterize the optimal solution of the spectral loss Spp directly through decopposooooiiionnn (ED), which leads to following braws First, because can be exponentially large, ED (O(N*) complexity) is not directly applicable in practice. Instead, adopted in practice with good performance limited training steps. Second, their theory only applics to the spectral loss and fails to characterize the widely adopted InfoNCE loss, which does not admit a closed- form solution. comparison, our message passing analysis naturally incorporates the InfoNCE loss, as shown in Section 2.2.2
analysis focuses the scalable first-order optimization methods that
Besides, Wen Li 2022 study the contrastive learning process, but only with single-layer networks under toy problems, which is hardly practical. TTin et al. (20TT) study the dynamics of non-contrastive learning, but only focus the predictor parameters. Different from theirs, we adopt a feature space view and analyze the overall evolution of sample features for contrastive learning in practice. Many works before like Wang Liu (202)) also intuitively discuss the behaviors of contrastive learning gradients. Instead, we provide a formal and systematic study of the learning dynamics the samples features by characterizing its effect the downstream performance. In particular, establish an intriguing connection between contrastive learning and message passing, which the first time that the two distinctive techniques from different areas are formally bridged together.
B PROOF OF MAIN THEORETICAL RESULTS
Proposition When the x-th row of defined as Fz fo we have:
Lnco Catign + Lunit (Lo)
(17)
Proof. We first prove the two equations 2a) and (q. and f(x).
2a
2bb, which states the relations between F
Latign (F) = Tr(FT LF)
2 Er,tt [l (xx)))
(2a)
Lunit FF) LSE(FFT)
|| = E, log E. exp(f(a) f(x')) E, E=1(()))
(2b)
13

---

Published as a conference paper at ICLR 2023
For the alignment loss, have
Tr(F" LF) 'T(FT (1 A)F)
Σ War Ann' (16) f(x')
[(x)) + Σ wn](a) f(x (xx)
E Aed f(x)) Σ> E Azw (x) f(x)
ΣΣ fax) Σ Σ (Aar (1x) f(x) ArzI(x') f(a))
ΣV
AZx f(x) S(x)) + ΣΣ (Aar's ( (xx) + A f(2)))
2 A -
1
E [1(x)) - f(xt)11
2
(19)
For the uniformity loss, have
LSE(FFT)+ ||111 = Tr(D log deg(Deeeeppppp
1/2 D D-1/211)) - 1F11
Σ w2 log Σ W (x)( ((x')) Σ w
(20)
= Eg log En exp(f(x) f(z')) Ex,l
Therefore, we have
Leco (0)= = -E (5(x)+([(x))))))))))))) [(x))) E, log Ea [exp(] [(x)))
= - E (sxx) +B
(f(x)) ((x)) + Eg log E (ex ((x')) - E. (=)))(
1 Be (z)) - fxx )|( F. log E. (exp([ s(x))) E,
Lalign (Fo) Lunit (F)).
(21)
Proposition 2. Denote the distance between feature matrix X and a subspace Q by dx(X)= inf Y| Y E D.. Assume that data augmentations are label preservvin.. Then, for samples in each class k. denoting their udjacency matrix Ag and their features at 1-th step ((()) the message passing in Eq. Aleads a smaller or equal distance to an one- one-iimensional subspace M
d S 11 d
(22)
where Age E [0,2] denotes the intra-class algebraic connectivity of & (Chung 1997), and M = {er @ yly E RR), and e1 the eigenvector corresponding to the largest eigenvalue of Ag-
Proof. Under the condition that the data augmentations are label-preserving, there are no edges between inter-class nodes. Thus, the propagation of the alignment update only happens within each class. We denote the propagation matrix Qs (1-2a) ( Then we have p((t++1) F( Since A is a symmetrically normalized adjacent matrix, we can denote the spectrum of A as I == 2 ... 2 An 2 -1. Then the spectrum of Q is = 1 - 2a + 2ad,, = 1, 2, .,.... N, and specifically, 4ı 1. We denote the corresponding orthonormal basis of Qx as [e....., en] since Qx a symmetric matrix.
Accordingly, we could decompose E:(!) according to the orthonormal basis as
F (1) = Σ
c & h4 hy "R..
(23)
1=1
14

---

Published as a conference paper at ICLR 2023
Therefore, the projection of (t) the space M fer yly Rc} gives the first component of the decomposition (Eq(23). i.e., 61 @ hy, and the residual is the remaining components orthogonal to ej, i.e., Di=2 Et @ h2s with the following distance,
(
Ees Eeoh,
Σ
14
(24)
Similarly, we can also deduce the distance of the updated features. Since Qe, = q,e,, we have
QF 8 hn
Σ es 6 (1.h4).
(25)
Thus, its projection to M is also 6₁ @ h1, and we have
(m(
Σ @ (gshs)
Σ
11:h/ 122.
(26)
2m
Combining Eqs 24 26
we arrive at
dF = (2((
Σ 1 11
Σ
(27)
Note that 1 - X2 is just the algebraic connectivity of G. Therefore, with Ø2 = - 2a + 2Ag, we have (F((t(tt) S (M) S (P(
(28)
By taking square root on this inequality, we arrive at the final result
dm((ktt)) s 1(1 20 d d
(29)
which completes the proof.
Proposition 3. The contrastive learning dynamics with the InfoNCE loss saturates when the real and the fake adjucency matrices agree, i.e., A Anoce which means that element-wise we have:
exp(folx) fo(x*)
VK, x E Xx Pa(x+(a)) px( Ex (xp( Te(x'))
(30)
That is, the ground-truth data conditional distribution Pa(x+1z) equals to the Boltzmann distribution
Po Po(x (x* (x) estimated from the encoder features.
Proof. Notice that the contrastive update takes the form of
p(++1) p(t) aV F(o Fto Laco (F)(t) p( + 2a(A (()))
(31)
Therefore, when A A' we have F(++)) F(O), which means that the contrastive learning dynamics with the InfoNCE loss saturates. At this time, we have
exp(S(x)T S(x))
A(i,)) = Anco E. (xp(f(x)) f(xx))
(32)
which implies the result in the proposition.
C EXPERIMENTAL DETAILS
C.I DETAILS OF MULTI-STAGE GRAPH AGGREGATION
For the experiments in Figure 2a we follow the default designs and settings of SimSiam (Chen & He 2021) on CIFAR-10, except that we remove its predictor to focus on the effect of multi-staee aggregation. We adopt ResNet-50 the backbone and pretrain for 100 epochs. During the training process, we use a memory bank to store the features of all samples from the last & epochs. To obtain the final aggregated feature, average the features from the last epochs. Afterwards, update
15

---

Published as a conference paper at ICLR 2023
the memory bank with the feature of the positive sample in the current epoch, and discard old features they are not from the last epochs.
For the experiments in Table Ta we also follow SimSiam's setting and consider image classification tasks on three benchmark datasets: CIFAR-10, CIFAR-I00 and ImageNet-100. We train ResNet- 18 and ResNet-50 backbones for 200 epochs with the default hyperparameters of SimSiam (Chen & He 2021). Specifically, we use the SGD optimizer with 256 batch size and 5e-5 weight decay. To implement the multi-stage variant SimSiam, after 50 warmup epochs, we replace the target feature in SimSiam (which is orginally the representation of positive sample) with the multi-stage aggeragated features over the last 2 epochs (current epoch and the last epoch).
C.2 DETAILS OF GRAPH ATTENTION EXPERIMENTS
We SimCLR (Chen al. 2020b)
baseline and ResNet-18 Ress the backbones. follow the backbone with a projection MLP to map the features to the 128-dimensional projection space. We consider image classification tasks CIFAR-10, CIFAR-100 and ImageNet-1000 There are two stages contrastive learning, i.e., self-supervieed pretraining and linear evaluation. We pretrain the encoder for 100 epochs on ImageNet- 100 and for 200 epochs on CIFAR-10 and CIFAR- 100. We use the LARS optimizer with cosine annealed learning rate schedule and 512 batch size. After the pretraining process, we train a linear classifier following the frozen backbones using the SGD optimizer. To obtain more accurate estimation of semantic similarity, we use the features before the projection layer and warmup the model for 30 epochs.
16


Reviews:
[{"review_id": "j_N-AOzyNG", "reviewer": "Reviewer_WHor", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies contrastive learning (CL) by looking the dynamics of (unconstrained) gradient descent for the contrastive loss and connecting it to message passing schemes on the augmentation graph.\n\n- The \"alignment update\" for CL is a message passing scheme on the augmentation graph, because each feature is updated as a weighted sum of neighborhood features in the augmentation graph\n- The \"uniformity update\" is equivalent to message passing on a \"fake affinity graph\" that is itself defined by the similarities induces by the current features.\n- Under this view, the contrastive update is a competition between two message passing rules and this helps characterize the \"equilibrium solution\"\n\nIt further elucidates connections/analogies between various techniques in contrastive learning and message passing methods on graph neural networks (MP-GNNs), like layer/batch normalization, feature collapse etc.\nInspired by these connections, the paper designs new contrastive learning variants that borrow techniques from GNNs like graph attention jumpy knowledge etc. that perform reasonably well on standard benchmarks.\n\n- Using graph attention mechanism for adaptive positives with a weighted alignment term $\\alpha_{f}(x, x^+) f(x)^{\\top} f(x^+)$, to avoid considering far off positives. Leads to ~1% improvement over vanilla evaluation.", "summary_of_the_review": "I think the connection between CL and GNNs could open up new avenues for ideas in contrastive learning. This paper showed one such connection through graph attentions that lead to a small benefit. Overall it would have nicer if more such connections/insights arose. I believe that the overall contribution of the paper is positive. However due to bugs in the derivations and proofs (listed above), that may or may not affect some results, I am hesitant to accept the paper immediately, and have thus assigned a score of weak reject. Happy to raise the score after further clarification from the authors.", "strength_and_weaknesses": "**Strengths**\n\n- The paper establishes many connections between contrastive learning updates and message passing methods that could potentially help both fields from advancements in the other.\n- It is also reasonably easy to follow and well written (comments towards the end).\n- The proposed method that utilizes graph attention mechanism perform slightly better than vanilla CL\n\n**Weaknesses**\n\n- While the paper makes many connections, most are derived from pretty straightforward or known calculations (the connections can still be valuable). If not technical, ideally the main contribution could have been to leverage the connections for new contrastive learning algorithms that do well. The only one that gets good performance (and improvements) is the graph attention idea.\n\n- Using multi-stage aggregation, i.e. using features from multiple epochs rather than just the current one, does avoid complete feature collapse without negatives. However the final performance is still far from CL; I think ~45% on CIFAR-10 is achievable with even a randomly initialized network, so I'm not sure how much to read into those results of ~50%.\n\n- Typos/bugs in derivations and proofs (that may be fixable)\n  - I fail to see why $LSE(FF^{\\top})$ is the same as $\\mathbb{E}\\_{x} \\log(\\mathbb{E}\\_{x'} \\exp(f(x)^{\\top}f(x^{+})))$ in Eq. (2b). It seems like it should be something like $Tr(D \\log(deg(D \\exp(D^{-1/2}FF^{\\top} D^{-1/2}))))$. The proof in Eq (20) seems wrong because the $w$ term vanishes from the $\\exp$ in the next step. Similar the expression for $A'\\_{\\text{exp}}$ below Eq (7) seems wrong and other places that use $FF^{\\top}$ kind of term by ignoring the $D$ factors. I'm not entirely sure if this affects other results (seems like it won't for uniform distribution), but it would be nice to get some clarification from the authors about this\n  - In Eq (8) should the first term be $(1+2\\alpha) I$ instead of $(1+\\alpha)I$? An so should Eq (9) not have the $(1-\\alpha)$ term? If not, then the Equilibrium distribution calculation will be incorrect, since $\\bar{A} = A'$ will give $F^{(t+1)} = (1-\\alpha) F^{(t)}$.\n\n- It might be worth mentioning in Proposition 3 that the equilibrium distribution is true only if $f$ is allowed to be arbitrarily expressive and that it ignores sample complexity considerations. Also, such a result of distribution recovery has been shown in prior work using different techniques, Theorem 1 in [1].\n\n- Other comments/questions\n  - Missing citation for GAT on page 1\n  - \"Empirically, we show that both techniques leads to clear benefits on benchmark datasets\": I would not say that there are clear benefits for both techniques\n\n\n[1] Zimmerman et al. Contrastive Learning Inverts the Data Generating Process", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ZKip7vDUuG", "reviewer": "Reviewer_KQoy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper revisits contrastive learning objectives from the feature propagation perspective. Specifically, by casting alignment and uniformity as two message propagation procedures on two respectively graphs, we can establish equivalence between contrastive learning and message passing graph neural networks. In this way, we can inherit existing techniques in graph neural networks to improve contrastive learning performance.", "summary_of_the_review": "I like the paper very much. I think the authors make a nontrivial contribution in understanding and improving contrastive learning.", "strength_and_weaknesses": "Strengths:\n+ This paper presents a very interesting perspective on understanding contrastive learning by formulating alignment and uniformity as message passing on graphs\n+ This paper is well-written and the idea is easy to follow\n+ The message passing formulation is supported by both theoretical analysis and empirical evidence\n\nWeaknesses:\nI only have several minor points:\n- The second paragraph in the introduction section is a little bit confusing. I would like to suggest the authors to elaborate on \"model distribution\" and \"data distribution\".\n- I wonder how we can understand the MoCo approach that maintains a memory bank for negatives by constructing similar message passing graphs?\n- In Page 1, GAT is not cited correctly.\n- In Page 5, reference (Eq. 30) should be (Eq. 11).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "M9_xLbLLxar", "reviewer": "Reviewer_7ZHM", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper casts a new understanding of contrastive learning, which interprets the alignment update and uniformity update as message-passing steps and bridges contrastive learning and message passing. It interprets some techniques used in contrastive learning from the perspective of message passing neural network and proposes some inspiration for the modification of contrastive learning, some empirical results also justify these changes inspired by MPNN.", "summary_of_the_review": "This paper is well-written and presents an innovative understanding of contrastive learning. It connects contrastive learning with MP-GNN in both theory and the analogy of problems/techniques. It also provides some ideas for new designs in contrastive learning inspired by the existing tools in MP-GNN. I think this is a solid paper and vote for accept.", "strength_and_weaknesses": "Strengths:\n1. This paper studies a significant problem about how data augmentation affects the learning dynamics of contrastive learning, although previous works have made some trials, this paper has a totally new perspective that is more interpretable. \n2. The analogy from MP-GNN is very interesting, such as the oversmoothing problem and feature collapse, this is inspiring for the future study of contrastive learning.\n3. It proposes some new designs which are also justified by the experimental results, the GAT-like attentive score weights is reasonable and promising.\n\nWeaknesses:\n1. The analogies and two new designs mainly focus on the alignment updates, while negative sampling for the uniformity updates is a well-known problem. I think the authors should also pay attention to the negative sampling part such as the hard negative sampling [1].\n\n\n[1] Contrastive Learning with Hard Negative Samples. Robinson et al., ICLR 2021.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster