Title: Learning Partitions from Context
Venue: NeurIPS 2024
Accepted: yes
Keywords: clustering

gradient-flow

token embeddings

partitions

complexity

sample complexity
Abstract: In this paper, we study the problem of learning the structure of a discrete set of $N$ tokens based on their interactions with other tokens. We focus on a setting where the tokens can be partitioned into a small number of classes, and there exists a real-valued function $f$ defined on certain sets of tokens. This function, which captures the interactions between tokens, depends only on the class memberships of its arguments. The goal is to recover the class memberships of all tokens from a finite number of samples of $f$. We begin by analyzing this problem from both complexity-theoretic and information-theoretic viewpoints. We prove that it is NP-complete in general, and for random instances, we show that on the order of $N\ln(N)$ samples, implying very sparse interactions, suffice to identify the partition. We then investigate the conditions under which gradient flow dynamics of token embeddings can reveal the class structure, finding that this is achievable in certain settings when given on the order of $N^2\ln^2(N)$ samples.

Learning Partitions from Context

Abstract

In this paper, we study the problem of learning the structure of a discrete set of N tokens based on their interactions with other tokens. We focus on a setting where the tokens can be partitioned into a small number of classes, and there exists a real-valued function f defined on certain sets of tokens. This function, which captures the interactions between tokens, depends only on the class memberships of its arguments. The goal is to recover the class memberships of all tokens from a finite number of samples of f . We begin by analyzing this problem from both complexity-theoretic and information-theoretic viewpoints. We prove that it is NP-complete in general, and for random instances, we show that on the order of N ln( N ) samples, implying very sparse interactions, suffice to identify the partition. We then investigate the conditions under which gradient flow dynamics of token embeddings can reveal the class structure, finding that this is achievable in certain settings when given on the order of N 2 ln 2 ( N ) samples.

1 Introduction

Modern machine learning systems are able to learn extremely complicated relations from data. They often rely on learned embeddings of discrete tokens in a continuous space. This is notably true for Large Language Models (LLMs) [9, 17, 25, 23] which encode their input by converting text into a sequence of discrete tokens that are embedded in a high dimensional embedding space and those embeddings are fed into, e.g., a transformer architecture [28] which allows predicting the next token. But also in the other domains, e.g., in vision, discrete embeddings are frequently used as a component of deep learning architectures [27, 10] as this enables capturing complex concepts that are often discrete.

It was observed that after training, these word embeddings exhibit many interesting structures. The most prominent example probably is the observation that the difference of the Word2Vec embedding vectors of the nouns 'king' and 'queen' approximately equals the difference of the embeddings of 'man' and 'woman' [21, 22]. Similarly, it was found that using word similarity as an inductive bias to structure latent spaces helps downstream performance [4]. Thus, a properly structured latent space seems to be an important ingredient to capture the intricate correlations in complex data.

A proper theoretical understanding of such complex models currently remains an elusive goal. However, there have been various attempts to understand various components of deep learning models. Many works investigated the behavior of feedforward-networks in particular focusing on shallow networks [2, 8] and asymptotic regimes [15, 29]. More recently, several works investigated transformer architectures (often focusing on the one layer case with linearized attention mechanism) [16, 1, 30, 12]. On a technical level [11] is closely related as they study the large depth limit of transformers through the lens of particle systems (however in their case time corresponds to depth, while in our case it corresponds to training time). A feature shared by many of those works is that

little structure is assumed on the input, i.e., fixed token embeddings are assumed and often those are even assumed to be isotropic Gaussian.

Here we instead focus on the dynamics of the token embeddings and study how these can recover structure present in the data. There is no ground truth target for the embeddings for large scale models used in deep learning and their embeddings need to capture a variety of nuanced correlations and relations that are hard to formalize. Therefore, we focus on a simplified problem that nevertheless shares important features with more complex real world settings. One general heuristic is that the embeddings contain information about the similarity of tokens. We focus here on the strongest form of similarity, namely equality. Indeed, our central assumption is that tokens can be clustered in a small number of groups such that tokens within a cluster behave exactly the same, i.e., they interact in the same manner with other tokens.

Then the central questions is what assumptions allow us to recover a hidden structure. The crucial feature of this setting is that we only get information about a token by its interaction with other tokens about which we also only learn through their interaction behavior. Moreover, those interactions are typically sparse, i.e., we observe only a small subset of all possible interactions. Note that this setting resembles observations made in the context of collective behavior where a global structure emerges from local interactions [13, 24, 3]. A related question was investigated in [6] where they study associative memory and also want to identify a hidden structure, however, they learn the class memberships directly through the interaction with a class embedding (and they train the interaction instead of the embedding). In [19] the dynamics of word embeddings and transformer layers was investigated when the data follows a topic model. This work shares the crucial feature that membership of a word in a certain topic is only transmitted through the co-occurrence with other words from the same topic. In contrast to their work, we here do not focus on learning class membership from token frequencies and in fact consider uniform token distribution. Instead, we view this problem as a logical inference problem: Given a set of facts about a set of tokens can the hidden structure of the tokens discovered.

We summarize the main contributions of the paper as follows:

We introduce a learning problem that shares important features with learning through the interaction behavior that is crucial for LLMs and complex systems.

We analyze this problem from a complexity-theoretic viewpoint where we show that it is in general hard, and from an information-theoretic viewpoint where we show, roughly, that for N different tokens in the alphabet order N ln( N ) samples are sufficient to identify the latent structure.

We then carefully investigate the gradient dynamics of token embeddings, finding local recovery of the cluster structure and global recovery for tensor-product functions on the tokens if we have more than N 2 ln( N ) samples for an alphabet with N tokens.

Notation. We write [ N ] = { 1 , . . . , N } for the set of the first N integers. The cardinality of a finite set A is denoted by | A | and we also denote the standard Euclidean norm of any vector v ∈ R d by | v | . The expressions λ max ( A ) and λ min ( A ) denote the largest and smallest eigenvalue of a symmetric matrix A respectively. We denote the uniform distribution over a set A by U ( A ) . For two subsets A,B ⊂ R d we denote by A + B = { a + b : a ∈ A,b ∈ B } their Minkowski sum. We denote the permutation group on k elements by S k . An overview of the used variable names can be found in Appendix A.

2 Setting and Motivation

In this section, we illustrate our problem with an example and define the setup more formally. Consider the set of all animals. Those can be grouped into classes such as mammals, birds, or reptiles (in fact there is a rich hierarchical structure which we ignore here). Those groups were conceived by findings sets of animals that share many properties. Once these groups are found, we can predict unobserved properties by first identifying the cluster to which an animal belongs and then predict that the property is shared with animals in the same cluster. Note that this is a specific instance of a general problem in scientific inference, where we want to uncover a hidden grouping of similar entities from sparse observations about these entities.

i

= 1

i

= 2

i

= 3

1

4

7

9

1

6

3

5

8

2

2

3

6

7

4

5

8

2

6

1

5

3

4

7

Here our main motivation, however, stems from the analysis of large language models where a similar problem arises implicitly during training. They are trained by next token prediction, so we do not expect them to learn structure by deductive reasoning such as cows are mammals, and mammals have lungs, so cows have lungs. Instead, their learning signal is whether a token can be replaced by another token for a given context. Thus, it is a natural question whether gradient descent-based training on token embeddings can uncover a hidden cluster structure of the data. Note that if the hidden structure is recovered, then generalization to unseen prompts is possible.

We now introduce our formal setup that captures key aspects of the discussion. We consider I sets of N 1 , N 2 , . . . , N I tokens or entities (such as words). For simplicity, we identify these with tokens from the set [ N i ] . For each of the sets [ N i ] there is a partition P i in K i classes which we can identify with the set [ K i ] . Then we can encode the partitions through maps Π i : [ N i ] → [ K i ] so that the partition is given by P i = (Π -1 i ( k i )) k i ∈ [ K i ] , i.e., Π i encode the class membership. We consider the map

This structure is illustrated in Figure 1 Now we assume that there is a function g : [ K 1 ] × [ K 2 ] × · · · [ K I ] → R which depends only on the classes. The map f = g ◦ Π extends this map to the tokens such that it only depends on the group a token belongs to. In the case where g (and thus f ) maps to { 0 , 1 } this can be interpreted as truth assignments, i.e., a statement consisting of a sequence ( n 1 , . . . , n I ) is true if f ( n 1 , . . . , n I ) = 1 and false otherwise and this case is the main motivation of our work. More generally, f could output the index of a suitable next token where in the 0 , 1 case 0 could correspond to a negation while 1 to an end of sentence token. Our goal is to learn the partitions P i or, equivalently, Π up to a permutation and thereby identify the hidden structure.

