# Understanding the Role of Nonlinearity in Training Dynamics of Contrastive
Learning

Yuandong Tian
Meta AI Research
<EMAIL_ADDRESS>
###### Abstract

While the empirical success of self-supervised learning (SSL) heavily relies
on the usage of deep nonlinear models, many theoretical works proposed to
understand SSL still focus on linear ones. In this paper, we study the role of
nonlinearity in the training dynamics of contrastive learning (CL) on one and
two-layer nonlinear networks with homogeneous activation
$h(x)=h^{\prime}(x)x$. We theoretically demonstrate that (1) the presence of
nonlinearity leads to many local optima even in 1-layer setting, each
corresponding to certain patterns from the data distribution, while with
linear activation, only one major pattern can be learned; and (2) nonlinearity
leads to specialized weights into diverse patterns, a behavior that linear
activation is proven not capable of. These findings suggest that models with
lots of parameters can be regarded as a _brute-force_ way to find these local
optima induced by nonlinearity, a possible underlying reason why empirical
observations such as the lottery ticket hypothesis hold. In addition, for
2-layer setting, we also discover _global modulation_ : those local patterns
discriminative from the perspective of global-level patterns are prioritized
to learn, further characterizing the learning process. Simulation verifies our
theoretical findings.

##  1 Introduction

Over the last few years, deep models have demonstrated impressive empirical
performance in many disciplines, not only in supervised but also in recent
self-supervised setting (SSL), in which models are trained with a surrogate
loss (e.g., predictive (Devlin et al., 2018; He et al., 2021), contrastive
(Chen et al., 2020; Caron et al., 2020; He et al., 2020) or noncontrastive
loss (Grill et al., 2020; Chen & He, 2020)) and its learned representation is
then used for downstream tasks.

From the theoretical perspective, understanding the roles of nonlinearity in
deep neural networks is one critical part of understanding how modern deep
models work. Currently, most works focus on linear variants of deep models
(Jacot et al., 2018; Arora et al., 2019a; Kawaguchi, 2016; Jing et al., 2022;
Tian et al., 2021; Wang et al., 2021). When nonlinearity is involved, deep
models are often treated as richer families of black-box functions than linear
ones (Arora et al., 2019b; HaoChen et al., 2021). The role played by
nonlinearity is also studied, mostly on model expressibility (Gühring et al.,
2020; Raghu et al., 2017; Lu et al., 2017) in which specific weights are found
to fit the complicated structure of the data well, regardless of the training
algorithm. However, many questions remain open: if model capacity is the key,
why traditional models like $k$-NN (Fix & Hodges, 1951) or kernel SVM (Cortes
& Vapnik, 1995) do not achieve comparable empirical performance, even if
theoretically they can also fit any functions (Hammer & Gersmann, 2003;
Devroye et al., 1994). Moreover, while traditional ML theory suggests
carefully controlling model capacity to avoid overfitting, large neural models
often generalize well in practice (Brown et al., 2020; Chowdhery et al.,
2022).

In this paper, we study the critical role of nonlinearity in the training
dynamics of contrastive learning (CL). Specifically, we show that even with
1-layer nonlinear networks, the nonlinearity plays a critical role by creating
many local optima, each corresponding to a pattern in the input. As a result,
the more nonlinear nodes in 1-layer networks with different initialization,
the more local optima are likely to be collected in the trained weights, and
the richer the resulting representation becomes. In contrast, in the linear
setting, contrastive learning becomes PCA under certain conditions (Tian,
2022), and only the most salient pattern (i.e., the maximal eigenvector of the
data covariance matrix) is learned while other less salient ones are lost,
regardless of the number of hidden nodes.

Based on this finding, we then extend our analysis to 2-layer ReLU setting
with non-overlapping receptive fields. In this setting, we prove the
fundamental limitation of linear networks: the gradients of multiple weights
at the same receptive field are always co-linear, preventing diverse pattern
learning. In contrast, with nonlinear activations, these weights can
specialize to different features, justifying why more parameters help.

Finally, we also characterize the interaction between layers in 2-layer
network: while in each receptive field, many patterns exist, those
contributing to global patterns are prioritized to learn by the training
dynamics. This _global modulation_ changes the eigenstructure of the low-level
covariance matrix so that relevant patterns are learned with higher
probability.

In summary, through the lens of training dynamics, we discover unique roles
played by nonlinearity which linear activation cannot do: (1) nonlinearity
creates many local optima for different patterns of the data (2) nonlinearity
enables weight specialization to diverse patterns. In addition, we also
discover a mechanism for how global pattern prioritizes which local patterns
to learn, shedding light on the role played by network depth. Preliminary
experiments on simulated data verify our findings.

Related works. Works (Wilson et al., 1997; Li & Yuan, 2017; Tian et al., 2019;
Tian, 2017; Allen-Zhu & Li, 2020) that analyze training dynamics mostly focus
on supervised learning. Different from (Saunshi et al., 2022; Ji et al., 2021)
that analyzes feature learning process in linear models of CL, we focus on the
critical role played by nonlinearity. Our analysis is also more general than
(Li & Yuan, 2017) that focuses on 1-layer ReLU network with symmetric weight
structure trained on sparse linear models.

##  2 Problem Setup

Notation. In this section, we introduce our problem setup of contrastive
learning. Let ${\bm{x}}_{0}\sim p(\cdot)$ be a sample drawn from the dataset,
and ${\bm{x}}\sim p_{\mathrm{aug}}(\cdot|{\bm{x}}_{0})$ be a augmentation view
of the sample ${\bm{x}}_{0}$. Here both ${\bm{x}}_{0}$ and ${\bm{x}}$ are
random variables. Let ${\bm{f}}={\bm{f}}({\bm{x}};\bm{\theta})$ be the output
of a deep neural network that maps input ${\bm{x}}$ into some representation
space with parameter $\bm{\theta}$ to be optimized. Given a batch of size $N$,
${\bm{x}}_{0}[i]$ represent $i$-th sample (i.e., instantiation) of
corresponding random variables, and ${\bm{x}}[i]$ and ${\bm{x}}[i^{\prime}]$
are two of its augmented views. Here ${\bm{x}}[\cdot]$ has $2N$ samples,
$1\leq i\leq N$ and $N+1\leq i^{\prime}\leq 2N$.

Contrastive learning (CL) aims to learn the parameter $\bm{\theta}$ so that
the representation ${\bm{f}}$ are distinct from each other: we want to
maximize squared distance $d_{ij}^{2}:=\|{\bm{f}}[i]-{\bm{f}}[j]\|_{2}^{2}/2$
between samples $i\neq j$ and minimize
$d_{i}^{2}:=\|{\bm{f}}[i]-{\bm{f}}[i^{\prime}]\|_{2}^{2}/2$ between two views
${\bm{x}}[i]$ and ${\bm{x}}[i^{\prime}]$ from the same sample
${\bm{x}}_{0}[i]$.

Many objectives in contrastive learning have been proposed to combine these
two goals into one. For example, InfoNCE (Oord et al., 2018) minimizes the
following (here $\tau$ is the temperature):

$\\!\\!\mathcal{L}_{nce}\\!:=\\!-\tau\sum_{i=1}^{N}\log\frac{\exp(-d^{2}_{i}/\tau)}{\epsilon\exp(-d^{2}_{i}/\tau)\\!+\\!\sum_{j\neq
i}\exp(-d^{2}_{ij}/\tau)}$ (1)

In this paper, we follow (Tian, 2022), which models contrastive learning with
a broad family of loss functions (including InfoNCE) as a coordinate-wise
optimization problem (see Corollary 1 in (Tian, 2022)):

$\displaystyle\text{(Min-player $\alpha$)}\quad\quad\quad\quad\quad\alpha_{t}$
$\displaystyle=\arg\min_{\alpha\in\mathcal{A}}\mathcal{E}_{\alpha}(\bm{\theta}_{t})-\mathcal{R}(\alpha)$
(2a) $\displaystyle\text{(Max-player
$\bm{\theta}$)}\quad\quad\quad\quad\quad\bm{\theta}_{t+1}$
$\displaystyle=\bm{\theta}_{t}+\eta\nabla_{\bm{\theta}}\mathcal{E}_{\alpha_{t}}(\bm{\theta})$
(2b)

where $\mathcal{A}$ is the feasible set of _pairwise importance_
$\alpha=[\alpha_{ij}]$ and the _energy_ function
$\mathcal{E}_{\alpha}(\bm{\theta}):=\mathrm{tr}\mathbb{C}_{\alpha}[{\bm{f}},{\bm{f}}]$
where $\mathbb{C}_{\alpha}[\cdot,\cdot]$ is the _contrastive covariance_
(Tian, 2022; Jing et al., 2022):

$\mathbb{C}_{\alpha}[{\bm{a}},{\bm{b}}]:=\sum_{i=1}^{N}\sum_{j\neq
i}\alpha_{ij}({\bm{a}}[i]-{\bm{a}}[j])({\bm{b}}[i]-{\bm{b}}[j])^{\top}-\sum_{i=1}^{N}\left(\sum_{j\neq
i}\alpha_{ij}\right)({\bm{a}}[i]-{\bm{a}}[i^{\prime}])({\bm{b}}[i]-{\bm{b}}[i^{\prime}])^{\top}$
(3)

For brevity, we define
$\mathbb{C}_{\alpha}[{\bm{x}}]:=\mathbb{C}_{\alpha}[{\bm{x}},{\bm{x}}]$. In
Eqn. 2, the max player $\bm{\theta}$ (Eqn. 2b) finds good representation to
maximize the distance of different samples, while minimizing the distance of
different views of the same sample, by taking one gradient step. On the other
hand, the min player $\alpha$ (Eqn. 2a) puts more emphasis on different sample
pairs which are similar in the representation space, regularized by the
regularizer $\mathcal{R}(\alpha)$.

In this paper, we mainly focus on how the nonlinearity affects representation
learning and treats $\alpha$ to be a fixed quantity. In short, we want to
understand the behavior of the gradient ascent rule (Eqn. 2b) of the
corresponding objective
$\max_{\bm{\theta}}\mathrm{tr}\mathbb{C}_{\alpha}[{\bm{f}}({\bm{x}};\bm{\theta})]$.
For this, first we give some intuitions of $\mathbb{C}_{\alpha}$:

######  Lemma 1 (Intuition of Contrastive Covariance).

When $\alpha_{ij}=1/2N(N-1)$ and $N\rightarrow+\infty$, then for any function
${\bm{g}}(\cdot)$,
$\mathbb{C}_{\alpha}[{\bm{g}}({\bm{x}})]\rightarrow\mathbb{V}_{{\bm{x}}_{0}\sim
p(\cdot)}\left[\mathbb{E}_{{\bm{x}}\sim
p_{\mathrm{aug}}(\cdot|{\bm{x}}_{0})}[{\bm{g}}({\bm{x}})|{\bm{x}}_{0}]\right]$
and is a PSD matrix.

######  Corollary 1 (Uniform $\alpha$, no augmentation and large batchsize).

With the condition of Lemma 1, if we further assume there is no augmentation
(i.e.,
$p_{\mathrm{aug}}({\bm{x}}|{\bm{x}}_{0})=\delta({\bm{x}}-{\bm{x}}_{0})$), then
$\mathbb{C}_{\alpha}[{\bm{g}}({\bm{x}})]\rightarrow\mathbb{V}[{\bm{g}}({\bm{x}})]$.

Please check Appendix (Supplementary Material) for detailed proofs. Note that
Lemma 1 rules out one possibility raised in (Jing et al., 2022) that
$\mathbb{C}_{\alpha}[{\bm{x}}]$ can contain negative eigenvalues. When
$\alpha$ is not uniform, $\mathbb{C}_{\alpha}[{\bm{x}}]$ can be regarded as a
more general version of variance and may have negative eigenvalues.

##  3 One-layer case

Now let us first consider 1-layer network with $K$ hidden nodes:
${\bm{f}}({\bm{x}};\bm{\theta})=h(W{\bm{x}})$, where
$W=[{\bm{w}}_{1},\ldots,{\bm{w}}_{K}]^{\top}\in\mathbb{R}^{K\times d}$ and
$h(x)$ is the activation function. The $k$-th row of $W$ is a weight
${\bm{w}}_{k}$ and its output is $f_{k}:=h({\bm{w}}_{k}^{\top}{\bm{x}})$. In
this case,
$\mathrm{tr}\mathbb{C}_{\alpha}[{\bm{f}}]=\sum_{k=1}^{K}\mathbb{C}_{\alpha}[f_{k}]$,
i.e., the summation of the contrastive covariance of $f_{k}$.

For convenience, we consider per-filter normalization
$\|{\bm{w}}_{k}\|_{2}=1$, which can be achieved by imposing BatchNorm (Ioffe &
Szegedy, 2015) at each node $k$ (Tian, 2022). In this case, optimization with
each filter ${\bm{w}}_{k}$ can be considered separately, yielding the
following smaller sub-problem:

$\max_{\|{\bm{w}}_{k}\|_{2}=1,1\leq k\leq
K}\mathrm{tr}\mathbb{C}_{\alpha}[{\bm{f}}]=\sum_{k=1}^{K}\max_{\|{\bm{w}}_{k}\|_{2}=1}\mathbb{C}_{\alpha}[h({\bm{w}}_{k}^{\top}{\bm{x}})]$
(4)

Now let’s think about, which parameters ${\bm{w}}_{k}$ maximizes the
summation? For the linear case, since
$\mathbb{C}_{\alpha}[h({\bm{w}}^{\top}{\bm{x}})]=\mathbb{C}_{\alpha}[{\bm{w}}^{\top}{\bm{x}}]={\bm{w}}^{\top}\mathbb{C}_{\alpha}[{\bm{x}}]{\bm{w}}$,
all ${\bm{w}}_{k}$ converge to the maximal eigenvector of
$\mathbb{C}_{\alpha}[{\bm{x}}]$ (a constant matrix), regardless of how they
are initialized and what the distribution of ${\bm{x}}$ is.

Therefore, the linear case will only learn the most salient single pattern.
All other features will be neglected due to the (overly-smooth) landscape of
the objective function. This is the winner-take-all effect and will miss many
patterns in the data.

Figure 1: Left: Summary of Sec. 3. (a) We analyze the dynamics of one-layer
network $h({\bm{w}}^{\top}{\bm{x}})$ under contrastive learning loss (Eqn. 6).
(b) With linear activation $h(x)=x$, then there is only one fixed point (PCA
direction). (c) Non-linear $h(x)$ creates many critical points. If $h$ is
ReLU, all locally maximal critical points are stable (Theorem 1). Right:
Convergence patterns (iteration $t$ versus iteration discrepancy
$\|{\bm{w}}(t+1)-{\bm{w}}(t)\|_{2}$) of Power Iteration (Eqn. PI) in latent
summation models, when $\|U^{\top}U-I\|_{2}$ is small but non-zero. In this
case, Theorem 2 tells there still exist local optima close to each
${\bm{u}}_{m}$.

In contrast, nonlinearity can change the landscape and create more local
maxima in $\mathbb{C}_{\alpha}[h({\bm{w}}^{\top}{\bm{x}})]$. In this paper, we
consider a general category of nonlinearity activations:

######  Assumption 1 (Homogeneity (Du et al., 2018)/Reversibility (Tian et
al., 2020)).

The activation satisfies $h(x)=h^{\prime}(x)x$.

Note that many activations satisfy this condition, including linear, ReLU,
LeakyReLU and monomial activations like $h(x)=x^{p}$ (with an additional
global constant). Under this assumption, we have:

$h({\bm{w}}^{\top}{\bm{x}})={\bm{w}}^{\top}h^{\prime}({\bm{w}}^{\top}{\bm{x}}){\bm{x}}={\bm{w}}^{\top}\tilde{\bm{x}}^{{\bm{w}}},$
(5)

where $\tilde{\bm{x}}^{{\bm{w}}}:={\bm{x}}\cdot
h^{\prime}({\bm{w}}^{\top}{\bm{x}})$ is the input data after nonlinear gating.
When there is no ambiguity, we just write $\tilde{\bm{x}}^{{\bm{w}}}$ as
$\tilde{\bm{x}}$ and omit the weight superscript.

Now let $A({\bm{w}}):=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]$, we have
$\mathbb{C}_{\alpha}[h({\bm{w}}^{\top}{\bm{x}})]={\bm{w}}^{\top}A({\bm{w}}){\bm{w}}$.
Considering the constraint $\|{\bm{w}}\|_{2}=1$, the learning dynamics can be
written as the following (here $P^{\perp}_{\bm{w}}:=I-{\bm{w}}{\bm{w}}^{\top}$
projects a vector into the complementary subspace spanned by ${\bm{w}}$):

$\dot{\bm{w}}=P_{{\bm{w}}}^{\perp}A({\bm{w}}){\bm{w}}$ (6)

Now the question is that: what is the critical point of the dynamics? In
linear case, the maximal eigenvector is the one fixed point; in nonlinear
case, we are looking for _locally_ maximal eigenvectors:

######  Definition 1 (Locally maximal eigenvector).

${\bm{w}}_{*}$ is called a locally maximal eigenvector of $A({\bm{w}})$, if
$A({\bm{w}}_{*}){\bm{w}}_{*}=\lambda_{*}{\bm{w}}_{*}$, where
$\lambda_{*}=\lambda_{\max}(A({\bm{w}}_{*}))$ is the distinct maximal
eigenvalue of $A({\bm{w}}_{*})$.

###  3.1 Two concrete generative models that lead to multiple local optima in
ReLU setting

To see why the nonlinear activation leads to many fixed points in Eqn. 6, we
first give two examplar generative models of the input ${\bm{x}}$ that show
Eqn. 6 has multiple critical points, then introduce a more general theorem. To
make the examples simple and clear, we assume the condition of Corollary 1
(i.e., uniform $\alpha$, no augmentation and large batchsize), and notice that
$\tilde{\bm{x}}^{{\bm{w}}}$ is a deterministic function of ${\bm{x}}$,
therefore
$A({\bm{w}}):=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{V}[\tilde{\bm{x}}^{{\bm{w}}}]$.
We also use ReLU activation $h(x)=\max(x,0)$.

Let $U=[{\bm{u}}_{1},\ldots,{\bm{u}}_{M}]$ be orthonormal bases with
${\bm{u}}_{m}^{\top}{\bm{u}}_{m^{\prime}}=\mathbb{I}(m=m^{\prime})$. Here are
two examples:

Latent categorical model. Suppose $y$ is a categorical random variable taking
$M$ possible values, $\mathbb{P}[{\bm{x}}|y=m]=\delta({\bm{x}}-{\bm{u}}_{m})$.
Then we have (see Appendix for detailed steps):

$A({\bm{w}})\big{|}_{{\bm{w}}={\bm{u}}_{m}}:=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{V}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{P}[y=m](1-\mathbb{P}[y=m]){\bm{u}}_{m}{\bm{u}}_{m}^{\top}$
(7)

Now it is clear that ${\bm{w}}={\bm{u}}_{m}$ is locally maximal eigenvector
(and thus a critical point) for any $m$.

Latent summation model. Suppose there is a latent variable ${\bm{y}}$ so that
${\bm{x}}=U{\bm{y}}$, where ${\bm{y}}:=[y_{1},y_{2},\ldots,y_{M}]$. Each
$y_{m}$ is a standardized Bernoulli random variable: $\mathbb{E}[y_{m}]=0$ and
$\mathbb{E}[y_{m}^{2}]=1$. This means that
$y_{m}=y_{m}^{+}:=\sqrt{(1-q_{m})/q_{m}}$ with probability $q_{m}$ and
$y_{m}=y_{m}^{-}:=-\sqrt{q_{m}/(1-q_{m})}$ with probability $1-q_{m}$. For
$m_{1}\neq m_{2}$, $y_{m_{1}}$ and $y_{m_{2}}$ are independent. Then we have:

$A({\bm{w}})\big{|}_{{\bm{w}}={\bm{u}}_{m}}:=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{V}\left[\tilde{\bm{x}}\right]=(1-q_{m})^{2}{\bm{u}}_{m}{\bm{u}}_{m}^{\top}+q_{m}(I-{\bm{u}}_{m}{\bm{u}}_{m}^{\top})$
(8)

which has a maximal and distinct eigenvector of ${\bm{u}}_{m}$ with a unique
eigenvalue $(1-q_{m})^{2}$, when $q_{m}<\frac{1}{2}(3-\sqrt{5})\approx 0.382$.
Therefore, different ${\bm{w}}$ leads to different locally maximal
eigenvectors.

