On the Importance of Contrastive Loss in Multimodal Learning
Authors: Yunwei Ren
Yuanzhi Li
Venue: Submitted to ICLR 2023 (2023)
Keywords: multimodal learning
contrastive learning

Abstract: Recently, contrastive learning approaches (e.g., CLIP (Radford et al., 2021)) have received huge success in multimodal learning, where the model tries to minimize the distance between the representations of diﬀerent views (e.g., image and its caption) of the same data point, while keep the representations of diﬀerent data points away from each other. However, from a theoretical perspective, it is unclear how contrastive learning can learn to align the representations from diﬀerent views eﬃciently, especially in cases where the data is not isotropic. In this work, we analyze the training dynamics of a simple multimodal contrastive learning model, and show that contrastive pairs are important for the model to eﬃciently balance the learned representations. In particular, we reveal a stage-wise behavior of the learning process: In the ﬁrst stage, the model aligns the feature representations using positive pairs and the condition number grows in this stage. Then, in the second stage, the model reduces the condition number of the learned representations using negative pairs.

Under review as a conference paper at ICLR 2023
ON THE IMPORTANCE OF CONTRASTIVE Loss IN MULTIMODAL LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Recently, contrastive learning approaches (e.g.. CLIP (Radford et al., 2021)) have received huge success multimodal learning, where the model tries to minimize the distance between the representations of different views (e.g., image and its caption) of the same data point, while keeping the representations of different data points away from each other. However, from theoretical perspective, it is unclear how contrastive learning can learn to align the representations from different views efficiently, especially in cases where the data this work, we analyze the training dynamics of simple multimodal contrastive learning model and show that contrastive pairs are important for the model efficiently balance the learned representations. In particular, we reveal stage- wise behavior of the learning process: the first stage, the model aligns the feature representations using positive pairs and the condition number grows in this stage. Then, in the second stage, the model reduces the condition number of the learned representations using negative pairs.
isotropic. In
INTRODUCTION
One of the exceptional abilities of humans is to associate data from different modalities (such as texts and images) together. For example, when we hear the words "white dog". we can immediately align it with the image of dog with white color. When we hear the loud sound of the engine, we can imagine an expensive sports car passing nearby.
Recently, in machine learning, multimodal learning methods training the model to align the data from different modules, has become an increasingly popular research direction, especially in deep learning (He & Peng (2017); Stroud et al. (2020); Radford et al. (2021); Ramesh et al. (2021); Xu et al. (2021); Jia et al. (2021); Wang et al. (2022b)). Among them, the recent work CLIP (Radford et al. (2021)) shows remarkable quality results on aligning the features of text and images. The contrastive learning based method CLIP empirically outperforms many existing non-contrastive approaches (Grill et al. (2020); Chen & He (2021); He al. (2020)). The major difference between the contrastive approach and other approaches is that contrastive loss not only requires the learned representations from the same pair of data (i.e. positive pairs) to be positively aligned, but also requires the data from different pairs (i.c. negative pairs) to be as negatively aligned as possible. In the paper, the authors also identify contrastive loss the most critical part to the success of CLIP.
Despite the empirical success of this contrastive learning based method, from theoretical perspec- tive, the most fundamental questions are still largely open: In particular, how do contrastive pairs help in this new multimodal learning approach? How can the non-convex contrastive loss be effi- ciently minimized to learn features from both modules?
Unlike the prior theoretical works on contrastive learning which mostly focus on extracting features from one module (e.g., Arora et al. (2019); Jing et al. (2022); Pokle et al. (2022); Tian et al. (2021); Wen & Li (2021)), one main technical challenge of analyzing contrastive learning in multimodal setting is how the model can be trained to align the feature representations fA. fB from modules A and B respectively. Due to the existence of negative pairs that emphasize negative correlations of fA and fR, it is unclear that the model still has incentives to align the features from different modules.

---