We assume that we are given data in the form of samples ( n s , f ( n s ) where n s = ( n s 1 , . . . , n s I ) and n s i ∈ [ N i ] . In other words, we try to learn the underlying structure from the interactions of a token with the other tokens, which is the same for every element of the partition. To simplify the notation and statements, we assume in the following that N i = N and K i = K for some N , K and all 1 ≤ i ≤ I . Our main interest concerns the case where N is large, i.e., there are many entities and K and I are small, i.e., the number of groups is small.

Let us summarize several features that this model shares with real world problem, such as learning suitable embeddings in latent space.

Hierarchical structures, i.e., groups of objects that share certain features, as discussed here, are abundant in language and science.

We only receive an indirect learning signal for the value of Π i ( n ) through its interaction with other tokens.

Interactions can be very complex, i.e., here the output depends on the interaction of I different tokens and ignoring parts of the context makes learning infeasible.

On the other hand, many important features are abstracted away, e.g.:

Here we assume that tokens from the same element of the partition interact in exactly the same way with other tokens while in reality there are many different partitions of the tokens depending on the broader context (e.g., we can group species by habitat, color, or, size each

resulting in different partitions) or there are exceptions, e.g., mammals generally do not lay eggs but the platypus does.

Many more complex notions of similarity or further properties of embeddings such as a vector space structure are not covered. Also, there can be many uninformative features.

We do not consider noisy data or errors in this work, which is crucial for real world applications.

3 Complexity-Theoretic and Information-Theoretic Analysis

We now study this learning problem in different settings. Let us first briefly discuss complexitytheoretic and information-theoretic properties of the learning problem to understand the general boundaries of this learning task. We first study the information-theoretic viewpoint, i.e., the question of how many samples are necessary to identify Π (and potentially g ). We focus on the case where we sample Π and the data-samples uniformly at random. To learn an unstructured map [ N ] I → R we generally need of the order N I ln( N I ) independent samples ( N I when sampling without replacement).

For the structured setting we show that if Π is drawn uniformly at random then generally order K I N ln( N ) samples are sufficient to learn f and the partition induced by Π . In other words, for every token n i and each of the K I classes Π -1 ( k ) we need of the order of ln( N ) samples n such that Π( n ) = k and n i = n i . In particular, for N ≫ K I any token will interact only with K I ln( N ) ≪ N other tokens, i.e., a very sparse subset of the other tokens.

̸

We require the following necessary condition for identifiability: For every k i = k ′ i there are k 1 , . . . , k i -1 , k i +1 , . . . , k I ∈ [ K ] such that

̸

Note that if this condition is indeed necessary because if it is not satisfied then it is not possible to distinguish Π -1 i ( k i ) and Π -1 i ( k ′ i ) . Clearly, we can generally only identify Π up to a permutation symmetry, i.e., we can only find ˜ Π such that there are permutations π i ∈ S K such that Π i = π i ˜ Π i . We have the following result.

Theorem 1. Assume that g : [ K ] I → R is a function satisfying the assumption (2) . Assume we randomly sample maps Π i such that Π i ( n i ) = k i with probability K -1 for all i , n i , and k i and such that (Π i ( n i )) i ∈ I,n i ∈ [ N ] are independent. Assume we are given S samples ( n , g ◦ Π( n )) where n ∼ U ([ N ] I ) . Then there is a constant N 0 ( I, K, η ) such that with probability at least 1 -2 e -η for N ≥ N 0 ( I, K, η ) and

we can recover Π and g up to permutations of [ K ] .

̸

This result is a special case of Theorem 6 in Appendix B which shows similar bounds for arbitrary maps Π that are not necessarily random. In the more general setting, there are additional dependencies on the size of the preimages Π -1 i ( k ) . Note that this dependency cannot be avoided because if there is a k such that | Π -1 ( k ) | = 1 and g ( k ) = 1 and g ( k ′ ) = 0 for k = k ′ then order N I samples are necessary to find n such that Π( n ) = k and thus Π . The general proof idea is to bound the probability that any fixed Π ′ = Π is compatible with the dataset. It turns out that it is possible to bound this probability in terms of the partitions distance of the partitions induced by Π and Π ′ . Then we are left with bounding the number of partitions, and we conclude with the union bound. We now show that this bound is essentially tight.

Theorem 2. Let g : [ K ] I → R be a function such that g (1 , k 2 , . . . , k I ) = g (2 , k 2 , . . . , k I ) for all k 2 , . . . , k I ∈ [ K ] except when k 2 = k 3 = . . . = k I = 1 . Assume that N is divisible by K and that | Π -1 i ( k ) | = N/K for all i ∈ [ I ] and k ∈ [ K ] . Given 3 ≤ S ≤ NK I -1 ln( N/K ) / 4 samples ( n s , g ◦ Π( n s )) where n ∼ U ([ N ] I ) i.i.d. Then the function Π is identifiable with probability at most 2 e -√ N/K .

The proof of this result can be found in Appendix B. Next, we emphasize that while typically a rather small number of samples is sufficient to learn Π it can generally be very hard to do this in practice. More concretely, we show that for I ≥ 3 even deciding whether there is a map

̸

Π = Π 1 ⊗ . . . ⊗ Π I : [ N ] I → [ K ] I such that f = g ◦ Π given access to samples of the form ( n s , t s ) is NP-complete. We show that this is true even if g is known, I = 3 and K = 2 .

Theorem 3. Consider the map g : { 0 , 1 } 3 →{ 0 , 1 } given by

Then it is an NP-complete problem to decide given samples of the form ( n s 1 , n s 2 , n s 3 , t s ) ∈ [ N ] 3 ×{ 0 , 1 } whether there is a map Π = Π 1 ⊗ Π 2 ⊗ Π 3 with Π i : [ N ] →{ 0 , 1 } such that t s = g ◦ Π( n s 1 , n s 2 , n s 3 ) for all samples.

The proof of this result can be found in Appendix C. Now that we established under what the conditions Π can in principle be learned, and clarified that this might be hard in general, we next discuss how we can find Π in practice. First, we remark that Theorem 3 rules out the existence of any general fast algorithms to learn Π . Given the combinatorial nature and the hardness of the problem, it is natural to reformulate the task as a constraint satisfaction problem which can then be solved using standard SAT solvers (see, e.g., [14] for a review). Indeed, we can introduce Boolean variables t i kn for i ∈ [ I ] , k ∈ [ K ] , and n ∈ [ N ] which encode whether Π i ( n ) = k and r k v for every k ∈ [ K ] I and v ∈ Im( f ) in the (finite) image of f that encode whether g ( k ) = v . It is then relatively straightforward to then express the conditions for the map Π as a constraint satisfaction problem which is satisfiable if and only if there are maps Π and g such that t s = g (Π( n s )) holds for all samples. We outline the construction in more detail in Appendix C. We leave the task of developing and studying efficient algorithms for the considered problem for future work because the main motivation of this paper is rather to understand how the complex statistical patterns can be extracted using simple gradient based algorithms. This will be investigated in the next section.

4 Analysis of Gradient Descent Dynamics

In this section, we investigate under what conditions the clustering induced by Π can be learned using gradient descent on token embeddings. Our main finding is that for uniformly random Π and S sufficiently large gradient descent can be used to uncover or at least preserve the cluster structure of the embeddings. This shows that while the general problem is NP hard typical random instances with sufficiently many samples can be solved with straightforward algorithms quickly. This is in spirit similar to the results found in, e.g., [5]. Let us start by introducing the setting.

Setting. We assume that we have token embeddings for each of the I sets [ N 1 ] to [ N I ] , i.e., we assume that there are vectors v ( i, n ) ∈ R D for some D and all i ∈ [ I ] , n ∈ [ N ] . Based on these embeddings, we assume that we are given a function ˆ f : R ID → R that transforms the embeddings into a prediction. We will abuse notation and write for n ∈ [ N ] I

i.e., we will suppress the map from tokens to embeddings in the notation.

Now we consider gradient descent for the embedding vectors using the least square loss on training samples, i.e., the loss of a sample ( n , t = f ( n )) is ( ˆ f ( n ) -f ( n )) 2 .

We assume that we are given a dataset D = { n 1 , . . . , n S } ∼ U ([ N ] I ) S . Then the empirical loss reads (the division by 2 is convenient for the gradient evaluation later)

We also define shorthands for certain concatenations of embeddings. For a sample n ∈ [ N ] I we denote the collection of embeddings by v ( n ) = ( v (1 , n 1 ) , . . . , v ( I, n I )) . Using the convention (5) we can then write ˆ f ( v ( n )) = ˆ f ( n ) .

Moreover, we define ¯ v ( i ) ∈ R DN as the concatenation of the vectors v ( i, 1) , . . . , v ( i, n ) , i.e., the combined embedding for the i -th slot and ¯ v ∈ R DNI as the concatenation of the vectors ¯ v ( i ) for 1 ≤ i ≤ I , i.e., all token embeddings concatenated. We consider the regularized loss given by

Note that the scaling by N/S (instead of usual 1 /S ) is natural because every token embedding v ( i, n ) occurs in approximately S/N of the samples and so the scaling ensures that the gradient of ˆ R λ with respect to the token embedding v ( i, n ) is of order one. Now, we consider the continuous gradient descent of the loss with respect to the embeddings, i.e., we consider (omitting the time variable from the notation)

This introduces a time dynamics on the token embeddings. We indicate the time dependence by v ( i, n, t ) but we drop t if not necessary. Our main goal is to understand which conditions ensure that the token embeddings v ( i, n, t ) and v ( i, n ′ , t ) converge to each other if Π i ( n ) = Π i ( n ′ ) as t →∞ . To investigate this we define the center of the class embeddings by

and we consider the deviations from the class centers given by

Thus the vectors δ ( i, n, t ) capture whether we recover the cluster structure, in particular if all norms | δ ( i, n, t ) | are small then we essentially recovered the hidden structure. Therefore, we define

Similarly, to the notation introduced before we consider for k ∈ [ K ] I the vector w ( k ) = ( w (1 , k 1 ) , . . . , w ( I, k I )) . As in (5) we abuse notation and write

Similar to ¯ v ( i ) and ¯ v we introduce ¯ w ( i ) as the concatenation of ( w ( k, i )) k ∈ [ K ] and ¯ w as the concatenation of ¯ w ( i ) .

Assumptions. Our first result for the gradient dynamics states that clusters are stable under the dynamics if the initial loss is sufficiently small. More precisely, this means that we assume that v ( i, n ) and v ( i, n ′ ) are close initially whenever Π i ( n ) = Π i ( n ′ ) , i.e., δ max (0) is small. In addition, we assume that | ˆ g ( k ) -g ( k ) | is small. To capture this we define

Then the result shows that δ max stays small for all times if S ≳ N 2 under mild additional assumptions. In other words, if we start from the correctly learned substructures and ˆ g ( k ) ≈ g ( k ) for all k ∈ [ K ] I then this remains true for all times. Note that while we phrase smallness as an assumption on the mean embeddings w ( i, k ) this is generally a consequence of δ max small and a small empirical loss ˆ L (¯ v ) . Let us now state the required assumptions.

Assumption 1. We assume that the map Π : [ N ] I → [ K ] I is approximately balanced which means that for all i ∈ [ I ] , k ∈ [ K ]

This assumption ensures that clusters are of approximately equal size. We have already seen in Section 3 that different cluster sizes increase the sample complexity of learning Π .

Assumption 2. We assume that there is a convex set Ω ⊂ R D and a constant M ′ ≥ 1 such that the following bound holds

Here it is convenient to introduce M = 16 M ′ 2 so that later certain errors in a Taylor approximation are bounded by M . We also assume that

This is a rather mild assumption. For C 3 functions ˆ f and Ω bounded this is always true. The next assumption entails a rigidity of approximate minimizers of the loss.

Assumption 3. We assume that for all embeddings w ( i, k ) ∈ R D for i ∈ [ I ] , k ∈ [ K ] that satisfy

the bound

holds for some positive constant ω 0 . Of course, the bound r max ≤ 1 could be replaced by any other constant.

Note that this condition can only hold if D ≤ K I -1 , i.e., the latent space dimension cannot be too large. The high level intuition of this assumption is essentially that (at least if ∑ k ( ˆ f ( w ( k )) -g ( k )) 2 is small) there is no direction v ∈ R D such that v · ∇ w ( i,k ) ˆ f ( w ( k )) ≈ 0 for all k such that k i = k , i.e., we cannot move one single embedding without changing the output ˆ f ( w ( k )) for at least one k . If this condition does not hold, then we cannot guarantee that ∑ k ( ˆ f ( w ( k )) -f ( k )) 2 is minimized for a unique w ( i, k ) (for all other embeddings fixed). This generally prevents concentration of v ( i, n ) . Note that this condition does not ensure that there is a unique minimizer ¯ w , in particular there could still be a rotationally invariant family of embeddings w ( i, k ) such that ˆ f ( w ( k )) = g ( k ) for all k ∈ [ K ] I . Finally, we need a further mild assumption that ensures that mean token embeddings w ( i, k ) stay bounded in some set if the loss is small. This can be achieved, e.g., if ˆ f → ∞ if | w ( k ) |→∞ .

Assumption 4. We assume that for all collections of mean embeddings w ( i, k ) ∈ R D for i ∈ [ I ] , k ∈ [ K ] that satisfy

there is a convex set Ω 0 ⊂ R d such that w ( i, k ) ∈ Ω 0 for all i ∈ [ I ] , k ∈ [ K ] . Again, the right-hand side of the bound r max ≤ 1 could be replaced by any other constant.

Results. The first stability theorem can then be stated as follows.

Theorem 4. Let Π : [ N ] I → [ K ] I be approximately balanced as stated in Assumption 1 Assume that the functions g : [ K ] I → R and ˆ f : R ID → R satisfy Assumption 3 for some ω 0 > 0 and Assumption 4 for some convex set Ω 0 . Assume that Assumption 2 holds for some M and the set Ω = Ω 0 + B 2 (0) . Then there are constants c 1 , C 2 , C 3 , C 4 > 0 depending on I , M , D , and ω 0 such that for all initial embeddings v ( i, n, t = 0) ∈ R D for i ∈ [ I ] and n ∈ [ N ] satisfying

the following holds with probability at least 1 -S -1 over the randomness of the dataset. When considering the gradient dynamics of the embeddings given by (8) the bound

In particular δ max ( t ) → 0 if r max ( t ) → 0 , i.e., all token embeddings for one fixed class converge to the same point.

and sample size holds and moreover

This result shows that for order N 2 ln( N ) samples and initialization sufficiently close to a global minimum the cluster structure remains stable. We do not provide conditions that ensure r max ( t ) → 0 which guarantees convergence to 0 loss and perfect recovery of the clusters.

Next we note that we cannot expect the clustering to be stable in general even if δ max (0) is arbitrarily small if initialization is not close to a minimum. This is true even in the simplest case where I = K = 1 , i.e, we consider gradient descent for a single function value. Then gradient descent does not necessarily converge to a global minimum and close by points do not necessarily stay close because gradient descent is not well posed. Let us clarify this by an example.

4

2

0

2

4

6

4

2

0

2

4

6

0

50

100

150

200

250

Gradient steps

0

1

2

3

4

Average distance all embeddings

Example 1. Assume I = K = D = 1 , N > 1 , ˆ f ( x ) = -x 2 + x 3 and f ( n ) = -2 . Consider the dataset D = { 1 , . . . , n } . Assume that v (1 , n, t = 0) ∼ N (0 , σ 2 ) for any σ 2 > 0 . Then the gradient dynamics introduced in (8) reads

We find that v (1 , n, t ) → 2 / 3 (which is a local minimum of ˆ f as t →∞ if v (1 , n, t = 0) > 0 . On the other hand v (1 , n, t ) →-1 if v (1 , n, t = 0) < 0 . So in this case δ max (0) = O ( σ 2 √ ln( N )) but δ max ( t ) → 5 / 3 as t →∞ . Slight modifications show that also δ max →∞ is possible.

The previous example shows that without additional assumption, we cannot expect to recover the structure of the data. Therefore, we impose additional restrictions on the function ˆ f . As apparent from Example 1 and also from the bound in Lemma 1 in Appendix E, it is the curvature of the function ˆ f that can push token embeddings of the same class apart. We therefore consider the function class of slow-wise linear functions defined as follows.

Definition 1. We call ˆ f slot-wise linear if for every v = ( v (1) , . . . , v ( I )) ∈ R D × I and any i ∈ [ I ] , α, β ∈ [ D ] the relation

holds.

Let us denote for v ∈ R D by ˜ v ∈ R D +1 the vector v where we append a 1. The most general slot-wise linear function is then of the form

where T : R ( D +1) I → R is a linear map. Note that this class covers linearized attention where the embeddings v ( i, n ) are split in three separate parts that are used to form key, query, and value vectors. For this function class we can show stronger clustering results.

Theorem 5. Let Π be approximately balanced, i.e., assume that Assumption 1 holds. Let ˆ f : R ID → R be a slot-wise linear. Assume that v ( i, n, t ) follow the gradient dynamics (8) and that v ( i, n, t ) ∈ Ω for all i ∈ [ I ] , n ∈ [ N ] and t > 0 for some convex set Ω . Assume that Assumption 2 holds with constant M for the set Ω . Let C 1 be the constant from Lemma 1 (which depends on I , D , M ). Assume that at initialization

Then there are constant C 2 , C 3 ≥ 0 depending on M , I , D , such that for

the bound

holds for all t ≥ 0 .

The high level summary of this result is that for order N 2 ln( N ) samples the clusters can be recovered up to an error of order λ -1 √ ln( S ) N/S . Note that the a-priori assumption that the gradient flow is restricted to the set Ω might appear difficult to guarantee in practice. However, we conjecture that the results extend to gradient dynamics clamped at the boundary. Moreover, in Lemma 2 we prove that the mean embeddings w ( i, k, t ) stay within a ball of radius R = O ( √ λ -1 ) for all times. This allows us to prove Theorem 8 which does not require any a-priori bounds on the evolution of v ( i, n, t ) but comes at the price that the constants C 1 , C 2 , and C 3 depend on λ so we cannot infer the explicit λ dependence. Let us make an important remark.

Remark 1 . While we state our results for a fixed function ˆ f this function could in principle be time-dependent, e.g., ˆ f could be given by a neural network, and we consider gradient descent not only on the token embeddings but also on the network parameters. The only requirement is that the assumptions hold uniformly for all times t . In particular, for Theorem 5 we only need to ensure that the derivatives of ˆ f stay uniformly bounded in time. This can, e.g., be guaranteed by clipping the parameters of the slot-wise linear map ˆ f .

Our results so far show that we can recover the right cluster structure in the sense that the embeddings from the same group cluster. However, this leaves open whether there is any non-trivial dynamics at all, i.e., all embeddings could cluster at the same point. This is in general not the case as can be seen from Corollary 1 which states that in the setting of Theorem 5 and for large times the dynamics of the cluster-means follow the equation

where (2 K ) -I ≤ α k ,i ≤ (2 /K ) I are positive numbers. This shows that w ( i, k ) follow generally a non-trivial dynamic (and this also justifies the scaling as this expression is of order 1). So in the generic case the cluster structure will be revealed if the numbers of samples is sufficiently large, however, there is no general guarantee that the clusters are well separated. As an illustration of this result we refer to Figure 2 where the clustering of the embeddings becomes apparent.

Proof idea and overview. Let use here give a quick summary of the main steps of the proof and where to find them. The first important ingredient is an expansion of the loss gradient. We Taylor expand the loss of a sample v ( n ) around the point w (Π( n )) to second order with remainder terms, the relevant calculations can be found in Appendix D (see Proposition 1 for the outcome). A second ingredient are concentration bounds for certain datapoint-statistics random variables and

random matrices. Those are derived in Section G with the necessary results collected in Theorem 9. Combining the Taylor expansion with the concentration result, we can extract the dominant terms of the expansion (see Appendix E). Moreover, we obtain such an expansion for ˙ w ( i, k ) (see Corollary 1) and thus the displacements ˙ δ ( i, n ) (see Corollary 2). This expansion can then be used to control ∂ t | δ ( i, n ) | 2 (see Lemma 1) which is sufficient to control δ max .

Discussion of assumptions Let us contemplate the differences and similarities to training token embeddings in neural networks.

For the first main result Theorem 4 we make minimal assumptions on ˆ f so this could in principle be a neural network applied to the token embeddings. The second result, Theorem 5, is more restrictive but covers subclasses of linearized attention.

An important feature of the results is that ˆ f itself could be time dependent (see Remark 1).

Differences to standard training of neural networks are:

We use continuous time gradient descent instead of stochastic gradient descent. This is a frequently used modification and in suitable limits those converge (see, e.g., [20]).

We use mean squared error, while sequence modelling usually relies on cross entropy loss. This simplification is frequently used in theoretical analysis, but it is expected that results generally extend to the non-convex cross entropy loss.

A more crucial difference is that the embedding space dimension in practice is usually chosen large to provide large representation capacity. Here D has to be rather small to allow a unique optimal solution of the token embeddings that allows us to recover the cluster structure.

5 Conclusion

In this paper, we considered a learning problem where we try to recover a partition of tokens from their interaction with other tokens. This can be seen as a toy problem for next token prediction in LLMs, but also more broadly as a problem in scientific inference. We studied this problem from different perspectives, namely from an information-theoretic, complexity-theoretic, and gradient descent based viewpoint. We found that order N ln( N ) samples are sufficient to recover the partition for N tokens, while we showed that N 2 ln( N ) samples are sufficient for gradient based methods. There are several natural open follow-up questions. First, there are some open questions regarding the tightness of our analysis of the gradient descent. In particular, it is a natural question whether already Ω( N ln( N )) samples are sufficient to control δ max which is the information-theoretic threshold and would be similar to the optimal results for matrix completion [18, 7]). Another interesting question for future research is whether Theorem 5 holds for standard initialization schemes for the token embeddings. Secondly, it is of interest to relax the notion of clustering of embeddings to more general notions that still allow recovering some structure but are also applicable to high dimensional latent spaces and potentially to multiple partitions and relations on the same tokens (e.g., tokens belonging to different clusters). Thirdly, it is a natural question whether this work can be connected more closely to empirical findings.

References

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research , 18(19):1-53, 2017.

William Bialek, Andrea Cavagna, Irene Giardina, Thierry Mora, Edmondo Silvestri, Massimiliano Viale, and Aleksandra M. Walczak. Statistical mechanics for natural flocks of birds. Proceedings of the National Academy of Sciences , 109(13):4786-4791, 2012.

Piotr Bojanowski, Edouard Grave, Armand Joulin, and Tomás Mikolov. Enriching word vectors with subword information. Trans. Assoc. Comput. Linguistics , 5:135-146, 2017.

Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for a ConvNet with Gaussian inputs. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning , volume 70 of Proceedings of Machine Learning Research , pages 605-614. PMLR, 06-11 Aug 2017.

Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent. CoRR , abs/2402.18724, 2024.

Emmanuel J. Candes and Terence Tao. The power of convex relaxation: Near-optimal matrix completion. IEEE Transactions on Information Theory , 56(5):2053-2080, 2010.

Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers) , pages 4171-4186. Association for Computational Linguistics, 2019.

Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021 , pages 12873-12883. Computer Vision Foundation / IEEE, 2021.

Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers, 2024.

Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021 , pages 5484-5495. Association for Computational Linguistics, 2021.