In both cases, the presence of ReLU removes the “redundant energy” so that
$A({\bm{w}})$ can focus on specific directions, creating multiple fixed points
that correspond to multiple learnable patterns. Furthermore, we can also prove
in certain circumstances, these fixed points can be stable. Here we introduce
the following notations: $\lambda_{i}({\bm{w}})$ and $\phi_{i}({\bm{w}})$ are
the $i$-th eigenvalue and unit eigenvector of $A({\bm{w}})$
($\phi_{1}({\bm{w}})$ is the largest),
$\lambda_{\mathrm{gap}}({\bm{w}}):=\lambda_{1}({\bm{w}})-\lambda_{2}({\bm{w}})$
is the eigenvalue gap, and $L$ is the Lipschitz constant of $A({\bm{w}})$:
$\|A({\bm{w}}^{\prime})-A({\bm{w}})\|_{2}\leq
L\|{\bm{w}}^{\prime}-{\bm{w}}\|_{2}$ over the unit sphere
$\|{\bm{w}}\|_{2}=1$.

######  Theorem 1 (Local stability of ${\bm{w}}^{*}$ for ReLU activation).

If ${\bm{w}}_{*}$ is an eigenvector of $A({\bm{w}}_{*})$, then (1)
${\bm{w}}_{*}$ is a critical point of Eqn. 6, (2) if
$\lambda_{\mathrm{gap}}({\bm{w}}_{*})>L$, ${\bm{w}}_{*}$ is stable.

###  3.2 Finding critical points with initial guess

The two examples can be computed analytically due to our specific choice on
nonlinearity $h$ and $\alpha$. In general, we might not find analytical forms
of any locally maximal eigenvector of $A({\bm{w}})$.

So do they really exist? The following theorem shows that if we could find an
“approximate eigenvector” of
$A({\bm{w}}):=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]$, then we could
show there exists a real one nearby. Specifically, we can construct a fixed
point using _Power Iteration_ (PI) (Golub & Van Loan, 2013), starting from
initial value ${\bm{w}}={\bm{w}}(0)$:

${\bm{w}}(t+1)\leftarrow
A({\bm{w}}(t)){\bm{w}}(t)/\|A({\bm{w}}(t)){\bm{w}}(t)\|_{2}$ (PI)

It turns out that, if the inner product
$c({\bm{w}}):={\bm{w}}^{\top}\phi_{1}({\bm{w}})$ is positive, then even
$A({\bm{w}})$ varies over $\|{\bm{w}}\|_{2}=1$, the iteration can still
converge to a fixed point ${\bm{w}}_{*}$:

######  Theorem 2 (Existence of critical points).

Let
$\mu({\bm{w}}):=.5(1+c({\bm{w}}))c^{-2}({\bm{w}})\left[1-\lambda_{\mathrm{gap}}({\bm{w}})/\lambda_{1}({\bm{w}})\right]^{2}$.
If $A({\bm{w}})$ is positive definite (PD), $c_{0}:=c({\bm{w}}(0))\neq 0$ and
there exists $\gamma<1$ so that:

$\sup_{{\bm{w}}\in B_{\gamma}}\mu({\bm{w}})+2\kappa
L^{2}(1+\mu({\bm{w}})c({\bm{w}}))+2L\lambda^{-1}_{\mathrm{gap}}({\bm{w}})\sqrt{\mu({\bm{w}})(1+\mu({\bm{w}})c({\bm{w}}))}\leq\gamma,$
(9)

here
$B_{\gamma}:=\left\\{{\bm{w}}:{\bm{w}}^{\top}{\bm{w}}(0)\geq\frac{c_{0}-c_{\gamma}}{1-c_{\gamma}}\right\\}$,
$c_{\gamma}:=\frac{2\sqrt{\gamma}}{1+\gamma}$ and $\kappa$ is the high-order
eigenvector bound defined in Appendix, then Power Iteration (Eqn. PI)
converges to a critical point ${\bm{w}}_{*}\in B_{\gamma}$ of Eqn. 6. If
$\lambda_{\mathrm{gap}}({\bm{w}}_{*})>L$, then ${\bm{w}}_{*}$ is stable.

######  Corollary 2.

When $A({\bm{w}})$ is not PD, Theorem 2 can still be applied to the PD matrix
$\hat{A}({\bm{w}}):=A({\bm{w}})-\lambda_{\min}({\bm{w}})I+\epsilon I$ with $L$
and $\kappa$ specified by $\hat{A}({\bm{w}})$, where $\epsilon>0$ is a small
constant.

Intuitively, with $L$ and $\kappa$ small, $c_{0}$ close to 1, and
$\lambda_{\mathrm{gap}}$ large. Eqn. 9 can always hold with $\gamma<1$ and the
fixed point exists. For example, for the two cases in Sec. 3.1, if
$U=[{\bm{u}}_{1},\ldots,{\bm{u}}_{M}]$ is only approximately orthogonal (i.e.,
$\|U^{\top}U-I\|$ is not zero but small), and/or the conditions of Corollary 1
hold roughly, then Theorem 2 tells that multiple local optima close to
${\bm{u}}_{m}$ still exist for each $m$ (Fig. 1). Note that the condition can
be further relaxed to handle more cases (e.g., $A({\bm{w}})$ changes a lot but
the largest eigenvector $\phi_{1}(A({\bm{w}}))$ doesn’t change much), and we
leave it for future work.

Relation to empirical observations. Now we know that there exist many local
optima in the dynamics (Eqn. 6). As a result, even if objective involving
${\bm{w}}_{k}$ are identical (Eqn. 4), each ${\bm{w}}_{k}$ may still converge
to different local optima due to their initialization. This naturally explains
the reason why larger model yields better performance: it collects more local
optima and some can be _useful_.

In addition, this also explains other empirical observations like lottery
ticket hypothesis (LTH) (Frankle & Carbin, 2019; Morcos et al., 2019; Tian et
al., 2019; Yu et al., 2020), recently also verified in contrastive learning
(Chen et al., 2021). In LTH, first a large network is trained and pruned to be
a small subnetwork, then resetting its parameters back to the initialization
and retraining yields comparable or even better performance, compared to the
large trained network, while retraining this small subnetwork with a different
initialization performs much worse. For LTH, our explanation is that the
subnetwork that is picked after pruning contains weights initialized _luckily_
, i.e., they are initialized to be close to useful local optima and converge
to them during training.

Given this intuition, a follow-up question naturally follows: what is the
probability of finding a specific local minimum, if the weights are
initialized randomly? For this, we need to study the volume of _attractive
basin_ for each local optimum of Eqn. 6. While Theorem 2 gives some hints, the
problem is very challenging in general and we leave for future work.

##  4 Two-layer setting

Now we understand how 1-layer nonlinearity learns in contrastive learning
setting. In practice, many patterns exist and most of them may not be relevant
for the downstream tasks.

A natural question arises: how does the network prioritizes which patterns to
learn? To answer this question, we analyze the behavior of 2-layer nonlinear
networks with non-overlapping receptive fields (Fig. 2(a)). The network has
both top-layer and lower-layer weights. In the lower layer, there are $K$
disjoint _receptive fields_ (abbreviated as _RF_) $\\{R_{k}\\}$, each has
input ${\bm{x}}_{k}$ and $M$ weight ${\bm{w}}_{km}\in\mathbb{R}^{d}$ where
$m=1\ldots M$. The output of the bottom-layer is denoted as ${\bm{f}}_{1}$,
$f_{1km}$ for its $km$-th component, and ${\bm{f}}_{1}[i]$ for $i$-th sample.
The top layer has weight $V\in\mathbb{R}^{d_{\mathrm{out}}\times KM}$. Define
$S:=V^{\top}V$. As the $(km,k^{\prime}m^{\prime})$ entry of the matrix $S$,
$s_{km,k^{\prime}m^{\prime}}:=[S]_{km,k^{\prime}m^{\prime}}={\bm{v}}_{km}^{\top}{\bm{v}}_{k^{\prime}m^{\prime}}$.

At each RF $R_{k}$, define $\tilde{\bm{x}}_{km}$ as an brief notation of gated
input $\tilde{\bm{x}}_{k}^{{\bm{w}}_{km}}:={\bm{x}}_{k}\cdot
h^{\prime}({\bm{w}}^{\top}_{km}{\bm{x}}_{k})$. Define
$\tilde{\bm{x}}_{k}:=[\tilde{\bm{x}}_{k1};\tilde{\bm{x}}_{k2};\ldots,\tilde{\bm{x}}_{kM}]\in\mathbb{R}^{Md}$
as the concatenation of $\tilde{\bm{x}}_{km}$ and finally
$\tilde{\bm{x}}:=[\tilde{\bm{x}}_{1};\ldots;\tilde{\bm{x}}_{K}]\in\mathbb{R}^{KMd}$
is the concatenation of all $\tilde{\bm{x}}_{k}$. Similarly, let
${\bm{w}}_{k}:=[{\bm{w}}_{km}]_{m=1}^{M}\in\mathbb{R}^{Md}$ be a concatenation
of all ${\bm{w}}_{km}$ in the same RF $R_{k}$, and
${\bm{w}}:=[{\bm{w}}_{k}]_{k=1}^{K}\in\mathbb{R}^{KMd}$ be a column
concatenation of all ${\bm{w}}_{k}$. Finally,
$P^{\perp}_{\bm{w}}:=\mathrm{diag}_{km}[P^{\perp}_{{\bm{w}}_{km}}]$ is a
block-diagonal matrix putting all projections together.

Figure 2: Our setting for 2-layer network. (a) We use $W$ for low-layer
weights and $V$ for top-layer weights. There are $K$ disjoint receptive fields
(abbreviated as _RF_) $R_{k}$, each with $M$ weight vectors in the low-layer,
denoted as ${\bm{w}}_{km}$. The activation function of hidden layer nodes is
$h(x)$ and can be linear or nonlinear. (b) Conditional independence in
Assumption 2: there exists a global categorical variable $z$. Given $z$,
variation in different RFs are assumed to be independent.

######  Lemma 2 (Dynamics of 2-layer nonlinear network with contrastive loss).

$\dot{V}=V\mathbb{C}_{\alpha}[{\bm{f}}_{1}],\quad\quad\quad\quad\dot{\bm{w}}=P_{{\bm{w}}}^{\perp}\left[(S\otimes{\bm{1}}_{d}{\bm{1}}_{d}^{\top})\circ\mathbb{C}_{\alpha}[\tilde{\bm{x}}]\right]{\bm{w}}$
(10)

where ${\bm{1}}_{d}$ is $d$-dimensional all-one vector, $\otimes$ is Kronecker
product and $\circ$ is Hadamard product.

Now we analyze the stationary points of the equations. If
$\mathbb{C}_{\alpha}[{\bm{f}}_{1}]$ has unique maximal eigenvector ${\bm{s}}$,
then following similar analysis as in (Tian, 2022), a necessary condition for
$(W,V)$ to be a stationary point is that $V={\bm{v}}{\bm{s}}^{\top}$, where
${\bm{v}}$ is any arbitrary unit vector. Therefore, we have
$S=V^{\top}V={\bm{s}}{\bm{s}}^{\top}$ as a rank-1 matrix and
$s_{km,k^{\prime}m^{\prime}}=s_{km}s_{k^{\prime}m^{\prime}}$. Note that
${\bm{s}}$, as a unique maximal eigenvector of
$\mathbb{C}_{\alpha}[{\bm{f}}_{1}]$, is a function of the low-level feature
computed by $W$.

On the other hand, the stationary point of $W$ can be much more complicated,
since it has the feedback term $S$ from the top level. A more detailed
analysis requires further assumptions, as we list below:

######  Assumption 2.

For analysis of two-layer networks, we assume:

  * •

Uniform $\alpha$, large batchsize and no augmentation. Then
$\mathbb{C}_{\alpha}[{\bm{g}}({\bm{x}})]=\mathbb{V}[{\bm{g}}({\bm{x}})]$ for
any function ${\bm{g}}(\cdot)$ following Corollary 1.

  * •

Fast top-level training. $V$ undergoes _fast_ training and has always
converged to its stationary point given $\mathbb{C}_{\alpha}[{\bm{f}}_{1}]$.
That is, $S={\bm{s}}{\bm{s}}^{\top}$ is a rank-1 matrix;

  * •

Conditional Independence. There exists $z$, a _global_ random variable taking
$C$ different categories, so that _given_ $z$, the input in each $R_{k}$ are
independent:

$\mathbb{P}[{\bm{x}}|z]=\prod_{k=1}^{K}\mathbb{P}[{\bm{x}}_{k}|z]$ (11)

Explanation of the assumptions. The _no augmentation_ condition is mainly
technical. Conclusion still holds if
$\mathbb{E}_{p_{\mathrm{aug}}}[{\bm{g}}({\bm{x}})|{\bm{x}}_{0}]\approx{\bm{g}}(\mathbb{E}_{p_{\mathrm{aug}}}[{\bm{x}}|{\bm{x}}_{0}])$
for ${\bm{g}}({\bm{x}}):=\tilde{\bm{x}}^{{\bm{w}}}$, i.e., augmentation swaps
with nonlinear gating. For conditional independence, intuitively $z$ can be
regarded as different type of global patterns that determines what input
${\bm{x}}$ can be perceived (Fig. 2(b)). Once $z$ is given, the remaining
variation resides within each RF $R_{k}$ and independent across different
$R_{k}$. Note that there exists many patterns in each RF $R_{k}$. Some are
parts of the global pattern $z$, and others may come from noise. We study how
each weight ${\bm{w}}_{km}$ captures distinct and useful patterns after
training.

With all the assumptions, we can compute the term
$\mathbb{C}_{\alpha}[\tilde{\bm{x}}]=\mathbb{V}[\tilde{\bm{x}}]$. Let
$p_{c}:=\mathbb{P}[z=c]$.

######  Lemma 3 (Close-form of variance under Assumption 2).

With Assumption 2, we have

$\displaystyle\mathbb{V}[\tilde{\bm{x}}]=\mathrm{diag}_{k}\left[L_{k}\right]+\sum_{c=0}^{C-1}p_{c}(1-p_{c})^{2}\Delta(c)\Delta^{\top}(c)$
(12)

where
$L_{k}:=\mathbb{E}_{z}\mathbb{V}[\tilde{\bm{x}}_{k}|z]\in\mathbb{R}^{Md}$ and
$\Delta(c):=\mathbb{E}[\tilde{\bm{x}}|z=c]-\mathbb{E}[\tilde{\bm{x}}|z\neq
c]\in\mathbb{R}^{MKd}$. In particular when $C=2$, the second term becomes
$p_{0}p_{1}\Delta\Delta^{\top}$, a rank-1 matrix. Here $\Delta:=\Delta(0)$ for
brevity.

Our Assumption 2 is weaker than orthogonal mixture condition in (Tian, 2022)
that is used to analyze CL, which requires the instance of input
${\bm{x}}_{k}[i]$ to have only one positive component.

###  4.1 The effect of global modulation in the special case of $C=2$ and
$M=1$

When $z$ is a binary variable ($C=2$) and without multiple weights per RF
($M=1$), it turns out that ${\bm{w}}_{k}$’s dynamics has close form. Let
${\bm{w}}_{k}$ represent ${\bm{w}}_{k1}$, the only weight at each $R_{k}$,
$\Delta_{k}:=\mathbb{E}[\tilde{\bm{x}}_{k}|z=1]-\mathbb{E}[\tilde{\bm{x}}_{k}|z=0]$
is the $k$-th section of $\Delta$. We now have:

######  Theorem 3 (Dynamics of ${\bm{w}}_{k}$ under conditional independence).

Let $d_{k}:={\bm{w}}_{k}^{\top}L_{k}{\bm{w}}_{k}\geq 0$ and $\lambda$ be the
maximal eigenvalue of $\mathbb{V}[{\bm{f}}_{1}]$, then (1)
$\lambda\geq\max_{k}d_{k}$, (2)
${\bm{s}}=Z^{-1}[{\bm{w}}_{k}^{\top}\Delta_{k}/(\lambda-d_{k})]$ is the
maximal unit eigenvector ($Z$ is the normalization constant), and (3) the
dynamics of ${\bm{w}}_{k}$ is given by:

$\dot{\bm{w}}_{k}=P^{\perp}_{{\bm{w}}_{k}}\left(s^{2}_{k}L_{k}+\frac{1}{Z^{2}(\lambda-
d_{k})}\Delta_{k}\Delta^{\top}_{k}\right){\bm{w}}_{k}$ (13)

There are several interesting observations. First, thanks to the nice property
of largest eigenvector of $\mathbb{V}[{\bm{f}}_{1}]$, the dynamics are
decoupled (i.e., $\dot{\bm{w}}_{k}=A_{k}(W){\bm{w}}_{k}$) and other
${\bm{w}}_{k^{\prime}}$ with $k^{\prime}\neq k$ only affects the dynamics of
${\bm{w}}_{k}$ through the matrix $A_{k}(W)$.

Second, while $L_{k}:=\mathbb{E}_{z}\mathbb{V}[\tilde{\bm{x}}_{k}|z]$ contains
multiple patterns (i.e., local optima) in $R_{k}$, the additional term
$\Delta_{k}\Delta_{k}^{\top}$, as the _global modulation_ from the top level,
encourages the model to learn $\Delta_{k}$ which is a discriminative feature
that separates the event of $z=0$ and $z=1$. Quantitatively:

######  Theorem 4 (Global modulation of attractive basin).

If the structural assumption holds:
$L_{k}({\bm{w}}_{k})=\sum_{l}g({\bm{u}}_{l}^{\top}{\bm{w}}_{k}){\bm{u}}_{l}{\bm{u}}_{l}^{\top}$
with $g(\cdot)>0$ a linear increasing function and $\\{{\bm{u}}_{l}\\}$
orthonormal bases, then for $L_{k}+c{\bm{u}}_{l}{\bm{u}}_{l}^{\top}$, its
attractive basin of ${\bm{w}}_{k}={\bm{u}}_{l}$ is larger than $L_{k}$’s for
$c>0$.

Therefore, if $\Delta_{k}$ is a locally maximal eigenvector of $L_{k}$ and
${\bm{w}}_{k}$ is randomly initialized, Thm 3 tells that
$\mathbb{P}[{\bm{w}}_{k}\rightarrow\Delta_{k}]$ is higher than the probability
that ${\bm{w}}_{k}$ goes to other patterns of $L_{k}$, i.e., the global
variable $z$ _modulates_ the training of the lower layer. This is similar to
“Backward feature correction” (Allen-Zhu & Li, 2020) and “top-down modulation”
(Tian et al., 2019) in supervised learning, here we demonstrate it in CL.

Finally, if RFs have diverse variances (i.e., small $d_{k}$ are large while
others are small), the weight updating pace can be much slower for RF with
small $d_{k}$, due to the factor $\frac{1}{\lambda-d_{k}}$ in both
$s_{k}^{2}L_{k}$ and $\frac{1}{Z^{2}(\lambda-
d_{k})}\Delta_{k}\Delta^{\top}_{k}$ and the fact that
$\lambda\geq\max_{k}d_{k}$. In this case, adding BatchNorm at each node
alleviates this issue, as shown in a modified version of Theorem 3. Please
check Appendix for details.

###  4.2 Why nonlinearity is critical: linear activation fails in general
cases ($C>2$ and $M>2$)

So far we have discussed the case of binary $z$ ($C=2$), in which the term
$\Delta_{k}$ essentially characterizes the single pattern at RF $R_{k}$. When
$C>2$, according to Lemma 3, there can be $C$ different discriminative
patterns $\Delta_{k}(c)$ at each RF $R_{k}$.

Since in each RF $R_{k}$, there are $M$ filters $\\{{\bm{w}}_{km}\\}$, it
would be ideal to have one filter to capture one distinct pattern
$\Delta_{k}(c)$. However, with linear activation,
$\tilde{\bm{x}}_{k}={\bm{x}}_{k}$ and both $L_{k}$ and $\Delta_{k}$ are
independent of ${\bm{w}}_{k}$. As a result, learning of diverse features never
happens no matter how large $M$ is:

######  Theorem 5 (Gradient Colinearity in linear networks).

With linear activation, $W$ follows the dynamics:

$\dot{\bm{w}}_{km}=s_{km}{\bm{b}}_{k}(W,V)$ (14)