Under review as a conference paper at ICLR 2023
In this paper, we make preliminary theoretical steps towards answering the fundamental theoretical questions of the importance of contrastive loss in multimodal learning. We assume the data from the two modules are of form A = A + AgE A and æ BzB Bg respectively, where AA, 2B are considered as the hidden signals, A, B linear transformations from the signal to the observation and AgC B&EB the noises. Similar linear models have also been used in previous works (Tian et al. (2021); Wen & Li (2021)) in the context of single-modal learning (A = B). The positive pair of the data shares the same signal AA = =B. and has different noises ÉA. and transformations A, B. The goal is to learn features fa, fs that align positive pairs while keeping representations of negative pairs away from each other.
Under this setting, we make the following contributions:
1. We consider the challenging (but more practical) setting where the features in A and B are inhomogeneous, that is, the condition number of A and B can be ((1.. Prior works (Jing et al. (2022); Tian et al. (2021); Wen & Li (2021)) only consider cases where A and B are exactly column orthonormal matrices even in the simpler single-modal setting (A = B).
2. We consider feature learners fA. fo with normalization, meaning that fA. fB are always normalized to have (expected) norm one during training. Output normalization plays critical role in the practical success of contrastive learning and is also employed in CLIP, but it is rarely considered in theory due the additional complexity of the division by norm
3. We analyze the learning process of stochastic gradient descent from random initialization. We prove that contrastive learning converges efficiently to nearly optimal solution, which indeed aligns the feature representation FA and fB.
4. We also demonstrate the importance of negative pairs by comparing with training only over the positive pairs: We prove that although the latter can also learn to align fA and fB, the features learned by contrastive learning with negative pairs is much more uniform, meaning that fA. fB can recover all the singular vectors of A and B and normalize them. On the other hand, without negative pairs, the learned representation is close to a rank one solution, meaning that fA, fa will only focus on the top singular direction of and B.
5. We also perform simulations and more practical experiments further support our theory.
2 RELATED WORKS
Multimodal learning Despite the empirical success of multimodal learning. there are very few theoretical results on this topic. The one most related to ours is Huang al. (2021), in which the authors show that, in certain cases, multimodal methods can provably perform better than single- modal models. However, the authors consider neither contrastive pairs nor the training dynamics.
Contrastive/Non-cootratii learning theory Another much richer line of research is about con- trastive and non-contrastive methods in the context of single-modal self-supervised learning. Start- ing from Arora et al. (2019), many recent works have provided various explanations on why the representations learned with contrastive learning are useful in downstream tasks (Chuang et al. (2020); Tosh et al. (2021); Nozawa & Sato (2021); Wang et al. (2022a); HaoChen al. (2021); Lee et al. (2021); Wang & Isola (2020)). These works mostly focus on the generalization aspect of the problem and do not consider training. Among them, Wang & Isola (2020) also study the problem using the notions of alignment and uniformity, and demonstrate that balanced representations bene- fit downstream tasks. However, they do not provide guarantees on training. Another related line of research is about non-contrastive learning. where the necessity of negative examples is questioned. In this line of research, the optimization problem does get considered as non-contrastive losses have trivial collapsed solutions. Tian et al. (2021) show that, under certain conditions, non-contrastive learning methods can learn non-collapsed solutions. Jing et al. (2022) show that, even with negative examples, contrastive learning can still suffer from another type of collapse, where the learned rep- resentations only span a now-dimensional subspace of the embedding space. In Pokle et al. (2022), the authors show that non-contrastive losses have many nollcollapsed bad minima that the training algorithm does not avoid. Another related work that takes optimization into consideration is Wen & Li (2021), in which the authors analyze the training dynamics of contrastive learning and show that,
2

---

Under review as a conference paper at ICLR 2023
with random augmentation, neural networks can learn features that suppressed by noises when no augmentations are used. Though these works do consider the optimization problem, they focus on the case where the features are uniform, and only Wen & Li (2021) considers output normalization. We compare our results with the most relevant works in the next paragraph.
Comparison with Jing et al. (2022); Tian et al. (2021); Wen & Li (2021) The dimensional collapse problem reported in Jing et al. (2022) is not real issue in our setting, since, in our case, the best the model can do is to recover the latent vector 2, up to some rotation. As result, it is natural low-dimensional subspace of the embedding space
for the learned representations to span only RR" Here, the point of choosing m > d is to make the optimization dynamics more regular, which is common strategy in the study of over-pare wver-paramettrizeed neural networks. The main difference between our work and the analysis in Tian et al. (2021) and Wen & Li (2021) is we do not assume the inputs are isotropic. In our setting, the condition number can be large (log d). When the condition number is 1, one can imagine that thanks to the symmetry, all directions will be learned simultaneously, and therefore, we do not need negative examples to prevent collapse (Tian et al. (2021)) or the negative examples do not play an important role in analysis (Wen & Li (2021)). On the other hand, when the condition number is larger than 1, we do need to use the negative examples to shrink the condition number, which corresponds to the second stage of our analysis.
3 PROBLEM SETUP
Similar to previous theoretical works (e.g., Lee al. (2021); Wen & Li (2021)), we consider linear data-generating model. Formally, we assume that the contrastive pairs (x T are constructed as
x Az+ + A = + B
(I)
where ==, E. EB are independent random variables following the uniforms distributions over (#1/vd)" /Ll/vvddrr and (=1/V respectively and A, B E Rdxr Ag, Be E Rdx(d-r) are matrices with AT A = B = diag(o²) and A A = B B = of for some a E RC and og E I... In words, we first sample the latent vectors == E R" independently and encode them with A, B to form the signal part of the input. Then, we sample the latent noises &A Rdd"", and encode them with Ag. Bgg to form the noise part of the input. Finally, we add the signal and noise parts together to obtain (xr 1. 2 ))). To generate a positive pair (æ A x ft). we use the same latent vector. That is, A = A + AgE and o = B² + B& Note that the latent noises here are still independent. We use a2 max and o² oiin to denote the maximum and minimum of of..... of oz.
respectively, We assume that (o /0 max [1, (d r)o (romi }
< c clog d for some small constant e > Our results can easily be generalized to settings where the dimensions of x and BB are not the same, since one can simply pad zeros at the end of each column of A and B.
One way of interpreting this model is to view each coordinate of the latent vector z as an indicator for the presence/absence of certain object, and the corresponding columns in A and B as the visual and word embeddings of this object, respectively.
Now, we describe our learner model. We consider (normalized) linear feature learners. Define
WIRA
W
WHE
BBB
fa(MA)
fB(øB)
V Bea a
VEn
AA
where WA. WB Bd are the trainable parameters. In words, first map the inputs (æ A. FB) into the embedding space RM using WA and WR, and then apply batch normalization to the out- puts. By saying the learned representations are aligned, we mean that fA and fn are close for positive pairs and far away from each other for negative pairs. Meanwhile, we say the learned representations are balanced if changing a small fraction of coordinates of z does not change the representation dramatically. See Section 4 for formal definitions.
One can easily verify that, in the population case, we have E
W X x 1| |[ A ||: 2/d+
||WT Asl /d. For notational simplicity, we write KA WTA, K WJB, KA
3

---

Under review as a conference paper at ICLR 2023
W and KB.-- Wr. These are the matrices that directly map latent vectors to their final representations. We also define N2 Klll + | ///.///)/d and NB KBBBl + K (BB.llr//dd Then our model can be rewritten as'
fa(wA) KAZA + K KASEA NA
K KgzB + K BBSEB NB
fn(an)
(2)
We initialize each entry of WA and Wn using iid Gaussian (0, 1/m). This scaling ensures the norm of outputs before initialization does blow up m oc. We train our model using gradient descent over the following contrastive loss L. First, we define?
exp(Ti fT) exp(TifÄ fh) + K E exp(Ti fn)
(
expl exp ft) exp(Ti fh) + K E A exp(Ti fF) (B x3)
where K is positive constant controlling the strength of negative samples and E (0, 1] is the inverse temperature. Then, we define the contrastive loss as
C CA CB == E log A x #h) log SB( 2/))
(3)
By the non-contrastive loss, we mean
- Ef
(4)
Formally, the training algorithm defined as follows.
Algorithm 3.1 (Training algorithm). Let L be L in the contrastive case and L in the non-contrastive
case. At each step, we first sample a batch of positive/negative pairs {(x) x n. x x ) use them to compute the empirical version of and update the weight matrices using WA <- WA 2 V and WB + WB T - V
In the non-contrastive case, always Ti 13, and we repeat the above update until gradient descent converges to an approximate stationary point. In the contrustive case, we first use a small Tt 1 (oly(d), the process for Ti =poly(d) iterations, switch to Tr 1, and run the process for another T2 Ti = -poly(d) iterations.
4 MAIN RESULTS
In words, our results say that though both contrastive and non-contrastive methods can align the representations, the representations learned by contrastive methods are more balanced. First, we need to define what do "alignment" and "balance" mean here. We still use SA and fn to denote the embeddings, but our definitions here will architecture-anootiic.. After some general discussion, also discuss these definitions in specific setting.
Definition 4.1 (Alignment). We define the alignment score
T'Lignn = E {LilsA - -fill | - #
f < < - ff +1llff +
Namely, Aiiggn the accuracy of classifying whether the input pair (x BB) is positive pair. We say that the learned representations are aligned T & 1.
Note that the notion of alignment introduced here is stronger than matching the positive pairs, which can be achieved by simply mapping all inputs to one single embedding. In that case, Faiiga will be 0 (or 0.5 if we choose to break ties randomly instead of using strict inequality).
Seee Section D for discussions the sample complexity.
WW use f as shorthand for fA (æ ( similarly for other combinations of A, B and +.
Note that, in the non-contrastive case, changing Tr only changes the learning
4

---

Under review as a conference paper at ICLR 2023
Definition 4.2 (Balance). We define the balance score
T Balance == وΣ 2
where Σ (fafā] a
Namely, Balance is the stable rank of the covariance matrix of the output embeddings. We say that the learned representations are balanced if Dalance 2 or for some a s
We make some short remarks on Definition 4.2. First, only use fA in this definition because the embeddings are well-aligned, fAÍÁ and f f should be approximately the same. Second, the stable rank is usually used small singular values (Rudelson & Vershynin (2007)). To see why this makes sense, note that | | Σx= / mm fik where Ki..... Km are the singular values of y,. If all nonzero K& are the same, then this recovers the actual rank. The intuition behind the use of rank is that the more independent latent variables the model learns, the higher the rank of the representations needs to be.
alternative to the actual rank because it is less sensitive to
In Jing et al. (2022), the authors also use rank to measure the degree of "dimensional collapse". Unlike their argument, here we only require I to be at least or, instead of ar,, for the representations to be called balanced because even if we can recover the underlying latent vectors, the rank is still at most r. Hence, it does not make much sense to expect the learned representations to span the entire embedding space. Finally, note that sufficient condition for I' Dalance 2 or is that the ratio the largest and r-largest singular values is at most va. In other words, after excluding those singular values that should be 0, the condition number is approximately 1.
Why balance representations are important? In our setting, one simple example of aligned but unbalanced representations is W A A and W diag(v) B -1 with 2) 1 and exp(-d) This model maps inputs whose latent vector is z to diag( up to
1/2 = some normalization, for both modules, whence has Aaiggn 1. Meanwhile, one can verify that this model has I BaBaancee/r and is sensitive to small changes in 21. This model can made balanced by replacing diag(u) with I, in which case the model directly recovers the latent vector As result, all changes in z will be reflected the final representation in a faithful way. Similar notions have also been studied in Wang & Isola (2020) under the name "uniforrity" in the context of self-supervised learning, and they also report that balanced representations lead to better performance in downstream tasks, though, unlike our result, they do not provide guarantees training. Still, this further suggests that learning balanced representations is reasonable and important goal.
2/r 2 O. The problem of this model is that it overly emphasizes 21
With these two definitions, we can now state our main results.
Theorem 4.3. Suppose that the network width m poly(d) is sufficiently large, the learning rate 1 = 1/ 3/poly(d) is sufficiently small, and we generate sufficiently, but still polynomially, many samples at each step to compute the loss". Suppose that, for some small constant e 0,
{ (d- r)o
0a max max
S clogd. d.
r mimn
min
(a) Non-contrastive loss. There exists E R." satisfying the above assumptions such that, after poly(d) iterations, Algorithm 3.1 will converge to an approximate stationary point, at which the learned representations are aligned but not balanced, that is, T Falign Ro but I Ppalance/T / O.
(b) Contrastive loss. There exists TO = 1/ and Ti = poly(d) such that, for any valid o, Algorithm 3.1 will reach a point after poly(d) iterations at which the learned representations are aligned and balanced, that is, Taiign & and T Belance/r ss 1.
We close this section with sufficient conditions for the learned representations to be aligned and balanced in our setting. First, in our specific setting, the most natural way for model to achieve Tilign & is to learn ||fa fBll = Ko(zA =B)I| for some Ko E mR with K, K, > 0.
"To make the proof cleaner, we write it in terms of gradient flow over population loss. See Section D for discussions on the discretization of gradient flow.
5

---

Under review as a conference paper at ICLR 2023
In this case, the distance between positive pairs is always close to 0 while the distance between negative pairs is positive. Recall that the output of our model is fA = AKA + KA.fAA)/AA and fn BK / Hence, have the following sufficient condition.
Lemma 4.4 (Sufficent condition for aligned representations). K > || A | >
| A B, and the condition number of a is upper bounded by poly(d), then the learned representations are aligned.
The first two conditions imply the signal parts dominate the outputted representations, the third condition gives the existence of Ko, and the final condition makes sure the smallest singular value Kj Ko of is still at least 1/ poly(d) after normalization.
Now consider the balance of the learned representations. Suppose that our representations already aligned, in the sense of Lemma 4.4. Then we have
= E KAZz KA
KAKA A
ΣS - = (SAAA] FA
KAll//d
|KAI
Recall the relationship between Fualance and the effective condition number of Σg. We have the following sufficient condition.
Lemma 4.5 (Sufficient condition for balanced representations). Suppose that our model is aligned, in the sense of Lemma 4.4. If the condition number of K A is close 1, then the learned representations are balanced.
Note that, since the columns of A orthonormal, even initialization, the condition number of K 7 KA is not close to 1. In other words, condition-numberrtee number-reducing stage is necessary for the model learn balanced representations.
TRAINING DYNAMICS AND PROOF OUTLINE
Our proof is based characterizing the dynamics of gradient descent the contrastive loss. We first choose a small inverse temperature 7² = 1/ poly(d) and run gradient descent for poly(d) iterations. This stage is called Stage 1. Then, we set 7ス poly(d) iterations. This stage is called Stage 2. The use of different тı is mainly for technical purposes since this gives cleaner separation between stages. We observe similar stage-wise behavior when using a uniform +?. See Figure for simulation results. We also report results for the non- contrastive loss there.
and run gradient descent for another
5.1 THE TRAINING DYNAMICS
Instead of tracking the parameters W and WB directly, will track KA, KA,E KB, KB.g, the matrices that directly map latent signals and noises to the final representations. For the case with contrastive pairs, can show that the dynamics of KA governed by the following equation
KA A
E (2-
T--SA(F a (p xb))
f NA
(
ft
KA
diag(o²)
N
d
2 A 1
SA(E) A x exp(T? fA fB)
f (ft
KA N
K
E
diag(o"
exp(T?2 ft f1²)
NA
Sp(xt
x B exp (T: fā fit)
f T NA
KA
K
E
(fx
diag
exp ft)
N
(5)
See Lemma A.3 for the calculation. The first term comes from the positive pairs and the second and third terms from the negative pairs. Within each term, the second part comes from the normalization. The equations for the other K- --matrices can be derived similarly. We rescale the gradients by 1/7? so that d K does not shrink with 77. For the non con-contrastive case, the equation is di
KA
KA
E
(ft A
diag(o
(6)
NA
N2 d
x

---

Under review as a conference paper at ICLR 2023
cccuracy (ion contrasbus)
SccurDcY (contrsttne T-12
accuracy (oontrattie. To
1.0
0G
02
0.5
120
22
122
0
20
60
120
berattoe
largend singule valuan Inanecnttred
lerged in singuler valles frantteciion.
larged wingulee veltees lconitaative.
7-
1-
30 40
120 22 60 30 120
22
60
300
8000015 50
Instlenstx SC
0e02ta15 24 SC
Figure 1: Simulation results. The first row reports the accuracies of different approaches on the classifying positive/negative pairs problem, and the second row reports the values of the largest singular values. From left to right, the columns correspond to the non-contrastive method, con- trastive method with T2 = throughout the entire process, and contrastive method with T; switch from 0.01 to I at the vertical dashed line. One can make several observations here. (a) All meth- ods can quickly attain near 100% accuracy. (b) Only contrastive methods will reduce the condition number to approximately 1. (c) Even when тi 1, we still have the stage-wise behavior, where the models first align the representations in Stage I, and then balance the representations in Stage 2.
Note that the RHS resembles the first term of (5). This is not a coincidence. We will establish the approximate equivalence between the non-contrastive approach and the contrastive approach with small inverse temperature T² (cf. Section 5.3 and Lemma B.9).
5.2 THE INFINITE- WIDTH DYNAMICS
The overall proof strategy is first characterize the dynamics of the infinite-width limit, which is much simpler compared to (5), and then control the error introduced by discretizing the infinite- width network using polynomially many neurons. This discretization is of the main technical challenges of the proof. In general, in order to track the infinite-width dynamics, an exponentially large network is needed (Mei et al. (2018)).
The basic idea is to use Taylor expansion around the infinite-width trajectory to factor out the first- order error terms and show that, either they drive the process towards the infinite-wide width trajectory or the eтror growth introduced by them is slower than the convergence rate.
Here, for ease of presentation, we will focus on the noiseless infinite-width dynamics and, in par- ticular, the evolution of the condition number. Recall that we use iid Gaussian to initialize the entries of WA and WB. Hence, orthogonal each other, at least at initialization. Moreover, one can verify that thanks to the sym- metry, this holds throughout the entire training procedure. Meanwhile, by symmetry, the corre- sponding quantities in modules A and B are always the same. Namely, when KJ KA = K = diag(k?) and K A KB = diag(k²) for some K, A E RT. As a result, in order to characterize the dynamics of K KB. it suffices to look at x²2 and R² One can show that, in this noiseless infinite-width limit, we have (cf. Lemma A.7)
the m - co limit, different columns of KA and KB B are
- we have
dt k = (1
A
A
P
02
(1 0-441115 (1-
R2
p
T
T
02
||<c||2
|
AC
2
||A||
3
|
2
2 = (1 115) di
|
2
R2
P
of (1 (1-3)
T
R-2
T
02
||>|||
|x|]
2
||<:]]
||A:||
||>=||?
7

---

Under review as a conference paper at ICLR 2023
[[<]]²). and T Ex-1 (R2/ ||<:|| yTp- The first term comes from the first term of (5) and the second term from the second and third terms of (5). When nig ~ Ri for all E [r], the model is aligned. When all n3
where S is (1) quantity depending on R and R, Tp tanh(r]
are roughly the same, the model is balanced.
5.3 STAGE 1
In this subsection, we describe the dynamics of gradient descent in Stage 1 and explain how we control the growth of the errors and condition number. For ease of presentation, we will mostly use the infth dynamics (7) instead of the finite-width one (5).
Equivalence of Stage 1 and non-contrastive methods Recall that we use a small 7? in Stage 1, T3 R O for all E [r], whence T is also approximately O. As a result, the second terms of (7) are approximately 0. In other words, only the positive pairs matter. Meanwhile, one can check that, in the anfinite-width limit, (6) corresponds to the first term of (7), up to some multiplicative factor. This gives the equivalence of the dynamics of Stage 1 and the non-contrastive method. See Lemma B.9 for more formal proof in the finite-width setting.
Now, we consider the contrastive loss. The main result of Stage 1 as follows.
Lemma 5.1 (Informal version of Lemma B.1), Under the assumptions of Theorem 4.3, the finite- width dynamics closely track the indinitewwidth ones throughout Stage I, which takes most poly iterations, and, at the end of Stage 1, we have, for any p. E [r] and s E =[d- -d,,
Rgll K2p Rd 1,
[||Ka. /Kg = O, Rg/ng S (Va).
(8)
Moreover, there exists a E R' such that maxp.q n
n2./n2 n2g = 2(va at the end of Stage 1.
In words, (8) says, at the end of Stage 1, KA KB in the relative sense, the noise-signal ratio is small, and the condition number is bounded by O(vd). By Lemma 4.4, the first two conditions of (8) imply that the learned representations are aligned. The proof of this lemma can be found in Section B. Basically, we couple the convergence of k2/ /N/ and the noise-signal ratio with the growth of discretization error and condition number. The main tool we use is the following nonlinear version of Gronwall's 's lemma.
Lemma 5.2. Let A₁ be a positive process. Let X, and Y, be defined as X, S A4 Y S 3A,X,Y, with Xo, Yo, B being positive. Then, for any T 2 O. we have Yr S Yo exp(
Here, X, represents the progress we have made and Yú the error we wish to control. In our case, X, is the maximum between 1 - f nG and the noise-signal ratio, and Y, the discretization error and condition number. This lemma says that, if the error growth rate decreases as we make more progress, then by coupling these two processes, we can make sure the error does not blow up.
5.4 STAGE 2
In Stage 2, rf no longer (1), and now the second term of (7) comes into play. We show that, in this stage, the model will reduce the condition number of K to approximately 1. By Lemma 4.5, this implies that the learned representations are balanced. Formally, we have the following lemma. The proof can be found in Section C.
Lemma 5.3 (Informal version of Lemma C.I). Under the assumptions of Theorem 4.3, the finite- width dynamics still closely track the infinite-widhh one throughout Stage 2. which again takes at most poly(d) iterations, and, at the end of Stage 2, we have, for any p.q q E [r] and s E [d -d-10 r).
||[K. /ng R0 0, 2///ng 1.
The way to control /2/k2? and the noise-signal ratio is similar to Stage 1. For the condition number, note that when 7/ = and & v K, the first term of (7) becomes close to 0 while the second term becomes -4(1 5) / ||<|| (T, -) Since Tр tanh(> //// [ with R3/ ||+ and T is a weight average of these Tp's, 's, this term will push R3/ ||<]]] towards their average and, consequently, reduces the condition number. To obtain a satisfactory convergence rate, suffices to consider the ratio n²? R directly. See Appendix C.5 for details. The discretization error is handled in Appendix C.2.
is positively correlated
8

---

Under review as a conference paper at ICLR 2023
EXPERIMENTAL RESULTS
11.0
10.9
10.5
0.56
10.0
200
The 25
(nommaizee)
#pach-3
0.55
spoch=10
eccuracy
0.0
spoch=200
0.50
csoch-200
0.5
epech 150
0.40
0.32
0.2
0.30
0.0
300
300
300
450
550
15
spach
Figure 2: Results of the MSCOCO experiments. The top row plots report the training loss, alignment scores, and balance scores during training, respectively. The bottom row figures plot the downstream accuracies and the largest 25 singular value of y at different epochs, normalized so that the largest one has value 1. One can see that the alignment score quickly reaches near 100%, and the balance score, as well as the downstream accuracy, increases gradually during training, which matches our theoretical analysis.
Besides the simulation results reported in Figure 1, we also conduct experiments on the MSCOCO- 2014 dataset (Lin al., 2014) using practical models. See Figure 2 for the results. For the text part, we use a pre-trained RoBERTa (Liu et al., 2019), followed by a 3-layer fully-connected network with batch norm between layers. For the image part, we use pre-trained ResNet101 (He et al., 2015), followed by the same layers. In both parts, the width of the fully-connected layers and the output dimension are 768. We freeze the pre-trained parts of the model and only train the cully-connected parts.
We measure the quality of the learned representation using its zero-shot performance on the MSCOCO-2014 validation set. Unlike common image classification datasets, images in the MSCOCO dataset usually have multiple labels, each corresponding to an object that appears in the image, and there are 80 categories in total. We regard prediction to be correct if it matches one label. The zero-shot classification is done compute the image embedding and the embeddings of prompts "This is [LABEL.AMME]" and use the prompt with the highest correlation with the image embedding as the prediction.
same way as Radford et al. (2021), Namely, we
7 CONCLUSION AND DISCUSSION
In this work, we study the role of contrastive pairs in multimodal learning, and show that contrastive pairs are important for the model learn representations that are both aligned and balanced. Our work extends previous results in several directions: First, we consider the more complicated mul- timodal learning problem. Meanwhile, our data generating model inhomogeneous, and we show that in this case, non-contrastive method contrastive method can learn all features. We also include output normalization in our analysis, a technique that is widely used in practice but is still under-studied in theory.
collapse to approximately rank-1 1 solution while
However, despite the complexity of the analysis, our model is still linear, which very different from contrasssevvv methods, do not consider
the models used in practice. Also, for the results more advanced training techniques such as Grill et al. (2020) and Chen He (2021). We leave the analysis of these more practical techniques for future work.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A Theoretical Analysis of Contrastive Unsupervised Representation Learning, February 2019.
URL http:htt////. /arxiviivvv....///v//.a.////aivv org/ abs/ 199.. 09229. arXiv: 1902..9929 [cs, stat].
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15750- 15758, June 2021.
Ching-Yar Yao Chuang, Joshua Robinson, Yen-Chen Lin, Antonio Torralba, and Stefanie Jegelka. De- biased Contrastive Learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 7765-8775. Cur- ran Associates, Inc., 2020, URL http:: hgtp:////pooeee........oo... file/ (alee/3c3dcce 63c3ddcc7b2 waw.aaaee22c19fffffe Pap.rrr.ppppppppp.....
nwwrips. pttpiiss.ccppppp/220000
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent a new approach to self-supervised learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp- 21271-21284 Curran Associates, Inc., 2020. URL https: : https: http://ptccccee.cceooepopoc .ew.eeeddddinnnnnnn.....e. neurips. cc cc/pappr/2020 /ele//iel//////eeee ff
00077701221111
42b1 7781 9 92w2958e-aaarr.eeee.
Jeff Z. HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable Guarantees for Self- Supervised Deep Learning with Spectral Contrastive Loss. arXiv:2106.00 04 156 [es, stat], August 2021. URL htwp::/t///o /atp:v///rrvvvva//
oww.//ii.....oooooooo...o.oooo.r aabs//////// /2106.166000.. 04156, arXiv: 2106.04156.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition, 2015. URL httss:/ /ttps//aaaiiii www.wii..../.vaaoo 151111122 215112221 03385.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9726-9735, 2020. doi: 10. (U09/CCVRR4200..22000000
X. He and Y. Peng. Fine-grained image classification via combining vision and language. In 2017 IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 7332-7340, Los Alamitos, CA, USA, jul 2017. IEEE Computer Society. doi: 10.1109/CVPR.2017.... /ttps://doi.. meceii.eeeccoe www.eeeeeeeeeeooceeeo...... www.ii..rrgl...0
URL Cww.///////22O 2017.
Yu Huang, Chenzhuang Du Zihui Xue, Xuanyao Chen, Hang Zhao, and Longbo Huang. What Makes Multi-Modal Learning Better than Single (Provably). In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 10944-10956. Curran Asso- ciates, Inc., 2021. URL https::t////.....t httt:p::/pproceediie 1aa3405a3f365c104422 315500
neur iprr.p..... cccppaper/2221/i1ee pdf.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision.
Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4904-4916. PMLR, 18-24 Jul 2021. URL htttps://prceeeeeeeeeteee /ttpp/edeceegc...//////p/ ..t.s..... mlr. press/v139// httl://jaalbbbhh...aa html.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding Dimensional Collapse in Contrastive Self-supervised Learning. In International Conference on Learning Representations, 2022.
Jason Lee, Qi Lei, Nikunj Saunshi, and JIACHENG ZHUO Predicting What You Already Know Helps: Provable Self-Supervised Learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 309-323. Curran Asso-
ciates, Inc., 2021. URL htt.h::t///// httpss:/ http:///wwo.eiiioooeeeeeeeeeeeee httparipp..pppppp//////tteee 02ee5aaeee0FF33990000999933333
10

---

Under review as a conference paper at ICLR 2023
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: Common Objects in Context. In David Fleet, Tomas Pajdla, Bernt Schicle, and Tinne Tuytelaars (eds.), Computer Vision ECCV 2014, pp- 740-755, Cham, 2014. Springer International Publishing. ISBN 978-3-319-10602-1.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zetrlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019. URL https: arxiv v. org/abs/
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33) 2018. ISSN 0027-8424. doi: (0003///aa.18065791111 URL https: /tw::///wwwww. pnas org/ Content/111/335 E66...
Kento Nozawa and Issei Sato. Understanding Negative Samples in Instance Discriminative Self- supervised Representation Learning. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:// svdPQ.
openreview. nwt./weeeeeeeeee htte://oormmiidddomed
Ashwini Pokle, Jinjin Tian, Yuchen Li, and Andrej Risteski. Contrasting the landscape of contrastive and non-contrastive learning, March 2022. URL http:://aaaet http:p:://xxx....///aa.. www.aasaaa0 arXiv:2203.15700 [es, stat].
22200.15222 15702.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning Transferable Visual Models From Natural Language Supervision. arXiv:2103. 00020 [cs], February 2021. URL http://arxive. wttpp/wwwo../b///....3333 0002.0000.. arXiv: 2103.00020.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8821 0821-8831. PMLR, 18 24 Jul 2021. URL
https: http::s//prcccednngs.. mlr. w.r..mmrr........ walls.pres//vee h http:9/rmmme22a.ahhaeeaaa html.
Mark Rudelson and Roman Vershynin. Sampling from large matrices: An approach through ge- ometric functional analysis. J. ACM, 54(4):21-es, jul 2007. ISSN 0004-5411. doi: 10.1145/
1255443.12554499 URL htwph:tt///cc. /tt::://oooooooo org/ 10. 11....155 1ee./1151111 1145/1255444311111111111 1255449.
Jonathan C. Stroud, Zhichao Lu, Chen Sun, Jia Deng, Rahul Sukthankar, Cordelia Schmid, and
David A. Ross. Learning video representations from textual web supervision, 2020. URL
https: //tssv//axvvv.///// org/
149377........
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynam- ics without contrastive pairs. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learn- ing Research, pp. 10268-10278. PMLR, July 2021. URL https: httpt:://pocceeeeeeeeeeeee pttpp/eeeceenssee//eeeeep.e...se mmr. html.
Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive learning, multi-view re- dundancy, and linear models, April 2021. URL http:://axixi wwwpxxii.orggbs////oeeaa 10150. arXiv:2008.1 10150 [cs, stat].
Tongzhou Wang and Phillip Isola. Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere. In Proceedings of the 37th International Conference on Machine Learning, pp. 9929-9939. PMLR, November 2020. URL https: //proceedings... mlr. www.mm.........l.m.mmm......m... /ttp9///aggk.../m// html. ISSN: 2640-3498.
Yifei Wang, Qi Zhang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Chaos is Ladder: A
New Theoretical Understanding of Contrastive Learning via Augmentation Overlap. In Interna- tional Conference on Learning Representations, 2022a. URL https: / httpp//eeweeennee.ccc net forum? d=CCgggVyee..

---

Under review as a conference paper at ICLR 2023
Zhecan Wang. Noel Codella, Yen-Chun Chen, Luowei Zhou, Jianwei Yang. Xiyang Dai, Bin Xiao, Haoxuan You, Shih-Fu Chang, and Lu Yuan. Clip-td: Clip targeted distillation for vision-language tasks, 2022b.
Zixin Wen and Yuanzhi Li. Toward Understanding the Feature Learning Process of Self-supervised 15134 [cs, stat], May 2021. URL http://arxivv.org
Contrastive Learning. abs/2105. 15134. arXiv: 2105.15134.
Mengde Xu, Zheng Zhang, Fangyun Wei, Yutong Lin, Yue Cao, Han Hu, and Xiang Bai. A simple baseline for zero-shot semantic segmentation with pre-trained vision-language model, 2021.
12

---

Under review as a conference paper at ICLR 2023
A GRADIENT CALCULATION
In this section, we compute the gradients and the equations governing some related quantities. We first consider the general finite-width case, and then the infinite-width case, in which we have simple formulas for many quantities of interests.
A.1 THE FINITE-W WIDTH CASE
The results in this subsection mostly brute-force calculations. First, prove the following auxiliary lemma.
Lemma A.I. For any x A and 2 B. have
AKT AcK A.E) /d
Vw exp(T) = -i exp(Ti fm)
-
xt(5n) - (SA.
NA
NA
Then we compute the gradients.
Lemma A.2. We have
AK AcK A.E
xA(fB)
Ww -
2
E
(2 SA(T) 2B) - SB( xB))
(fA
NA
Nad
2
(AK) + AcK /Va.
KT²
SA 1: x exp(1 exp(T? fm)
+
E
(t - -(ft
NA 2
exp( ( f
VIKAII + IK
(AKT + AgK /Va VA + +IAA&l
KT.
SB expl ( fA FB) eXp(Ti fH)
E
x (AFB - (fA
NA x a
The formula for V Wa can be obtained by interchanging the roles of A and B.
Instead of tracking WA and WB directly, we are going to track KA, KB, aA and K The dynamics of them is governed by the following equation. As direct corollary of Lemma A.2, have the following.
Lemma A.3. We have
d dt
KA
(2- SA(E xT) - (2 (KBz* + -KBEEB)(Zt)
KA N )
2
(ft.
diag(o?)
NANB
#
SA(E)) x expl (T+ ft fB)
ff(lt)]
K A
KT
E
(f)
diag(o²)
+ 00
e f f
NA
N
SB(w TB) exp(Tiff f
f+(x-)
KA diag(o")
KT 2
E
(fA.
exp S 1
NA
N
and
d
T2 NA T
(ft fE) KAE/Va
KAS
E
(2 (2- SA(x), 2 Sp(at,
(B(E
d
VILK + IKAAE
KT?
S SA(I) xt (1) exp(TT ft fT)
KAe/Va
E
fm(Et) VI + KACl?
NA at m
exp( f
KT
SB(I 2 exp TifA f
S (f%. KAe/Vo
E
N A
exp(T? 15t f
VIIKA 11 + +I|KAS|!?
13

---

Under review as a conference paper at ICLR 2023
We can rewrite the above result as follows.
Corollary A.4. Define
Qo :=
E {(2
SA(T
m3) SB(at. - (ft: {3)
SA(I xH) exp(TiSS fB)
+ K
R
(f.
A =
cXp(TifA f3)
Sn( xT) exp (
(5x.5t)
+ K
E
exp(T. fB)
"A
and
Q1 = E {( -SA(() A x3)) S xB)) z x++))
SA A x B exp (i fB. (
KE
expl +p fĀ 53)
SB T exp(=] S x (x
KE
exp(Ti SB)
Q 1 {
SA A x (n i #
Et TT
SA 2t exp 7if fB.
KE
E (
exp(Ti S ST)
Sn(# A exp(TiSA S ct( Td
KE
cxp( fA fD)
Q1EE = E ( SA x1) Sn A x (x Tx
SA v exp T S {x
KE
exp(T f fT)
SB(T x exp exp(Ti5A f (x
KE
exp( fā)
Q2 E ((2
SA ( 2T) SB SB x TT df Sn T
SA(EA x exp f
KE
exp(T? ft)
SB( x exp(=?fA f
KE
( d
exp(T? {T
We have
d
KB
KB. Qr.e diag(o²)
A Qo diag o²).
KA
Q1 diag(o
dt
NANBd
NAd
N d
KA
KA Qlea diag(o"
K
KB
o diag(o²) + NAN
+
Qo diag(o")
di
NANad
N
d
KB
KBÆ
Q20 + kA
KA.
QTe, oe
+ NANBd
Qoo
dt
NANBd
Nºd
d
KA
KA
K KB.
K
TTe oe +
Q o + N
Qoa
dt
NANBd
NANBd
We are interested each column of K and KB, whose dynamics given by the next lemma. The next lemma also decompose the dynamics along the radial and tangent directions.
14

---

Under review as a conference paper at ICLR 2023
Lemma A.5. For any p [r] and q [d - r/. we have
| = KAlpp K Q1pp)
KKlpp, K B QI..mp)
+p 1111 ||[ | 1|
op +
20o p
dt
NA
Nd
N d
||abbl]² (k ((Qi
([KBl KK A.E A oz +2
[K Blll
0g +2
200 p
dt
NAN
NAN
d
=2 (KAs K
KBQIAAM
L
K KA K B.
K AKAgllll N
di [[K.&e/lll
NAN
0g
2
+2
NAd
06
2
+2
Q00
||[KB (KBskg [K KAQLEo
d
[K K KA.
19)
0g +2
NAN O6 +
Q0o
di
NAN
N3d
and
KBQ
KI.CQL
of [|(abpll
KAl
I
KAlp
(AAp)
[ Alll
NAN of ||KB]p
KAQ
KA..QL.
K KBIp
I
KBl
(Blpp)
dt
[Blpl
NANBd
[KBQT.e
K B.
K KA,el
I
KA.el
[KA.lg)
+
dt
A|Llll
[KA.el
NANBd
')
[K [KAQT. T.egh
KB.l dt
I
(BBelO)
KB.ele
NAN
Finally, as a simple corollary of Lemma A.1, we have the following result on the gradients of the non-contrastive loss.
Lemma A.6. For the non-contrastive loss (4), we have
(AKT +
AcK A.E) N2d
x (fB
VwaC=
(ft.
NA
As corollary, we have, in this case,
f+(
K A (S.
KA
diag(o
dt
NA
N2 d
(
KA. N
kA. = E.
f
dt
NA
OMITTED PROOF OF THIS SUBSECTION
Proof of Lemma A.I. We compute
Vw exp(Ti fí) 7? exp VE) (ft
(WA $
= VE
KAll + /KA.el///Va
VWA W (T SE)
= - f3)
KAll
(KACIE/ //d/
V 1A + KA
-x (f)
KAlE + KKAel
15

---

Under review as a conference paper at ICLR 2023
For the first term, we have Vw (W AA fB) = (fFB) For the second term, we have
VWAy / KAll + KA 2
Vwa
KAll + VWa ||A&el
AKT + AcK A,E
V KAll + KK
| All; 1|p + ||
| F |KA&l
Hence,
=X(FE)T
WA exp(Ti fB) = exp(T fB)
V |K + |KA, .l..el//Vv
AK + AACKAE A.E
- exp((i fB)
|K 1 + +|Ka.ell'
(AK) AeK A.E) /Va VIKAali + +||KA
of explTF fB)
T - (F)
KAKl 1/ + KA
Proof of Lemma A.2. We compute
V Vw ( A SA A 3
V LA - E
exp(T? fB)
= E
VWA
SA + a
a B)
exp(TiSx 51) K Eg- exp(Te TT f
1
Vw exp( fT)
EE
SA +
A. BB) eXp(T? f3 + K E,- exp(T? ( fA fB)
VWA eXp(T f3] + KE,- Vw, exp( fB)
+ E
eXp(TifAA fD) + K E,- (T. fB)
By Lemma 1, the first term is
1
VWA exp(1 ft)
- EE-
SA A hT) explTi ft) KE exp(Ti f
(AKT AeK A.E /Va
E
x xA( - (f)
NA
VIKAll + KAEl
and the second term is
Vwa exp(7 f + KE, Vw exp(TiSā SE)
exp(TIST
f + KR explTIf fB)
[AKT + ACK /Va.
E
(xx(fb) - (b)
NA
VIKKB + -KAAE
(AKT + ACK A. /Va
KT²
SA(E), BB) exp(Fi f eXp(TifX f1)
E
T (F)
NA
VI + +1KAE
Thus,
(AKT + ACKXe) /Va (x((TE) - -(fA.
CA =
TI
E
(1- (1- (
(*A(F))-
NA
VILKa + +IKA
(AK) ACK x. /Va
KT? +
SA TT) exp( fā)
x*(FE) (fX
NA
exp(T2FA) f1
VIKAll + ||KAEll
16

---

Under review as a conference paper at ICLR 2023
Then, for Vw LB, we compute
WWALB= = -E Vw Sp SB (1 x x
SB( A A.
Vw exp(TiSA 1D)
=-
SB o A a exp(T fA fB)
KE E 2 exp(T1 fA 5B}
VWa exp( (7 fā) KVw, Ba exp(TE ft)
E
exp(T? fī fD) + KR exp(Ti SB
Again, by Lemma A. the first term is
VW eXp(Ti fū)
E
SB A:
BB)
exp(Tifā fD) K KR exp(T? 2 fā fB)
(AK] Ag K A.E
/vd
7
E
xA(F
fB)
NA
| 2 + KKAE
and the second term is
V W exp (T. 5 + KVw, 2 exp(TI SB)
E
exp( (TP(Tif (D) + K KB,,, exp(TiSA SB)
[AKT + ACK IVa
E
x (A( - (F) SB( TB))
NA
VIIKAlA +IKAAl
(AK A /Va VIIKAli + I| +IIKA
KT? E
SB at)) exp(T?5A 5t)
x7(fE) T FF)
+
NA
eX fB)
Thus,
(AKT + AeK /Va
Vw LB
7 E.
(1-Sp(xx).)))) (BLUA)
(.(
- -(FA
NA
VIKKAl IKAE
(AKT ACK A.E /Va
T
Sn x B) exp(T?/f ft)
xx( (FI.
NA
exp( fA fB)
v |IKAI IK
Combine these formulas together and we complete the proof.
A.2 THE INFINITE- CASE
Now we consider the noiseless infinite-width dynamics. The results of this subsection will not be used in the proof. mainly serves as a way to give intuitions on how the dynamics looks. As we have discussed in the main text, in this noiseless infitite-width case, it suffices to track ni? = ||| K 1 1|
and 23 {KK Alp, [K Blp)-
Lemma A.7. In the noiseless indimite-witth case, we have
dt K
d 2
(1 (1-
||<:|]
R22
p
|x]
2
||<=|²
Rip
2
07 4 (1-5)
||=_|
A2 T
|G/||
R2
T
02
d
22
(1-8)
x
||
2
A:2
o7 (1 (1-5) s)
nop
T,
R2
T
di p
||<:||
||<:||
2
||<=|]
2
||>_||
||==||
17

---

Under review as a conference paper at ICLR 2023
Proof. First, note hat
(KBAz ,2 (fā fn)
(2 2 R2
||=]|² /d
||::|] /d
This implies that (a) (fā B does depend the actual value of 2t, and (b) if we flip the signs of 2p simultaneously, then the value of (fA- remain unchanged. To compute SA. then need to take expectation over z We compute
E exp ( fn)
II k=1 *
E exp
Ti
2 2
&
II cosh
Zo. T 2
||]||
/d
||x>
Again, it does not depend on the actual value of 2+. One can conduct similar calculation for SB and, consequently, we have SA = SB = for some Š that depends on k and & but not on z Then, we can rewrite (5) as
d KA = (1- dt
fĀ(Z))
K A
E
NA (fX
diag(o²)
N d (fn(zt)
exp(T? fn)
KA
} diag(o")
2KS
E
(fx
EXp(T? ? f fn)
NA
N2
= 2 -2(1-8)
K B
|
A
(o²)
||AG|]
||A
||=||
2KŠ
T²² (z 2
12 d
KBz (2+)T
(2t,2 ( KA
7 1 ||=|||
E
exp
|x]|
||
||<l /d ||a011
diag(
exp
Note that
7 (2t
)RO
2 +
exp
z
||=]|
T 2x
22 2 2 d |/ 12
}
{p- 4 E
exp
2p 2
E
exp
||
2p
T. |||s||2
T2 ||<=|2
-p = sinh
II cosh
-fp ZoTp
Meanwhile, note that
+? (a+
22
(zt, 2=)
62
T? (zt 2 122 d ||]
cxp
2 2/
|
||< 1 /dd
|
/d
k=1
2n
Tk
||::1|
1=1
7.7.
18

---

Under review as a conference paper at ICLR 2023
Thus,
d dl
K 2 (1-1)
K
||A]|2
KA
diag(o²)
||<=|?
||~=||
||==|
2K 2KSZe
KB
KA f ding(o²)
diag ([Tk)ke(r)
-
||<=|²
||x||²
exp
= -2 5)
(
KB
||A:|2
KA
|||=||
diag(o²)
||~=||
||==||
- 2 ( -)
S)
K B
diag (ITWke())
KA T ||=||²
diag(o²).
||>=|/
As corollary, we have
d dt
KAlp 2 -1(15)
||:11
[Kapp
o7
-2-2 2 (11-3)
(KKB) ||x||²
Tp
[KAle I
[K
og.
||<||2
||ai|]
||<]||
||~:11
Hence,
d
d
dt kp = 2 [Kalp= dt
2
(Kapp)
R2
||:: 1 ni
(
= 4 (1 (1 51
o7 2 ( (1-
TT
og.
||0e||
||+b||
||<|]|
||0|||
||/o|
By symmetry, for KB,
B, we have
d dt
KBlp = (1-5)
[Kalp
||
Blp
0z 2 2011 2111)) .)
(IKr))
T,
[K
oz.
||a/|]
||
|||a||
||<_|/
||x]||
Then, compute
d R3 4 (1-5)
(
||
o/ - =1--5)
(
ng
Tp
R2 T
o² p
dt
||<s||
||=
||<=||
||||||
||a||²
B STAGE 1
In this section, we show that the following hold:
(a) KA* K after Stage 1.
(b) The noise-signal ratio is small after Stage 1.
(c) The condition number is O(Vd) in Stage 1.
(d) The trajectory is still close to the infinite-width one in Stage 1.
We formalize the main results of Stage 1 bellow.
Lemma B.1 (Stage 1). Under the assumption of Theorem 4.3. We can choose sufficiently (poly-
nomially) large m and a sufficiently (inverse polynomially) small Ti that may depend on the 8' that appear in this lemma so that the following statement holds.
Let Ti be the earliest time all the following hold:
KAlp KBIP)
2 1-18. S_,
Vp E [r].
KA < <ŌN/S:
Vp 1rl.g [d r].
Alpl
19

---

Under review as a conference paper at ICLR 2023
where 8_,6N/S E 1/ppoll poly(d) are two given parameters. We have Ti I poly(d). Moreover any time E [0, we have Ko = maxp,1ll1) || || || - ((vd) and
[|[KA]
(1 AA/B) KBbpll
KA.s/el (1 SA/B) K KB.e
p r),
KK A.
S dgg
e [d-]] -rl, r],
K 1,1)ll
max
{K
(Klla)
C, & ÓAB.L:
vp # E
max
[Kolll K (
(Kr.elo. ( : C, C (A. sg
Ap g. ( -],
(10)
where 5 AA/B: 6g AAB.L, 6,,1 1 1/ poly(d) are given parameters.
Basically, the conditions in (10) mean that the norm the corresponding columns are roughly the same, and the columns of all these matrices approximately orthogonal to each other. Both of them are true the infinite-width limit, and by some standard concentration argument, one can make all these errors to be arbitrarily poleosllpoyynmmiall small, at least at initialization. Note that, as simple corollary of (10), we have
NA NB (1 (dABB
For notational simplicity. we also define
K A. K ||
(KAl (BBD)}
= max pEp1
PN/S = max peir gEld-r]
Characterizing the dynamics in Stage 1 is relatively straightforward. We will in Section B. that all those Q-matrices have simple forms. The main tool we use to control the condition number and the discretization error the following nonlinear version of Gronwall's lemma.
Lemma B.2. Let A, be positive process. Let X, and Y, be defined as
X.S A,X,, aaxxY.
with Xo, Yo: being positive. Then, for any T 2 have YT & Yo exp(oxx))
Remark. Here, X, represent the progress have made and Y, the error. In our case, X, is the maximum between
KAlp KBIp) and the noise-signal ratio, and Y, the discretization emror This lemma says that, if the error growth rate depends on the progress, then by coupling these two processes, we can make sure the error does not blow up. The point of this lemma is that, with coupling, we not need a very tight estimation on the convergence time nor the error growth rate.
Proof. The solution of this ODE system is given by
Xy Xo exp (-/-" A, att)
aXx /" A exp (-/² A. ds) dt
Yr = Yo exp
Note that
S A exp (-//" As ds dt dt=
-d=--f 1- dexp (-f² As ds
exp (-5² A A,
Hence,
Yr exep aXo (1 exp (-/ Ar dt )) < Yo exp (aXo).
1-exp
Ardt)
20

---

Under review as a conference paper at ICLR 2023
The organization of this section is follows. In Section B.1, derive estimations for the Q- matrices defined in Corollary A.4 and use them to simplify the equations governing the training dynamics. In Section B.2, we estimate the rate which 1 AAlp KBlp and the noise-signal ratio converge to 0 and the growth rate of the condition number. In Section B.3, we estimate the growth rate of the discretization error. Then, Section B.4, we prove Lemma B.I, the main lemma of Stage 1, Finally, we prove the negative result for non-contrastive learning in Section B.5.
B. ESTIMATIONS FOR Q AND THE DYNAMICS
Thanks to Lemma A.5, in order to analyze the dynamics, it suffices to estimate the Q-matrices. In this subsection, we derive estimations for them and use these estimations to simplify the equations in Lemma A.5. Recall that, in Stage 7 is small. Hence, exp(Tē ( = 1 O(T²). With this approximation, one can derive the following estimation for SA and SB.
Lemma B.3 (Estimations for 5). In Stage I, we have
1
SA 1+K +0, (TH) # (T
and the same is also true for SB. Here, (z(r²) means (TT)) and this does not depend &A nor ÊB-
The proof of this lemma and all following lemma is deferred to the end of this subsection. Note that derive a slightly finer estimation for the noise-related part. This additional PNys will be used cancel with terms like [ |pll /|[ A.,e)lll, at the cost of Ko factor, in later analysis. We emphasize here that O₂ does not depend on the noises so that later we argue Eg (O,(r())) 0. With this lemma, we now derive estimations for Qı, QI, e Q2 and Qo. respectively.
Lemma B.4 (Estimations for Q2). In Stage I, we have
2K
Q = I K La + (dT,)
Lemma B.5 (Estimations for Qi.e and Q2). In Stage 1, we have
max { 11 Q 1.6 |
+0 pNISbe..).
The above two lemmas say that Stage 1, Q is approximately diagonal, and Q1,e and Q2 can more or less be ignored. Lemma B.6 (Estimations for Qo). In Stage I, have
2K (KA, KB) Qo 1+K NANBd 0
The proof of Lemma B.6 is essentially the same as the proof of Lemma B.4 so we omit it. With these three lemmas, we can now simplify Lemma A.5 as follows.
Corollary B.7. In Stage I, for any p [r] and q E [d r). have
4K
of
([KAlp) [KB)p
(KA, KB)
di |K Ab1l
[[[A]b||
1 K
Bd
NANBd
+
NANBd (VBABB+ + + ki
d
||
4K
of
(KA,
||[ +0
( ot
(drt' + voa/n)
|[[ [|AA.llll]
dt
1+ NAd
NANBd
NANBd
The formulas for K B and K B. can be obtained by interchanging the roles A and B.
21

---

Under review as a conference paper at ICLR 2023
Basically, the two parts of the first term of dt ||pl correspond to the signal and normaliza- tion, respectively. For d independent and have mean 0. Corollary B.8. In Stage I, we have
ae Alelll the signal term is because the noises in x and o are
d Kap
+ K NANBd (1
2K
of
(I-l (Kaio) [KB), =0 ( of NANBd (fe + / (A1B)
[Kal
dt
d KA.E =0
oe
=0
aed (.nnoo..)) nodg,11)
dt
NAN
Interchange the roles of A and B and obtain the formulas for K B and K B,
Note that, after normalization, the normalization terms cancel with each other, and the signal terms drive Kalp and A toward each other. Moreover, without the normalization terms, different A are approximately independent of each other. This allows us to maintain the orthogonality between different columns.
Note that the above results also imply the following lemma.
Lemma B.9. The dynamics of the non contraateet method and the Stage I dynamics are equivalent, up multiplicative constant and some higher-order terms.
OMITTED PROOFS OF THIS SUBSECTION
Proof of Lemma B.3. First, we write
(Kazt, Knz)) NANB
(Kaz*, (KAzt,KBeEi) KB.EA) + (KA.EE KB=)) (KA.E KB.E.)
(5)
NANB
NANB
For the second term, we write
Kaz KBCC
[[
[[Kell
Σ
(KAli, (B..)) 2 Et
NANB
NANBd
[[Kll
jE(r)-((ar))
For each summand, we have
(B,l,l Alll S O(pN/s), (KATi K KB&B)) :))....).
and 1 E 1. Hence, WaaztKBeet =0(d8;, The same is also true for NaNn
KA (KAet and the third term. Therefore.
Nad
(KAzt, KB = =ddde..IN/s)
(ft,
NANB
Then, we compute
KAzt KBz (1)
exp (ft. fE)) exp "
NANB
=1+0, (T?)
(7, PP//))
Thus,
1
SA 1+ K
+0, (Tt) + d8 d7:
The proof for Sg is essentially the same.
Proof of Lemma B.4. Recall that
Q = E (( SA +
A. - x x [x]
SA ( : xt expl f J = (2t)Td
KE
exp( f f
SB(x B exp 7 (T. f1h x+ +x (x xx
KE
exp(T2 f
22

---

Under review as a conference paper at ICLR 2023
Since SA = (1+ 0 and SB = (1+ we have
2
2K
2 2-5 -
=2
1+K + 0
1+K
As a result,
2K
K
K
Q1
E {z (zt)
B (z )
d}
B( )2-) :0 (dr,)
1+ K
1+ K
I+K
Note that E (2+) and E. (-+(2+))
Ia/d. Therefore,
2K
Q
la + + (dv)).
1+ K
Proof of Lemma B.5. Recall that
QLaa == E {2 SA(2) - #z)) ( df
SA(EA xx) expl exp(Ti fn) 5x
KB
exp(T fT)
SB(# TB) exp(=? fB)
KE
exp(Ti fA fB)
Note that if some quantity X does not depend on &. then E [XE] = 0. Hence, by Lemma B.3, we have
Qi,a = =0 ( pNysoe,
The proof for Qlea and Q2 is essentially the same.
Proof of Corollary B.7. Recall that
d
K [Apppl = 2 ([K alps [K KnQuI)) NANBd
(Kaap [K, (.eepl)
og + [|[ 1lpl
og + +2
N2 Qoo
dt
NAd
d Tr
[kbbb)
dt
We now estimate these terms bye-by-one. For Tı, we have
([K p)p (KQQi)) (Kahn (KB))) 1Qth + Σ ([K Alp: [K])))[Q]ee.
2K
(Kalp, KBlp) + 0 (7, dr = 0 (7. dn ko AB.1)
1 + K
2K
([Klpp [K + (T)d=²)
1 1 K
Hence,
( d dt
4K
2
Ti
p
(Kapps (KBb) (Kal)(Kah)) +0 +0
1+ K NAN Bd
For T2, compute
d [[Aa)|)])
20 d-r
T2
dt
=
NANBd
[Klpp (KB.le)) (Q1.alk.p
[[B.g]
=0(1)
[Alpll ?
[KAlp: KBal)) 1Q1.cake
NAN
k=1
=0
NANBd dd
PN/S 12)
23

---

Under review as a conference paper at ICLR 2023
Finally, for Ts. compute
T3 dt
[ [[Kappp²)
2
11 Abol
2K (Ka, Ka)
+ +0 (dTtt)
og
Na
1 K NANBd
4K
of
(1+ (DAIB)
(KA. KB) NA
Kalpll +0
NAN do?
1+ K NAN
4K
of (K),KB)
of
(VOA/D + rid) mip
#
1+ NAN
NANBd
NANnd
Combine these together and we obtain
4K
4K
of
K Alpll
P
(K (KBy)
dt
1+ K N
NANBd
I+ K NANBd NANBd
of
of
2d 2 PNISS 2.112
+
NANBd (VAAIB- + +zdd) R2
NANBd
4K
of
(Kapp ( KBlp)
(K KB) NANBd
|||K||||
1+ K NANBd
||[AApl|
#0 of
(VANIB + rid) Ngg
NANBd
Interchange the roles of A and and we obtain the formula for |||KB]pl Now we consider KA.g- We write
(Ka
[K Qisaa))))
og + ([K A.gg. [. V2le) of + N2d
||[ A.elg
K K
=
Qoo
dt
NAN
NBd
Tj
K (KAelll)))
dt
1=1
For Tı, compute
Tı (
[KAssll?))
K KA 4.
KBIc IQ1.aalak
=
dt
NANBd
km1
20
[KK]]ll
[|
K (K
1Quealee
NANBd
[[
k=
20
=0111)
II[KAElol|* Σ
Ko PNIS
PNYSOE.1
N
Ndd
+0
o NA
|||
Similarly, one can show that this bound also holds for T2. Finally, for T3, by Lemma B.6, we have
( d dt
= KA.
(1+ VSAIB) ( [T NANBd + Od+?)) of
2K (KA. KB)
T3
NANBd
4K
of (KA, KB)
1+ K NANBd NANBd [A..l
dF
(d + (0A/B) [[[Kscaale]
NANBd
24

---

Under review as a conference paper at ICLR 2023
Combine these together and we obtain
d
4K of (KA, KB)
|| ( of
NANBd (dri + voA/B)
[[[ [IKasbolit
K
dt
1+ K Bd NANBd
Interchange the roles of A and B and we obtain the formula for K B,C-
Proof of Corollary B.8. We write
d
K K of ||A|pll] NANBd
KApp
-KA
[KAl))
di
KB..QI.5nle of
+
I
[alp
(KA)))
[K Abll
NANBd
d
d + +T²
Tı
KAlp
KAlp
dt
dt
For Tı, we have
d Ti
[Qikkp
of
dt
Kalp
=
Σ
I KAlp
(Alp)
K KB|}
||| A|[l
NA
[KBBpl [Q1
KAlp
(KAl))
[BB)
P
|| |Kaa 1
NANdd
Σ
I- KAD
(KAb)) ))
BBk
[KB]all [Qı]kpp
o
|Alllll
NAN
kpp
For the first term, by Lemma B.4, we have
K [ (1 (
2K I+K 0 I 1 (dri' (DAAB)
2K
KK Allll
Also by Lemma B.4, for each summand in the second term, we have
[KB | [Qile
+0 (r,dro).
K Allll
Therefore,
Ti (Klp)) ( dt
2K
of
I- [KAlp (KKAio))
[KBp
1 1+ K NA
+ NANBd Bd
K KB) 0 (Fddroo (8A/B)
For T2, by Lemma B.5, we have
d To
1 1 KAlp [KAb)] ")
KB.ll [Qiealee
KAlp)
K BB.E
di
=
NANBd
NANBd
k=1
Combine these together, and we obtain
d
KAP = 1 K
2K
of
I KAP (KAlP)
KBIp
dt
NBd
d-r 0
+
P
Σ
(7, + VoAIB)
KBI +
Σ
7, (iddPh/sSc..
KB.l
NA
Bd
NA
km
kml
2K
I A
(KAlz)
K BB +
(T.0PRo d
1+ K N NANB
NaNVd
25

---

Under review as a conference paper at ICLR 2023
Now, consider K..g- Again, we write
d dt
1 KA.E (KAtEl)) ))
1KnQisn i of
KAEW
|||K A.llol NANBd
[ BIB.Q221 og
+ I
-T-Ke -AKa
(K (KARTA)
NA
d KAA.ll
d + +T2
Ti
K KA.Elo
di
di
For the first term, by Lemma B.5, have
( d
r
(KASl)) "
KBl QQ1galg. of NA
Tı KA.eg
KAcl
KBl
dt
||[K
+
Ko
of pN/ISoee
NAN
PN/S
k=1
o²
+0
NA 7 2rokoss
The same bound also hold for Ty. In fact, we can have slightly sharper bound for it because we no longer have B| K A. Combine these and we obtain
d
K KA.g = +0
of 7ie (ddkoodg.)
dt
NaNBd
Proof of Lemma B.9. The proof of Corollary B.8 and Corollary B.7, mutatis mutandis, yields
2K
KB
2K KA (KA: KB)
KA =
diag(o?)
diag(o²)
dt
1 K NA
1 K N NANBd
+
max
dd + (5ABB)
|KAAlF)
NA
KA dt
2K
kA KB)
oe +
of Nad
1 (KAell))
dii
1+ K N NANBd
Recall from Lemma A.6 that, in the non-contrastive case, we have
(
KA A
KA =
diag(o²)
d
NA
N
KB
(KA. KB)
KA
d
diag (
NANBd
NA N
KA = (K KB) kA
NANBd N o
di
Note that they exactly the same, except for 2K/(1 + K) factor and some higher order error terms.
B.2 CONVERGENCE RATE AND CONDITION NUMBER
In this subsection, estimate the rate at which p- and PN/S converge to 0 and the growth rate of the condition number. The basic idea is to use the estimations we have derived Corollary B.7 and Corollary B.8 approximate the finite-width dynamics with the infinite-wiee width ones.
26

---

Under review as a conference paper at ICLR 2023
Lemma B.10 (Convergence rate of p-). In Stage 1, we have, for any p E [r].
d KAlpp KBlp)
4K
of
(1+ KAlp: KBlx) (1
Kalp KBlp)
dt
1 K NANBd
# p
(R.dRo + /DA/B)
NANBd
Note that 1 + K (Klpp KBlp) rate.
= (1). Hence, (Kam K Bln) will converge to at linear
Now we consider the noise-signal ratio. For some technical reason, instead of characterizing the dynamics of PN/S. consider
PN/S = [ + K
Note that always have
K A |
2 ((1) (r) A
PAYs 2 ((1)
1aiT [Klppl]
In other words, PN/S can be bounded by PN/S S O dna PNYS). Lemma B.11 (Convergence rate PN/S). In Stage we have
d Je RNis S 1 K NAN
4K
o2
min
PNI + O o
( + VoAlB) 08v1s)
NANBd
This lemma says the noise-signal ratio decreases exponentially fast.
Lemma B.12 (Growth rate of the condition number). Define Pp/o ||[KAppl |I[[K In Stage 1, we have
0p/4 S 1 K NA (p- + min {DN/S, 11) Pvt9 + 0
16K
o max
o Fuax
(VOA/B + t/d) (V Ppia)
N NANBd
Note that error growth slows down as p- and PN/S decrease. This allows Lemma B.2 (cf. Section B.4).
OMITTED PROOFS OF THIS SUBSECTION
Proof of Lemma B. By Corollary B.8, we have
2K
1) =
oe NANBd
(o voA1e))
KAlp
KBLp
1
KKAp KBW
+ K Nd
2K
1 (Klp (Bb
)
K KAln: KBB)
+ K NAN
+
NANBd
Hence, by symmetry, we have
d
[Klpp
KBlp)
4K
K NANBd (I+ KAlpp (KBp) (1 KAlpp [Kb))
di
+
(iddoo + voA1B)
NANBd
27

---

Under review as a conference paper at ICLR 2023
Proof of Lemma B.11. Similar to the proof of Corollary B.7 and Corollary B.8, we have
d
KB
B.
K A
KA =
Q1 diag(o")
+
Q1.cn diag(o²)
+
Qu diag(o")
dt
NANBd
NANBd
N d
KB
2K
Ia + O (dtt)) diag(o")
Nd
1+K"
KB
I
(T. PNYSSE,L) diag(o")
NANBd
KA
2K
KB) NANBd
+ 0d diag(o²)
+
Nād
1+K
2K
KB
2K
KA (Ka, KB)
1 K NAN diag(o²)
NANBd diag(o²)
1
+ +
(r? + ( [KAllF)
NANBd
Define
KA* + KB. Then, by symmetry, we have
d
K
2K
(1- (K A. KB)
K diag(o²) =00
22 max NANBd
(d7? + v5A/B) [KK]1))
dt
+ K NANBd
NANBd
Hence,
d
|
4K
1
(I- (KA, AB)
(1- NANBd
(K. K diag(o?))
#0
- NAN
(dit + (oAIB) |K
dt
1 K
NANBd
4K
min
(KA, KB) NANBd
|K| 0
max NANBd
02
(dri+ + vdAIB) ||<|?P)
1 K NAN
Similarly, we also have
d
KB
KBS Q20 + Ti oe +
kA N KB.E NANBd
KAy== = NANBd
Quoद
dt
NANBd
Kn
NANBd ( oz
(T. PN:sS,,)))
+ KAE 2K KA,KB))
N I+K NANBd =drii)
oe
2K KA (KA, KB)
oz 22+0
of
d ad
Nad dri KK A.ellF
+
o.
|K
1+ K Nad NANBd
NANBd
r
2K KA (KA, KB)
2
dT² |K |KA.s|
de
1+ K Nad NANBd
.2
Therefore,
d
K A
4K of (KA, KB)
|k +0
N2 d7² | (KAclPP))
of
1+ K Nad NANBd
Then, compute
dt PNys= = |K
d
|KAl
PNIS
|K||
|
4K of (KA.KB)
of d P311
PANS +
PNIS
1 K Nad NANnd
Nad
4K
22
(KA,
mmx (dr? + (0AIB))
-Piss 1 NANBd
NANBd
NANBd
4K
2 miin
t NANBd
(dT? + (BAIB) pPN1s)
S-1+ NAN N
PN/S +
28

---

Under review as a conference paper at ICLR 2023
Proof of Lemma B.12. For notational simplicity, define Pr/a := ||[K))p]| / By Corol- lary B.7, we have
d. || pll Pp/4
d.
dr
Pp/9
[|[ lllll
[[ Al
4K
o
([Kpp. pkpp)
(K KB) NANBd
of (Vonnntr8d) + rēd) Prio
Paia +
1+ K Nd
||[Ka|]l'
NANBd (Vois+rid)))
OZ 1+ K NANBd
4K
(Kalqs (Blg)
(KA.KB)
o
PPll
||[K)|l
NANBd
+
NaNBd
4K
of
(IKAL KBlp)
(KA. KB)
1 K NA of
NANBd Pr/a
4K
(KATo (Blu)
(KA, KB)
I+KNANDO
NANBd
Pp/g
VoAIB++H + rid) Po/g
NAN
Now we bound
[KAlp (B))) - (KA,KB) (N)N.d). Clear that this term is bounded by 2. 1 1 p_. For the second term, have
Meanwhile, by definition, have
KAlpp (KB))
KA. KB)
|KAl KBll
|Klle |
KAlk: KBlk)
NA
KAlF
KBIFF
NANBd
(1 (1 + min 1bx
Ki
(1+ v (oA/B)
|
k=1
+ p_ min 1)) (1 v (OABB)
Combine these together and we obtain
(K Kn) m) NAN
(KApps KBLp
2 + 2mi {ix/s, 1] + VoABB.
The same is also true for q. Thus,
16K
+ K NANBd (p_ + +min {pNYS, 11) Pa/a = ( m
2
NANBd (V5A/B + + (2d)) Poio)
Po/9 I
B.3 CONTROLLING THE DISCRETIZATION ERROR
In this subsection, we estimate the growth rate of the eπrоrs described in (10). As in the previous subsections, the proofs are deferred to the end of this subsection.
First, we consider the relative difference between [|[ A| and ||BB)pl' Instead of directly control the difference, we define
||l) and PB/A,P [| } ||
|| })pl1
PA/B,²
Note that PA/B.p + PB/A,, 2 2, with equality obtained iff [|| [KAppl [|[[ Meanwhile, at initialization, this quantity can be made arbitrarily close to 2. Hence, it suffices to control the growth of this quantity. The reason we consider PA/B.P + PB/A,p is to leverage the symmetry. Similarly, we also define
|| and PB/A =
[ BK. K elg 2
PA/BE =
KB.elol
and analyze PA/AE + PB/AE
29

---

d Under review as a conference paper at ICLR 2023
Lemma B.13 (Difference of diagonal terms). In Stage 1, have
of
dt (PA/B. PP/AA) 0
(5% (aa/t ByBB+i
NANBd
di ( 0
max
(6% ( + (fA/B+=/d))
NA
Note that the RHS are higher-order terms, which implies the relative difference of the norms barely grows.
Now we consider the condition number of A. Unlike KA. for the noises, the ss for different coordinates are the same. Hence, have the following simple bound on the growth rate of óg
Lemma B.14 (Condition number of K K. Define Pe,via = A.Ep/lllll/ /l|[K we have
Stage I,
oe
(drž + voa/1))
PEnt So
NANBd
Again, the RHS is a higher-order term that can be made arbitrarily small by choosing small oî and maintaining small AA/B-
Then, we consider the orthogonality conditions.
Lemma B.15 (Orthogonality between signals). For any p + q, define
8.. =
K Appp KBlg)
+ <KBBp KAlo)² + (Klpp KAlo) + (KBlp KBlg
In Stage I, we have
8 dt 8.P
0
max NANBd) p p-pppg +
om NANBd (Fe + /
Recall that p- converges to 0 at sufficiently fast rate. As a result, by Lemma B.2, 8 ò1.p. will blow up. Mcanwhile, since 81.p.q can be made arbitrarily small at initialization, this implics that we make sure it still small at the end of Stage 1. Finally, we consider the orthogonality conditions between the signals and noises and between noises. The proof follows the same spirit.
Lemma B.16 (Orthogonality between signals and noises). For any E [r] and q [d - r). define
1LEA. = ( (KAel)) + + (Krnp, (KA.l))
and define 81,5 8..gmpp similarly. Stage we have
d i. dt
LAPD 50
of NAN
p- p-8ieangg
o NANBd
( dd (8AAB)
(n. +
For any q, 8 E ld - r), define
8 81...99
+ K (K [K (Kaels) KA (..g..)
In Stage 1, we huve
81.. 50
ti 2d dd 1
di
NAN
30

---

Under review as a conference paper at ICLR 2023
OMITTED PROOFS OF THIS SUBSECTION
Proof of Lemma B.13. Note that we cannot directly use Corollary B.7 as the error term contains ó AA/B B. the quantity wish to control. However, by the proof of it, we still have
d KAbp|
4K
of
4K (KA. KB) of [
<[Kapp, (KB))
di
I+
N Bd
I+K
NANBd
N2d
of
#
+2 d
NANBd'
Interchange the roles of A and B and we obtain
d
4K of
4K (KA, KB) of
Bkpppp
K
KBBp|l
(K (KBp)
dt
1+
NAN Bd
1+ NANBd
Ngd
#
+? drip
NAN Bd
For notational simplicity, define PA/B, = ||(KApll
/B]ll
Then, we compute
d [
4
PA/B5*
|
ppl
PA/B.P
||plp
4K
of
([KApp. (B)))
4K (KA. KB) of PAIB.
1 K NANBd
IKAl KAlpl
PAIB.
1+ K
NANBd Nºdd
4K
of
([KA)) (KB
4K (KA. B)) of PA/B.
1+ K NANBd
|BBlll
PA/B.
1+K NANBd Nyd
NAN Bd Tid #
4K
of
(K (Kal))
1 K NA BBd
||Ka)pll'
[BBpl
PA/B2
4K
(KA, KB)
1
1+K NANBd
N2d
Nzd
PA/B.P
of NAN NANBOTTO Rid
+
Then, by symmetry, we have
(PA/B + PB/App
4K
of
(KBl)
(DA/B. - PB/AA.)
1 K NANBd
||[A|p|| I|[ ||[KBbl1
4K
(KA, KBi
1
(DA/BRr (PA/B>-PB/AA)
1+K
NANBd
Nad
NEd
+
NANBdd -id
00
P
" (A/B+T/d)
NANnd
The above proof, mutatis mutandis, yields the result for PA/B.E.q PB/A.5.4
31

---

Under review as a conference paper at ICLR 2023
Proof of Lemma B.14. By Corollary B.7, we have
||(K d
Pe.e/y
PE.vly
||| A.Elall
4K
oe
(KA. BB)
NANBd Ps.p/g =
oe
(di + (
1 K NA
NaNBd
4K of
KA,,Bii #
of
(uti + (OAIB)
PS.ll9
BANNNN NANBd
NAN
oe NA
(drî + vDAIB))
+0
Proof of Lemma B.15. By Corollary B.8, we have
d dt
K Alpp (Kll))
2K
o
(K (alp: (Kal)((X (IKal, [K (KAr))
=
1+ K
Nd
0
p
(ridro + (d/
NANBd
2K
of
[KBlpp (KBp))
Kalpi (Blg)
1 K
NAN Bd
2K
of
KAlp KBlg)
(1-(KAlp: (KAlp) (KBpp
1 K NANBd
NANBd (ne toootd//orp /d/0//B)
Interchange the roles of p, and A, B and we obtain
KApp: pt dt
d
KByy =
2K
d NAN
( (KAPp. KAla)
(Kapp [KB)]]
+ K
2K
(KAl K KBlg)
(1- (Kalg: [K (Bgg)
K NAN
0
NANBd (r.Pko + +d/5A//)
Therefore,
KAlp [Blg]
2K
of
( (KKBlp (BBg)
((KAl, (Kalp (KB1)
KK
dt
+
K NA
2K
0Z
((iKAb. KAlo)
(IKAlo (KBa)
1 K NANDd
22
0
max
(KAlp: (Kn).) (Ka)):(KA)))
NAN
22 maxx
+ 0
(:d 6
NANBd
32

---

Under review as a conference paper at ICLR 2023
Interchange the roles of and and obtain
dt KBlp
Kalg)
2K
og
(BBBp (KBl)) - (Kolpp [K KAls)
+ K N Bd
+
2K
of
((KAlp. Kala)
-(KBlp [K (Kla))
1 K NANBd
2 max
P- (KBlp) KAly)
NANBd
max
(T + /d/2//1)
NAN
Similarly, we compute
d dt
K Kalm
K al,))
2K
of
((Knlp [K (Kla) - -(KKA)p (Kll) (K
=
+ K NA
of
#
(T.dRo + /d/5//1)
NANBd
2K
of
(KBlp) KAls)
(KAlp. Kals)
1+ NAN Bd
2K
of
[KAlp,[Kle)] KAle)
(1-(KA)) (Kalp KBlp
1+ NAN
of
(2dd + -d/5///)
+
NANBd
2K
of
KK KAla)
(KAl) (All))
1 K
Nd
2 maax
P (KAlm Kalo) +0 (
NANBd +okko+d/o//00 )
NANnd
Then, by interchanging the roles of and g. obtain
d
I+ K Bd ((KAl) (Kl.)) - -(KK)) (Kalo)
2K
of
K Allp KAli
di
=
o² o NANBd
+
(KApp: Kala) 0
(g, + +dv5A/B)
NANBd
Add them together and we get
d
2K
KAlpp
KAl.)
p
[Blpp. (Alp)
KKlpp KAlg)
dt
1 K Nd
2K
of
(KApp KBlp)
<(KAlp: (Klp))
1+ NANBd
(KAle, (Kle) +0 ( o
(1.dPR + ddda,11
+
NAN
NANBd
Interchange the roles of and and we obtain
d
2K
+ NANBO ((KAI) (Blp) ([BBpp (Bly)
of
dt (KBlp,(BB,
2K
((KEB) [[KB))][Aa))) -(KKBpp (KBl))
+
1 NANBd
2 max
+
P- (Kn)))(Kal))) +
(7id2 (i
NANBd
NANBd
33

---

Under review as a conference paper at ICLR 2023
For notational simplicity, define 1
Kappp (Blg).
Ze = [KBpp
Z3 =
(IK |p) [ ala). and Z4 ( (k (Knlo).
Also define Gp
2K
Then, we can
I+K NANBd'
summarize the above equations as
Z1
G-pp
Gg
0
Gg
Gp
71
d
Z2
Gp
G4
Gp
Gg
7z2
Zz
G
Gp
Gp
G
0
Z3
Z4
Gp
Gg
-Gp
Ga_
[44]
Z1
0
Fuax
Ze
0
FHa NAN
(Ti + (1/0//1)
NANBd
P-
Z3
Z4
Note that the eigenvalues of the first matrix is -2G, 2G, -2Gp. -2G, and Namely, negative semi-definite. Thus,
dt 2 00
d
P NaNnd P- 1-1Z12
o
max
( + /v5A/B)
NAN
Proof of Lemma 16. By Corollary B.8,
KAlp: KA
2K
of
( (KBlpp, (AA.)))
[KAlp (A
(Kalp) KBlp)
dt
1 + K Nd
of
(Ti + (
#
NANBd
and
d (AAle = +0
of
K KAlps di
NANBd o d' Kode..)
Therefore,
2K
of
[Kll: (A.,)))
KBlp. (K..l - (Kalp, (AAsl))
di
1+ K
Ad
of
o NANBd
+
NANBd P- (KApps KALla) +0
(of + /d/oA/B)
Similarly, we also have
d
2K
of
K [BBBpp[KA.
(Kalp: KA
[KBpp KA.sly)
dt
1 + K NA
+
of
P- KKBB
KAcly) +
NANND (F:dro (1////////))))
2
NAN
34

---

Under review as a conference paper at ICLR 2023
Thus,
d
(Klp, [KA.l + (KBlp (AAk)
dt
4K
og
2 2 [KByp, [KAEl
(KApp) KA.e)
1 K NANBd
of
P- ((KATI KAck + (BBpp (A 2)
Nd
max
Te + /v/AABB
NANBd
of
(KBBp, [A.6 2
0
NANBd
P-
KKlpp KA +
+
max NANBd
(o,dRo + +d/dan /A/B)
For the orthogonality between noises, by Corollary B.8, we have
o KA dt KA +0 NANBd
Te d kodg,1
Clear that this bound holds for all other combinations. Thus,
so dt SLEPA
d 8 SBLEPASO
( of
NANBd did Fd V/ )LEOA))
B.4 PROOF THE MAIN LEMMA OF STAGE
Proof of Lemma B.I. First, recap the estimations we have derived in previous subsections and introduce some notations. By Lemma 10 and Lemma B.1 we have
d dt s -1((1
mmax NAN (2,dRo + (5A1B)
min
+ O
NANBd 2
(dr? + OA/B d
dt PNIS S -1(1)
PNIS +0
Nd
NANBd
Define p = max (p-//x/s) to be the indicator of the progress we have made. We have
di p < -Q(1)
d
22
min
+ +0
max
(5idemo + d +dv
(11)
NAN
NAN
Let Ru = max,,4 have
|||k|p|
[|[ ll be the condition number at time t. By Lemma B. we
d Ru S 0(1)
a max
NAN N Nd pRo-
(12)
dt
Now we consider the discretization errors, define
8 0A/B= pE(r)-eE(d-r] {PA/B. PB/A.P 2, PA/B.ER + PB/A.G 2].
Note that time t, the first condition of (10) holds with ò replaced by O(6A/B). Meanwhile, by Lemma B31,, we have
A/B
0
0 2
(62
+ (aa/++7.))
(13)
dt
N
NANBd
35

---

Under review as a conference paper at ICLR 2023
Let 6g.n0 (t) be the smallest number such that the second condition of (10) holds at time t. By Lemma B.14, we have
d so
(
of (d + voa/e)
8 dt
(14)
NANBd
Then, define
81 max دT (iLp. p1 Enpks 8.g.K.1 p I
Clear that the last conditions hold time when dAB,L and , are replaced by vor(0) (t). By
Lemma B.15 and Lemma B.16, we have
81 5.50 o NAN
(
pôd +0
o max (NANDd (
(15)
Now, we are ready show that the errors do blow up in Stage 1. Note that for all these ö's, can make them arbitrarily inverse-polyoomially small by choosing sufficiently large m.
First, we consider A Note that the dependence of the RHS of (13) on 8 is quadratic. Hence, by making the initial value of 8 A A/BB B. the RHS can be made be dominated the 7e-related terms. Hence, A S SO NaNnd " As we will see later, Ti = poly(d). Therefore, by choosing a sufficiently small r², we can make A/ B remain small throughout Stage 1.
o²
Then, we consider ,g. As we have argued earlier, the RHS of (14) can be made arbitrarily small, so that 6g,no remains small in Stage 1.
Now, consider the condition number Ru and 81. For 81, by our previous argument, the second term of the RHS of (15) can be merged into the first term, by choosing a sufficiently large m and a sufficiently small 7.. The same is also true for (11). Hence, for these quantities, we have
d dt S 2111
miin NANBd P
Ro S 50(1)
max NANBO Bd
pho:
81 S 0(1) NA P
dt
Hence, by Lemma B.2, we have
Ro S (0) exp
0
② mux
p(0)
81 80) exp 0
o max
p(0)
o
2 min
(d-rjo ram
Note that p0) 0() and PN/s
Therefore,
exp
(o 2
max
P(0)
< exp
(o(1)
max 0(1)
max
1,
(d - r)o (d-r)o²)
1)
< exp
1 (oggd) Vd.
T min
ro T
min
In other words, both Rg and 81 at most grow vd times.
Finally, we derive an upper bound on Ti to complete the proof. Similar to the proof for the con- dition number, one can show that ||[K)ppl* 1 can at most grow vd times in Stage 1, As result, 1/(NANBd) is lower bounded by some 1/ poly(d). Thus, by (11), Ti poly(d).
L
B.5 NEGATIVE RESULTS
Lemma B.17. There exists a E R satisfying the assumptions of Theorem 4.3 such that, at the end of Stage 1, the condition number of KA is 2(1)
Proof. We choose = r and of clog d, =3 = Clear that this satisfies the condition of Theorem 4.3. Note that it suffices to consider the infinite-wide case, since, as we have proved earlier, the finite-width trajectory tracks the infinite-width one. By Lemma B. we have
-
(KAlp KB.) di
d
1+ K NANDd (1+ (KA KBB)) (I-(KA)), (KBp)
4K
of
36

---

Under review as a conference paper at ICLR 2023
By the proof of Lemma 1.1 we have
4K
of
(KaAlp. KBlP)
(K KB)
Pp/9 ~
Pp/9
1+ K NA
NN NANBC
4K
og
(IKAJO (BAA)
(K A KB)
1 NANBd
NANBd
Pp/e
Note that, in the wnfihite-width case, we have
KAli. KBBI)
2 (Al22 KBl2) =
K KAld: KBld
Therefore,
(Klpp KBY)
(KA.Ka) NaNnd
2 Oa and
[KA KBl9)
(KA.Kg) NaNnd S 0 for any
q 2 2. Hence,
4K
of
AK KB)
P1/2
1 K NAN
(KKa)I,(K)) (n))
NANBd
P1/2
4K
of
1-
x² ||=]|/
(KKa);(Kaa))
(d (d- 1)n
( A)2, (K)))) P1/2-
1 K NAN
||a1|
For notational simplicity, define X1 =1
X2 = -KKAla (BB22) Y
P1/2. A
4K
Then we have
NaNgd
X2 -2AX2.
(-of A) and
First, by Gronwall's lemma, have X1(T) I exp
T
X²(T) 2 exp
-2 0
A
2 x2
2/ (T).
As result, when X, reaches 1/d, we have X₂ 211). Let Ti be the time X₃ reaches 1 1/d and T2
the time (22) = (2(2))/2. On [Ti, T2), we have
41 0())oAAp12
By Gronwall's lemma, in order for X2 to half, we need exp(-2 A) = 1/22. Hence,
P1/2(T2) 2 P1/2(Ti) exp (a(1)of (")] A > 2(( = 2()
(7)
C STAGE 2
In this section, we show that, throughout Stage 2, the discretization error and the noise-signal ratio still remain small, and, at the end of Stage 2, the condition number is close to 1. Formally, prove the following.
Lemma C.1 (Stage 2). Suppose that the beginning of Stage 2, have Ko S Vd and all errors mentioned in (16) are sufficiently smalf. Let Crarget > be a constant. Let T₁ be the earliest time that
K } S Crarget: Vp. E [r].
K Allll
'By Lemma B. this condition indeed holds.
37

---

Under review as a conference paper at ICLR 2023
We have Ti poll(dd) Moreover, throughout Stage 2, we have
K ppll
max
1 (Kllp) KBLp)
S 8_
Vp E [r],
|| Bln 1|
|| Aspll
max
S 6_
Vp. E [d-rl} r],
[[[ B.S.pll
[|[KA, &llll
[Kclelll
max
E,, (A,
S dN/S3
ap 1pell. [d - -r],
max
[Kclps (DDla)
C B] s S.AB...
Ap # q [ [r],
[Kooo (KD.).)) < Ka.slo (B.SSu) : (C.DE dc. Ap g. ( -r,
max
(16)
where the ö's are some small 1/ poly(d) values.
The rest of this section organized as follows. We derive estimations for the Q-matrices Sec- tion C.I. In Section C.2, we maintain the last two conditions of (16). In Section C.3, we handle the first two conditions of (16). In Section C.4, we deal with the noise-signal ratio. We estimate the convergence rate in Section C.5, Finally, we prove Lemma C.1 in Section C.6.
C.I. ESTIMATIONS FOR Q
As in Stage 1, we estimate the Q-matrices in this subsection. The analysis here will be more com- plicated than the one in Section B.1 since o? is longer close 0, and we cannot simply approximate S and SB with (1 + K)-1 However, the idea is still fairly straightforward. We split all terms into the infinite-width part and the discretization error part. Then we Taylor expand the corresponding function around the hhintte-width part to factor out the first-order error terms. Then, we evaluate and simplify these first-order terms with Eg- and Eg+.
First, we need the following lemma which gives closed-form formulas for some expectations we will encounter later.
/ E [r]. For any NANB
Lemma C.2. Define (2t = xz A2== == Ek RE 2x and Tp = tanh
p * 9 E [r], we have
E
exp
(x
AA2
II
cosh
R2//d NANB
=: Zo.
N NANB
k=1
(zt 2 km
E
exp
2/p
ZoT pp
NANB
E
exp
(z_,=~)
22
2p =
Zo ToT
NANB
Then, we derive estimations for SA and Sp. There are two types of errors we need to consider. The first one comes from the noises and the second one from the non-diagonalness of K J K B. Similar to Lemma B.4, we deal with them separately. The next lemma handles the first type of error.
Lemma C.3 (Estimations for S). Define
(KAz*, KAZ
R.,- == exp ( (2t, A2
E4,- = exp
exp
NANB
NANB KA.EE KB diag( (elkgg(TTkkeee) (TT)EE()))**
KAzt, KBBEEB)
84.5
8E+.T+
NANB
NANB
ZA = E E+.,, ZB. : E E_Et,
E+,
E+.+
É+.+
SA = E., KZAA
S = E+, + KKB.''
$:= E.. K
38

---

Under review as a conference paper at ICLR 2023
In Stage 2, we have
SA SA + S(1 - 5) (ô+.e+ + +g+t.t de+,T+) (d8 (ōAR, +dNYs) de, ÖN/S):
SB 55 + S(1 - S) (6+.et + - dT+.++) + O (d (ÕAB, + +N/s)) gg, dN/5),
-
and
S SA A expl ff fB). SAL+I= + SE+- (f+.E-
+ - -5 + bct.+) - 5)0E+,T+)
exp(fĀ 53)
E+it
E+.+
(d) (dAB.1 + +dv/s) 6g.LŌNIs)
SB( FB) exp(fA fā fB)
F + St_,t (1+ 5g- + - (ó+.st + +o+++t) - (1 55)5++s+)
exp(fA fB)
E+,
E+i
= 0 (SAB.L dN/s) 1g.15
Then, we consider the error comes from the non-diagonalness of K K B. Lemma C.4 (Further estimations for S). Define
(Kak, [Kb))
(K [K [BBb) 2Bb)2&T,,2 TT
8 8+.- Σ
27 21. 2..4=-
Σ
NA
NA
i71
itj
||Ai|1
E. = = BB, =exp
NANB
In Stage 2, we have
SA(*t)= 8 1 (1 -0) 8+.I+) + (d
SD(=t)= $ 1 (1 (5) 8T+,t)) + +0 (1P ( 1B
SAE+.- SE+-- (1- S5+,t - (1 5)0+TT+ + _,,-) + 0 +
E+.- En
5gE-.
SE-+ (1- $8+,+ - -1- 5)8++ + +8_..) (f AB
Exti
En
With the above two lemmas in hand, we can now derive estimations for the Q-matrices. Lemma C,5 (Estimations for Q1). Define K KAB K A KK and KBA K A. In Stage 2, for any + E [r]. we have
[Q1p..r = 2(1 5)(1 Tp) + (d AB,
K KAnppa + [K n.
[Q1]gg = -Š(1 -)
(2 (2- TT- Ta) TT)
NANBd
KABBp. 3 (2T Tp - Tg))
- -(1-5) 5)
NANBd
K AABLR
-1-
(2 2 S(T, + +T)) (1 - () + -12) + (odddo (d&
N NAN
In particular, we have
KOOAB.I <0
SAB. Vd
Q Qu S0
d
Note that the diagonal term has a zero-order term, i.e., is not proportional to some error. That the signal term. On the other hand, all off-diagonal terms depend on [K KABBp (p * q). Recall that the dynamics of KA and KB can be described using these Q-matrices. Therefore, for any p * g. can have equations of form dr [KA R G( K (ABlp.. (A,,),) where G is some complicated matrix. By carefully analyzing G. we can then derive bounds for the off-diagonal entries using Gronwall's lemma.
Similar things also hold for Q1.s and Q2. The difference here is that for these two matrices, we do not have signal terms.
39

---

Under review as a conference paper at ICLR 2023
Lemma C.6 (Estimations for Q1,)). In Stage 2, for any p E [d and q E [r], have
Blg [ (A + + (d (ŌAB. + dd/s)) (á/1)nss)
Queaale -(1 -11 S) (1 (1 S)T;
NAN
Alg K kB =0 (d (6AB.L + ōx/s) 1ó/s))
[Qı.alea = $) (1 - 11 S (1- (1-
NANBd
In particular, we have
max (1Q1.5alp.: QQis I 10 (6NN/555,_))
Lemma C.7 (Estimations for Q2). In Stage for any p, E (d r}, we have
Q22p.. = + (d (6AB.1 + +N/S) ōx1s)
Lemma C.8 (Estimations for Qo). In Stage 2, we have
Qo =
Q1IK. + O O88 + 8_ + kndōw/soe,
Finally, we use these estimations to simplify the formulas for the norms. We do not consider the tangent movement here since the situation is trickier there, and we will handle them in later subsec- tions.
Corollary C.9 (Dynamics of the norms). In Stage 2, we have
20
[|[K
Alll
B|]ll
- [|[ 1pl N2d d
[K <Appll==
(Kalp:
KBbp)
(Q1p. +2
Q0o
NAd
of NANBd Kodo AB,IL
d dt
2 ||[[ A
( [f dōc
(
||
Qook +
N2d
NANBd
The formulas for [K Blpl and
[| B, can be obtained by interchanging the roles of A and B.
OMITTED PROOFS OF THIS SUBSECTION
Proof of Lemma C.2. First, we compute
(2t,2 2 )2 } II
R 2 NANB
R2/d
(-
R2//d NANB
E
exp
exp
exp
+ exp
NANB
2
NANB
F
k2//d
cosh
NANB
k=
Similarly, we also have
exp
(2t,2
2p
E
cxp
22
2p
II
E
cxp
RZ 2n NANB
NA
N NA
Zp
kop
Each factor in 1Rx is still cosh
#2 NaNB
For the first term, we have
R 2p NANB
R22t/Va
RR-t /Va
E
exp
pp
exp
E
exp
2vd ap
NA
2vd of
NANB
R2 /Vd
sgn 2 Vd
R2./d
12//d
sinh
sinh
sinh
Vd
NANB
NANB
NANB
Therefore,
(zt x
k NANB
R
E
NANB
R2 sinh
2/
2
cosh
[NAN
T.TP
T=Tp
kpp
The above calculation, mutatis mutandis, also yields the last identity.
40

---

Under review as a conference paper at ICLR 2023
Proof of Lemma C.3. First, write
(KAz + KAÆEA KB= + K KBBEB))
ff
NANB
(Kazt, 3 NANB
(Aaz KB.EEB)
(KA&EA Kg²
+1 18yss)
+ +0
NANB
NANB
Also note that the middle terms are (g ÓN//S) Then, we compute
exp (Sz (B) = exp (1 1 (Kaz*. KB
(KA KB²
+ (185g, (x1s)))))))
NANB
NANB
Similar results also hold for other combinations of +. With the notations defined in this lemma, can write these results as
exp (ft fn) = E+,- 64.5 + -
8 (
exp (J fB) = E_,+ (1 + - + --) + + (P0g, 8x1s).
exp (ft fB) = E+,+ (1 + S+.e+ + )c+.+) + 0 ( 1811))
To compute SA and SB, we then need take expectations over the negative examples. Note that by taking expectation over &B. the term E,. ô+.e- becomes 0. Therefore, we have
E exp (fz fB) = E [E+.- (1+ de+ + (o (166 1
Unfortunately, the same argument does not apply to de+.- since both F,,-- and de+.- depend on However, is still possible to further simplify the expression. First, we write
E [E.. - { (1 O(dō (1) (
E -5. det (0 18//s)
Recall Lemma C.2. Then, compute
KA KB KBEE- [Bt--=1) x-
KAGEt. Kn=
E E+. 6c --.-- E E
NANB
NANB
x
(K K B (iag((T))ee))))= Z.de+.T+
Zo
NANB
Hence,
E exp (ft fn) = E (E,, +, } E (E 8c+, 10 =0(
83/s)
ZA, + Zede+.T+ + Ocd 8g, (N/S) + O
ZA. + Zede+.I+ - =0( (dAB.1 + dN/S) 85.1
Similarly, we also have
exp (ft fā) = BAC + Z + ( (dΛB, + SN/s) 6g 5N/S)
A
Recall that exp (ft 11) E+.+ (1 8 + be+ + ( (d55_187/3)) Hence, have
E+ 8+ + 0g+.+)
(A
+ (d3 (DAB.L + ŌN/S) de11))
E4,4 84.e + 8et,)) KZ 0. + KZ.det.T
=3A 1SA((044 (c+.+) (1 -A)6
+ + + +)t+.)) + 0 (d (AAB, + 88/S) 8g, 16155)
5A + 5(1 5) (6+.et + + d5+.T4) + (d8 (DAB.L dN/s) d 8g.18N/1))
41

---

Under review as a conference paper at ICLR 2023
Similarly, also have
SB( a A zB) SB+
(6+.s+ + + - +++++++=+000 (BAB.1+0 +NN/s) (g.1))
Then, compute
S (a i / (B) exp (f SB) = (S^ + +1111 - (6+ + g BeH.T))) E+.- (1 (1+00.6- g- + 8gt.-)
exp(FA fB)
E+st (1 + + + detit)
#0 (d (ÉAB,L + (N/s)) 6g.1
SAE+.- (1 (1+ + + -S (6+.et + )g+,t) -))6++T+)
E+it
+0 (d3 (iAB,L + bN/s) 6g,15
SAE+. + SE-,- (5+.8- + 33 (ō+.e+ + +etttt) - (1- -))
Ext, E
+ (d (SAB,1 + 8x/s) (g,15815)
and
SB( 2 B) exp(ff fB)
SBE_,t + SE-,t (1+ + + (0+.E+ + (g+.+) (1- 5)8T+
exp( x f
E++t
(df (BAB.1 + +8N/s) 18/s))
Proof of Lemma C.4. We write
K xt KB²
([Kle,, KBlk)
(Ka), (KBli) NANB
2k + Σ
20 =: =:: + +J+.-.
NANB
NANB
inl
To- In other words, I+++ and I_.- do not depend on the actual value of =²_ Also note that I,,- is bounded by A(d)AB.)). Then, we compute
Note that, special case, we have L+,+ = I_,-
E+.- cxp((+.-) (1+ (1 + (d' and E+. exp(To) (1 54,+ + +0 (d
Take expectation over = and we obtain
E E+, E exp((++-) E. = =
(oxp(T+,))+,.- + (
By Lemma C.2, the first term is Ze. For the second term, we compute
KABBis E {exp(1+.-)=5) z
exp(( (xxp(4+-))8e (8)
Σ
NANB
KABla
NANB
-2,4.7+
where the second line again comes from Lemma C.2. Hence, we have
E E+,- Z. Z. + O (d6 ABBL))
Then, for SA. we have
exp(lo) (1+5+.t)
Fx
(x(zt)= = E+,+ + K Eg. Et.- exp(To) (1 5+,,) + KK7 + KZ.
= - (1-58..+ (1- - -S)0...T
55,,+ + +0 (dB5)
= (1 -5) (+.44) = (df
42

---

Under review as a conference paper at ICLR 2023
Similarly, also have
SB(z*) = 011 (1+(1-5)( 81+1t) = = 8A.
Then, compute
exp(L+,-) 1 [
SAB+.- = 5 5(115(111111 (1+ (1+ - E+,+
3+, -
= (d
exp(L+,-) 1 11
eexp((+.-) (1-55
S5+,+ - - (1- 5)0+ + + + 00 (d (BBL)
exp(lo)
Similarly, also have
FBE (exp(
(1-58+, S8+,+ -1- + +8,+) +0 0 AB,
E+++ exp(1o)
Proof of Lemma C.5. Recall that
Q1: 11: -E{(2 (22 SA( BB) (B m²²) 2 (2t)
SA a 1 2 exp(fx fB) 2 (z+)) a
KE SB( 2B) exp(pxx fB) x (x
KE
exp(fA fB)
exp(f* fB)
Note that there is & here other than the ones in the coefficient. As result, all terms contain 8+.6 and alike are O. Hence, we have
Q = R ((
SA(M - (1s TB)) = =t(2+)
SAA+.-
===x)== 2 d
(BB=+tz*(27 (x
K E
K E
E+it
E+s+
(Q(1) + +
Now estimate each of these three terms separately. We also deal with the diagonal and off- diagonal terms separately. By Lemma C.4, we have
((Qi)))))) E (2 5 ; (1)
=E E {2 $ (1 11-11)))) (1- ),.T+) 3 (1+ (1+(1-5)(+... 3T+h))) 1 O AB..)
= 2( - 5) = 0 (d28A
Note that we use the fact that all these 8's have mean 0. Also by Lemma C.4, we have
{
(SE,-
(1 (1-55, 55,, - (1 5y + ób..)) &p x + (d282A.))
T2((Qinna) -K
Eo
ŠK
E. B 27
Eo
ŠK
{B
(11) + 2p x 0 (185
E.
Note that 4, __._ and 54,T+ only contain cross terms of form 2, + with + j. Hence, the second term is O. Meanwhile, by Lemma C.2, the first term is
ŠK
SKZ.T, E.
E D. Ex. 25 4
-(1 -SYT,
E.
As a result,
T2 T2((Qi)o.p) = --11-5)7, + + (1852 ABBL))
43

---

Under review as a conference paper at ICLR 2023
Similarly, one can show that T3([Q1ppp) = -(11 -1-1 Sy1, = (B8A also holds. Thus,
[Q1p. = 2 -)) - ( AB.
Now, we consider the off-diagonal terms. For notational simplicity, we define kAB = K2KB. For any p p q, we compute
T=((Qlpa)== = - ( A xT)) + (B xBB) 3/ 2 #
=-1(1 - S) EE { (26 - 3+,r+ - 81+3 2/ 2 d 0 (A
[KAB K (2 T. T2) B
= =-5(1-5) - 5)
Σ
2 = 0 (d 82 AB.L).
NANB
Clear that the summand is nonzero only if (i.j) = (p, or (i. = (q. Hence,
Ti ((( - [ABB p.g + [AB 4.8 (2-Tp TT) = =0 (d (B..)..
NA
Then, for T2, we compute
SA(D)) A. B exp(fx -f)
T2 T2((Q1ln.)) = -K K
+
exp(Fx .f1)
2p 2g
SÉ++- (1- 58 55+.+ (1- 11)58) + 8. (b..)) 2/ 2 d + (d AB.L)
-E
En
- K E {
[SE.. (-$b.,+ (1 (1-5) -S-)8=++++ + +b+.-) 3p 2
+ 0 0 (12332 (B..)
Eu
KABlia E
{(5E)-(4 (-5xt 2 - - -(1-5)2 2 + x) +
K
3/ t
NANB
Eo
0 (R.)...))
Again, the summand is nonzero only if (i,j) (p, q) (i,j) (q,p). By Lemma C.2, we have
Te((Qu
SK [KABPp
(1-5 E (exp((+,))2 + E
2y zgd}
Eo
NANB
ŠK KABleR
((-3- (1 (1- - Byr;) E (exp((44.) xp + d [(ep(((+.-)) +0 # (d 8AB.1)
E. NANB
ŠKZ.
KABBpą
((-5- Syr,) T2 + TTTT))
E.
NANBd
ŠKZ. g.P. ((-5 (1 - SyT,, Tp + 1) + (d2
Eo
NANBd
KAB (T, - 1) - (1 - 5) KK AB [KB
=-
(-5T, - (1 (1- -YTT +1) ( (...).
NANBd
NANBd
44

---

Under review as a conference paper at ICLR 2023
Similarly, for T3, we have
Ts( pp9g))
SB T exp(FA ft)
p Rq d
exp(Sī fiB
SS_.+ (1-56 - - ((1-5)8r++4 + + zy ] + + (d (dB))
= =-KE E
Eo
5K
KABLiS E. B ii ((-3- - - -1--)(1))) S)T. z] + - 2/))
 + +0 (1 AB.))
Eo
NANB
ŠK [KABPPa
(-5- ZcTa/d ZcTpTa/d)
Eo NANB
SK KABBER
((33 (1 -1-5)1a) ZoT /d + Zo/d) = =0(d8 8 (AB,1)
Eo NANB
K KABIp&
(T, 1) - (1--)-1-11 5) K AB KAB 4.
--1 3) 3)
ST - (1 - -5yr,) + O ABBL)
NAd
N NA Bdd
Combine these together and we obtain
KAB p.g + KAB 4.P (2 -2-22 Tp TT))
QQ1pp. SS(1 5)
NA
K KABlu. 3 (2ITTg - - Tg)
- -11 5)
NANBd
K KAnlan aR
(1 5)
(2 S(T, + +Tg) (1 -TT TT)) + +0 0 (d AB.L)
N NAdBd
Proof of Lemma C.6. We write
Q1. =
: { SA(a aB) Sn a))) + d}
KKE SA A. t exp(f f
A xT) exp(fa ft) ( exp(St fū)
(= d
KE
exp(Sx fū)
T; (Q1.en) + To (Qi,s) + T3 (Q1,8))
We will use the fact that if some quantity X does not depend on E, then [E simplify these terms. For T1, by Lemma C.3, we have
Tr ((Q1,ell.g) = - E. (S x + Sa xB) 15t) d}
=-5(1 $)E E (255..+ + 25 det,T+ 8r+.et) (EBlp=t
+0 O (d (DAB.L + +5N1s) 6x//)
5(1 - 5)E E ( (20+..+ - -5 15 + (d8 (jAB.L + dx11)) )g.18)15)
Recall that
K Ka,kes NANB
K K 1E
8
2, LEb,. 81+.e
Σ
NANB
i.j
Hence, we can further rewrite Ti ([Q1,epla))
K KKKKll (2- (2-T)) ((ttEtt))))
Ti ((Q1.eolp)) = - -(
Σ i.j
NANB
- 0 (DAB.L + &N/s) ò d21s))
45

---

Under review as a conference paper at ICLR 2023
Note that the summand is nonzero only if q and p. Thus,
Ti ((Qu,))la)) S((1 -) (2- Tn) [[ [K (B&))) + (d (SAB.1 + +8N1s) de, 12/s))
NANBd
Similarly, also have
KŠ
E { 0.5 d + (de (SAB,1 + (N/s)) 6g.15N15)
T2 ((Q1.5mlp.g)
F.
KS
T3 ((Q1,5))..))
B. -i. (o= - S8 E+ (1-S)ōr+ee) (5Bb=t
E
- (d (bAB,L + +NN7s) óg.
Then, we write
K A K K.
E BB 8+. 2 Σ
E (E. [ Enl
NANB
1.j
If j # p. clear that that summand is O. If 9, then by flipping the sign of = simultaneously, flip the sign of E+,- 2t z.. Therefore, when + g, the summand is also O. Thus,
5K2. [KAlo- [K (REW) + (d8 (SAB, + +NN1S) dg, 8N/9) T2 ((Qi,el)a)
E.
NANBd
([KAlg Alg: K K.ele) (18 (6AB.1 + dx/s) (515)))))
-1
N NAN
Similarly, for T3, we compute
E
B E E_,+ (6- - -56. - -60.5.-1111
(1- (6Elpztd
K KA,ess E. (B-i,t (=](5T) - Szz (EBli -11 -1- (T)) (5)l)))))))) (5BIp=
NANB
KK ( (zj (17- - -52t - (1 - (1-5)(1= TT:);) 2
5zz
NANB
[ [ K.B (Š(I, - 1- 112) + T
= Ze
NAN
Thus,
T3 ((Q1salpa) -(1-)) K Klg [K
-1))) +
NANBd
- # (d (DAB.1 ÖN/s) 1g.18)1s)
Combine these together, rearrange terms and we obtain
Aly [K KB + 0 (d (BAB.) +N/s) 1x/)))
[Qi,gnln - (1+ 5 5)T;
NAN
To obtain the formula for Q1,ea suffices to interchange the roles of A and B.
Proof of Lemma C.7. Recall that
Q2 := E ((2 SA x SB =b)) ( T d} -KE
SA A. x cxp(fx f3)
SB(x $B xt)) exp(fx ft) (T))) } exp(fx ftt
-KE
( d
exp f f
=Q2)(Q2)) + T2 (2))2) + (Q)).)
46

---

Under review as a conference paper at ICLR 2023
We now estimate each of these three terms. Again, the strategy is to leverage the symmetry of the distribution & to argue that some part of the expectation is 0. For T1, we have
Tı (1Q22)))))) S(1 5) { + 2 6g+.T+ Ör+ (Bbll() - =0 (d (DAB,L + SN/S) 6g.15N/5)
Note that none these 0's depends on both E and En- Therefore, the first term is O. Similarly, for T2, have
(2((Q2))))) = KS
E
(B_.-(++.-- + - (d+ + )c+.+) -( (1- -3)55-14)) (5bb15111d
E
O (6AR.1 + ÓN/S) 0g.10N/S)
= +0 =0 (d (SAB, + +ŌN/S) dc, ÖN1S).
The same is also true for T3. Thus,
Q21pa = +0 (d8 (DAB,) + SN/S) 6g 1x/s)
Proof of Lemma C.8. Recall from Corollary A.4 that Qo is defined as
Q0:=-
E
{(2
SA A x Sn =B)) (ft fb))
2x #B
SA + A: exp( fm)
+ K
E
ff)
eXp(T? f fT)
SB A: x exp(T f
+ E
K
(fx.5D)
eXp(Ti fD)
We write
(Kazt, KB= 0 (Kodō
(ft
NANB
KAAKS
Σ
2, 2 + (
NANB
Therefore, have
K KABBis
Qo =-
[Q ki + (NoddN/sOE..)
NA B
K2
Q1kkA = O (68 - d_ + kodōN//oe..).
|
Proof of Corollary C.9. Recall from Lemma A.5 that
d di |||a|p1
2 ( (K Alps [K BBBBB1 Q1lp))
(KKapps [ Q1,colp)
of
NA Ndd 07 + N3
Qoo
NANnd
Σn T
d Tt.
[[Kabbll))
By Lemma C.5, we have
d di
KAlp [B
Ti
K
=
(Kapp [Ka
IQ 1Qi
NAN
1=1
20 [[[ Ablll [BBpll]
[Alpp (Blp) [Qilp.p
NAd
+
N NANBd Kodo2 A
AB,
47

---

Under review as a conference paper at ICLR 2023
By Lemma C.6, we have
d
Alll [[
T2
=
(Kapp, (B.e)t)
[Q Q
dt
NA
k=1
0
p
do
N/SO E.L
NANBd
Combine these together and we obtain
d dt
207
[|[K|]ll Aplll
(Q + 1221[Ka]] 1pll
|
(Kapp (KBp)
N2d d Qoo
NAd
o NANBd
Nodo
AB.L
To get the formula for d the noises, we write
([K)pllf' it suffices to interchange the roles of A and B. Similarly, for
at
2
{ (K..g (ia
2 2 2 (2.e
2 K Qoo3
K eelll
og
dt
NANBd
NANBd
og Nº d
dt (Ka.llll
d
T2
By Lemma C.6, we have
( d dt
K KA KB)akl NANBd
Ti
|||
= Σ k=1
KA KBlk) 1Qi.aaludo
K AA
d82 (
NANBd
By Lemma C.7, we have
d Ti
K A. 11llll] [Bekkl (KaAch (Knela) 1Q2lk
K
dt
NA
k=1
|[| dt (dAB.1 + oN|s) +
NANnd
Combine these together and we get
d |K ||
2 21[[ 2|[[1l
Q003 =0 ( K KA.S/ul (
dō2
de
N2d
NANBd
C.2 MAINTAINING THE ORTHOGONALITY
In this subsection, we control the growth of ÓAB.L and & First, we derive the equations that govern the evolution of the off--iagonal terms.
48

---

Under review as a conference paper at ICLR 2023
Lemma C.10. In Stage 2. for any p + 9 [r], we have
Kig/ 0p + Rig
dt (KK alp [Kn)s) [Q1)a.p
NANBd
QQ1po NANnd
K KAlp (KBl))
NaNnd
1
+
KAlps (Kag)
(KAl (Kol))
NANBd
NANBd
Pmax
+ 0
NANBd dRodAB, AB, (Vo= (55ttAn.)))
Lemma C.11.
In Stage 2. for any p + 9 [r], we have
Kg/ O
K
KAlp KATs) = =QQilq
+ +1Qile +QileaaNNe
di
N Bd
NAN
(Q NANBd
QQ0y NANBd
KAlps KAle)
Qilag NANBd +
Qilpoop
KAlp [Blg)
KBlp. Kalg
NANBd
of
NAN drog ABBL (Vo_+ +6B.))
-
and
Kg// K
Kpp o2/
K Blpp
K KB)g)
= 1Q1).
+ 1Q
dt
NBd
N NANBd
Q1Ip.op
QQian NANBd
KBIp KBla)
NANBd
QQinn
Q 1
KB KAlg
+
(Bla)
NANBd
NANBd
maxx Bd
(Vo- 35_+8ABD +DAB.))
+
AB.L
Now, we ready to control the off-diagonalness. The proof is similar to the one of Lemma B. To provide intuitions, we first consider an idealized case. That is, we assume here the noise-signal ratio is 0 and KA B. and explain how to control (iKAl:: (B)a) > for p # q E [r] under these assumptions. In this case, we have
d di
1/ppop0 + tkppngg NANBd
Alp KBly)
[Qilgæ
and
K ABBpg N NANBd
QQip. ~ -11 5) (2- Tp - T + +5(2 - Tp - -2,)²)
Note that the coefficient negative. As a result,
K Alpp K (BB,) will move towards
When KA- KB is not exactly true, the situation is trickier because need to deal with the middle four terms of the RHS of 4 ( K Alp K KB The idea is to view all these off-diagonal de errors as a whole and show that their sum is non-increasing, up to some higher-order terms.
Lemma C.12 (Orthogonality between signals). For any + q, define
.p.0 == (KK lpp, (Kal)) + (KKlm [KAl,)² (KAn)) +
( (KKA [K (ala) + (KKal, (Knl)) (
(
49

---

Under review as a conference paper at ICLR 2023
In Stage 2, we have
d 8 pa S
NANBd no (d + VAB. (
o
dt
Note that the LHS is of order 8% AB, and for the RHS, the only term that can potentially have the same order is the 6- (0-VoL 01p.0 related term. We will show later that d o(ŌAB,1) whence this also a higher order term.
Then, consider the orthogonality between signals and noises. Lemma C.13 (Orthogonality between signals and noises). For any E [r] and q E [d - r), define
[KBp, KA 2
6g.1...4 = KAlpp KA
KAlps (B +
[Kpp (B
In Stage 2, we have
86 di
d p.g S O
L.g
NANBA droge (6E.1 + v=-))
O muax
Finally, we deal with the orthogonality between the noises.
Lemma C.14 (Orthogonality between noises). In Stage for any p, E r), have
o
( (K..ll =1 +0
do
NAN Bd
OMITTED PROOF OF THIS SUBSECTION
Proof of Lemma C.10. Recall that
d
I- [Ka]p (KAlp) ')
[
KB.
o
KAlp
di
KAppl
NANBA of NAN
I KBIy (Kalo) ")
K [
QIA.Q1.EAll [[[ nlall
KB =
di
[[[B]]]l
Then, we write
[KApp T
I- KB]y [ (KBla) 1)
1 o²
KAlps dt
K
K BBg
K Blgll NA
I- KK[ (Knl.))
[KA [
o² NAN
KApp
BBggl
For the second term, by Lemma C.6, we have
KAh (I- KB] [KB))) )
K IKA.Q1.EAl
[ n/ll
d-r
||| ||[K,.]la|| (Qi.sa |
KAlp, KAskk)
(KAlp, KBly)
KBlg KAcle)
0 [d82 3Rv1s)
Hence,
K [nl, = =(Kr (I (I- KKK (Kal)))
K KAQTI
K Allp: dt
K
N NANBd Bd
o2 max
+ 0
d52 152_58315
NAN
50

---

Under review as a conference paper at ICLR 2023
For the first term, we have
(KBla) ))
K KAQIl4
KBly (Bly
--
K Bl
|| 1Q1ggh
Σ
KAlP KAlk)
(Klpp (BBl))
[K (Alk)
K|
When & $ {p. we have and Hence,
Kalp KAJa
S O(dAB.L).
[KB [ KAlk)
S ((0ABBL).
||| KA]all/ [|] S No. Meanwhile, by Lemma C.5, we also have |[Q1] 11) O(5AB,1).
Σ kE(p.9)
( (Klpp (KAl)
(
(KAp
KBls)
(K (ATA))
[| (KAK la 1Q1lg.k
K Blull
S 0 (dRo
Therefore,
d KAlp: dt
o NANBd
KBlg)
[|]]ll 1Q11.
(K [KAlm KB(I)
(KBlor ( KAln))
[|[K
([KAlp: KAlo)
-(KAlp .Blaa
(K Kale))
||| Algll [Q1 g.g K
#0 (dko0OAB, # (áōž, 13y1))
For the first term, we have
(KAlm (Kn).) (KAa))(KK)))) (
(Kalny) (Kab)) (Kanl:(Kalb))
||a||ll [Q1]ap
[[BB]]]]
Kpl KppQ21ep (1+ v=_)
= (1 S' 1ABB1))
Rig
Kop Qulae + (HOSAB.LVO_)
+ #0 (knd AB.L).
Kiq
For the second term, have
(( (KApps KAla)
KKApp (B)g)
[KB)g (KAl))
[Ka]|l| [Q
|| [|KBlgl
(KApp (Kala)
(Kalp: (KBy))
(1+ (Qula. (1 + vō_)
KAIps (KAla)
KK Alp (KB.)
[Q11a + 20000 (oLV8))
Thus,
o2
K KAlps dt
KB)g]
NANBd
Rip
Q21ap + (KODAB_LVO + (KO0)AB.L)
Kig
((KAlm (KAl.) (KKlp (KAp))
(Qihan + (01Vo_)
+
+ (dd1000AB... +0 (do2_LO%713
Kpl xplQilae + ((1 [KATp (KAla)
(Kal: K (Rla)) [Q1
Kig
+ 0 (dRodAB. (V0_+6AB.1)
51

---

Under review as a conference paper at ICLR 2023
Hence,
Kpl 1plQ1
of
+ ((KAln (Klo)) - -(KA), (Kll))
Q 1
K AAlm di
K kBl)
Kiq
NAN
N NBd
+ 0
max
dRoōAB,I
[Vō_+5AB,)
NANBd
Interchange the roles of A, B, P. 9. replace Q1 with Qi, and we obtain
K
KApp
K KBlg
Rig
of
+ ( [KBlp. (Bla)
(K (Klpp
KBlg))
Quip.
dt
Kip NANBd
NANBd
+ T
4 (V=+dAB..l) Combine these together and we complete the proof.
NaNnd dRodAB.
Proof of Lemma C.III Similar to the proof of the previous lemma, compute
K KAlm
K All))
AA T
I
(KA)))
K KBQQle
K KB..Qi
of
dl
[[
[[(a)
NAN
of NANBd
= KAlq
I
Kalp
(AAp)
K Allll
NAN Bd
max
dog
N/S
Again, we have
(KAlp) 1)
K B 1u1
KAle
- [KAl
K 1
[[KBBpl [11p.
-[al
(KAl))
[K
||(AIpl
[[[KB]all [Q1)..
KAlg
I (I- KAlp
(Klp)) [I-KAlp
KBiu
||[KAppl
Σ KAlg (I-(KAl,, KKlp (KAlo)) )) [k)]
[[K]] [Q
K
Apll
ke (p-a)
Kig 1211e
((Kllpr KAli)
(KKlpp KAl:)) 1Q11p + (dRoōAB.L (Vo_ + +
Therefore,
KAlp KAlg)
=Q11a. N NAN
Kg/KKp Kpp
+ ( ([KBlp, Kals) - (Kalp: (ATo))
[Q 1Qilp
dt
NANBd
+ ((BB+6AB..))
NANBd dr
Interchange the roles of p, and obtain
K [KAp, di [Alp)
d =1Q N NANBd
Kp/Big
((KAlor [K (Kalp (Kalo)
Q11a
NANBd
2
dreg ARQdAB,I 1/_+5AB..1
+
NANBd
52

---

Under review as a conference paper at ICLR 2023
Combine these together and we get
Kg/ 0 N Bd
Kpp o N NA
K Alps
KAla)
(Qila
+Qi1p
dt
QiIp. NANBd
Q Q N NAN LANBd
KAlp (Als)
QQ1
Q Qilnn
KAlps [Bl
KBlp Kaly)
NANBd
NANBd
-
drgp A AB,
NaNnd
Interchange the roles of A, B, replace
Q with Q and obtain the formula for
d de
K K
Proof of Lemma C.12. First, we consider the (Qulp.e-rllaeed related terms, by Lemma C.5, we have
Q1Ip. 5(1 5) K ABB.g + |K BAlg (2 TTpp
Tp TTg)
NANBd
KABBP S (2I,Tq - Tp Tg)
(1 $)
NANBd
K KABLUP U.P
(1 5)
(2- S(T, + +Tg) - (1 - (T + +T2)) + +0 (d& 8AB.L)
NAd
0 K
NANBd
Hence,
K NANBd
o3 =0 o VV
(
Qulba
i ABBLL
AB
NANBd Nd
The same bound also hold for other (Qilp--related related terms. The important thing here is that we have and additional I/vd factor.
Now, ready prove the result. For notational simplicity, define Z1 [KAlp: K K
7/ K KBIp K KAls). Zo KA (Ala)). and Za
KBBp
K KB)g] Also define Gp
[Q /N and similarly for Gg. Then, we can write the results of Lemma C.10 and Lemma C. 1.11u as
Z1
Gp
Gg
G
Gp
Z1
d
Zz
0
Gp
Gg
G
Gg
7z
dt
Za
G
Gp
G p
Go
0
Z3
_4
G
G
-Gp
G
Z4
1
22 muax
(vo- (V0_+8AR.))
+
VANANBd
AB.
+
NAN
drgg AB.
The eigenvalues of the first matrix -2Gp 2G -2G, 2G, and 0. For the first three eigen- values, note that
1Q 1Q
maax 2 Vā NAN
Gp = NANBd
Hence, the signals will dominate the noises, in particular, the first term on the second line, and push |Z] toward O. Now consider the eigen-pair (0, (1, 1,1,,,)), 1)), for which will the actual
53

---

Under review as a conference paper at ICLR 2023
form of (Qilp.g we obtained in Lemma C.5. We have
d
Kg Ro + Kp/ NANBd
Σ Z4 2(Q (Q11a.) dt i=1
mmx NANBd A [V0.+5BB.1)
+
By Lemma C.5, we have
[Q1pp. p.g + +Q11..
- K KABlpg + [KBApp (2 (2-Tz- T2 + 5 Tn TT =0 ( AB.L)
NANBd
Then, we write
[KABbe + KBA = ||[Ka|lll| ||| 21 + [|
KAlpl Zz
= = + tin
+ Vo))
Hence,
2 2-Tr T2 T 55( - T, Tzg² Qu1pa + 1 = (1-5) Ki (Z) Z2)
NA Bd
+ 0 6 + +
To convert Z1 + 7-2 to Ei=1 Z, Note that we have
+ Z2 Zz Zo = ( [KApp - [K [KA) - (Kalp) so_
Therefore,
1 Ripty
[Qipp + [Q1lg. =
2 NANBd T2- -Tg
T + 5(2 - T, - TTg²)
Σ Ze
i=1
+
(8
22
A + V v0_
Thus,
d
2.
(1 5)
5)x²0a))))) Fg/ Kgp + kip/tigg
NANBd (2- - T,,
Kipigg
Σ 22
dt
NANBd
NANnd No (882) (BB_1 + (
2
Note that the coefficient of the first term negative. Combine this with the previous bound for and we complete the proof.
|
Proof of Lemma C.13. Recall that
d
K KAlp = [ -- Kalp [-KKalp (KAl,)) ')
KBQu
KB.CQI.sal K
di
K KA
NAN
d K
A.Elo = -- KAKKo (KARl)) ')
[K [KDQLeak
Q2l
di
[|
NA
We now compute
d dt
KKapp
K
K Ka.ll)) ar d (K (eAel)) and (KAlp=
separately. First, we write
d dt
KAlp;
(AAlo)
KA
[ [--KAli
(KA)>)
KBQ of
AKKll NANBd [ [KB.Q1
(I-KAl, (KAP))
of [[ APpll
+ KA,cly
NAd
TT
d ((
KAlp (KAsla)
+T2
dd dt
KAlps KAAS)) ))
dt
54

---

Under review as a conference paper at ICLR 2023
Then, compute
(
K almo [K (Aell))
Ti
dt
KAiso (KKss K
(KA Kalp)
[Kalp- (BB)))
||[ (Q
p
||[Kappl
NANBd [| QQ)kkp
of NANBd
Σ
KAAEEo KB)K) - (KAsoo (Alp))
(K (Kalp: (Ka)))
+
kpp
- ((AAel (nnl))
- (Kail (Kab))
QQ1pp
+0 (
mmx
dridgg, (Be.1 + vo.)).
NB
NANBd
and, by Lemma C.6, we have
T2 (( (Jjiaal Kalp ([Aal))KK)))))
(KAel)))
K KB.e)a (11, KA T (I- (Kaly (KKA))")
of NAN
Σ
K 1ppll
k=1
o² NA N/
dō
Combine these together and we get
d K KAl KAcl)
= ((KA.elo
(BBp) (KA.Ee (Kab))
1
di
NaN
NA dmodg,1 (De,1 + vo=))
+ 0
Then, we compute
d (Klpp. de (KA.
We write
d
1 I-|KA,l (KARlo)) 7)
[K [K 1Q1Ealg. of ||[KA
KAlpi KK dt
KAlp
NANBd
Kalp (I-[KAel, (KAela (Kasla))
K [Kne]a(Q2]
[KA
NANBd
k=1
d ((1 [KAlp
d (AAc))) )) +T2 ((KA)),
KAAlo )
Tı
dt
dt
For T1, when = p. by Lemma C.6. we have
KAlp (I [-- KAAsl (KAsSO)) [KB] (Q1.e
|
-((KAlp: [[KK)):(Knb)>>>>>> - (IKalr (Taali-KK.tl (Keee) (Ka. (pl,))
||[KB|yl|
1Q la
KA.elul
[[[KB)pl NANBd
- -(1-5)
(1 (1 + ( ((_L)
(K Kaslg
#0 (d3 (BAB. + +8N/)) 6g/5))
= ( |||||||| (KEP KA +0 (d (20ddd())) (DAB,) + &N/s)) 581/))
NANBd
When & *P, p. we have
[ K-KK (KAela) )
1R) @1 Q1.exla. = O (100ABB_15..))
K KAlp
K (KAe) Igll
55

---

Under review as a conference paper at ICLR 2023
Hence,
d Ti (KKATo: dt
KAssl )) =-0
||
B|l
oe NANBd
[KBlp. KA
NANBd
o max
d Ko (ÓAB._ + +6N1s) og, ōN/S)
NANnd
Now we consider T2. By Lemma C.7, we have
T2 ( KA dt
=0
o
d'
EN/S (DAB.) + )w1s)
NANBd
Therefore,
d KAlp: dt
KAsla
(AAs)===0 = -
||[K)p||
o (Kllp KA
NANBd NAN
d Fua NAN nd
+
d no (SAB.L + 5N/s)) gτ )g.1
Thus,
Q P1ppp
K KAll (A.elo = (KaAEl (Blp)
[K [K KAT:)
di
NANBO
of NANBd
KB'p KAglg
NAN
Puax NANBd
disgde,1 (6E,1 + v6.))
Similarly, one can show that
d
K
(A
= ((KA.lo KAlp)
[K [BBp)
Q QuIp.po NANBd
dt
||| Al
oe NANBd
KAlp (KAss))
NANBd
(5e, + +vo.).
+
NANBd disndg,
For notational simplicity, define
1Qilpp07 NANBd
of
X KAlp: KAcla).
== [Klppt (AA.ls))
C1
Cz
NANBd NAN
Then, we can write
d x [ -c - Ju Y
X y G O
( _
NaNnd dmoden_ (dE.) + v²-))
+
Ci - C
-C
The first matrix is negative semi-definite, whence
d de (x2 Y2) so
02 max NAN Bd d (ōe,1 + vo-))
Since the roles of kAE and KB, are interchangeable, the same bound also holds for (KA)) [K Knel)) and ( (Kaapp, (K....)
56

---

Under review as a conference paper at ICLR 2023
Proof of Lemma C.14.
We write
d KK A.Elp df
I K.yy (KBEla) 1)
[K Qignn K.elll
Ka. [K
of N
K (BBll
= AA
KAAQIeR i of n.elll NANnd
AK
I KKKKa
(KBEl))
Kacp
(
II- (-iKa (KBell) ")
[ K2 oz [ NANBd
d Ti (( ((KAe)
- dt
(KBello) +T2 ( (KAEl
d
K
dt
Then, we compute
Ti
K (K..El
K knell))
di
=Σ r
[A]x QQ.gn
of
=
(KKKAsl KAl$)
KKA.pp: (B
(K KAlk)
[.
NA
Note that, the first part of each summand is bounded by (( For the second part, by Lemma C.6 we also have
[I a&&ll (Q1,5bla.
[|[K) (OCI) ([KBpp. [ KAgga)) NA
+ + (d (6AB.) +8v7s) o5 55.5))
=
Hence,
d KAAPP dt
T1
=0
do2 NANBd
Now we consider Tz. By Lemma C.7, We have
T2 KKA dt (KBEll)
d-r
T (I-IKBE)o (KBElA) ")
||Ka,elll| 12lg. of
KA
K KA,gh
K
B.glll
NANBd
o² O(1)[Q2la. NANBd
22
= +
d' (AAB.1 + ENIS) 65,15
NANBd
Combine this together and we obtain
d K KK di
0²
K (...lo))
= =0
do 1
NA
Similarly, derive the same bound for
K
KBalg
and complete the proof.
C.3 MAINTAINING KA KB
In this subsection, we show that (K KB)
~ 1, ||[KAp|l ~ [[KBlp] and also KA,ell z [| throughout Stage 2. The proofs are similar to the corresponding ones
in Stage 1.
57

---

Under review as a conference paper at ICLR 2023
Lemma C.15. In Stage 2, we have
dt KAlp KBly) = = of 11p
=0 (0 o T
(1- KK (Alpp
(1-(KAlp) KBly)
dRog AB.
NAN
NANBd
Lemma C.16. Define PA/B.y = ||[Ka|p]|* /||K/| and PB/A. = ||[KB)pl| /|| In Stage 2, we huve
d (PA/B', + PB/Aa) dt
21Q11
(Kal (nl)) (2-PA/B.r PA/B.B -PBAAP)
NANBd
+
NANBd AB.L +82)
Lemma C.17. For any p.g E [d- - rl, define PE,A/A.pia
= ||AA / /I[ and
In Stage 2, we have PE.A/B.//a || ell [[
T
PE.A/AD// = +0
d8 d5
NANBd
d dt (PEAA/BO// + Pc.B/A =0 =0
oe
NaNBd (do + +62)
OMITTED PROOF OF THIS SUBSECTION
Proof of Lemma C. 151 First, we write
d
KBBp (I-(KAlp (KA)o)) ))
K BQQIpp 1lp
of
K KAlps KBBp dt
|1 K AAppl NA
+ KB, T
7(1 (I (I-IKAl [Kalp (K (KAll))
[K 1LeEh
o
[[A)p]]
Ndd
((
d dt
KAlp (KTp)) +T2
((d
T;
K KAlp KBlp)
For Tı, we compute
Ti ((0 K KAlp (KB)))
- Blp T
= K
I [-lKAlp
[KAlp)
K Blp
|KB)ppl [Q
of
[[ Alpl
NANBd
+ Σ KB
(I-[Kllp
(Klp)
KBB
[[K
[11]es kapp
o
K K
NANB
kpp
Alpp KBIp) 2 || [Q
of NANBd
1
||Klplll
1 0
dRo02 NANnd (BOOABBLL. AAB.).
max
For T2, we compute
Tz ((d d dt
d-r
[K (etg)l/Q1..ukepp of [||a]pp
KAlp
KBLp)
= Σ KBlp I KAlp
[- [KATp)
NAN
k=
= 0 ( o² 2
NANBd 1
d82
Therefore,
2) K KAlp (Klp) = (1- (KApp (BB/)/
=(1-
[||l QiImp of NANgd
+0 ( o drin $2
dt
K [KA)p
NNNN AB,
58

---

Under review as a conference paper at ICLR 2023
Then, by symmetry, we have
di ( (Kalm KBlp)
= (1
K (kpps (BBp)
2) IQilm
[|[ B Blll
[KAbll
of NANBd
|||
[|
oi2 Tuax
+
NANBd drind AB,L)
Ω
of 11pp
NAN
(1-( (Kam (KA1))
+ +0
o F NANNd
dran 02
AB.1)
Proof of Lemma C.16. Similar to Stage 1, we define PA/B.P = ||<b|l| [[[ By Lemma A.5, have
dt PA/B.p = de K Alll
d
at [KB|pp
PA/B.P
K Blpl
KBlpl
2 K 1p:
KnQ 1p)
Op [[K Alp ( of
PA/B.p Qoo p
Na
KBlpl
NA KB)p xlpp
N2d
Bpp.. (AQTTp) of +2 BBppp [ KI.Al1
PA/B,P
op
+ 2 Ngd 2007
NAN Bd [
NA K A
K lp:
1nQUpp)
((A)p: K KB Q1.emlm
= 2
O +2
NAN
KBBpp
NANBd KBBpl
2 (K (KAQTp)
NANBd |Kallll of +2 ((KB)p (I,g
PA/B,P
NaN of
1
2Quo
N2d NZd) PA/B.
Then, by Lemma C.5 and Lemma C.6, we have
([K KBQ))p ([KA)ps (KB))) 1Q1).pp +0 (dr22 AB
([ (KBBpp [ T1p) = ([KApp, K (BBp) (Q1 + O (dr (i AB.1)
(K Alp, [B. 11. +0 (dng NI
([K Blp [KA& +0
Thus,
dt PA/B. = QQlppo0
d
([Alpp [KBP)
(1 (A-AABB.))
NANBd
||[Klpl1
2
+ 2Q000
PA/B,P +0
NANBd drond2 d AB..)
N2d
NEd
Interchange the roles of A, B and we get
d
Qilp,00 NANDd
(KAppt
[KBlp)
(1 (1-PB/Ap))
di PB PB/App
||[A)p|1
1
PB/A,p +0 ( -
+ 2Q00?
drago A AB..)
N2d
Nad
NANDd
59

---

Under review as a conference paper at ICLR 2023
Hence,
d dt
Q1 p
1 A./
(PA/B= + PB/A.p) = NANBd Bd
{[Kapps (KBp)
[[[K
[[[a)pp]
o max
+ +QQ0
(PA/BR PB/AAp +0
NANBd d
Nad
N
AB.1
2QQ 2QQ1
KAlpp KBlp)
(2 PA/B.P PB/A.)
NANBd
max NANBd
(drgd
AB,
Proof of Lemma C.17. For notational simplicity, we drop the subscript & in the proof. Recall from Corollary C.9 that
| [[[[K,Elal]
(
o K A.eall
Qooद=
dt
Năd
NAN
Hence, for any p. 9 E [d r], we have
d dt PA/A.//1
a |[[
A.
[[[K& 2 A.|Kll
A||Kl PA/AN//
2
+0
d8 NANBd F.L
=
Nad
NA d8c ))
2
Nad Qoo? +
PA/A.//
of
=0
do2 C.1
NANBd
Similarly, have
d
K A
K L
de
de
dt PA/Bal
PA/B.//o
(.KKelll1 2
2
of
Q00 P
10
do
Nad
NAN
2 Pook
+0
NANBd do2
PA/B,p/u
NZ
2
2
QOoEPA/BAN/0 #0
NANBd do2
N2
Ngd
By symmetry, we also have
o2
2
Quoe PB: A.g/p P +0
NANBd do
dt PB/A.//9 =
N2d
N2d
Hence,
2
of
dt (PA/Balg + PB/AA/D)
2no (PA/B/// +0
d82
N22
22.
NANBd
=1
(do2_) + +82)
NANBd
60

---

Under review as a conference paper at ICLR 2023
C.4 CONTROLLING THE NOISE-SIGNAL RATIO
In this subsection, show that the noise-signal ratio remains small throughout Stage 2. Lemma C.18. Let [| Alpll be the smallest among all [|[KA]llll (Illelee))) For any q E [r], in Stage 2, we have
d KA.slal SO of BA (d8 + +d + + aL ((owss0e..)
K 2
dt ||[K
||[K>>>
Proof.
Recall from Corollary C.9 that
20
[[[KA)p
[[[ ( (Kapp: (Kml)) (Qih + +2 21[[K Alpll
200
NA
N2 d
Kadō AB.1 AB.LL
NANBd
2 | A.
of [ Al
di ||
Qoo2 +
do E.
NA
Since the condition number of K is bounded by vd, it suffices to consider the smallest [|[ Appl for which we have
20
20
K = NAd
Q11p.p
[
Qo
N2d
NANBd (Kndō
AB,L + +6_)
20
20 1Q1).p [A|p1
Σ
Qilkst
||(a|>|
NANBd
N2d
||]l|
k=
ph NANBd
d dE + +6))
+
where the last line comes from Lemma C.8. By Lemma B.4, [Q1lp.p negative correlated with xzz As result, we have
d
KAll| -0
NAN (BB1 d + +6-)
o2
dt
For K KA,e, we simply have
[| ( NANDd Bd
d
(d62 1 + dd AB + + dl dPōwisbe.l) || A|
di
Thus,
d || KA.E/yl
||[||| KA [[ 11 di
de
dt |||ppll
K A|pl
[|[KA|]''
||[ Alpll
[| |||K|||1
5o NAN
(d82 + d8o0 + + + d +d& dd
61

---

Under review as a conference paper at ICLR 2023
C.5 ESTIMATING THE CONVERGENCE RATE
In this subsection, we estimate how fast the condition number will become close to 1.
Lemma C.19. Suppose that ||[Kalrl is the largest and ||[Kal the smallest among all [ {|] In Stage 2, we have
d [Kabl
4(1 $)o
||K)pp
S
(Tp TT))
de [K
NANBd
||Ka)gl
NANBa (dB + 8_ - dl d801500, mo
Corollary C.20 (Convergence rate). Suppose that |[ Alpll is the largest and ||| Algll the smallest among all [I[[K Alall For any constant > 1, it takes at most poly(d) amount of time for
|||ll ||| Algll to become smaller than
OMITTED PROOF OF TIIIS SUBSECTION
Proof of Lemma C.19. By Corollary C.9, Lemma C.8 and Lemma C.5, we have
d dt
20
20
K A|rl
11
[|[ [|A)pll
Qo [|[K)|ll]
NAd
NAd
22
+ 0
NAN (Nidd
Pp"
AB.L + +6_)
4(1 $103
(1 Txp ||
4(1 - 4)0-5)00
Σ N2
(1 T:))
||Kabb|f
N NAN
Nºd
||x]|
k=
#
NANBd (LBA +8_ ddNN/s0E_.)
4(1 Soog
+0 BBd (T) --]]I[Kallll [[[Ab|l
p
(d
AB. 1 && ddd +d6xy/6b.11)
NAN
where T Ek=1
x² |<<
Tp. Then, we compute
d
|KAppl|
[KAppl Apll di [KAgg
de
dt
|| |Alll
[[[K Jall
[[[KAJll Aall
||[K Alrll
4(1 - 5) NANBd
(0)) Tp 7)
oy (T, -+)
|||K|] |n
| [KAllll
0 NA 2 (& + + + d' d
Since [[AApp]] is the largest, [|[Ka], is the smallest, Tp is positively correlated with || Alpll and 7 is weighted average of Tp. we have
of (T - 1) -02 (Tg (Ta- - 7) 2 og (Tp - Tq)
Thus,
d |Apppl
4(1 510 min
|KA|p]
(T, -Ty)
df [|
NANBd
||(A|go1 2
2 +
P
AB,1 (682 ABBL1 + - d
NANBd
62

---

Under review as a conference paper at ICLR 2023
Proof of Corollary C.20. Recall that
[I[KA|pl
Tp tanh
Kp
tanh
+ 0 +dN1s))
NANBd
||Kali ||
Since Ko & vd. we have (11) for any 1.1/2. Therefore,
K ppl
K |//l
S 1/2. Note that tanh'(z) = ttanh(2))
K AAllll
+ 0 + &N/5) 2 12 2
T - Tg 20 Ω
K |Alli
Then, by Lemma C.19, we have
d K p|pl
4(1 50 min
|[Ka|pll
S
(Tp -Tg)
dt ||| [ Alqll
NA
||
of
0
(dORA
+ 8_ + d' Kg
NANBd
min NANBa
||AAppl
||(all.1
By the proof of Lemma C.19, the largest [|[ ||pl is non-increasing. Hence, NAd is upper bounded by some poly(d). Thus, it takes at most poly(d) for [Ka]]]| [|]]]] to become smaller than c.
C.6 PROOF THE MAIN LEMMA OF STAGE 2
Proof of Lemma C.I. The polynomial bound on the convergence time has been proved in Corol- lary C.20. For the errors, recall from Lemma C.12, Lemma C.13 and Lemma C.14 that
dt
d
8
50
NAN
2
Ko
(d26%BB + +dAR, V8_ + 5-) v Vol
d
o2 NAN nd
F.d (0E.1 + vo-))
DE
50
drog8
dt
K K
(B
10
do
NAN
Recall from Lemma C.15, Lemma C.16, and Lemma C.17 that
o
dt KAlp KBbp =
NANBd (1 KAlpp KBIp)
+0
NAN AB.L)
drggo 2
dt (PA/B PB/A.)) 50
NANBd (dra0d
AB.L + (82)
d
de PEA///// 50
NANBd d52 E.L
of (d52 + +02)
dt (PE.A/Bm// + P..B/AA/)) 50
NANBd
(17)
By Lemma C.18, we have
d KA,e/ 5O dt ||[K|p||
(d82 + d888BB. + + doNysSo.L.)
NANBd
|||
Note that on the RHS of these equations, the only terms whose order may potentially be smaller than or equal LHS are the 8_ -_-reltedd terms. However, (17), can make sure 8_ is most 81.5 As result, the orders of the RHS are all greater than the orders of the LHS, which implies that these errors can at most double within poly (d) time they are sufficiently small at the beginning of Stage 2.
n
63

---

Under review as a conference paper at ICLR 2023
D FROM GRADIENT FLOW TO GRADIENT DESCENT
Converting the above gradient flow argument to gradient descent one standard. All our esti- mations can tolerate inverse polynomially large error. Since the all quantities of interest here are polynomially and inverse polynomially bounded, at each step of gradient descent, one can al- ways make the GF-to-GD discretization error sufficiently (inverse polynomially) small by choose a sufficiently (inverse polynomially) small learning rate and generating sufficiently (polynomially) many samples. Since the times need for Stage 1 and Stage 2 both polynomial, this also implies a polynomial sample complexity.
64


Reviews:
[{"review_id": "lEuOmJG716w", "reviewer": "Reviewer_SHU1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "The paper set out to investigate the learning dynamics of contrastive learning (e.g.CLIP) in, investigate how contrastive learning learn to align the representations from different views efficiently.", "summary_of_the_review": "This paper is well-written and original, it provides an interesting view in investigating the training dynamic of contrastive learning with alignment and balance of feature representations, however, the experiment seems only done on simulated data and linear models.", "strength_and_weaknesses": "Strengths:\n(a) The paper provided a fresh view that the alignment and balance of representation play a important role in the training dynamics of contrastive learning, and the training can be decoupled into two stages.\n(b) The paper is backed by a series of theory-grounded definitions and proof.\nWeaknesses:\n(a) The experiment part only covers simulation results, not real data, the model only cover linear models.", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VHUfmZo7_-Q", "reviewer": "Reviewer_KDMd", "summary": "", "questions": "", "limitations": "", "rating": 5, "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 theoretically studied the joint-embedding training dynamics in a linear setting - which may be suitable for a multimodal environment. They studied both contrastive and non-contrastive cases and showed that contrastive negative pairs are essential for preventing representations from becoming a rank-one solution. The analysis largely follows Tian et al. (2021), Wen & Li (2021), and Jing et al. (2022) using gradient flows. Lastly, the authors provide a numerical simulation that verifies their claims. ", "summary_of_the_review": "This is a well-written theory paper studying contrastive learning dynamics but only limited to the small scope and seems not to provide real insight. ", "strength_and_weaknesses": "Strength:\n1. The authors considered the normalization step in the contrastive/non-contrastive loss when conducting the analysis of their dynamics. This is a big improvement over previous theoretical works (Tian et al. (2021), Wen & Li (2021), Jing et al. (2022)) on this topic.\n2. The phenomenon of the two-stage training dynamics is novel. Most previous theoretical work ignores training dynamics and only focuses on generalization bounds that are unable to explain such phenomena.\n\n\nWeaknesses:\n1. The non-contrastive setting is known to collapse. In fact, all successful non-contrastive methods like BYOL/DINO rely on asymmetric architecture. This is already proven by Tian et al. (2021). The authors derive an obvious conclusion (align->1, balance->0) that does not provide any insight into the problem. \n2. The overall claim on contrastive loss is shallow. The authors show that under perfect conditions, align-> one and balance->same rank as input. For linear settings, this is obvious that the optimal solution exists. However, the analysis seems impossible to scale to nonlinear settings as the definition of `balance` becomes trivial for nonlinear networks.\n3. There is no discussion on possible follow-up based on this work. For example, based on theoretical discovery, Tian et al. (2021) and Jing et al. (2022) propose the corresponding solution. It is unclear how this theory can boost empirical advancement.\n ", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "L0ihBycQX0K", "reviewer": "Reviewer_kjwZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "In this paper, the authors examine the role of contrastive loss across modalities (such as in CLIP).  They demonstrate that the contrastive loss promotes the learning of “aligned” and “balanced” representations.  The authors investigate the learning dynamics of the optimization process (under specific listed conditions) and show that the dynamics can be interpreted in the form of two stages.", "summary_of_the_review": "Given the novelty of the work, as a comprehensive step towards understanding multimodal contrastive losses, I am inclined to support this paper for acceptance.  This is predicated on the correctness of the listed mathematics - I willingly confess that a large caveat of my review is that I was not able to carefully check all the math or the claims in this paper.  Therefore, I will provide a marginal acceptance and look forward to extended discussions with the AC.", "strength_and_weaknesses": "The strength of the paper lies in its supposed mathematical rigor and derivations.  A weakness of this paper is that it is limited to a linear data-generating model, and therefore the conclusions are unable to directly translate to understanding multimodal contrastive losses used in practice.\n\nI enjoyed how the authors included both contrastive and non-contrastive loss formulations.  The authors analyze the non-contrastive loss (such as BYOL) in comparison with the contrastive loss to demonstrate why contrastive losses are better (in terms of going beyond alignment only).", "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": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Reject