Irene Giardina. Collective behavior in animal groups: theoretical models and empirical studies. HFSP J , 2(4):205-219, Aug 2008.

Carla P. Gomes, Henry Kautz, Ashish Sabharwal, and Bart Selman. Chapter 2 satisfiability solvers. In Frank van Harmelen, Vladimir Lifschitz, and Bruce Porter, editors, Handbook of Knowledge Representation , volume 3 of Foundations of Artificial Intelligence , pages 89-134. Elsevier, 2008.

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 37822-37836. Curran Associates, Inc., 2022.

Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.

Raghunandan H. Keshavan, Andrea Montanari, and Sewoong Oh. Matrix completion from a few entries. IEEE Transactions on Information Theory , 56(6):2980-2998, 2010.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 19689-19729. PMLR, 23-29 Jul 2023.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. CoRR , abs/1902.06015, 2019.

Tomás Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. In Yoshua Bengio and Yann LeCun, editors, 1st International Conference on Learning Representations, ICLR 2013, Scottsdale, Arizona, USA, May 2-4, 2013, Workshop Track Proceedings , 2013.

Tomás Mikolov, Wen-tau Yih, and Geoffrey Zweig. Linguistic regularities in continuous space word representations. In Lucy Vanderwende, Hal Daumé III, and Katrin Kirchhoff, editors, Human Language Technologies: Conference of the North American Chapter of the Association of Computational Linguistics, Proceedings, June 9-14, 2013, Westin Peachtree Plaza Hotel, Atlanta, Georgia, USA , pages 746-751. The Association for Computational Linguistics, 2013.

OpenAI. GPT-4 technical report, 2023.

P. Romanczuk, M. Bär, W. Ebeling, B. Lindner, and L. Schimansky-Geier. Active brownian particles. The European Physical Journal Special Topics , 202(1):1-162, 2012.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR , abs/2302.13971, 2023.

Joel A. Tropp. User-friendly tail bounds for sums of random matrices. Foundations of Computational Mathematics , 12(4):389-434, 2012.

Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 6306-6315, 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research , 25(49):1-55, 2024.

Supplementary Material