where
${\bm{b}}_{k}(W,V):=\mathbb{C}_{\alpha}\left[{\bm{x}}_{k},\sum_{k^{\prime},m^{\prime}}s_{k^{\prime}m^{\prime}}{\bm{w}}^{\top}_{k^{\prime}m^{\prime}}{\bm{x}}_{k^{\prime}}\right]$
is a linear function w.r.t. $W$. As a result, (1) $\dot{\bm{w}}_{km}$ are co-
linear over $m$, and (2) If there exists one $m$ with $s_{km}\neq 0$, then
from any critical point with distinct $\\{{\bm{w}}_{km}\\}$, there exists a
path of critical points to identical weights (${\bm{w}}_{km}={\bm{w}}_{k}$).

In this situation, non-linearity plays a key role, by pushing weights
$\\{{\bm{w}}_{km}\\}$ at the same RF $R_{k}$ towards diverse features. To see
this, we consider a simple example of two weights ${\bm{w}}_{k1}$ and
${\bm{w}}_{k2}$ with the dynamics
$\left[\begin{array}[]{c}\dot{\bm{w}}_{k1}\\\
\dot{\bm{w}}_{k2}\end{array}\right]=\left[\begin{array}[]{cc}P^{\perp}_{{\bm{w}}_{k1}}&0\\\
0&P^{\perp}_{{\bm{w}}_{k2}}\end{array}\right]\left[\begin{array}[]{cc}A_{k1}({\bm{w}}_{k1})&A_{k12}({\bm{w}}_{k1},{\bm{w}}_{k2})\\\
A_{k21}({\bm{w}}_{k2},{\bm{w}}_{k1})&A_{k2}({\bm{w}}_{k2})\end{array}\right]\left[\begin{array}[]{c}{\bm{w}}_{k1}\\\
{\bm{w}}_{k2}\end{array}\right]$, where $A_{k1}$ and $A_{k2}$ are symmetric
matrices and
$A_{k12}({\bm{w}}_{k1},{\bm{w}}_{k2})=A_{k21}^{\top}({\bm{w}}_{2},{\bm{w}}_{1})$.
In the nonlinear setting, from Lemma 2 and Lemma 3, we have
$A_{k1}({\bm{w}}_{k1})=\mathbb{E}_{z}[\mathbb{V}[\tilde{\bm{x}}_{k}^{{\bm{w}}_{k1}}|z]]+\text{topdown-
terms}$,
$A_{k2}({\bm{w}}_{k2})=\mathbb{E}_{z}[\mathbb{V}[\tilde{\bm{x}}_{k}^{{\bm{w}}_{k2}}|z]]+\text{topdown-
terms}$, and
$A_{k12}({\bm{w}}_{k1},{\bm{w}}_{k2})=\mathbb{E}_{z}[\mathrm{Cov}[\tilde{\bm{x}}_{k}^{{\bm{w}}_{k1}},\tilde{\bm{x}}_{k}^{{\bm{w}}_{k2}}|z]]$.
If
$A_{k12}({\bm{w}}_{k1},{\bm{w}}_{k2})\propto{\bm{w}}_{k1}{\bm{w}}_{k2}^{\top}$,
then we have $\dot{\bm{w}}_{kl}\sim
P_{{\bm{w}}_{kl}}^{\perp}(A_{kl}({\bm{w}}_{kl})+I){\bm{w}}_{kl}$ for $l=1,2$.
From our analysis in Sec. 3, with nonlinear activation,
$\dot{\bm{w}}=P_{{\bm{w}}}^{\perp}A_{k}({\bm{w}}){\bm{w}}$ has multiple fixed
points, and ${\bm{w}}_{k1}$ and ${\bm{w}}_{k2}$ can converge to one of them
based on their initializations. This achieves specialization to diverse
features. See Appendix for detailed analysis.

Figure 3: Experimental setting (Sec. 5). When generating input, we first
randomly pick one generator (e.g., **C*B*A*D*) from a pool of $G$ generators,
generate the sequence by instantiating wildcard * with an arbitrary token, and
then replace the token $a$ of sequence with an embedding vector ${\bm{u}}_{a}$
to form the input ${\bm{x}}$. Inputs from the same generator are treated as
positive pairs, otherwise negative pairs for contrastive loss.

##  5 Experiments

Setup. To verify our finding, we perform contrastive learning with a 2-layer
network on a synthetic dataset containing token sequences, generated as
follows. From a pool of $G=40$ generators, we pick a generator of length $K$
in the form of **C*B*A*D* (here $K=10$) and generate EFCDBAACDB by sampling
from $d=20$ tokens for each wildcard *. The final input ${\bm{x}}$ is then
constructed by replacing each token $a$ with the pre-defined embedding
${\bm{u}}_{a}\in\mathbb{R}^{d}$. $\\{{\bm{u}}_{a}\\}$ forms a orthonormal
bases (see Fig. 3). The data augmentation is achieved by generating another
sequence from the same generator.

While there exists $d=20$ tokens, in each RF $R_{k}$ we pick a subset
$R_{k}^{\mathrm{g}}$ of $P<d$ tokens as the candidates used in the generator,
to demonstrate the effect of global modulation. Before training, each
generator is created by first randomly picking $5$ receptive fields, then
picking one of the $P$ tokens from $R_{k}^{\mathrm{g}}$ at each RF $R_{k}$ and
filling the remaining RFs with wildcard *. Therefore, if a token appears at
$R_{k}$ but $a\notin R^{\mathrm{g}}_{k}$, then $a$ must be instantiated from
the wildcard. Any $a\notin R^{\mathrm{g}}_{k}$ is noise and should not to be
learned in the weights of $R_{k}$ since it is not part of any global pattern
from the generator.

We train a 2-layer network on this dataset. The 2-layer network has $K=10$
disjoint RFs, within each RF, there are $M=\beta P$ filters. Here $\beta\geq
1$ is a hyper-parameter that controls the degree of _over-parameterization_.
The network is trained with InfoNCE loss and SGD with learning rate $2\times
10^{-3}$, momentum $0.9$, and weight decay $5\times 10^{-3}$ for $5000$
minibatches and batchsize $128$. Code is written in PyTorch and a single
modern GPU suffices to run the experiments.

Evaluation metric. We check whether the weights corresponding to each token is
learned in the lower layer. At each RF $R_{k}$, we know $R_{k}^{\mathrm{g}}$,
the subsets of tokens it contains, as well as their embeddings
$\\{{\bm{u}}_{a}\\}_{a\in R^{\mathrm{g}}_{k}}$ due to the generation process,
and verify whether these embeddings are learned after the model is trained.
Specifically, for each token $a\in R_{k}^{\mathrm{g}}$, we look for its best
match on the learned filter $\\{{\bm{w}}_{km}\\}$, as formulated by the
following per-RF score $\chi_{+}(R_{k})$ and overall matching score
$\bar{\chi}_{+}\in[-1,1]$ as the average over all RFs (similarly we can also
define $\bar{\chi}_{-}$ for $a\notin R^{\mathrm{g}}_{k}$):

$\chi_{+}(R_{k})=\frac{1}{P}\sum_{a\in
R^{\mathrm{g}}_{k}}\max_{m}\frac{{\bm{w}}^{\top}_{km}{\bm{u}}_{a}}{\|{\bm{w}}_{km}\|_{2}\|{\bm{u}}_{a}\|_{2}},\quad\quad\quad\quad\bar{\chi}_{+}=\frac{1}{K}\sum_{k}\chi_{+}(R_{k})$
(15)

$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
---|---|---|---|---
Linear | ReLU | Linear | ReLU | Linear | ReLU | Linear | ReLU
1 | $0.95\\!\pm\\!0.00$ | $0.31\\!\pm\\!0.23$ | $0.79\\!\pm\\!0.02$ | $0.75\\!\pm\\!0.11$ | $0.67\\!\pm\\!0.03$ | $0.70\\!\pm\\!0.06$ | $0.60\\!\pm\\!0.01$ | $0.68\\!\pm\\!0.05$
2 | $0.95\\!\pm\\!0.00$ | $0.61\\!\pm\\!0.13$ | $0.81\\!\pm\\!0.01$ | $0.90\\!\pm\\!0.09$ | $0.70\\!\pm\\!0.00$ | $0.88\\!\pm\\!0.04$ | $0.63\\!\pm\\!0.01$ | $0.93\\!\pm\\!0.02$
5 | $0.96\\!\pm\\!0.00$ | $0.93\\!\pm\\!0.06$ | $0.85\\!\pm\\!0.01$ | $0.99\\!\pm\\!0.02$ | $0.73\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.66\\!\pm\\!0.01$ | $0.99\\!\pm\\!0.01$
10 | $0.96\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.86\\!\pm\\!0.01$ | $1.00\\!\pm\\!0.00$ | $0.77\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.68\\!\pm\\!0.01$ | $1.00\\!\pm\\!0.00$
Table 1: Overall matching score $\bar{\chi}_{+}$ (Eqn. 15). (a) When $P=1$, linear model works well regardless of the degree of over-parameterization $\beta$, while ReLU model requires large over-parameterization to perform well; (b) When each $R_{k}$ has multiple patterns ($P>1$) related to generators, ReLU models can capture diverse patterns better than linear ones in the over-parameterization region $\beta>1$. In contrast, linear models are much less affected by over-parameterization. Each setting is repeated 3 times and mean/standard derivations are reported. $\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
---|---|---|---|---
NoBN | BN | NoBN | BN | NoBN | BN | NoBN | BN
1 | $0.35\\!\pm\\!0.17$ | $0.23\\!\pm\\!0.06$ | $0.51\\!\pm\\!0.02$ | $0.24\\!\pm\\!0.11$ | $0.54\\!\pm\\!0.04$ | $0.38\\!\pm\\!0.04$ | $0.54\\!\pm\\!0.03$ | $0.56\\!\pm\\!0.04$
2 | $0.39\\!\pm\\!0.18$ | $0.33\\!\pm\\!0.12$ | $0.59\\!\pm\\!0.05$ | $0.41\\!\pm\\!0.07$ | $0.58\\!\pm\\!0.07$ | $0.56\\!\pm\\!0.02$ | $0.56\\!\pm\\!0.01$ | $0.71\\!\pm\\!0.05$
5 | $0.53\\!\pm\\!0.06$ | $0.46\\!\pm\\!0.06$ | $0.66\\!\pm\\!0.06$ | $0.60\\!\pm\\!0.08$ | $0.65\\!\pm\\!0.09$ | $0.74\\!\pm\\!0.02$ | $0.60\\!\pm\\!0.01$ | $0.90\\!\pm\\!0.03$
10 | $0.56\\!\pm\\!0.16$ | $0.46\\!\pm\\!0.15$ | $0.78\\!\pm\\!0.02$ | $0.80\\!\pm\\!0.03$ | $0.70\\!\pm\\!0.04$ | $0.92\\!\pm\\!0.01$ | $0.63\\!\pm\\!0.01$ | $0.97\\!\pm\\!0.02$
Table 2: The effect of BatchNorm (BN) with ReLU activation in the presence of
non-uniformity $\zeta$ of the input data. The non-uniformity is set to be
$\zeta=10$. BatchNorm yields better matching score $\bar{\chi}_{+}$ in the
over-parameterization region (large $\beta$) and multiple candidate tokens per
RF (large $P$), i.e., the shaded region. More results in Appendix. Figure 4:
Visualization of learned weights with $P=3$ ($3$ local patterns related to
generators at each RF) and $\beta=5$ (5x over-parameterization). Each of the
$K=10$ subfigures corresponds to a RF ($R_{0}$-$R_{9}$). In each subfigure,
the left panel is the learned weight by ReLU, while the right panel is from
linear activations. 15 rows corresponds to $M=\beta P=15$ weights and each
weight is $d=8$ dimensional. With ReLU activation, learned weights clearly
capture the 3 candidate tokens within $R^{\mathrm{g}}_{k}$ at each RF $R_{k}$,
while linear activation cannot.

###  5.1 Results

Linear v.s ReLU activation and the effect of over-parameterization (Sec. 4.2).
From Tbl. 1, we can clearly see that ReLU activation achieves better
reconstruction of the input patterns, when each RF contains many patterns
($P>1$) and specialization of filters in each RF is needed. On the other hand,
when $P=1$, linear activation works better. ReLU activation clearly benefits
from over-parameterization ($\beta>1$): the larger $\beta$ is, the better
$\bar{\chi}_{+}$ becomes. In contrast, for linear activation, over-
parameterization does not quite affect the performance.

Global modulation (Sec. 4.1). As shown in Fig. 4, the learned weights indeed
focus on the token subset $R_{k}^{\mathrm{g}}{}$ that receives top-down
support from the generators and no noise token is learned. We also verify that
quantitatively by computing $\bar{\chi}_{-}$ over multiple runs, which is
provided in Appendix.

The effect of BatchNorm (BN) (Sec. 4.1). When the scale of input data varies a
lot, BatchNorm starts to matter in discovering features with low magnitude
(Sec. 4.1). To model the scale non-uniformity, we set
$\|{\bm{u}}_{a}\|_{2}=\zeta$ for $\lfloor d/2\rfloor$ tokens and
$\|{\bm{u}}_{a}\|_{2}=1/\zeta$ for the remaining tokens. Larger $\zeta$
corresponds to higher non-uniformity across inputs. Tbl. 2 shows BN yields
higher $\bar{\chi}_{+}{}$ in the presence of high non-uniformity (e.g.,
$\zeta=10$) when the network is over-parameterized ($\beta>1$) and there are
multiple candidates per $R_{k}$ ($P>1$), a setting that is likely to hold in
real-world scenarios.

Quadratic loss versus infoNCE loss. In Appendix, we show that quadratic
contrastive loss underperforms InfoNCE, while the trend of linear/ReLU and
over-parameterization remains similar. One hypothesis is that due the presence
of min-player $\alpha$, contrastive covariance $\mathbb{C}_{\alpha}$ focuses
on important sample pairs, leading to low-rank structure. A formal study about
the pairwise importance $\alpha$ is left as one of the future works.

## References

  * Allen-Zhu & Li (2020) Allen-Zhu, Z. and Li, Y.  Backward feature correction: How deep learning performs deep learning.  _arXiv preprint arXiv:2001.04413_ , 2020.
  * Arora et al. (2019a) Arora, S., Cohen, N., Golowich, N., and Hu, W.  A convergence analysis of gradient descent for deep linear neural networks.  In _International Conference on Learning Representations_ , 2019a.  URL https://openreview.net/forum?id=SkMQg3C5K7.
  * Arora et al. (2019b) Arora, S., Khandeparkar, H., Khodak, M., Plevrakis, O., and Saunshi, N.  A theoretical analysis of contrastive unsupervised representation learning.  _arXiv preprint arXiv:1902.09229_ , 2019b.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Bunch et al. (1978) Bunch, J. R., Nielsen, C. P., and Sorensen, D. C.  Rank-one modification of the symmetric eigenproblem.  _Numerische Mathematik_ , 31(1):31–48, 1978.
  * Caron et al. (2020) Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A.  Unsupervised learning of visual features by contrasting cluster assignments.  In _NeurIPS_ , 2020.
  * Chen et al. (2020) Chen, T., Kornblith, S., Norouzi, M., and Hinton, G.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pp. 1597–1607. PMLR, 2020.
  * Chen et al. (2021) Chen, T., Frankle, J., Chang, S., Liu, S., Zhang, Y., Carbin, M., and Wang, Z.  The lottery tickets hypothesis for supervised and self-supervised pre-training in computer vision models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 16306–16316, 2021.
  * Chen & He (2020) Chen, X. and He, K.  Exploring simple siamese representation learning.  In _CVPR_ , 2020.
  * Chowdhery et al. (2022) Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A. M., Pillai, T. S., Pellat, M., Lewkowycz, A., Moreira, E., Child, R., Polozov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier-Hellstern, K., Eck, D., Dean, J., Petrov, S., and Fiedel, N.  Palm: Scaling language modeling with pathways.  _CoRR_ , abs/2204.02311, 2022.  doi: 10.48550/arXiv.2204.02311.  URL https://doi.org/10.48550/arXiv.2204.02311.
  * Cortes & Vapnik (1995) Cortes, C. and Vapnik, V.  Support-vector networks.  _Machine learning_ , 20(3):273–297, 1995.
  * Devlin et al. (2018) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Devroye et al. (1994) Devroye, L., Gyorfi, L., Krzyzak, A., and Lugosi, G.  On the strong universal consistency of nearest neighbor regression function estimates.  _The Annals of Statistics_ , 22(3):1371–1385, 1994.
  * Du et al. (2018) Du, S. S., Hu, W., and Lee, J. D.  Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced.  _arXiv preprint arXiv:1806.00900_ , 2018.
  * Fernández (2000) Fernández, F. M.  _Introduction to perturbation theory in quantum mechanics_.  CRC press, 2000.
  * Fix & Hodges (1951) Fix, E. and Hodges, J. L.  Nonparametric discrimination: consistency properties.  _Randolph Field, Texas, Project_ , pp. 21–49, 1951.
  * Frankle & Carbin (2019) Frankle, J. and Carbin, M.  The lottery ticket hypothesis: Finding sparse, trainable neural networks.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rJl-b3RcF7.
  * Golub & Van Loan (2013) Golub, G. H. and Van Loan, C. F.  _Matrix computations_.  JHU press, 2013.
  * Goriely & Hyde (1998) Goriely, A. and Hyde, C.  Finite-time blow-up in dynamical systems.  _Physics Letters A_ , 250(4-6):311–318, 1998\.
  * Grill et al. (2020) Grill, J.-B., Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., Doersch, C., Pires, B. A., Guo, Z. D., Azar, M. G., et al.  Bootstrap your own latent: A new approach to self-supervised learning.  _NeurIPS_ , 2020.
  * Gühring et al. (2020) Gühring, I., Raslan, M., and Kutyniok, G.  Expressivity of deep neural networks.  _arXiv preprint arXiv:2007.04759_ , 2020.
  * Hammer & Gersmann (2003) Hammer, B. and Gersmann, K.  A note on the universal approximation capability of support vector machines.  _neural processing letters_ , 17(1):43–53, 2003\.
  * HaoChen et al. (2021) HaoChen, J. Z., Wei, C., Gaidon, A., and Ma, T.  Provable guarantees for self-supervised deep learning with spectral contrastive loss.  _NeurIPS_ , 2021.
  * He et al. (2020) He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 9729–9738, 2020.
  * He et al. (2021) He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R.  Masked autoencoders are scalable vision learners.  _arXiv preprint arXiv:2111.06377_ , 2021.
  * Ioffe & Szegedy (2015) Ioffe, S. and Szegedy, C.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International conference on machine learning_ , pp. 448–456. PMLR, 2015.
  * Jacot et al. (2018) Jacot, A., Gabriel, F., and Hongler, C.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Ji et al. (2021) Ji, W., Deng, Z., Nakada, R., Zou, J., and Zhang, L.  The power of contrast for feature learning: A theoretical analysis.  _arXiv preprint arXiv:2110.02473_ , 2021.
  * Jing et al. (2022) Jing, L., Vincent, P., LeCun, Y., and Tian, Y.  Understanding dimensional collapse in contrastive self-supervised learning.  _ICLR_ , 2022.
  * Kawaguchi (2016) Kawaguchi, K.  Deep learning without poor local minima.  _NeurIPS_ , 2016.
  * Li & Yuan (2017) Li, Y. and Yuan, Y.  Convergence analysis of two-layer neural networks with relu activation.  _Advances in neural information processing systems_ , 30, 2017.
  * Lu et al. (2017) Lu, Z., Pu, H., Wang, F., Hu, Z., and Wang, L.  The expressive power of neural networks: A view from the width.  _Advances in neural information processing systems_ , 30, 2017.
  * Morcos et al. (2019) Morcos, A., Yu, H., Paganini, M., and Tian, Y.  One ticket to win them all: generalizing lottery ticket initializations across datasets and optimizers.  _Advances in neural information processing systems_ , 32, 2019.
  * Oord et al. (2018) Oord, A. v. d., Li, Y., and Vinyals, O.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2018.
  * Raghu et al. (2017) Raghu, M., Poole, B., Kleinberg, J., Ganguli, S., and Sohl-Dickstein, J.  On the expressive power of deep neural networks.  In _international conference on machine learning_ , pp. 2847–2854. PMLR, 2017.
  * Saunshi et al. (2022) Saunshi, N., Ash, J., Goel, S., Misra, D., Zhang, C., Arora, S., Kakade, S., and Krishnamurthy, A.  Understanding contrastive learning requires incorporating inductive biases.  _arXiv preprint arXiv:2202.14037_ , 2022.
  * Thompson et al. (1990) Thompson, J. M. T., Stewart, H. B., and Turner, R.  Nonlinear dynamics and chaos.  _Computers in Physics_ , 4(5):562–563, 1990.
  * Tian (2017) Tian, Y.  An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis.  In _International Conference on Machine Learning_ , pp. 3404–3413. PMLR, 2017.
  * Tian (2022) Tian, Y.  Deep contrastive learning is provably (almost) principal component analysis.  _arXiv preprint arXiv:2201.12680_ , 2022.
  * Tian et al. (2019) Tian, Y., Jiang, T., Gong, Q., and Morcos, A.  Luck matters: Understanding training dynamics of deep relu networks.  _arXiv preprint arXiv:1905.13405_ , 2019.
  * Tian et al. (2020) Tian, Y., Yu, L., Chen, X., and Ganguli, S.  Understanding self-supervised learning with dual deep networks.  _arXiv preprint arXiv:2010.00578_ , 2020.
  * Tian et al. (2021) Tian, Y., Chen, X., and Ganguli, S.  Understanding self-supervised learning dynamics without contrastive pairs.  In _International Conference on Machine Learning_ , pp. 10268–10278. PMLR, 2021.
  * Wang et al. (2021) Wang, X., Chen, X., Du, S. S., and Tian, Y.  Towards demystifying representation learning with non-contrastive self-supervision.  _arXiv preprint arXiv:2110.04947_ , 2021.
  * Wilson et al. (1997) Wilson, C. L., Blue, J. L., and Omidvar, O. M.  Training dynamics and neural network performance.  _Neural Networks_ , 10(5):907–923, 1997.
  * Yu et al. (2020) Yu, H., Edunov, S., Tian, Y., and Morcos, A. S.  Playing the lottery with rewards and multiple languages: lottery tickets in rl and nlp.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1xnXRVFwH.

