Feature Collapse
Authors: Thomas Laurent
James von Brecht
Xavier Bresson
Venue: ICLR 2024 poster (2024)
Keywords: deep learning theory
representation learning
neural collapse
Primary area: learning theory
TL;DR: We prove that a network trained on a synthetic data model learns interpretable and meaningful representations in its first layer.

Abstract: We formalize and study a phenomenon called *feature collapse* that makes precise the intuitive idea that entities playing a similar role in a learning task receive similar representations. As feature collapse requires a notion of task, we leverage a synthetic task in which a learner must classify  `sentences' constituted of $L$ tokens. We start by showing experimentally that feature collapse goes hand in hand with generalization. We then prove that, in the large sample limit,  distinct tokens that play identical roles in the task receive identical local feature representations in the first layer of the network. This analysis shows that a neural network trained on this task provably learns interpretable and meaningful representations in its first layer.

Published as a conference paper at ICLR 2024
FEATURE COLLAPSE
Thomas Laurent', James H. von Brecht, Xavier Bresson Loyoll Marymount University, tlaurent 1mu edu
2 National University of Singapore, xaviercs edu. ag
ABSTRACT
We formalize and study a phenomenon called feature collapse that makes precise the intuitive idea that entities playing similar role learning task receive sim- ilar representations. As feature collapse requires a notion of task, we leverage a synthetic task which a learner must classify 'sentences' constituted of L. tokens. We start by showing experimentally that feature collapse goes hand in hand with generalization. We then prove that, in the large sample limit, distinct tokens that play identical roles in the task receive identical local feature in the first layer of the network. This analysis shows that a neural network trained on this task provably learns interpretable and meaningful representations in its first layer.
INTRODUCTION
Many machine learning practices implicitly rely the belief that good generalization and transfer learning requires good features. Despite this, the notion of "good features' remains vague and carries many potential meanings. One definition is that fentures/reprreenntatioss should only encode the information necessary to do the task at hand, and discard any unnecessary information as noise. For example, two distinct patches of grass should map essentially identical representations even if these patches differ in pixel space. Intuitively, a network that gives the same representation to many distinct patches of grass has learned the 'grass' concept. We call this phenomenon feature collapse, meaning that a learner gives the same feature to entities that play similar roles for the task at hand. This phenomenon captures the common intuition, often confirmed by practice, that the ability to learn good representations in early layers is essential for the empirical success of neural networks.
Broadly speaking, we conduct theoretical investigation into how such representations are learned in the early layers of neural network. To make progress we adopt a common approach in theoretical deep-learning community. One starts with a synthetic data model exhibiting clear latent structure, and then prove that specific neural network architecture successfully uncovers this latent structure during training. For example, recent work in representation learning (Damian et al. [022 Mousavi- Hosseini al. 2022) leverages data model (the multiple-index model) with latent structure defined by low dimensional subspace. In this specific setting, the first layer of fully connected neural network provably learns this low-dimensional subspace.
Our work follows along the same lines. Specifically, we start with a data model that asks a learner to classify 'sentences' comprised of L tokens. Some of these tokens play identical roles in the sense that replacing one with another does not change the label of the sentence, and this equivalence of tokens defines the latent structure in the data. We then consider neural network containing shared embedding module, optionally followed by a LayerNorm module, and then a final linear classifica- tion head. We show that this network, when equipped with the LayerNorm module, provably learn the equivalence of these tokens in its first layer. To do so, we consider the large sample limit and derive analytical formulas for the weights of the network trained on the data model. Under some symmetry assumptions the task, these analytical formulas reveal:
(i) If the network includes LayerNorm then feature collapse takes place. The neural network learns to give the same embedding to tokens that play identical roles.
(ii) the network does not include LayerNorm then feature collapse fails. The network does not give the same embedding tokens that play identical roles. Moreover, this failure stems from the fact that common tokens receive large embeddings and rare tokens receive small embeddings.

---

Published as a conference paper at ICLR 2024
Finally, conduct experiments that show feature collapse and generalization go hand in hand. These experiments demonstrate that for the network to generalize well it is essential for tokens playing identical roles to receive the same embeddings.
In summary, our main contributions are as follow:
We study how network learns representations in its first layer. To do so, we make the notion of" "good features' mathematically rigorous via synthetic data model with latent structure.
We derive analytical formulas for the weights of a two-layer network trained on this data model. These analytical formulas show that, when equipped with a LayerNorm module, the network provably learns interpretable and meaningful representations in its first layer.
The remainder of the paper proceeds as follows: In subsection we discuss related works; In section 2 we describe the data model and present a set of visual experiments to illustrate the main ideas used in the paper; In section 3 we present the three theorems that constitute the main results of the paper; Finally, section outlines a set of additional experiments performed in the appendix.
1.1 RELATED WORKS
Our work most closely resembles recent work theoretical representation learning and, some extent, the recent literature on neural collapse. The works by Damian et al. (2022) and Mousavi- Hosseini al. (2022) consider synthetic data model, the multiple-index model, to investigate representation learning. In this model, the learner must solve a regression task with normally dis- tributed inputs E IR and targets y = g((u) x)...., (u,, for some function g R" - R, some vectors U1...., U, E Rd, and with & d. The target y therefore solely depends on the projection of x on the low-dimensional subspace spanned by u1,..., u, These works prove that fully connected, two-layer neural network learns, in its first layer, the low dimensional subspace. The behavior of fully connected networks trained on the multiple-indes model, or on related data models, has also been studied in various other works, including Ba et al. (2022). Bietti et al. (2022). Abbe et al. (2022) and Parkinson et al. (203)). Both this line of investigation and our work prove, in the appropriate sense, that network uncovers latent structure. Nevertheless, our data model differs quite substantially from the multi-index model. Moreover, we do not study a fully connected net- work. Our network has shared embedding module, and this allows to tie good features to notion of semantic equivalence.
The phenomenon we study in this but in detail is quite different. In a pioneering work, Papyan et al. (2020) conducted a series of experiments that revealed that a well-trained network gives identical representations, in its last layer, to training points that belong to the same class. In K-class classification task we therefore see the emergence, in the last layer, of K vectors coding for the К classes. Additionally, these K vectors 'point' in 'maximally opposed directions. This phenomenon, coined neural collapse, has been studied extensively since its discovery (e-g. Mixon et al. ( (2020); Lu Steinerberger (2020); Wojtowytsch et al. (2020); Fang al. (2211); Zhu et al. (2021); Ji et al. (2021); Tirer & Bruna (2022); Zhou et al. (2022) ). To emphasize the difference with feature collapse, note that neural collapse refers to phenomenon where all training points from the same class receive the same representation at the end of the network. Unlike feature collapse, this does not provide any indication that the network has learned good representations in its early layers, or that the neural network has uncovered some latent structure beyond the one encoded in the training labels.
collapse, has superficial similarity to neural collapse
2 A TALE OF FEATURE COLLAPSE
We begin by more fully telling the empirical tale that motivates our theoretical investigation. To make progress we adopt common approach in theoretical deep-learning and leverage synthetic data model exhibiting a clear latent structure. The model generates sequences of length L from some underlying set of latent variables that encode the K classes of classification task. To make it concrete we describe this generative process using NLP terminology like 'sentences' and 'words', but of course do intend it realistic language model. Figure [ illustraate the basic idea. The left side of the figure depicts a vocabulary of n = 12 word tokens and ne 3 concept tokens
V (oottoo, cheese, carrots, chicken....
and
C = (vegetable, dairy, meat)
2

---

Published as a conference paper at ICLR 2024
12 wordh partitioned Into concepts
varinlss
Sentences
pork
lodice
fratle
chicker
lettens
veggie
pages
demta
batter,
lectoos
bettery
seal, daoy daky. vage, 14AM
Xi
Chilken
pock
Inaioe
cheon
perk
Lea
dairy
ment dtrtuutte
Edie dortuctinn
Au
perk
Hezk, diyy rest, daty
diviee
purk
Figure 1: Data model with parameters set to L = 5, fw = 12, ne = 3, K=
= 3.
with the 12 words partitioned into the 3 equally sized concepts. sentence x E V is a sequence of L words (L 5 on the figure), and latent variable C. is sequence of L concepts. The latent variables generate sentences. For example
z == [dairy, veggie, meat, veggie, dairy] generates
X =
[cheese carrot, pork, potato, butter]
with the sentence on the right obtained by sampling each word at random from the corresponding concept. The first word represents random sample from the dairy concept (butter, cheese, cream, yogurt) according to the dairy distribution (square box at left), the second word represents a random sample from the vegetable concept (potato, carrol, leek, lefluce) according to the vegetable distribu- tion, and so forth. At right, figure [ depicts a classification task with K = categories prescribed by the three latent variables Z1, Z2, Z3 C²,, Sentences generated by the latent variable z share the same label k, yielding a classification problem that requires a learner to classify sentences among K categories.
We use two similar networks to empirically study if and when the feature collapse phenomenon occurs. The first network x +> hw, u (x), depicted on the top panel of figure 2 starts by embedding each word in a sentence by applying x The matrix W to the one-hot representation of each word. then concatenates these d-dimensional embeddings of each word into a single vector. Finally, it applies linear transformation U to produce a K- dimensional score vector = hw; u (x) with one entry for each of the K classes. The d x nw embedding matrix W and the K x Ld matrix U of linear weights are the only learnable parameters, and the network has no nonlinearities. The second network x + hw, u (x). depicted at bottom, differs only by the application of a LayerNorm module to the word em- beddings prior to the concatenation. For simplicity we use a LayerNorm module which does not contain any learnable parameters; the module sim- ply removes the mean and divides by the standard deviation of its input vector. As for the first network, the only learnable weights are W and U.
DAp
Iaterr
Iok
The task depicted on figure [ and the networks depicted on figure 2 pro- vide clear way of studying how interpretable and meaningful representa- tions are learned in the first layer of a network. For example, the four words bulter, cheese, cream and yogurt clearly play identical role for the task at hand (replacing one with another does not change the label of the sentence). As consequence we would expect the embedding layer of the network to map them to the same representation. Similarly, the words belonging to the vegetable concepts should receive same representation, and the words belonging to the meat concept should receive same representation. If this takes place, we say that feature collapse has occurred.
Figure 2: Networks
If feature collapse occurs, this will also be reflected in the second layer of the network. To see this, partition the linear transformation
U1.1
U1,2
u1,2
U2
-1222
U2,1
U-
(1)
uK,1
uK22
3

---

Published as a conference paper at ICLR 2024
into its components Ug,e E R. Suppose for example that 2,,6 veggie, meaning that the latent variable Zy contains the veggie concept in the 7th position. If W properly encodes concepts then we expect the vector u to give strong response when presented with the embedding of word that belongs to the veggie concept. So we would expect Ug,e to align with the embeddings of the words that belong to the veggie concept, and feature collapse would occur in this manner as well.
In the remainder of this section we conduct experiments that visually illustrate the feature collapse phenomenon and the formation of interpretable representations in networks h and ht. These exper- iments also show how layer normalization plays key role in the feature collapse phenomenon. In our experiments we use the standard cross entropy loss fyy, k), with y R2 and I k I K, and then minimize the corresponding regularized empirical risks
K Repl
Remp (W,U)
1
1
ΣΣ
A (hwwo (xi,s) k)) + U 2
K Thapl
1
K
Papl
A |
Ro (W,U)
=
(hiww (xx,)) , k
(3)
K Rapl
of each network via stochastic gradient descent. The X denote the i-th sentence of the k- th category the training set, and so each of the K categories has Nspl representatives. For the parameters of the architecture, loss, and training procedure, we use an embedding dimension of = 100, a weight decay of À = 0.001, a mini-batch size of 100 and a constant learning rate 0.1, respectively, for all experiments. The LayerNorm module implicitly regularizes the ma- trix W so we do not penalize it in equation (3). The codes for our experiments are available at /ttps://qiite com. xbressson/ feature collapse
Remark: Without weight decay (i.e. = 0), the above objectives typically do not have global minima. We therefore focus our theoretical investigation the case À >>> which is analytically more tractable. In appendix A we provide an empirical investigation of the case without weight decay to show that both cases (i.e. À >> and X 0) exhibit the same behavior in practice.
2.1 THE UNIFORM CASE
We start with an instance of the task from figure D with parameters n. 3, n., 1200, == 15, == 1000, and with uniform word distributions. So each of the 3 concepts (say veggie, duiry, and meal) contain 400 words and the corresponding distributions (the veggie distribution, the dairy dis- tribution, and the meat distribution) are uniform. We form K = 1000 latent variables Zi,... ... Z1000 by selecting them uniformly at random from the set C²,, which simply means that any concept se- quence [21, ...., =z] has equal probability of occurrence. We then construct training set by generating Nspll = 5 data points from each latent variable. We then train both networks h, h" and evaluate their generalization performance; both achieve 100% accuracy on test points.
Since both networks generalize perfectly, we expect them to have learned good representations. To confirm this, start by visualizing in figure learnable parameters W, U of the network lw: after training. The embedding matrix W contains Th = 1200 columns. Each column is a vector in 100 and cor- responds to word embedding. The top panel of figure 3 depicts these 1200 word embeddings after dimensionality reduction via PCA. The top singular values 0 34.9, 02 34.7 and 03 = 0.001 associated with the PCA indicate that the word embeddings essentially live in a dimensional subspace of RR 100 and the PCA paints an accurate picture of the distribution of word embed- dings. We then color code each word embedding accorded to its concept, so that all embeddings of words within a concept receive the same color (say all veggie words in green, all dairy words in blue, and so forth). As the figure illus- trates, words from the same concept receive nearly identical embeddings, and these embeddings form equilateral triangle or two-dimensional simplex. We therefore observe collapse of features into a set of n. = 3 equi-angular vectors at the level word embeddings. The bottom panel of figure illlstrates collapse for the parameters U of the linear layer. We partition the matrix U into vectors Ug, E R 100 via [1] and visualize them once again with PCA. As for the word embeddings, the singular values of the PCA (01 34.9, d2 34.6 and Ø3 0.0003) reveal
Figure 3: W & U
4

---

Published as a conference paper at ICLR 2024
that the vectors Ug,6 essentially live in two dimensional subspace of IR 100, We color code each u according to the concepts contained in the corresponding latent variable (say Ug, is green Z,,e veggie, and so forth). The figure indicates that vectors Ux,e that correspond cept collapse around a single vector. A similar analysis applied to the weights of the network liyu tells the same story, provided examine the actual word features (i.e. the embeddings after the LayerNorm) rather than the weights W themselves.
same
In theorem and (see section 3 we prove the correctness of this empirical picture. We show that the weights of / and ht collapse into the configurations illustrated on figure 3 in the large sample limit. Moreover, this limit captures the empirical solution very well. For example, the word embeddings in figure 3 have norm equal to 1.41 + 0.13, while we predict a nomn of 1.42214 theoretically.
2.2 THE LONG-TAILED CASE
At superficial glance it appears if the LayerNorm module plays essential role, both net- works h and h", in the previous experiment, exhibit feature collapse and generalize perfectly. To probe this issue further,
continue investigation by conducting similar experiment (keeping n. = 3, Thp = 1200, I = 15, and = 1000) but with non-uniform, long-tailed word distributions within each of the ne concepts. For concreteness, say the veggie concept contains the 400 words
potato, lettuce,
arugula, parsnip,
achojcha
where achojcha a rare vegetable that grows in the Andes mountains. We form the veggie distribu- tion by sampling potato with probability C/L, sampling lettuce with probability C/2, and so forth down to achojcha that has probability C/400 of being sampled (C' is chosen so that all the probabil- ities sum to 1). This "1/i" power law distribution has a long-tail, meaning that relatively infrequent words such as arugula or parsnip collectively capture a significant portion of the mass. Natural data in form of text or images typically exhibit long-tailed distributions (Salakhutdinov et al.) 2011 Zhu et al., 2014 Liu et al. 2019; Feldman, 2020; Feldman & Zhang 2020). For instance, the frequencies of words in natural text approximately conform to the "1/1" power law distribution, also known as Zipf's s law (Zipí] 1935). which motivates the specific choice made in this experiment. Many datasets of interest display some form of long-tail behavior, whether at the level object oc- currences in computer vision the frequency of words topies in NLP, and effectively addressing these long-tail behaviors is frequently a challenge for the learner.
To investigate the impact of variables ,,,..., Z1000 uniformly at random as before. We then use them to build two distinct training
word distributions, first randomly select the latent
We build large training set by generating Nspl 500 training points per latent variable and a small training set by generating napl = 5 training points per latent variable. We use the "1/i" power law distribution when sampling words from concepts in both cases. We then train h and h² on both training sets and evaluate their generalization performance. When trained on the large training set, both are 100% accurate test time (as they should be the large training set has 500, 000 total samples). A significant difference emerges between h and h" when trained on the small training set. The network achieves test accuracy of 45% while h² remains 100% accurate.
We once again visualize the weights of each network to study the relationship between generaliza- tion and collapse. Figure aa) depicts the weights of lww (via dimensionality reduction and color coding) after training the large training set. The word embeddings the linear weights uk, on the right sub-panel. Words that belong to the same concept still receive
the left sub-panel and
(a) h trained on the large training set. Test acc. 100%
(b) h trained on the small training set. Test acc. 43%
(c) h" trained on the small training set. Test acc. 100%
Figure 4: Visualization of matrices W (left in each subfigure) and U (right in each subfigure)
5

---

Published as a conference paper at ICLR 2024
embeddings that are aligned, however, the magnitude of these embeddings depends upon word fre- quency. The most frequent words in a concept (e.g. potato) have the largest embeddings while the least frequent words (e.g. achojcha) have the smallest embeddings. In other words, observe 'directional collapse' of the embeddings, but the magnitudes do not collapse. In contrast, the linear weights U mostly concentrate around three well-defined, equi-angular locations; they collapse in both direction and magnitude. A major contribution of our work (c.f. theorem 2 in the next section) theoretical insight that explains the configurations observed in figure a), and in particular, explains why the magnitudes of word embeddings depend on their frequencies.
Figure Tb) illustrates the weights of hw.u after training on the small training set. While the word embeddings exhibit similar pattern as in figure A a). the linear weights us remain dispersed and fail to collapse. This leads to poor generalization performance (45% accuracy at test time).
To summarize, when the training set is large, the linear weights ug,e collapse correctly and the network hw.u generalizes well. When the training set is small the linear weights fail to collapse, and the network fails to generalize. This phenomenon can be attributed to the long-tailed nature of the word distribution. To see this, say that 2 = veggie, dairy, veggie, the hth latent variable for the sake of concreteness. With only Mspl 5 samples for this latent variable, we might end up in a situation where the 5 words selected to represent the first occurrence of the veggie concept have very different frequencies than the five words selected to represent the third occurrence of the veggie concept. Since word embeddings have magnitudes that depend on their frequencies, this will result in serious imbalance between the vectors ug,1 and ug,3 that code for the first and third occurrence of the veggie concept. This leads to two vectors Ug,1, 3 that code for the same concept but have different magnitudes (as seen on figure (b)) so features do not properly collapse. This imbalance results from the 'noise" introduced by sampling only 5 training points per latent variable. Indeed, if napl = 500 then each occurrence of the veggie concept will exhibit a similar mix of frequent and rare words, up. 1 and 3 will have roughly same magnitude, and full collapse will take place (c.f. figure a)). Finally, the poor generalization ability of hw.u when the training set is small really stems from the long-tailed nature of the word distribution. The failure mechanism occurs due to the relatively balanced mix of rare and frequent words that occurs with long-tailed data. If the data were dominated by few very frequent words, then all rare words combined would just contribute small perturbations and would not adversely affect performance.
meat, dairy] represents
We conclude this section by examining the weights of the network hiv. after training on the small training set. The left panel of figure 4cc) provides a visualization of the word embeddings after the LayerNorm module. These word representations collapse both in direction and magnitue;; they do not depend on word frequency since the LayerNorm forces vectors to have identical magnitude. The right panel of fig depicts the linear weights Ug,e and shows that they properly collapse. As consequence, hww, generalizes perfectly (100% accurate) even with only Rapl = 5 sample per class. Normalization plays crucial role by ensuring that word representations do not depend upon word frequency. In turn, this prevents the undesired mechanism that causes hw.y to have uncollapsed linear weights u when trained on the small training set. Theorem in the next section proves the correctness of this picture. The weights of the network h* collapse to the 'frequency independent' configuration of figure (c) in the large sample limit.
THEORY
Our main contributions consist in three theorems. In theorem we prove that the weights of the network fw.u collapse into the configurations depicted on figure when words have identical fre- quencies. In theorem 2 we provide theoretical justification of the fact that, when words have distinct frequencies, the word embeddings of hw.u must depend on frequency in the manner that figure 4a) illustrates. Finally, in theorem 3 we show that the weights of the network hww exhibit full collapse even when words have distinct Trequencies. Each of these theorems hold in the large Napt limit and under some symmetry assumptions on the latent variables. All proofs are in the appendix.
Notation.
The set of concepts, which up to now was = (veggie, dairy, meat), will be represented in this section by the more abstract C = .nc).. We let Se nw/ne denote the number of words per concept, and represent the vocabulary by
V { (0,1)) N2 1 1 a na and 1 S 3 S 8c

---

Published as a conference paper at ICLR 2024
So elements of V are tuples of the form (a, 9) with 1 I I n. and 1 $c, and we think of the tuple (ax, 5) as representing the 3th word of the at concept. Each concept a E C comes equipped [0, 1] over the words within it, so that (3) is the probability of selecting the 3th word when sampling out of the o concept. For simplicity we assume that the word distributions within each concept follow identical laws, so that
with a probability distribution Pa {1, {I,..., se)
(3(3) = uB for all (a,8)E) EV
for some positive scalars H3 > O that sum to 1. We think of ug as being the 'frequency" of word (a, 9) in the vocabulary. For example, choosing Ps 1/8. gives uniform word distribu- tions while pa x 1/8 corresponds to Zipf"s law. We use X' == V denote the data space and 2 == c² to denote the latent space. The elements of the data space X' correspond to sequences x = [[ Bi)...., (oz,
Sr.]] of L words, while elements of the latent space Z correspond to se- quences [11 ar. of L concepts. For given latent variable we write x ~ DD to indicate that the data point x was generated by z (formally D2 : X + [0, 1] is a probability distribution).
Word embeddings, LayerNorm, and word representations. We ((..3) Rd to denote the embedding of word (a, (3) V. The collection of all ((..)) determines the columns of the matrix W E Rdxa. These embeddings feed into a LayerNorm module without learnable parameters:
sp(v) me olv)
where mean(v) (v) = ヒニΣ os and o²(v) = d (1; (c--mean(v))+
So the LayerNorm module converts word embedding W(,, B) E Rdd into vector (( 3)) E RRd, and we call this vector a word representation.
Equiangular vectors. We call collection of n. vectors f Inn E Rd equiangular if
ne
and faa) {- 11/(ne otherwise
i = a'
Σ
fa 0
(4)
hold for all possible pairs a, a' E [ne] of concepts. For example, three vectors f2, fa E RR 100 are equiangular exactly when they have unit norms, live in two dimensional subspace of 10000 and form the vertices of an equilateral triangle in this subspace. This example exactly corresponds to the configurations in figure 3 and (up scaling factor). Similarly, four vectors fI, f2, fa f4 E R R0000 equiangular when they have unit norms and form the vertices of a regular tetrahedron. We will sometimes require f....., fn. E Rd to also satisfy (fa, la) = for all a E [ne]. in which case we say f1, fi...., fn, E Rdd form collection of mean- zero equiangular vectors.
Collapse configurations. Our empirical investigations reveal two distinct candidate solutions for the weights (W, U) of the network hw, U and h u- We therefore isolate each of these possible candidates as a definition before turning to the statements of our main theorems. We begin by defin- ing the type of collapse observed when training the network hwę u with uniform word distributions (see figure 3 for visual illustration of this type of collapse).
Definition 1 (Type-I Collapse). The weights (W,U) of the network tw.v form a type-I collapse configuration if and only if the conditions
There exists e 2 so that W(..3) = cfa for all (a, B) E V.
ii) There exists c' 2 0 so that Ug,t = d' fa for all (k, E) satisfying z,,t and all a E C.
hold for some collection fi.... In. E IRd of equiangular vectors.
Recall that the network hw.y exhibits collapse well, up the fact that the word representations p(( collapse rather than the word embeddings themselves. Additionally, the LayerNorm also fixes the magnitude of the word representations. We isolate these differences in the next definition.
Definition 2 (Type-Il Collapse). The weights (W, U) of the network liw. form a type-ll collapse configuration if and only if the conditions
i) (()() = Vd for all (a. E V.
ii) There exists e 2 so that ug,e cfa for all (k, satisfying zk,e = a and all C. hold for some collection fi.... fn. E IRd of mean-zero equiangular vectors.
7

---

Published as a conference paper at ICLR 2024
Finally, when training the network kw, with non-uniform word distributions (c.f. figure 4a) observe collapse in the direction of the word embeddings ((a,)) 8) but their magnitudes depend upon word frequency. We therefore isolate this final observation as Definition 3 (Type-III - Collapse). The weights (W, U) of the network fw.c form rype-lll collapse configuration if and only if
i) There exists positive scalars r3 2 0 so that (a. = rrgfa
for all (a, E V.
ii) There exists e 2 that u,,e" cfa for all (k, E) satisfying 2 a and all C.
hold for some collection fi...., Inn E d of equiangular vectors.
In a type-Ill collapse we allow the word embedding ((a, 8) to have frequency-depence magni- tude rg while type-I collapse we force all embeddings have the same magnitude; this makes type-I collapse a special case of type-Ill collapse, but not vice-versa.
3.1 PROVING COLLAPSE
Our first result proves that the words embeddings W(a and linear weights us. exhibit type-I collapse in appropriate large-sample limit. When turning from experiment (c.f. figure 3 to theory we study the true risk
K
1 K
Σ DaD (x)) + À (IW11 + ((U))E)
R(W,U) U)
(5)
rather than the empirical risk Remp UW and place a symmmetry assumption on the latent variables. Assumption (Latent Symmetry). For every k E [K], r E [L]. & e [L]. and [ne] the identities
(ne - 1)" i/ - 1)r-1
R [K] dist(Zz, zx) = T and ZN:E a
(6)
if a
hold, with dist (zs, zx) denoting the Hamming distance between a pair (Zy: Ze ) of latent variables.
With this assumption in hand may state first main result
Theorem 1 (Full Collapse of h). Assume uniform sampling μĄ = 1/sc for each word dis- tribution. Let T 2 O denote the unique minimizer of the strictly convex function H
log 1 1-
- + mhe (
1)e-n) (me )))=")) + At where 7j
Assume Z1, .... ZK
KL are mutually distinct and satisfy the symmetry assumption Then any (W,U) in type-I collapse configuration with constants vT/n= and d' (T/(KL) is a global minimizer of (5)
n.-I
We also prove two strengthenings of this theorem in the appendix. First, under an additional techni- cal assumption on the latent variables 21,..., ZK prove its converse; any (W, U) that minimizes (5) must be in a type-I collapse configuration (with the same constants c, c'). This additional as- sumption mild but technical, so we state it in appendix C We also prove that if > nu then R(W, U) does not have spurious local minimizers; all local minimizers are global (see appendix H.
The symmetry assumption, while odd a first glance, is both needed and natural. Indeed, a type-I collapse configuration is highly symmetric and perfectly homogeneous. We therefore expect that such configurations could only solve an analogously 'symmetric' and 'homogeneous' optimization problem. In our case this means using the true risk (5) rather than the empirical risk (2). and im- posing that the latent variables satisfy the symmetry assumption. This assumption means that all latent variables play interchangeable roles, at intuitive level, that there is no 'preferred" latent variable. To understand this better, consider the extreme case meaning that all latent variables in 2 are involved in the task. The identity (6) then holds by simple combinatorics. We may therefore think of (6) as an equality that holds in the large K limit, so is neither impossible nor unnatural. We refer to appendix [o more discussion about assumption []]
n and {z....., [x] = Z,
While theorem I proves global optimality of type-I collapse configurations in the limit large Napl and large K, these solutions still provide valuable predictions when K and napl have small to
8

---

Published as a conference paper at ICLR 2024
moderate values. For example, in the setting of figure 3 (Aspl 5 and K 1000) the theorem predicts that word embeddings should have a norm
VT/ng 1.42214 (with T obtained by minimizing (t) numerically). By experiment we find that, on average, word embeddings have 50 and 'apl = 100 (and keeping n, 3. nhw = 1200, L = 15) the theorem predicts that words embeddings should have norm 0.61602. This compares well against the values 0.61 + 0.06 observed experiments. The idealized solutions of the theorem capture their empirical counterparts very well.
norm 1.41 with standard deviation 0.13. To take another example, when
For non-uniform Ng we expect hw, U to exhibit type- collapse rather than type-I collapse. Addi- tionally, in our long-tail tail experiments, we observe that frequent words (i.c. large p4)) receive large embeddings. We now prove that this is the case in our next theorem. To state it, consider the following system of Sc equations
L The - 1 + exp
ns ne-
= p4
for all 1 1535
(7)
E
( (T2)² Ln!
(8)
for the unknowns (c. ...... ). If the regularization parameter À is small enough, namely X² has unique solution. This solution defines the magnitudes of the word embeddings. The left hand side of (7) is an increasing function of ra. so pa < pa implies F3 < ray and more frequent words receive larger embeddings.
TT Ei==1 1g.. then
Theorem 2 (Directional Collapse of h). Assume X² < (L/n 2 K = n and {zi..... 2x} = Z. Suppose (W, U' in a rype-lll collapse configuration for some constants (c, :. Then (W,U) is critical point of the true risk ( if and only if (c,
solve the system 7-
Essentially this theorem shows that word embeddings must depend on word frequency and so feature 2 network exhibiting
collapse fails. Even the fully-sampled case K n' and ; 2 type-I collapse is never critical if the word distributions are non-uniform. While we conjecture global optimality of the solutions in theorem @ under appropriate symmetry assumptions, we have no proof of this yet. The bound on À is the natural one for theorem 2 for if À is too large the trivial solution (W, U = (0, 0) is the only one. In our experiments, À satisfies this bound.
Our final theorem completes the picture; it shows that normalization restores global optimality of fully-collapsedd configurations. For the network hW, o with LayerNorm, we use the appropriate limit
K
RJ (WWWU)
Σ
E
DDDDr (x), k)] + 2 1
K
2
1
of the associated empirical risk and place no assumptions on the sampling distribution.
Theorem (Full Collapse of le"). Assume the non-degenerate condition μ. > 0 holds.
Lei 2 0 denote the unique minimizer of the strictly convex function H²
log
1 1 -
K + Dr
(1+ (ne 1ine-1ee 1))= (
=22 where 1T" = n.-1 VKLjd
Assume Zi...., ZK
are mutually distinct and satisfy assumption 1 Then any (W, U) in a type-ll collapse configuration with constant e = K/VKK is a global minimizer of (9)
As for theorem we prove the converse under an additional technical assumption on the latent variables. Any (W, U) that minimizes (9p must be in a type-Il collapse configuration with T/VKL. The proof and exact statement can be found in section F of the appendix.
ADDITIONAL EXPERIMENTS
Our theoretical investigation of feature collapse uses a simple synthetic data model and basic net- work. These simplifications allow us to rigorously prove that feature collapse occurs this setting. The first section of the appendix provides preliminary evidence that the feature collapse phenomenon also occurs in more complex settings, which are beyond the reach of our current analytical tools. In particular, experimentally observe feature collapse in more complex data models that involve a deeper hierarchy of latent structures. We also investigate the feature collapse phenomenon in transformer architectures in both classification setup and the usual wext-word-prediche setup.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENT
Xavier Bresson is supported by NUS Grant ID R-252-000-B97-133. The authors would like to express their gratitude the reviewers for their feedback, which has improved the clarity and con- tribution of the paper.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staireast property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782-4881 PMLR, 2022.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305. / 2023.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35:9768-9783,
2022.
Noam Chomsky. Three models for the description of language. IRE Transactions information theory, 233):113-124, 1956.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory, pp. 5413-5452. PMLR, 2022.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer- peeled model: Minority collapse imbalanced training. Proceedings of the National Academy of Sciences, 118(43))e210309) 118, 2021.
Vitaly Feldman. Does learning require memorization? short tale about long tail. In Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, pp. 954-959, 2020.
Vitaly Feldman and Chiyuan Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. Advances in Neural Information Processing Systems, 33:2881- 2891, 2020.
Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie Su. An unconstrained layer-peeled perspective neural collapse. arXiv preprint arXiv:2110.02796, 2021.
Yoon Kim, Chris Dyer, and Alexander M Rush. Compound probabilistic context-free grammars for grammar induction. arXiv preprint arXiv: 10225, 2019.
Thomas Laurent and James Brecht. Deep linear networks with arbitrary loss: All local minima global. In International conference on machine learning, pp. 2902-2907. PMLR, 2018.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better down- stream: Implicit bias matters for language models. In International Conference on Machine Learming. pp. 22188 18882221444---111111111111111111 PMLR, 2023.
Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang, Boqing Gong, and Stella Yu. Large- scale long-tailed recognition in an open world. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 2537-2546, 2019.
Jianfeng Lu and Stefan Steinerberger. Neural collapse with cross-entropy loss. arXiv preprint arXiv:2012.08445.. 2020.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint arXiv:2011.1166 2020.
10

---

Published as a conference paper at ICLR 2024
Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv:2209. 14863, 2022.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40): 24652-24663, 2020.
Suzanna Parkinson, Greg Ongie, and Rebecca Willett. Linear neural network layers promote learn- ing single-and multiple-index models. arXiv preprint arXiv:2305. 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188)9, 2019.
Ruslan Salakhutdinov, Antonio Torralba, and Josh Tenenbaum. Learning to share visual appearance for multiclass object detection. In CVPR 2011, pp. 1481-1488. IEEE, 2011.
Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural col- lapse. In International Conference Machine Learning, pp. 21478-21505. PMLR, 2022.
Stephan Wojtowytsch et al. On the emergence of simplex symmetry in the final and penultimate layers neural network classifiers. arXiv preprint arXiv:2012.05440, 2020.
Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference on Machine Learning, pp. 271 77179-27202. PMLR, 2022.
Xiangxin Zhu, Dragomir Anguelov, and Deva Ramanan. Capturing long-tail distributions of object subcategories. Proceedings of the IEEE Conference on Computer Vision and Pattern Recogni- fion, pp. 915-922, 2014.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu A ge- ometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34:29820 29834, 2021.
George K Zipf. The psycho-biology of language. 1935.

---

Published as a conference paper at ICLR 2024
Appendix
In section A we conduct empirical investigation of the feature collapse phenomenon in settings beyond the reach of our current analytical tools. The remaining sections are all devoted to the proofs of the three theorems that constitute the main results of the paper.
Section B provides formulas for the networks bw. and hw.u depicted on figure of the main paper, and formula for the distribution D2, X - [0, 1] underlying the data model depicted on figure [f the main paper. We also use this section to introduce various notations that our proofs will rely on.
Section C is devoted to the symmetry assumptions that we impose on the latent variables. We start with depth discussion of assumption[] |] from the main paper. This assumption is required for theorem 1 and to hold. We then present and discuss an additional technical assumption on the latent variables (c.f. assumption
BB that we will use to prove the converse of theorems aand
Whereas sections B and C are essentially devoted to notations and discussions, most of the analysis
occurs in section D and G We start by deriving sharp lower bound for the unregularized
risk in section D Theorem from the main paper, as well as its converse, are proven in section
Theorem and its converse are proven in section F Finally we prove theorem 2 in section G
We conclude this appendix by proving in section that if d > min(nw, L). then the risk associ- ated to the network lw:u does not have spurious local minimizers; local minimizers are global. This proof follows the same strategy that used in Zhu et al. (202I)
A FURTHER EMPIRICAL INVESTIGATIONS
In this section, show empirically that the feature collapse phenomenon is not limited to the simple controlled setting where were able prove it. In particular, show that feature collapse occurs in the absence of weight decay and when the LayerNorm has learnable parameters. We also show that feature collapse occurs transformer architectures, and also when the classification task is replaced with language modeling task. Finally, we show that feature collapse occurs in data models involving a deeper hierarchy latent structures, such Context Free Grammar (CFG).
A.1 EXPERIMENTS WITHOUT WEIGHT DECAY
We start by reproducing the experiments depicted in Figure and but without weight decay. Specifically, we set = 0 in equations and All other parameters defining the networks and data model remain the same. To train the networks, we perform 5 million iterations batch size of 100 and learning rate of 0.1. After training, the empirical losses for all networks are below 10-4, The outcomes of these experiments depicted in Figures 5 These figures are virtually identical to Figures 3 and in the main paper. In other words, the absence or presence of weight decay does affect our main qualitative findings.
stochastic gradient descent with
Figure 5
(a) h trained on the large training set. Test acc. 100%
(b) l trained the small training set. Test acc. 66%
h² trained the small
training set. Test acc. 100%
Figure 6: Same experiments in Figure 4 but with no weight decay. Note that the result qualitatively similar.
12

---

Published as a conference paper at ICLR 2024
A.2 LAYERNORM WITH LEARNABLE PARAMETERS
When the LayerNorm module has no learnable weights, the word embeddings must lie on sphere of constant radius. This con- straint aids collapse. A natural question is whether collapse still occurs when the LayerNorm has learnable parameters. To reproduce the experiment that corre- sponds to Figure c) but allow the LayerNorm module to have learnalle weights. The result of this experiment is depicted Figure 7 and one can clearly observe that feature collapse does occur this setting as well.
answer this question,
Figure 7: Same experiment as in Figure [ but with learnable weight in the LayerNorm
A.3 CLASSIFICATION EXPERIMENTS WITH TR TRANSFORMERS
In this set of experiments, we train a transformer on the classification task depicted on figure | The transformer has layers, 8 heads, 512 dimensions, and we absolute positional embeddings (as in GPT-2). A classification token is appended to each input sentence, and this classification token used in the last layer predict the category. The network is trained with Adam (constant learning rate of 10-4, weight decay of 0.1, 31 = 0.9 and B2 0.95) during epochs on a training sets containing 0.5 million sentences. For the data model,
use The n 1200, L 15, K = 1000 as in the main paper. In figure 8 () and g (b) we display the word embeddings via dimensionality reduction and color coding. These are the word embeddings obtained before addition of the positional embeddings, and before going through the first transformer layer. Figure ((a corresponds to the case in which the words are uniformly distributed, and Figure K(b) corresponds to the long-tail case. We observe that the word embeddings, in both the uniform and long-tail case, are properly collapsed.
A.4 LANGUAGE MODELING EXPERIMENTS WITH TRANSFORMERS
In this set of experiments, we train a transformer to predict the next token on sentences generated by the data model depicted on Figure [ We use the GPT-2 architecture (Radford et al.) 2019) with 2 layers, 8 heads, and 512 dimensions. The training set contains million sentences generated by our data model with parameters ne = 3, n = 1200, L = 15, = 1000, and with uniform word distributions. We perform a single epoch through the training set and AdamW with same parameters as above. On figure (c) we display the word embeddings via dimensionality reduction and color coding, and we observe that they they properly collapsed.
A.5 EXPERIMENTS WITH CONTEXT FREE GRAMMAR
The data model presented in the paper extends to one with a deeper hierarchy of latent structures. Recall that words partitioned into concepts and that the latent variables are sequences of con- cepts. We can further partitioned the latent variables into 'meta-concepts' and create deeper latent variables that are sequences of 'meta-concepts'. We iterate this process obtain hierarchy of any depth. Such a data model particular instance of a Context Free Grammar (Chomsky) 1956), which generates sentences with probabilistic trees and are widely used understand natural
(a) Classification task (uniform word distribution)
(b) Classification task (long-tail word distribution)
(c) Language modeling task (uniform word distribution)
Figure 8: Experiments with transformers.
13

---

Published as a conference paper at ICLR 2024
language models (e.g. Kim al. (2019); Allen- & Li (2023): ;
Liu et al. (2023). In Figure 9 we
provide an illustration of a simple depth context free grammar.
We ran experiments with context free grammar of depth 4, meaning that we have words, concepts, meta-concepts and meta--eta-conceplee We used a deep neural network with ReLU nonlinearitics and LayerNorm module each layer. The architecture of the neural network was chosen to match that of the context free grammar, see Figure T0 In Figure 11 we plot the activations after each of the three hidden layers and readily observe the expected feature collapse phenomenon. All segments of the input sentence that correspond to same concept, meta-concept, or meta- receive the same representations in the appropriate layer of the network (layer for concepts, layer 2 for meta- concepts, and layer 3 for meta-meta- metta-meta-concpps)) This shows that the feature collapse phenomenon is a general one.
Details of the Context Free Grammar: We used the following context free grammar for our experi- ment. We choose K 100 categories. Each category generates a sequence of meta-meta-conceptss of length 8 by choosing uniformly at random among 5 possible sequences of meta-meta-conceells Each meta--etta-ooneet then generates a sequence of meta-concepts of length 8, again by choos- ing uniformly at random among 5 possible sequences of meta-concepts. Each meta-concept then generates a sequence of concepts of length 8 by choosing among 5 possible sequences of concepts. Finally, each concept generates a sequence of words of length 8 by choosing among 5 possible sequences of words. At level ( the sequences of words have an overall length of 8 4096.
Class Index (1.2...))
Oovel B)
Sequence of mata-concapts
[Jevel
Sequence concepts
[esel
ponasio
penostor
passoees
ginenss
10g
PL
06
m1
Seguence of words
Bleeel D]
Figure 9: Probabilistic Context Free Grammar of depth 3. The class index generates a sequence of meta-concepts. Each meta-concept further generates sequence of concepts. Finally each concept generates a sequence of words. The process by which a token from one level generates a sequence of tokens in the level below is random. For example, the meta-concept 722 5 might generate the sequence of concepts [3,, So, fa] = [4, 1, with probability 1/3, the sequence [84, 3o, Sel = [2, 5, 3] with probability 1/3, and the sequence [5a, $5, 36] = [3,5, 5, 5] with probability 1/3.
Wectors in 2t Bevet 3]
Vector in (level 2]
Frs
Vector in Rr (Jowl 1]
S K
An
XrT
Vectors In
(level
Emshedding
PM
017
Sequence of words
Figure 10: Neural network architecture matching the context free grammar from Figure Each block consists in MLP followed by layer normalization. The vectors are concatenated before being fed to a block. If feature collapse occurs, then two sequences of words [01, 02, 0g) generated
by the same concept 31 should have almost identical representation by in level of the network. Similarly, two sequences of words [an, ..,..., agg generated by the same meta-concept Y1 should have almost identical representation č in level 2 of the network.
14

---

Published as a conference paper at ICLR 2024
Level representations
Level representations
Level represenaations
color coded according to concepts
oslor-caded according metiommtt--cnnceph
-15
Figure 11: Results of an experiment ran on context free grammar of depth (so deeper than the one depicted Figure 1). At every level, each token generates in the level below sequence of length chosen uniformly at random among possible sequences. After training the network, we generate 1000 test sequences, feed them to the network, and visualize via PCA the representations obtained each layer (i.e. we plot the vectors b,, é₁ and d,). The vectors color coded accord- ing to the concept, meta-conceph, and meta-meta-conen representations corresponding the first concepts, meta-concepts, and meta-meta-concepts. We clearly observe the feature collapse phenomenon. Note that since we are dealing with network of depth 4, the vectors d, are not the output, but the last hidden representations.
that generated them, and we keep only
B PRELIMINARIES AND NOTATIONS
B.1 FORMULA FOR THE NEURAL NETWORKS
Recall that the vocabulary is the
y = [(, : S no and 1 5 5 seh.
and that think of the tuple (a, E V representing the 8th word of the a concept. The data space is X = VL, and sentence x E A' is a sequence of L words:
((1, (o1, BL)] 1 1 oz S n. and 1 S 8r 8c-
The two neural networks h, ht studied this work process such a sentence x E X' multiple steps:
1. Each word (ar, Be) of the sentence is encoded into a one-hot vector.
2. These one-hot vectors multiplied by matrix W to produce word embeddings that live in a d-dimensional space.
3. Optionally (i.e. in the case of the network h"), these word embeddings go through Lay- erNorm module without learnable parameters.
4. The word embeddings concatenated and then goes through a linear transformation U.
We now formalize these 4 steps, and in the process, the notations which will rely in all our proofs.
Step 1: One-hot encoding. Without loss of generality, choose the following one-hot encoding scheme: word (a, B) E 1/ receives the one-hot vector which has a in entry (a - 1188 + B and everywhere clse. To formalize this, define the one-hot encoding function
C(a, 3) = efa-1)sc+8
(10)
where e; denotes the 5th basis vector of Rw, The one-hot encoding function Ç can also applied to sequence of words. Given a sentence [(a), (1d,,,8).... Bi)...., (a,)...,....)) BI)] E X' we let
((x) = ((o), ((02) 92)
((a),8)) E x
(11)
and maps sentences to Pew X L matrices.
15

---

Published as a conference paper at ICLR 2024
Step 2: Embedding. The embedding matrix W has nw columns and each of these columns be- longs to R.. Since Ç(a, 3) denote the one-hot vector associated to word (a, B) E V, we define the embedding of word (a, by
(o. == W ((a, 3) R
(12)
Due (10), this means that () is the 5th column of W, where (a 1)se + 3. The embedding matrix can therefore be visualized as follow (for concreteness we choose ne = 3 and The 12 as in figure the main paper):
W ((.)) ((2)) W(1,3) W(1.4)
((.1)) ((3.)) ((33)) W(3 (2,1) (2,2) W(2.3) W(2,4)
Emheddings of the words the 1 concege.
Eibbeddings the words the 2nd concept.
Embeddings of and concept.
Given a sentence x = ((01,8)).
(aL, E X, appealing to (11 and 12
we find that
W[(x)= = ((an.9)) W(a).02)
E RR
(13)
and therefore WC(x) is the matrix that contains the d-dimensional embeddings of the words that constitute the sentence x X.
Step 3: LayerNorm. Recall from the main paper that the LayerNorm function yp Rd Rd is
defined by
(Pv))
mean(y)la where o(v)
mean (v)
1
Σ Us and (v)
1 Σ (Us mcan(v))
We will often apply this function column-wise matrix. For example if V is the x matrix
V
Vi
V2
Vmn
then
v(V) = (l(vi) vlv2)
(p(mm)
Applying yp
13
gives
x <(W((x))=
o ((Wrocite))
x ((Wrrrssu))
E RaxL
(14)
and so to (W((x) contains the word representations of the words from the input sentence (recall that by word representations we mean the word embeddings after the LayerNorm).
Step 4: Linear Transformation. Recall from the main paper that
U1,1
U122
U1
U2.1
1222
U2
Kx
(15)
uK.1
UK
where each vector u belongs to IRd. The neural networks fw. and hūv. are then given by the formula
(w) U Vec [W [x(()]
(16)
(xx)((x) U Vec ppl (W((x))
(17)
16

---

Published as a conference paper at ICLR 2024
where Vec R Rd is the function that takes as input d L matrix and flatten it out into a vector with dL entries (with the first column filling the first d entries of the vector, the second column filling the next d entries, and so forth). It will prove convenient to gather the L vectors u that constitute the Ethh row of U into the matrix
Üx =
1g,11 11222
E Rdx1
(18)
With this notation, we have the following alternative expressions for the networks hww and wwo
Ü1. w((x))
U1, P(W ((x((()))
U2. w ((xx))
Ü2, 10 s( ((x)) ((W((x)))
bw,o(x)= =
and lx)x(x==
=
(19)
(Üx, w((x))
(Ük, o (W<(x)))
where ( denote the Frobenius inner product between matrices (see next subsection for defini- tion).
Finally, we use Ü to denote the matrix obtained by concatenating the matrices UK. that is
= Ú1 U2
Üx] E RRdXL
(20)
The matrix Ü, which is nothing but a reshaped version of the original weight matrix U R will play crucial role in our analysis.
B.2 BASIC PROPERTIES OF THE FROBENIUS INNER PRODUCT
We recall that the Frobenius inner product between two matrices A, E mx is defined by
(A, = -10==E2 Bi
(=1==1
and that the Frobenius norm of a matrix A E m given by ||Al|= ((A, (A,A) A) In the course of our proofs, will constantly appeal to the following property of the Frobenius inner product, so we state a lemma once and for all.
Lemma A. Suppose A Rmxn, B E Rmxr and C E Brxn, Then
(A, BC) = (BF
and
= (ACT, ,B)
Proof. The Frobenius inner product can be expressed (A, B)F Tr(AT B), and have
(A. BC) Tr Tr ((D" A) c) =(BT SSTACCY A,
Using the cyclic property of the trace, we also get
(A, BC) = Tr(A BC) TT TI ((ACT)T T B) =ACT,,
B.3 THE TASK, THE DATA MODEL, AND THE DISTRIBUTION Dz:
Recall that C (1,...,,.) n.) represents the set of concepts, and that Z C is the latent space. We aim to study a classification task in which the K classes are defined by K latent variables
Z1
24 E
17

---

Published as a conference paper at ICLR 2024
We write X Des to indicate that the sentence X' is generated by the latent variable Zμ Z (see figure of the main paper for visual illustration).. Formally, Dz is a probability distribution on the data space X, and now give the formula for its p.d.f. First, recall that μa >> 0 stands for the probability of sampling the 3th word of the ah concept. Let us denote the 4 latent variable by
Zk ] Z where 1 I t S ne. The probability of sampling the sentence
Zk [ 2k,1 2k.2
x
(a1. B1), (a2, 52)
(aL, X
according to Dz, is then given by the formula
Dr, ((x))
1 (az=s,1)) M3 t=1
II
So a sentence x has a non- zero probability of being generated by the latent variable z only if its words match the concepts in Zk. If this is the case, then the probability of sampling x according to Dz, is simply given by the product of the frequencies of the words contained in
Note that Da., ((x)) > O if and only if [.g.1..... , Zt.r.] =
We use 2 to denote the support the distribution Dzu, that
X = {x E X Dz, (x) > 0}
and we note that if the latent variables Zi,.... 2x mutually distinct, then X, n. 0 for all j * k. Since the latent variables define the K classes of our classification problem, we may alternatively define N3 by
X4 = fx E X' belongs to the htth category)
To each latent variable z 2k,1 2k.2
we associate matrix L
Zk
e
R
(21)
In other words, the matrix Z provides a hot representation of the concepts contained in the latent variable Zg. Concatenating the matrices Zi,..., ZK gives the matrix
Z Z2 Z2
Zk] E R
(22)
which is reminiscent of the matrix t defined by 20
We encode the way words are partitioned into concepts into a 'partition matrix' P E KR. For example, if we have 12 words and 3 concepts, then the partition matrix is
1 1
0
0
0 1
0
0
0
0
E
(23)
indicating that the first words belong to concept 1, the next 4 words belongs to concept 2, and so forth. Formally, recalling that ((a,9 8) is the the one-hot encoding of word (a, E V. the matrix P is defined the relationship
P P((a, 3) = ea for all (a, B) E V.
(24)
Importantly, note that the matrix P maps datapoints their associated latent variables. Indeed, if = [(a1 (az, is generated by the latent variable (meaning that x E Xx,,, then we have that
P((x)= P (1)1,18)) (2)2)8))
C(al,
Bai
G
Pal = Zs (25)
where the last equality is due to definition of the matrix Zy-
Another important matrix for our analysis will be the matrix Q e R In the concrete case where we have 12 words and concepts, this matrix takes the form
P22
44
0
Q
0
0
0
14
12
13
44
0
0
0
0
R
(26)
0
0
P1
P2
Hg
P4
and, in general, it is defined by the relationship
Q ((a: (3)
118 Cex
for all (a, 3)
(27)
18

---

Published as a conference paper at ICLR 2024
C SYMMETRY ASSUMPTIONS ON THE LATENT VARIABLES
In subsection we provide an in depth discussion of the symmetry assumption required for theo- rems and 3 hold. In subsection C.2 we present and discuss the assumption that will be needed to prove the converse of theorems and 3.
C.I SYMMETRY ASSUMPTION NEEDED FOR THEOREM 1 AND 3
To better understand the symmetry assumption I from the main paper, let us start by considering the extreme case
K K=n n and
{Z1, Z........ ZK] 2,
(28)
meaning that 21...., ZKK are mutually distinct and represent all possible latent variables in 2. In this case, we easily obtain the formula
(L. Ki [K] (2s) 21) and Z3. = =1 11- (1
(nc -
1)"
(29)
where dzst(z;, zı) is the Hamming distance between the latent variables 2, and 21. To see this, note that the left side of (29) counts the number of latent variables z, that differs from 21 r locations and agrees with 21 at the last location é = L. This number is clearly equal to the right side of 22 ( since we need to choose positions out of the first L positions, and then, for each chosen
position f, we need to choose a concept out of the ne - 1 concepts that differs from Z1, A similar reasoning shows that, if 2 # a, then
Kis E [K] dist(z,, 21) and 23, = =a = - -(--1) (ne 1)"
(30)
where the term (+=)) arises from the fact that we only need choose r- positions, since z and 2, differ in their last position é L. Suppose now that the random variables 21,..., 2x selected uniformly at random from 2, and say, for the sake of concreteness, that
K
n
so that 21..... Z represent 20% of all possible latent variables (note that |2|= -=]=n)) n²). Then (9
30 should be replaced by
{s E [K] dist(zj, 21) = and 2j, = =21,1
~ 1/1. L -
(ne -1)"
(31)
{s E [K] dist(zj, 21) = and 25 = =a> ~
1 L 1
(n. 1)"-1 for a 21
(32)
where the equality only holds approximatively due the random choice of the latent variables. In the above example, we chose 1 as our 'reference' latent variables and we 'froze' the concept appearing in position T.. These choices were clearly arbitrary. general, when K large, we have
K (1-11 (no  1)"
if if
{i [K] dist(z,, zb) = and Zj.e aafl
(33)
K
(n. (
1)--1 if if + a
and this approximate equality hold for most & [K] r [L], [L], and [nc]. The symmetry assumption [ from the main paper requires (33 to hold not approximatively, but exactly.
For convenience we restate below this symmetry assumption:
Assumption A (Latent Symmetry). For every [K],
[L], E [E], and a E [md] the identities
K [--1] (no 1)"
zk
{i E [K] dist(2,, 21) = and Zj,E = =
no
=
(34)
K
(x-))
(1-))(mo -1)-11
if + a
hold.
19

---

Published as a conference paper at ICLR 2024
To be clear, if the latent variables Z1...... ................................ Zg selected uniformly at random from 2, then they will only approximatively satisfy assumption A Our analysis, however, is conducted in the idealized case where the latent variables exactly satisfy the symmetry assumption. Specifically, we show that, in the idealized case where assumption is exactly satisfied, then the weights W and U of the network given by some explicit analytical formula. Importantly, as it is explained in the main paper, our experiments demonstrate that these idcalized analytical formula provide very good approximations for the weights observed in experiments when the latent variables are selected uniformly at random.
In the next lemma, we isolate three properties which hold for any latent variables satisfying as- sumption A Importantly, when proving collapse, we will only rely on these three properties - we will never explicitly need assumption A We will see shortly that these three properties, in essence, amount to saying that all position E TT. and all concepts a E [ne] plays interchangeable roles for the latent variables. There are no 'preferred' & or a, and this is exactly what will allow us to derive symmetric analytical solutions.
Before stating our lemma, let us define the 'sphere' of radius centered around the k*h latent variable
(,) = {i E [K] dist 2x) =
for r, k E [7]
(35)
With this notation in hand we may now state
Lemma B. Suppose the latent variables Zi..... Zk satisfy the symmetry assumption A Then
Zi...., ZK satisfies the following properties:
(i) |s,(i] = (k)( for all r E [L] and all j, E [K].
(ii) The equalities
K
K 1n.
and Zz KL The
Σ
Zn
II
Ince
The
k=1
hold, with In. denoting the ne x identity matrix.
(iii) There exists 0 0,,...,,,
O2 > 0 and matrices AI,...,AA E BR.xxL xL such that
Zk
Σ
Z,= 22 -, A
[k)
JES-(R)
holds for all r E [L], all j E [K], and all k E [K].
We will prove this lemma shortly, but for now let us start by getting some intuition about properties (i). (ii) and (iii). Property (i) is transparent: it states that all latent variables have the same number of "distance-r neighbors'. Recalling how matrix Zμ was defined (c.f. (1)). we see that the first identity of (ii) is equivalent to
K ne
(k
[K] Zk,e a]]=
for all & [L] and all [nc].
(36)
This means that the number of latent variables that have concept a in position & is equal to K/ K/ne- In other words, each concept is equally represented each position é. We now turn the second identity of statement (ii). Recalling the definition (22) of matrix Z, we see that Zz E R is diagonal matrix since each column of Z contains single nonzero entry. One also easily that the ah entry of the diagonal is
[zz"] a.a = [(k, E) E [K] x [L] 24, = =a3]
which is the total number of times concept appears in the latent variables. Overall, the identity Z7T KL Tn, is therefore equivalent to the statement ne
KL
k[k, E) [K] x [Z] 24,E = atf- =
for all a E [ne]
ne
and it is therefore direct consequence of 36
Property (iii) is harder to interpret. Essentially it is a type of mean value property that states that summing over the latent variables which mean value property plays a key role in our analysis.
distance of 2μ gives back 2.. We will see that this
To conclude this subsection, we prove lemma B
20

---

Published as a conference paper at ICLR 2024
Proof of lemma B we clearly have that
We start by proving statement (i). Since (k) {i E [K] dist(zz, zx) r),
The
[s,(k)
Σ
{
[K]
dist z4) - and Zj,t a
(37)
We then use identity
33 and Pascal's 's rule to find
K
L - 1 r r-
(no -1)--1)
3 ) (ne 1)"
K
L 1-1)
[A)] (ne 1)
+
2
K
((-11)
-(-))
L (1-11
(nc - -1)²
|Z|
K
L
(ne -1)
(38)
Z
which clearly implies that (k)) (S) for all j,k [K] and all E [Z].
We now turn to the first identity of (ii). As previously mentioned, this identity is equivalent to 33).
Choose k such that Zx,e # a. Then any any latent variable 2; with 2j is at least at distance
1 of ze and we may write
I. =
Ki E [K] 2/0 - a
Σ
{s [K] dist(z,, 2x) = r and 23.6 = a}}
(39)
1=1
K [--1)
Σ
(ne (nc 1)
(40)
n
(--1)
which is equal to K K/n, according to the binomial theorem. The second identity of (ii), as mentioned earlier, is a direct consequence of the first identity.
We finally turn to statement (iii). Appealing to (38), we find that,
{
[K] dist(zz, 24) and 2j, = }}
(1=11) (me -1)-
(1-1
L-r (+)
[K)
K (nte -1)-
L
=k,e= = a. On the other hand, if zn.e F ar we obtain
{ j []] tist(zz,2z) 2ι) and Zj,e = a)
[ r- (nc 1)r-1
K I- 11
L-1
p-
1
(k)]
K
(n. 1)r
Ne
[2
Fix E [Z] and assume that z3,e = a², We then have
1 The
Σ Czye =
Σ
{ E S,(k) :2,e zje
cal Ca
SS (1))
[k))
(E)
{ [K] : dist(zj, 2x) = and Zj. = =a
Co
(k))
a=1
L-r
Σ Ca
Ca. +
ne L
astt
L
r
1
r
e
C
+
e
L
The
- 1
N1e 1
L
a=1
r
1
r 1 no
On
+
ne
T
Nee 1I
Recalling that 24,E = a*, the above implies that
e
1
Σ
Re
ea,
1
L Inc
(41)
[s,(k:]
The -
1L
To
(ES.(K)
21

---

Published as a conference paper at ICLR 2024
Finally, recalling that
Z
2z
1
we see that (41) can be written in matrix format as
Z
Σ
Z,
n2
Z4
1n. 1f
(k)
The
Re
JES-(K)
and therefore the scalars 9, and the matrices A, appearing in statement (iii) are given by the formula 8,
Th I and A, ne-1 1n,1% 1..
n--1E
C.2 SYMMETRY ASSUMPTION NEEDED FOR THE CONVERSE OF THBOREM
AND
In this subsection we present the symmetry assumption that will be needed to prove the converse of theorem [ndd This assumption, for small values of K.
will shortly see, is quite mild and is typically satisfied even
For each pair of latent variables (z,, zk) we detine the matrix
(() = Z,(Z, - -x)² T E R
We also define
A= A X.. There exists a, b s.t. A aln. + bln, 1
(42)
which is the set of matrices whose diagonal entries are equal to some constant and whose off- diagonal entries are equal to some possibly different constant. We may now state our symmetry assumption.
Assumption B. Any positive semi-definite matrix A E Re that satisfies
(A, p(Ak) - r(('k)) =
Vj, J', K E [K] s.t. dist(z,, za) = dist(zy, zz)
(43)
must belongs A.
Note that 143 can be viewed as a linear system of equations for the unknown A E R Rn. with one equation for each quadruplet (j, j', k²) satisfying dist(z,, 2x) = dist(Zy. 2w). To put differently, each quadruplet (j, j k²) satisfying dist(z,, z dist(zy', Zue adds one equation to the system, and our assumption requires that we have enough of these equations so that all positive semi-definite solutions are constrained
live in the set A. Since a symmetrie matrix has (n, + 1)ne/2 2 distinct entries, we would expect that (n, + 1)n./2 quadruplets should be enough to fully determine the matrix. This number of quadruplets easily achieved even for small values of K. So assumption B is quite mild.
The next lemma states that assumption B is satisfied when K = In light of the above discussion this is not surprising, since the choice K = n leads to system with a number of equations much larger than (ne + 2//e/2 The proof, however, 2 quadruplets to determine the entries of the matrix A. The 2-2 arises from the fact A is a 2 dimensional subspace, and therefore (n. + 1)ne// 2 equations are 'enough' to constrain A to be in A.
instructive: it simply handpicks (n. + 1)n./2
Lemma C. Suppose K = n and (2i....,2y) = 2. Then Zi Zi...., zK satisfy the symmetry assumption B
Proof. Let A C be positive semi-definite matrix that solve satisfies 43) We use Ca to denote the th column of C. Since . . 2k} = Z, we can find i, k E [K] such that
2, 2, 1] Z
z, 1,,,,,,..,,11
Z 1,.,, E
22

---

Published as a conference paper at ICLR 2024
Using lemma A and recalling the definition 21 of the matrix Zro we get
( A, ((J)
=
= (C',, CZi) (CZj, CZi)
(122, 02) + (T. 1)(0), (102, 03) + (I. (a)
= (C2, C2) - (c2, c3)
Similarly we obtain that
< A, ((i, (c2, C2) (C2, c4)
Since dist(z;, zj)
dist (Zi, zn) = 1, and since satisfies 43 we must have
A, I(t)) ) = - A, r(ta)
which in turn implies that
2 (c2, 03) (c2, ca) A2,4
This argument easily generalizes to show that all off-diagonal entries of the matrix A must be equal to some constant R.
We now take care of the diagonal entries. Since (Z1 Zk} Z, find j' k E [K] such that
Z [1, 1, ..,,, E
Zj' 2 22
Z [3,3, ..,3....3......... E
As before, compute
A, (( = -CZz, CZr) CZe) L(ci,
=l(c,,))) L(ci, L(ei, ei) - Lb
where we have used the fact that the off diagonal entries are all equal to b. Similarly we obtain
< A, r()' L(c22, c2) Lb
Since dist(zy, 2y) dist(z), xΦ) L, we must have ( (A,T (A,T x)) which
implies that A1,1 A2,2- This argument generalizes show that all diagonal entries of A equal.
D SHARP LOWER BOUND ON THE UNREGULARIZED RISK
In this section we derive a sharp lower bound for the unregularized risk associated with the network hw.
Ro(W,U =
D ((),,,(())))))
(44)
K
where E: R.K R is the cross entropy loss
exp (3/k)
fly, = - log
for e R
==1 exp (u;)
The tt entry of the output y hw.u (x) of the neural network, according formula 19 given by
8/4 = (Ür,
Ü W <(x))
23

---

Published as a conference paper at ICLR 2024
Recalling that Xx is the support of the distribution Dz. : :x= [0, 1], we find that the unregularized risk can be expressed as
Ro
(()w,o(x)) D2, (x)
K
k= XxX
(l,
log
Dzo (x)
Σx ((,,
xxXx
E
log
1+ Σ -0-(ti--ii -e- jtk
Dea (x)
XEX,
where we did the slight abuse of notation of writing Zzn (x) instead of Day ((x)). Note that data points x that belongs class k is correctly classified by the the network hww if and only if
( ÜUS W ((xx)) > { W (((x)) for all j + k
With this in mind, we introduce the following definition: Definition A (Margin). Suppose Ng. Then the margin between data point x and class is
Mw,u(x.)) = (Ux (Uk-U,yh
U W((x))
With this definition in hand, the unregularized risk can conveniently be expressed as
Ro(W, U) K
ΣΣ log 1
Σ 37k
Nuw (x..))
Dzs (x)
(45)
k=
and data point Xr is correctly classified by the network and only if the margins Nww (x, j) are all strictly positive (for j + k). We then introduce a definition that will play crucial role in our analysis.
Definition B (Equimargin Property). If
dist(Zk. 2z) (iis(zw,ZZ)) = Rx.((x.J) = Www (x', Vx E Nx Vxd'Vx' E X
then say that (W, U') satisfies the equimargin property.
To put simply, (W,U) U) satisfies the equimargin property if the margin between data point Xx and class j only depends on dist(Zx. Zz). We denote by & the set of all the weights that satisfy the equimargin property
((WW)) : (W,U) satisfies the equimargin property)
(46)
and by A the set of weights for which the submatrices Ür defined by (18) sum to 0,
K
N (W.U): ΣG-0
(47)
A=1
We will work under the assumption that the latent variables 21,.--, z satisfy the symmetry sumption A According to lemma B [k)] then doesn't I depend on k, and so we will simply use |S,] to denote the size of the set (k). Lemma also states that
Zu
Σ
Zj = Zx + Ar
[K)(K)
jES,(R)
for matrices A,,..., Ar. and some scalars
01,..., 0r. O. We these scalars define
g(x) = log
Σ
(48)
r=1
and we note that g: RR + R is a strictly increasing function. With these definitions in hand we may state the main theorem of this section.
24

---

Published as a conference paper at ICLR 2024
Theorem D. the latent variables satisfy the symmetry assumption A then
Ro(W.U) U) o( (0(--,, (U, r)
for all (W,U) ENnE
(49)
Ro U) > (0((0,
for all (W,U) E
(50)
We recall that the matrices Ü, Q. and where defined in section B (c.f. (20). and 22 The
remainder of this section is devoted to the proof of the above theorem.
D.1 PROOF THE THEOREM
We will use two lemmas to prove the theorem. The first one (lemma D below) simply leverages the strict convexity of the various components defining the unregularized risk Ro- Recall that f I R is strictly convex, and if the strictly positive scalars ,r > 0 sum to 1, then
Σ PiVi S E p.f(v.)
(51)
and that equality holds if and only if v1 V2 .... nn. For this first lemma, the only property
we need on the latent variables is that [k)( = |S, = |, for all j, k E [K] and all E [4].
Define the quantity
(w. = K |
Σ
Σ
Mx.u(x.)) Dau (x)
(52)
k= jES, X
which should be viewed as the averaged margin between data points and classes which are distance of one another. We then have the following lemma:
Lemma D. for all j, k E [K] and all r E [4], then
L
U) log
1+ S Tw.o6))
Σ
for all (W,U)e&
(53)
L
R) log
1+ |
Σ
nw.u(r)
for all (W,U) &
(54)
Proof. Using the strict convexity of the function R- defined by
f(V1...., Uk-1, ukt1...., UK) = log (1+ Σ
e 1/
jz
obtain
K
Ro(W,U)=
log 1 Σe e [((x)
Dze (x)
k=
xexxx
jfk
1
log
1 -+
Exex, Dax
jyk
and equality holds if and only if, for all [K], we have that
Mxx = M)
for all x, y E Xx and all j i k
(55)
We then let
m)(kk
Σ x
My( D)D (x)
25

---

Published as a conference paper at ICLR 2024
and use the strict convexity of the exponential function to obtain
K
K
log
1+ Σ
(k.)))
=
log
1+
Di(ke))
K
K
k=1
1=1
r (ES((R)
K
log
1+
Sr1
(k())
K
1=1
r=1
jES.(R)
K
L
log
1+
Σ
Exs,, (K.)))))
K
r=1
Moreover, equality holds and only if, for all l E [K] and all r E [L]. we have that
DK(k, i) () for all i,j E (k)
(56)
We finally set
Mi
().
St
jES,(K)
and use the strict convexity of the function f(v1....,
UL. = log 1 + ++E (5
to get
1
Σ log
1
Σ |s
Is,le-micr.)) wick.r)
2 log 1+ Σ
s * Ek Ti(ker)
K k=
r=1
Moreover equality holds if and only if, for all & E [K] and all r E [L], we have that
Ti(kr) = (i(,') r) for all k. K' e [K] and all r E [Z]
(57)
Importantly, note that
K r)
K
Di(k
Σ (x, Day (x)
K
K S,
(k)
XEXx
which is precisely how (r) defined (c.f.
(52). To conclude the proof, remark that
conditions (55. 56 and (57 are all satisfied if and only if (W, U') satisfies the equi-margin property.
We now show that, if assumption holds, Nww (r) can be expressed in a simple way. Lemma E. Assume that the latent variables satisfy the symmetry assumption A Then
0r (w: K (U.
for all (W,U) E
(58)
Proof. We let
X
Σ XEX
((x)D2x (x)
and note that the averaged margin can be expressed
K
(r.u(r)
1 1
Σ
Σ
Σ
(x, Den (x)
K
|Srl
k= (k), xe
Σ (Ux -la-ūp Cg, ((x)) F (x)
K S
k=ljes, (k) XEdx
1
Σ Σ
(Un- -l--ū, Ú, Un,X,) F
K
S
(k)
1
Σ
(Xx)
1 11 M (Un,WX,),
(59)
K
K (k))
26

---

Published as a conference paper at ICLR 2024
Let
1
if dist(Zz: 2,)
0
otherwise
and rewrite the second term in 59 as
K
1 1 EE
WU,,
(2) (U,, W
K S.
K
St
(M)
(r)
(xx,
K
S
K
Wix,
K
St
(EI)ES.(K)
K
ÜK W
Σ JES,(M)
Xj
K
S.
k=
F
Combining this with (59) we obtain
(r)
ÜR. W
X
Σ
X,,
(60)
S,
jES-(R)
From formula (66, we see that row a of the matrix Q is given by the formula
QT ea
3) ug.
(61)
9=1
We then write 2ạ = [ai,..., onl] and note that the gth column of X can be expressed as
X
Clar, 8)us QTear.
(62)
B=1
From this we obtain that
X 7T
and therefore (60 becomes
K Σ
Rw.u(r) =
Ür WQ (zx [Sr
Σ 2,)
K
k=1
jESr(K)
F
K
ÜK WQ (D.24 + +A,) ) );
K
k=1
where we have used the identity Z 1
|s,1 EjEs,( Z;= = 0, + A, to obtain the second equality.
Finally, we use the fact that EB Úx = 0 to obtain
(w.v(r)
0, Σ
Ür. WQT Zz )r K
0r
(U,
K
k=1
Combining lemma D and B
concludes the proof of theorem D
27

---

Published as a conference paper at ICLR 2024
E PROOF OF THEOREM 1 AND ITS CONVERSE
In this section we prove theorem under assumption A and its converse under assumptions and
We start by recalling the definition of type-I collapse configuration.
Definition C (Type-I Collapse). The weights (W, U) of the network hww form a type-l collapse consiguration if and only if the conditions
There exists e 2 that W(a.3) cfa for all (a, EV..
ii There exists d' 2 so that Ug,F = e fa for all (k:, () satisfying žk,t = a and all a E C.
hold for some collection fi.....fn. E IRd of equiangular vectors.
It will prove convenient to reformulate this definition using matrix notations. Toward this goal, we define equiangular matrices as follow:
Definition D. (Equiangular Matrices) A matrix & E the relations
Rdxx. is said to be equiangular if and only if
81n, =
and
8
3
The Inc
1
1 1
The ne
hold.
Comparing the above definition with the definition of equiangular vectors provided in the main paper, we easily see that matrix
8
=
f1
f2
fn. E dRd.
equiangular if and only if its columns fi..... ino E Rdd equiangular. Relations (i) and (ii) defining type-I collapse configuration can now be expressed in matrix format as
W = c P and t- 8 Z
for some equiangular matrix 8
where the matrices Z and are given by formula 22 and We then let
2 = (W,U): There exist equiangular matrix 3 such that
W = c and Ü=
The
(63)
KL
and note that 52/ simply the set of weights (W, U) which are in a type-I collapse configuration now state the main theorem of this section.
with constant and e' = /////(K) Theorem E. Assume uniform sampling μs 1/se for each word distribution. Let T 2 0 denote the unique minimizer of the strictly convex function
=log and let c TT/mmm- Then we have the following:
(t) log (1- K
K
1iee-11ee-) At where
The
v Thpp KL.
(i) If the latent variables Zi,..., ZKK are mutually distinct and satisfy assumption A then Si arg min R.
.....,,K
(ii) If the latent variables Zi are mutually distinct and satisfy assumptions then
A and
2 arg min R
Note that (i) states that any U) E 22 is minimizer of the regularized risk - this corresponds theorem T from the main paper. Statement (ii) assert that any minimizer of the regularized risk must belong S2 this is the converse of theorem | The remainder of this section is devoted to the proof of theorem We will assume uniform sampling
A3 1/s. for all 3 [sc]
everywhere this section - all lemmas and propositions are proven under this assumption, even when not explicitly stated.
28

---

Published as a conference paper at ICLR 2024
E.1
THE BILINEAR OPTIMIZATION PROBLEM
From theorem D it is clear that the quantity
(o.wQTz), F
plays an important role our analysis. In this subsection we consider the bilinear optimization problem
maximize (Ū.wQ*z) 2)
(64)
(IWII + +10111 = c2 nhw
subject to
(65)
where è E R is some constant. The following lemma identifies all solutions of this optimization problem.
Lemma Assume the latent variables satisfy assumption A Then (W,U) U) is a solution of the optimization problem 644 (65) if and only if belongs to the set
By (W, There exist matrix F E [gdxne with F The
such that W = cP and Ú
FZ
(66)
KL
Note that the set B is very similar the set 22 that defines type-I collapse configuration (c.f. In particular, since an equiangular matrix has ne columns of norm 1, it always satisfies || | Res and therefore have the inclusion
92
2 C By
(67)
The remainder of this subsection is devoted to the proof of the lemma.
First note that the lemma trivially true if 0, so we may assume * for the remainder of the
proof. Second, we note that since μ4 = 1/8.. then the matrices P and Q defined by (23) and (26)
are scalar multiple of one another. We may therefore replace the matrix Q appearing by
wich leads to
maximize
(U,WP*Z)
(68)
1 (IWII + 1101122)
subject to
6² nhw
(69)
2
(69) and have objective value equal to We now show that any (W,Ū) E By satisfies the constraint
Se 2 Lm
Claim A. If (W,Ū) E B,, then
2 [111++10013 c² nw
1
and
0 (U.WPTZ)
= / KIn
Proof. Assume (W,U) E By From definition 23 of the matrix P. we have PpT sel and therefore
- || = 2IFF1 = (FP,FP) = ( FF) c² 8. | = è 8e n, c² mh
where we have used the fact that nww nec Using ZzT KL I from lemma B we obtain
no. KT.
FZ
= (FZZ", F)r
|
KL
ne
As consequence we have
| c Nuw
||FZ|| à new
KL
and, using PPT = 8oIn, one more time,
(0.WPTz) = =
The
(FZ, FPPFPP" = ???
The
(FZ, F7) =2 8c v KTn
KL
KI.
29

---

Published as a conference paper at ICLR 2024
We then prove that W and U must have same Frobenius norm if they solve the optimization problem.
Claim B. If (W, U) is a solution of 68
then
1101 = =.² fh
(70)
Proof. We prove it by contradiction. Suppose (W.U) U) is solution of 66
65
with || / *
|| Since the average of || 11 and || 1/% is equal to e > 0 according to the constaint, there must then exists of 0 such that
| = e'ww te
and
|1013 c -
Let
n
&n Ú
Wo
V cn +
W
and
Üg V c2n
and note that
|Wlll = |100| c² now
and therefore (Wo, U0) clearly satisfies the constraint. We also have
(Uo. WoP 2) - An cn
(0.wPTZ), > (0,weTZ),
since # 0 and therefore (W, Ú) can not be maximizer, which is contradiction.
As consequence of the above claim, the optimization problem 68 169 is equivalent to maximize
(Ü, (U,WP') WWP22 z)
(71)
subject to [W]] = o² nww and | c² the
(72)
We then have Claim C. (W, Ú) is solution of 77
72 then (W,Ū) By
Note that according the first claim, all (W, Ú) E B have same objective value, and therefore, according to the above claim, they must all be maximizer. As consequence, proving the above claim will conclude the proof of lemma F
Proof of the claim. Maximizing 71 over Ú first gives
WP'2 2 0 =
(73)
PT WWPTZ|F Z||F
and therefore the optimization problem
reduces to
maximize
W P' Z W | | e² Thu
subject to
The Using ZZT KL / from lemma
B we then get
KL ne
| PT D (WP
(WP" n7. WPT),
| 1 1|
and therefore the problem further reduces to
maximize
IWPT112
subject to
| c² Tbu
The KKT conditions for this optimization problem are W p = wW
(74)
W 1 2 n
(75)
30

---

Published as a conference paper at ICLR 2024
where E R is the Lagrange multiplier.
Assume that (W,U) Ú) is solution of the original optimization problem 71
72
Then, according
to the above discussion, W must satisfy 74 PP" sefn, gives
75
Right multiplying (74 by pr, and using
8c WP =WPp
So either v = Se or WP = 0. The latter not possible since the choice WPT = leads to an objective value equal to in the original optimization problem (71) have / = Sc, and equation (74) becomes
72 We must therefore
1 W
WPTP P
(76)
8c
which can obviously be written
W cFP
by setting F = c3c
1. WPP WP.. S Since W satisfies
75 must have
2 Thw W = 2/IPP (FP, @ = scl1P112
(77)
and ||1113 nw/se Sec Ne.
According to (73 Ú bust be scalar multiple of the matrix
WPT Z (cFP) = cs. FZ
Using the fact that zz KL and || we then obtain that
KL || KL
[F (FZ, = (FZZT, FFF==
(78)
and equation 73 becomes
WPTZ
F7 KKT.
i = c /nw
|| WWPTZ|e = c cnn
(79)
which concludes the proof.
E.2 PROOF OF COLLAPSE
Recall that the regularized risk associated with the network hww is defined by
R(W,U) = Ro U) + (I 1101)))
(80)
and recall that the set of weights in type-I collapse configuration is 2 { (W,U) U): There exist equiangular matrix 3 such that
W = c and
Ó =
3 2 KL
The
(81)
This subsection devoted the proof of the following proposition. Proposition A. We have the following:
(i) the latent variables .......% Z.K mutually distinct and satisfy assumption A then there exists e E such that
Q. C arg min R
(ii) If the latent variables ZI..... ZK are mutually distinct and satisfy assumptions and then any (W, U) that minimizes R. must belong Di. for some E RR.
31

---

Published as a conference paper at ICLR 2024
This proposition states that, under appropriate symmetry assumption, the weights of the network hw. do collapse into a type-I configuration. This proposition however does not provide the value of the constant involved in the collapse. Determining this constant will be done in the subsection
We start with simple lemma. Lemma G. Any global minimizer of 100 must belong to N.
Proof. Let (W*,U* be global minimizer. Define = K -k=1 U and
Un = iU; B U2 BB ... Uk - B]
From the definition of the unregularized risk we have RW ;U)) Ro Ro(W":U²) and therefore
K
1 (R(W";U)) R(W";U*))
A
(1UE - B1I - 1Ug1T
K
2
K
k=1
K
A
(11B1² - (B,UE)))
2 K
R=1
A
2 B, 1 x )
1
||
2
1111
So B must be equal to zero, otherwise we would have R(W", Uo) < R(W*.U")
The next lemma bring together the bilinear optimization problem from subsection E.1 and the sharp lower bound on the unregularized risk that we derived in section D
Lemma H. Assume the latent variables satisfy assumption A Assume also that (W*, U*) is global minimizer of (80 and let e RR be such that
(1IW 11 + I0* = =e² n 2
Then the following hold: (i) Any (W, U) that belongs to N i n.n. n. is also global minimizer of (80
(ii) If NENDE + 0, then (W*,U0 U*) must belong to N E
Proof. Recall from theorem that
(W U)
(U,WQzZ),)
for all (W,U) ENnE
(82)
Ro (2,,UU o( (-(U,
for all (W,U) e ge
(83)
We start by proving (i). If
By then we have
Ro U*) 29
( ('',
[because (W*,U") E N due to lemma G
29 -- (Ū, wQ" 2)
[because (W,U) E By and g is increasing]
= UW,
[because (W,U) N (E)
Since (W,U) E By we must have 1 (II || 1111)) o² ng = (IIW* -| + [1U) There-
fore RWW.)) < R(W", and (W,U) is minimizer.
We now prove (ii) by contradiction. Suppose that WW,U") &WWNEEES This must mean that (W* U²) E
since it clearly belongs to W. If ( U*) & E then the first inequality in the above computation is strict according to (83 If (W* U*) & B! then the second inequality is strict because is strictly increasing.
32

---

Published as a conference paper at ICLR 2024
The above lemma establishes connections between the set of minimizers of the risk and the set NWW n B. The next two lemmas shows that the set N.W. n B. is closely related to the set of collapsed configurations 22 In other words of minimizers and the set of type-l collapse configurations.
the set & n. n n. as bridge between the set
Lemma 1. If the latent variables satisfy the symmetry assumption then
2 CENNOB
Proof. We already know from (67) that 22 B We show that D E. Suppose (W,U) 0% Then there exists an equiangular matrix 8 E Iddx such that
W 3 P
and
ú=d32
where c' -CV nw/ L)... Recall from 25 that
P((x) Zs
for all x E Xz.
Consider two latent variables
Zi (0i,..., and 2j [o [o a'p
and assume x is generated by 2k. meaning that x Xx. We then have
Nw.u(x (Üx Ú, [((x))
=cd (3 Zv - 3 2 P((x))
cd (8 Zp 3 2 Zu)
cd' Σ
for faas fire
t=1
cd
L
Σ
faz Far )p)
Since fi...., ......,,n Inn. are equiangular, we have
Σ faz fae = -(6 dist(2z,2x)) 24))
dist dist(zz, (2,;))))))))) 2x) = L nr
dist (eott(Z,,22 2z)-
nc -
Therefore
Dw.u(x,1) = cd. The The
dist(zz; 2x)
and it clear that the margin only depends on dist(zy, zx), and therefore (W,U) U) satisfies the equimargin property. Finally we show that S2d C N. Suppose (W,U) E St. From property (ii) lemma B we have
K
Z
1n
17
The
Therefore,
K
K
K
Ü
d'
3 7/p
d'
8 n, 1 =
ne
1=1
where we have used the fact that 3 1,, = 0. Lemma J. If the latent variables satisfy assumptions 1 and then
2 EnNNB!
33

---

Published as a conference paper at ICLR 2024
Proof. From the previous lemma know that Di C E &NNNB need to show that
&NNnE, C 2
Let (W,U) EENNN n By Since (W,U) belongs to By there exists matrix F Rd with | ne such that
W c P and
U = FZ
(84)
where d' = cv Rw//l (KL. Our goal is show that F is equiangular, meaning that it satisfies the two relations
1 1 ne
F1 0 and FTF
Inc
(85)
The
N1e
The first relation is easily obtained. Indeed, using the fact that (W,U) E W together with the identity
Zs
De 1n 1 (which hold due lemma B
K.
we obtain
K
0
Uk
Σ
FZz = -
F1 F1 1.
A==0
k=
We then note that the matrix Fln,, 17 is the zero matrix if and only if Fl,, = 0.
We now prove the second equality of (85 Assume that x Ng. Using the fact that P((x) Zk together with (84), we obtain
Mw.
(Üx Üg. W((x))
=cd (F Ex - F En, P((x))
cd (F ZB F Zzj Zk)
= =cd (F" F(Zx - Zz), Za)
= =cd' FTF, ((k.)) )
(86)
We recall that the matrices
(k.)) = Zx(Zx - 2,)² R nnee
are precisely the ones involved in the statement of assumption B Since (W,U) E E, the margins must only depend on the distance between the latent variables. Due to (86). we can be express this
as
FFFF, r(skk) - (,"FF,
V',kj''' E [K] s.t. dist 2x) = dist (zy', Zxr)
Since the FT F. is clearly positive semi-definite, we may then use assumption conclude that FTF F A. Recalling definition (42) of the set A, we therefore have
FTF In. + 1n. 1 for some a, b E RR. To conclude our proof, we need to show that
(87)
The and
(88)
The
Combining 87 with the first equality of (85 we obtain
0 F In. = In, + 1n, 1s 1n, = (a + bne)
(89)
Combining with the fact that 1 ne. we obtain
The | Tr(FTF) F) na bb)
(90)
The constants a, b IR according to
89 and 90 must therefore solve the system
a + bn = a la+b =
and easily check that the solution of this system is precisely given by 88
34

---

Published as a conference paper at ICLR 2024
We conlude this subsection by proving proposition A
Proof of Proposition A
Let (W*, U*) be global minimizer of R. and let e E IR be such that
1 [I + 11
c² nww
2
If the latent variables satisfies assumption
A we can use lemma ttogethr with the first statement of
lemma to obtain
D c CENNNBA By C arg min R,
which is precisely statement (i) of the proposition.
We now prove statement (ii) of the proposition. If the latent variables satisfies assumption and then lemma asserts that
2 EENNE The set S2/ is clearly not empty (because the set of equiangular matrices is not empty), and we may therefore use the second statement of lemma Ho to obtain that
(W* U*) E &NNNB 22
E.3 DETERMINING THE CONSTANT e
The next lemma provides explicit formula for the regularized risk of network whose weights are in type-I collapse configuration with constant c.
Lemma K. Assume the latent variables satisfy assumption to 2% then
A If the pair of weights (W.U) belongs
R(W,U)= = log (1-
K 1
(ne - 1)e mmnee))) + A
(1-
+
I nhe
1
(91)
where 1)
KL
From the above lemma is clear that if the pair (W,U) E S2. minimizes R. then the constant minimize the right hand side of (9T). Therefore combining lemma K with proposition concludes the proof of theorem
Remark In the previous subsections, we only relied on relations (i), (ii) and (iii) of lemma B to prove collapse. Assumption A was never fully needed. this section however, in order to determine the specific values of the constant involved in the collapse, we will need the actual combinatorial values provided by assumption A
The remainder of this section is devoted to the proof of lemma K
Proof of lemma K Recall from 45
that the unregularized risk can be expressed as
Ro(,,
Σ
log
1+ Σ
Des (x)
J
We also recall that the set Q2! is given by
22 = { (W,U). U): There exist equiangular matrix 3 such that
Phw W = c and Ú=c
(92)
KL
and that PÇ(x) Zk for all E Xs (see equation 25 from section B Consider two latent variables
2g [ai,................... -0,,...,,..) az.)
and nj [o ... a'r.)
35

---

Published as a conference paper at ICLR 2024
and assume x is generated by zy. meaning that Xx. Mw. (x,j)
{Üx Ü, Wc(())
c2
Nhw
(3 7x
3 7 P((x))
KL.
c²
Tho
( Zx
3 Zz. Zk) F
KL
0²
Thw
fas
logs far
KL
1
e²
The KL
L
fa!; faie )F)
V
Since fi.....fee.............. fn. are equiangular, we have
fot fas )p
( 1i-dist(zj, zx))
dist(zj. 2k) L
ne
dist Zz. zz).
no
The
Therefore
Nw.o (x,j) =²
Tue Te KL ne
dist(zz. 2R)
Letting
May nr KL n
we therefore obtain
1
Ro(W,U)
log
1+ 1+
Dec (x)
K
1= XEX,
jfh
log
1
Σ
we dis)
(93)
jæk
where we have used the quantity inside the log does not depends on x. We proved in section equation (38) that if the latent variables satisfy assumption A then
(see
K
(
L
S. = r
(me
1)"
no
Using this identity we obtain
Σ
deidis((e,
{j dist(zj, 2x) = r}] e weer
K
(ne -1)- -weer
n
K
L
(ne 1) e- -wer
n!.
r=0
K
11 (ne - 1)e ve-wrr)
n
n!
where we have used the binomial theorem to obtain the last equality. The above quantity does not depends on k therefore (33 can be expressed as
K nt
K ni
{1+
(n. 1)e 1)e- )e)
Ro( log
We then remark that the matrix 3 P has nw columns, and that each of these columns has norm 1. Similarly, the Z has KL columns of length 1. We therefore have
1 N/
+ + - (21 PI + à The
13 2122 = 2n
KL
To conclude the proof simply remark that nn
36

---

Published as a conference paper at ICLR 2024
F PROOF OF THEOREM 3 AND ITS CONVERSE
this section we prove theorem 3 under assumption A and its converse under assumptions and We start by recalling the definition of type-ll collapse configuration.
Definition E (Type-Il Collapse). The weights (W, U) of the network lww form a type-ll collapse configuration and only if the conditions
i) (o) Va for all (). E.
ii) There exists c 2 so that ug,E = cfa for all (k, satisfying zn,e - a and all hold for some collection f1. f....., Inne E RR of mean-zero equiangular vectors.
As in the previous section we will reformulate the above definition using matrix notations. Toward this aim we make the following definition:
Definition F. (Mean Zero Equiangular Matrices) A matrix 3 E Bdxx is said to be a mean-zero equiangular matrix if and only if the relutions
1 8 =
8 BIn =0
and
Inc
1
1
Ne
The
hold.
Comparing the above definition with the definition of equiangular vectors provided in the main paper, we easily see that & is mean-zero equiangular matrix if and only if its columns are mean- zero equiangular vectors. Relations (i) and (ii) of definition Fan be conveniently expressed as
((W) 3 and
c= Z
for some equiangular matrix 3. We then set
22 = { (W.U) U) There exist mean-zero equiangular matrix 3 such that
* ((W) = Vd and -=c3 2
(94)
and note that ! is simply the set of weights (W, U) which are in type-II collapse configuration. We now state the main theorem of this section.
Theorem F. Assume the non con-degenerate condition μ4 minimizer of the strictly convex function
holds. Let 0 denote the unique
(t) = =log (1 (1-
K (1- n
+
II
K
(1+e
+
tt 2
where 17 = The
Te -1 v LKL
and let c TLVKL. Then we have the following:
(i) If the latent variables Zı,... ZKK are mutually distinct and satisfy assumption A then
Ω C arg min R
(ii) If the latent variables 2,,.... ZKK are mutually distinct and satisfy assumptions and then
Qfd arg min R
Note that statement (i) corresponds to theorem of the main paper, whereas statement (ii) can be viewed as its converse. To prove F we will follow the same steps than the previous section. The main difference occurs in the study of the bilinear problem, We will assume
will see in the next subsection.
pg everywhere this section - all lemmas and propositions are proven under this assumption, even when not explicitly stated.
Before to go deeper in our study let us state a very simple lemma that expresses the regularized risk R" associated with network h in term of the function Ro defined by equation 44
37

---

Published as a conference paper at ICLR 2024
Lemma L. Given pair of weights (W, U), we have
R² (W,WU) = Ro ((W), 17 + 1101
(95)
Proof. Recall from section B that
hw. (x) U Vec (W{())
( U Vec (p (W((x))]
Note that since ((a, () is hot vector, obviously have that up (W((a, + (W) <(a, 3).
Therefore the the network h" and are related as follow:
h (x) U Vec ( (W((x))] U Vee Vee(y((W))
As consequence, the regularized risk associated with the network hig can be expressed as
1
K
(())),(()))))) + X 1101/
RTWWUU) =
Σ
Ex
K
R=1
1
((hhp((w),) k)] - 1
Ex
K
Ro( (0 (W), U) +)+ 2 1111
where Ro is the unregularized risk defined in (44)
F.1
THE BILINEAR OPTIMIZATION PROBLEM
Let
Range(y) = (V E d There exist W E Rdxx= such that V= =WW)W)
and consider the optimization problem
maximize (U.Vo*z)
(96)
subject V Range(p) and - KLC
(97)
where the optimization variables are the matrix V Rdx and the matrix Ó E RaxKL.
Lemma M. Assume the latent variables satisfy assumption A Then (V,U) is a solution of the optimization problem 96 if and only if it belongs the set
I.ll = (V.U): There exist matrix F E F such that V VdFP and i FZ C
(98)
where denotes the set of matrices whose columns have unit length and mean zero, that
== [ddx..
12 F = and the columns of have unit length).
The remainder of this subsection devoted the proof the above lemma.
We start by showing that all (V,U) E !!! have same objective values and satisfy the constraints.
Claim D. If (V.U) E B! then
| P = KLC,, and (0.vorz),- eVdKL =oadk
V E Range({p)
Proof. Assume (V,U) E B. Since the columns P are one hot vectors in "R, the columns of FP have unit length and mean zero. Therefore the columns of V have norm equal to vd and mean Therefore V E Range(y).
38

---

Published as a conference paper at ICLR 2024
Using Zz" KL I from lemma B
together with the fact that || | ne since its columns have
ne unit length, obtain
KL
FZ (F = =FZZZ, =
F
KL
(99)
The
As consequence we have || o² KL. Finally, note that
PQ
Ince
as can clearly be seen from formulas 23 and
26
We therefore have
(Ü, (ū.vQTz) - - )
We then prove that
Claim E. (V, is a solution of 96) 977. then (V,Ú) B
Note that according the first claim, all (V.Ū) E B have same objective value, and therefore, according to the above claim, they must all be maximizer. As consequence, proving the above claim will conclude the proof of lemma M
Proof of the claim. Maximizing 96)
97 over Ú first gives
VQ ||Q
i cVKL
(100)
and therefore the optimization problem reduces to
maximize subject
IIV 2 V E Range(p)
(101)
(102)
Using the fact that ZZT I we then get
KL
|IQ 21 = (VQ" Z.Q"=.VQQ=' = (VQ" 22" " ne
IIQ 1/
(103)
and so the problem further reduces to
maximize IIQ''1 subject to V E Rango(p)
(104)
(105)
Let us define
Va V((a, V((a,,)
In other words ((a is the jth column of V, where j (a 1)s, B. The KKT conditions for the optimization problem (104) - (I05) then amount to solving the system
VQ" VD, + 1a X
(106)
(V(o,)), 1a)
for all (o, 3) [| 2 for all (a. 3) EV
(107)
(108)
for D, some X diagonal matrix of Lagrange multipliers for the constraint (108) and À R"
vector of Lagrange multipliers for the mean zero constraints. Left multiplying the first equation by 1f and using the second shows A = Or and so it proves equivalent to find solutions of the reduced system
VQ VD
(109)
(V(a,)), 1s)= = 0
for all (a, 3) E v for all (a,
(110)
|; 2 d
(111)
39

---

fo Published as a conference paper at ICLR 2024
instead. Recalling the identity QÇ(a, B) p4ca (see 27 in section B obtain
"'Q ng Q" ea
and so right multiplying (109 by ((a,3) 3) gives
VQ" e, v(a, 3)
(a,
for all (a, EV
M4
where we have denoted by v( B) the Lagrange multiplier corresponding to the constraint Define the support
E. = {B [sc] v(a, B) *0} and E: fa E * 0} of the Lagrange multipliers.
then imposing the norm constraint IIV e.l v(a, B) Vd.
gives
up
and so IVQ" eall > af E since v(a, 9) for some 3 E [8e] by definition. This implies that the relation
VQ" ea
8) Vd
for all
(a, 3) E x [sc]
IVQ" eall
must hold. As consequence there exist mean-zero, unit length vectors f1, ..... fn, (namely the normalized VQ" e.) so that
V(a, 8) Vd
holds for all pairs (a, 3) with a E E. Taking a look at (26). we casily see that its th row of the matrix Q can be written Q" Eg ( (), and therefore
VQ ea = Σ
M3V 3)
Σ ( = Vdr, SE(se]
Σ
vd
SE(x,]
BE|xc]
holds as well. a & E then VQ" ca = 0 since the corresponding Lagrange multiplier vanishes. therefore follows that
IIQQ
IIQQ"e. = Σ ||fa|1
d
|=|
nE(ne]
aEE
and so global maximizers of (104) zero, unit-length vectors fj .... L. so that
105 must have full support. In other words, there exist mean-
V(a Vdl..
(112)
holds. Equivalently V = Vd for some F F. We then recover Ô using (100).
= KVLL VQ
FPQ Z
FZ
= KL
[FPQ KL
(113)
|IQQ=Z||
|F
where we have used the fact that PQ" In.. To conclude the proof, use the fact [[
LKL, was shown in 99
F.2 PROOF OF COLLAPSE
Recall from lemma
that the regularized risk associated with the network hiv. can be expressed
(Ro()(W) R² R"(W,U) = Ro ((W), U + L
(114)
and recall that the set of weights in type-Il collapse configuration is 2 = { (W,U) U) There exist mean-zero equiangular matrix 3 such that
((W) vdzP P and Ú=c3zi 2 (115)
This subsection is devoted to the proof of the following proposition.
40

---

Published as a conference paper at ICLR 2024
Proposition B.
We have the following:
(i) the latent variables Z1..... KK are mutually distinct and satisfy assumption A then there exists E I. such that
D C arg min R
(ii) If the latent variables Zi....,. ZK are mutually distinct and satisfy assumptions and then any (W, U) that minimizes R² must belong Didd for some e IR
As in the previous section, we have the following lemma. Lemma N. Any global minimizer
(14 must belong N.
The proof is identical to the proof of lemma G The next lemma bring together the bilinear opti- mization problem from subsection F.I and the sharp lower bound the unregularized risk that derived in section
Lemma O. Assume the latent variables satisfy assumption A Assume also that (W*, U**) is global minimizer of (T14) and let E R. be such that
|U01 = KLC²
The the following hold:
(i) Any (W, U) that satisfies
(c))), E ENEnB
is also global minimizer of R
(ii) WWNDEB B + 0, then
((W"),U")NNEEE
Proof. Recall from theorem D that
(U.vQ" 2 2),)
Ro(V,U)= -
for all (V,U) NNE for all (V,U) e Nn
(116)
Ro(V.U) (0
(117)
We start by proving (i). Define V+ ((W*), and assume that V, W are such that (oW) = V
and (V,U) E NE n B.. Then have
Ro(s(W"))) Ro(V*,U"
2 g (VQ)))
[because (V*,0*) E
2 g (U,VQZ)
[because (V,U) E B
Ro(V,U)
[because (V,U) E [N
Ro(((WW)U))
Since | KL2 J have RW(WUU) RD (W* (U*) and therefore (W, U) is minimizer.
We now prove (ii) by contradiction. Suppose that (((W*)), U*) Nn'nB!'. This must mean that
(p(W"), U*) &EN
since it clearly belongs to M. If (pW(")) ",U*") U²) & E then the first inequality in the above computation strict according to 117 If (P(W)), UT) B!f then the second inequality is strict because is strictly increasing.
The next two lemmas shows that the set E A n B!! is closely related to the of collapsed configurations D2! In order to states these lemmas, the following definition will prove convenient
" = {V(VU) There exist mean-zero equiangular matrix 3 such that
V = Vd3P P and Ú=c 2 (118)
41

---

Published as a conference paper at ICLR 2024
Note that (W,U) Q and only if (v(W), D.
Also, in light of (98 the inclusion
2 11
is obvious. We now prove the following lemma.
Lemma P. If the latent variables satisfy the symmetry assumption A then
T CENNOB
Proof. The proof is almost identical to the one of lemma ! We repeat it for completeness. We Jf
already know that 22 C B We the show that S C E. Suppose (V,U) E S Then there exists mean-zero equiangular matrix 3 E Rdx such that
V vdzp P
and = Z
Recall from 25 that P((x) = Zμ for all x E Kt. Consider two latent variables
Zi ..,,..., asl and 2j [o', ./z) and assume x is generated by Zg. meaning that Xx. We then have
M = (Üx 19--- (Ux=-,,((xx) V(x)
cvdd (3 Zv 3 13,3 P((x))
=cvd (3 Zi 3 Z Zn)
=Vd Σ
fax laz far )
= =dvd dist 2x))
From the above computation is clear that the margin only depends on dist(z,, zς), and therefore (V, U' satisfies the equimargin property.
Finally we show that 2 C. Suppose (V,U) E s²" Using the identity Σ Zz we obtain
K 1 no
I
K
K
K -8 3 1 17 0
Σ
UR
3 Zs
The
1=11
1=1
where we have used the fact that 3 Ine = 0.
Finally, we have the following lemma.
Lemma Q. If the latent variables satisfy assumptions A and
then
T.' = EENN
Proof. The proof, again, is very similar to the one of lemma D From the previous lemma we know that Ω C CENNDB! !! so need to show that
EnN n B c.
Let (V.U) e EnNN B Since (V, belongs to B! there exists a matrix F gdxne whose columns have unit length and mean 0 such that
Vd and U=cFz
Our goal is to show that F is mean-zero equiangular matrix, meaning that it satisfies the three relations
1 F = O.
F Fl,, 0
and
FT
I
1n IT
(119)
nhe ne
42

---

Published as a conference paper at ICLR 2024
We already know that the first relation is satisfied since the columns of F have mean 0. The sec- ond relation is easily obtained. Indeed, using the fact that (V,U) E together with the identity Ex= Z me
K 1,, 17, (which hold due to lemma B we obtain
K
K FZR
K F1,
U
1,
ne
k=
k=0
which implies Fln, 0.
We now prove the third equality of 119 together with (84) we obtain
Assume that E Xg. Using the fact that P((x) Z
Nyu (x,
((,x())))()
=cvd (F Zz Zz. F P((x)) =cvd (F Za 1F -F, Zi)
=cvd (FT F(4 - -7,)) 7x)
=cvd FTF, ((k.)
(120)
Since (V,U) E the margins must only depend on the distance between the latent variables. Due to (20). we can be express this as
FFFF, r(ik))
- ( (
Vj,k,j' E [K] s.t. dist 2n. = dist(Z,7 Zer)
Since the FT FTF A. Recalling definition (42) of the set A, we therefore have
F is clearly positive semi-definite, we may then use assumption B to conclude that
FTF In. + In. 1
(121)
for some a, IR. To conclude our proof, need to show that
The
and
(122)
The
Combining 121 with the first equality of 119.. we obtain
0 FT F In. = I,, + 1,, 1T, Inc = (a + b
Since the columns of F' have unit length, the diagonal entries of FT F. must all be equal to and therefore (1T) implies that a + = 1. The constants a, E R, according must therefore solve the system
a + bne a aa+b
bn
and one can easily check that the solution of this system is precisely given by 122
We conlude this subsection by proving proposition
Proof of Proposition B Let (W*,U*) be global minimizer of R and let e E IR be such that
|| KL
We first prove statement (i) of the proposition. If the latent variables satisfies assumption A then lemma asserts that
T c CENNBB"
Assume (W.U) e Q. This implies that (v(W), U) E D.d and and therefore (+(W). U) EnN n !li We can then use lemma to conclude that (W, U) is global minimizer of R*
43

---

Published as a conference paper at ICLR 2024
We now prove statement (ii) of the proposition. If the latent variables satisfies assumption A and then lemma asserts that
Ted = EENNNE
The set Si is clearly empty (because the set of mean-zero equiangular matrices empty), and we may therefore use the second statement of lemma Co obtain that
(g(W"))U" E &nWN which in turn implies (W* U*) E
52
F.3 DETERMINING THE CONSTANT
The next lemma provides an explicit formula for the regularized risk of network hww whose weights are in type-Il collapse configuration with constant c.
Lemma R. Assume the latent variables satisfy assumption D then
A If the pair weights (W, U) belongs
K n
K (1 (ne 1)e 1)e-" VKK., n
(VKL.c)
R' R'(WUU) log
1
+
(123)
where 1)*
d
ne
KL
Combining lemma with proposition B concludes the proof of theorem F
Proof of lemma R
We recall that
Ro(W,U)
Σ
Σ
log 1+ Σ
Bw.o(kx))
Das (x)
XEX2
and
S2 = { (W,U) U): There exist mean-zero equiangular matrix 3 such that
<(W) = Vd $ P and 1 0=c82) 2) (124)
Consider two latent variables
Zi [a1, (a.,...,.. .0i........... o. and 2, - a')]
and assume x E 24. Using the identity PC(x) = ZB we then obtain
M (x()),xxxx)) (Üx -U)+((W((xx))
-ovd (3 Z - 3 2 P((x))
cvd (8 2z - 3 3 Zk) r
cvd d Σ
Jar - loy far )
8=1
ovd
L-> fass Fece )
L
E=1
cvd n
dist(z,,, 2x))
no
Letting w" Vd
therefore obtain
n--I
K
Ro(W,U) U)
1
log
1+ Σ e -ow dist(z, Dzn (x)
K
XEX
jzk
log
1 ++
é
-ew" dist(
(125)
K
k=1
44

---

Published as a conference paper at ICLR 2024
where we have used the quantity inside the log does not depends Re (+) (ne 1)" then obtain obtain
Using the identity |S,]
Σ
| dist(zz, r]] e-ww"r
j&k
K
(ne -1)² c c-wwr
n
K
K
1. ()
(ne - 1)" e-wwrr
T
1
K
K
(ne 1)e-ow))
n
n
where we have used the binomial theorem to obtain the last equality. The above quantity does not depends k, therefore (125) can be expressed
K K n
(1+ (ne -)ee =0))))
(W U) log
n.
We then remark that the matrix 3 Z has KL columns, and that each of these columns has norm 1. We therefore have
|
|c 2 2 KL
for all (W,U) 2
To conclude the proof we simply note that ~- VKLn"
G PROOF OF THEOREM 2
This section is devoted to the proof of theorem P from the main paper, which we recall below for convenience.
Theorem 2 (Directional Collapse of h). Assume K = n and {z,..... zK} Z. Assume also that the regularization parameter À satisfies
(126)
Thc
Finally, assume that (W, U) is in a type-lll collapse configuration for some constants C,T1,-..,," Ts. O. Then (W, U) is critical point of R. if and only if (c,ri,..., vx.) solve the system
A r3
)) 41 No 1
ne 14 exp
CT3)
for all 1 1SB5 Sc
(127)
L
2
Tmm
(128)
3=1
At the end of this section, also show that if 149
holds, then the system (150)
151 has
unique solution (see proposition Di in subsection G.2)
simply need evaluate the gradient of the risk on weights (W, U) which are in a type-III collapse configuration. Setting this gradient to zero will then lead system for the constants c, G. defining the configuration. While conceptually simple, the gradient computation is quite lengthy.
The strategy prove theorem 2 is straightforward:
We start by deriving formulas for the partial derivatives of Ro with respect the linear weights u and the word embeddings (a.9)) As we will see, ORo/ /uu and @R//wwaa.e) roles, In order to observe this symmetry, the following notation will prove convenient:
(a,(a plays symmetric
K
$ (W,
(m==(.,0)) (IU=K) - (.w:(()) Dz, (x)
(129)
1=1 x
45

---

Published as a conference paper at ICLR 2024
where
(Ox,W)
W. (x)
Lk'=1 ((Ux)W((x))
R
We may now state the first lemma of this section:
Lemma S. The partial derivatives of Ry with respect to Ug, and (a,3) are given by
ÖRa
(W,U)
(a..6),(3.)) W(a,8)
a=1 B=1
8Ro
(W,U) Σ
(0.3),00,,) ug,e
(a,
Proof. Given K matrices Vi,..., VK E RRowXLL we define
f(V. VK)
ΣΣD
1
Σ E
(VVi, ..... (Vx, ((x)) :k) Dan (x)
K
k= xxXx
where "(y......'''' k) is the cross entropy loss
exp (yn)
fy. .......:. .y.......ki... k)
log
kkrm exp (yW))
The partial derivative of f with respect to the matrix V, can easily be found to be
0f av, (V1 VK)
e (V)c((x))
1(==)
((x) Des (x)
(130)
_k=re c(Vui<(x)))))
xxE
We then recall from 19 that the hth entry of the vector y = Nw.u is
yk (Üx, W (((()) = <WWTUr {(x))
and so the unregularized risk can be expressed in term of the function f:
R)(WU 1 K K
Σ Σ ((
(W) Dzs (x)
x=IxEX,
= f((WT Ü...., (T(x)
The chain rule then gives
URa (W,U)
K Σ
0f (WTU,.... ,.TU)) )
Ü,
(131)
aw
TV
j=1
0Ru
Of oo,
( Ü ................................ wT
(W,U) W
(132)
aú,
Using formula 130 for 05/0v, and the notation
(
(x..u(())
Σ c(WTUpr((()))
we can express
131 and 132
as follow
OR. aw
(WU)
K
K
-lw,c(()) ((x) Dan (xx)
K
1=1
k=
1xx
RRo
K
(1(2 - -...0x)) ((x) (x)(x))
(W,)) = w
aü,
K
k=
XEE
46

---

Published as a conference paper at ICLR 2024
We now compute the partial derivative of Ro with respect uze- Let og E R be the eth basis vector. We then have
ORo Oujee
aRo ei
(W,U)=
=-
(W, U)
Σ Σ - (WC(x) oz) (x) k=1xEM
K
Recall from 13 that W{(x) the matrix that contains the d-dimensional embeddings of the words that constitute the sentence x E X. So W((x) et is simply the embedding of the 4th word of the sentence x, and we can write as
mee
WC(x) e,
Σ
()) (a,
We therefore have
ORo Ouje
k
(W,U) = K
Σ
Σ
(I(==)) -xw,(())
Day (x)
k= xEdx
B=1
(1(==) - (.wx((x)) (() (xx) W(a.8)
K
8=1
#= X
Re
(a..)... (W,U) ((..))
which is the desired formula.
We now compute the gradient with respect W(a.3)- Recalling that <(a, () is the one hot vector associate with word (a, B), have
ORo (W,U) = ORo
(W. ((a,8)
Ow OW(a,)) (a,
aw
1
Σ (1(==4) - (.mo(()) (Ü, ((x))((a,))) Dzo (x) k1
K
Recall that the 9th column of ((x) is the one-hot encoding of the Øth word the sentence x. There- fore, the bth entry of the vector ((x))((a,,)) B) E R² is given by the formula
((a,
s(a,))(
{
if xt (a, 8) otherwise
As a consequence
Ü, (()
Σ
()= Us
1=1
which leads to
ŌRo
1 K
K
(W,U)
Σ (1y=-4- -.w.o(x))
Σ
(e==(a,9))
Day (x)
(w
K
j=1
k=
XE
E=1
1 K
K
(1y==)- - (x)xx()) [x==(a,9))
Day (x) Uie
K
E=
k=IxEd'
K
(a.,))... uge
which is the desired formula.
47

---

Published as a conference paper at ICLR 2024
G.1 GRADIENT THE RISK FOR WEIGHTS IN TYPE- III COLLAPSE CONFIGURATION
In lemma S we computed the gradient of the risk for any possible weights (W,U) and for any possible latent variables ;1;-..;2K: In this section will derive formula for the gradi- ent when the weights are in type-Ill collapse configuration and when the latent variables satisfy {zi..... ................................ zk} Z. We start by recalling the definition of type- collapse configuration.
Definition G (Type-III Collapse). The weights (W,U) of the network fw.u form rype- collapse configuration if and only if
There exists positive scalars rg that W(,, = ra for all (a, V.
ii) There exists с 2 so that uy, = cia for all (k, (), satisfying 2 = a and all hold for some collection fi.....fng. E _Rd of equiangular vectors.
We also define the constant R and the sigmoid IR + RRa as follow:
and
(x) = 1 ve(t7))
(133)
The
The goal of this subsection is to prove the following proposition.
Proposition C. Suppose K = n! and (Zi,...,24 2x Z. If the weights (W, U) are in rype-
collapse configuration with constants c, "1....,,, Tsc 2 0, then
ORo (W,U)- du..
1 1
Σ pa ((rrs) 8-1
u
n!
0 RRo
(W.U) = (1+7)) pgo(ers)
ow (a.
ne
rg
((a,9)
Importantly, note that the above proposition states that θ Ro Oug,e and Us,e are aligned, and that 0 (R//W((a.))) and W(o.3) aligned.
We start by introducing some notations which will make these gradient computations easier. The latent variables 21...., ZK will be written
2₁
2k,1 2k.2
LK 2
where 1 1 Zn,E S nh. We remark that any sentence x generated by the latent variable 2₃ must be of the form
x =
(1) (E),L,(1)]
for some (B1,.... 3r.) [ne]4, and that this sentence has probability 118, "., M4,, of being sampled. In light of this, we make the following definitions. For every B (B1..., BL) E [m]]] we let
X [(=,1,, .... (ZK, (1.)] X
(134)
n|33 := m1831 ... n(3r,] E [0,
(135)
where we have used yu[3c] instead of 14,, in order to avoid the double subscript. With these definitions at hand we have that
Dz, (xx,g) [M[3] i
otherwise
We are now ready to prove proposition C We break the computation into four lemmas. The first one simply the notations that we just introduced in order express (a,,))(kk)) in convenient
format.
Lemma T. The quantity i (1ah,,") U) can be expressed as
K
① (.),)))) (W,U)
Σ 1 (88=39))
Σ ==1
1 (=)e=ar)) (x,,)0(X,,,)) 111)..
K
48

---

Published as a conference paper at ICLR 2024
Proof. Using the above notations, we rewrite (W, follow:
((t,))),(kk))W.. (W.U)
1
Σ x
(xt=(ar, (000))((())))))) ((i=)) (x.w,((x)) Dz, (x)
K
((s)e.8)) (1(3-8) (x Da, (x)
K BE
(zae=a)) [ (16=1) (x.,,.((33,00) p/g)
j=1 BE|U
Σ
Σ
(x (11j=1) - mmmotsss)) - m13)
K
BE(nE]
j=1
To conclude the proof we simply remark that Ej (xya=ar) (j=k] = (32,==0))
The following notation will be needed in our next lemma:
8(xx, a') 1 a
for a, a' [me]
(136)
-y i
where we recall that 7 = 1/(ng - 1). We think of 8(ax, a²) as 'biased Kroecker delta' on the concepts. Importantly, note that if fi.....ff. fn. are equiangular, then (fa. ò(a,a')
which is the motivation behind this definition. We may now state our second lemma. Lemma U. Assume K = n and {Zi,..., ,,,]] = 2. Assume also that the weights (W,U) are in a rype-Ill collapse configuration with constants c, r..,,.... 2 Then
M exp rB. 8(2),e,
(x),))(3),g))
where v(x)) e" + -
I (C((r)))
7
for all j,k E [K] and all B ($1,.... BL) E [mc]"
Proof. Recalling that xj,3 = [[2,,1 ..., (zj, BI.),, we obtain
(Ü. (((x),)))
Σ (use, w (=3..18)) Σ (ole r feg Σ rpe 8(zk,e 24,e)
We then have
(U,,,(())))))
exp C El Te, 8(zk,e ,a,e))
W. (x =
Σx= (Cc.
K
Σ exp ( (t The 8(zw ,,5,))
K
I exp (crp, 8(2k,e .2.,))
Σx - Il exp (erg, 8(#w.e =z.e)
Since (zi....,Z,, [ = Z. the latent variables Z = [ 2x,2] achieve all possible tuples Ja', ... a/y1 [_ The bottom term can therefore be expressed as
EI exp (cre, 8(zk,e (z,e) W=1e=1
exp (erg, 0(ah, 242)) ((b)))2)))
exp (crg) o(ab,
exp
0 0/
Re Σ exp (crs,6(ah (o 0/=1
49

---

Published as a conference paper at ICLR 2024
Recalling the definition of 8(a; a²), we find that
exp (cro, 2.)) exp(cr8.)+ + Σ exp (- CT
Σ af=1
Nhe
(-n cr =b(Crg)) (137) (-me-
exp(cr3,) (ne - 1) exp
1
We now find a convenient expression for the term appearing between parenthesis in the statement of lemma T
Lemma V. Assume K n and (zı...., 2x} = Z. Assume also that the weights (W, U) are in a type-lll collapse configuration with constants c, ........ Tso 2 Then
Σ [2y (x, - 5(zu,e, o(era)) j=1
(138)
for all k E [K],E [L], a* [nd] and all B (BI...., [no]"
Proof. For simplicity we are going to prove equation (38) in the case = 1. Using the previous
lemma we obtain
I exp (cr3, I (crr,)
Σ (e
[
Since the latent variables z, 1 2,, achieve all possible tuples [ai, ,aL E [nc]", we can rewrite the above as
Ie exp (crg,ōlae, =x.))
1 (m
I (((rr))
01
a1=10
02=1
01=11
exp (erp,0la ((r,, 22.1)) I1 exp (cra,((a,
The
Il (e(rg))
a==1
01==1
exp
crp, 5(a", 2x,1))
ne
ne
Σ
II exp (cra,b(ar,2u.))) 8=2
(139)
Il (CT
a1=1
We then note that
nee
nc
ne Σ exp ((TeD(aeszzee)) a/=1
II exp (aa,
01
0==1
and, repeating computation (37). , we find that
ne
exp ((ae, (((rg,)
Going back to 119
we therefore have
exp (cr
exp (ers,((at, =x.1))
(a ( (x),))
I (crr))) v(crp,)
1=1
and so
exp(era, M(rr,
K 1 (3), (x),,)
if 2k.1 a*
1 1
(140)
exp(-2or9)) v(cro))
if #k.1 + a*
1=1
50

---

Published as a conference paper at ICLR 2024
We now manipulate the above formula. Recalling that y 1/(nc 1). and recalling the definition of (z), we get
ex
ex e²
1
1 + 7el1-r1 = o(x)
(141)
(x)
1
? e
and
7a
yx
(11760ltt)))) -
-7
((x)
ex ee
which concludes the proof.
Our last lemma provides a formula for the quantity ((ar (W,U) when the weights are in type-III collapse configuration.
Lemma W. Assume K = ne and (21.-..,,2g = Z. Assume also that the weights (W, U) are in a type-Ill collapse configuration with constants c, 11, ................................ ..... 2 Then
$ (a.3)) (W,U) ((rp)) 0(
(142)
n
for all l E [K],e L E [L], a E [nd] and [sc]
Proof. Combining lemmas T and and recalling that K n we obtain
$ (1k,8)) (WUU) = Σ
1
[au.1=0)) (0==8+)
Σ
( (.m.(((s))) m13)
ni
BEln
1=1
Σ
1,,=801 (8(z)e)ae) o(ers)) u1s)
m BElNEl
)(E,,,,0*)
(0:=0 a(crg)
nf.
pE(n
Choosing { = 1 for simplicity we get
Se
Σ BE(ng]
1,,=8-0 (Crr)) n|
1 o(ers)) p1B1]a183) -. -.13r]]
Br.ml
(crgg) /185)11122 ... [1[11
1==1
9==1
(
which concludes the proof.
We now prove the proposition.
Proof of proposition obtain
Combining lemmas
S and W and using the fact that W(a.3) rgfas
ØRo Juge
The (W,U)
(0.9), W(..))
B=1
U3 a(era) o(ene) rgta
nod
a= 8=1
1
%e
ng a rg
Σ
a)
I
P=1
51

---

Published as a conference paper at ICLR 2024
Using the fact that Σo= fa 0 get
Σ a fa - fanne
Σ OF2
fane + 7 fec. - 7 Σ fa
lo - (1 +7) feu.
(143)
Using the fact that Ugee = CfaN. we then get
aRo Jus,
(W,U)
Σ
us (cTs)) no
(1+ fee.
The
B=1
1+9 7
Σ B=1
pe o(rrg) rB
uk.e
n
c
which is the desired formula.
Moving to the other gradient get
ORo
K L
(W.U)-
1 (k,0)) (W,0)) us.
(a.3))
k=
E=1
Mg n
alcra) o(En),)) c fine
K 8(zz, faue
14 (Crr)) e
Σ
n
Since the latent variables Zy = [2x,1 have, fixing f = 1 for simplicity,
1...
L achieve all possible tuples (a "af .ayl E [nd]", we
K
8(2 a) fee.: =
Σ a( fa; 1-1 Σ
8(a a) Jaj
(144)
x/=1
a;=1
Repeating computation 143) shows that the above is equal to n 1--1 (1 + 7) far We then use the fact that ((a rafa obtain
8Ro (a,8))
(W,U)
Uı ((rr)) cl n 11 (1+ fa
n
143 o(erg) cl (nt 1 +) W(a,8) (1 o(rr)) cL (1+ W( T3
n
149
nhe
which is the desired formula.
G.2 PROOF OF THE THEOREM AND STUDY OF THE NON-LINEAR SYSTEM
In this subsection we start by proving theorem 2 and then we show that the system (55)) has unique solution if the regularization parameter À is small enough.
151
Proof of theorem Recall that the regularized risk associated with the network hww is defined by
R(WUU) Ro U) +
(1|W +
(145)
= Ro, A fle
No
(
112
(146)
2=1
E=1
52

---

Published as a conference paper at ICLR 2024
and therefore (W, U' is critical points if and only if
0Ru
ORu (W,U) À ((a.9)
(W,U) À Ug.6
and
Ōug,e
dW(a.))
According to proposition C if (W.)) is in a type-lll collapse configuration, then the above equations becomes
1 1 1+ n
L(1 + p43 (CT())
Σ pe a(ers)
ux,t = À up, and
(o,88)
((,,))
ne
T3
So (W, U) is critical if and only if the constants r
Ts. and satisfy the 8, + equations
1 7
na a(cra)ra=
(147)
The
L(1 + u o(ers) =
for all 3 E [sc]
(148)
rg
From the second equation have that
B₂ r Le
(1+ ps ocrrs) rg
Using this we can rewrite the first equation as
1 9c
Th A r
=
which simplifies
8c
F8
Ln
Le
c
which is the desired equation (see We now rewrite the second equation as
151
A rg
T
L c (1+)) 4g
then recall that (x)) := T+yal1+5)) and therefore
The (crs)
The (1 + ye (tt))a))))) no - + exp
Tee
CT
+
11
and therefore the second equation can be written as
À r3
ne 1 exp ( The crg )) -a.
exp
L
me
We prove that if the regularization parameter À is small enough then the system has unique solution.
Proposition D. Assume μn 2 p2 2 ... 2 Mo.
and
L
Scc
(149)
L+1
NCc
B=1
Then the system 8.
equations
r3
Ne
))
L
ne 1 eex
CT3
113
for all 1 S 3 s
(150)
1h
2
In
(151)
3=1
has a unique solution (c. rs.)) E Rict! Moreover this solution satisfies r A r2 .... 2
Ts. 0.
53

---

Published as a conference paper at ICLR 2024
Proof. Letting Pa rg/c, the system is equivalent to
L
g(c,
4g Ane
for all B E [se]
(152)
Σ Am1
2 L
(153)
for the unknowns (c, P1, .2..... , where
Poc)
g(c,x) = x (1+ (1+ yel 7e(t)))e)) /(1+7)
and 7= 1/ - 1)
Note that
(c, x) 2 + ++ 7e(1_)) (1 (1+ +) 1 0g dr
V(c,x) E R x (0, +o0)
and therefore x + g(c, x) is strictly increasing on [0, +00). Also note that we have
0), 0, lim g(c,x) too
So x > g(.xz) is bijection from [0, + t00) to +00) as well bijection from (0, +o0) to (0, (x). Recall that My E (0, +oo) for all B E [sc]- Therefore given E R and , E [sc]. the equation
L
g x) =
Anc 1/8
has unique solution in (0, +oo) that we denote by ds(c). In other words, the function Φs(() is implicitly defined by
((()))(c) Ane uB.
(154)
Also, since g(0, = æ, we have
L
6s(()
An. 14
Claim F. The function 09 : [0, +00) (0, +00) continuous, strictly decreasing, and satisfies oo ()(C)= = O.
Proof. We first show that c++ (c(c) is continuous. Since (c, 2 for all x 2 0, we have
x2 0g
g(c,z2) g(c, 21) Jan dr
(c. z)dx
1dr x2 F1
for all and all xy Φ1 2 0.
As consequence, for all C1, C2, we have (ds(c2) - db((c)) I 1g(c1, dg(e2)) o((a))) [g(c,, ((l))))))=(((,,)((0))) g (s(c2))
(155)
where have used the fact that g(c1, ()(())) -e = g(e2, (c(2))) From 155 it is clear that the continuity of x() implies the continuity of @z(c).
We now prove that d3 is strictly decreasing on [0, +o)). Let 0 S e₁ < C2. Note that for any x >> 0,
the function g(c,x) x) is strietly increasing on (00) Since (c)
0 therefore have
J mg gecn, ds((e)) < ( (c())
g )n((2))
Ane
Since + g(c, .r) is strictly increasing for all the above implies that dh (c2) ds((i).
Finally we show that lim (c) = 0. Since thg decreasing and non-negative on [0, +00), the limo-stoc Φp(c) A is well defined. We obviously have ds(c) > A for all > 0. Since x + g(e, x)) is increasing we have
L
U3 = g(c, Φs(c)) 2 A)
An
But the function g(c, A) is unbounded for all A 0. Therefore must have A 0.
54

---

Published as a conference paper at ICLR 2024
System
152
153
is equivalent to
Pg Qg((c)
for all B [sc]
(156)
(dp(c)) = Lnn
(157)
0=1
Define the function
(cc) = (os(c))
B=1
Then & clearly inherits the properties of the 0s'ss is continuous, strictly decreasing, and satisfies
2
(())
148
and
lim e-+00
((c) = 0
Anc
D=1 1
Therefore, if
I.
2
L L-1
PB
An
B=1
then there is unique 2 satisfying (57 Since x + g(c,x) is increasing, equation 152
implies that the corresponding P3 's satisfy P1 2 P2 ... Ps. >
H No SPURIOUS LOCAL MINIMIZER FOR R(WU))
In this section we prove that if d min KL), then R(W, U) does not have spurious local minimizers; all local minimizers are global. To do this, we introduce the function
f: Rd +
define as follow. Any matrix E g can be partition into K submatrices Vi RdxL according
== IVi 1/2
Vxl
where V4 E RdXL
(158)
The function f is then defined by the formula
S(V) = 1 Σ Σ
((x.)()))) p... (Vxi<(x)) ;k) Des (x)
K
Am XEXi
where E(y)..... yK: k:) denotes the cross entropy loss
exp (yn)
".......y::: k) log
Et'm exp (yAr)
We remark that f is clearly convex and differentiable, We then recall from the vectory = fw., is
19 that the kthh entry of
yk = - ÜK (Us.
W w((x)) - ((x))
Recalling that Ú [0
Ük]. we then see that the risk can be expressed as
R(W, U) f(WT() À (1WIi++10 (
(159)
The fact that R(W,U, U) does not have spurious local minimizers come from the following general theorem.
55

---

Published as a conference paper at ICLR 2024
Theorem G. Let g: Rmxx 7 R be convex and differentiable function. Define
= :-
(1AII-- +BIF))
where A RRxx and B E Rdxx
2
and assume À and > min(m,n), Rdxm IR 1 + R is also a global minimizer.
n.. Then any local minimizer (A, B) of the function 4 :
The above theorem directly apply to 1.59 and shows that the risk R(W, U) does not have spurious local minimizers when
and min (.n,, L)).
The remainder of the section is devoted to the proof of theorem G We will follow the exact same
steps as Zhu et al. (2021), and provide the proof mostly for completeness (and also show how
the techniques from Zhu et al. (2021) apply to our case). Finally, we refer to Laurent & Brecht 2018 for a proof of theorem
the case X O.
Proof of theorem G To prove the theorem it suffices to assume that d > m without loss of general- ity. To see this, note that the function g(D) g(D") is also convex and differentiable and note that (A,, B) is a local minimum of
9(A B) + (|A + (I))))
if and only if it is local minimum of
"(B A) + 2 (|A| + |[
So the theorem for the case d > follows by appealing to the case d > with the function g-
So we may assume d > m. Following Zhu et al. (022)) we define the function tp mxxx - by
(D) = (D) ||
where || D|| denote the nuclear norm of D. We then have:
Claim G. For all E Rdxm and E Rd we have that (AAA B) S (4
Proof. This is direct consequence of the inequality
1
| |A" B|| S (11-11 [BI).))
-
that we reprove here for completeness. Let A UEVT the compact SVD of A" B. That is Σ Rexr, U E Rmxr V E Raxr, and r is the rank of A B. We then have
1
1
||AT B|| Tr(E) Tr(UTA AV (AU, BV) < (IAAI +IV
s (IA + +1B1))
Computing the derivatives of yp gives
0e (A, B) BB [V)(A" AA
0ip aB
and
(A, B) A V9(A B) + AB (160)
04
So (A, B) a critical point of jp if and only if
XA = B B))
(161)
AB = - (9(AT
(162)
Importantly, from the above we get
AAT BBT E dxd
(163)
56

---

Published as a conference paper at ICLR 2024
which implies that A and B have same singular values and same left singular vectors. Let U R be the orthonormal matrix containing the eigenvectors of AA" = BB". From this matrix we can construct SVD for both A and B:
A UEAVA and
B=
where ΣA E Rdxo and Eg Rdxx have the same singular values. From this get the SVD of AT
B² VVE
(164)
and it transparent that,
||A" B1| -| |
(165)
In particular this implies that if (A, B) is critical point of ip then we must have (AA, B)
B). This also implies that
(AVa(A (Va(A" B), - - -
(166)
Using this together with the fact that the nuclear norm is the dual of the operator norm, that is ||C]], = Gllls (G, C) ps we easily obtain:
Claim H. Suppose (A, B) is a critical point of ip which satisfies is a global minimizer of y.
Vg(AT op I X, then (A, B)
Proof. For any matrix C jpux have
||AT Bll. + 11 Vg(
1
B),
S sup (G,C) ||C|| G1
and therefore - Vg(A" E B|I- This implies that A is global min of s. The fact - that (oA, B) = w'AT B) (because (A, B) is critical point :)) together with Claim G then implies
that (A, B) is global minimizer of
We now show that all local min (A, of oo with ker(A") must satisfy (/V (A]BB)]. S X
Claim I. Suppose (A is critical point of yp which satisfies
(i) (er(AT) +
(ii) ()(ATBB]l >
Then (A, B) not local min.
Proof. We follow the computation from Zhu et al. 2021 Let (A, B) be critical point of Since According to (ii) these kernels are non trivial and we may choose unit vector z E Id that belongs to them. We then consider the perturbations
AAT = BBT, we must have that ker(A") = ker (AA) = ))) = -er(("")..
dA = za dB= = =bT
where a E m and b E R are unit vectors to be chosen later. Note that since and b are unit vectors we have ||dA|| = || = Moreover, the columns of dA and dB are clearly in the BT dB This implies that all the
kernel of AT and BT, therefore A dA ATd Bd dAA "cross terms' disappear when expanding the expression:
(A + EddA (B edB) A 2d dB = abb
We also have
|| + =dAl =1AA| + ||adA|
||A11/ + +
57

---

Published as a conference paper at ICLR 2024
and similarly, |B + edB| || + t?. We then get
(A + cdA, B cdB) = o (A + =dA)² (B edB)) + X
2 (II. edAIl + B + edBl
= 9 + A 2 (1A1) |BIb:) + X
A
[ (ab + 2 (|Al] + -11B1))) + >
= +(A, B) + 8²
O(e*)
Let G B) E mxx" We want to choose the unit vectors and b that makes (G.ab") ab"), as negative as possible. The best choice is to choose -a and b to be the first left and right singular vectors of G since this give the negative of the best rank-one approximation of G. So we choose E R.th and b E R" such that Gb a and therefore
(ab" T G)F Tr(ba' G) = Tr(a" = -0
which gives
c + EdA, B + EdB) = () Vo(A"B))l +2) + 0(5t)
and (ii) implies that (A, B) is not a local min.
Combining the previous two claims we can easily prove the theorem. Indeed, if m, then the kernel of A" is nontrivial and (i) is always satisfied. As a consequence, if (A, B) is a local min of sp, then Vg(A" B)llop B)(lo)))lo I X, and therefore (A, B) must be a global min of sp
58


Reviews:
[{"review_id": "OGYZLEw52O", "reviewer": "Reviewer_pmX6", "summary": "The paper delves into the concept of \"feature collapse,\" where entities with similar roles in a learning task are given analogous representations. To understand this, the authors use a synthetic task where a learner classifies 'sentences' made of L tokens. Their experiments reveal a direct relationship between feature collapse and generalization. They demonstrate that, with sufficient data, distinct tokens with the same task roles have identical feature representations in a network's first layer. The research conclusively proves that neural networks, when trained on this task and equipped with a LayerNorm module, develop interpretable and meaningful representations in their initial layer. The paper's key contributions include defining 'good features' mathematically and deriving analytical formulas for a two-layer network's weights.", "questions": "Questions:\n1. Maybe a simpler and more intuitive experiment to demonstrate this “feature collapse” behavior, is to assign two embeddings (to be learned) to one token, but during training this token is randomly mapped to these two embeddings, and see if these two embeddings become similar after training.\n2. Does the condition assumed on λ (in Thm. 2) align with the choice of λ in experiments? \n3. Can we empirically verify the feature collapse in deeper layers?\n4. I have concerns about the difficulty of learning equivalence from the data by the model itself without special treatments, including data augmentations and introducing explicit invariance into models. I think the root is: if the optimization objective is not the equivalence, the network could ignore the equivalence and just memorize the data. Another related and broader phenomenon is the failure to learn “A = B” equals to “B = A” in language models [1]. In computer vision, semantically equivalent inputs (pixels, small image patches, whole images, etc.) may not lead to the same feature, for both early layers and deep layers/final outputs. This is easy to understand, because even for contrastive learning where the network is forced to learn invariant features under augmentations, different views (augmented images) cannot guaranteed to lead to the same feature after training.\n\n[1] 'The Reversal Curse: LLMs trained on \"A is B\" fail to learn \"B is A\"' 2023.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strength:\n1. Paper is well organized\n2. The settings of experiments and assumptions of theoretical analysis are carefully designed.", "weaknesses": "Weakness:\n\nIn general, I am not sure how the results could be useful.\n* For the theory part, it would be helpful to highlight core technical challenges the authors met and addressed in the proofs, or any new proof techniques the authors developed.\n* For the experiment part, I am not sure how it could help design or explain applications. I totally agree that semantically similar instances would lead to similar representations. But is it benign or not for the network to assign exactly the same representation to equivalent instances? If it is benign, how to ensure it happens in practice? If not, then how to avoid it? It would be helpful for the authors to give some convincing scenarios that people care about this phenomenon.\n\nIn addition, I suppose there should be a section for conclusions and limits.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5JSfihLDeX", "reviewer": "Reviewer_nQiF", "summary": "The paper investigates feature collapse and \"good\" features learned at the first few layers. The good features should be invariant for entities serving the same concept for tasks at hand. The authors construct a synthesized task to investigate how feature collapse works. The paper reveals three kinds of feature collapse through synthetic experiments and theories and can be roughly summarized as follows: (1) Type-I feature collapse for models w/o LayerNorm, which collapses in both magnitudes and directions and will happen with uniform distribution. (2) Type-II feature collapse for models w/ LayerNorm will still collapse in both directions in magnitudes under long-tail distribution due to the help of LayerNorm. (3) Type-III feature collapse for models w/o LayerNorm, which collapses in directions but not magnitudes and will happen with long-tail distribution.", "questions": "Please see the weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper proposes a view to understand the feature collapse with good definitions of \"good\" features. The understanding of LayerNorm functionality in the feature collapse, which is a key component in the modern Transformer architecture, would help the community further understand language models.\n2. The paper's presentation is extremely clear and easy to follow. The synthetic experiments in Section 2 are sound, plausible, and offer reasonable intuitions.", "weaknesses": "My major concern is if the Type-II feature collapse still holds when LayerNorm has trainable parameters. The reason is that LayerNorm w/ trainable parameters will not force the word features to have a fixed magnitude. A simple experiment would suffice. \n\nWhile I tend to vote for acceptance, I am not an expert in this field, so I will keep my score modest.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m4LugWYh5S", "reviewer": "Reviewer_ZbKw", "summary": "The manuscript introduces the concept of \"Feature Collapse,\" which states that input features with analogous semantic meanings would receive identical representations. This phenomenon is examined under two distinct scenarios: the first being a 2-layer network, and the second incorporating LayerNorm on the 2-layer network. Empirical evaluations reveal that in a uniform setting, both network configurations attain feature collapse in terms of representations and weights. However, in a non-uniform context, the integration of LayerNorm becomes essential for the complete feature collapse. The study concludes with theoretical findings that prove feature collapse.", "questions": "1. The work devotes to study word representations, would similar phenomenon happens in vision as well (e.g., VisionTransformer uses image patches, does image patch with similar sematic also receive identical representations)? \n2. The work primarily studies that features of words with similar semantics would collapse, how about words that sharing similar syntactic in a sentence? Would something similar happen as well?\n\n(The reviewer is not asking for empirical evidence, just curious to hear the authors' thoughts.)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The work is well presented and all notions included in the paper are discussed clearly.\n2. The concept of feature collapse is interesting and intuitive.\n3. The paper potentially offers an explanation regarding the impact of LayerNorm.", "weaknesses": "1. While the reviewer appreciates the feature collapse concept and believes this phenomenon is not limited to the simple network settings used in the manuscript, the absence of empirical validation within more complex networks or datasets constrains the soundness and significance of the findings.\n2. The authors differentiate their work from the Neural Collapse (NC) concept by emphasizing that, unlike NC, Feature Collapse inherently suggests enhanced generalization and transfer learning capabilities. But once again, this claim lacks strong empirical evidence, the 2-layer network used in the paper is not comprehensive enough to demonstrate this point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JwhLYTFcqU", "reviewer": "Reviewer_xgCm", "summary": "This paper defines the notion of \"feature collapse\" which relates to the intuition that entities that play a similar role should receive similar representations. This notion is made precise for a toy NLP task. Here, labels are invariant to substitutions of input tokens within certain equivalence classes. The paper then studies, both theoretically and empirically, under what conditions tokens within such an equivalence class are assigned similar representations, i.e. \"feature collapse\" has occurred. The paper shows that for this toy task, feature collapse corresponds to stronger generalization in the presence of tail tokens. The paper also empirically shows that feature collapse only occurs when the network includes layer norm. Otherwise, the norm of representations depends on the frequency of the tokens in the training data. A major contribution is a theoretical proof that explains why the magnitudes of the learned representations depend on their frequencies.", "questions": "* Is there any setting where feature collapse may *not* be desirable? For example, consider a task with a significant degree of label noise and with limited data samples to learn from. Intuitively, features that are observed more frequently during training may be more reliable predictors of the label. Therefore, to the extent that the magnitude of word embeddings relates to their salience in making predictions, it seems that feature collapse may not be strictly desirable in such a setting.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The proof that token embeddings depend on token frequencies in the proposed setting may provide a useful theoretical insight for future work to build on. The proofs in the paper are non-trivial.\n* The paper may help explain a reason why layer normalization can improve generalization.", "weaknesses": "* While the key contributions of the paper are theoretical, the motivation for the theoretical results would have been clearer if a stronger connection could be made to more realistic settings. For example, the connection between feature collapse and generalization outside the context of the toy task does not seem to be well established. Additionally, the degree to which feature collapse sufficiently occurs or does not occur in the context of more realistic models and tasks is not clear. While perhaps not strictly necessary for such a theoretical paper, addressing such questions would have improved the motivation for this work.\n* The role of the L2 regularization term in equation 2 could be clarified. The authors state \"The regularization terms play no essential role apart from making proofs easier; the empirical picture remains the same without weight decay\". Could the authors include the experimental results that support this claim? This also seems to diminish (although not completely) the key theoretical contribution. It is not surprising that in the presence of L2 regularization, feature magnitudes would depend on their frequency in the training data. It would be more interesting if we had theory that elucidated why the learned embeddings for rarer tokens have lower magnitude in the absence of explicit regularization.\n\nMinor typos:\n\n* Intro paragraph 1 - \"gives same features\"\n* Intro contribution bullet 1 - \"how a network learn representations\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)