This supplementary material is structured as follows. We first review the notation used in the paper in Appendix A. Then we provide the proofs for the information-theoretic results in Section 3 in Appendix B and the proof of Theorem 3 and a reduction to a constraint satisfaction problem in Appendix C. The proofs of the main results concerning the gradient flow rely on a careful Taylor expansion of the loss gradient. This expansion can be found in Appendix D and bounds for this expansion are derived in Appendix E. Based on these bounds, we can prove our main results in Appendix F. An important ingredient in bounding the Taylor expansion are concentration results for the dataset statistics that can be found in Appendix G. Finally, we review some results on random matrices in Appendix H which are necessary for the concentration bounds and we review the definition of Kronecker products in Appendix I.

A Overview of Notation used

Let us here collect important notation used throughout the paper and the proofs.

General notation.

Numbers up to n : [ n ] = { 1 , . . . , n }

Eigenvalues of a matrix: λ i ( A )

Largest eigenvalue: λ max ( A )

Operator norm of a matrix: ∥ A ∥

Notation used in the learning problem.

Number of slots: I

Number of tokens for each slot: N

Number of classes: K

Number of samples: S

Map defining the partition in subclasses: Π = Π 1 ⊗ . . . ⊗ Π I : [ N ] I → [ K ] I

Function on classes: g : [ K ] N → R

Induced function on tokens: f : [ N ] I → R , f = g ◦ Π

Notation used in the gradient descent analysis.

Dimension of latent space: D

Embedding for token n of slot i : v ( i, n )

Token embeddings for a sample n : v ( n ) = ( v (1 , n 1 ) , . . . , v ( I, n I ))

Token embeddings of a slot i : ¯ v ( i ) = ( v ( i, n )) n ∈ [ N ]

All token embeddings: ¯ v = (¯ v ( i )) i ∈ [ I ]

Mean of cluster embeddings: w ( i, k )

Sample version: w ( k ) = ( w (1 , k 1 ) , . . . , w ( I, k I ))

Mean token embeddings for slot i : ¯ w ( i ) = ( w ( i, k )) k ∈ [ K ]

All mean embeddings: ¯ w = ( ¯ w ( i )) i ∈ [ I ]

Displacements from cluster center: δ ( i, n ) = v ( i, n ) -w ( i, Π i ( n ))

Displacements of a sample: δ ( n ) = ( δ (1 , n 1 ) , . . . , δ ( I, n I ))

Displacements of all tokens for slot i : ¯ δ ( i ) = ( δ ( i, n )) n ∈ [ N ]

Function on token embeddings: ˆ f : R DI → R , identified with ˆ f : [ N ] I → R , ˆ f ( n ) = ˆ f ( v ( n ))

Function on classes: ˆ g : [ K ] I → R , ˆ g ( k ) = ˆ f ( w ( k ))

Regularization: λ

B Proofs for Information-Theoretic Results

In this section, we provide the proofs and additional results for Section 3. Let us first start to state the general information-theoretic bound that handles the case where | Π -1 i ( k ) | might be of arbitrary size.

Theorem 6. Let g : [ K ] I → R satisfy (2) and there is a projection Π = Π 1 ⊗ . . . ⊗ Π I : [ N ] I → [ K ] I . Assume there is L > 0 such that for every k ∈ [ K ] I we have Π -1 ( k ) ≥ N I /L . Moreover, assume for all i ∈ [ I ] and k ∈ [ K ] the bound | Π -1 i ( k ) | ≥ M , i.e., every class has at least M members. Assume we are given S samples { ( n s , g ◦ Π( n s )) : s ∈ [ S ] } where n ∼ U ([ N ] I ) . If

for some η ≥ 2 then with probability at least 1 -6 e -η we can recover Π and g up to permutations of [ K ] . In other words, for any map Π ′ = Π ′ 1 ⊗ . . . ⊗ Π ′ I and g ′ such that g ◦ Π( n s ) = g ′ ◦ Π ′ ( n s ) . There are permutations π i ∈ S K such that Π i = π i ◦ Π ′ i and the corresponding relation holds for g , g ′ .

A few remarks to explain this result are in order.

Remark 2 . The scaling of S might appear slightly complicated, so let us comment on this. We are mostly interested in the case where N large and K stays bounded. In addition, we are primarily interested in the regime where L ≤ CK I bounded and M ≥ cN/K (which holds for random Π ), i.e., the sampling probability of each class ¯ k ∈ [ K ] I is of similar size. In this case, the first term in the condition (31) for S stays constant as N →∞ . The second term dominates, and we see that we need only O ( N ln( N )) samples to identify Π and g . This is the setting studied in Theorem 1 in the main text. The result is essentially tight in this limit (up to the dependence on I ) as stated in Theorem 2. Note that as remarked in the main text, the dependence on L cannot be avoided. Indeed, consider the extreme case where Π i ( n ) = 1 for n = 1 and Π i ( n ) = 2 for n > 1 and g ( k 1 , . . . , k I ) = 1 iff k 1 = . . . = k I = 0 and g ( k 1 , . . . , k I ) = 1 otherwise. Then we have L = N I and we also need of order N I samples to sample the point (0 , 0 , . . . , 0) ∈ [ N ] I which is necessary to identify Π .

Proof. Let us first introduce some notation. We denote the set of samples

Consider two partitions P , Q of a set [ N ] . We denote by P A for A ⊂ [ N ] the restriction of a partition to a subset (i.e., { P ∩ A : P ∈ P} ). The partition distance is defined by

i.e., the minimal number of elements that need to be removed such that the partitions agree.

We call Π ′ compatible with the datapoints D if there is a g ′ such that g ′ ◦ Π ′ = g ◦ Π on the data. Now the general strategy is to consider any other candidate map Π ′ and upper bound the probability that is compatible, i.e., for some function g ′ the functions f ′ = g ′ Π and f agree on the data. We will then conclude by applying the union bound over maps Π ′ . Thus, we need to prove an upper bound on the number of partitions with partition distance at most ∆ and a lower bound on the error probability for a given ∆ . Let us start with the latter. Denote by

the partition generated by Π i . Consider any other map Π ′ = Π ′ 1 ⊗ . . . ⊗ Π ′ I : [ N ] I → [ K ] I with corresponding partitions P ′ i . We define for any such Π ′ the quantity

and in particular D ( P i , P ′ i ) ≤ ∆ . Let P 1 = { P 1 , . . . , P K } where P k = Π -1 1 ( k ) . Let A be a set of maximal size such that P A 1 = P ′ A 1 . After relabeling we can assume that P ′ 1 = { P ′ 1 , . . . , P ′ K } and

Assume now that

P k ∩ A = P ′ k ∩ A . By composing Π ′ 1 with a permutation we can also assume that P ′ k = Π ′-1 1 ( k ) . Moreover,

The same applies to all partitions P i and P ′ i . Next, we claim that if f = g ◦ Π and f ′ = g ′ ◦ Π ′ agree on all samples then with probability K I e -S 2 I L over the randomness of the samples g = g ′ . Define E k := Π ′-1 ( k ) ∩ Π -1 ( k ) . Using (37) and the assumption on Π we have that for k ∈ [ K ] N

The probability that none of the S samples is in E k can then be bounded as follows

Applying the union bound over k we find that

̸

Now we bound the probability that g ′ = g is compatible, i.e., g ◦ Π( n s ) = g ◦ Π ′ ( n s ) for all s ∈ [ S ] . Let now n 1 ∈ [ N ] . Denote by E n 1 ⊂ [ N ] I the set of all vectors n = ( n 1 , n 2 , . . . , n I ) such that

̸

Wenowlower bound the size of E n 1 under the assumption that there is k ∈ [ K ] such that n 1 ∈ P ′ k ∩ A c (where A is as above a set such that P A 1 = P ′ A 1 ). Then Π ′ 1 ( n 1 ) = k and (by minimality of A c ) we find n 1 / ∈ P k and thus Π 1 n 1 = ¯ k = k for some ¯ k . By assumption, we can find k 2 , . . . , k I such that g ( k, k 2 , . . . , k I ) = g ( ¯ k, k 2 , . . . , k I ) . Consider now any vector n = ( n 1 , n 2 , . . . , n I ) such that n i ∈ Π -1 i ( k i ) ∩ Π ′ -1 i ( k i ) . Then, for any such n

̸

̸

we conclude that all such n are in E n 1 By assumption we have

The sets E n 1 are disjoint so we conclude that E = ∪ n 1 ∈ [ N ] E n 1 satisfies

Now we can upper bound the probability that D is compatible with Π ′ and g by

Clearly the same reasoning applies to any other index instead of 1.

The next step is to upper bound the number of such candidate partitions. The number of partitions such that D ( P 1 , P ′ 1 ) ≤ ∆ can be bounded by ( NK ) ∆ , i.e., we ∆ times select one of N tokens and assign it to another (or the same) class. This bound can be applied to all indices 1 ≤ i ≤ I and by the union bound, we find that

This implies that and therefore

̸

Thus we can bound the the probability that any of the maps Π ′ (there are more such maps than partitions because there is a label assigned to each class, but consistency with the data depends only on the underlying permutation) with 0 < ∆(Π ′ ) ≤ M/ 2 is consistent with the data is bounded by

Here the first term corresponds to the upper bound for the probability that (after permutation) g = g ′ is compatible with the data and the second term bounds the probability that any other function is compatible with the data and I accounts for the fact that any i ∈ [ I ] can be arg max i D ( P i , P ′ i ) .

̸

We now consider the remaining case that there is an index i ∈ [ I ] such that D ( P i , P ′ i ) ≥ M/ 2 where we can assume w.l.o.g. that i = 1 . As above, let P 1 = { P 1 , . . . , P K } and P ′ 1 = { P ′ 1 , . . . , P ′ K } such that P k ∩ A = P ′ k ∩ A and A is a set of maximal size. We now claim that we can find an index k = ¯ k such that

Indeed, first assume that there is an index k such that | P ′ k ∩ P k | ≤ M/K . Since | P ′ k | ≥ M there is another index ¯ k such that | P ′ ¯ k ∩ P k | ≥ M/K . By maximality of A we have P k ∩ A = P ′ k ∩ A = P k ∩ P ′ k and moreover

because otherwise exchanging P ′ k and P ′ ¯ k would allow picking a larger A . By our assumption, we conclude from here that

̸

Thus, in this case | P ′ ¯ k ∩ P ¯ k | ≥ M/ (2 K ) and | P ′ ¯ k ∩ P k | ≥ M/K and (50) holds. Assume now to the contrary that there is no index k such that | P ′ k ∩ P k | < M/K . The assumption | A c | ≥ M/ 2 implies that there is k such that | A c ∩ P ′ k | ≥ M/ (2 K ) and by minimality of A c we find A c ∩ P ′ k ∩ P k = ∅ and thus there is ¯ k = k such that | P ′ k ∩ P ¯ k | = | A c ∩ P ′ k ∩ P ¯ k | ≥ M/ (2 K 2 ) . This finishes the proof of (50).

̸

We now fix k and ¯ k such that | P k ∩ P ′ k | ≥ M/ (2 K 2 ) and | P ¯ k ∩ P ′ k | ≥ M/ (2 K 2 ) . Then, by assumption, we can find k 2 , . . . , k I ∈ [ K ] such that k = ( k, k 2 , . . . , k I ) and ¯ k = ( ¯ k, k 2 , . . . , k I ) satisfy g ( k ) = g ( ¯ k ) .

Moreover | Π -1 ( k ) | ≥ N I /L . Now for every k i there is k ′ i such that | Π -1 i ( k i ) ∩ Π ′ -1 i ( k ′ i ) | ≥ K -1 | Π -1 i ( k i ) | . Thus, there is k ′ ∈ [ K ] I such that