##  Appendix A Proofs

###  A.1 Problem Setup (Sec. 2)

See 1

###### Proof.

Let ${\bm{g}}[i]:={\bm{g}}({\bm{x}}[i])$ and
${\bm{g}}^{\prime}[i]:={\bm{g}}({\bm{x}}[i^{\prime}])$ be the $i$-th and
$(i+N)$-th samples of ${\bm{g}}({\bm{x}})$. Therefore, both ${\bm{g}}$ and
${\bm{g}}^{\prime}$ has $N$ samples. Note that when $\alpha_{ij}\equiv
1/2N(N-1)$, for the first term we have:

$\displaystyle\sum_{i=1}^{N}\sum_{j\neq
i}\alpha_{ij}({\bm{g}}[i]-{\bm{g}}[j])({\bm{g}}[i]-{\bm{g}}[j])^{\top}=\frac{1}{N(N-1)}\sum_{i=1}^{N}\sum_{j=1}^{N}({\bm{g}}[i]-{\bm{g}}[j])({\bm{g}}[i]-{\bm{g}}[j])^{\top}$
(16) $\displaystyle=$
$\displaystyle\frac{N}{N-1}\left[\frac{1}{N}\sum_{i=1}^{N}{\bm{g}}[i]{\bm{g}}^{\top}[i]-\left(\frac{1}{N}\sum_{i=1}^{N}{\bm{g}}[i]\right)\left(\frac{1}{N}\sum_{i=1}^{N}{\bm{g}}[i]\right)^{\top}\right]$
$\displaystyle=$
$\displaystyle\frac{N}{N-1}\left[\mathbb{E}_{\mathrm{sample}}[{\bm{g}}{\bm{g}}^{\top}]-\mathbb{E}_{\mathrm{sample}}[{\bm{g}}]\mathbb{E}_{\mathrm{sample}}[{\bm{g}}^{\top}]\right]$

where
$\mathbb{E}_{\mathrm{sample}}[{\bm{g}}]:=\frac{1}{N}\sum_{i=1}^{N}{\bm{g}}[i]$
is the sample mean. Similarly for the second term, we have:

$\displaystyle\sum_{i=1}^{N}\left(\sum_{j\neq
i}\alpha_{ij}\right)({\bm{g}}[i]-{\bm{g}}^{\prime}[i])({\bm{g}}[i]-{\bm{g}}^{\prime}[i])^{\top}=\frac{1}{2N}\sum_{i=1}^{N}({\bm{g}}[i]-{\bm{g}}^{\prime}[i])({\bm{g}}[i]-{\bm{g}}^{\prime}[i])^{\top}$
(17) $\displaystyle=$
$\displaystyle\frac{1}{2}\mathbb{E}_{\mathrm{sample}}[({\bm{g}}-{\bm{g}}^{\prime})({\bm{g}}-{\bm{g}}^{\prime})^{\top}]$
(18)

Therefore, when $N\rightarrow+\infty$, we have

$\mathbb{C}_{\alpha}[{\bm{g}}]\rightarrow\mathbb{V}[{\bm{g}}]-\frac{1}{2}\mathbb{E}[({\bm{g}}-{\bm{g}}^{\prime})({\bm{g}}-{\bm{g}}^{\prime})^{\top}]$
(19)

Since ${\bm{x}}\sim p_{\mathrm{aug}}(\cdot|{\bm{x}}_{0})$ are augmented views
of the same input ${\bm{x}}_{0}\sim p(\cdot)$, splitting ${\bm{x}}$ that
contains $2N$ samples into ${\bm{x}}$ and ${\bm{x}}^{\prime}$ (each containing
$N$ samples), then we have:

$\displaystyle\mathbb{E}[({\bm{g}}-{\bm{g}}^{\prime})({\bm{g}}-{\bm{g}}^{\prime})^{\top}]$
$\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}_{{\bm{x}}_{0}\sim
p(\cdot)}[\mathbb{E}_{{\bm{x}},{\bm{x}}^{\prime}\sim
p_{\mathrm{aug}}(\cdot|{\bm{x}}_{0})}\left[({\bm{g}}-{\bm{g}}^{\prime})({\bm{g}}-{\bm{g}}^{\prime})^{\top}|{\bm{x}}_{0}]\right]$
(20) $\displaystyle=$ $\displaystyle\mathbb{E}_{{\bm{x}}_{0}\sim
p}\left[\mathbb{V}_{{\bm{x}}\sim
p_{\mathrm{aug}}}[{\bm{g}}({\bm{x}})|{\bm{x}}_{0}]\right]$ (21)

Using the law of total variation, we now have:

$\mathbb{C}_{\alpha}[{\bm{g}}({\bm{x}})]\rightarrow\mathbb{V}[{\bm{g}}({\bm{x}})]-\mathbb{E}_{p}[\mathbb{V}_{p_{\mathrm{aug}}}[{\bm{g}}({\bm{x}})|{\bm{x}}_{0}]]=\mathbb{V}_{p}[\mathbb{E}_{p_{\mathrm{aug}}}[{\bm{g}}({\bm{x}})|{\bm{x}}_{0}]]$
(22)

∎

##  Appendix B One-layer model (Sec. 3)

###  B.1 Computation of the two example models

Here we assume ReLU activation $h(x):=\max(x,0)$, which is a homogeneous
activation $h(x)=h^{\prime}(x)x$. Note that we consider $h^{\prime}(0)=0$.
Therefore, for any sample ${\bm{x}}$, if ${\bm{w}}^{\top}{\bm{x}}=0$, then we
don’t consider it to be included in the active region of ReLU, i.e.,
$\tilde{\bm{x}}^{{\bm{w}}}={\bm{x}}\cdot
h^{\prime}({\bm{w}}^{\top}{\bm{x}})=0$.

Let $z$ be a hidden binary variable and we could compute $A({\bm{w}})$ (here
$p_{0}:=\mathbb{P}[z=0]$ and $p_{1}:=\mathbb{P}[z=1]$):

$\mathbb{V}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{V}_{z}[\mathbb{E}[\tilde{\bm{x}}^{{\bm{w}}}|z]]+\mathbb{E}_{z}[\mathbb{V}[\tilde{\bm{x}}^{{\bm{w}}}|z]]=p_{0}p_{1}\Delta({\bm{w}})\Delta^{\top}({\bm{w}})+p_{0}\Sigma_{0}({\bm{w}})+p_{1}\Sigma_{1}({\bm{w}})$
(23)

where
$\Delta({\bm{w}}):=\mathbb{E}[\tilde{\bm{x}}|z=1]-\mathbb{E}[\tilde{\bm{x}}|z=0]$
and $\Sigma_{z}({\bm{w}}):=\mathbb{V}[\tilde{\bm{x}}|z]$.

Latent categorical model. If ${\bm{w}}={\bm{u}}_{m}$, let
$z:=\mathbb{I}(y=m)$. This leads to
$\Sigma_{1}({\bm{u}}_{m})=\Sigma_{0}({\bm{u}}_{m})=0$ and
$\Delta({\bm{u}}_{m})={\bm{u}}_{m}$. Therefore, we have:

$A({\bm{w}})\big{|}_{{\bm{w}}={\bm{u}}_{m}}:=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{V}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{P}[y=m]\left(1-\mathbb{P}[y=m]\right){\bm{u}}_{m}{\bm{u}}_{m}^{\top}$
(24)

Latent summation model. If ${\bm{w}}={\bm{u}}_{m}$, let
$z:=\mathbb{I}(y_{m}>0)$, then we can compute
$\Delta({\bm{u}}_{m})=y_{m}^{+}{\bm{u}}_{m}$,
$\Sigma_{1}({\bm{u}}_{m})=I-{\bm{u}}_{m}{\bm{u}}_{m}^{\top}$ and
$\Sigma_{0}({\bm{u}}_{m})=0$. Therefore, we have:

$A({\bm{w}})\big{|}_{{\bm{w}}={\bm{u}}_{m}}:=\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{{\bm{w}}}]=\mathbb{V}\left[\tilde{\bm{x}}\right]=(1-q_{m})^{2}{\bm{u}}_{m}{\bm{u}}_{m}^{\top}+q_{m}(I-{\bm{u}}_{m}{\bm{u}}_{m}^{\top})$
(25)

See 1

###### Proof.

(1) It is clear that if $A({\bm{w}}_{*}){\bm{w}}_{*}=\lambda_{*}{\bm{w}}_{*}$,
then $P^{\perp}_{{\bm{w}}_{*}}A({\bm{w}}_{*}){\bm{w}}_{*}=0$ so ${\bm{w}}_{*}$
is a critical point of Eqn. 6.

(2) For any unit direction $\|{\bm{u}}\|_{2}=1$ so that
${\bm{u}}^{\top}{\bm{w}}_{*}=0$, consider the perturbation
${\bm{v}}=\sqrt{1-\epsilon^{2}}{\bm{w}}_{*}+\epsilon{\bm{u}}$. Since
$\|{\bm{w}}_{*}\|_{2}=1$ we have $\|{\bm{v}}\|_{2}=1$.

Now let’s compute $P_{{\bm{v}}}^{\perp}A({\bm{v}}){\bm{v}}$. First, we have:

$\displaystyle P_{{\bm{v}}}^{\perp}$ $\displaystyle=$ $\displaystyle
I-{\bm{v}}{\bm{v}}^{\top}=I-\left(\sqrt{1-\epsilon^{2}}{\bm{w}}_{*}+\epsilon{\bm{u}}\right)\left(\sqrt{1-\epsilon^{2}}{\bm{w}}_{*}+\epsilon{\bm{u}}\right)^{\top}$
(26) $\displaystyle=$ $\displaystyle
I-{\bm{w}}_{*}{\bm{w}}_{*}^{\top}-\epsilon({\bm{u}}{\bm{w}}_{*}^{\top}+{\bm{w}}_{*}{\bm{u}}^{\top})+O(\epsilon^{2})$
(27) $\displaystyle=$ $\displaystyle
P^{\perp}_{{\bm{w}}_{*}}-\epsilon({\bm{u}}{\bm{w}}_{*}^{\top}+{\bm{w}}_{*}{\bm{u}}^{\top})+O(\epsilon^{2})$
(28)

So we have:

$\displaystyle P_{{\bm{v}}}^{\perp}A({\bm{w}}_{*}){\bm{v}}$ $\displaystyle=$
$\displaystyle
P^{\perp}_{{\bm{w}}_{*}}A({\bm{w}}_{*}){\bm{v}}-\epsilon({\bm{u}}{\bm{w}}_{*}^{\top}+{\bm{w}}_{*}{\bm{u}}^{\top})A({\bm{w}}_{*}){\bm{v}}+O(\epsilon^{2})$
(29) $\displaystyle=$ $\displaystyle
P^{\perp}_{{\bm{w}}_{*}}A({\bm{w}}_{*})\epsilon{\bm{u}}-\epsilon\lambda_{*}{\bm{u}}+O(\epsilon^{2})$
(30) $\displaystyle=$ $\displaystyle
P^{\perp}_{{\bm{w}}_{*}}(A({\bm{w}}_{*})-\lambda_{*}I)\epsilon{\bm{u}}+O(\epsilon^{2})$
(31)

The previous derivation is due to the fact that
$P^{\perp}_{{\bm{w}}_{*}}A({\bm{w}}_{*}){\bm{w}}_{*}=0$,
${\bm{u}}^{\top}A({\bm{w}}_{*}){\bm{w}}_{*}=0$ and
$P^{\perp}_{{\bm{w}}_{*}}{\bm{u}}={\bm{u}}$. Therefore, for
$P_{{\bm{v}}}^{\perp}A({\bm{v}}){\bm{v}}$, we can decompose it to two parts:

$\displaystyle P_{{\bm{v}}}^{\perp}A({\bm{v}}){\bm{v}}$ $\displaystyle=$
$\displaystyle
P_{{\bm{v}}}^{\perp}A({\bm{w}}_{*}){\bm{v}}+P_{{\bm{v}}}^{\perp}(A({\bm{v}})-A({\bm{w}}_{*})){\bm{v}}$
(32) $\displaystyle=$ $\displaystyle
P^{\perp}_{{\bm{w}}_{*}}(A({\bm{w}}_{*})-\lambda_{*}I)\epsilon{\bm{u}}+P_{{\bm{v}}}^{\perp}(A({\bm{v}})-A({\bm{w}}_{*})){\bm{v}}+O(\epsilon^{2})$
(33)

Therefore, since ${\bm{u}}^{\top}{\bm{w}}_{*}=0$, we have:

$\displaystyle{\bm{u}}^{\top}P^{\perp}_{{\bm{w}}_{*}}(A({\bm{w}}_{*})-\lambda_{*}I)\epsilon{\bm{u}}$
$\displaystyle=$
$\displaystyle{\bm{u}}^{\top}(I-{\bm{w}}_{*}{\bm{w}}_{*}^{\top})(A({\bm{w}}_{*})-\lambda_{*}I)\epsilon{\bm{u}}$
(34) $\displaystyle=$
$\displaystyle\epsilon{\bm{u}}^{\top}(A({\bm{w}}_{*})-\lambda_{*}I){\bm{u}}\leq-\lambda_{\mathrm{gap}}({\bm{w}}_{*})\epsilon+O(\epsilon^{2})$
(35)

and since $\|{\bm{u}}\|_{2}=\|{\bm{v}}\|_{2}=1$ and
$\|P_{\bm{v}}^{\perp}\|_{2}=1$, we have:

$|{\bm{u}}^{\top}P_{{\bm{v}}}^{\perp}(A({\bm{v}})-A({\bm{w}}_{*})){\bm{v}}|\leq\|(A({\bm{v}})-A({\bm{w}}_{*})){\bm{v}}\|_{2}\leq
L\|{\bm{v}}-{\bm{w}}_{*}\|_{2}\leq L\epsilon+O(\epsilon^{2})$ (36)

Therefore, we have:

${\bm{u}}^{\top}P_{{\bm{v}}}^{\perp}A({\bm{v}}){\bm{v}}\leq-(\lambda_{\mathrm{gap}}({\bm{w}}_{*})-L)\epsilon+O(\epsilon^{2})$
(37)

When $\lambda_{\mathrm{gap}}({\bm{w}}_{*})>L$ and we have
${\bm{u}}^{\top}P_{{\bm{v}}}^{\perp}A({\bm{v}}){\bm{v}}<0$ for any
${\bm{u}}\perp{\bm{w}}_{*}$ and sufficiently small $\epsilon$. Therefore, the
critical point ${\bm{w}}_{*}$ is stable. ∎

###  B.2 Finding critical points with initial guess (Sec. 3.2)

Notation. Let $\lambda_{i}({\bm{w}})$ and $\phi_{i}({\bm{w}})$ be the $i$-th
eigenvalue and unit eigenvector of $A({\bm{w}})$ where $\phi_{1}({\bm{w}})$ is
the largest. We first assume $A({\bm{w}})$ is positive definite (PD) and then
remove this assumption later. In this case,
$\lambda_{1}({\bm{w}})\geq\lambda_{2}({\bm{w}})\geq\ldots\geq\lambda_{d}({\bm{w}})>0$.
Let $c({\bm{w}}):={\bm{w}}^{\top}\phi_{1}({\bm{w}})$ be the inner product
between ${\bm{w}}$ and the maximal eigenvector of $A({\bm{w}})$.

Consider the following Power Iteration (PI) format:

$\tilde{\bm{w}}(t+1)\leftarrow
A({\bm{w}}(t)){\bm{w}}(t),\quad\quad\quad\quad{\bm{w}}(t+1)\leftarrow\frac{\tilde{\bm{w}}(t+1)}{\|\tilde{\bm{w}}(t+1)\|_{2}}$
(38)

Along the trajectory, let $\phi_{i}(t):=\phi_{1}(A({\bm{w}}(t)))$ be the
$i$-th unit eigenvector of $A({\bm{w}}(t))$ and $\lambda_{i}(t)$ to be the
$i$-th eigenvalue. Define $\delta{\bm{w}}(t):={\bm{w}}(t+1)-{\bm{w}}(t)$,
$\delta A(t):=A({\bm{w}}(t+1))-A({\bm{w}}(t))$, and

$c_{t}:=c({\bm{w}}(t))=\phi^{\top}_{1}(t){\bm{w}}(t),\quad
d_{t}:=\phi^{\top}_{1}(t){\bm{w}}(t+1)$ (39)

Then $-1\leq c_{t},d_{t}\leq 1$ since they are inner product of two unit
vectors.

See 2

###### Proof.

Note that if $c_{0}<0$, we can always use $-\phi_{1}({\bm{w}})$ as the maximal
eigenvector and follow Lemma 9, and notice that
$\|{\bm{w}}-{\bm{w}}(0)\|_{2}=\sqrt{2(1-{\bm{w}}^{\top}{\bm{w}}(0))}$, so

$\|{\bm{w}}-{\bm{w}}(0)\|_{2}\leq\frac{\sqrt{2(1+\gamma)(1-c_{0})}}{1-\sqrt{\gamma}}\quad\quad\Longleftrightarrow\quad\quad{\bm{w}}^{\top}{\bm{w}}(0)\geq\frac{c_{0}-c_{\gamma}}{1-c_{\gamma}}$
(40)

By Theorem 1, when $\lambda_{\mathrm{gap}}({\bm{w}}_{*})>L$, ${\bm{w}}_{*}$ is
stable. ∎

Remarks. Note that Lemma 9 assumes that along the trajectory
$\\{{\bm{w}}(t)\\}$, $\mu_{t}+\nu_{t}\leq\gamma$ holds. In Theorem 2, this can
not be assumed true until we prove that the entire trajectory is within
$B_{\gamma}$.

##  Appendix C Two layer case (Sec. 4)

See 2

###### Proof.

The output of the 2-layer network can be written as the following:

$f_{2l}=\sum_{k}v_{lk}h({\bm{w}}_{k}^{\top}{\bm{x}}_{k})$ (41)

For convenience, we use ${\bm{f}}_{1}:=[h({\bm{w}}_{k}^{\top}{\bm{x}}_{k})]$
to represent the column vector that collects all the outputs of intermediate
nodes, and ${\bm{v}}^{\top}_{l}$ is the $l$-th row vector in $V$.

According to Theorem 1 in (Tian, 2022), the gradient descent direction of
contrastive loss corresponds to the gradient ascent direction of the energy
function $\mathcal{E}_{\alpha}(\bm{\theta})$. From Eqn. 25 of that theorem, we
have (here we omit $1/2$ which is a global constant):

$\frac{\partial\mathcal{E}}{\partial\theta}=\sum_{l}\mathbb{C}_{\alpha}\left[\frac{\partial
f_{2l}}{\partial\theta},f_{2l}\right]$ (42)

Therefore, for $V=[v_{ik}]$ we have:

$\displaystyle\dot{\bm{v}}_{i}=\frac{\partial\mathcal{E}}{\partial{\bm{v}}_{i}}$
$\displaystyle=$ $\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[\frac{\partial
f_{2l}}{\partial{\bm{v}}_{i}},f_{2l}\right]$ (43) $\displaystyle=$
$\displaystyle\mathbb{C}_{\alpha}\left[{\bm{f}}_{1},{\bm{v}}^{\top}_{i}{\bm{f}}_{1}\right]$
(44) $\displaystyle=$
$\displaystyle\mathbb{C}_{\alpha}\left[{\bm{f}}_{1},{\bm{f}}_{1}\right]{\bm{v}}_{i}$
(45)

So we have $\dot{\bm{v}}_{i}=\mathbb{C}_{\alpha}[{\bm{f}}_{1}]{\bm{v}}_{i}$,
or $\dot{V}=V\mathbb{C}_{\alpha}[{\bm{f}}_{1}]$.

Now we compute $\partial\mathcal{E}/\partial{\bm{w}}_{k}$:

$\displaystyle\dot{\bm{w}}_{k}=\frac{\partial\mathcal{E}}{\partial{\bm{w}}_{k}}$
$\displaystyle=$ $\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[\frac{\partial
f_{2l}}{\partial{\bm{w}}_{k}},f_{2l}\right]$ (46) $\displaystyle=$
$\displaystyle\sum_{l}\mathbb{C}_{\alpha}[v_{lk}h^{\prime}({\bm{w}}_{k}^{\top}{\bm{x}}_{k}){\bm{x}}_{k},{\bm{v}}_{l}^{\top}{\bm{f}}_{1}]$
(47) $\displaystyle=$
$\displaystyle\sum_{l}v_{lk}\mathbb{C}_{\alpha}[\tilde{\bm{x}}_{k},{\bm{v}}_{l}^{\top}{\bm{f}}_{1}]$
(48) $\displaystyle=$
$\displaystyle\sum_{l}v_{lk}\mathbb{C}_{\alpha}\left[\tilde{\bm{x}}_{k},\sum_{k^{\prime}}v_{lk^{\prime}}h({\bm{w}}_{k^{\prime}}^{\top}{\bm{x}}_{k^{\prime}})\right]$
(49) $\displaystyle=$
$\displaystyle\sum_{k^{\prime}}\left(\sum_{l}v_{lk}v_{lk^{\prime}}\right)\mathbb{C}_{\alpha}\left[\tilde{\bm{x}}_{k},\tilde{\bm{x}}_{k^{\prime}}\right]{\bm{w}}_{k^{\prime}}$
(50) $\displaystyle=$
$\displaystyle\sum_{k^{\prime}}s_{kk^{\prime}}\mathbb{C}_{\alpha}\left[\tilde{\bm{x}}_{k},\tilde{\bm{x}}_{k^{\prime}}\right]{\bm{w}}_{k^{\prime}}$
(51)

where $S=[s_{kk^{\prime}}]=V^{\top}V=\sum_{l}{\bm{v}}_{l}{\bm{v}}_{l}^{\top}$.
Let ${\bm{w}}:=[{\bm{w}}_{1};\ldots;{\bm{w}}_{K}]$ and it leads to the
conclusion. When $M>1$, the proof is similar. ∎

Figure 5: Decomposition of the variance term $\mathbb{V}[\tilde{\bm{x}}]$.

See 3

###### Proof.

Use variance decomposition, we have:

$\mathbb{V}[\tilde{\bm{x}}]=\mathbb{E}_{z}\mathbb{V}[\tilde{\bm{x}}|z]+\mathbb{V}_{z}\mathbb{E}[\tilde{\bm{x}}|z]$
(52)

Remember that $\tilde{\bm{x}}_{km}$ is an abbreviation of gated input:

$\tilde{\bm{x}}_{km}:=\tilde{\bm{x}}_{k}^{{\bm{w}}_{km}}:={\bm{x}}_{k}\cdot
h^{\prime}({\bm{w}}^{\top}_{km}{\bm{x}}_{k})$ (53)

By conditional independence, we have

$\mathrm{Cov}[\tilde{\bm{x}}_{km},\tilde{\bm{x}}_{k^{\prime}m^{\prime}}|z]=0\quad\quad\quad\forall
k\neq k^{\prime}$ (54)

This is because $\tilde{\bm{x}}_{km}$ and
$\tilde{\bm{x}}_{k^{\prime}m^{\prime}}$ are deterministic functions of
${\bm{x}}_{k}$ and ${\bm{x}}_{k^{\prime}}$ and thus are also independent of
each other.

Let

$\tilde{\bm{x}}_{k}:=\left[\begin{array}[]{c}\tilde{\bm{x}}_{k1}\\\
\tilde{\bm{x}}_{k2}\\\ \ldots\\\
\tilde{\bm{x}}_{kM}\end{array}\right]\in\mathbb{R}^{Md}$ (55)

and $L_{k}:=\mathbb{E}_{z}\mathbb{V}[\tilde{\bm{x}}_{k}|z]$. Then we know that
$\mathbb{E}_{z}\mathbb{V}[\tilde{\bm{x}}|z]=\mathrm{diag}_{k}[L_{k}]$ is a
block diagonal matrix (See Fig. 5).

On the other hand, $\mathbb{V}_{z}\mathbb{E}[\tilde{\bm{x}}|z]$ is a low-rank
matrix:

$\mathbb{V}_{z}\mathbb{E}[\tilde{\bm{x}}|z]=\mathbb{E}_{z}\left[(\mathbb{E}[\tilde{\bm{x}}|z]-\mathbb{E}[\tilde{\bm{x}}])(\mathbb{E}[\tilde{\bm{x}}|z]-\mathbb{E}[\tilde{\bm{x}}])^{\top}\right]$
(56)

Let ${\bm{q}}_{c}:=\mathbb{E}[\tilde{\bm{x}}|z=c]$ and
${\bm{q}}_{-c}:=\mathbb{E}[\tilde{\bm{x}}|z\neq c]$, then we have:

$\displaystyle\mathbb{E}[\tilde{\bm{x}}|z=c]-\mathbb{E}[\tilde{\bm{x}}]$
$\displaystyle=$
$\displaystyle{\bm{q}}_{c}-\sum_{c}p_{c}{\bm{q}}_{c}=(1-p_{c})\left({\bm{q}}_{c}-\sum_{c^{\prime}\neq
c}\frac{p_{c^{\prime}}}{1-p_{c}}{\bm{q}}_{c^{\prime}}\right)$ (57)
$\displaystyle=$
$\displaystyle(1-p_{c})\left({\bm{q}}_{c}-\sum_{c^{\prime}\neq
c}\mathbb{P}[z=c^{\prime}|z\neq c]{\bm{q}}_{c^{\prime}}\right)$ (58)
$\displaystyle=$ $\displaystyle(1-p_{c})({\bm{q}}_{c}-{\bm{q}}_{-c})$ (59)

Therefore, we have:

$\displaystyle\mathbb{V}_{z}\mathbb{E}[\tilde{\bm{x}}|z]$ $\displaystyle=$
$\displaystyle\mathbb{E}_{z}\left[(\mathbb{E}[\tilde{\bm{x}}|z]-\mathbb{E}[\tilde{\bm{x}}])(\mathbb{E}[\tilde{\bm{x}}|z]-\mathbb{E}[\tilde{\bm{x}}])^{\top}\right]$
(60) $\displaystyle=$
$\displaystyle\sum_{c}p_{c}(1-p_{c})^{2}({\bm{q}}_{c}-{\bm{q}}_{-c})({\bm{q}}_{c}-{\bm{q}}_{-c})^{\top}$
(61) $\displaystyle=$
$\displaystyle\sum_{c}p_{c}(1-p_{c})^{2}\Delta(c)\Delta^{\top}(c)$ (62)

where

$\Delta(c):=\Delta(c;W):={\bm{q}}_{c}-{\bm{q}}_{-c}=\left[\begin{array}[]{c}\Delta_{11}(c)\\\
\ldots\\\ \Delta_{KM}(c)\end{array}\right]\in\mathbb{R}^{KMd}$ (63)

and

$\Delta_{km}(c):=\Delta_{km}(c;{\bm{w}}_{km}):=\mathbb{E}[\tilde{\bm{x}}_{km}|z=c]-\mathbb{E}[\tilde{\bm{x}}_{km}|z\neq
c]$ (64)

We can see that $\mathbb{V}_{z}\mathbb{E}[\tilde{\bm{x}}|z]$ is at most
rank-$C$, since it is a summation of $C$ rank-1 matrix.

In particular, when $C=2$, it is clear that $\Delta(0)=-\Delta(1)$ and thus
$\Delta(0)\Delta^{\top}(0)=\Delta(1)\Delta^{\top}(1)$ and
$\sum_{c}p_{c}(1-p_{c})^{2}=p_{0}p_{1}^{2}+p_{1}p_{0}^{2}=p_{0}p_{1}$. Hence
the conclusion. ∎

See 3

###### Proof.

Since $M=1$, each receptive field (RF) $R_{k}$ only output a single node with
output $f_{k}$. Let
$d_{k}:={\bm{w}}^{\top}_{k}L_{k}{\bm{w}}_{k}=\mathbb{E}_{z}\mathbb{V}[f_{k}|z]\geq
0$, $D:=\mathrm{diag}_{k}[d_{k}]$ is a diagonal matrix, and
${\bm{b}}:=[b_{k}]:=[{\bm{w}}_{k}^{\top}\Delta_{k}]\in\mathbb{R}^{K}$, Then

$\mathbb{V}[{\bm{f}}_{1}]=D+p_{0}p_{1}{\bm{b}}{\bm{b}}^{\top}$ (65)

is a diagonal matrix plus a rank-1 matrix. Since
$p_{0}p_{1}{\bm{b}}{\bm{b}}^{\top}$ is always PSD,
$\lambda=\lambda_{\max}(\mathbb{V}[{\bm{f}}_{1}])\geq\lambda_{\max}(D)=\max_{k}d_{k}$.
Then using Bunch–Nielsen–Sorensen formula (Bunch et al., 1978), for largest
eigenvector ${\bm{s}}$, we have:

$s_{k}=\frac{1}{Z}\frac{b_{k}}{d_{k}-\lambda}$ (66)

where $\lambda$ is the corresponding largest eigenvalue satisfying
$1+p_{0}p_{1}\sum_{k}\frac{b^{2}_{k}}{d_{k}-\lambda}=0$, and
$Z=\sqrt{\sum_{k}\left(\frac{b_{k}}{d_{k}-\lambda}\right)^{2}}$. Note that the
above is well-defined, since if $k^{*}=\arg\max_{k}d_{k}$ and $b_{k^{*}}\neq
0$, then $\lambda>\max_{k}{d_{k}}=d_{k^{*}}$. So $b_{k}/(d_{k}-\lambda)$ won’t
be infinite.

So we have:

$\displaystyle\dot{\bm{w}}_{k}$ $\displaystyle=$
$\displaystyle\sum_{k^{\prime}}s_{k}s_{k^{\prime}}\mathbb{C}_{\alpha}[\tilde{\bm{x}}_{k},\tilde{\bm{x}}_{k^{\prime}}]{\bm{w}}_{k^{\prime}}$
(67) $\displaystyle=$
$\displaystyle\sum_{k^{\prime}}s_{s}s_{s^{\prime}}(L_{k}\mathbb{I}(k=k^{\prime})+p_{0}p_{1}\Delta_{k}\Delta^{\top}_{k^{\prime}}){\bm{w}}_{k^{\prime}}$
$\displaystyle=$ $\displaystyle
s^{2}_{k}L_{k}{\bm{w}}_{k}+p_{0}p_{1}s_{k}\Delta_{k}\sum_{k^{\prime}}s_{k^{\prime}}\Delta^{\top}_{k^{\prime}}{\bm{w}}_{k^{\prime}}$
$\displaystyle=$ $\displaystyle
s^{2}_{k}L_{k}{\bm{w}}_{k}+\frac{p_{0}p_{1}b_{k}}{Z^{2}(d_{k}-\lambda)}\Delta_{k}\sum_{k^{\prime}}\frac{b^{2}_{k^{\prime}}}{d_{k^{\prime}}-\lambda}$
$\displaystyle=$ $\displaystyle
s^{2}_{k}L_{k}{\bm{w}}_{k}+\frac{1}{Z^{2}(\lambda-
d_{k})}\Delta_{k}\Delta^{\top}_{k}{\bm{w}}_{k}$

∎

See 4

###### Proof.

Since
$L_{k}({\bm{w}})=\sum_{l}g({\bm{u}}_{l}^{\top}{\bm{w}}){\bm{u}}_{l}{\bm{u}}_{l}^{\top}$,
we could write down its dynamics (we omit the projection $P^{\perp}_{\bm{w}}$
for now):

$\dot{\bm{w}}=L_{k}({\bm{w}}){\bm{w}}=\sum_{l}g({\bm{u}}_{l}^{\top}{\bm{w}}){\bm{u}}_{l}{\bm{u}}_{l}^{\top}{\bm{w}}$
(68)

Let $y_{l}(t):={\bm{u}}_{l}^{\top}{\bm{w}}(t)$, i.e., $y_{l}(t)$ is the
projected component of the weight ${\bm{w}}(t)$ onto the $l$-th direction,
i.e., a change of bases to orthonormal bases $\\{{\bm{u}}_{l}\\}$, then the
dynamics above can be written as

$\dot{y}_{l}=g(y_{l})y_{l}$ (69)

which is the same for all $l$, so we just need to study $\dot{x}=g(x)x$.
$g(x)>0$ is a linear increasing function, so we can assume $g(x)=ax+b$ with
$a>0$. Without loss of generality, we could just set $a=1$.

Then we just want to analyze the dynamics:

$\dot{y}_{l}=(y_{l}+b_{l})y_{l},\quad\quad\quad b_{l}>0$ (70)

which also includes the case of $L_{k}+c{\bm{u}}_{l}{\bm{u}}_{l}^{\top}$, that
basically sets $b_{l}=b+c$. Solving the dynamics leads to the following close-
form solution:

$\frac{y_{l}(t)}{y_{l}(t)+b_{l}}=\frac{y_{l}(0)}{y_{l}(0)+b_{l}}e^{b_{l}t}$
(71)

Figure 6: The one-dimensional dynamics (Eqn. 70) ($b_{l}>0$). There exists a
stable critical point $y_{l}=-b_{l}$ and one unstable critical point
$y_{l}=0$. When $y_{l}(0)>0$, the dynamics blows up in finite time
$t_{l}^{*}=\frac{1}{b_{l}}\ln\left(1+\frac{b_{l}}{y_{l}(0)}\right)$.

The 1-d dynamics has an unstable fixed points $y_{l}=0$ and a stable one
$y_{l}=-b_{l}<0$. Therefore, when the initial condition $y_{l}(0)<0$, the
dynamics will converge to $y_{l}(+\infty)=-b_{l}$, which is a finite number.
On the other hand, when $y_{l}(0)>0$, the dynamics has _finite-time blow-up_
(Thompson et al., 1990; Goriely & Hyde, 1998), i.e., there exists a critical
time $t^{*}_{l}<+\infty$ so that $y_{l}(t^{*}_{l})=+\infty$. See Fig. 6.

Note that this finite time blow-up is not physical, since we don’t take into
consideration of normalization $Z(t)$, which depends on all $y_{l}(t)$. The
real quality to be considered is $\hat{y}_{l}(t)=\frac{1}{Z(t)}y_{l}(t)$.
Fortunately, we don’t need to estimate $Z(t)$ since we are only interested in
the ratio:

$r_{l/l^{\prime}}(t):=\frac{\hat{y}_{l}(t)}{\hat{y}_{l^{\prime}}(t)}=\frac{y_{l}(t)}{y_{l^{\prime}}(t)}$
(72)

If for some $l$ and any $l^{\prime}\neq l$,
$r_{l/l^{\prime}}(t)\rightarrow\infty$, then $y_{l}(t)$ dominates and
$\hat{y}_{l}(t)\rightarrow 1$, i.e., the dynamics converges to ${\bm{u}}_{l}$.

Now our task is to know which initial condition of $y_{l}$ and $b_{l}$ makes
$r_{l/l^{\prime}}(t)\rightarrow+\infty$. By comparing the critical time we
know which component $l$ shoots up the earliest and that
$l^{*}=\arg\min_{l}t^{*}_{l}$ is the winner, without computing the
normalization constant $Z(t)$.

The critical time satisfies

$\frac{y_{l}(0)}{y_{l}(0)+b_{l}}e^{b_{l}t^{*}_{l}}=1$ (73)

so

$t^{*}_{l}=\frac{1}{b_{l}}\ln\left(1+\frac{b_{l}}{y_{l}(0)}\right)$ (74)

It is clear that when $y_{l}(0)$ is larger, the critical time $t^{*}_{l}$
becomes smaller and the $l$-th component becomes more advantageous over other
components.

For $b_{l}>0$, we have:

$\frac{\partial t^{*}_{l}}{\partial
b_{l}}=\frac{1}{b_{l}^{2}}\left[\frac{b_{l}/y_{l}(0)}{1+b_{l}/y_{l}(0)}-\ln(1+b_{l}/y_{l}(0))\right]<0$
(75)

where the last inequality is due to the fact that $\frac{x}{1+x}<\ln(1+x)$ for
$x>0$. Therefore, larger $b_{l}$ leads to smaller $t^{*}_{l}$. Since adding
$c{\bm{u}}_{l}{\bm{u}}_{l}^{\top}$ with $c>0$ to $L_{k}$ increase $b_{l}$, it
leads to smaller $t^{*}_{l}$ and thus increases the advantage of the $l$-th
component.

Therefore, larger $b_{l}$ and larger $y_{l}(0)$ both leads to smaller
$t^{*}_{l}$. For the same $t_{l}^{*}$, larger $b_{l}$ can trade for smaller
$y_{l}(0)$, i.e., larger attractive basin. ∎

Remark. To quantify the probability that a random weight initialization leads
to convergence of ${\bm{u}}_{l}$, we setup some notations.

Let the event $E_{l}$ be “a random weight initialization of ${\bm{y}}$ leads
to ${\bm{y}}\rightarrow{\bm{e}}_{l}$”, or equivalently
${\bm{w}}\rightarrow{\bm{u}}_{l}$. Let $Y_{l}$ be the random variable that
instantiates the initial value of $y_{l}(0)$ due to random weight
initialization. Then the convergence event $E_{l}$ is equivalent to the
following: (1) $Y_{l}>0$ (so that the $l$-component has the opportunity to
grow), and (2) $Y_{l}+\epsilon_{l}$ is the maximum over all $Y_{l^{\prime}}$
for any $l^{\prime}\neq l$, where $\epsilon_{l}$ is an advantage ($>0$) or
disadvantage ($<0$) achieved by having larger/smaller $b_{l}$ due to global
modulation (e.g., $c$). Therefore, we also call $\epsilon_{l}$ the _modulation
factor_.

Here we discuss about a simple case that $Y_{l}\sim U[-1,1]$ and for
$l^{\prime}\neq l$, $Y_{l}$ and $Y_{l^{\prime}}$ are independent. In this
case, for a given $l$, $\max_{l^{\prime}\neq l}Y_{l^{\prime}}$ is a random
variable that is independent of $Y_{l}$, and has cumulative density function
(CDF) $F_{\max}(x):=\mathbb{P}[\max_{l^{\prime}\neq l}Y_{l^{\prime}}\leq
x]=F^{d-1}(x)$, where $F(x)$ is the CDF for $Y_{l}$.

Then we have:

$\displaystyle\mathbb{P}[E_{l}]$ $\displaystyle=$
$\displaystyle\mathbb{P}\left[Y_{l}>0,Y_{l}+\epsilon_{l}\geq\max_{l^{\prime}\neq
l}Y_{l^{\prime}}\right]$ (76) $\displaystyle=$
$\displaystyle\int_{0}^{+\infty}\mathbb{P}\left[\max_{l^{\prime}\neq
l}Y_{l^{\prime}}\leq
Y_{l}+\epsilon_{l}\Big{|}Y_{l}=y_{l}\right]\mathbb{P}[Y_{l}=y_{l}]\mathrm{d}y_{l}$
(77) $\displaystyle=$
$\displaystyle\int_{0}^{+\infty}F^{d-1}(y_{l}+\epsilon_{l})\mathrm{d}F(y_{l})$
(78)

When $Y_{l}\sim U[-1,1]$, $F(x)=\min\left\\{\frac{1}{2}(x+1),1\right\\}$ has a
close form and we can compute the integral:

$\mathbb{P}[E_{l}]=\mathbb{P}\left[Y_{l}>0,Y_{l}+\epsilon_{l}\geq\max_{l^{\prime}\neq
l}Y_{l^{\prime}}\right]=\left\\{\begin{array}[]{cc}\frac{1}{2}&\epsilon_{l}>1\\\
\frac{1}{d}\left[1-\left(\frac{1+\epsilon_{l}}{2}\right)^{d}\right]+\frac{\epsilon_{l}}{2}&0\leq\epsilon_{l}\leq
1\\\
\frac{1}{d}\left[(1+\frac{\epsilon_{l}}{2})^{d}-(\frac{1+\epsilon_{l}}{2})^{d}\right]&-1<\epsilon_{l}<0\end{array}\right.$
(79)

We can see that the modulation factor $\epsilon_{l}$ plays an important role
in deciding the probability that ${\bm{w}}\rightarrow{\bm{u}}_{l}$:

  * •

No modulation. If $\epsilon_{l}=0$, then $\mathbb{P}[E_{l}]\sim\frac{1}{d}$.
This means that each dimension of ${\bm{y}}$ has equal probability to be the
dominant component after training;

  * •

Positive modulation. If $\epsilon_{l}>0$, then
$\mathbb{P}[E_{l}]\sim\frac{\epsilon_{l}}{2}$, and that particular $l$-th
component has much higher probability to become the dominant component,
independent of the dimensionality $d$. Furthermore, the stronger the
modulation, the higher the probability becomes.

  * •

Negative modulation. Finally, if $\epsilon_{l}<0$, since $1+\epsilon_{l}/2<1$,
$\mathbb{P}[E_{l}]$ decays exponentially w.r.t the dimensionality $d$.

See 5

###### Proof.

In the linear case, we have $\tilde{\bm{x}}_{km}={\bm{x}}_{k}$ since there is
no gating and all $M$ shares the same input ${\bm{x}}_{k}$. Therefore, we can
write down the dynamics of ${\bm{w}}_{km}$ as the following:

$\displaystyle\dot{\bm{w}}_{km}$ $\displaystyle=$
$\displaystyle\sum_{k^{\prime},m^{\prime}}s_{km,k^{\prime}m^{\prime}}\mathbb{C}_{\alpha}[\tilde{\bm{x}}_{km},\tilde{\bm{x}}_{k^{\prime}m^{\prime}}]{\bm{w}}_{k^{\prime}m^{\prime}}$
(80) $\displaystyle=$
$\displaystyle\sum_{k^{\prime},m^{\prime}}s_{km,k^{\prime}m^{\prime}}\mathbb{C}_{\alpha}[{\bm{x}}_{k},{\bm{x}}_{k^{\prime}}]{\bm{w}}_{k^{\prime}m^{\prime}}$
(81)

Now we use the fact that the top-level learns fast so that
$s_{km,k^{\prime}m^{\prime}}=s_{km}s_{k^{\prime}m^{\prime}}$, which gives:

$\displaystyle\dot{\bm{w}}_{km}$ $\displaystyle=$ $\displaystyle
s_{km}\sum_{k^{\prime},m^{\prime}}s_{k^{\prime}m^{\prime}}\mathbb{C}_{\alpha}[{\bm{x}}_{k},{\bm{x}}_{k^{\prime}}]{\bm{w}}_{k^{\prime}m^{\prime}}$
(82) $\displaystyle=$ $\displaystyle
s_{km}\mathbb{C}_{\alpha}\left[{\bm{x}}_{k},\sum_{k^{\prime},m^{\prime}}s_{k^{\prime}m^{\prime}}{\bm{w}}^{\top}_{k^{\prime}m^{\prime}}{\bm{x}}_{k^{\prime}}\right]$
(83)

Let
${\bm{b}}_{k}(W,V):=\mathbb{C}_{\alpha}\left[{\bm{x}}_{k},\sum_{k^{\prime},m^{\prime}}s_{k^{\prime}m^{\prime}}{\bm{w}}^{\top}_{k^{\prime}m^{\prime}}{\bm{x}}_{k^{\prime}}\right]$
be a linear function of $W$, and we have:

$\dot{\bm{w}}_{km}=s_{km}{\bm{b}}_{k}(W,V)$ (84)

Since ${\bm{b}}_{k}$ is independent of $m$, all $\dot{\bm{w}}_{km}$ are co-
linear.

For the second part, first all if $W^{*}$ is a critical point, we have the
following two facts:

  * •

Since there exists $m$ so that $s_{km}\neq 0$, we know that
${\bm{b}}_{k}(W^{*})=0$;

  * •

If $W^{*}$ contains two distinct filters
${\bm{w}}_{k1}=\bm{\mu}_{1}\neq{\bm{w}}_{k2}=\bm{\mu}_{2}$ covering the same
receptive field $R_{k}$, then by symmetry of the weights, ${W^{\prime}}^{*}$
in which ${\bm{w}}_{k1}=\bm{\mu}_{2}$ and ${\bm{w}}_{k2}=\bm{\mu}_{1}$, is
also a critical point.

Then for any $c\in[0,1]$, since ${\bm{b}}_{k}(W)$ is linear w.r.t. $W$, for
the linear combination $W^{c}:=cW^{*}+(1-c){W^{\prime}}^{*}$, we have:

${\bm{b}}_{k}(W^{c})={\bm{b}}_{k}(cW^{*}+(1-c){W^{\prime}}^{*})=c{\bm{b}}_{k}(W^{*})+(1-c){\bm{b}}_{k}({W^{\prime}}^{*})=0$
(85)

Therefore, $W^{c}$ is also a critical point, in which
${\bm{w}}_{k1}=c\bm{\mu}_{1}+(1-c)\bm{\mu}_{2}$ and
${\bm{w}}_{k2}=(1-c)\bm{\mu}_{1}+c\bm{\mu}_{2}$. In particular when $c=1/2$,
${\bm{w}}_{k1}={\bm{w}}_{k2}$. Repeating this process for different $m$, we
could finally reach a critical point in which all
${\bm{w}}_{km}={\bm{w}}_{k}$. ∎

###  C.1 Specialization of weights in Nonlinear setting

Following Lemma 2 and Lemma 3, we have the following dynamics:

$\displaystyle\left[\begin{array}[]{c}\dot{\bm{w}}_{k1}\\\
\dot{\bm{w}}_{k2}\end{array}\right]$ $\displaystyle=$ $\displaystyle
P^{\perp}_{k}M({\bm{w}}_{k1},{\bm{w}}_{k2})\left[\begin{array}[]{c}{\bm{w}}_{k1}\\\
{\bm{w}}_{k2}\end{array}\right],$ (90)

where $P^{\perp}_{k}:=\left[\begin{array}[]{cc}P^{\perp}_{{\bm{w}}_{k1}}&0\\\
0&P^{\perp}_{{\bm{w}}_{k2}}\end{array}\right]$ is a projection matrix, and
$M({\bm{w}}_{k1},{\bm{w}}_{k2})$ is defined as:

$\displaystyle M({\bm{w}}_{k1},{\bm{w}}_{k2})$ $\displaystyle:=$
$\displaystyle\left[\begin{array}[]{cc}s_{k1}^{2}A_{k}({\bm{w}}_{k1})&s_{k1}s_{k2}A_{k}({\bm{w}}_{k1},{\bm{w}}_{k2})\\\
s_{k1}s_{k2}A^{\top}_{k}({\bm{w}}_{k1},{\bm{w}}_{k2})&s^{2}_{k2}A_{k}({\bm{w}}_{k2})\end{array}\right]$
(93) $\displaystyle+$
$\displaystyle\sum_{c=0}^{C-1}p_{c}(1-p_{c})^{2}\left[\begin{array}[]{cc}s_{k1}^{2}\Delta_{k1}(c)\Delta_{k1}^{\top}(c)&s_{k1}s_{k2}\Delta_{k1}(c)\Delta_{k2}^{\top}(c)\\\
s_{k1}s_{k2}\Delta_{k2}(c)\Delta_{k1}^{\top}(c)&s^{2}_{k2}\Delta_{k2}(c)\Delta_{k2}^{\top}(c)\end{array}\right]$
(96)

Here
$\Delta_{km}(c)=\mathbb{E}[\tilde{\bm{x}}_{km}|z=c]-\mathbb{E}[\tilde{\bm{x}}_{km}|z\neq
c]$, and $A_{k}$ is the following:

$A_{k}({\bm{w}},{\bm{w}}^{\prime})=\mathbb{E}_{z}[\mathrm{Cov}[\tilde{\bm{x}}_{k}^{{\bm{w}}},\tilde{\bm{x}}_{k}^{{\bm{w}}^{\prime}}|z]]$
(97)

with the abbreviation
$A_{k}({\bm{w}}):=A_{k}({\bm{w}},{\bm{w}})=\mathbb{E}_{z}[\mathbb{V}[\tilde{\bm{x}}_{k}^{{\bm{w}}}|z]]$.
Note that if the activation is linear, then
$A_{k}({\bm{w}},{\bm{w}}^{\prime})=\mathbb{E}_{z}[\mathbb{V}[{\bm{x}}_{k}|z]]$,
which is independent of ${\bm{w}}$ and ${\bm{w}}^{\prime}$.

Case 1: $A_{k}({\bm{w}},{\bm{w}}^{\prime})={\bm{w}}{\bm{w}}^{\prime\top}$.
This is a simplest case. Then since
$P^{\perp}_{\bm{w}}A_{k}({\bm{w}},{\bm{w}}^{\prime})=0$, only the global
modulation term $\Delta_{k}$ will have impact on the final dynamics. Depending
on the initialization of ${\bm{w}}_{k1}$, the weight may converge to different
patterns specified by $\Delta_{k1}$, which unlike in the linear case, is now a
function of ${\bm{w}}_{k1}$. Same for ${\bm{w}}_{k2}$. This is the
specialization behavior of different weights. On the other hand, whether or
not this assumption can be achieved by any distribution is left for future
work.

Case 2: Spherical Gaussian. In (Tian, 2017), when the input is spherical
Gaussian (i.e., ${\bm{x}}\sim N(0,I)$) and the nonlinearity is ReLU, then we
could compute $A_{k}$ analytically (see Theorem 1 in (Tian, 2017)):

$A_{k}({\bm{w}},{\bm{w}}^{\prime}){\bm{w}}^{\prime}=\frac{\pi-\theta}{2\pi}{\bm{w}}^{\prime}+g(\theta){\bm{w}}$
(98)

where $\theta$ is the angle between ${\bm{w}}$ and ${\bm{w}}^{\prime}$ and
$g(\theta)$ is certain scalar function of $\theta$. Therefore, we have:

$P^{\perp}_{\bm{w}}A_{k}({\bm{w}},{\bm{w}}^{\prime}){\bm{w}}^{\prime}=\frac{\pi-\theta}{2\pi}P^{\perp}_{\bm{w}}{\bm{w}}^{\prime}$
(99)

If we check back with the original dynamics, this corresponds to the case that
$A_{k}({\bm{w}},{\bm{w}}^{\prime})=\frac{\pi-\theta}{2\pi}I$. Therefore, the
cross term $A_{k}({\bm{w}},{\bm{w}}^{\prime})$ is “weaker” than diagonal term
$A_{k}({\bm{w}})$, which is $\frac{1}{2}I$. The weaker cross term allows more
independent (and less correlated) evolvement of ${\bm{w}}_{k1}$ and
${\bm{w}}_{k2}$ from global modulation $\Delta_{k1}$ and $\Delta_{k2}$, in
particular when their angle $\theta$ is large, i.e., distinct weights. In
contrast, when the activation is linear, $A_{k}=I$ everywhere and the gradient
is perfectly co-linear (see Theorem 5).

We leave more realistic and more formal treatment of the interplay among
different weights $\\{{\bm{w}}_{km}\\}$ at the same receptive field $R_{k}$ as
the future work.

##  Appendix D Analysis of Batch Normalization

We consider BatchNorm right after ${\bm{f}}$:
$f^{\mathrm{bn}}_{k}[i]=(f_{k}[i]-\mu_{k})/\sigma_{k}$, where $\mu_{k}$ and
$\sigma_{k}$ are the batch statistics computed from BatchNorm on all $2N$
samples in a batch:

$\displaystyle\mu_{k}$ $\displaystyle:=$
$\displaystyle\frac{1}{2N}\sum_{i}f_{k}[i]+f_{k}[i^{\prime}]$ (100)
$\displaystyle\sigma^{2}_{k}$ $\displaystyle:=$
$\displaystyle\frac{1}{2N}\sum_{i}(f_{k}[i]-\mu_{k})^{2}+(f_{k}[i^{\prime}]-\mu_{k})^{2}$
(101)

When $N\rightarrow+\infty$, we have $\mu_{k}\rightarrow\mathbb{E}[f_{k}]$ and
$\sigma^{2}_{k}\rightarrow\mathbb{V}[f_{k}]={\bm{w}}_{k}^{\top}\mathbb{V}[\tilde{\bm{x}}_{k}]{\bm{w}}_{k}$.

Let $\tilde{\bm{x}}^{\mathrm{bn}}_{k}:=\sigma_{k}^{-1}\tilde{\bm{x}}_{k}$ and
$\tilde{\bm{x}}^{\mathrm{bn}}:=\left[\begin{array}[]{c}\tilde{\bm{x}}^{\mathrm{bn}}_{1}\\\
\tilde{\bm{x}}^{\mathrm{bn}}_{2}\\\ \ldots\\\
\tilde{\bm{x}}^{\mathrm{bn}}_{K}\end{array}\right]$. When computing gradient
through BatchNorm layer, we consider the following variant:

######  Definition 2 (mean-backprop BatchNorm).

When computing backpropagated gradient through BatchNorm, we only backprop
through $\mu_{k}$.

This leads to a model dynamics that has a very similar form as Lemma 2:

######  Lemma 4 (Dynamics with mean-backprop BatchNorm).

With mean-backprop BatchNorm (Def. 2), the dynamics is:

$\dot{V}=V\mathbb{C}_{\alpha}[{\bm{f}}^{\mathrm{bn}}_{1}],\quad\quad\quad\quad\dot{\bm{w}}=\left[(S\otimes{\bm{1}}_{d}{\bm{1}}_{d}^{\top})\circ\mathbb{C}_{\alpha}[\tilde{\bm{x}}^{\mathrm{bn}}]\right]{\bm{w}}$
(102)

###### Proof.

The proof is similar to Lemma 2. For $\dot{V}$ it is the same by replacing
${\bm{f}}_{1}$ with ${\bm{f}}^{\mathrm{bn}}_{1}$, which is the input to the
top layer.

For $\dot{\bm{w}}$, similarly we have:

$\displaystyle\dot{\bm{w}}_{k}=\frac{\partial\mathcal{E}}{\partial{\bm{w}}_{k}}$
$\displaystyle=$ $\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[\frac{\partial
f_{2l}}{\partial{\bm{w}}_{k}},f_{2l}\right]$ (103) $\displaystyle=$
$\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[v_{lk}\frac{\partial
f^{\mathrm{bn}}_{1k}}{\partial{\bm{w}}_{k}},\sum_{k^{\prime}}v_{lk^{\prime}}f^{\mathrm{bn}}_{1k^{\prime}}\right]$
(104) $\displaystyle=$
$\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[v_{lk}\frac{\partial
f^{\mathrm{bn}}_{1k}}{\partial{\bm{w}}_{k}},\sum_{k^{\prime}}v_{lk^{\prime}}(f_{1k^{\prime}}-\mu_{k^{\prime}})\sigma^{-1}_{k^{\prime}}\right]$
(105)

Note that
$\mathbb{C}_{\alpha}[\cdot,\mu_{k^{\prime}}\sigma^{-1}_{k^{\prime}}]=0$ since
$\mu_{k^{\prime}}$ and $\sigma_{k^{\prime}}$ are statistics of the batch and
is constant. On the other hand, for $\partial
f^{\mathrm{bn}}_{1k}/\partial{\bm{w}}_{k}$, we have:

$\frac{\partial
f^{\mathrm{bn}}_{1k}}{\partial{\bm{w}}_{k}}=\frac{1}{\sigma_{k}}\left(\frac{\partial
f_{1k}}{\partial{\bm{w}}_{k}}-\frac{\partial\mu_{k}}{\partial{\bm{w}}_{k}}\right)-\frac{f_{1k}^{\mathrm{bn}}}{\sigma_{k}}\frac{\partial\sigma_{k}}{\partial{\bm{w}}_{k}}$
(106)

Note that

$\frac{\partial\mu_{k}}{\partial{\bm{w}}_{k}}=\mathbb{E}_{\mathrm{sample}}[\tilde{\bm{x}}_{k}]$
(107)

where $\mathbb{E}_{\mathrm{sample}}[\cdot]$ is the sample mean, which is a
constant over the batch. Therefore
$\mathbb{C}_{\alpha}\left[\cdot,\partial\mu_{k}/\partial{\bm{w}}_{k}\right]=0$.
For mean-backprop BatchNorm, since the gradient didn’t backpropagate through
the variance, the second term is simply zero. Therefore, we have:

$\displaystyle\dot{\bm{w}}_{k}$ $\displaystyle=$
$\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[v_{lk}\sigma_{k}^{-1}\frac{\partial
f_{1k}}{\partial{\bm{w}}_{k}},\sum_{k^{\prime}}v_{lk^{\prime}}f_{1k^{\prime}}\sigma_{k^{\prime}}^{-1}\right]$
(108) $\displaystyle=$
$\displaystyle\sum_{l}\mathbb{C}_{\alpha}\left[v_{lk}\sigma_{k}^{-1}\tilde{\bm{x}}_{k},\sum_{k^{\prime}}v_{lk^{\prime}}{\bm{w}}_{k^{\prime}}^{\top}\tilde{\bm{x}}_{k^{\prime}}\sigma_{k^{\prime}}^{-1}\right]$
(109) $\displaystyle=$
$\displaystyle\sum_{k^{\prime}}s_{kk^{\prime}}\mathbb{C}_{\alpha}[\sigma^{-1}_{k}\tilde{\bm{x}}_{k},\sigma^{-1}_{k^{\prime}}\tilde{\bm{x}}_{k^{\prime}}]{\bm{w}}_{k^{\prime}}$
(110)

Let $\tilde{\bm{x}}^{\mathrm{bn}}_{k}:=\sigma_{k}^{-1}\tilde{\bm{x}}_{k}$ and
$\tilde{\bm{x}}^{\mathrm{bn}}:=\left[\begin{array}[]{c}\tilde{\bm{x}}^{\mathrm{bn}}_{1}\\\
\tilde{\bm{x}}^{\mathrm{bn}}_{2}\\\ \ldots\\\
\tilde{\bm{x}}^{\mathrm{bn}}_{K}\end{array}\right]\in\mathbb{R}^{Kd}$. The
conclusion follows. ∎

######  Corollary 3 (Dynamics of ${\bm{w}}_{k}$ under conditional independence
and BatchNorm).

Let

$\displaystyle L^{\mathrm{bn}}_{k}$ $\displaystyle:=$
$\displaystyle\mathbb{E}_{z}[\mathbb{V}[f^{\mathrm{bn}}_{k}|z]]=\sigma^{-2}_{k}\mathbb{E}_{z}[\mathbb{V}[f_{k}|z]]=\sigma^{-2}_{k}L_{k}$
(111) $\displaystyle d^{\mathrm{bn}}_{k}$ $\displaystyle:=$
$\displaystyle\sigma^{-2}_{k}d_{k}$ (112)
$\displaystyle\Delta^{\mathrm{bn}}_{k}$ $\displaystyle:=$
$\displaystyle\sigma^{-1}_{k}\Delta_{k}=\mathbb{E}[\tilde{\bm{x}}_{k}^{\mathrm{bn}}|z=1]-\mathbb{E}[\tilde{\bm{x}}_{k}^{\mathrm{bn}}|z=0]$
(113)

and $\lambda^{{\mathrm{bn}}}$ be the maximal eigenvalue of
$\mathbb{V}[{\bm{f}}^{\mathrm{bn}}_{1}]$. Then we have

  * •

(1) $\lambda^{\mathrm{bn}}\geq\max_{k}d^{\mathrm{bn}}_{k}$;

  * •

(2) For $\lambda^{\mathrm{bn}}$, the associated unit eigenvector is

${\bm{s}}^{\mathrm{bn}}:=\frac{1}{Z^{\mathrm{bn}}}\left[\frac{{\bm{w}}_{k}^{\top}\Delta^{\mathrm{bn}}_{k}}{\lambda^{\mathrm{bn}}-d^{\mathrm{bn}}_{k}}\right]\in\mathbb{R}^{K},$

where $Z^{\mathrm{bn}}$ is the normalization constant;

  * •

