O N THE I MPORTANCE OF C ONTRASTIVE L OSS IN M ULTIMODAL L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 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 a theoretical perspective, it is unclear how contrastive learning can learn to align the representations from different views efficiently, 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 efficiently balance the learned representations.
In particular, we reveal a stagewise behavior of the learning process:
In 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.
I NTRODUCTION
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 a 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.
; Radford et al.
(2021); Ramesh et al.
; Xu et al.
; Jia et al.
; 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.
; Chen & He (2021);
He et al.
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 it also requires the data from different pairs (i.e. negative pairs) to be as negatively aligned as possible.
In the paper, the authors also identify contrastive loss as the most critical part to the success of CLIP.
Despite the empirical success of this contrastive learning-based method, from a theoretical perspective, 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 efficiently 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.
; Jing et al.
; Pokle et al.
; Tian et al.
Wen & Li (2021)), one main technical challenge of analyzing contrastive learning in a multimodal setting is how the model can be trained to align the feature representations f A , f B from modules A and B respectively.
Due to the existence of negative pairs that emphasize negative correlations of f A and f B , it is unclear that the model still has incentives to align the features from different modules.
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 x A =
Az A + A ξ ξ A and x B = Bz B + B ξ ξ B , respectively, where z A , z B are considered as the hidden signals, A , B linear transformations from the signal to the observation and A ξ ξ A , B ξ ξ B the noises.
Similar linear models have also been used in previous works (Tian et al.
; Wen & Li (2021)) in the context of single-modal learning ( A = B ).
The positive pair of the data shares the same signal
z A = z B , and has different noises ξ A , ξ B and transformations A , B .
The goal is to learn features f A , f B that align positive pairs while keeping representations of negative pairs away from each other.
Under this setting, we make the following contributions:
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.
; 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 ).
We consider feature learners f A , f B with normalization, meaning that f A , f B are always normalized to have (expected) norm one during training.
Output normalization plays a critical role in the practical success of contrastive learning and is also employed in CLIP, but it is rarely considered in theory due to the additional complexity of the division by norm.
We analyze the learning process of stochastic gradient descent from random initialization.
We prove that contrastive learning converges efficiently to a nearly optimal solution, which indeed aligns the feature representation f A and f B . 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 f A and f B , the features learned by contrastive learning with negative pairs is much more uniform , meaning that f A , f B 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 f A , f B will only focus on the top singular direction of A and B . 5.
We also perform simulations and more practical experiments to further support our theory.
R ELATED W ORKS
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 et al.
(2021), in which the authors show that, in certain cases, multimodal methods can provably perform better than singlemodal models.
However, the authors consider neither contrastive pairs nor the training dynamics.
Contrastive/Non-contrastive learning theory Another much richer line of research is about contrastive and non-contrastive methods in the context of single-modal self-supervised learning.
Starting 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 et al.
Lee et al.
; 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 benefit 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 representations only span a low-dimensional subspace of the embedding space.
In Pokle et al.
, the authors show that non-contrastive losses have many non-collapsed 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, with random augmentation, neural networks can learn features that are 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.
; Wen & Li (2021)
The dimensional collapse problem reported in Jing et al.
(2022) is not a real issue in our setting, since, in our case, the best the model can do is to recover the latent vector z , up to some rotation.
As a result, it is natural for the learned representations to span only a low-dimensional subspace of the embedding space R m .
Here, the point of choosing m ≫ d is to make the optimization dynamics more regular, which is a common strategy in the study of over-parameterized neural networks.
The main difference between our work and the analysis in Tian et al.
and Wen & Li (2021) is we do not assume the inputs are isotropic.
In our setting, the condition number can be as large as Θ(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.
P ROBLEM S ETUP
Similar to previous theoretical works (e.g., Lee et al. (2021); Wen & Li (2021)), we consider a linear data-generating model.
Formally, we assume that the contrastive pairs ( x + A , x − B ) are constructed as
[Equation]
where z ± , ξ − A , ξ − B are independent random variables following the uniforms distributions over {± 1 / √ d } r , {± 1 / √ d }
d − r and {± 1 / √ d } d − r , respectively, and A , B ∈ R d × r , A ξ , B ξ ∈ R d × ( d − r ) are matrices with A ⊤ A
= B ⊤ B = diag( σ 2 ) and A ⊤ ξ
A ξ = B ⊤ ξ B ξ = σ 2 ξ I d − r for some σ ∈ R r + and σ ξ ∈ R + .
In words, we first sample the latent vectors z ± ∈ R r independently and encode them with A , B to form the signal part of the input.
Then, we sample the latent noises ξ + A , ξ − B ∈ R d − r , and encode them with A ξ , B ξ to form the noise part of the input.
Finally, we add the signal and noise parts together to obtain ( x + A , x − B ) .
To generate a positive pair ( x + A , x + B ) , we use the same latent vector.
That is, x + A =
Az + + A ξ ξ + A and x + B = Bz + + B ξ ξ + B .
Note that the latent noises here are still independent.
We use σ 2max and σ 2min to denote the maximum and minimum of σ 21 , . . .
, σ 2 r , σ 2 ξ , respectively.
We assume that ( σ 2max /σ 2min ) max (cid:110) 1 , ( d − r ) σ 2 ξ / ( rσ 2min ) (cid:111)
≤ c log d for some small constant c > 0 .
Our results can easily be generalized to settings where the dimensions of x A and x B 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 a 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
[Equation]
, where W A , W B ∈ R d × m are the trainable parameters.
In words, we first map the inputs ( x A , x B ) into the embedding space R m using W A and W B , and then apply batch normalization to the outputs.
By saying the learned representations are aligned, we mean that f A and f B 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 x A (cid:13)(cid:13)
W ⊤ A x A (cid:13)(cid:13) {+
+} =
(cid:13)(cid:13)
W ⊤ A A (cid:13)(cid:13) 2 F /d + (cid:13)(cid:13) W ⊤ A
A ξ (cid:13)(cid:13) 2 F /d .
For notational simplicity, we write K A = W ⊤ A A , K B = W ⊤ B B , K A , ξ =
W ⊤ A A ξ , and K B , ξ = W ⊤ B B ξ .
These are the matrices that directly map latent vectors to their final representations.
We also define N 2 A
= ( ∥ K A ∥ 2 F + ∥ K A , ξ ∥ 2 F ) /d and N 2 B = ( ∥ K B ∥ 2 F + ∥ K B , ξ ∥ 2 F ) /d .
Then our model can be rewritten as
[Equation]
We initialize each entry of W A and W B using iid Gaussian N (0 , 1 /m ) .
This scaling ensures the norm of outputs before initialization does not blow up as m → ∞ .
We train our model using gradient descent over the following contrastive loss L .
First, we define
[Equation]
where K is a positive constant controlling the strength of negative samples and τ t ∈ (0 , 1] is the inverse temperature.
Then, we define the contrastive loss as
[Equation]
By the non-contrastive loss, we mean
[Equation]
Formally, the training algorithm is 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 + A , x + B , x − A , x − B ) } Ni =1 , use them to compute the empirical version of ˜ L , and update the weight matrices using W A ←
W A − τ − 2 t η ∇ W A ˜ L and W B ← W B − τ − 2 t η ∇ W B ˜ L .
In the non-contrastive case, we always use τ t = 1 3 , and we repeat the above update until gradient descent converges to an approximate stationary point.
In the contrastive case, we first use a small τ t = 1 / poly( d ) , run the process for T 1 = poly( d ) iterations, switch to τ t = 1 , and run the process for another T 2 − T [-
-]1 = poly( d ) iterations.
M AIN R ESULTS
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 f A and f B to denote the embeddings, but our definitions here will be architecture-agnostic.
After some general discussion, we also discuss these definitions in our specific setting.
Definition 4.1 (Alignment) .
We define the alignment score as
[Equation]
Namely, Γ Align is the accuracy of classifying whether the input pair ( x A , x B ) is a positive pair.
We say that the learned representations are aligned if Γ Align ≈ 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, Γ Align will be 0 (or 0 .
5 if we choose to break ties randomly instead of using strict inequality).
Definition 4.2 (Balance) .
We define the balance score as
[Equation]
Σ f [Equation] = E x A (cid:8)
[Equation]
A f ⊤ A
[Equation]
Namely, Γ Balance is the stable rank of the covariance matrix of the output embeddings.
We say that the learned representations are balanced if Γ Balance ≥ αr for some α ≈ 1 .
We make some short remarks on Definition 4.2.
First, we only use f A in this definition because if the embeddings are well-aligned, f A f ⊤ A and f B f ⊤ B should be approximately the same.
Second, the stable rank is usually used as an alternative to the actual rank because it is less sensitive to small singular values (Rudelson & Vershynin (2007)).
To see why this makes sense, note that ∥ Σ f ∥ 2 F / ∥ Σ [-(cid:107) 
-]∥ 22 = (cid:80) mk =1 κ 2 k / max k ∈ [ m ] κ 2 k , where κ 1 , . .
, κ m are the singular values of Σ f .
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.
In Jing et al.
(2022), the authors also use rank to measure the degree of “dimensional collapse”.
Unlike their argument, here we only require Γ Balance to be at least αr , instead of αm , 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 a sufficient condition for Γ Balance ≥ αr is that the ratio of the largest and r -largest singular values is at most √ α .
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 = diag( ν ) A − 1 and W ⊤ B = diag( ν ) B − 1 with ν 1 = 1 and ν 2 = · · · = ν r = exp( − d ) .
This model maps inputs whose latent vector is z to diag( µ ) z , up to some normalization, for both modules, whence it has Γ Align = 1 .
Meanwhile, one can verify that this model has Γ Balance /r ≤ 2 /r ≈ 0 .
The problem of this model is that it overly emphasizes zand is sensitive to small changes in z 1 .
This model can be made balanced by replacing diag( µ ) with I r , in which case the model directly recovers the latent vector z .
As a result, all changes in z will be reflected in the final representation in a faithful way.
Similar notions have also been studied in Wang & Isola (2020) under the name “uniformity” 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 on training.
Still, this further suggests that learning balanced representations is a reasonable and important goal.
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 / poly( d ) is sufficiently small, and we generate sufficiently, but still polynomially, many samples at each step to compute the loss 4 .
Suppose that, for some small constant c > 0 ,
[Equation]
(a) Non-contrastive loss.
There exists a σ ∈ R 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, Γ Align ≈ 1 but Γ Balance /r ≈ 0 .
(b) Contrastive loss.
There exists τ 20 = 1 / poly( d ) and T 1 = poly( d ) such that, for any valid σ , Algorithm 3.1 will reach a point after poly( d ) iterations at which the learned representations are aligned and balanced, that is, Γ Align ≈ 1 and Γ Balance /r
≈ 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 a model to achieve Γ Align ≈ 1 is to learn ∥
f
A − f B ∥ ≈ ∥ K [-
-]0 ( z A − z B ) ∥ for some K 0 ∈ R m × r with K ⊤ 0 K 0 ≻ 0 .
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 f A
= ( K A z A + K A , ξ ξ A ) /N
A and f B =
( K B z B + K B , ξ ξ B ) /N
B .
Hence, we have the following sufficient condition.
Lemma 4.4 (Sufficent condition for aligned representations) .
If ∥ K A ∥ F ≫ ∥ K A , ξ ∥ F , ∥ K B ∥ F ≫ ∥ K B , ξ ∥ F , K A ≈ K B , and the condition number of K 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 K 0 , and the final condition makes sure the smallest singular value
K [-(cid:62) 
-]⊤ 0 K 0 of is still at least 1 / poly( d ) after normalization.
Now we consider the balance of the learned representations.
Suppose that our representations are already aligned, in the sense of Lemma 4.4.
Then we have
[Equation]
Recall the relationship between Γ Balance and the effective condition number of Σ f .
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 K A is close to 1 , then the learned representations are balanced.
Note that, since the columns of A are not orthonormal, even at initialization, the condition number of K ⊤ A K A is not close to 1 .
In other words, a condition-number-reducing stage is necessary for the model to learn balanced representations.
T RAINING D YNAMICS AND P ROOF O UTLINE
Our proof is based on characterizing the dynamics of gradient descent over the contrastive loss.
We first choose a small inverse temperature τ 2 t = 1 / poly( d ) and run gradient descent for poly( d ) iterations.
This stage is called Stage 1.
Then, we set τ 2 t = 1 and run gradient descent for another poly( d ) iterations.
This stage is called Stage 2.
The use of different τ t is mainly for technical purposes since this gives cleaner separation between stages.
We observe similar stage-wise behavior when using a uniform τ 2 t .
See Figure 1 for simulation results.
We also report results for the noncontrastive loss there.
5. T HE TRAINING DYNAMICS
Instead of tracking the parameters W A and W B directly, we will track K A , K A , ξ , K B , K B , ξ , the matrices that directly map latent signals and noises to the final representations.
For the case with contrastive pairs, one can show that the dynamics of K A are governed by the following equation
[Equation]
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 /τ 2 t so that dd t K A does not shrink with τ 2 t .
For the non-contrastive case, the equation is
[Equation]
5. T HE INFINITE - WIDTH DYNAMICS
The overall proof strategy is to 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 infinitewidth network using polynomially many neurons.
This discretization is one 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.
The basic idea is to use Taylor expansion around the infinite-width trajectory to factor out the firstorder error terms and show that, either they drive the process towards the infinite-width trajectory or the error 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 particular, the evolution of the condition number.
Recall that we use iid Gaussian to initialize the entries of W A and W B .
Hence, in the m → ∞ limit, different columns of K A and K B are orthogonal to each other, at least at initialization.
Moreover, one can verify that thanks to the symmetry, this holds throughout the entire training procedure.
Meanwhile, by symmetry, the corresponding quantities in modules A and B are always the same.
Namely, when m → ∞ , we have
K ⊤ A K A = K ⊤ B K B = diag( κ 2 ) and K ⊤
A K B = diag(ˆ κ 2 ) for some κ , ˆ κ ∈ R r .
As a result, in order to characterize the dynamics of K A , K B , it suffices to look at κ 2 and ˆ κ 2 .
One can show that, in this noiseless infinite-width limit, we have (cf.
Lemma A.7)
[Equation]
where ˜ S is a Θ(1) quantity depending on ˆ κ and κ , T p = tanh( τ 2 t ˆ κ 2 p / ∥ κ ∥ 2 ) , and ˜ T
= (cid:80) rk
=1 (ˆ κ 2 p / ∥ κ ∥ 2 ) T p .
The first term comes from the first term of (5) and the second term from the second and third terms of (5).
When κ 2 p ≈ ˆ κ 2 p for all p ∈
[ r ] , the model is aligned.
When all κ 2 p are roughly the same, the model is balanced.
5. S TAGE
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 infinite-width dynamics (7) instead of the finite-width one (5).
Equivalence of Stage 1 and non-contrastive methods Recall that we use a small τ 2 t in Stage 1,
T p ≈ 0 for all p ∈ [ r ] , whence ˜ T is also approximately 0 .
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 infinite-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. for a more formal proof in the finite-width setting.
Now, we consider the contrastive loss. The main result of Stage 1 is as follows.
Lemma 5.1 (Informal version of Lemma B.1) .
Under the assumptions of Theorem 4.3, the finitewidth dynamics closely track the infinite-width ones throughout Stage 1, which takes at most poly( d ) iterations, and, at the end of Stage 1, we have, for any p, q ∈
[ r ] and s ∈ [ d − r ] ,
[Equation]
Moreover, there exists a σ ∈ R r such that max p,q κ 2 p /κ 2 q = Ω( √ d ) at the end of Stage 1.
In words, (8) says, at the end of Stage 1, K A ≈ K B in the relative sense, the noise-signal ratio is small, and the condition number is bounded by O ( √ d ) .
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 ˆ κ 2 p /κ 2 p 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 lemma.
Lemma 5.2.
Let A t be a positive process.
Let X t and Y t be defined as ˙ X t ≤ − A t X t , ˙
Y t ≤ βA t X t Y t , with X 0 , Y 0 , β being positive.
Then, for any T ≥ 0 , we have Y T ≤ Y 0 exp( βX 0 ) .
Here, X t represents the progress we have made and Y t the error we wish to control.
In our case, X t is the maximum between 1 − ˆ κ 2 p /κ 2 p and the noise-signal ratio, and Y t 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. S TAGE
In Stage 2, τ 2 t is no longer o (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 A 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.1) . Under the assumptions of Theorem 4.3, the finitewidth dynamics still closely track the infinite-width 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 ∈
[ r ] and s ∈ [ d − r ] ,
[Equation]
The way to control ˆ κ 2 p /κ 2 p and the noise-signal ratio is similar to Stage 1.
For the condition number, note that when τ 2 t = 1 and ˆ κ ≈ κ , the first term of (7) becomes close to 0 while the second term becomes − 4(1 − ˜ S ) κ 2 p / ∥ κ ∥ 2 ( T p − ˜ T ) σ 2 p .
Since T p = tanh( τ 2 t ˆ κ 2 p / ∥ κ ∥ 2 ) is positively correlated with ˆ κ 2 p / ∥ κ ∥ 2 and ˜ T is a weight average of these T p ’s, this term will push ˆ κ 2 p / ∥ κ ∥ 2 towards their average and, consequently, reduces the condition number.
To obtain a satisfactory convergence rate, it suffices to consider the ratio κ 2 p /κ 2 q directly.
See Appendix C.5 for details.
The discretization error is handled in Appendix C.2.
E XPERIMENTAL RESULTS
Besides the simulation results reported in Figure 1, we also conduct experiments on the MSCOCO2014 dataset (Lin et al., 2014) using more 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 a pre-trained ResNet
+}
(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 fully-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 a prediction to be correct if it matches one label.
The zero-shot classification is done in the same way as in Radford et al. Namely, we compute the image embedding and the embeddings of prompts “This is a [LABEL NAME]”, and use the prompt with the highest correlation with the image embedding as the prediction.
C ONCLUSION AND D ISCUSSION
In this work, we study the role of contrastive pairs in multimodal learning, and show that contrastive pairs are important for the model to learn representations that are both aligned and balanced .
Our work extends previous results in several directions: First, we consider the more complicated multimodal learning problem.
Meanwhile, our data generating model is inhomogeneous, and we show that in this case, non-contrastive method can collapse to an approximately rank- 1 solution while 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.
However, despite the complexity of the analysis, our model is still linear, which is very different from the models used in practice.
Also, for the results on non-contrastive methods, we do not consider more advanced training techniques such as Grill et al.
and Chen & He (2021).
We leave the analysis of these more practical techniques for future work.
A G RADIENT C ALCULATION
In this section
[Equation]
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. T HE FINITE - WIDTH CASE
The results in this subsection are mostly brute-force calculations.
First, we prove the following auxiliary lemma.
Lemma A.1.
For any x + A and x − B , we have
[Equation]
Then we compute the gradients.
Lemma A.2.
We have
[Equation]
The formula for ∇ W B L can be obtained by interchanging the roles of A and B .
Instead of tracking W A and W B directly
[Equation]
we are going to track K A , K B , K A , ξ and K B , ξ .
The dynamics of them is governed by the following equation.
As a direct corollary of Lemma A.2, we have the following.
Lemma A.3.
We have
[Equation]
diag( σ 2 ) −
Kτ 2 t
E x ± A , x + B (cid:26) S B ( x + A , x + B ) exp( τ 2 t f − A · f + B ) exp( τ 2 t f + A · f + B ) (cid:18) f
+ B ( z − ) ⊤
N A − (cid:10) f − A , f + B (cid:11)
K
A N
A d (cid:19)(cid:27) diag( σ 2 ) , and d d t K A , ξ = τ 2 t
N A
E x + A , x + B   (cid:0) 2 − S A ( x + A , x + B )
− S B ( x + A , x + B ) (cid:1)   f + B ( ξ + A ) ⊤ − (cid:10) f
+ A , f + B (cid:11) K A , ξ / √ d (cid:113) ∥ K A ∥ 2 F + ∥ K A , ξ ∥ 2 F    
σ 2 ξ − Kτ 2 t
N A
E x + A , x ± B  
S A ( x + A , x + B ) exp( τ 2 t f + A · f − B ) exp( τ 2 t f + A · f + B )   f − B ( ξ + A ) ⊤ − (cid:10) f
+ A , f − B (cid:11) K A , ξ / √ d (cid:113) ∥ K A ∥ 2 F + ∥ K A , ξ ∥ 2 F     σ 2 ξ − Kτ 2 t N A
E x ± A , x + B   S B ( x + A , x + B ) exp( τ 2 t f − A · f + B ) exp( τ 2 t f + A · f + B )   f + B ( ξ − A ) ⊤ − (cid:10) f − A , f + B (cid:11) K A , ξ / [Equation]
We can rewrite the above result as follows.
Corollary A.4.
Define
[Equation]
[Table]
[Equation]
We have
[Equation]
We are interested in each column of K A and K B , whose dynamics is given by the next lemma.
The next lemma also decompose the dynamics along the radial and tangent directions.
Lemma A.5.
For any p ∈ [ r ] and q ∈ [ d − r ] , we have
[Equation]
and
[Equation]
Finally, as a simple corollary of Lemma A. [Equation] we have the following result on the gradients of the non-contrastive loss.
Lemma A.6.
For the non-contrastive loss (4) , we have
[Equation]
As a corollary
[Equation]
we have
[Equation]
in this case,
[Equation]
O MITTED PROOF OF THIS SUBSECTION
Proof of Lemma A.1.
We compute
[Equation]
For the first term, we have ∇ W A (cid:10) W
⊤ A x + A , f − B (cid:11) =
x + A ( f − B ) ⊤ . For the second term, we have ∇ W A
[Equation]
Hence,
[Equation]
Proof of Lemma A.2.
We compute
[Equation]
By Lemma A.1, the first term is
[Equation]
and the second term is
[Equation] [Equation]
=
[Equation]
[Equation]
Then, for ∇ W A L B , we compute
[Equation]
Again, by Lemma A.1, the first term is
[Equation]
and the second term is
[Equation] [Equation]
ξ ∥ 2 F     .
Combine these formulas together and we complete the proof.
A. T HE INFINITE - WIDTH CASE
Now we consider the noiseless infinite-width dynamics.
The results of this subsection will not be used in the proof.
It mainly serves as a way to give intuitions on how the dynamics looks.
As we have discussed in the main text, in this noiseless infinite-width case, it suffices to track κ 2 p = ∥
[ K A ] p ∥and ˆ κ 2 p = ⟨ [ K A ] p , [ K B ] p ⟩ .
Lemma A.7.
In the noiseless infinite-width case, we have
[Equation]
Proof.
First, note hat
[Equation]
This implies that (a) (cid:10) f
+ A , f + B (cid:11) does not depend on the actual value of z + , and (b) if we flip the signs of z ± p simultaneously, then the value of (cid:10) f
+ A , f − B (cid:11) remain unchanged.
To compute S A , we then need to take expectation over z − .
We compute
[Equation]
Again, it does not depend on the actual value of z + .
One can conduct similar calculation for S B and, consequently, we have S A ≡ S B ≡ ˜ S for some ˜ S that depends on κ and ˆ κ but not on z + .
Then
[Equation]
we can rewrite (5) as
[Equation]
Note that
[Equation]
κ ∥
[Equation]
(cid:33) =
[Equation]
{ p = q } Z
[Equation]
p .
Meanwhile, note that
E x [Equation] , x ± B
[Equation]
Thus,
[Equation]
As a corollary
[Equation]
we have d
[Equation]
Hence, d
[Equation]
σ
[Equation]
By symmetry, for K B , we have
[Equation]
Then [Equation] we can compute
[Equation]
B S TAGE
In this section
[Equation]
we show that the following hold:
(a) K A ≈ K B after Stage 1.
(b) The noise-signal ratio is small after Stage 1.
(c)
The condition number is O ( √ d ) 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 a sufficiently (polynomially) large m and a sufficiently (inverse polynomially) small τ 2 t that may depend on the δ ’s that appear in this lemma so that the following statement holds.
Let T 1 be the earliest time all the following hold:
[Equation]
where δ − , δ N/S ∈ 1 / poly( d ) are two given parameters.
We have T 1 ≤ poly( d ) .
Moreover, at any time t ∈ [0 , T 1 ]
[Equation]
we have κ 0 :=
max p,q ∈
[ r ] ∥ [ K A ] p ∥ / ∥
[ K A ] q ∥ ≤ O ( √ d ) and ∥ [
[Equation]
A
[Equation]
, ⊥ , δ ξ, ⊥ ∈ 1 / poly( d ) are given parameters.
Basically, the conditions in (10) mean that the norm of the corresponding columns are roughly the same, and the columns of all these matrices are approximately orthogonal to each other.
Both of them are true in the infinite-width limit, and by some standard concentration argument, one can make all these errors to be arbitrarily inverse-polynomially small, at least at initialization.
Note that, as a simple corollary of (10), we have
[Equation]
For notational simplicity, we also define
[Equation]
Characterizing the dynamics in Stage 1 is relatively straightforward.
We will see in Section B.1 that all those Q -matrices have simple forms.
The main tool we use to control the condition number and the discretization error is the following nonlinear version of Gronwall’s lemma.
Lemma B.2.
Let A t be a positive process.
Let X t and Y t be defined as
[Equation]
with X 0 , Y 0 , α being positive.
Then, for any T ≥ 0 , we have Y T ≤ Y 0 exp( αX 0 ) .
Remark .
Here, X t represent the progress we have made and Y t the error.
In our case, X t is the maximum between 1 − (cid:68) [ K A ] p , [ K B ] p (cid:69) and the noise-signal ratio, and Y t the discretization error.
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 do 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
[Equation]
Note that
[Equation]
Hence,
[Equation]
The organization of this section is as follows.
In Section B.1, we derive estimations for the Qmatrices defined in Corollary A.4 and use them to simplify the equations governing the training dynamics.
In Section B.2, we estimate the rate at which 1 − (cid:68) [ K A ] p , [ K B ] p (cid:69) 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, in Section B.4, we prove Lemma B.1, the main lemma of Stage 1.
Finally, we prove the negative result for non-contrastive learning in Section B.5.
B. E STIMATIONS 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 1, τ 2 t is small.
Hence, exp( τ 2 t (cid:10) f
+ A , f − B (cid:11) )
= 1 ± O ( τ 2 t ) .
With this approximation, one can derive the following estimation for S A and S B .
Lemma B.3 (Estimations for S ) .
In Stage 1, we have
[Equation]
and the same is also true for S B .
Here, O z ( τ 2 t ) means O ( τ 2 t ) and this does not depend on ξ A nor ξ B .
The proof of this lemma and all following lemma is deferred to the end of this subsection.
Note that we derive a slightly finer estimation for the noise-related part.
This additional ρ N/S will be used cancel with terms like ∥ [ K A ] p ∥ / ∥ [ K A , ξ ] q ∥ , at the cost of a κ 0 factor, in later analysis.
We emphasize here that O z does not depend on the noises so that later we can argue E ξ (cid:8)
O z ( τ 2 t ) ξ (cid:9) = 0 .
With this lemma, we now derive estimations for Q 1 , Q 1 , ξ , Q 2 and Q 0 , respectively.
Lemma B.4 (Estimations for Q 1 ) .
In Stage 1, we have
[Equation]
Lemma B.5 (Estimations for Q 1 , ξ and Q 2 ) .
In Stage 1, we have
[Equation]
The above two lemmas say that in Stage 1, Q 1 is approximately diagonal, and Q 1 , ξ and Q 2 can more or less be ignored.
Lemma B.6 (Estimations for Q 0 ) .
In Stage 1, we have
[Equation]
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 1, for any p ∈ [ r ] and q ∈ [ d − r ] , we have
[Equation]
The formulas for K B and K B , ξ can be obtained by interchanging the roles of A and B .
Basically, the two parts of the first term of dd t ∥
[ K A ] p ∥ 2 correspond to the signal and normalization, respectively.
For dd t ∥
[ K A , ξ ] q ∥ 2 , the signal term is 0 because the noises in x + A and x + B are independent and have mean 0 .
Corollary B.8.
In Stage 1, we have
[Equation]
Interchange the roles of A and B and one can 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 [ K A ] p and [ K A ] q toward each other.
Moreover, without the normalization terms, different [ K A ] p 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-contrastive method and the Stage 1 dynamics are equivalent, up to a multiplicative constant and some higher-order terms.
O MITTED PROOFS OF THIS SUBSECTION
Proof of Lemma B.3.
First, we write
[Equation]
N
[Equation]
For the second term, we write
[Equation]
For each summand, we have ∥ [ K B , ξ ] j ∥ / ∥ [ K A ] i ∥ ≤ O ( ρ N/S ) , (cid:68) [ K A ] i , [ K B , ξ ] j (cid:69) = ±
O ( δ ξ, ⊥ ) , and | z
+ i ξ − B ,j d | ≤ 1 .
Hence, ⟨ K A z + , K B , ξ ξ − B ⟩ N A N B = ± O ( dδ ξ, ⊥ ρ N/S ) .
The same is also true for ⟨ K A , ξ ξ + A , K B z − ⟩ N
A N B d and the third term.
Therefore,
[Equation]
Then, we compute
[Equation]
Thus,
[Equation]
The proof for S B is essentially the same.
Proof of Lemma B
[Equation]
Recall that
[Equation]
Since S A = (1 + K ) − 1 ±
O ( τ 2 t ) and S B = (1 + K ) − 1 ±
O ( τ 2 t ) , we have [Equation]
−
[Equation]
As a result,
[Equation]
Note that E (cid:8) z − ( z + [Equation] (cid:9)
= {+
+}
and E (cid:8) z
+ ( z + ) ⊤ (cid:9) =
I d /d .
Therefore,
[Equation]
Proof of Lemma B.5.
Recall that
[Equation]
Note that if some quantity X does not depend on ξ , then E { X ξ } = 0 .
Hence, by Lemma B.3, we have
[Equation]
The proof for Q 1 ,ξ A and Q 2 is essentially the same.
Proof of Corollary B.7.
Recall that
[Equation]
We now estimate these terms one-by-one.
For T 1 , we have
[Equation]
Hence,
[Equation]
T 2 , we compute
T
[Equation]
Finally, for T 3 , we compute
[Equation]
Combine these together and we obtain
[Equation]
K
[Equation]
(cid:33) .
Interchange the roles of A and B and we obtain the formula for dd t ∥ [ K B ] p ∥ 2 .
Now we consider
K A , ξ .
We write
[Equation]
For T [Equation] we compute
[Equation]
Similarly, one can show that this bound also holds for T 2 .
Finally, for T 3 , by Lemma B.
[Equation]
we have
T
[Equation]
Combine these together and we obtain
[Equation]
Interchange the roles of A and B and we obtain the formula for K B , ξ .
Proof of Corollary B
[Equation]
We write
[Equation]
For T 1 , we have
[Equation]
For the first term, by Lemma B.4, we have [Equation] [Equation] p,p
[Equation]
Also by Lemma B.4, for each summand in the second term, we have
[Equation]
Therefore,
[Equation]
For T 2 , by Lemma B.5, we have
[Equation]
Combine these together
[Equation]
and we obtain
[Equation]
Now, we consider K A , ξ .
Again, we write
[Equation]
For the first term
[Equation]
by Lemma B. [Equation] we have
[Equation]
The same bound also hold for T 2 .
In fact, we can have a slightly sharper bound for it because we no longer have ∥ [ K B ] k ∥ / ∥ [ K A , ξ ] q ∥ .
Combine these and we obtain
[Equation]
Proof of Lemma B.9.
The proof of Corollary B.8 and Corollary B [Equation] 7, mutatis mutandis , yields
[Equation]
Recall from Lemma A.6 that, in the non-contrastive case
[Equation]
we have [Equation] Note that they are exactly the same, except for a 2 K/ (1 + K ) factor and some higher order error terms.
B. C ONVERGENCE RATE AND THE CONDITION NUMBER
In this subsection, we estimate the rate at which 1 − ρ − and ρ N/S converge to 0 and the growth rate of the condition number.
The basic idea is to use the estimations we have derived in Corollary B. and Corollary B.8 to approximate the finite-width dynamics with the infinite-width ones.
Lemma B.10 (Convergence rate of ρ − ) .
In Stage 1, we have, for any p ∈
[ r ] ,
[Equation]
Note that 1 + (cid:68) [ K A ] p , [ K B ] p (cid:69) = Θ(1) .
Hence, (cid:68) [ K A ] p , [ K B ] p (cid:69) will converge to 1 at a linear rate.
Now we consider the noise-signal ratio.
For some technical reason, instead of characterizing the dynamics of ρ N/S , we consider
[Equation]
Note that we always have
[Equation]
In other words, ρ N/S can be bounded by ρ N/S ≤ O (cid:0) dκ 0 r ˆ ρ N/S (cid:1) .
Lemma B.11 (Convergence rate of ρ N/S ) .
In Stage 1, we have
[Equation]
This lemma says the noise-signal ratio decreases exponentially fast.
Lemma B.12 (Growth rate of the condition number) .
Define ρ p/q = ∥
[ K A ] p ∥ 2 / ∥
[ K A ] q ∥ 2 .
In
Stage 1, we have
[Equation]
Note that error growth slows down as ρ − and ˆ ρ N/S decrease.
This allows us to use Lemma B.
(cf.
Section B.4).
O MITTED PROOFS OF THIS SUBSECTION
Proof of Lemma B.10.
By Corollary B.8, we have
[Equation]
Hence, by symmetry
[Equation]
we have
[Equation]
Proof of Lemma B.11.
Similar to the proof of Corollary B.7 and Corollary B.8, we have
[Equation]
Define K = K A + K B .
Then, by symmetry [Equation] we have
[Equation]
Hence, d
[Equation]
Similarly, we also have
[Equation]
Therefore
[Equation]
Then, we compute
[Equation]
Proof of Lemma B.12.
For notational simplicity, define ρ p/q := ∥ [ K A ] p ∥ 2 / ∥
[ K A ] q ∥ 2 .
By Corollary B.7, we have
[Equation]
Now we bound (cid:68) [ K A ] p , [ K B ] p (cid:69) − ⟨ K A , K B ⟩ / ( N A N B d ) .
Clear that this term is bounded by 2 .
Meanwhile, by definition, we have (cid:68) [ K A ] p , [ K B ] p (cid:69) = 1 ± ρ − .
For the second term, we have
[Equation]
Combine these together and we obtain
[Equation]
The same is also true for q .
Thus, ˙
[Equation]
B. C ONTROLLING THE DISCRETIZATION ERROR
In this subsection, we estimate the growth rate of the errors 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 ∥ [ K A ] p ∥ 2 and ∥ [ K B ] p ∥ 2 .
Instead of directly control the difference, we define
[Equation]
Note that ρ A / B ,p + ρ B / A ,p ≥ 2 , with equality obtained iff ∥
[ K A ] p ∥ 2 = ∥
[ K B ] p ∥ 2 .
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 ρ A / B ,p + ρ B / A ,p is to leverage the symmetry.
Similarly, we also define
[Equation]
and analyze ρ A / A ,ξ + ρ B / A ,ξ .
Lemma B.13 (Difference of diagonal terms) .
In Stage 1, we have
[Equation]
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 K A , ξ .
Unlike K A , for the noises, the σ ’s for different coordinates are the same.
Hence, we have the following simple bound on the growth rate of δ ξ,κ 0 .
Lemma B.14 (Condition number of K A , ξ ) .
Define ρ ξ,p/q := ∥
[ K A , ξ ] p ∥ 2 / ∥
[ K A , ξ ] q ∥ 2 .
In Stage 1, we have
[Equation]
Again, the RHS is a higher-order term that can be made arbitrarily small by choosing a small τ 2 t and maintaining a small δ A / B .
Then, we consider the orthogonality conditions.
Lemma B.15 (Orthogonality between signals) .
For any p ̸ = q , define
[Equation]
In Stage 1, we have
[Equation]
Recall that ρ − converges to 0 at a sufficiently fast rate. As a result, by Lemma B.2, ˆ δ ⊥ ,p,q will not blow up.
Meanwhile, since ˆ δ ⊥ ,p,q can be made arbitrarily small at initialization, this implies that we can make sure it is 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 p ∈ [ r ] and q ∈ [ d − r ] , define
[Equation]
and define ˆ δ ⊥ ,ξ B ,p,q similarly.
In Stage 1, we have
[Equation]
For any q, s ∈ [ d − r ] , define
[Equation]
In Stage
[Equation]
we have
[Equation]
O MITTED PROOFS OF THIS SUBSECTION
Proof of Lemma B.13 .
Note that we cannot directly use Corollary B.7 as the error term contains δ A / B , the quantity we wish to control.
However, by the proof of it, we still have
[Equation]
Interchange the roles of A and B and we obtain
[Equation]
For notational simplicity, define ρ A / B ,p = ∥ [ K A ] p ∥ 2 / ∥
[ K B ] p ∥ 2 .
Then, we compute ˙
ρ A / B ,p = dd t ∥
[ K A ] p ∥
[Equation]
A ] p
[Equation]
=
[Equation]
Then, by symmetry, we have d d
[Equation]
=
[Equation]
−
[Equation]
± O (cid:32) σ [Equation] N A N B dτ
[Equation]
t d (cid:33) ≤ O (cid:32) σ 2 p N A N B d (cid:16) δ 2 A / B + τ
[Equation]
d (cid:17)(cid:33) .
The above proof, mutatis mutandis , yields the result for ρ A / B ,ξ,q + ρ B / A ,ξ,q .
Proof of Lemma B.14.
By Corollary B.7, we have
[Equation]
Proof of Lemma B.15.
By Corollary B.8, we have
[Equation]
Interchange the roles of p
[Equation]
q and A , B and we obtain
[Equation]
Therefore,
[Equation]
[ K B ] p , [ K B ] q (cid:69) − (cid:68) [ K A ] p , [ K B ] q (cid:69)(cid:17) + 2 K 1 + K σ
[Equation]
q
N A N B d (cid:16)(cid:68) [ K A ] p , [ K A ] q (cid:69) − (cid:68) [ K A ] p , [ K B ] q (cid:69)(cid:17) ±
O [Equation]
N A N [Equation]
d (cid:19) ρ [Equation]
(cid:68) [ K A ] p , [ K [Equation] [Equation]
(cid:69) ± O (cid:18) σ [Equation] N A N [Equation]
d (cid:16) τ [Equation] t d
[Equation]
κ 0 + d (cid:113) δ A /
B (cid:17)(cid:19) .
Interchange the roles of p and q and we obtain
[Equation]
Similarly
[Equation]
we compute
[Equation]
Then, by interchanging the roles of p and q
[Equation]
we obtain (cid:28) [Equation] A N
[Equation]
(cid:16)
[Equation]
(cid:17)(cid:19) .
Add them together and we get d d t
[Equation]
[ K A ]
[Equation]
[ K A
[Equation]
=
[Equation]
Interchange the roles of p and q and we obtain d d t (cid:68) [ K B ] p , [ K B ] q (cid:69) =
2 K 1 + K
[Equation]
N
[Equation]
N
[Equation] d (cid:16)(cid:68) [ K A ] p , [ K B ] q (cid:69) − (cid:68) [ K B ] p , [ K B ] q (cid:69)(cid:17) + [Equation] K 1 + K
[Equation]
N
[Equation]
N
[Equation]
d (cid:16)(cid:68) [ K
[Equation]
[Equation]
[ K
[Equation]
[Equation]
(cid:69) − (cid:68) [ K
[Equation]
] p [Equation] [ K B ] [Equation]
(cid:69)(cid:17) ± O (cid:18) σ 2max
[Equation]
For notational simplicity, define Z 1 = (cid:68) [ K A ] p , [ K B ] q (cid:69) , Z 2 = (cid:68) [ K B ] p , [ K A ] q (cid:69) , Z 3 = (cid:68) [ K A ] p , [ K A ] q (cid:69) , and Z 4 = (cid:68) [ K B ] p , [ K B ] q (cid:69) . Also define G p = 2 K 1+ K σ 2 p N A N B d .
Then, we can summarize the above equations as
[Equation]
Note that the eigenvalues of the first matrix is − 2 G p − 2 G q , − 2 G p , − 2 G q and 0 .
Namely, it is negative semi-definite.
Thus,
[Equation]
Proof of Lemma B.16.
By Corollary B.8,
[Equation]
d
[Equation]
and (cid:28)
[Equation]
(cid:33) . Therefore, d d t (cid:68) [ K A ] p , [ K A , ξ ] q (cid:69) =
[Equation]
Similarly, we also have d d t (cid:68)
[ K B ] p , [ K A , ξ ] q (cid:69) =
2 K 1 + K σ 2 p
[Equation]
[Table]
[Equation]
Thus,
[Equation]
For the orthogonality between noises
[Equation]
by Corollary B. [Equation] we have
[Equation]
Clear that this bound holds for all other combinations.
Thus,
[Equation]
B. P ROOF OF THE MAIN LEMMA OF S TAGE
Proof of Lemma B.1.
First, we recap the estimations we have derived in previous subsections and introduce some notations.
By Lemma B.10 and Lemma B.11, we have d
[Equation]
Define ρ := max (cid:8) ρ − , ˆ ρ N/S (cid:9) to be the indicator of the progress we have made.
We have
[Equation]
[Equation]
Let ˆ κ 0 :=
max p,q ∥
[ K A ] p ∥ 2 / ∥
[ K A ] q ∥ 2 be the condition number at time t .
By Lemma B.12, we have
[Equation]
[Equation]
Now we consider the discretization errors, define
[Equation]
Note that at time t , the first condition of (10) holds with δ A / B replaced by O (ˆ δ A / B ) .
Meanwhile, by Lemma B.13, we have
[Equation]
Let ˆ δ ξ,κ 0 ( t ) be the smallest number such that the second condition of (10) holds at time t .
By
Lemma B.14, we have
[Equation]
[Equation]
Then, define
[Equation]
Clear that the last two conditions hold at time t when δ AB , ⊥ and δ ξ, ⊥ are replaced by (cid:113) ˆ δ ⊥ ( t ) .
By
Lemma B.15 and Lemma B.16, we have
[Equation]
Now, we are ready to show that the errors do not blow up in Stage 1.
Note that for all these δ ’s, we can make them arbitrarily inverse-polynomially small by choosing a sufficiently large m .
First, we consider ˆ δ A / B .
Note that the dependence of the RHS of (13) on ˆ δ A / B is quadratic.
Hence, by making the initial value of ˆ δ A / B , the RHS can be made to be dominated the τ 2 t -related terms.
Hence, ˆ δ A / B ≤ O (cid:16) σ 2max N A N B d τ 2 t dT 1 (cid:17) .
As we will see later, T 1 = poly( d ) .
Therefore, by choosing a sufficiently small τ 2 t , we can make ˆ δ A / B remain small throughout Stage 1.
Then, we consider ˆ δ ξ,κ 0 .
As we have argued earlier, the RHS of (14) can be made arbitrarily small, so that ˆ δ ξ,κ 0 remains small in Stage 1.
Now, we consider the condition number ˆ κ 0 and ˆ δ ⊥ .
For ˆ δ ⊥ , 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 τ 2 t .
The same is also true for (11).
Hence, for these quantities, we have
[Equation]
Hence, by Lemma B.2, we have
[Equation]
Note that ρ − (0) =
O (1) and ˆ ρ N/S ≤ ( d − r ) σ 2 ξ rσ 2min .
Therefore,
[Equation]
In other words, both ˆ κ 0 and ˆ δ ⊥ can at most grow √ d times.
Finally, we derive an upper bound on T 1 to complete the proof.
Similar to the proof for the condition number, one can show that ∥ [ K A ] p ∥ 2 can at most grow √ d times in Stage 1.
As a result, 1 / ( N A N B d ) is lower bounded by some 1 / poly( d ) .
Thus, by (11), T 1 ≤ poly( d ) .
B. N EGATIVE RESULTS
Lemma B.17.
There exists a σ ∈ R r satisfying the assumptions of Theorem 4.3 such that, at the end of Stage 1, the condition number of K A is d Ω(1) .
Proof.
We choose d = r and σ 21 =
c log d , σ 22 = · · · = σ 2 d = 1 .
Clear that this satisfies the condition of Theorem 4.3.
Note that it suffices to consider the infinite-width case, since, as we have proved earlier, the finite-width trajectory tracks the infinite-width one.
By Lemma B.10, we have d
[Equation]
(cid:69)(cid:17) (cid:16)
[Equation]
− (cid:68) [ K A ] p , [ K B ] p (cid:69)(cid:17) .
By the proof of Lemma B.12, we have
[Equation]
Note that, in the infinite-width case, we have
[Equation]
Therefore, (cid:68) [ K A ] p , [ K B ] p (cid:69) − ⟨ K A , K B ⟩ N A N B d ≥ 0 a and (cid:68) [ K A ] q , [ K B ] q (cid:69) − ⟨ K A , K B ⟩ N A N B d ≤ 0 for any q ≥ 2 .
Hence, ˙ ρ 1 / 2 ≥
[Equation]
[ K B ] 2 (cid:69) (cid:33) ρ 1 / 2 .
For notational simplicity, define X 1 = 1 − (cid:68) [ K A ] 1 , [ K B ] 1 (cid:69) , X 2 = 1 − (cid:68) [ K A ] 2 , [ K B ] 2 (cid:69) , Y = ρ 1 / 2 , A = 4 K 1+ K 1 N A N B d .
Then we have ˙ X 1 ≤
[Equation]
21 AX 1 , ˙
[Equation]
First, by Gronwall’s lemma, we have X 1 ( T ) ≤ exp (cid:16) − σ 21 (cid:82) T
A (cid:17) and
X [Equation] ( T [-) ≥ exp (cid:32) − 2 (cid:90) T 
-][Equation]
≥
[Equation] ( T ) .
As a result, when X 1 reaches 1 /d , we have X 2 = Ω(1) .
Let T 1 be the time X 1 reaches 1 /d
and Tthe time X 2 ( T 2 )
=
X 2 ( T 1 ) / 2 .
On [ T 1 , T 2 ] , we have ˙
ρ 1 / 2 ≥ Ω(1) σ 21 Aρ 1 / 2 .
By Gronwall’s lemma, in order for X 2 to half, we need exp( − 2 (cid:82) T 2 T 1 A ) = 1 / 2 .
Hence,
[Equation]
C S TAGE
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, we prove the following.
Lemma C.1 (Stage 2) .
Suppose that at the beginning of Stage 2, we have κ 0 ≤ √ d and all errors mentioned in (16) are sufficiently small 5 .
Let c Target > 1 be a constant.
Let T 1 be the earliest time that
[Equation]
We have T 1 ≤ poly( d ) .
Moreover, throughout Stage 2, we have
[Equation]
[Table]
[Equation]
(16) where the δ ’s are some small 1 / poly( d ) values.
The rest of this section is organized as follows.
We derive estimations for the Q -matrices in Section C.1.
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. E STIMATIONS FOR Q
As in Stage 1, we estimate the Q -matrices in this subsection.
The analysis here will be more complicated than the one in Section B.1 since now τ 2 t is no longer close to 0 , and we cannot simply approximate S A and S B 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 infinite-width part to factor out the first-order error terms.
Then, we evaluate and simplify these first-order terms with E z − and E z ± .
First, we need the following lemma which gives closed-form formulas for some expectations we will encounter later.
Lemma C.2.
Define ⟨ z + , z − ⟩ ˆ κ 2 :=
(cid:80) rk {+=
+}
ˆ κ 2 k z
+ k z − k and T p :=
tanh (cid:16) ˆ κ 2 p /d N A N B (cid:17) , p ∈
[ r ] .
For any p ̸ = q ∈ [ r ] , we have
[Equation]
Then
[Equation]
we derive estimations for S A and S B .
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 ⊤ A 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
[Equation]
˜
[Equation]
In Stage
[Equation]
we have
[Equation]
and
[Equation]
Then
[Equation]
we consider the error comes from the non-diagonalness of K ⊤ A K B .
Lemma C.4 (Further estimations for S ) .
Define ˜
[Equation]
+ , − = (cid:88)
[Equation]
In Stage 2, we have
[Equation]
With the above two lemmas in hand, we can now derive estimations for the Q -matrices.
Lemma C.5 (Estimations for Q 1 ) .
Define K AB = K ⊤
A K B and K BA
= K ⊤ B K A .
In Stage 2, for any p ̸ = q ∈ [ r ] , we have [ [Equation] ] p,p
[Equation]
[Equation]
In particular, we have
[Equation]
Note that the diagonal term has a zero-order term, i.e., it is not proportional to some error.
That is the signal term.
On the other hand, all off-diagonal terms depend on [ K AB ] p,q ( p ̸ = q ).
Recall that the dynamics of K A and K B can be described using these Q -matrices.
Therefore, for any p ̸ = q , we can have equations of form dd t [ K AB ] p,q ≈ G ([ K AB ] p,q , [ K AB ] q,p ) , 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 Q 1 ,ξ and Q 2 .
The difference here is that for these two matrices, we do not have signal terms.
Lemma C.6 (Estimations for Q 1 ,ξ ) .
In Stage 2, for any p ∈ [ d − r ] and q ∈ [ r ] , we have
[Equation]
In particular, we have
[Equation]
Lemma C.7 (Estimations for Q 2 ) .
In Stage 2, for any p, q ∈ [ d − r ] , we have
[Equation]
Lemma C.8 (Estimations for Q 0 ) .
In Stage 2, we have
[Equation]
Finally
[Equation]
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 subsections.
Corollary C.9 (Dynamics of the norms) .
In Stage
[Equation]
we have
[Equation]
The formulas for ∥ [ K B ] p ∥ and ∥ [ K B , ξ ] q ∥ can be obtained by interchanging the roles of A and B [Equation]
O MITTED PROOFS OF THIS SUBSECTION
Proof of Lemma C
[Equation] First
[Equation]
we compute
[Equation]
(cid:18) ˆ κ
[Equation]
k /d
N
[Equation]
N
[Equation]
(cid:19) .
Similarly, we also have
[Equation]
N
[Equation]
N
[Equation]
Each factor in (cid:81) k
̸ = p is still cosh (cid:16) ˆ κ 2 k /d N A N B (cid:17) .
For the first term, we have
E
[Equation]
Therefore,
E
[Equation]
The above calculation, mutatis mutandis , also yields the last identity.
Proof of Lemma C.3.
First, we write
[Equation]
Also note that the middle terms are O ( d 2 δ ξ, ⊥ δ N/S ) .
Then, we compute
[Equation]
Similar results also hold for other combinations of ± .
With the notations defined in this lemma, we can write these results as
[Equation]
To compute S A and S B , we then need to take expectations over the negative examples.
Note that by taking expectation over ξ − B , the term E + , − δ + ,ξ − becomes 0 .
Therefore, we have
[Equation]
Unfortunately, the same argument does not apply to δ ξ + , − since both E + , − and δ ξ + , − depend on z − .
However, it is still possible to further simplify the expression.
First, we write
[Equation]
Recall Lemma C.2.
Then, we compute
[Equation]
Hence,
[Equation]
Similarly, we also have
[Equation]
Recall that exp (cid:0) f
+ A · f + B (cid:1) =
E
+ , + (1 + δ + ,ξ + + δ ξ + , + ) ± O (cid:16) d 2 δ ξ, ⊥ δ 2 N/S (cid:17) . Hence, we have S A ( x + A , x + B ) = E + , + (1 + δ + ,ξ + + δ ξ + , + )
[Equation]
+
[Equation]
+
[Equation]
+
[Equation]
+
[Equation]
AB
[Equation]
⊥
[Equation]
S
[Equation]
δ
[Equation]
⊥ δ N/S (cid:1) =
˜
[Equation]
1 −
[Equation]
[Equation]
+ ,
[Equation]
−
[Equation]
δ
[Equation]
+ ˜
[Equation]
[Equation]
δ
[Equation]
δ
[Equation]
δ ξ
[Equation]
δ
[Equation]
Similarly, we also have
[Equation]
Then, we compute
[Equation]
and
[Equation]
Proof of Lemma C.4.
We write
[Equation]
Note that, as a special case, we have I + , + = I − , − = I 0 .
In other words, I + , + and I − , − do not depend on the actual value of z ± .
Also note that I + , − is bounded by O ( dδ AB , ⊥ ) .
Then, we compute
[Equation] and E +
[-( I + , − ) (cid:16) 1 + ˜ δ + , − ± O (cid:0) d 2 δ 2 AB , ⊥ (cid:1)(cid:17) and E + , + = exp( I 0 ) (cid:16) 
 + ˜ δ + , + ± O (cid:0)-][Equation]
+ = exp [Equation]
Take expectation over z − and we obtain
[Equation]
By Lemma C.2, the first term is Z c .
For the second term, we compute
[Equation]
where the second line again comes from Lemma C.2.
Hence, we have
[Equation]
Then, for ˜ S A , we have
[Equation]
Similarly, we also have
[Equation]
Then, we compute
[Equation]
Similarly, we also have
[Equation]
Proof of Lemma C.5.
Recall that
[Equation]
Note that there is no ξ here other than the ones in the coefficient.
As a result, all terms contain δ + , ξ + and alike are 0 .
Hence, we have
[Equation]
Now we estimate each of these three terms separately.
We also deal with the diagonal and offdiagonal terms separately.
By Lemma C.4, we have
[Equation]
Note that we use the fact that all these δ ’s have mean 0 .
Also by Lemma C.4, we have T
[Equation]
[Equation]
Note that ˜ δ + , + , ˜ δ + , − and ˜ δ + ,T + only contain cross terms of form
z
+ i z ±
j with i ̸ =
j .
Hence, the second term is 0 .
Meanwhile, by Lemma C.2, the first term is
[Equation]
As a result
[Equation]
Similarly, one can show that T 3 ([ Q 1 ] p,p ) =
− (1 − ˜ S )
T p ±
O (cid:16) d 2 δ 2 AB , ⊥ (cid:17) also holds.
Thus,
[Equation]
Now, we consider the off-diagonal terms.
For notational simplicity, we define K AB = K ⊤ A K B .
For any p ̸ = q , we compute
[Equation]
Clear that the summand is nonzero only if ( i, j ) =
( p, q ) or ( i, j ) =
( q, p ) .
Hence,
[Equation]
Then, for T 2 , we compute
[Equation]
Again, the summand is nonzero only if ( i, j ) =
( p, q ) or ( i, j ) =
( q, p ) .
By Lemma C.2, we have
T 2 ([ Q 1 ] p,q ) =
− ˜
SK ˜ E
[Equation]
Similarly, for T 3 , we have
[Equation]
Combine these together and we obtain
[Equation]
Proof of Lemma C.6.
We write
[Equation]
We will use the fact that if some quantity X does not depend on ξ , then E { X ξ } = 0 to simplify these terms.
For T 1 , by Lemma C.3, we have
[Equation]
S A ( x + A , x + B ) + S B ( x + A , x + B ) (cid:1) [ ξ + B ] p z + q d (cid:9) = − ˜ S (1 − ˜ S ) E (cid:8) (2 δ + , ξ + + 2 δ ξ + , + − δ ξ + ,T + − δ T + , ξ + ) [ ξ + B ] p z + q d (cid:9) ± O (cid:0) d 3 (cid:0) δ AB , ⊥ + δ N/S (cid:1) δ ξ, ⊥ δ N/S (cid:1) =
− ˜ S (1 − ˜ S ) E (cid:8) (2 δ + , ξ + − δ T + , ξ + ) [ ξ + B ] p z + q d (cid:9) ± O (cid:0) d 3 (cid:0) δ AB , ⊥ + δ N/S (cid:1) δ ξ, ⊥ δ N/S (cid:1) . Recall that δ + ,ξ + = (cid:88) i,j [ K ⊤ A K B , ξ ] i,j N A N B z
+ i [ ξ + B ] j , δ T + , ξ + = (cid:88) i,j [ K ⊤ A K B , ξ ] i,j N A N B T i z
+ i [ ξ + B ] j .
Hence, we can further rewrite T 1 ([ Q 1 ,ξ B ] p,q ) as
T 1 ([ Q 1 ,ξ B ] p,q ) = − ˜ S (1 − ˜ S ) (cid:88)
i,j [ K ⊤ A K B , ξ ] i,j N A N B (2 − T i ) E (cid:8) z + i [ ξ + B ] j [ ξ + B ] p z + q d (cid:9) ±
O (cid:0) d 3 (cid:0) δ AB , ⊥ + δ N/S (cid:1) δ ξ, ⊥ δ N/S (cid:1) .
Note that the summand is nonzero only if i = q and j = p .
Thus,
[Equation]
Similarly, we also have
[Equation]
Then, we write
[Equation]
If j ̸ = p
[Equation]
clear that that summand is 0 .
If i ̸ = q , then by flipping the sign of z ± q simultaneously, we flip the sign of ˜ E + , − z
+ i z
+ q .
Therefore, when i ̸ = q , the summand is also 0 .
Thus,
[Equation]
Similarly, for T 3 , we compute
[Equation]
Thus
[Equation]
− (1 − ˜ S ) ⟨ [ K A ] q , [ K B , ξ ] p ⟩ N [Equation] N
[Equation]
AB
[Equation]
⊥
[Equation]
S [Equation] Combine these together, rearrange terms and we obtain [ Q 1 , ξ B ] p,q = − (1 − ˜ S ) (cid:16) 1 + ˜ S + (1 − ˜ S ) T 2 q (cid:17) ⟨ [ K A ] q , [ K B , ξ ] p ⟩ N
A N B d ± O (cid:0) d
[Equation]
δ AB , ⊥ + δ N/S (cid:1) δ ξ, ⊥ δ N/S (cid:1) .
To obtain the formula for Q 1 , ξ A , it suffices to interchange the roles of A and B .
Proof of Lemma C.7.
Recall that
[Equation]
[Equation]
We now estimate each of these three terms.
Again, the strategy is to leverage the symmetry of the distribution of ξ to argue that some part of the expectation is 0 .
For T 1 , we have
[Equation]
Note that none of these δ ’s depends on both ξ + A and ξ + B .
Therefore, the first term is 0 .
Similarly, for
T 2 , we have
[Equation]
The same is also true for T 3 .
Thus,
[Equation]
Proof of Lemma C.8.
Recall from Corollary A.4 that Q 0 is defined as
[Equation]
We write
[Equation]
Therefore
[Equation]
we have
[Equation]
[-= κ 2 k (cid:107) κ (cid:107) 
-]
Proof of Corollary C.9.
Recall from Lemma A.5 that
[Equation]
[Equation]
By Lemma C.
[Equation]
we have T
[Equation]
By Lemma C.6, we have
[Equation]
Combine these together and we obtain
[Equation]
To get the formula for d d t ∥ [ K B ] p ∥ [Equation] it suffices to interchange the roles of A and B .
Similarly, for the noises
[Equation]
we write
[Equation]
By Lemma C.6, we have
[Equation]
By Lemma C [Equation]
[Equation]
we have
[Equation]
Combine these together and we get
[Equation]
C. M AINTAINING THE ORTHOGONALITY
In this subsection [Equation] we control the growth of δ AB , ⊥ and δ ξ, ⊥ .
First, we derive the equations that govern the evolution of the off-diagonal terms.
Lemma C.10.
In Stage 2, for any p ̸ = q ∈ [ r ] , we have
[Equation]
Lemma C.11.
In Stage 2, for any p ̸ = q ∈
[ r ] , we have
[Equation]
and d
[Equation]
Now, we are ready to control the off-diagonalness.
The proof is similar to the one of Lemma B.15.
To provide intuitions, we first consider an idealized case.
That is, we assume here the noise-signal ratio is 0 and K A = K B , and explain how to control (cid:68) [ K A ] p , [ K B ] q (cid:69) for p ̸ = q ∈ [ r ] under these assumptions.
In this case, we have
[Equation]
and [ [Equation] Note that the coefficient is negative. As a result, (cid:68) [ K A ] p , [ K B ] q (cid:69) will move towards 0 .
When K A =
K B is not exactly true, the situation is trickier because we need to deal with the middle four terms of the RHS of dd t (cid:68) [ K A ] p , [ K B ] q (cid:69) [-
-]. The idea is to view all these off-diagonal 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 p ̸ = q , define ˆ δ ⊥ ,p,q :=
(cid:68) [ [Equation] A
[Equation]
In Stage 2, we have
[Equation]
Note that the LHS is of order δ 2 AB , ⊥ and for the RHS, the only term that can potentially have the same order is the
[Equation]
-related term.
We will show later that δ − =
o ( δ AB , ⊥ ) , whence this is also a higher order term.
Then, we consider the orthogonality between signals and noises.
Lemma C.13 (Orthogonality between signals and noises) .
For any p ∈ [ r ] and q ∈ [ d − r ] , define
[Equation]
In Stage 2, we have
[Equation]
Finally, we deal with the orthogonality between the noises.
Lemma C.14 (Orthogonality between noises) .
In Stage 2, for any p, q ∈ [ d − r ] , we have
[Equation]
O MITTED PROOF OF THIS SUBSECTION
Proof of Lemma C.10.
Recall that
[Equation]
Then
[Equation]
we write (cid:28)
[Equation]
For the second term, by Lemma C.
[Equation]
we have
[Equation]
Hence,
[Equation]
For the first term, we have
[Equation]
When k / ∈ { p, q } , we have (cid:12)(cid:12)(cid:12)(cid:68)
[ K A ] p , [ K A ] k (cid:69)(cid:12)(cid:12)(cid:12) ≤
O ( δ AB , ⊥ ) , (cid:12)(cid:12)(cid:12)(cid:68)
[ K B ] q , [ K A ] k (cid:69)(cid:12)(cid:12)(cid:12) ≤
O ( δ AB , ⊥ ) , and ∥ [ K A ] k ∥ / ∥ [ K B ] q ∥ ≤ κ 0 .
Meanwhile, by Lemma C.5, we also have | [ Q 1 ] q,k | ≤ O ( δ AB , ⊥ ) .
Hence,
[Equation]
k
[Equation]
[Equation]
(cid:69) (cid:68) [
[Equation]
[Equation]
[Equation]
(cid:69)(cid:17) ∥ [ K A ] k ∥ [ Q 1 ] q [Equation]
k ∥ [
[Equation]
Therefore,
[Equation]
For the first term [Equation] we have
[Equation]
For the second term, we have
[Equation]
Thus, (cid:28)
[Equation]
=
[Equation]
Hence,
[Equation]
Interchange the roles of A , B , p, q , replace Q 1 with Q ⊤ [Equation] and we obtain
[Equation]
Combine these together and we complete the proof.
Proof of Lemma C.11.
Similar to the proof of the previous lemma, we compute
[Equation]
Again, we have
[Equation]
Therefore,
[Equation]
Interchange the roles of p
[Equation]
q
and we obtain
[Equation]
Combine these together and we get
[Equation]
Interchange the roles of A , B , replace Q 1 with Q ⊤ 1 , and we obtain the formula for dd t (cid:68) [ K B ] p , [ K B ] q (cid:69) .
Proof of Lemma C.12.
First, we consider the [ Q 1 ] p,q -related terms, by Lemma C.5, we have
[Equation]
Hence,
[Equation]
The same bound also hold for other [ Q 1 ] p,q -related terms.
The important thing here is that we have and additional 1 / √ d factor.
Now, we are ready to prove the result.
For notational simplicity, define Z 1 = (cid:68) [ K A ] p , [ K B ] q (cid:69) ,
Z 2 = (cid:68) [ K B ] p , [ K A ] q (cid:69) , Z 3 = (cid:68) [ K A ] p , [ K A ] q (cid:69) , and Z 4 = (cid:68) [ K B ] p , [ K B ] q (cid:69) .
Also define G p :=
[ Q 1 ] p,p σ 2 p / ( N A N B d ) and similarly for G q .
Then, we can write the results of Lemma C.10 and
Lemma C.11 as
[Equation]
The eigenvalues of the first matrix are − 2 G p , − 2 G q , − 2 G p − 2 G q and 0 .
For the first three eigenvalues, note that
[Equation]
Hence, the signals will dominate the noises, in particular, the first term on the second line, and push ∥ Z ∥ toward 0 .
Now we consider the eigen-pair (0 , (1 , 1 , 1 , 1)) , for which we will use the actual form of [ Q 1 ] p,q we obtained in Lemma C.5.
We have
[Equation]
By Lemma C
[Equation]
5, we have
[Equation]
Then, we write
[Equation]
Hence,
[Equation]
To convert Z 1 + Z 2 to (cid:80) 4 i =1 Z i .
Note that we have
[Equation]
Therefore,
[Equation]
Thus,
[Equation]
Note that the coefficient of the first term is negative.
Combine this with the previous bound for ∥ Z ∥ , and we complete the proof.
Proof of Lemma C.13.
Recall that
[Equation]
A ] q ∥
[ K A , ξ ] q ∥ + [ K B , ξ Q 2 ] q ∥ [ K A , ξ ] q ∥ (cid:33) σ 2 ξ N A N B d.
We now compute (cid:68) dd t [ K A ] p , [ K A , ξ ] q (cid:69) and (cid:68) [ K A ] p , dd t [ K A , ξ ] q (cid:69) separately.
First, we write (cid:28) d d t [ K A ] p , [ K A , ξ ] q (cid:29) = [ K A ,
[Equation]
] q ⊤ (cid:18) I
[Equation]
(cid:16) [
[Equation] [Equation]
(cid:17)
[Equation]
(cid:19) [ K
[Equation]
1 ] p
[Equation]
[ K A ] p ∥ σ 2 p N A N
[Equation]
d + [ K A , ξ ] q ⊤ (cid:18)
I − [ K A ] p (cid:16) [ K A ] p (cid:17) ⊤ (cid:19) [ K B , ξ Q 1 ,ξ B ] p ∥ [ K
A ] p ∥ σ 2 p
N A N B d =: T(cid:18)(cid:28) d d t [ K A ] p , [ K A , ξ ] q (cid:29)(cid:19) + T(cid:18)(cid:28)
d d t [ K A ] p , [ K A , ξ ] q (cid:29)(cid:19) .
Then, we compute
[Equation]
and, by Lemma C.6, we have
[Equation]
[Equation]
Combine these together and we get
[Equation]
Then
[Equation]
we compute (cid:68) [ K A ] p , dd t [ K A , ξ ] q (cid:69) .
We write (cid:28)
[Equation]
T(cid:18)(cid:28) [ [Equation] (cid:29)(cid:19)
[Equation] (cid:18)(cid:28) [ K
[Equation]
[Equation]
[Equation]
(cid:29)(cid:19) . For T 1 , when k = p , by Lemma C.6, we have [
[Equation]
[Equation]
[Equation]
A ]
[Equation]
p =
[Equation]
=
[Equation]
δ N/S (cid:1) .
When k ̸ = p , we have [ K A ] p ⊤ (cid:18) I − [ K A , ξ ] q (cid:16) [ K A , ξ ] q (cid:17) ⊤ (cid:19) [ K B ] k
[ Q 1 , ξ A ] q,k ∥ [ K
A , ξ ] q ∥ =
O ( κ 0 δ AB , ⊥ δ ξ, ⊥ ) .
Hence,
[Equation]
Now we consider T 2 .
By Lemma C. [Equation] we have
[Equation]
Therefore,
[Equation]
Thus
[Equation]
Similarly
[Equation]
one can show that
[Equation]
For notational simplicity
[Equation]
define
[Equation]
Then, we can write d
[Equation]
The first matrix is negative semi-definite, whence
[Equation]
Since the roles of K A , ξ and K B , ξ are interchangeable, the same bound also holds for (cid:68) [ K A ] p , [ K B , ξ ] q (cid:69) and (cid:68) [ K A ] p , [ K B , ξ ] q (cid:69) .
Proof of Lemma C.14.
We write
[Equation]
Then
[Equation]
we compute
[Equation]
Note that, the first part of each summand is bounded by O ( δ ξ, ⊥ ) .
For the second part, by Lemma C. we also have
[Equation]
Hence,
T
[Equation]
Now we consider T 2 .
By Lemma C
[Equation]
7, We have
[Equation]
Combine this together and we obtain
[Equation]
Similarly, one can derive the same bound for (cid:68) dd t [ K A , ξ ] p , [ K B , ξ ] q (cid:69) and complete the proof.
C. M AINTAINING K A ≈ K B
In this subsection, we show that (cid:68) [ K A ] p , [ K B ] p (cid:69) ≈ 1 , ∥ [ K A ] p ∥ ≈ ∥ [ K B ] p ∥ , and also ∥ [ K A , ξ ] q ∥ ≈ ∥
[ K B , ξ ] q ∥ throughout Stage 2.
The proofs are similar to the corresponding ones in Stage 1.
Lemma C.15.
In Stage 2, we have d
[Equation]
⊥
[Equation]
Lemma C.16.
Define ρ A / B ,p :=
∥
[ K A ] p ∥ 2 / ∥ [ K B ] p ∥ 2 and ρ B / A ,p := ∥ [ K B ] p ∥ 2 / ∥
[ K A ] p ∥ 2 .
In Stage 2, we have
[Equation]
Lemma C.17.
For any p, q ∈ [ d − r ] , define ρ ξ, A / A ,p/q = ∥ [ K A , ξ ] p ∥ 2 / ∥
[ K A , ξ ] q ∥ 2 and ρ ξ, A / B ,p/q = ∥ [ K A , ξ ] p ∥ 2 / ∥ [ K B , ξ ] q ∥ 2 .
In Stage 2, we have
[Equation]
A ,
[Equation]
O MITTED PROOF OF THIS SUBSECTION
Proof of Lemma C.15.
First, we write (cid:28) d
[Equation]
For T 1 , we compute
T(cid:18)(cid:28) d d
[Equation]
k
[Equation]
⊤ (cid:19) [ K B ] k ∥
[ K B ] k ∥
[ Q 1 ] k,p ∥
[ K
A ] p ∥ σ 2 p
N A N B d = (cid:18) 1 − (cid:68) [ K A ] p , [ K B ] p (cid:69) 2 (cid:19) ∥ [ K B ] p ∥
[ Q 1 ] p,p ∥
[ K A ] p ∥ σ 2 p
N A N B d ± O (cid:18) σ 2max
N A N B ddκ
δ 2 AB , ⊥ (cid:19) .
For T 2 , we compute
T(cid:18)(cid:28)
d d t [ K A ] p , [ K B ] p (cid:29)(cid:19) =
d − r (cid:88) k
= [ K B ] p ⊤ (cid:18)
I − [ K A ] p (cid:16) [ K A ] p (cid:17) ⊤ (cid:19)
[ K B , ξ ] k [ Q 1 ,ξ B ] k,p ∥ [ K
A ] p ∥ σ 2 p
N A N B d = ± O (cid:18) σ 2 max
N A N B ddδ 2 N/S δ 2 ξ, ⊥ (cid:19) .
Therefore, (cid:28) d d t [ K A ] p , [ K B ] p (cid:29) =
(cid:18) 1 − (cid:68) [ K A ] p , [ K B ] p (cid:69) 2 (cid:19) ∥ [ K B ] p ∥
[ Q 1 ] p,p ∥
[ K A ] p ∥ σ 2 p
N A N B d ± O (cid:18) σ 2max
N A N B d dκ 0 δ 2 AB , ⊥ (cid:19) .
Then, by symmetry, we have
[Equation]
Proof of Lemma C.16.
Similar to Stage [Equation] we define ρ A / B ,p = ∥ [ K A ] p ∥ 2 / ∥
[ K B ] p ∥ 2 .
By
Lemma A.
[Equation]
we have
[Equation]
Then, by Lemma C
[Equation]
5 and Lemma C [Equation]
[Equation]
we have
[Equation]
Thus,
[Equation]
[-N A N B ddκ 
 δ 2 AB , ⊥ (cid:19) .-] Interchange the roles of A , B and we get
[Equation]
[-N A N B ddκ 
 δ 2 AB , ⊥ (cid:19) .-] Hence,
[Equation]
Proof of Lemma C.17.
For notational simplicity
[Equation]
we drop the subscript ξ in the proof.
Recall from
Corollary C.
that
[Equation]
Hence, for any p, q ∈
[ d − r ] , we have
[Equation]
Similarly, we have
[Equation]
By symmetry
[Equation]
we also have
[Equation]
Hence,
[Equation]
C. C ONTROLLING THE NOISE - SIGNAL RATIO
In this subsection
[Equation]
we show that the noise-signal ratio remains small throughout Stage 2.
Lemma C.18.
Let ∥ [ K A ] p ∥ be the smallest among all {∥ [ K A ] k ∥} k ∈ [ r ] .
For any q ∈ [ r ] , in Stage 2, we have
[Equation]
Proof.
Recall from Corollary C
[Equation]
9 that
[Equation]
Since the condition number of K A is bounded by √ d [Equation] it suffices to consider the smallest ∥ [ K A ] p ∥ , for which we have
[Equation]
where the last line comes from Lemma C.8.
By Lemma B.4, [ Q 1 ] p,p is negative correlated with κ 2 p .
As a result, we have
[Equation]
For K A , ξ , we simply have
[Equation]
Thus, d
[Equation]
C. E STIMATING THE CONVERGENCE RATE
In this subsection, we estimate how fast the condition number will become close to 1 .
Lemma C.19.
Suppose that ∥ [ K A ] p ∥ is the largest and ∥ [ K A ] q ∥ the smallest among all {∥ [ K A ] k ∥} k ∈ [ r ] .
In Stage 2, we have
[Equation]
Corollary C.20 (Convergence rate) .
Suppose that ∥ [ K A ] p ∥ is the largest and ∥ [ K A ] q ∥ the smallest among all {∥ [ K A ] k ∥} k ∈ [ r ] .
For any constant c > 1 , it takes at most poly( d ) amount of time for ∥ [ K A ] p ∥ 2 / ∥
[ K A ] q ∥ 2 to become smaller than c .
O MITTED PROOF OF THIS SUBSECTION
Proof of Lemma C.19.
By Corollary C [Equation] [Equation] Lemma C.8 and Lemma C.
[Equation]
we have d d
[Equation]
where ˜ T = (cid:80) rk =1 κ 2 k ∥ κ ∥ 2 T p .
Then, we compute
[Equation]
Since ∥ [ K A ] p ∥ is the largest, ∥ [ K A ] q ∥ is the smallest, T p is positively correlated with ∥ [ K A ] p ∥ , and ˜ T is a weighted average of T p , we have σ
[Equation]
p (cid:16) T [Equation] ˜ T (cid:17) − σ
[Equation]
q (cid:16) T q − ˜ T (cid:17) ≥ σ 2 q ( T p − T q ) .
Thus, d d t
[Equation]
Proof of Corollary C.20.
Recall that
[Equation]
Since κ 0 ≤ √ d , we have ∥ [ K A ] p ∥ 2 / ∥ [ K A ] F ∥ 2 ≤ 1 / 2 .
Note that tanh ′ ( z ) = 1 − tanh 2 ( z ) = Ω(1) for any z ≤ 1 . 1 / 2 .
Therefore,
[Equation]
Then, by Lemma C.19, we have
[Equation]
By the proof of Lemma C.19, the largest ∥ [ K A ] p ∥ 2 is non-increasing.
Hence, N A N B d is upper bounded by some poly( d ) .
Thus, it takes at most poly( d ) for (cid:13)(cid:13)
[ K A ] 2 p (cid:13)(cid:13) / ∥ [ K A ] q ∥ to become smaller than c .
C. P ROOF OF THE MAIN LEMMA OF S TAGE
Proof of Lemma C.1.
The polynomial bound on the convergence time has been proved in Corollary C.20.
For the errors, recall from Lemma C.12, Lemma C.13 and Lemma C.14 that
[Equation]
[-ˆ δ ξ, ⊥ ,p,q ≤ O (cid:18) σ 2max N A N B ddκ 
 δ 2 ξ, ⊥ (cid:16) δ ξ, ⊥ + (cid:112) δ − (cid:17)(cid:19) , d d t (cid:68) [ K A , ξ ] p , [ K B , ξ ] q (cid:69) = ±-]Recall from Lemma C.
[Equation]
Lemma C.
[Equation]
and Lemma C.17 that
[Equation]
By Lemma C.18, we have
[Equation]
Note that on the RHS of these equations, the only terms whose order may potentially be smaller than or equal to LHS are the δ − -reltaed terms.
However, (17), we can make sure δ − is at most δ 1 . 5 AB , ⊥ .
As a 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 if they are sufficiently small at the beginning of
Stage 2.
D F ROM G RADIENT F LOW TO G RADIENT D ESCENT
Converting the above gradient flow argument to a gradient descent one is standard.
All our estimations can tolerate an 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 always 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 are both polynomial, this also implies a polynomial sample complexity.