Define A i = Π ′-1 i ( k ′ i ) ∩ Π -1 i ( k i ) and A 1 = P k ∩ P ′ k and ¯ A 1 = P ¯ k ∩ P ′ k . Define A = A 1 × A 2 × . . . × A I and ¯ A = ¯ A 1 × A 2 × . . . × A I . Note that by (50)

Clearly Π ′ is constant (and equal to ( k, k ′ 2 , . . . , k ′ I ) ) on A ∩ ¯ A and therefore also g ′ ◦ Π ′ is constant for any g ′ . On the other hand, Π( A ) = { k } is constant and Π( ¯ A ) = ¯ k is constant but

̸

for ¯ n ∈ A and ¯ n ∈ ¯ A . Now Π ′ can only be consistent with the data (for any g ′ ) if there is no sample in A or no sample in ¯ A , i.e., and thus

where we used the union bound in the last step. Then we find using (54) and (56)

Finally, we can bound the number of partitions by the number of maps Π i : [ N ] → [ K ] which implies that

So we get the upper bound

P D ( ∃ Π ′ such that ∆(Π ′ ) ≥ M/ 2 and Π ′ compatible with D ) ≤ 2 IK NI e -MS KI +1 L . (59) Combining (49) and (59) and using the union bound we find

P D ((Π , g ) not identifiable up to permutations ) ≤ P D ( ∃ Π ′ with ∆(Π ′ ) > 0 is compatible with D )

Note that clearly the assumptions imply M ≤ N/K < N and thus e -S/ (2 I L ) ≤ e -S ∆ / (2 I LN ) for ∆ ≤ M/ 2 . We then find (using the simple bound IK I ≤ (2 K ) I ≤ ( NK ) I ∆ that

for

and η > 2 where we bound the geometric sum by twice its largest term. The first summand in (60) can be bounded by

The proof of Theorem 1 is now a direct consequence of the previous result because when assuming that a uniformly random map Π is chosen we can estimate the quantities M and L in the previous theorem with high probability.

Proof of Theorem 1. We observe that | π -1 i ( k i ) | ∼ Bin( N,K -1 ) . Applying a Chernoff bound on the tail of the binomial variable we obtain

By the union bound we get

if N ≥ N 0 ( η, K, I ) . Note that M ≥ N/ (2 K ) implies that L ≤ ( N/M ) I ≤ (2 K ) I . Assuming that Π is such that the bounds for M and L hold we can apply Theorem 6 and find that for N ≥ N 0 sufficiently large (depending on I , K , η ) and

(we bounded ln( NK ) ≤ 2 ln( N ) ) the maps Π and g are identifiable with probability at least 1 -e -η . Here we used that as N →∞ the term indicated above is dominating in (31). Now the union bound over the bad events ends the proof.

for

Proof of Theorem 2. As before, we note D = { n 1 , . . . , n S } . Assume there is n 1 such that Π 1 ( n 1 ) = 1 and such that

̸

Then Π ′ = Π is compatible with D where Π ′ i = Π i for i ≥ 2 and Π 1 ( n ) = Π ′ 1 ( n ) for n = n 1 and Π ′ 1 ( n 1 ) = 2 = 1 = Π 1 ( n 1 ) (by assumption on g ). Let us denote

To estimate the probability of the event ⋃ n 1 { A n 1 ∩D = ∅} we use Poissonization to make the events independent. Consider datasets ˜ D whose distribution is generated by first sampling ¯ S ∼ Poi(2 S ) and then conditional on ¯ S sample a dataset ˜ D as before, i.e., n s ∼ U ([ N ] I ) for s ∈ [ ¯ S ] . Then we find that

̸

and those events are independent for n 1 = n ′ 1 . Thus, we find that

By independence we now find that

where we used the upper bound for S and (1 -x ) ≤ e -x . Assume D s ∼ U ([ N ] I ) s and define

̸

Note that p S is an upper bound on the probability that Π is identifiable as explained above. We have shown that

Clearly p s is decreasing in s . This implies that

A Chernoff bound for the Poisson distribution reads

which implies with λ = 2 S and x = S that

We thus conclude that

for S ≥ 3 (implying 1 -(2 /e ) S ≥ 1 / 2 ).

Then

̸

̸

C Proofs of Complexity-Theoretic Analysis and Constraint Satisfaction Reduction

In this section, we provide the proof of Theorem 3 and a general reduction to a constraint satisfaction problem. We start with the proof of Theorem 3. An important property we will use frequently in the proof is that T has the property that for all values x 1 and x 2 there is x 3 ( x 1 , x 2 ) such that T ( x 1 , x 2 , x 3 ( x 1 , x 2 )) = 0 . Indeed, x 3 ( x 1 , x 2 ) = 0 except for x 1 = x 2 = 1 where x 3 (1 , 1) = 0 has this property.

Proof. We reduce the problem to 3SAT. We consider an arbitrary formula

on ℓ variables with m clauses C i . We denote f = g ◦ Π . To improve readability, we consider a symbol set S instead of [ N ] which is used for all 3 slots. We proceed in three steps. First we show that we can write down a set of equations that ensures that two symbols s 0 , s 1 ∈ S satisfy Π i ( s j ) = j for all i (if f = g ◦ Π ). Suppose the following relations hold

Then it is easy to conclude that Π i ( s j ) = j holds. Indeed, suppose that Π 1 ( s 1 ) = 0 . Then the lower part of the previous display combined with the definition of g imply Π 2 ( s 0 ) = Π 3 ( s 1 ) = Π 2 ( s 1 ) = Π 3 ( s 0 ) = 1 . But this is a contradiction to f ( s 1 , s 0 , s 0 ) = 0 . Thus Π 1 ( s 1 ) = 1 and the same reasoning implies Π i ( s 1 ) = 1 . Then the second equation in (81) directly implies Π i ( s 0 ) = 0 .

We now consider a symbol set

will encode the variables in the formula ϕ and

are auxiliary variables for the clauses. In addition, we need further auxiliary variables

We now add a set of equations that will ensure that for some value x i ∈ { 0 , 1 }

This can be achieved by adding the following relations for all i

̸

Note that the these relations indeed imply that Π a ( X i ) = Π b ( ¯ X i ) for a = b , a, b ∈ { 1 , 2 , 3 } . This implies then that (86) holds. We add the similar relations for C i,k which again ensure that

̸

̸

Now we encode the clauses of the formula. Consider a clause C l of the form x i 1 ∨ x i 2 ∨ x i 3 . Then we add the relations

where

For a given choice of Π we set x i = Π 1 ( X i ) . We now claim that the relations in the last display can hold if and only if x i 1 ∨ x i 2 ∨ x i 3 evaluates to true. Suppose the formula evaluates to false. Then

for i = 1 , 2 . Using (88) we find that Π 1 ( ¯ C l, 1 ) = Π 2 ( ¯ C l, 2 ) = 0 which implies

for any Π 3 ( o l ) . Therefore, the equations cannot all hold.

Suppose to the contrary that the formula evaluates to true. Let us first assume that Π 1 ( X i 1 ) = 1 . Then Π 1 ( ¯ X i 1 ) = 0 and we can set Π 3 ( C l, 1 ) = 0 which ensures that the first equation is satisfied (for any Π 2 ( ¯ X i 2 ) ). By definition of g we can also find a value c l, 2 such that for c l, 2 = Π 3 ( C l, 2 ) the relation f ( ¯ X i 2 , ¯ X i 3 , C l, 2 ) = f ( x i 2 , x i 3 , c l, 2 ) = 0 holds. Now Π 1 ( ¯ C l, 1 ) = 1 -Π 3 ( C l, 1 ) = 1 ensures that for any value of ¯ c l, 2 = Π 2 ( ¯ C l, 2 ) we can choose Π 3 ( o l ) such that the relation f ( ¯ C l, 1 , ¯ C l, 2 , o l ) = f (1 , ¯ c l, 2 , o l ) = 1 holds. The same reasoning applies for Π 1 ( X i 3 ) = 1 and a similar argument applies if Π 1 ( X i 2 ) = 1 . For clauses containing negations the same construction works except that ¯ X i j has to be replaced by X i j for the negated variables in equation (89).

Putting everything together, we have shown that for a given formula ϕ there is a set of relations of the form f ( n s ) = t s where n s ∈ S 3 given by (81), (87) (plus similar equations for C l,i ), and (89) which has the following properties. For an assignment x = ( x 1 , . . . , x n ) ∈ { 0 , 1 } n such that ϕ ( x ) = 1 evaluates to true, there is a map Π such that for all s the relations t s = f ( n s ) hold and where Π 1 ( X i ) = x i . On the other hand, if for some Π all the relations t s = f ( n s ) hold, then the Boolean variables x i = Π( X i ) will satisfy the formula ϕ . This ends the proof.

We now show how the problem of finding Π and g can be generally expressed as a constraint satisfaction problem. Recall that we want to find (if they exist) for a given K a projection map Π : [ N ] I → [ K ] I and g : K I → R such that for all given samples ( n s , t s ) the relation t s = g ◦ Π( n s ) holds. The general strategy is to introduce Boolean variables encoding the maps Π and g and then express all conditions as suitable constraints for these variables. We first introduce Boolean variables t i kn for i ∈ [ I ] , k ∈ [ K ] , and n ∈ [ N ] which are 1 if the token n is in cluster k , i.e., Π i ( n ) = k and 0 otherwise. Then the expression

for i ∈ [ I ] and n ∈ [ N ] is true if and only if n is assigned to to at least one cluster. In addition, we consider variables r k v for every k ∈ [ K ] I and v ∈ Im( f ) in the (finite) image of f . These variables shall encode whether g ( k ) = v is true or not. Then the constraints

for v = v ′ and all k ensure that each cluster is assigned at most one value v . Finally, we add for every datapoint ( n s , f ( n s )) and every k the constraint

This ensures that if Π( n s ) = k then this cluster must be assigned the value f ( n s ) .

We then consider the constraint satisfaction problem consisting of the ∧ of all the conditions in (94), (95), and (96). Then any satisfying assignment gives rise to a map Π and g such that g ◦ Π( n 2 ) = t s . Indeed, we set g ( k ) = v if s k v = 1 if such a v exists and arbitrarily otherwise. Note that for every satisfying assignment and every k there is at most one such v so g is well-defined. Moreover, we set Π i ( n ) = k for any k such that t i nk = 1 and at least one such k exists. On the other hand, we can easily construct a satisfying assignment given Π and g so that there are no solutions Π and g if the constraint satisfaction problem has no solution.

Note that we did not ensure that each token is assigned only to a single cluster but this can be achieved in post-processing or by adding additional constraints (such as ¬ t i k 1 n ∨ ¬ t i k 1 n ).

and therefore

Then relation (89) implies

̸

D Taylor Expansion of the Loss Gradient

The goal of this and the following section is to lay the groundwork for the proof of the main results.

The general strategy of our proofs is to Taylor expand the loss of each term around the mean token embedding w ( i, Π i ( n i )) to first order plus a remainder term. We can then extract the dominating terms of these expansions using concentration results for the datapoint statistics. It turns out that the linearized dynamics has favorable properties, while the remainder terms can be bounded. In this section, we derive the expansion of the loss gradient while the required bounds, in particular Theorem 7 and Lemma 1, are derived in the next section, Appendix E (they rely on concentration results which are deferred to Appendix G).

As pointed out above, the goal of this section is to Taylor expand the sample loss for one sample, where we expand the loss around the point w (Π( n )) . Let us first introduce some notation. Recall that we denoted by v ( n ) = ( v (1 , n 1 ) , . . . , v ( I, n I )) and by ¯ v ( i ) the concatenation of the embeddings ( v ( i, n )) n ∈ [ N ] . We define similarly

and we denote by ¯ δ ( i ) ∈ R DN the concatenation of δ ( i, 1) , . . . , δ ( i, n ) , i.e.,

and ¯ δ ∈ R IDN as the concatenation of ¯ δ (1) , . . . , ¯ δ ( I ) .

Consider a data-point n with Π( n ) = k . We now consider the derivative of the mean squared error of this term, i.e.,

We use Greek-indices for the latent space dimension. Fix an i ∈ [ I ] as the slot with respect to which we take the derivative. We then find for α ∈ [ D ]

Here we introduced the shorthand h ( v ( n )) for this function which also depends on i , k , and α . Now we estimate this function by Taylor expanding it around the point

Then we find the following expansion to second order

where R ( j 1 ,β 1 ) , ( j 2 ,β 2 ) denotes the remainder terms. We remark that if we assume that v ( i, n ) ∈ Ω for some convex set Ω and all i ∈ [ I ] , n ∈ [ N ] then (by convexity of Ω ) also w ( i, k i ) ∈ Ω and

Let us now (abusing notation again, it will be clear from context which function we refer to) write

We introduce some more quantities to get a concise representation of the Taylor expansion. We define the matrices D i,j L ( k ) ∈ R D × D containing the derivatives of a function L with respect to w ( j 1 , n j 1 ) and w ( j 2 , n j 2 ) , i.e., we consider

Similarly we define the vector D j L ( k ) ∈ R D by

For the diagonal entries D j,j we need a more fine-grained decomposition. Note that by the product rule we have

Using the notations we introduced we can therefore write (recall ˆ g ( k ) = ˆ f ( w ( k )) )

Collecting finally all remainder terms as

we can thus summarize the Taylor expansion result as follows

̸

Based on the expansion (110) we now want to get an expression for the gradient of the total loss on the entire dataset. We decompose the dataset as follows

̸

We find the following expansion

̸

We now rewrite or bound the four summands. First, we relate those expressions to the datapoint statistics matrices. We define

Note that if Π i ( n ) = k i then D k ,n,i = ∅ . We also define similarly

and

Then the first term equals

This implies

The second term is similarly given by

The third term is given by

̸

̸

where ¯ δ was introduced in (98). Let us summarize those findings as a proposition.

Proposition 1. Assume that the bound (15) holds for some Ω and v ( i, n ) ∈ Ω for all i ∈ [ I ] , n ∈ [ N ] . Then the following expansion holds

̸

Proof. The result follows from the calculations above.

E Bounds for the Loss Gradient

The goal of this section is to extract the asymptotically (with high probability) dominating terms of the expansion from the previous section as S →∞ . The crucial observation is that the appearing averages can be essentially replaced by their expectation. This is a consequence of concentration properties of the datapoint statistics. These properties will be derived in Appendix G where they are collected in Theorem 9. Here, we just prove the result conditional on the bound in Theorem 9. Recall that we defined δ max = max i ∈ [ I ] ,n ∈ [ N ] | δ ( i, n ) | in (11) as the maximal deviation norm and r max = max k ∈ [ K ] I | ˆ g ( k ) -g ( k ) | in (13) as the maximal residual. Finally, we introduce the notation

Note that B k ,i n ∼ Bin( S, p ( k , i )) if Π i ( n ) = k i and B k ,i n = 0 otherwise. In particular, E B k ,i n = S · p ( k , i ) . Note that if Π is approximately balanced if the following bounds hold

̸

̸

Theorem 7. Assume that the bounds in Theorem 9 hold and assume ˆ f satisfies Assumption 2 for some set Ω and v ( i, n ) ∈ Ω for i ∈ [ I ] and [ N ] . Then we obtain the following expansion for the loss gradient

Here the error terms are bounded by

Proof. The proof is a bit technical and essentially combines the assumptions and the concentration results in a straightforward fashion. First we remark that for n ∈ Π -1 i ( k ) we have as stated before that

and B k ,i n = 0 otherwise. Then Proposition 1 implies that (123) holds with the definitions of the error terms given below. We now define and bound the error terms in the decomposition. Note that we have

where we either bound both terms using Assumption 2 or the second term by r max and the first one by Assumption 2 together with M ′ ≤ M . This then implies that the error term capturing the fluctuations of the occupation statistics

can be bounded by

̸

Here we used that there are K I -1 non-vanishing terms in the sum (if Π i ( n ) = k i we have B k ,i n = E B k ,i n = 0 ). Next, we consider the self interaction error terms

For E S, 1 we get (similar to above) and using that the operator norm of a D × D matrix A is bounded by D max | a ij | ) the bound