(3) the dynamics of ${\bm{w}}_{k}$ is given by:

$\dot{\bm{w}}_{k}=\left[(s^{\mathrm{bn}}_{k})^{2}L^{\mathrm{bn}}_{k}+\frac{1}{(Z^{\mathrm{bn}})^{2}(\lambda^{\mathrm{bn}}-d^{\mathrm{bn}}_{k})}(\Delta^{\mathrm{bn}}_{k})(\Delta^{\mathrm{bn}}_{k})^{\top}\right]{\bm{w}}_{k}$
(114)

###### Proof.

Similar to Theorem 3. ∎

Remarks In the presence of BatchNorm, Lemma 3 still holds, since it only
depends on the generative structure of the data. Therefore, we have

$\sigma^{2}_{k}\rightarrow\mathbb{V}[f_{k}]={\bm{w}}^{\top}_{k}\mathbb{V}[\tilde{\bm{x}}_{k}]{\bm{w}}_{k}=d_{k}+p_{0}p_{1}\left({\bm{w}}^{\top}_{k}\Delta_{k}\right)^{2}$

and thus

$d^{\mathrm{bn}}_{k}=\sigma_{k}^{-2}d_{k}\rightarrow\frac{d_{k}}{d_{k}+p_{0}p_{1}\left({\bm{w}}^{\top}_{k}\Delta_{k}\right)^{2}}=\frac{1}{1+p_{0}p_{1}\left({\bm{w}}^{\top}_{k}\Delta_{k}/\sqrt{d_{k}}\right)^{2}}$

becomes more uniform, normalizing the effect of magnitude of the input. As a
result, $\lambda^{\mathrm{bn}}-d^{\mathrm{bn}}_{k}$ becomes almost constant
across different receptive field $R_{k}$ and won’t lead to slowness of feature
learning.

##  Appendix E More experiments

We also do more experiments to further verify our theoretical findings.

Overall matching score $\bar{\chi}_{+}{}$ and overall irrelevant-matching
score $\bar{\chi}_{-}$. As defined in the main text (Eqn. 15), the matching
score $\chi_{+}(R_{k})$ is the degree of matching between learned weights and
the embeddings of the subset $R^{\mathrm{g}{}}_{k}$ of tokens that are allowed
in the global patterns at each receptive field $R_{k}$. And the overall
matching score $\bar{\chi}_{+}{}$ is $\bar{\chi}_{+}{}$ averaged over all
receptive fields:

$\chi_{+}(R_{k})=\frac{1}{P}\sum_{a\in
R^{\mathrm{g}}_{k}}\max_{m}\frac{{\bm{w}}^{\top}_{km}{\bm{u}}_{a}}{\|{\bm{w}}_{km}\|_{2}\|{\bm{u}}_{a}\|_{2}},\quad\quad\quad\quad\bar{\chi}_{+}=\frac{1}{K}\sum_{k}\chi_{+}(R_{k})$
(115)

Similarly, we can also define _irrelevant-matching score_ $\chi_{-}(R_{k})$
which is the degree of matching between learned weights and the embeddings of
the tokens that are NOT in the subset $R^{\mathrm{g}{}}_{k}$ at each receptive
field $R_{k}$. And the overall _irrelevant-matching_ score $\bar{\chi}_{-}$ is
defined similarly.

$\chi_{-}(R_{k})=\frac{1}{P}\sum_{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}a\notin
R^{\mathrm{g}}_{k}}}\max_{m}\frac{{\bm{w}}^{\top}_{km}{\bm{u}}_{a}}{\|{\bm{w}}_{km}\|_{2}\|{\bm{u}}_{a}\|_{2}},\quad\quad\quad\quad\bar{\chi}_{-}=\frac{1}{K}\sum_{k}\chi_{-}(R_{k})$
(116)

Ideally, we want to see high overall matching score $\bar{\chi}_{+}$ and low
overall irrelevant-matching score $\bar{\chi}_{-}{}$, which means that the
important patterns in $R^{\mathrm{g}}_{k}$ (i.e., the patterns that are
allowed in the global generators) are learned, but noisy patterns that are not
part of the global patterns (i.e., the generators) are not learned. Tbl. 3
shows that this indeed is the case.

Non-uniformity $\zeta$ and how BatchNorm interacts with it. Tbl. 5 shows that
BN with ReLU activations handles large non-uniformity (large $\zeta$) very
well, compared to the case without BN. Note that in the real-world scenario,
features from different channels/modalities indeed will have very different
scales, and some local features that turn out to be super important to global
features, can have very small scale. In such cases, normalization techniques
(e.g., BatchNorm) can be very useful and our formulation justifies it in a
mathematically consistent way.

Selectively Backpropagating $\mu_{k}$ and $\sigma_{k}^{2}$ in BatchNorm. In
our analysis of BatchNorm, we assume that gradient backpropagating the mean
statistics $\mu_{k}$, but not variance $\sigma_{k}^{2}$ (see Def. 2). Note
that this is different from regular BatchNorm, in which both $\mu_{k}$ and
$\sigma_{k}^{2}$ get backpropagated gradients. Therefore, we test how this
modified BN affects the matching score $\bar{\chi}_{+}{}$: we change whether
$\mu_{k}$ and $\sigma_{k}^{2}$ gets backpropagated gradients, while the
forward pass remains the same, yielding the four following variants:

$\displaystyle f^{\mathrm{bn}}_{k}[i]$
$\displaystyle:=\frac{f^{\mathrm{bn}}_{k}[i]-\mu_{k}}{\sigma_{k}}$ (Vanilla
BatchNorm) $\displaystyle f^{\mathrm{bn}}_{k}[i]$
$\displaystyle:=\frac{f^{\mathrm{bn}}_{k}[i]-\text{stop-
gradient}(\mu_{k})}{\sigma_{k}}$ (BatchNorm with backpropated $\sigma_{k}$)
$\displaystyle f^{\mathrm{bn}}_{k}[i]$
$\displaystyle:=\frac{f^{\mathrm{bn}}_{k}[i]-\mu_{k}}{\text{stop-
gradient}(\sigma_{k})}$ (BatchNorm with backpropated $\mu_{k}$) $\displaystyle
f^{\mathrm{bn}}_{k}[i]$
$\displaystyle:=\frac{f^{\mathrm{bn}}_{k}[i]-\text{stop-
gradient}(\mu_{k})}{\text{stop-gradient}(\sigma_{k})}$ (BatchNorm without
backpropating statistics)

As shown in Tbl. 6, it is interesting to see that if $\sigma^{2}_{k}$ is not
backpropagated, then the matching score $\bar{\chi}_{+}{}$ is actually better.
This justifies our BN variant.

Quadratic versus InfoNCE loss. Tbl. 4 shows that quadratic loss (constant
pairwise importance $\alpha$) shows worse matching score than InfoNCE. A high-
level intuition is that InfoNCE dynamically adjusts the pairwise importance
$\alpha$ (i.e., the focus of different sample pairs) during training to focus
on the most important sample pairs, which makes learning patterns more
efficient. We leave a comprehensive study for future work.

Overall matching score $\bar{\chi}_{+}{}$
---
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
Linear | ReLU | Linear | ReLU | Linear | ReLU | Linear | ReLU
1 | $0.95\\!\pm\\!0.00$ | $0.31\\!\pm\\!0.23$ | $0.79\\!\pm\\!0.02$ | $0.75\\!\pm\\!0.11$ | $0.67\\!\pm\\!0.03$ | $0.70\\!\pm\\!0.06$ | $0.60\\!\pm\\!0.01$ | $0.68\\!\pm\\!0.05$
2 | $0.95\\!\pm\\!0.00$ | $0.61\\!\pm\\!0.13$ | $0.81\\!\pm\\!0.01$ | $0.90\\!\pm\\!0.09$ | $0.70\\!\pm\\!0.00$ | $0.88\\!\pm\\!0.04$ | $0.63\\!\pm\\!0.01$ | $0.93\\!\pm\\!0.02$
5 | $0.96\\!\pm\\!0.00$ | $0.93\\!\pm\\!0.06$ | $0.85\\!\pm\\!0.01$ | $0.99\\!\pm\\!0.02$ | $0.73\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.66\\!\pm\\!0.01$ | $0.99\\!\pm\\!0.01$
10 | $0.96\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.86\\!\pm\\!0.01$ | $1.00\\!\pm\\!0.00$ | $0.77\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.68\\!\pm\\!0.01$ | $1.00\\!\pm\\!0.00$

Overall matching score $\bar{\chi}_{-}{}$
---
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
Linear | ReLU | Linear | ReLU | Linear | ReLU | Linear | ReLU
1 | $0.06\\!\pm\\!0.00$ | $-0.06\\!\pm\\!0.03$ | $0.12\\!\pm\\!0.01$ | $0.02\\!\pm\\!0.03$ | $0.17\\!\pm\\!0.01$ | $0.01\\!\pm\\!0.01$ | $0.22\\!\pm\\!0.00$ | $-0.00\\!\pm\\!0.00$
2 | $0.08\\!\pm\\!0.01$ | $-0.01\\!\pm\\!0.02$ | $0.19\\!\pm\\!0.00$ | $0.02\\!\pm\\!0.02$ | $0.23\\!\pm\\!0.01$ | $0.00\\!\pm\\!0.00$ | $0.27\\!\pm\\!0.01$ | $0.01\\!\pm\\!0.02$
5 | $0.16\\!\pm\\!0.02$ | $0.02\\!\pm\\!0.02$ | $0.24\\!\pm\\!0.02$ | $-0.00\\!\pm\\!0.00$ | $0.27\\!\pm\\!0.01$ | $0.00\\!\pm\\!0.00$ | $0.30\\!\pm\\!0.01$ | $0.04\\!\pm\\!0.01$
10 | $0.23\\!\pm\\!0.00$ | $-0.00\\!\pm\\!0.00$ | $0.27\\!\pm\\!0.01$ | $0.00\\!\pm\\!0.00$ | $0.28\\!\pm\\!0.01$ | $0.02\\!\pm\\!0.00$ | $0.32\\!\pm\\!0.01$ | $0.12\\!\pm\\!0.02$

Table 3: Overall matching score $\bar{\chi}_{+}$ (Eqn. 115) and irrelevant-
matching score $\bar{\chi}_{-}$ (Eqn. 116). This is an extended table of Tbl.
1. (a) When $P=1$, linear model works well regardless of the degree of over-
parameterization $\beta$, while ReLU model requires large over-
parameterization to perform well; (b) When each $R_{k}$ has multiple local
patterns that are related to the global patterns ($P>1$) related to
generators, ReLU models can capture diverse patterns better than linear ones
in the over-parameterization region $\beta>1$ and stay focus on relevant local
patterns that are related to the global patterns (i.e., low
$\bar{\chi}_{-}{}$). In contrast, linear models are much less affected by
over-parameterization. Each setting is repeated 3 times and mean/standard
derivations are reported.

Overall matching score $\bar{\chi}_{+}{}$
---
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
Linear | ReLU | Linear | ReLU | Linear | ReLU | Linear | ReLU
1 | $0.92\\!\pm\\!0.07$ | $-0.02\\!\pm\\!0.19$ | $0.60\\!\pm\\!0.04$ | $0.31\\!\pm\\!0.22$ | $0.45\\!\pm\\!0.03$ | $0.42\\!\pm\\!0.07$ | $0.38\\!\pm\\!0.01$ | $0.51\\!\pm\\!0.06$
2 | $0.96\\!\pm\\!0.01$ | $0.37\\!\pm\\!0.29$ | $0.64\\!\pm\\!0.07$ | $0.61\\!\pm\\!0.05$ | $0.48\\!\pm\\!0.02$ | $0.63\\!\pm\\!0.11$ | $0.43\\!\pm\\!0.02$ | $0.64\\!\pm\\!0.02$
5 | $0.95\\!\pm\\!0.02$ | $0.67\\!\pm\\!0.29$ | $0.68\\!\pm\\!0.04$ | $0.65\\!\pm\\!0.12$ | $0.52\\!\pm\\!0.04$ | $0.80\\!\pm\\!0.05$ | $0.48\\!\pm\\!0.01$ | $0.69\\!\pm\\!0.04$
10 | $0.96\\!\pm\\!0.01$ | $0.86\\!\pm\\!0.06$ | $0.68\\!\pm\\!0.08$ | $0.80\\!\pm\\!0.16$ | $0.52\\!\pm\\!0.01$ | $0.82\\!\pm\\!0.04$ | $0.50\\!\pm\\!0.01$ | $0.74\\!\pm\\!0.02$

Overall matching score $\bar{\chi}_{-}{}$
---
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
Linear | ReLU | Linear | ReLU | Linear | ReLU | Linear | ReLU
1 | $0.06\\!\pm\\!0.01$ | $-0.10\\!\pm\\!0.02$ | $0.12\\!\pm\\!0.02$ | $0.01\\!\pm\\!0.01$ | $0.20\\!\pm\\!0.02$ | $0.02\\!\pm\\!0.02$ | $0.24\\!\pm\\!0.02$ | $0.00\\!\pm\\!0.00$
2 | $0.08\\!\pm\\!0.00$ | $-0.04\\!\pm\\!0.01$ | $0.14\\!\pm\\!0.04$ | $0.01\\!\pm\\!0.01$ | $0.21\\!\pm\\!0.03$ | $0.00\\!\pm\\!0.01$ | $0.29\\!\pm\\!0.02$ | $0.01\\!\pm\\!0.01$
5 | $0.12\\!\pm\\!0.01$ | $-0.01\\!\pm\\!0.01$ | $0.23\\!\pm\\!0.02$ | $0.01\\!\pm\\!0.02$ | $0.31\\!\pm\\!0.02$ | $0.00\\!\pm\\!0.00$ | $0.33\\!\pm\\!0.01$ | $0.03\\!\pm\\!0.02$
10 | $0.18\\!\pm\\!0.03$ | $-0.00\\!\pm\\!0.00$ | $0.28\\!\pm\\!0.03$ | $0.00\\!\pm\\!0.00$ | $0.37\\!\pm\\!0.00$ | $0.01\\!\pm\\!0.01$ | $0.38\\!\pm\\!0.02$ | $0.05\\!\pm\\!0.02$

Table 4: Overall matching score $\bar{\chi}_{+}$ and overall irrelevant-
matching score $\bar{\chi}_{-}$ (Eqn. 15) using quadratic loss function rather
than InfoNCE. The result using InfoNCE is shown in Tbl. 1, with all
experiments setting being the same, except for the loss function. While we see
similar trends as in Sec. 5.1, quadratic loss is not as effective as InfoNCE
in feature learning. $\bar{\chi}_{+}{}$ with nonuniformity $\zeta=2$
---
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
NoBN | BN | NoBN | BN | NoBN | BN | NoBN | BN
1 | $0.35\\!\pm\\!0.23$ | $0.30\\!\pm\\!0.17$ | $0.68\\!\pm\\!0.12$ | $0.66\\!\pm\\!0.10$ | $0.79\\!\pm\\!0.02$ | $0.71\\!\pm\\!0.05$ | $0.96\\!\pm\\!0.01$ | $0.76\\!\pm\\!0.09$
2 | $0.63\\!\pm\\!0.06$ | $0.60\\!\pm\\!0.10$ | $0.89\\!\pm\\!0.05$ | $0.82\\!\pm\\!0.02$ | $0.97\\!\pm\\!0.02$ | $0.91\\!\pm\\!0.05$ | $0.99\\!\pm\\!0.01$ | $0.94\\!\pm\\!0.01$
5 | $0.80\\!\pm\\!0.00$ | $0.80\\!\pm\\!0.00$ | $0.99\\!\pm\\!0.02$ | $0.98\\!\pm\\!0.02$ | $1.00\\!\pm\\!0.00$ | $0.99\\!\pm\\!0.01$ | $0.99\\!\pm\\!0.00$ | $0.98\\!\pm\\!0.01$
10 | $0.97\\!\pm\\!0.06$ | $0.96\\!\pm\\!0.06$ | $1.00\\!\pm\\!0.00$ | $0.99\\!\pm\\!0.00$ | $1.00\\!\pm\\!0.00$ | $0.99\\!\pm\\!0.00$ | $0.99\\!\pm\\!0.00$ | $0.99\\!\pm\\!0.00$
$\bar{\chi}_{+}{}$ with nonuniformity $\zeta=5$
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
NoBN | BN | NoBN | BN | NoBN | BN | NoBN | BN
1 | $0.34\\!\pm\\!0.14$ | $0.22\\!\pm\\!0.07$ | $0.56\\!\pm\\!0.03$ | $0.36\\!\pm\\!0.02$ | $0.49\\!\pm\\!0.04$ | $0.45\\!\pm\\!0.01$ | $0.62\\!\pm\\!0.01$ | $0.63\\!\pm\\!0.08$
2 | $0.57\\!\pm\\!0.06$ | $0.33\\!\pm\\!0.12$ | $0.61\\!\pm\\!0.01$ | $0.50\\!\pm\\!0.09$ | $0.65\\!\pm\\!0.14$ | $0.68\\!\pm\\!0.04$ | $0.72\\!\pm\\!0.13$ | $0.86\\!\pm\\!0.02$
5 | $0.56\\!\pm\\!0.06$ | $0.50\\!\pm\\!0.10$ | $0.89\\!\pm\\!0.06$ | $0.90\\!\pm\\!0.04$ | $0.60\\!\pm\\!0.15$ | $0.96\\!\pm\\!0.02$ | $0.80\\!\pm\\!0.19$ | $0.97\\!\pm\\!0.00$
10 | $0.63\\!\pm\\!0.06$ | $0.56\\!\pm\\!0.06$ | $0.82\\!\pm\\!0.22$ | $0.97\\!\pm\\!0.02$ | $0.64\\!\pm\\!0.09$ | $0.98\\!\pm\\!0.01$ | $0.81\\!\pm\\!0.20$ | $0.98\\!\pm\\!0.00$
$\bar{\chi}_{+}{}$ with nonuniformity $\zeta=10$
$\beta$ | $P=1$ | $P=3$ | $P=5$ | $P=10$
NoBN | BN | NoBN | BN | NoBN | BN | NoBN | BN
1 | $0.35\\!\pm\\!0.17$ | $0.23\\!\pm\\!0.06$ | $0.51\\!\pm\\!0.02$ | $0.24\\!\pm\\!0.11$ | $0.54\\!\pm\\!0.04$ | $0.38\\!\pm\\!0.04$ | $0.54\\!\pm\\!0.03$ | $0.56\\!\pm\\!0.04$
2 | $0.39\\!\pm\\!0.18$ | $0.33\\!\pm\\!0.12$ | $0.59\\!\pm\\!0.05$ | $0.41\\!\pm\\!0.07$ | $0.58\\!\pm\\!0.07$ | $0.56\\!\pm\\!0.02$ | $0.56\\!\pm\\!0.01$ | $0.71\\!\pm\\!0.05$
5 | $0.53\\!\pm\\!0.06$ | $0.46\\!\pm\\!0.06$ | $0.66\\!\pm\\!0.06$ | $0.60\\!\pm\\!0.08$ | $0.65\\!\pm\\!0.09$ | $0.74\\!\pm\\!0.02$ | $0.60\\!\pm\\!0.01$ | $0.90\\!\pm\\!0.03$
10 | $0.56\\!\pm\\!0.16$ | $0.46\\!\pm\\!0.15$ | $0.78\\!\pm\\!0.02$ | $0.80\\!\pm\\!0.03$ | $0.70\\!\pm\\!0.04$ | $0.92\\!\pm\\!0.01$ | $0.63\\!\pm\\!0.01$ | $0.97\\!\pm\\!0.02$
Table 5: The effect of BatchNorm (BN) with ReLU activation in the presence of non-uniformity $\zeta$ of the input data. The non-uniformity is set to be $\zeta=2$ (top), $\zeta=5$ (middle) and $\zeta=10$ (bottom). For small non-uniformity, BN doesn’t help much. For larger nonuniformity, BN yields better matching score $\bar{\chi}_{+}$ in the over-parameterization region (large $\beta$) and multiple tokens per RF (large $P$), i.e., the red region. $\beta$ | $P$ |  $\mu_{k}$ no backprop |  $\mu_{k}$ backprop
---|---|---|---
$\sigma^{2}_{k}$ no backprop |  $\sigma^{2}_{k}$ backprop |  $\sigma^{2}_{k}$ no backprop |  $\sigma^{2}_{k}$ backprop
$1$ | $1$ | $0.31\pm 0.24$ | $0.23\pm 0.06$ | $0.30\pm 0.24$ | $0.23\pm 0.06$
$3$ | $0.65\pm 0.09$ | $-0.03\pm 0.01$ | $0.64\pm 0.07$ | $0.24\pm 0.11$
$5$ | $0.61\pm 0.06$ | $-0.00\pm 0.00$ | $0.62\pm 0.02$ | $0.38\pm 0.04$
$10$ | $0.66\pm 0.06$ | $0.53\pm 0.05$ | $0.70\pm 0.03$ | $0.56\pm 0.04$
$2$ | $1$ | $0.36\pm 0.13$ | $0.27\pm 0.06$ | $0.63\pm 0.11$ | $0.33\pm 0.12$
$3$ | $0.78\pm 0.01$ | $0.00\pm 0.01$ | $0.80\pm 0.02$ | $0.41\pm 0.07$
$5$ | $0.78\pm 0.02$ | $0.22\pm 0.21$ | $0.77\pm 0.07$ | $0.56\pm 0.02$
$10$ | $0.83\pm 0.08$ | $0.72\pm 0.06$ | $0.80\pm 0.04$ | $0.71\pm 0.05$
$5$ | $1$ | $0.63\pm 0.06$ | $0.43\pm 0.12$ | $0.67\pm 0.06$ | $0.46\pm 0.06$
$3$ | $0.90\pm 0.06$ | $0.45\pm 0.07$ | $0.90\pm 0.03$ | $0.60\pm 0.08$
$5$ | $0.90\pm 0.01$ | $0.72\pm 0.01$ | $0.88\pm 0.06$ | $0.74\pm 0.02$
$10$ | $0.88\pm 0.01$ | $0.88\pm 0.04$ | $0.92\pm 0.03$ | $0.90\pm 0.03$
$10$ | $1$ | $0.63\pm 0.12$ | $0.37\pm 0.15$ | $0.70\pm 0.17$ | $0.46\pm 0.15$
$3$ | $0.95\pm 0.05$ | $0.72\pm 0.06$ | $0.94\pm 0.05$ | $0.80\pm 0.03$
$5$ | $0.98\pm 0.02$ | $0.84\pm 0.05$ | $0.96\pm 0.06$ | $0.92\pm 0.01$
$10$ | $0.90\pm 0.01$ | $0.97\pm 0.02$ | $0.89\pm 0.03$ | $0.97\pm 0.02$
Table 6: The effect of backpropagating different BN statistics under
nonuniformity $\zeta=10$. Backpropagating the gradient through the sample mean
$\mu_{k}$ but not the sample variance $\sigma^{2}_{k}$ gives overall good
matching score $\bar{\chi}_{+}{}$, justifying our setting of mean-backprop
BatchNorm (Def. 2).