The error term E S, 2 could also be absorbed in other terms later on but since it is not dominant, we just bound it. We can obtain similarly to our treatment of E S, 1 the bound

Next, we consider the interaction error term

̸

We note that E ( A k ,i,j n 1 ,n 2 ) = c for some constant c if Π i ( n 1 ) = k i and Π j ( n 2 ) = k j and 0 otherwise. This implies that if Π i ( n ) = k i

here we used

Recall that and

̸

which follows from the definition (10) in the last step. If Π i ( n ) = k i then this expression is clearly zero as the corresponding row of A k ,i,j is zero. Thus, we find that

̸

Here we used the submultiplicativity of the operator norm of Kronecker-products stated in (281). Now we use | ¯ δ ( j ) | ≤ √ Nδ max , the concentration bound (252), and bound ∥ D i,j L ( k ) ∥ ≤ DM similar as before to find

Finally, we bound the Taylor expansion remainder term

where h = h n,i,α was introduced in (100).

Assuming the bounds (15) and (16) we find

Indeed, the derivatives of h can be decomposed by the product rule in 4 terms which each can bounded using (15) and (16). Then we can bound the remainder term as follows

Recalling that D n,i defined (112) satisfies |D n,i | = B i n we obtain using the concentration bound (253) from Lemma 9

If we consider the gradient descent dynamics we obtain the following expansion for the time derivative ˙ w ( i, n ) . Let us from now on absorb the dependence on I and D in generic constants C .

Corollary 1. Under the same assumptions as in Theorem 7 we get

where E w can be bounded by

for S ≥ N ln( N ) .

Proof. The result follows from Theorem 7, the gradient dynamics (see (8)), and the definition of w ( i, n ) (see (9)). We also use the relation ∑ n ∈| Π -1 i ( k ) | δ ( i, n ) = 0 (already stated in (139)) to conclude that terms involving δ ( i, n ) cancel. The condition S ≥ N ln( N ) allows us to bound ln( S ) N/S ≤ 2 to control E T . Note that here we use that E S, 1 and E S, 2 can be bounded by E I . Note that we absorbed all terms involving δ max into the dominating term.

We therefore get the following bound for ˙ δ ( i, n ) .

Corollary 2. Under the same assumptions as in Theorem 7 we get

where E δ can be bounded by

for S ≥ N ln( N ) .

Proof. This follows from Theorem 7, Corollary 2, and the relation

Now we are in the position to control the time evolution of | δ ( i, n ) | 2 . We obtain the following relation Lemma 1. Under the same assumptions as in Theorem 7 the following bound holds

where C 1 is a constant depending on D , I , and M , and ω is defined by

Proof. We have

Recall that by (122) we have

We consider ω as defined in the statement of the lemma where ω ≥ 0 follows because we sum positive semi-definite rank one matrices. Then we find using the lower bound on the spectrum

where we used for the second to last contribution that we sum over K I -1 terms which is cancelled by the (2 /K ) I factor. We finally bound the error term by

|

δ

(

i, n

)

·

E

δ

(

i, n

)

|

This ends the proof.

F Proofs of the Main Results

In this section, we prove our main results, Theorem 4 and Theorem 5. In addition, we state and prove Theorem 8 which is a variant of Theorem 5 without a-priori bound on the embeddings. Essentially, the proof strategy for both results is to rely on the groundwork from the previous two sections, in particular on Lemma 1. Indeed, we first verify that the conditions of Lemma 1 hold for all times and then application of this Lemma allows us to control the evolution of δ max for all times.

Proof of Theorem 4. The proof proceeds in three steps. First we use the monotonicity of the loss to deduce that r max ( t ) can be bounded in terms of r max (0) and δ max ( t ) and δ max (0) . Then we show that choosing the variables as in the statement of the theorem allows us to bound all the terms in Lemma 1 in Appendix E. Then we apply Lemma 1 and deduce the decay of the maximum of the displacements | δ ( i, n, t ) | . Here we need to check carefully that the assumptions of the lemma are satisfied for all t .

First we assume the conclusions on the concentration properties from Theorem 9 in Appendix G hold which occurs with probability at least 1 -S -1 over the randomness of the training data.

Let us consider any time T > 0 and we assume that v ( i, n, t ) ∈ Ω for all i ∈ [ I ] , n ∈ [ N ] and 0 ≤ τ ≤ T . This allows us in particular to apply the Assumption 2 up to time T for embeddings v ( i, n ) .

We now implement the first step where we bound r max ( t ) for all 0 ≤ t ≤ T in terms of r max (0) , δ max ( t ) , and δ max ( t ) . The idea is to upper bound the initial sample loss ˆ L ( t = 0) and lower bound the loss ˆ L ( t > 0) . First we note that by a first order Taylor expansion and using Assumption 2

This implies that

Next, we derive a lower bound on the loss. Let k be such that | ˆ f ( k ) -f ( k ) | = r max . By assumption

and using concentration bounds (e.g., by summing (254) over n and using the lower bound for S ) we find that

Thus we can lower bound the loss using the same Taylor expansion as before by

Since gradient descent does not increase the loss the bounds (161) and (164) together imply for t ≥ 0

This implies

Then we get the bound

Now we choose C 2 and C 3 such that the initialization condition (20) becomes

and

Assume in addition that for all times 0 ≤ t ≤ T the bound δ max ( t ) ≤ ¯ δ holds. Using (167) we then find that for 0 ≤ t ≤ T

and similarly and therefore

The next step is to bound the various error terms appearing in Lemma 1. Using the last display we can bound for all i ∈ [ I ] and n ∈ [ N ] using Assumption 2 for t ∈ [0 , T ]

Moreover, we can estimate for 0 ≤ t ≤ T

Now we observe that for x = a ln 2 ( a ) and a ≥ 1 we find (using ln( a ) ≤ a ) the bound

Hence we consider

Note that tracking only the dependence on K and N (and using K ≤ N ) this condition reads

Then we find using (175) combined with the observation that ln( S ) / √ S is decreasing for S ≥ e 2 for S ≥ S 0 the bound

We note that similarly for x = a ln( a ) we get

Thus we consider similarly

where ¯ δ was introduced in (168). Note that tracking only N and K this bound becomes

Then, similar to before, (using monotonicity of ln( S ) /S for S ≥ e ) we find for S ≥ S 1

After all these preliminary estimates we can now finish the proof.

Let us define T ∈ R + 0 ∪{∞} to be the maximal time such that v ( i, n, t ) ∈ Ω for all i ∈ [ I ] , n ∈ [ N ] and 0 ≤ t ≤ T and the bound (168) holds for δ max ( t ) where we for convenience recall

We want to show that T = ∞ . We assume that T < ∞ and prove by contradiction that this does not hold. By (172) we know that for t ≤ T the bound r max ( t ) ≤ 1 holds and therefore Assumption 4 implies that w ( i, k, t = 0) ∈ Ω 0 and since moreover δ max ( t ) ≤ 1 for t ≤ T we find that v ( i, n, t = 0) ∈ Ω 0 + B 1 (0) ⊂⊂ Ω . By continuity of v ( i, n, t ) in t we conclude that there is ε > 0 such that v ( i, n, t ) ∈ Ω for all 0 ≤ t ≤ T + ε , i ∈ [ I ] , and n ∈ [ N ] . Thus, we can apply Lemma 1 for t = T and then find using Assumption 3 and the bounds (173), (174), (178), and (182)

If | δ ( i, n, t ) | ≥ 3 / 4 δ max we find | δ ( i, n, t ) | 2 ≥ δ max ( t ) 2 / 2 and we conclude that

Now, if δ max ( T ) ≤ ¯ δ/ 2 we conclude by continuity that there is ε > 0 such that δ max ( t ) < ¯ δ for t ∈ [ T, T + ε ] . On the other hand, if δ max ( t ) > ¯ δ we conclude that d d t 1 2 | δ ( i, n, T ) | 2 < 0 for all i , n such that | δ ( i, n, t ) | ≥ 3 / 4 δ max . This in particular implies that δ max ( t ) < δ max ( T ) for t ∈ [ T, T + ε ] and some ε > 0 ( δ max is non-increasing at T ). This is a contradiction, and we conclude that T = ∞ .

Finally, we prove the decay of δ max . Assume that for t ≥ T the bound min(1 , r max ( t )) ≤ R holds. Note that the function δ max ( t ) is not necessarily differentiable but its left and right derivative exist (as the maximum of finitely many differentiable functions). Then we obtain, similar to (184), for the right derivative the bound

For

the previous bound simplifies to

which implies that for t ≥ T by Gronwall's Lemma

Thus in finite time we achieve

This ends the proof. Note that we also get an exponential rate of convergence.

Remark 3 . The exponent 9 / 2 for the lower bound of S in (181) is not tight because we could use that r max is small, but this provides only a small improvement that does not justify the additional technicalities.

Before proving and stating Theorem 8 let us first prove that the weight decay λ > 0 allows us to derive a-priori bounds on the time evolution as stated in the following lemma.

Lemma 2. Let Π be approximately balanced, i.e., assume that Assumption 1 holds. Let M 1 > 0 be such that

Let 0 < λ ≤ 1 be a fixed number and assume that v ( i, n, t ) follow the gradient dynamics (8) and are initialized such that

and assume that the dynamics exists for all times. Then for all i ∈ [ I ] , k ∈ [ K ] , and all times

Proof. Note that by assumption

Then we conclude that

Thus we conclude that

This ends the proof.

Now we state a refined version of Theorem 5 which does not require an a-priori bound but instead proves boundedness of the embeddings v ( i, n, t ) for all times using Lemma 2.

Theorem 8. Let Π be approximately balanced, i.e., assume that Assumption 1 holds. Assume that the function ˆ f : R ID → R is slot-wise linear. Assume that there is M 1 > 0 such that

Let 0 < λ ≤ 1 be a fixed number and assume that for

Assumption 2 holds with Ω = B R (0) and some M > 0 . Note that M depends on R and thus on K and λ . Let C 1 be the constant from Lemma 1. Assume that v ( i, n, t ) follow the gradient dynamics (8) and are initialized such that

Then there are constants C 2 , C 3 ≥ 0 depending on M , I , D , λ , and C 1 such that for

the bound

holds for all t ≥ 0 .

Remark 4 . We emphasize again that since M might depend on λ we obtain no explicit rate in λ . However, the S and N dependence is the same as in Theorem 5.

Proof. The general strategy of the proof is similar to the proof of Theorem 4, but the proof is slightly simpler, as it is easier to obtain a-priori estimates on the evolution of v ( i, n, t ) . We assume that the conclusion of Theorem 9 holds, which occurs with probability at least 1 -S -1 over the randomness of D .

Now let T ≥ 0 be the largest time such that for all 0 ≤ t ≤ T and all i ∈ [ I ] and n ∈ [ N ] we have v ( i, n, t ) ∈ Ω and the bound

(we assume w.l.o.g. that C 1 ≥ 1 ) holds for 0 ≤ t ≤ T . By assumption, those relations hold at t = 0 (note that δ max ( t ) ≤ 2 max i,n | v ( i, n ) | ). We can bound for t = T using the a-priori estimate (193) from Lemma 2 and find

Thus, by continuity, there is ε > 0 such that v ( i, n, t ) ∈ B R (0) for some ε > 0 and all t ∈ [0 , T + ε ] . Now we can apply Lemma 1 for the interval [ T, T + ε ] . We note that by slot-wise linearity we have D i,i ˆ f = 0 and ω ≥ 0 . Therefore, we get

Now (with the same reasoning as in the proof of Theorem 4) there is a constant C 2 depending on D , I , and M such that for

the bound

holds. Clearly, we get for a suitable constant C 2

Now we can continue to bound (204) for S satisfying (205) and (208) and using (202) as follows

Now, if δ max ( T ) ≤ λ/ (8 C 1 ) then there is ε > 0 such that δ max ( t ) ≤ λ/ (4 C 1 ) holds for t ∈ [0 , T + ε ] . Thus, we assume that δ max ( T ) ≥ λ/ (8 C 1 ) . Let i ∈ [ I ] , n ∈ [ N ] be any index such that | δ ( i, n, T ) | = δ max ( T ) . Then we conclude that

This implies that

Similarly, we get for the bound

holds. And we find

from which we conclude that there is ε > 0 such that δ max ( t ) ≤ λ/ (4 C 1 ) holds for 0 ≤ t ≤ T + ε . In either case, we get a contradiction and thus T = ∞ . In particular, we can apply Lemma 1 for all times t ≥ 0 . Suppose now that for some t

then we conclude from Lemma 1 that

Thus we conclude from Gronwall's inequality that

This ends the proof.

Proof of Theorem 5. The proof is the same as the proof of Theorem 8 above with the only exception that we do not need to show that v ( i, n, t ) ∈ Ω but this holds by assumption which makes the proof strictly simpler.

G Concentration of Datapoint Statistics Matrices

In this section, we prove high probability concentration bounds for certain matrices and vectors capturing the statistics of the dataset. They are used in Section E to extract the asymptotically dominating contribution of the loss gradient and therefore of the gradients of the embeddings under gradient flow. All concentration bounds derived in this section are a simple consequence of the general and standard matrix concentration bound stated in Lemma 7.

We start with the matrices A ij ∈ R N × N with entries

the matrix counting the appearance of a pair of tokens in slot i and j (this is closely related to the matrices A k ,i,j which we will consider below). Then we have the following result.

Lemma 3. Let A i,j ∈ R N × N be as defined above. Assume that S ≥ 12 N . Then the following bound holds for η > 0

This implies

Proof. As in the proof of Theorem 2, we use Poissonization, i.e., we consider dataset ˜ D generated by first sample ¯ S ∼ Poi( S ) and then D ¯ S ∼ U ([ N ] I ) ¯ S . A uniform sample n ∼ U ([ N ] I ) satisfies n i = n 1 and n j = n 2 with probability

and the entries are independent. We can now apply Lemma 7 where we use η ln( S ) as η in the statement of the lemma and find using S ≥ N that

Therefore the distribution of A i,j is

Let us define

Note that P ( ¯ S = S ) = S S e -S /S ! ≥ 1 / (3 √ S ) . This implies

Here we used S ≥ 12 N in the last step. The relation (219) follows because ∥ E A i,j ∥ = NS/N 2 = S/N .

Next we prove a similar but slightly more involved result for the matrices that capture a similar statistics as A i,j but which in addition only consider n such that Π( n ) = k . We thus consider any k ∈ [ K ] I and any indices i, j ∈ I which we consider fixed for now. Consider the sets N ( i, k ) = Π -1 i ( k i ) , N ( j, k ) = Π -1 j ( k j ) ⊂ [ N ] . For a given dataset D = { n 1 , . . . , n S } we define the matrices A k ,i,j ∈ R N ( i, k ) ×N ( j, k ) by

for n 1 ∈ N ( i, k ) , n 2 ∈ N ( j, k ) . Thus, the entries are the number of datapoints that are mapped to k by Π and whose entries i and j are equal to n 1 and n 2 . We emphasize that this essentially agrees with the definition in (115) given in Section D (except that the dimensions of the matrices do not agree, we address this below the lemma). Similar reasoning as in the previous result, Lemma 3, gives the following statement.

Lemma 4. Let A k ,i,j ∈ R N ( i, k ) ×N ( j, k ) be as defined above. Assume that for all i ∈ [ I ] and k ∈ [ K ] the bound N/ (2 K ) ≤ | Π -1 i ( k ) | ≤ 2 N/K holds and let S ≥ 24 N . Then the following bound holds for η > 0

Note that in particular for K ≥ 2 we find for S ≥ 24 N the bound

Proof. The proof is essentially the same as the proof of Lemma 3 with some additional notational complications. Again, we consider a dataset ˜ D generated by first sampling ¯ S ∼ Poi( S ) and then D ¯ S ∼ U ([ N ] I ) ¯ S . Note that a uniform sample n ∼ U ([ N ] I ) satisfies Π( n ) = k , n i = n 1 , and n j = n 2 for n 1 ∈ N ( i, k ) and n 2 ∈ N ( j, k ) with probability

and thus the distribution of A k ,i,j is

for all n 1 ∈ N ( i, k ) and n 2 ∈ N ( j, k ) and the entries are independent. Note that the assumption N/ (2 K ) ≤ | Π -1 i ( k ) | ≤ 2 N/K implies that

which implies

Now we apply Lemma 7 where we use η ln( S ) as η in the statement of the lemma and find that

Let us set

Here we used S ≥ 24 N in the last step.

We also consider the matrix A k ,i,j ∈ R IN × IN already defined in (115), which we equivalently obtain by embedding A k ,i,j suitable, i.e, we have

Since A k ,i,j is a submatrix of A k ,i,j we find

so that Lemma 4 applies to A k ,i,j .

We also need a simpler concentration statement for the frequency of datapoints n s such that n i = n . We define the vectors B i ∈ R N by

We need a simple upper bound on the vectors B i .

Lemma 5. The following bound holds for η > 0

Proof. Note that B i n ∼ Bin( S, N -1 ) . The variance of a Ber( p ) variable is p (1 -p ) ≤ p . Bernstein's one-sided inequality then reads

Now we apply this bound with t = max(2 √ η ln( S ) SN -1 , 4 / 3 η ln( S )) which implies

As before for the A matrices we also need an extension of the previous result to a setting where we in addition require Π( n ) = k for a given k and n ∈ [ N ] . Recall that N ( i, k ) = Π -1 i ( k i ) . Then we consider the vector B k ,i ∈ R N ( i, k ) given by

Again, this agrees with the definition of B given in (114). The following Lemma holds.

Lemma 6. Let B k ,i ∈ R N ( i, k ) be as defined above. Assume that for all i ∈ [ I ] and k ∈ [ K ] the bound | Π -1 i ( k ) | ≤ 2 N/K holds. Then the following bound holds for η > 0 and n ∈ N ( i, k )

Note that viewing I and K as constant we get for S ≥ N · ln( N ) the bound

Proof. The proof is along the lines of Lemma 5 but slightly more technical. Note that the entries of B k ,i are distributed according to Bin( S, p ( k , i )) where

Since the variance of a Ber( p ) variable is p (1 -p ) ≤ p Bernstein's inequality then implies

Now we apply this bound with t = max(2 √ η ln( S ) Sp ( k , i ) , 4 / 3 η ln( S )) which implies

P ( | Bin( S, p ( k , i )) -S · p ( k , i ) | ≥ max(2 √ η ln( S ) Sp ( k , i ) , 4 / 3 η ln( S )) ) ≤ 2 S -η . (247) Applying (245) we find