##  Appendix F Other lemmas

######  Lemma 5 (Bound of $1-d_{t}$).

Define

$\mu({\bm{w}}):=\frac{1+c({\bm{w}})}{2c^{2}({\bm{w}})}\left(\frac{\lambda_{2}(A({\bm{w}}(t))}{\lambda_{1}(A({\bm{w}}(t))}\right)^{2}=\frac{1+c({\bm{w}})}{2c^{2}({\bm{w}})}\left[1-\frac{\lambda_{\mathrm{gap}}(A(t))}{\lambda_{1}(A(t))}\right]^{2}\geq
0$ (117)

and $\mu_{t}:=\mu({\bm{w}}(t))$. If $c_{t}>0$ and $\lambda_{1}(t)>0$, then
$1-d_{t}\leq\mu_{t}(1-c_{t})$.

###### Proof.

We could write $d_{t}$:

$\displaystyle d_{t}$ $\displaystyle=$
$\displaystyle\frac{\phi^{\top}_{1}(t)\tilde{\bm{w}}(t+1)}{\|\tilde{\bm{w}}(t+1)\|_{2}}=\frac{\lambda_{1}(t)\phi^{\top}_{1}(t){\bm{w}}(t)}{\sqrt{\sum_{i}\lambda^{2}_{i}(t)\left(\phi^{\top}_{i}(t){\bm{w}}(t)\right)^{2}}}$
(118) $\displaystyle\geq$
$\displaystyle\frac{\lambda_{1}(t)c_{t}}{\sqrt{\lambda_{1}^{2}(t)c_{t}^{2}+\lambda_{2}^{2}(t)(1-c_{t}^{2})}}=\frac{1}{\sqrt{1+\left(\frac{\lambda_{2}(t)}{\lambda_{1}(t)}\right)^{2}\left(\frac{1}{c_{t}^{2}}-1\right)}}$
(119) $\displaystyle=$
$\displaystyle\left[1+\left(\frac{\lambda_{2}(t)}{\lambda_{1}(t)}\right)^{2}\left(\frac{1}{c_{t}^{2}}-1\right)\right]^{-1/2}$
(120) $\displaystyle\geq$ $\displaystyle
1-\frac{1}{2}\left(\frac{\lambda_{2}(t)}{\lambda_{1}(t)}\right)^{2}\left(\frac{1}{c_{t}^{2}}-1\right)=:1-\mu_{t}(1-c_{t})$
(121)

The first inequality is due to the fact that
$\sum_{i>1}\lambda^{2}_{i}(t)\left(\phi^{\top}_{i}(t){\bm{w}}(t)\right)^{2}=1-c^{2}_{t}$
(Parseval’s identity). The last inequality is due to the fact that for $x>-1$,
$(1+x)^{\alpha}\geq 1+\alpha x$ when $\alpha\geq 1$ or $\alpha<0$ (Bernoulli’s
inequality). Therefore the conclusion holds. ∎

######  Lemma 6 (Bound of weight difference).

If $c_{t}>0$ and $\lambda_{i}(t)>0$ for all $i$, then
$\|\delta{\bm{w}}(t)\|_{2}\leq\sqrt{2(1+\mu_{t}c_{t})(1-c_{t})}$

###### Proof.

First, for ${\bm{w}}^{\top}(t+1){\bm{w}}(t)$, we have (notice that
$\lambda_{i}(t)\geq 0$):

$\displaystyle{\bm{w}}^{\top}(t+1){\bm{w}}(t)$ $\displaystyle=$
$\displaystyle\frac{\sum_{i}\lambda_{i}(t)\left(\phi^{\top}_{i}(t){\bm{w}}(t)\right)^{2}}{\sqrt{\sum_{i}\lambda^{2}_{i}(t)\left(\phi^{\top}_{i}(t){\bm{w}}(t)\right)^{2}}}$
(122) $\displaystyle\geq$
$\displaystyle\frac{\lambda_{1}(t)c_{t}^{2}}{\sqrt{\lambda^{2}_{1}(t)c_{t}^{2}+\lambda^{2}_{2}(t)(1-c_{t}^{2})}}\geq\left[1-\mu_{t}(1-c_{t})\right]c_{t}$
(123)

Therefore,

$\|{\bm{w}}(t+1)-{\bm{w}}(t)\|_{2}=\sqrt{2}\sqrt{1-{\bm{w}}^{\top}(t){\bm{w}}(t+1)}\leq\sqrt{2(1+\mu_{t}c_{t})(1-c_{t})}$
(124)

∎

######  Lemma 7.

Let $\delta A=A^{\prime}-A$, then the maximal eigenvector
$\phi_{1}:=\phi_{1}(A)$ and $\phi^{\prime}_{1}:=\phi_{1}(A^{\prime})$ has the
following Taylor expansion:

$\phi^{\prime}_{1}=\phi_{1}+\delta\phi_{1}+\mathcal{O}(\|\delta A\|_{2}^{2})$
(125)

where $\lambda_{i}$ is the $i$-th eigenvalue of $A$,
$\delta\phi_{1}:=\sum_{j>1}\frac{\phi^{\top}_{j}\delta
A\phi_{1}}{\lambda_{1}-\lambda_{j}}\phi_{j}$ is the first-order term of
eigenvector perturbation. In terms of inequality, there exist $\kappa>0$ so
that:

$\|\phi^{\prime}_{1}-(\phi_{1}+\delta\phi_{1})\|_{2}\leq\kappa\|\delta
A\|_{2}^{2}$ (126)

###### Proof.

See time-independent perturbation theory in Quantum Mechanics (Fernández,
2000). ∎

######  Lemma 8.

Let $L$ be the minimal Lipschitz constant of $A$ so that
$\|A({\bm{w}}^{\prime})-A({\bm{w}})\|_{2}\leq
L\|{\bm{w}}-{\bm{w}}^{\prime}\|_{2}$ holds. If $c_{t}>0$ and
$\lambda_{i}(t)>0$ for all $i$, then we have:

$|d_{t}-c_{t+1}|=\Big{|}\left(\phi_{1}(t)-\phi_{1}(t+1)\right)^{\top}{\bm{w}}(t+1)\Big{|}\leq\nu_{t}(1-c_{t})$
(127)

where

$\nu({\bm{w}}):=2\kappa
L^{2}(1+\mu({\bm{w}})c({\bm{w}}))+2L\lambda^{-1}_{\mathrm{gap}}(A{\bm{w}}(t))\sqrt{\mu({\bm{w}})(1+\mu({\bm{w}})c({\bm{w}}))}\geq
0$ (128)

and $\nu_{t}:=\nu({\bm{w}}(t))$.

###### Proof.

Using Lemma 7 and $\|{\bm{w}}(t+1)\|_{2}=1$, we have:

$|d_{t}-c_{t+1}|\leq|\delta\phi^{\top}_{1}(t){\bm{w}}(t+1)|+\kappa
L^{2}\|\delta{\bm{w}}(t)\|_{2}^{2}$ (129)

where

$\delta\phi_{1}(t):=\sum_{j>1}\frac{\phi^{\top}(t)_{j}\delta
A(t)\phi_{1}(t)}{\lambda_{1}(t)-\lambda_{j}(t)}\phi_{j}(t)$ (130)

and $\delta A(t):=A(t+1)-A(t)$. For brevity, we omit all temporal notation if
the quantity is evaluated at iteration $t$. E.g., $\delta{\bm{w}}$ means
$\delta{\bm{w}}(t)$.

Now we bound $|\delta\phi_{1}^{\top}{\bm{w}}(t+1)|$. Using Cauchy–Schwarz
inequality:

$\displaystyle|\delta\phi_{1}^{\top}{\bm{w}}(t+1)|$ $\displaystyle=$
$\displaystyle\Bigg{|}\sum_{j>1}\left(\frac{\phi^{\top}_{j}\delta
A\phi_{1}}{\lambda_{1}-\lambda_{j}}\right)\left(\phi^{\top}_{j}{\bm{w}}(t+1)\right)\Bigg{|}$
(131) $\displaystyle\leq$
$\displaystyle\sqrt{\sum_{j>1}\left(\frac{\phi^{\top}_{j}\delta
A\phi_{1}}{\lambda_{1}-\lambda_{j}}\right)^{2}}\sqrt{\sum_{j>1}\left(\phi^{\top}_{j}{\bm{w}}(t+1)\right)^{2}}$
(132) $\displaystyle\leq$
$\displaystyle\frac{1}{\lambda_{\mathrm{gap}}(A)}\sqrt{\sum_{j>1}\left(\phi^{\top}_{j}\delta
A\phi_{1}\right)^{2}}\sqrt{\sum_{j>1}\left(\phi^{\top}_{j}{\bm{w}}(t+1)\right)^{2}}$
(133)

Since $\\{\phi_{j}\\}$ is a set of orthonormal bases, Parseval’s identity
tells that for any vector ${\bm{v}}$, its energy under any orthonormal bases
are preserved: $\sum_{j}(\phi^{\top}_{j}{\bm{v}})^{2}=\|{\bm{v}}\|_{2}^{2}$.
Therefore, we have:

$\displaystyle|\delta\phi_{1}^{\top}{\bm{w}}(t+1)|$ $\displaystyle\leq$
$\displaystyle\frac{1}{\lambda_{\mathrm{gap}}(A)}\|\delta
A\phi_{1}\|_{2}\sqrt{1-d^{2}_{t}}$ (134) $\displaystyle\leq$
$\displaystyle\frac{L}{\lambda_{\mathrm{gap}}(A)}\|\delta{\bm{w}}(t)\|_{2}\sqrt{1-d^{2}_{t}}$
(135)

Note that using $-1\leq d_{t}\leq 1$ and Lemma 5, we have:

$\sqrt{1-d_{t}^{2}}=\sqrt{1+d_{t}}\sqrt{1-d_{t}}\leq\sqrt{2(1-d_{t})}\leq\sqrt{2\mu_{t}(1-c_{t})}$
(136)

Finally using bound of weight difference (Lemma 6), we have:

$\displaystyle|d_{t}-c_{t+1}|$ $\displaystyle\leq$ $\displaystyle 2\kappa
L^{2}(1+\mu_{t}c_{t})(1-c_{t})+L\lambda_{\mathrm{gap}}^{-1}\sqrt{2(1+\mu_{t}c_{t})(1-c_{t})}\sqrt{1-d_{t}^{2}}$
(137) $\displaystyle\leq$ $\displaystyle\nu_{t}(1-c_{t})$ (138)

Here $\nu_{t}:=2\kappa
L^{2}(1+\mu_{t}c_{t})+2L\lambda^{-1}_{\mathrm{gap}}(A(t))\sqrt{\mu_{t}(1+\mu_{t}c_{t})}$.
∎

######  Lemma 9.

Let $c_{0}:=c({\bm{w}}(0))={\bm{w}}^{\top}(0)\phi_{1}(A({\bm{w}}(0)))>0$.
Define local region $B_{\gamma}$:

$B_{\gamma}:=\left\\{{\bm{w}}:\|{\bm{w}}-{\bm{w}}(0)\|_{2}\leq\frac{\sqrt{2(1+\gamma)(1-c_{0})}}{1-\sqrt{\gamma}}\right\\}$
(139)

if there exists $\gamma<1$ so that

$\sup_{{\bm{w}}\in B_{\gamma}}\mu({\bm{w}})+\nu({\bm{w}})\leq\gamma,$ (140)

then

  * •

The sequence $\\{c_{t}\\}$ increases monotonously and converges to $1$;

  * •

There exists ${\bm{w}}_{*}$ so that
$\lim_{t\rightarrow+\infty}{\bm{w}}(t)={\bm{w}}_{*}$.

  * •

${\bm{w}}_{*}$ is the maximal eigenvector of $A({\bm{w}}_{*})$ and thus a
fixed point of gradient update (Eqn. 6);

  * •

For any $t$,
$\|{\bm{w}}(t)-{\bm{w}}(0)\|_{2}\leq\frac{\sqrt{2(1+\gamma)(1-c_{0})}}{1-\sqrt{\gamma}}$.

  * •

$\|{\bm{w}}_{*}-{\bm{w}}(0)\|_{2}\leq\frac{\sqrt{2(1+\gamma)(1-c_{0})}}{1-\sqrt{\gamma}}$.
That is, ${\bm{w}}_{*}$ is in the vicinity of the initial weight
${\bm{w}}(0)$.

###### Proof.

We first prove by induction that the following _induction arguments_ are true
for any $t$:

  * •

$c_{t+1}\geq c_{t}>0$;

  * •

$1-c_{t}\leq\gamma^{t}(1-c_{0})$;

  * •

${\bm{w}}(t)$ is not far away from its initial value ${\bm{w}}(0)$:

$\|{\bm{w}}(t)-{\bm{w}}(0)\|_{2}\leq\sqrt{2(1+\gamma)(1-c_{0})}\sum_{t^{\prime}=0}^{t-1}\gamma^{t^{\prime}/2}$
(141)

which suggests that ${\bm{w}}(t)\in B_{\gamma}$.

Base case ($t=1$). Since $1\geq c_{0}>0$, $\mu({\bm{w}})\geq 0$, and
$A({\bm{w}})$ is PD, applying Lemma 6 to $\|{\bm{w}}(1)-{\bm{w}}(0)\|_{2}$, it
is clear that

$\|{\bm{w}}(1)-{\bm{w}}(0)\|_{2}=\|\delta{\bm{w}}(0)\|_{2}\leq\sqrt{2}\sqrt{(1+\mu_{0}c_{0})(1-c_{0})}\leq\sqrt{2(1+\gamma)(1-c_{0})}$
(142)

Note that the last inequality is due to $\mu_{0}\leq\gamma$. Note that

$1-c_{1}=1-d_{0}+d_{0}-c_{1}\leq
1-d_{t}+|d_{0}-c_{1}|\leq(\mu_{0}+\nu_{0})(1-c_{0})\leq\gamma(1-c_{0})$ (143)

and finally we have $c_{1}\geq 1-\gamma(1-c_{0})\geq c_{0}>0$. So the base
case is satisfied.

Inductive step. Assume for $t$, the induction argument is true and thus
${\bm{w}}(t)\in B_{\gamma}$. Therefore, by the condition, we know
$\mu_{t}+\nu_{t}\leq\gamma$.

By Lemma 6, we know that

$\|{\bm{w}}(t+1)-{\bm{w}}(t)\|_{2}=\|\delta{\bm{w}}(t)\|_{2}\leq\sqrt{2(1+\mu_{t}c_{t})(1-c_{t})}\leq\sqrt{2(1+\gamma)(1-c_{0})}\gamma^{t/2}$
(144)

Therefore, we know that ${\bm{w}}(t+1)$ also satisfies Eqn. 141:

$\displaystyle\|{\bm{w}}(t+1)-{\bm{w}}(0)\|_{2}$ $\displaystyle\leq$
$\displaystyle\|{\bm{w}}(t)-{\bm{w}}(0)\|_{2}+\|\delta{\bm{w}}(t)\|_{2}$ (145)
$\displaystyle\leq$
$\displaystyle\sqrt{2(1+\gamma)(1-c_{0})}\left[\sum_{t^{\prime}=0}^{t-1}\gamma^{t^{\prime}/2}+\gamma^{t/2}\right]$
(146) $\displaystyle=$
$\displaystyle\sqrt{2(1+\gamma)(1-c_{0})}\sum_{t^{\prime}=0}^{t}\gamma^{t^{\prime}/2}$
(147)

Also we have:

$\displaystyle 1-c_{t+1}$ $\displaystyle=$ $\displaystyle
1-d_{t}+d_{t}-c_{t+1}\leq 1-d_{t}+|d_{t}-c_{t+1}|$ (148) $\displaystyle\leq$
$\displaystyle(\mu_{t}+\nu_{t})(1-c_{t})\leq\gamma(1-c_{t})$ (149)
$\displaystyle\leq$ $\displaystyle\gamma^{t+1}(1-c_{0})$ (150)

and thus we have $c_{t+1}\geq 1-\gamma(1-c_{t})\geq c_{t}>0$.

Therefore, we have

$1-c_{t}\leq\gamma^{t}(1-c_{0})\rightarrow 0$ (151)

thus $c_{t}$ is monotonously increasing to $1$. This means that:

$\lim_{t\rightarrow+\infty}c_{t}=\lim_{t\rightarrow+\infty}\phi_{1}^{\top}(t){\bm{w}}(t)\rightarrow
1$ (152)

Therefore, we can show that ${\bm{w}}(t)$ is also convergent, by checking how
fast $\|\delta{\bm{w}}(t)\|_{2}$ decays:

$\|\delta{\bm{w}}(t)\|_{2}\leq\sqrt{2(1+\mu_{t}c_{t})(1-c_{t})}\leq\sqrt{2(1+\gamma)(1-c_{0})}\gamma^{t/2}$
(153)

By Cauchy’s convergence test,
${\bm{w}}(t)={\bm{w}}(0)+\sum_{t^{\prime}=0}^{t-1}\delta{\bm{w}}(t^{\prime})$
also converges. Let

$\lim_{t\rightarrow+\infty}{\bm{w}}(t)={\bm{w}}_{*}$ (154)

This means that $A({\bm{w}}_{*}){\bm{w}}_{*}=\lambda_{*}{\bm{w}}_{*}$ and thus
$P_{{\bm{w}}_{*}}^{\perp}A({\bm{w}}_{*}){\bm{w}}_{*}=0$, i.e., ${\bm{w}}_{*}$
is a fixed point of gradient update (Eqn. 6). Finally, we have:

$\|{\bm{w}}(t)-{\bm{w}}(0)\|_{2}\leq\sqrt{2(1+\gamma)(1-c_{0})}\sum_{t^{\prime}=0}^{t-1}\gamma^{t^{\prime}/2}\leq\frac{\sqrt{2(1+\gamma)(1-c_{0})}}{1-\sqrt{\gamma}}$
(155)

Since $\|\cdot\|_{2}$ is continuous, we have the conclusion. ∎