To conclude that (243) holds, we just need to show that for S ≥ N ln( N ) the expression S/ ( N ln( S )) is bounded which follows by monotonicity of S/ ln( S ) in S .

Recall that we introduced B k ,i ∈ R [ N ] in (114) and then

For our analysis we want to summarize all those concentration bounds in one result. To simplify the statement, we only consider our main regime of interest.

Note that by assumption

Theorem 9. Assume that the bound N/ (2 K ) ≤ | Π -1 i ( k ) | ≤ 2 N/K holds for all i ∈ [ I ] and k ∈ [ K ] . Assume that

Then with probability at least 1 -S -1 the following bounds hold simultaneously for all k ∈ [ K ] I , i, j ∈ [ I ]

Proof. Applying Lemma 3 with η = 4 we obtain that the first bound holds with probability at least 1 -S -2 (here we used S ≥ 12 N ). Similarly, we obtain for S ≥ ( K/ 2) I -1 N

Setting η = 5 in Lemma 4 we get that the second bound holds (by the union bound) with probability at least 1 -3 I 2 K I S -3 for all i, j ∈ [ I ] and k ∈ [ K ] I . For S ≥ αN ln( N ) and N ≥ α we find

which implies for η = 2 and S ≥ 3 N ln( N )

and thus applying Lemma 5 with η = 2 implies that the third bound holds with probability at least 1 -S -2 . Finally, we find for S ≥ 3( K/ 2) I -1 N ln( N ) and η = 3 the bound

Lemma 6 implies that the last bound holds with probability at least 1 -2 IK I S -3 . All bounds simultaneously then hold with probability at least

H Spectral bounds for Random Matrices

In the derivation of the concentration bounds in the previous section, we needed concentration bounds for random matrices whose entries follow a Poisson distribution. In this section we provide the required result. Lemma 7 below should be folklore, but we did not find an exact reference so we provide a proof based on standard concentration results for random matrices. Let us first state the general result.

Theorem 10 (Corollary 3.7 in [26]) . Let X k be a sequence of independent random symmetric matrices with dimension d . We assume that there is a function g : (0 , ∞ ) → [0 , ∞ ] and symmetric matrices A k such that

Then for all t ∈ R the following bound holds

We now apply the previous concentration bound to our specific setting of interest.

Lemma 7. Consider a random matrix A ∈ R d 1 × d 2 whose entries are independent random variables with A ij ∼ Poi( λ ) for some λ > 0 . Then the following bound holds for all η ≥ 0

Proof. To bound the norm of this non-symmetric matrix we use the usual approach to consider

where E i,j is the matrix whose entry ( i, j ) is 1 and all other entries vanish and N i,j ∼ Poi( λ ) . Let X ij = ( N i,j -λ )( E i,j + E j,i ) . Note that by induction one directly finds

In any case we get ± ( E i,j + E j,i ) k ≤ E i,i + E j,j in the sense of symmetric matrices and we set A ij = E i,i + E j,j . Now we obtain for any θ ∈ R the relation

Note that the first summand is invariant under θ →-θ while the term in absolute values changes its sign. This implies that

Define

Consider the index set then we can write

Here we used the moment generating function of the Poisson distribution and in the last step that λ ≥ 0 and then for θ ≥ 0

Thus we infer from (268) and (269) that

using that A k ij = A ij in the second step. Note that here

which implies

Thus we can apply Theorem 10 and get

If max( d 1 , d 2 ) λ ≤ 1 we apply this bound with t = η +1 and θ = 1 and find

If max( d 1 , d 2 ) λ ≥ 1 we set θ = √ max( d 1 , d 2 ) λ -1 ≤ 1 and t = ( η +1) √ max( d 1 , d 2 ) λ and find using that for 0 ≤ θ ≤ 1 the bound e θ -1 -θ ≤ θ 2 holds

Thus we can combine both cases and obtain the relation

We observe that by (271)

Thus the same reasoning shows that

and since ∥ A -E A ∥ = ∥ Q ∥ = max | λ i ( Q ) | the claim follows by applying the union bound over (277) and (279).

I Kronecker Products

Let us here state some basic properties of the Kronecker product of matrices. For two matrices A ∈ R n × m and B ∈ R p × q the Kronecker product is defined by

We will need the property that the operator norm satisfies

Note that we use the notation v ⊗ w = vw ⊤ to denote the outer product. Formally we have vw ⊤ = v ⊗ w ⊤ but it is convenient and common to drop the transposed here.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: See main results and Section 3.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See discussions in Section 2 and Section 4 and Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All assumptions are stated and complete proofs can be found in the sumpplementary material.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
This paper studies a learning problem where we are given sequences of tokens and the task is to predict a label. At each step of the sequence, there is a different clustering of the tokens into classes and the output label only depends on the sequence of classes corresponding to the tokens. The task is then to estimate the clusterings at each of the time steps. The authors present several information-theoretic and algorithmic results about this problem and propose an embedding-based method for learning the clusterings.

STRENGTHS:
The authors propose an interesting and novel algorithmic problem.

The analysis of Section 3 is interesting. The authors derive some interesting limits regarding what is achievable in this problem.

WEAKNESSES:
The usage of the word "interact" in the abstract is rather vague.

The paper is poorly written with many language errors. Especially the paragraph from lines 69 to 81 is incomprehensible. 

In Theorem 1, the $M$ is not defined. Should this be $I$?

The proposed embedding approach seems somewhat unnatural for this discrete problem.

Gradient descent seems like a rather crude approach.

QUESTIONS:
Is assuming knowledge of $k$ realistic?

On line 132, it is written that order $N^I$ samples are needed to learn an unstructured map. Don't you mean $N^I\cdot I\log N$? For a uniform discrete random variable with $m$ possible values, it takes $O(m\log m)$ samples before all possible variables are encountered. 

From section 4, it seems that the method needs at least $O(DMI)$ memory. Is this scalable?

LIMITATIONS:
Assumption 1 is quite strong. In practice, one often encounters highly imbalanced clusterings.

The paper does not include experiments on real data.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis article studies the properties of tokens, i.e word embeddings in NLP, that are grouped in a small number of clusters. It proposes and analyzes a relatively simple model which is basically a composition of a clustering and a real function, but that shares many similarities with real world complex models. The authors tackle the problem from both complexity theoretic and information theoretic viewpoints (section 3), where they showed that $O(N\\log(N))$ samples are sufficient to recover the clustering map pi, and $O(N^2\\log(N))$ are needed in gradient-based methods. They also show that tokens belonging to the same class interact in approximately the same way with other tokens, and that their embeddings, i.e their vector representation, have the same dynamics (section 4).\n\nSTRENGTHS:\n1. The article gives very important insights on why word embeddings used in Large Language Models are successful. The authors made a solid theoretical study about the stated problem, and they covered it from different viewpoints. They also proved that we can recover the partition of data with a very good precision and with a relatively small number of data samples.\n2. This article can be a good starting point to investigate on the theoretical interpretability of word embeddings in LLMs for instance.\nAlso, the paper is coherent and the authors interpreted and commented all their results in a clear manner so that researchers from different backgrounds can read it without being lost in the formulas.\n\nWEAKNESSES:\n1. The authors also show that it is possible in theory to recover the partition map pi, but they don’t provide a clear algorithm or method on how to do so when their assumptions are satisfied. They also specified that it can be a very hard problem to determine this partition, which reduces the applicability of their findings, unfortunately.\n2. Some assumptions made in this paper are quite restrictive, for example, the authors stated that one of the necessary conditions to get a stable cluster structure is to have $S > N^2 \\log(N)$, which can be still a high number of samples: take the example of OpenAI’s tiktoken tokenizer “base200a” that has more than $2.10^5$ in its vocabulary size.\n3. This work lacks experimentation as these findings can be applied to some simple real-world examples, like to an encoder-only model that performs sentiment analysis ($f \\in \\{0, 1\\}$), but the authors kept it theoretical.\n\nQUESTIONS:\nI am curious why the authors couldn’t apply their approach to recover the clustering function $\\Pi$ for a simple real-world example, for example with an Encoder-only transformer (e.g. with one attention block) that performs some binary classification for instance.\n\nLIMITATIONS:\nThe authors have properly addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new learning problem: learning the partitions of tokens given sample sequences of tokens. The authors first study the problem from an information-theoretical perspective, where $\\tilde{O}(N)$ samples are sufficient to recover the partition for an alphabet of $N$ tokens. Then, they investigate the gradient dynamics of token embeddings and show a sample complexity $\\tilde{O}(N^2)$.\n\nSTRENGTHS:\n-  The problem formulation is novel and clear. Though it simplifies many empirical settings, it shares some similarities with them and can serve as a starting point for theoretical analysis.\n- The assumptions and main results are clearly stated. I appreciate the explanations provided after each assumption and result. While the notations, including several shorthands, are somewhat complex, they are generally easy for readers to follow. I also appreciate the efforts to condense the main results and proof ideas in the main text.\n- I appreciate the inclusion of several hard case examples (including Theorem 3 and Example 1) that illustrate the hardness of the problems.\n\nWEAKNESSES:\n- I appreciate the solid proof provided by the authors. However, it is not technically novel to the community. In Section 3, the paper assumes that the sizes of clusters are almost even, under which it can be expected that the sample complexity bound can be improved from $O(N^I)$ to $\\tilde{O}(K^IN)$. In Section 4, most of the analysis follows the usual path in analyzing dynamics. If there is more technical novelty, I would suggest emphasizing it explicitly.\n\n-  The problem setting is a bit restricted. It could be interesting to explore extensions such as: \n\n1) Allowing $f$ to be a function that depends not only on clusters but also on token embedding values, that is $f(n_1,\\cdots,n_I)=f(n_1,\\cdots,n_I,k_1,\\cdots,k_l)=g_{\\Pi}(n_1,\\cdots,n_I)$. An example is that when $n_1\\in k_1,n_2\\in k_0$, $f(n_1,n_2)=n_1$.  \n\n2) Introducing some dependency between tokens and potentially allowing one item to belong to several clusters.\n\n\nMinor Nits:\n- Page 2 line 74. “such that”?\n- Page 3, Line 117. Duplicates of “such as” and “e.g.,”\n- Page 4, line 135 “need of the order”\n- Page 5, equation (7), it should be $||u||$? u should be a vector.\n\nQUESTIONS:\n1. Page 4, Line 145, what is $N_0$?\n2. Page 5 in the section on analysis of gradient descent. Firstly, I do not see where you use $\\bar{v}(i)$ after defining it on Line 187. Secondly, what is $u$? I do not find its definition anywhere. One guess is that $u$ is essentially $v(n)$ from the context that follows. If so, I do not understand why scaling by $N/S$ is reasonable. Another guess is that $u$ is $\\bar{v}(i)$, then it seems odd when you discuss ‘samples’ because $\\bar{v}(i)$ is created as the concatenation across all possible N choices in the $i$-th slot, which essentially means having all possible samples at once.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper defines a learning problem for functions that depend only on a set of unknown partitions of the data. It establishes sample complexity, computational hardness and guarantees for GD-based algorithms under additional assumptions.

STRENGTHS:
The model is simple and clean, the presentation is clear without embellishments. The progression of guarantees is natural and well-motivated.

WEAKNESSES:
The assumptions required for proving GD are many and some appear to be strong and unverifiable (?).

QUESTIONS:
Is it possible to (approximately) verify whether a given data set satisfies the assumptions you need to have the GD guarantee?

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-17