Learning Partitions from Context
Authors: Simon Buchholz
Venue: NeurIPS 2024 poster (2024)
Keywords: clustering
gradient-flow
token embeddings
partitions
complexity
sample complexity
Primary area: learning_theory

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
Simon Buchholz
Department for Empirical Inference
Max Planck Institute for Intelligent Systems
Tübingen Al Center Tübingen, Germany
sbuchholz@tue. mpg.
Abstract
In this paper, study the problem of learning the structure of discrete set of tokens based on their interactions with other tokens. We focus on a setting where the tokens can be partitioned into 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 the class memberships of its arguments. The goal is to recover the class memberships of all tokens from finite number of samples of f- We begin by analyzing this problem from both complexity-thcoreti and ioormationn-hhoreee NP-complete in general, and for random instances, we show that on the order of N In(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 N2 (NV) samples.
viewpoints. We prove that it is
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 25 23 which encode their input by converting text into a sequence of discrete tokens that are embedded in high dimensional embedding space and those embeddings are fed into, e.g., 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.
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 Word2 Vec embedding vectors of the nouns 'king' and 'queen' approximately equals the difference of the embeddings of 'man' and 'woman' 21 122 Similarly, it was found that using word similarity as an inductive bias to structure latent spaces helps downstream performance 44. Thus, a properly structured latent space seems to be an important ingredient capture the intricate correlations in complex data.
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
particular focusing on shallow networks [2] [ and asymptotic regimes 29]. More recently, several works investigated transformer architectures (often focusing on the one layer case with linearized attention mechanism)
[16 [ 30 [2]. On technical level [1I] is closely related as they study the large depth limit of transformers through the lens of particle systems (however in their case time corresponds depth, while in our case it corresponds to training time). A feature shared by many of those works is that
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

little structure is assumed on the input, i.e., fixed token embeddings are assumed and often those are even assumed 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 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 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 small number of groups such that tokens within a cluster behave exactly the same, i.c., they interact in the same manner with other tokens.
Then the central questions is what assumptions allow us 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 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 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-oceurrence 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 logical inference problem: Given a set of facts about 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 -hhooreicc viewpoint where we show that it is in general hard, and from an information-thhome for N different tokens in the alphabet order N In(N) samples are sufficient identify the latent structure.
viewpoint where we show, roughly, that
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 22 In(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 finite set A is denoted by |A] and we also denote the standard Euclidean norm of any vector E R by lu]. The expressions Amax((A (A) and Amin (A) denote the largest and smallest eigenvalue of a symmetric matrix A respectively. We denote the uniform distribution over a set by (AA). For two subsets A, E B} their Minkowski sum. We denote the
A, B C Rdd we denote by A + B {a + b permutation group on h: elements by 65. 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 specific instance of a general problem scientific inference, where we want to uncover a hidden grouping of similar entities from sparse observations about these entities.

---

i= 2
i ==
3
3
8
4
2
3
8
7
9
1
6
6
6
2
5
5
Figure 1: Illustration of the setting for / 3 different groups clustered in 3, 2. and 3 subgroups respectively. Samples consist of one element of each group, the dashed lines indicate samples (1,3, 1) and (3, 7, 7,
Here main motivation, however, stems from the analysis of large language models where similar problem arises implicitly during training. They are trained by next token prediction, so we do not expect them learn structure by deductive reasoming 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, is a natural question whether gradient descent- training on token embeddings can uncover a hidden cluster structure of the data. Note that if the hidden structure recovered, then generalization to unseen prompts possible.
We now introduce our formal setup that captures key aspects of the discussion. We consider sets of Ni, N2. ..... Ny tokens or entities (such as words). For simplicity, we identify these with tokens from the set [N:]. For each of the sets [N;] there is a partition P; in K, classes which we can identify with
the set [ Then can encode the partitions through maps II : [N+] [K,] so that the partition
is given by P, (IT, (k:))).E(K.)) i.e., IT, encode the class membership. We consider the map
II I1, ... @ i..,, I[( n1)
(1)
This structure is illustrated in Figure Now we assume that there a function g [K1] x [Kg] x ... [Ky] + R which depends only the classes. The map f 9 II extends this map to the tokens such that it only depends on the group token belongs to. In the case where g (and thus f) maps to {0, 1) this can be interpreted (ni....,ny nr) is true if fn),....,, ny) = 1 and false otherwise and this case the main motivation of our work. More generally, f could output the index of suitable next token where in the 0, 1 case could correspond to a negation while 1 to an end of sentence token. Our goal is to learn the partitions P. or, equivalently, II up a permutation and thereby identify the hidden structure.
truth assignments, i.e., statement consisting of sequence
We assume that we are given data in the form of samples (n", (n²,f(nu fn=") where n' = (ni..... nf) and nf e [N.]. In other words, try to learn the underlying structure from the interactions of 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; N and K; K for some N. K and all 1 < i < I. Our main interest concerns the case where 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 II,(n) through its interaction with other tokens.
Interactions can be very complex, i.e., here the output depends on the interaction of 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 the tokens depending the broader context (c.g., group species by habitat, color, size each
3

---

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
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-Theoretie and Information-Theoreke Analysis
We now study this learning problem in different settings. Let us first briefly discuss complexity- theoretic and information-theare boundaries of this learning task. We first study the information-theoretic viewpoint, i.e., the question of how many samples are necessary to identify II (and potentially g). We focus on the case where we sample II and the data-samples uniformly at random. To learn an unstructured map [N]] > we generally need of the order N In(N4) independent samples (N4 when sampling without replacement).
properties of the learning problem to understand the general
For the structured setting we show that if II drawn uniformly at random then generally order K'N N (n) samples are sufficient to learn f and the partition induced by IT. In other words, for every token n; and each the K" classes II (k)) we need of the order of In(N) samples n such that [[n)) = k and
n4. In particular, for N >> K' any token will interact only with K In(N) < N other tokens, i.e., a very sparse subset of the other tokens.
We require the following necessary condition for identifiability; For every Ki + K2 there are ki,..., ki-1, ... ki [K] such that
g(ki...., , 1
ki Kett,...,,Ay +
,,, ks+1, kr)-
(2)
Note that if this condition is indeed necessary because if it is not satisfied then it is not possible to distinguish II," (k:) and II, (kt). Clearly, we can generally only identify II up to a permutation symmetry, i... We have the following result.
only find II such that there are permutations T. G such that II, æ,IL.
Theorem 1. Assume that g [K]' RR is function satisfying the assumption 2 Assume we randomly sample maps II; such that (T,)) = ke with probability K- for all i, ni and kG and such that (1(4)))eE
are independent. Assume we are given S samples (n, g o II(n)) where n ~ "( Then there a constant (I, K,n) such that with probability at least 2e for N > No[I, K,t) and
S 2114 22'
we can recover 11 and g up to permutations [K].
This result is a special case of Theorem 6 in Appendix B which shows similar bounds for arbitrary maps II that are not necessarily random. In the more general setting, there are additional dependencies on the size of the preimages I1, (k:). Note that this dependency cannot be avoided because if there 1 and g(k) 1 and g(k') 0 for k * k' then order N samples
is & such that (k)] are necessary to find n such that I1(m) = & and thus II. The general proof idea is to bound the probability that any fixed II' * 11 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 II and 11'. Then 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. be a function such that g(1, k2....,ky kT) g22,,2g k2....,kk k2) for all K1 1. Assume that N is divisible by K and that |11, (k)] = N/K for all i E [T] and l E [K]. Given 3 < S < InI-/In(NKK3
k2.... Ky E [K] except when 2 = =k3
samples
(n",g onn [[(")) where n ~ U([N)) ii.dd Then the function II is identifiable with probability at most
2e NN/K
The proof of this result can be found in Appendix B Next, we emphasize that while typically 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 map

---

=I 11 .... . 8I [N]' - [K]' such that f = go II given access to samples of the form (n", complete We show that this is true even if g known, I and K" 2.
Theorem 3. Consider the map 9 {0,1) 13= + {0, 1 given by g(ki, k2, k2) 1 Rk
(4)
Then it is an NC-complete problem to decide given samples of form (ni, nÿ, nz,4") E [[]s x {0. 1} whether there is map I1 111 @ I12 @ II3 with II, [NV] + {0, 1) such that t" go [I(n nz. ng))
for all samples.
The proof of this result can be found in Appendix C Now that we established under what the conditions II can in principle be learned, and clarified that this might be hard in general, we next discuss how we can find IT in practice. First, we remark that Theorem 3 rules out the existence of any general fast algorithms to learn II. Given the combinatorial nature and the hardness of the problem, it is natural to reformulate the task as 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 tinm for i E [7]. k E [K]. and E [N] which encode whether (n) = l and Tke for every & E [K] and v E (f) in the (finite) image of ʃ that encode whether g = It is then relatively straightforward to then express the conditions for the map IT as a constraint satisfaction problem which is satisfiable if and only if there are maps IT and g such that t g (Tn))) holds for all samples. We outline the construction in more detail in Appendix [ We leave the task of developing and studying efficient algorithmss 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 II can be learned using gradient descent on token embeddings. Our main finding is that for uniformly random II 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 start by introducing the setting.
Setting. We assume that we have token embeddings for each of the sets [Nr] to [Nr], i.e., we assume that there are vectors v(i, n) E IR D for some D and all i E [T]. n E [M]. Based on these embeddings, assume that we are given a function f f: R - R that transforms the embeddings into a prediction. We will abuse notation and write for n E [N]
f(n) f(r(1, mi).... v(I, nz)),
i.., 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.c., the loss of a sample (n, f = f(n)) is (f(n) f(n))?
We assume that we given dataset D {n!.... ..... n ~ ((N]').. Then the empirical loss reads (the division by 2 is convenient for the gradient evaluation later)
1 (((mm))((nn)))
(6)
We also define shorthands for certain concatenations of embeddings. For a sample n E [N]] we denote the collection of embeddings by v(n)= = (v(1, "(,),.,,,,IIIn))) Using the convention (5) we can then write f(v(n)) = f(n).
Moreover, we define 0(1) E DN as the concatenation of the vectors v(i, 1)...., e(i, n), i.e., the combined embedding for the i-th slot and IR. DNI as the concatenation of the vectors o(i) for 1 < i < I, i.e. all token embeddings concatenated. We consider the regularized loss given by
N [()) + 2 s
X 2
(x)
5

---

Note that the scaling by N/S (instead of usual 1/S) is natural because every token embedding v(i, m) occurs in approximately S/N of the samples and 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., consider (omitting the time variable from the notation)
N d ((b) Au(i, n). S du(i,n) This introduces time dynamics on the token embeddings. We indicate the time dependence by v(t, n, 1) but we drop / 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') co. To investigate this we define the center of the class embeddings by
d (i, n) do n)
R^
(8)
wii k.
Σ EEll ((k)
vii, n,,
II, (k)
and we consider the deviations from the class centers given by
8(1, n, () - v(i, n, () w(i, (n),
(10)
Thus the vectors ó(i, n, t) capture whether we recover the cluster structure, in particular if all norms 16(i,n, ((t,,,,)) t) are small then we essentially recovered the hidden structure. Therefore, we define
mmax (t) max max (ö(i, n, t)l-
(11)
sE|l] xE[N]
Similarly, to the notation introduced before we consider for Io E [K]] the vector w(k) (w(1,ki),.... wT,, k.)). As in (5) abuse notation and write
=
g(k) = f(w(k)) - ki),...... w k1)).
(12)
Similar to o(t) and ö we introduce wi) as the concatenation of (w(k, [){kk] and w as the concatenation (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 w(i, n²) are close initially whenever (1) [I,)'',, i.e., 80ax(() is small. In addition, we assume that lg(k) g(k) is small. To capture this we define
T Tiax (t) max (9(w(k,)) ()) g(k)].
(13)
KE[K]
Then the result shows that 6max stays small for all times if 2 N2 under mild additional assumptions. In other words, we start from the correctly learned substructures and g(k) g(k) for all & E [K]' then this remains true for all times. Note that while we phrase smallness as an assumption on the mean embeddings wiii ke) this generally a consequence of 8max small and a small empirical loss L(t). Let us now state the required assumptions.
Assumption 1. We assume that the map 11: [N]' that for all E [T], h E [K]
[K]] is approximately balanced which means
N 2K
2N K
S II
(14)
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 1I.
Assumption 2. We assume that there is convex set Q RD and constant M' such that the following bound holds
VM
sup AEEA
sup max (1f(0)). f(v)). 0,, 10,, 00, SM' in
(15)
4
Here is convenient introduce M 16M so that later certain errors a Taylor approximation are bounded by M. We also assume that
VM 4
max [g(k)) kE K]]
S M
(16)

---

This is rather mild assumption. For C functions f and St 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) RD for i [T], E [K] that satisfy
Tmax
max (f(w(())) - g(k) 1
(17)
k KK
the bound
Wp min Amin
Σ
(u(te))((((k)))
V Vulss)((w((k)) >0
(18)
ke(K] me
holds for some positive constant wo. Of course, the bound Tmax S 1 could be replaced by any other constant.
Note that this condition can only hold D - k!-1 i.e., the latent space dimension cannot be too large. The high level intuition of this assumption essentially that (at least if En (f(w(())) 2 small) there is no direction v E RD such that v (u((s),f(w((k)) se for all such that K; = k. i.e., we cannot If this condition does hold, then we cannot guarantee that (f(((w())) - f(k))² is minimized for 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 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 E [K] 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 wo [k)) XC.
single embedding without changing the output f(w for at least one &.
if
Assumption 4. We assume that for all collections of mean embeddings w(i, k) E TRD for i [7] k E [K] that satisfy
Tmax
max kE(K]
g(()]
(19)
there is convex set Do C Rdd such that w(i, k) E So for all i E (T], k E [K]. Again, the right-hund side of the bound Tmax < 1 could be replaced by any other constant.
Results. The first stability theorem can then be stated as follows.
Theorem 4. Let II: [N] [K] be approximately balanced as stated Assumption Assume
that the functions [K] + R and f IDID - IR satisfy Assumption 3 for some wn > O and Assumption 4 for some convex Sg. Assume that Assumption 2 holds for some M and the = 24 + (0). Then there are constants C1, C2. C3, CA O depending on T, M, D, and wo such that for all initial embeddings u(i, n,t 0) E Th. for E [T) and [N] satisfying
8 6max (0) S CCK--1112
Tmax (0) S CgK-31/2
(20)
and sample size
S2 C1 max 322 In) INNN)
(21)
the following holds with probability least 1- over the randomness of the dataset. When considering the gradient dynamics of the embeddings given by (8) the bound
R lim sassuprr mupp (t) S
(22)
t0000
holds and moreover
lim (tax < 1113112
31/2
(23)
10000
In particular 8max - if Tmmx (1) - O, i.e., all token embeddings for one fixed class converge to the same point.
7

---

This result shows that for order Nº In(N) samples and initialization sufficiently close to global minimum the cluster structure remains stable. We do not provide conditions that ensure T'ax (t) which guarantees convergence to loss and perfect recovery of the clusters.
Next note that cannot expect the clustering be stable in general even if mmax (0) is arbitrarily small if initialization is not close to a minimum. This is true even in the simplest case where I K i.e, we consider gradient descent for 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.
6-
Average distance embeddings.
2-
2
1-
50
150
200
250
Gradient stops
Figure 2: Simulation of the setting in Theorem 5 with N = 1000, - x 0. 100.000. (left) trajectories of 50 randomly sampled tokens from 6 different classes. (right) Average distance of token embeddings within a class for different classes (colored) and average distance between all pairs of embeddings (black).
Assume 1 = K= D = 1, N 1, f(x) = -x² x and f(n) -2 Consider the dataset = {I..... .,...... n... Assume that v n, 0) o(0,0²) for any o² > 0. Then the gradient dynamics introduced in reads
Example 1.
6(1, n, 1) = i' (v(11 ())()(((1,))))))))))))))))() n, ()) - (n)) n, [1 n, ()] small. (24)
We find that vLL, n, 2/3 (which local minimum of f t if v(1, n, t 0) 0. On other hand u1,, n, -1 (11, n, / = 0) < 0. So in this case (0) 0 /ln(N)) but (1) + 5/3 as oo. Slight modifications show that also δmax co is possible.
The previous example shows that without additional assumption, cannot expect recover the structure the data. Therefore, we impose additional restrictions on the function f. As apparent from Example and also from the bound in Lemma in Appendix E. is the curvature of the function 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 slot-wise linear if for every = (v(1),... ... v(T) E Dx/ and any i [T], a,3 E [D] the relation
d de (i) dus(i) f(v)
(25)
holds.
Let us denote for E RD by 5 D+1 the vector e where we append a 1. The most general slot-wise linear function is then of the form
f(v) (()(1) @ (()) 0(T))
(26)

---

where T: R(D+1) + R is 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 lI be approximately balanced, i.e., assume that Assumption 1 holds. Let f IR + R. be a slot-wise linear. Assume that v(i, n. t) follow the gradient dynamics (8) and that v(i. n,,) t) E S for all E [ E [N] and > for some convex set S2. Assume that Assumption 2 holds with constant M for the set Q. Let C₁ be the constant from Lemma Z depends on I, D. M). Assume that initialization
no 100,=00) 00] 8C1
(27)
Then there are constant C22 C3 depending on M, I, D, such that for
C2 N2K-1
In (N/A),), NK In(N/A))
$ max
(28)
the bound
(mmax(0)e N1/8 4C2K(--1/22
N V VIn(S) S
6max (t) S max
(29)
A
holds for all 2 O.
The high level summary of this result is that for order 22 In( 22In samples the clusters can be recovered Note that the a-priori assumption that the gradient flow is
up to an error of order X x "..IInnSN/// restricted to the set might appear difficult guarantee in practice. However, we conjecture that the results extend to gradient dynamics clamped at the boundary. Moreover, in Lemma we prove that the mean embeddings wi, k, t) stay within a ball radius = O(V for all times. This allows us to prove Theorem 8 which does not require any a-priori bounds on the evolution of v(ū,, n, () but comes at the price that the constants C1. C2, and C3 depend on so we cannot infer the explicit dependence. Let us make important remark.
Remark 1. While we state our results for a fixed function ff this function could in principle be time-dependent, e.g.. f could be given by 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 f. In particular, for Theorem 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 far show that
recover the right cluster structure in the sense that the embeddings from the same group cluster. However, this leaves open whether there any non-trivial dynamics at can be seen
all, i.e., all embeddings could cluster at the same point. This is in general not the from Corollary [ which states that in the setting of Theorem B and for large times the dynamics of the cluster-means follow the equation
V (i ((k))² Aw k) +0 (Vints) N
úb(i,, k) RE[K] No;
Σ
O
(30)
where (2K) -1 < 0k.s S (2/K) are positive numbers. This shows that w(i, k) follow generally 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 quick summary of the main steps of the proof and where find them. The first important ingredient an expansion of the loss gradient. We Taylor expand the loss of sample vnn) around the point w (T)) to second order with remainder terms, the relevant calculations can be found in Appendix [ (see Proposition [1 for the outcome). A second ingredient are concentration bounds for certain datapoint-statisto random variables and

---

random matrices. Those are derived in Section Gith 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 R). Moreover, we obtain such an expansion for (i(i, k) (see Corollary
and thus the displacements ((i,n) (see Corollary 2). This expansion can then be used to control 0g(5(2, n²| (see Lemma which is sufficient to control 6max-
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 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 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 has to be rather small to allow unique optimal solution of the token embeddings that allows us recover the cluster structure.
Conclusion
In this paper, we considered 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 problem in scientific inference. We studied this problem from different perspectives, namely from an drformation-theoretie descent based viewpoint. We found that order N In(N) samples sufficient to recover the partition for N tokens, while we showed that 22 22 (n(N) samples are sufficient for gradient based methods. There 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 natural question whether already RNN In(N)) samples are sufficient control 8max which is the information-theoreet threshold and would be similar to the optimal results for matrix completion [8 77). 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 (c.g., tokens belonging different clusters). Thirdly, it is natural question whether this work can be connected more closely to empirical findings.
complexity-theoree
and gradient
References
[1] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In seirth- Conference on Neural Information Processing Systems, 2023.
[2] Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(19):1-53 2017.
[3] William Bialek, Andrea Cavagna, Irene Giardina, Thierry Mora, Edmondo Silvestri, Massi- miliano Viale, and Aleksandra M. Walczak. Statistical mechanics for natural flocks of birds. Proceedings of the National Academy of Sciences, (09(13):4786-470) 2012.
[4] Piotr Bojanowski, Edouard Grave, Armand Joulin, and Tomás Mikolov. Enriching word vectors with subword information. Trans. Assoc. Comput. Linguistics, 5:135-146, 2017.
10

---

[5] Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for ConvNet with Gaussian inputs. Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th Inter- national Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 605-614. PMLR, 06-11 Aug 2017.
[6] Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent. CoRR, abs/2402.18824, 2024.
[7] 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.
[8] Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Pro- cessing Systems, volume 31. Curran Associates, Inc., 2018
[9] Jacob Devlin, Ming-Wai 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 (Long and Short Papers), pages 4171-4188 Association for Computational Linguistics, 2019.
[10] Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis. IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pages 12873-12883 Computer Vision Foundation / IEEE, 2021.
[11] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective transformers, 2024.
[12] Mor Geva, Roci 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 Punta Cana, Dominican Republic, 7-11 November, 2021, pages 1484-5495. Association for Computational Linguistics. 2021.
Natural Language Processing, EMNLP 2021, Virtual Event /
[13] Irene Giardina. Collective behavior in animal groups: theoretical models and empirical studies. HFSP J, 2(4):205-219, Aug 2008.
[14] Carla P. Gomes, Henry Kautz, Ashish Sabharwal, and Bart Selman. Chapter 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
[15] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bisachi and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[16] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 77822-37836. Curran Associates, Inc., 2022.
[17] 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 Lavri Thomas Wang. Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
[18] Raghunandan H. Keshavan, Andrea Montanari, and Sewoong Oh. Matrix completion from few entries. IEEE Transactions Information Theory, 56(6))9900-299) 2998 2010.
11

---

[19] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. 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.
[20] Song Mei, Theodor Misiakiewiez, and Andrea Montanari. Mcan-field theory of two-layers neural networks: dimension-free bounds and kemnel limit. CoRR, aws///9/..... 0tts/1992.06015, 2019.
[21] 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, Scoutsdale, Arizona, USA, May 2-4, 2013, Workshop Track Proceedings, 2013.
[22] 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.
[23] OpenAL. GPT-4 technical report, 2023.
[24]
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.
[25] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anm Lachaux, Timo- thé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.
[26] Joel A. Tropp. User-friendly tail bounds for sums of random matrices. Foundations of Compu- tational Mathematics, 12(4):389-433) 2012.
[27] 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, V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 6306-6315 6615, 2017.
[28] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[29] Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. H. Wallach, 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.
[30] 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.
12

---

Supplementary Material
This supplementary material is structured follows. We first review the notation used in the paper in Appendix [ Then we provide the proofs for the cnermation-theoreed results in Section [ in Appendix
B and the proof of Theorem [ and reduction to 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 [ and bounds for this expansion are derived in Appendix [ Based on these bounds, we can prove our main results in Appendix F An important ingredient in bounding the Taylor expansion
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
Overview of Notation used
Let us here collect important notation used throughout the paper and the proofs.
General notation.
Numbers up to n: [n] [l,...,m]
Eigenvalues of matrix: (A)
Largest eigenvalue: Amax(A)
Operator norm of a matrix: ||]||]
Notation used in the learning problem.
Number of slots:
Number of tokens for each slot: N
Number of classes: K
Number of samples: S
Map defining the partition in subclasses: = 111 @... & 1y [N]' [K]
Function on classes: g [K] + R
Induced function on tokens: f [N]' fR
Notation used the gradient descent analysis.
Dimension of latent space: D
Embedding for token n of slot i: u(i,n) n)
Token embeddings for sample n: v(n) = (1(1,n1), v(I, ni))
Token embeddings slot i: a(i) = (u(i, me)
All token embeddings: & (o(1))
Mean of cluster embeddings: w(i,k)
Sample version: w(k) = (wo11 ki),..., w k1))
Mean token embeddings for slot w(ii = k)),kk]
All mean embeddings:
(w (()))et
Displacements from cluster center: 8(i, n) = v(i, n) (I(1))
Displacements of a sample: 8(n) (8(1, ,.,,,,,)),,88 ,0,,,,,,,,)))))
Displacements of all tokens for slot 8(1) (5(i, n) mE(N]
Function on token embeddings:
IR, identified with f [N]' - R. i(n)
f(v(n))
Function on classes: 9 [K]| R. 5kk f (w)(k)
Regularization: X
13

---

B Proofs for Information-Theoreiis Results
In this section, we provide the proof's and additional results for Section Let first start to state the general information-thewe
bound that handles the case where [IT (k) might be of arbitrary size.
1I @.. ................ .III [N'' + [K]'. Assume there is L >> 0 such for every k E [K] we have II (k) 2 /L. Moreover, assume for all E [4] and & E [K] the bound [11, (k)l 2 M, i.e., every class has at least M members. Assume we are given S samples {(n", .9= l1((11)))) )) 8 E [S]} where n ~ ~U((N) If
Theorem Let g [K' + R satisfy ( and there is a projection
(KL+II. (M
S 822 2 max
max (In(K)IN) m) 2 LN max(2In(NVKII.))
(31)
for some 1i 2 2 then with probability at least - " we can recover 11 and g up permutations of [K]. In other words, for any map 11' 11,, @ ... @ 11, and g' such that g [l(m²)) g' o (1' There are permutations π. E GK such that 11, π. 11, and the corresponding relation holds for g. g'.
A few remarks explain this result are in order.
Remark 2. The scaling of S might appear slightly complicated, let comment this. We mostly interested in the case where N large and K stays bounded. In addition, we are primarily interested in the regime where S CK! bounded and M 2 eN/K (which holds for random II), i... the sampling probability of each class &o E [K] is of similar size. In this case, the first term in the condition (31) for S stays constant as N -> oc. The second term dominates, and we see that we need only O( (NN) samples to identify II and g. This is the setting studied in Theorem [ 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 (n)
_ for n = 1 and II, (1, = for n > and glk),..., kz)) = iff k- = ky 0 and g(ki,... kr) = 1 otherwise. Then we have L = N and we also need of order N' samples to sample the point (0, 0, ........). 0) E [ which necessary to identify IT.
Proof. Let us first introduce some notation. We denote the set of samples
[n n 5) D
(32)
Consider two partitions P. of a set [N]. We denote by Pd for A [N] the restriction of partition to subset (i.c., {P n A: P P)). The partition distance is defined by
D(PI, P2) min(1A9] PA PP A C [N]],
(33)
ie., the minimal number of elements that need to be removed such that the partitions agree.
We call IT' compatible with the datapoints D if there is g' such that g' o II' go IT on the data. Now the general strategy is to consider any other candidate map IT' and upper bound the probability that is compatible, i.e., for some function g' the functions f' = g'll and agree the data. We will then conclude by applying the union bound over maps II' Thus, we need to prove an upper bound on the number of partitions with partition distance at most A and a lower bound on the error probability for a given A. Let us start with the latter. Denote by
P. = (1,'(A) l E [K]]
(34)
the partition generated by 11, Consider any other map II' = II .... @ 11, [N]' + [K]* with corresponding partitions P. We define for any such 11' the quantity
A(II') max D(P,,P!) xE[7]
(35)
Assume now that
M A = A(TT' = D(PI, <
(36)
2
and in particular D(P,, P!) < A. Let P1 (PI.... Pk) where Pk I1 (k). Let A be set of
maximal size such that PA PlA After relabeling we can assume that P/ [PI.... Pk ) and
14

---

PA Pk A. By composing II; with permutation we can also assume that Pk = Moreover,
|P = 1P n 2 P |4 2 Pa - |P2122 - 1Pu122
(37)
The same applies to all partitions P. and P/ Next, we claim that if f g II and = o IT' agree
on all samples then with probability K'e शे.. over the randomness of the samples g = g'. Define Ex = 11' (k) n II (k)). (k). Using (37) and the assumption on II we have that for &o [K]
N'
E
II' (k) II (k)) 2 1-
(38)
L
The probability that none of the S samples is in ER: can then be bounded as follows
Ex
(1 2TT SE e
I(D Ex 0) Applying the union bound over we find that
ZT.
(39)
NT
(II' is consistent with DD for some g' $ g) S K'e ziz.
(40)
Now we bound the probability that g' = is compatible, i.e., go ll(nn) go (1²)))) ) for all E [S]. Let now ny E [M]. Denote by En, C [N]' the set of all vectors n (11, ......, n1) such that
g o l1(m)) + (I)
(41)
We now lower bound the size of Ex, under the assumption that is k [K] such that n1 E P n (where A is as above set such that Pi PlA). Then II' (mz) k and (by minimality of A9) find nt & Pk and thus 1I, = # k for some k. By assumption, we can find k ky such that (kc .g...., ..) * g(k, ......, , ky). Consider now any vector n (n1 Re...,,... such that ni IT," (k)) I (k:). Then, for any such n
go 11())= = g(k, ka,...,,ki) kI) + g(k, k2...., kT) = g o (1
(42)
we conclude that all such n are in E,, By assumption we have
N
[N]
(k))
(43)
TTT
im2
This implies that
N1-1 L
|mn1 12
(44)
and therefore
1
P(Em) 2 2'
(45)
The sets E,, are disjoint so we conclude that F. = UniE[N] Fr, satisfies
1 A
P(E) = Σ P(E,,) | LL
(46)
2L
Ti E[E
Now we can upper bound the probability that D is compatible with 11' and g by
PD (D compatible with 11', 9) S Po(E 0) (1 2L
(1-
A
S
Se
ZEN
(47)
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 (PI, Pi) & A can be bounded by (NK)A i.c., 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 < and by the union bound, we find that
[[(P).... Pf): max D(P.. & AH & (NK)AA
(48)
15

---

Thus we can bound the the probability that any of the maps II' (there are more such maps than partitions because there a label assigned to each class, but consisteney with the data depends only on the underlying permutation) with 0 < A(TI') S M/2 is consistent with the data is bounded by
M/2
Σ (NK) (e F C + K'K K' (49)
(3II' with A(I)) S M/2 consistent with D) S
Here the first term corresponds to the upper bound for the probability that (after permutation) g o 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 E [4] can be arg max, (P,,, P!).
We now consider the remaining case that there is an index i E [7] such that D(P;, P.) 2 M/2 where
we can assume w.l.o.g. that i = 1. As above, let Pi = {P...... Pk) and Pf = [Pi- p..... Pk } such
that PA nA P nA and A is a set of maximal size. We now claim that we can find an index h + k such that
M 2K
M
IP
and P
(50)
2K
Indeed, first assume that there is an index k such that |P n Pal S M/K. Since |P]| 2 M there is another index & such that PP 2 M/K. K. By maximality of A we have PANA= = PANA = P&nP and moreover
P + P N Pa
P P
(51)
because otherwise exchanging P/ and P/ would allow picking a larger A. By our assumption, conclude from here that
M
M
M
P2
Pa + PE P&
(52)
K
2K
2K
Thus, in this case IP/ n 2 M2(2 and PP n 2 M/K and 50) holds. Assume now to the contrary that there is no index K such that |P P < M/K. The assumption |44 2 M/2 implies that there & such that | n PI 2 M/(2K) and by minimality of A we find A A n Ps- and thus there is k of k such that [ P&l = A of (50).
M/(2)²). This finishes the proof
We now fix & and & such that |P& n P > M/(2K3) and IP PI 2 M/(2K3). Then, by assumption, we can find Kk E [K] such that & (k, k ......,,k) k7) and k (k, k2....,kki kr) satisfy g(k) * g(k).
Moreover [ Z /L/L. Now for every ky there is KG such that 11, (k:) n 11, (K)) 2
K K-I [. Thus, there is k' E [K]' such that
i 2=2
N
I (k)) (kx) 2 K-1+1
I (k:)
2
(53)
KI-1L
1=
Define A; = III (k() n 11, (k:) and A = P& PK and Ài = PB n P% Define A = Ai A2 x
x A, and Ā = A x A2 x
x Ay. Note that by (50)
M
A A N
(54)
2K
Clearly II' is constant (and equal (ke, K2 .g..., . k4)) on A n A and therefore also g' II' is constant for any g.. On the other hand, II(A) {k} constant and II(A)= is constant but
g o 1l(n) = g(k) + g(x)) g o 11(
(55)
for ñ E A and ñ E A. Now II' can only be consistent with the data (for any g') if there is no sample in A or no sample in A, i.c., and thus
PD((II'.g) gg) consistent with D for g¹) - Pp(DD A 0 or D Ā 0) D)
(56)
0) < Pp A = +P ( == )
16

---

where we used the union bound in the last step. Then we find using (54 and (56)
PD ((II' ,') consistent with D for a g') S 11
M
8M
2K1+IT
2
2xT
(57)
Finally, we can bound the number of partitions by the number of maps II, [N] + [K] which implies that
I[P...... Pr: : Ps partition of [AV in at most classes} S ) KNI.
(58)
So we get the upper bound
(D such that A(TT') 2 M/2 and II' compatible with D) S 21K KfiZ..
(59)
Combining (49) and (59) and using the union bound we find
PD((II,) g) not identifiable up to permutations) S PP 3I' with A(TI') > 0 compatible with D)
M/2
MS s 22 Ie KILLE
+1 (NK)
SA (e (e K'eett
(60)
Note that clearly the assumptions imply M N/K < N and thus e- (-8/(21)) I e -SSA LN) for < M/2. We then find (using the simple bound IK' & (2K)' & (NK) that
M/2
M/2
(NK) 14 (e LENN + L'Ke 1
2
e (2ln(NK) ViN I 4e -
(61)
for
$2 max(44l(NX))22)) LN, (n
(62)
and v) >> 2 where we bound the geometric sum by twice its largest term. The first summand in can be bounded by
660
MS NI
K 2e 2nn(KNII
KIFTI S 20-7 71
(63)
for
4 III LN
2 M
S S2 max
(64)
M
The proof of Theorem | is now direct consequence of the previous result because when assuming that a uniformly random map II is chosen we can estimate the quantities M and L in the previous theorem with high probability.
Proof of Theorem L We observe that |& (kx)) Bin(NN Applying Chernoff bound on the tail of the binomial variable we obtain
IP | (ka)
N
(1 -1))
Se N/K =
(65)
K
By the union bound we get
IP
min KE(K
|. (x:) < N\/2K)
S K''e - In(K) N4 <e
(66)
if N > No K, 7). Note that M 2 N2K2K) implies that L < (N/M) < (2K)'. Assuming that II is such that the bounds for M and hold we can apply Theorem 6 and find that for N 2 No sufficiently large (depending on K, 7) and
S 22 Nl((N)
(67)
(we bounded In (NK) < (ln()) the maps II and g are identifiable with probability at least e - Here we used that as N - oo the term indicated above is dominating (3). Now the union bound over the bad events ends the proof.
17

---

n²)) Proof of Theorem As before, we note D = =nn'. .....n²), Assume there is n2 such that II (n1)
1 and such that
Dn [m1) x _1," (1) x ... x 1 (1) = =0
(68)
Then IT' *1 IT is compatible with D where IT: = IT; for 2 and [T = IT, (n) for n + n and
(m1) = Th (n1) (by assumption on g). Let denote
An1 {mI} x (1) x .... x 17 (1)
(69)
Then
P(n" An,) NK
(70)
To estimate the probability of the event Un, (An, nD = 0} we use Poissonization to make the events independent. Consider datasets D whose distribution is generated by first sampling S Poi(2S) and
then conditional on S sample dataset D as before, i.e., n U([N)) for E [5]. Then we find that
28 NKI-1
|An nDI ~ Poi
(71)
and those events are independent for # n/1. Thus, we find that
Po([An, nD =0) =e x2
(72)
By independence now find that
P
n
(An, nD * 0}
= II
Po An, nD # 0)
(LElT,'(1)) (1)
MIE (1)
1 - NRT N/K
$ (1 - 1--e-In(N/RR) N/K (1-
VK VN
N/K
SE-VNIK
(73)
where we used the upper bound for S and (1-x) Se Assume D. u([NT'> and define
P. = PD.
n
(An, DD, + 0]
(74)
ml (1)
Note that ps is an upper bound the probability that IT is identifiable explained above. We have shown that
B(PPo122)) S NNKK
(75)
Clearly pa decreasing in s. This implies that
B(ppo((2)) 2 ps PP(Poi(2S) 2 S). A Chernoff bound for the Poisson distribution reads
(76)
P(Poi(A) xx) < (- 2
(77)
which implies with X 2S and x S that
P(Poi((S) S) (2Se)s -28
s
(78)
58
We thus conclude that
ps S 2c N//K
(79)
for S 2 3 (implying - (2/e)s > 1/2).
18

---

C Proofs of Complexity-Theoreti Analysis and Constraint Satisfaction Reduction
In this section, we provide the proof of Theorem 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 I1 and 22 there is 23 (21, such that
T(r1, .2, (11, x2)) O. Indeed, æa 22(21 x2) except for æ1 = 22 where 2z 1) has this property.
Proof. We reduce the problem to 3SAT. We consider an arbitrary formula
<(x) - ^ Ci(x)
(80)
on variables with m clauses C.. We denote f= g o IT. To improve readability, we consider symbol set S instead of [N] which is used for all slots. We proceed in three steps. First we show that we can write down a set of equations that ensures that two symbols so, $1 E S satisty I1,(s;) = for all (if f g II). Suppose the following relations hold
(so 50, 81) = [ 81, 9o) f(31, 80, 80) 0,
(81)
f(so, $1, $1) = f(1 80. $1) = f(s1, 80) = 1.
Then it is easy conclude that [1;(s;) j holds. Indeed, suppose that 11, (s1) 0. Then the lower part of the previous display combined with the definition of imply (l2(5)) = [I1(s1)
(s1) = (s) = But this a contradiction f(s1, Sg. su) -0. Thus II (s1) 1 and the
same reasoning implies (s1) = 1. Then the second equation in 81 directly implies (so) = 0.
We now consider a symbol set
= (5o, 8 WWU.UUUUUUUTTTTT
(82)
where
X' (x....... .x)
(83)
will encode the variables in the formula Φ and
C {C1,1. C1 C2,1, O2.2..... Cm,,, Cm,23
(84)
are auxiliary variables for the clauses. addition, we need further auxiliary variables
X = (X1,...,Xo, Xe),
Č = (1,,,, C2,, O2....., Cm Cm.), and O .o. ).
C1,2,
(85)
We now add a set of equations that will ensure that for some value E (0, 1)
11 (X,) 11, (Xx 34
(86)
(X,, TT, (X,) -x [3, = Tt.
This can be achieved by adding the following relations for all
f(X,, $1) = f(, = f(11, X,,X,) = 1,
(87)
f(s1, X,) f(X; $1. X,) = f(x 81, Xi) = 1.
Note that the these relations indeed imply that II, * (X,) for a + b, a, e (1,2.3). This implies then that (86) holds. We add the similar relations for CiA which again ensure that
(C, (C,,)) 4 11; (C.,.).
(88)
Now we encode the clauses of the formula. Consider clause Ct of the form Fi, V Fi, V Tig. Then add the relations
f(X,, Xizi C1.1) = 0 f(Xiz) Xis, C1a) 0
(89)
f(Cr1,, O12,0r) 1.
19

---

For a given choice of II we set 2; = 11, (X;). We now claim that the relations in the last display can hold if and only if Fi, Ii2 V Fig evaluates to true. Suppose the formula evaluates to false. Then
(X, = (X,, = =
(90)
and therefore
[x,(x,,= (X,(xx,)) = (1,(5,)) =
(91)
Then relation
89
implies
[C,()),) =
(92)
for = 2, Using 88 we find that II (Ct,) = (Ct.2) = 0 which implies
f(Cr1, C1,,,) 01) f( 0, oi) 0
(93)
for any (ot) Therefore, the equations cannot all hold.
Suppose to the contrary that the formula evaluates to true. Let us first assume that 11 (X,,) = 1. Then 11, (Xx,) 0 and we can set 1l (Ct,)) = 0 which ensures that the first equation is satisfied (for any (x,,,;)) )). By definition of g we can also find value Q, such that for c22 = (Ct,)) the ensures
relation/ f Xig: C1.2) = fZ Tag q,2) = 0 holds. Now II (Cr11 = 1 - (Cr that for any value of 6122 = (C2( we can choose (o:) such that the relation f(Ci.1, C2,,, 01) = 01,, 01) holds. The same reasoning applies for 11, (xxa) = and a similar argument applies if (X,,, ) 1. For clauses containing negations the same construction works except that X;, has to be replaced by Xi, for the negated variables in equation (89).
Putting everything together, we have shown that for given formula & there is a set of relations of the form s(n²) = where n² E S given by (TT) (87) (plus similar equations for Ct,). and (89) which has the following properties. For an assignment x (#1,..., ..n) E 13 such that <(x)= = 1 evaluates to true, there is a map 11 such that for all s the relations t = f(n") hold and
where (X,, z₁. On the other hand, if for some 11 all the relations t" = f(n") hold, then the Boolean variables x; = I1(X,) will satisfy the formula ø. This ends the proof.
We now show how the problem of finding II and g can be generally expressed satisfaction problem. Recall that we want to find (if they exist) for a given K a projection map I1: [N]' + [K] and g : K' - R such that for all given samples (n²,, (t) the relation t go (Il((n)) holds. The general strategy is to introduce Boolean variables encoding the maps IT and and then express all conditions as suitable constraints for these variables. We first introduce Boolean variables
constraint
tin for i E [7], & E [K]. and n E [N] which are I if the token n is in cluster k i.e., [n) = k and 0 otherwise. Then the expression
in V E2 V... E
(94)
for E [7] and n E [N] is true if and only n is assigned to to at least one cluster. In addition, we consider variables Tik for every & E [K] and v E Im([) in the (finite) image of f. These variables shall encode whether g(k) = v is true or not. Then the constraints
Tk V -Tke'
(95)
for e * vl and all l ensure that each cluster is assigned at most one value v. Finally, we add for every datapoint f(")) and every l the constraint
-t
kinng V... V -hing V Vraf(nn)-
(96)
This ensures that if ITn")) = Is then this cluster must be assigned the value f(n²).
We then consider the constraint satisfaction problem consisting of the A of all the conditions in 994 (95), and (6). Then any satisfying assignment gives rise to a map 1I and g such that g o 11(n²) t.. Indeed, we set g(k) satisfying assignment and every there is at most one such u so g is well-defined. Moreover, we set (n) = k for any i such that t = 1 and at least one such & exists. On the other hand, we can easily construct a satisfying assignment given II and so that there are no solutions II and g if the constraint satisfaction problem has no solution.
if sku if such u exists and arbitrarily otherwise. Note that for every
Note that we did not ensure that each token is assigned only to a single cluster but this can be achieved in post-procssiing or by adding additional constraints (such -th, V -th -
20

---

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 proof's is to Taylor expand the loss of each term around the mean token embedding wi(, (n.)) 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 Lemmall II 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 Taylor expand the sample loss for sample, where we expand the loss around the point ([][))).. Let us first introduce some notation. Recall that we denoted by v(n) (v(1, ni)),..., ((I,n)) and by o(() the concatenation of the embeddings (u(1, ne)] We define similarly
(()) = (8(1, ................ 8(1, n2))
(97)
................................ and we denote by 8(1) E RDN the concatenation of 8(1, 1)...., 8(i,n), i...
(i) = ()),,n))ne(M)
(98)
and 8 E RIDN as the concatenation of 8(1, . 8(7).
Consider a data-point n with [[())= = k. We now consider the derivative of the mean squared error of this term, i.e.
[(n) = =
f 9(k)))
(99)
We Greek-indices for the latent space dimension. Fix E [4] as the slot with respect which we take the derivative. We then find for a E [D]
(i n2,) (f(o(m)) - ((k)) 2 d
d
dva (i, ni) (f(v(n)) - s(w)))
-(aad(...)) f((())) (f(o(m))- h(v(n)).
(100)
Here we introduced the shorthand (n))))) for this function which also depends on i, k,, and a. Now we estimate this function by Taylor expanding it around the point
w (k) (w(1, k.),..., wr(T k,)))
(101)
Then we find the following expansion to second order
(x((n) = hkw(k)) + Σ
Σ dwa(j. k,) ((w(k)))og() nz)
jEl RE|D
(102)
Σ
Σ
RuLLo ( (1)(k))) (w)() (i,, n dd (i2. n2))
didhe (.D2E(D)
where RG..9.).2.dd denotes the remainder terms. We remark that if we assume that u(i, n)) Ω for some convex set Ω and all E [T], n E [N] then (by convexity of ②2) also w(i, é, ky) Ω and
d duis (ii, njj)) duys (i2, n h(v)
d
max went
(103)
Let us now (abusing notation again, it will be clear from context which function we refer to) write
L(k) = 2
(f(w()) g(k)))
(104)
21

---

We introduce some more quantities to get a concise representation of the Taylor expansion. We define the matrices Di,5 (kk E RDD containing the derivatives of function L with respect to w(i nz and w(22, /2 i.e. we consider
DD L(k)) = dwa kij dws((i. k(a))
L(k).
(105)
Similarly we define the vector Dj E ERDD by
(1,L(H))a dw, (j. kj) L(k).
(106)
For the diagonal entries Dj we need a more ane-ggained decomposition. Note that by the product rule we have
d
d L(k) (dwa (j. k)) dw (i k,) ((co(k))) (f(w(())) g(())
dw (j, k5) dw (j, k)
(107)
di k.) f(wo(k)) (dwo(i, k.)) f(w((k)))
d
Using the notations we introduced we can therefore write (recall ÿ(k) f(w(k))) L,5 = (P),j0(k)) (0(k) g(k)) + (2ū(k) @ D)
(108)
Collecting finally all remainder terms
Ra = Σ EE 2a.60 (n (w)k))& (f1, n,, )69 (22, nj2)
(109)
we can thus summarize the Taylor expansion result as follows
D.L(n) D,L(A) + ((D),,((k)) (Ā(k) - g(k)) + Diū(k) @ ((i,
Σ [(P))L(4))) - R(i,
(110)
Based on the expansion 110 now want to get an expression for the gradient of the total loss on entire dataset. We decompose the dataset as follows
Dk. {n" x E D D[[n(")) = k, n n).
(111)
Note that if (n) + k2 then Dk.n. = D. We also define similarly
D {n² E D n = n...
(112)
We find the following expansion
d ((0) = do(i,n) n)
d
Σ L(n') Σ Σ
d L(n)
duli, m)
KE(K] nEDki.. du n1)
Σ Σ D,I(K) ((Dr,)(k)) (0(k) g(k)) + Değ(k) @ Dig(k)) 5(i, n.)
kE(K] DiDiim
Σ
Σ
((D. + Σ R(i,n)
KE[K]' DED
1#1
nED
(113)
We now rewrite bound the four summands. First, relate those expressions to the datapoint statistics matrices. We define
Bk. ID
In D: 11()) k. n, =
(114)
and
A k.i. n,n
= In E Diemiinn = n2) = I(n E D: [[n) = k. n, = n, n,= =n
(115)
22

---

Then the first term equals
Σ ΣΣ
Σ D2I(k)=
E
B D,I(K)).
(116)
KE MEDA...
KE[K]
The second term similarly given by
Σ Σ
(()) (Ø(k) - g(k)) D.ğ(k) @ ((
kE[K]] nEDi...
(117)
Σ B (k),)(())) (0(k) g(k)) Dig(k) Dig(k) ((2,)) kE(K]'
The third term is given by
Σ Σ [()i)L())))(
ΣΣΣ
A (D),L())))((.n')
KE[K] nEDiin jzs
kE|K j WE[N]
Σ Σ KE[ jyt
(Ahrhd & Lis 8(1)
nD:(n+1)D
(118)
where 5 was introduced in (98) Let us summarize those findings proposition.
Proposition 1. Assume that the bound (55) holds for some Q and v(i, n) E for all [4] n [N].
Then the following expansion holds
d dv(i,n)
L(0)= = Σ D.I(k)
RE[K]
+
Σ
Bhi ((De,)(k)) (Ø(k) - g(b)) + Dig(t) @ D.)(k)) 8(1,n)
kE(K
(119)
Σ
Σ
(Ahhbd 60 Dey T(k)) 8(1) nD((++1DD
kE(K
tXi
Σ
R(i,n)
Proof. The result follows from the calculations above.
E Bounds for the Loss Gradient
The goal of this section extract the asymptotically (with high probability) dominating terms of the expansion from the previous section as 5 + co. The crucial observation is that the appearing averages can be essentially replaced by their expectation. This is consequence of concentration properties of the datapoint statistics. These properties will be derived in Appendix G where they Here, we just prove the result conditional on the bound Theorem 9
are collected in Theorem /
Recall that we defined 8max max [N][N] jō(i, n)) in (1] the maximal deviation norm and = ma (Ø( g(k) in (13) as the maximal residual. Finally, we introduce the notation
III (K))
p(k, i) = N' (k.))
(120)
Note that Bl... Bin(S, p(k, 1)) if (m) = k; and B = 0 otherwise. In particular, EB S p(k, i). Note that if II is approximately balanced if the following bounds hold
ITT (k)) N' I (k,))
11, (k2)
Np(k, N
-II J
(121)
N
This implies
1
(2)
2
< Np(k,
(122)
2K
K
23

---

Theorem 7. Assume that the bounds in Theorem 9 hold and assume f satisfies Assumption 2 for some set S2 and v(i, n) for [1] and [N). Then we obtain the following expansion for the loss gradient
N
(
S du n)
E
Np(k, (Ik)
kE (eKK/kk (m)
(123)
E
Np(k; i) (b(k)) - g(k)) + -iū(()) @ Diū(k))) 8(1,
+ E (i, n) + (i,
(i,n) n) + ES2 (i. + E (i, + E (i.
Here the evror terms are bounded by
F. n)] 1 4M min(1, mmax) (2K) /I((5)
N
(124)
5
|02 ni, < 4DM min(1, (2K) (2K) 11-112
|I((S)
6
max
(125)
5
| (i, m)] < 4DM (2K) 11/2
V (In(S)
mmax.
(126)
S
N2
IE' (i, n)] 6DM (1-1/2 in(s)
8max
(127)
max (1+11) In(S)
F (i, m) s MT D8
(128)
The proof is a bit technical and essentially combines the assumptions and the concentration results in a straightforward fashion. First we remark that for E (k, we have as stated before that
Proof.
Bp S S-p(k, i)
(129)
and = 0 otherwise. Then Proposition limplies 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
d
d
((D)(())) du. k.)
L(k)
du. (j. k,) g (k))
g(k)) < M mnin(1,/max)) (130)
where either bound both terms using Assumption 2 or the second term by Tmax and the first one by Assumption toogetherr with M' S M. This then implies that the error term capturing the fluctuations of the occupation statistics
E²(in))=
N
Σ KE[K]'
EDD D2I(K)
(131)
S
can be bounded by
V [n(()) N
M)),K, In(s) N S
E (i, < 4M min(1, (2K)
(132)
Here we used that there are K!-1 non-vanishing terms in the sum (if [1, A k; we have Bl E. 0). Next, we consider the self interaction error terms
N
ES. (i,
Σ
Bk.. n
B.B(..)) n (ū(k) - 9(k)) (D.,((()) 8(1,
(133)
S
KE[K]
N
Σ
(Bh.
EBA (Dig(k) @ (.
(134)
S
kE[K]'
24

---

For ESS we get (similar to above) and using that the operator norm of a D x D matrix A is bounded by D Daxx [ay) the bound
1 (i, & 4DM min(1, Tmax) (2K) V (n((S)
N
max
S
(135)
N
S C(I, K, M, min(1, m'ax mmax)
In(S)
S
The enror term S could also be absorbed in other terms later on but since it is not dominant, we just bound it. We obtain similarly to treatment of FSS11 the bound
(i. n)] & 4DM (2K) 1)/2
[m(s)
N
max
S
(136)
S C D) maxx V [I(S)
5
Next, we consider the interaction error term
(i.
Σ Σ kE(K!
(A k.dd @ DiyL(k)) 8(1) nD((++1jD
(137)
S
jzi
We note that TE( kkk. 11. for some constant (I, = k; and IT (12) k. and 0 otherwise.
This implies that if (1) = k;
( [EA [EA @ Dik)L(k)) 8(1)
Σ
L(K))(i,')= =
(138)
DD((++))D
m'nr'((y)
here we used
Σ
n(i
(139)
which follows from the definition (10) in the last step. IT, (m) * k, then this expression is clearly zero as the corresponding row of A kad is zero. Thus, we find that
E (i,
N
Σ
Σ
E. & L(k)( 8(1) nD((x+1)D ((A
=
S
kE(K]]
jæi
(140)
N
K 1--
max A kki.j k...j
EARA hii
[()) max ((()]-
Here we used the submultiplicativity the operator norm of Kronecker-product stated in (81).
Now we use [Ø())| s Vm6maxx the concentration bound 252 and bound [ s DM similar before to find
N
LE' (i, n)) = 6DM (2K) (1-)/2 In())
8
(141)
S
Finally, we bound the Taylor expansion remainder term
N E ET(i,n)=
Σ R(i,n).
(142)
S
nED...
Recall that
Ro (i(in) = Σ Σ
(k(k)) (1 (1) (i. B
paá
(i,,,,))8 16, (Jz, ni2))
(143)
and
d 2 dus, (51, n2,) du (1z. n2,))
d
max weft
h(v)
(144)
25

---

where h = h was introduced in (1OO).
Assuming the bounds (5) and (66 find
S max 1 d
v
M
VM
M
RR
wEn 2 du (1) n,,) de (J2. nzz)) h(w)
M.
(145)
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
[ia n)) M
Σ
Σ
168, (i,, nz 16 (2. n2,)1 S MID Σ n8G). (146) JE(M
Recalling that D defined 253 from Lemma
(122)
satisfies (Da.j| = B we obtain using the concentration bound
E (i,
N 5
Σ n) nEDMa
S
N B max S D5 max
(1+
In(S)N
S
(147)
If we consider the gradient descent dynamics we obtain the following expansion for the time derivative vi(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
d(i, k)
Σ
Np(k, ((L((() Awii, k)) E" (1,
(148)
KE
where E" can be bounded by
E. Ei, (i, C mmin(1, Fmax
) ((-11)/2
CKUU-1)2 In(S)) /In(S) (/)()
N
N2
s 8 + C82 C2max
(149)
for 2 (IN)...
Proof. The result follows from Theorem 7 the gradient dynamics (see (8). and the definition of w(inn) (see ()) We also use the relation Σn '(k) ò(i, n) (already stated in (13)) to
conclude that terms involving 8(i,n) n) cancel. The condition 2 In(N) allows us to bound In(S)N/S 2 to control ET. Note that here use that ES and E can be bounded by E Note that we absorbed all terms involving 6max into the dominating term.
We therefore get the following bound for ó(i, Corollary 2. Under the same assumptions as in Theorem we get
ô(i,n) Σ
Ny(k, i) (0(()) ,(k))) + Dig(k) @ D:9(s)) 0(im))
KEIKT' K.=n
(150)
x8(i,n) n) (i,
where E can be bounded by
N
N2
E (i,n)) S (in (1-1)//2 [n(() + CKU
8max + C82
S
max
(151)
for S In(N).
26

---

Proof. This follows from Theorem
7 Corollary 2 and the relation
8(i, n) = ö(i, n) - wi (n)).
(152)
Now we in the position to control the time evolution of n))2. We obtain the following relation Lemma 1. Under the same assumptions as Theorem the following bound holds
de 2 8(1, 5(i,n)) (
W + 2 sup max
[ Po, ;n). U rmax r))
(2K
+ C1 min(1, I Fmax ) In(s) mmax
(153)
S
N2
-K In(e))/II((SS)))))) In(S))
6 C16
S
where C₁ is a constant depending on D, I, and M, and w is defined by
min Amin
Σ
Dig(k) @ Diū(k)
2 0.
(154)
R,S
KEIKT &
Proof.
We have d
dt 8(1, n)12 5(i,, ((((m)
Σ (D(K)))(((()) (kk)))(i,n)) (Di,g(k)) ö(i,n)
kE(K) kin
(155)
Σ
(p(k, D.ğ(k) s D.0(k))(i,n)
kE[K] kinn
5(2, m) E (i,m) 11)(
Recall that by (122) we have
2
S Np(k,
(156)
2K
K
We consider w as defined in the statement of the lemma where w 2 0 follows because we sum positive semi-definite rank matrices. Then we find using the lower bound the spectrum
ni,
Σ
Diğ(k) & Diū(k) 6(i,n)
KE|K] =
(157)
2 (0(i, 2Amin
Σ
2 w|8(i, Dig(k) @ Diū(k)
KE[K] R.=n
Thus find d
dt 2 5(1
-18((,
w (2K) 18(2, n²)²
(158)
+ 2
(0mmmaxx(( sup max DP D2. f(r un [| 8(2, n))² - 8ii (i) EEER
where we used for the second to last contribution that sum over K!-1 1-1 terms which is cancelled by the 12/K²² factor. We finally bound the error term by 18(i,)) (i,n)
N Smax -CK(U-1)/2
N2 V In(()
S CVD min max )K( (1-1/) (n(5)
1)/2
82 8 C83 max (159)
S
N
.2 (I((S)
< C Cmmi(1, K
In(s)
6mxx + CK(1-11/2
8 + C83 -
S
27

---

This ends the proof.
F Proofs of the Main Results
In this section, we prove our main results, Theorem and Theorem [] 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 to rely on the groundwork from the previous two sections, in particular on I. Indeed, we first verify that the conditions of Lemma[hole for all times and then application of this Lemma allows us to control the evolution mmax for all times.
Proof of Theorem The proof proceeds in three steps. First use the monotonicity of the loss to deduce that r mmax (t can be bounded in terms (0) and 5max(t) and maxx Then we show that choosing the variables as in the statement of the theorem allows to bound all the terms in Lemma Tin in Appendix [E] Then we apply Lemma [ and deduce the decay of the maximum of the displacements (ō(é, n, t.]- Here we need to check carefully that the assumptions of the lemma
satisfied for all
First we assume the conclusions the concentration properties from Theorem 9 in Appendix hold which occurs with probability at least S over the randomness of the training data.
Let us consider any time T > 0 and we assume that v(i, n,, () E Ω for all i E [4], n E [N] and T T. This allows us in particular apply the Assumption up time for embeddings n).
We now implement the first step where we bound r'ma (1) for all < l < T in terms of m'mx (0), m (1), and ma)))). The idea upper bound the initial sample loss -(t = 0) and lower bound the loss L(I 0). First we note that by first order Taylor expansion and using Assumption
[f(v(n)) f- f(w((I())) s Σ M (1, M MI6
(160)
RET
This implies that
( Σ (f(v(n)) ((I(())))²
2 ((I()))) + Σ Mij(i,
Σ [[(([[[n))) NED
1
( IIM IM8
SET
(161)
Next, we derive a lower bound on the loss. Let & be such that (f(k) - [(k)) = -maax- By assumption
N
|I
[
(162)
2K
and using concentration bounds (e.g., by summing (254) over n and using the lower bound for S) we
find that
1 S
D In D: [[n) k)
(163)
2
(2K)
Thus we can lower bound the loss using the same Taylor expansion as before by
1 S L(e(()) 4 (2KY (max((,, Tiax (1) (()mmax(()))²
(164)
Since gradient descent does not increase the loss the bounds l66 and (144) together imply for t 2 0
1 S(r (rmax (0) + IM max: (0))2 ((v(O))) > ((o(t))
2
(165)
1 S
(max(0, (max(() (t) M16 (t)))²
4 (2KY7
28

---

This implies
Tmax (t) - 2(2 + max (0)) + IMomaa (t).
(166)
Then we get the bound
Tmax (1) S (22 + IM max(6max (0), max (t)))
(167)
Now we choose C2 and C3 such that the initialization condition 20 becomes
(2K()))) 62 ((2 1)
WD
wi
max (0) I min
=:
(168)
and
Wp
Timx (0) S min 64MD(2K))u/22' 8(2K)
(169)
Assume in addition that for all times 0 s t s T the bound m (t) S holds. Using 167 we then find that for 0 S ST
Tmax (1) S (2K))122
TM
we 64TA72(22)33/22)
WD
313/222
82D(2K) 2
(170)
and similarly
T mmax (t) S 4(2K)1/2
+ IM
S1
(171)
8(2K)/22 8(2K)1/1M
and therefore
Wp SMD(2K)27
Tiax (1) S min
(172)
The next step is to bound the various error terms appearing in Lemma [ can bound for all E [1] and E [N] using Assumption 2 for & E [0,
Using the last display we
2 sup Vy en
max D m [ Un)) |/ (1))5(i, n, 2) MD
DMD( 18(i, n, t)l
W0
wy
5(i, n.. 12
8(2K)
(173)
Morcover, we can estimate for 0 T
C max (t) < C182 muax (t)
Wo
w0 8(2K)
2 (f).
(174)
8C, (2K)
Now we observe that for x = a (n) and a 2 we find (using In(a) S a) the bound
In(r) (()
In(a²) 3
(175)
Vx Valn(a)
Va
Hence consider
24 C2K)3 S S.
(2PC9(2K)3N wp
In
(176)
wg
Note that tracking only the dependence on K and (and using < N) this condition reads
5 22
(177)
29

---

Then we find using (755) combined with the observation that In(S)/VS is decreasing for S 2 e² for
S 2 So the bound
N CKUU11//N(()) S
N
S CCK(T-1/22 1)/2 (n()))))))(
So
3CKK1/2N
(178)
24Ci (2K))))aa
1 wo
8 (2K)²*
We note that similarly for x a aln(a) we get
In
2 In(a)
2 a
(179)
x
a In(a)
Thus we consider similarly
S2S 2 16
(2 (2-16'N(2K) C7
In
(180)
w
w/0
where 5 was introduced in 168 Note that tracking only N and K this bound becomes
$ CN (1/2)
(181)
Then, similar before, (using monotonicity of (n for S e) we find for S 2 Si
N C C max(1, (axxx) Vi((s)
N
S 12K//2 max(1, ( ) S₁
5
(182)
wo 5 max (t)
S
16 (2K)²
After all these preliminary estimates we can finish the proof.
Let us define T E IRJ U ] to be the maximal time such that v(i, n, I) E S for all i E [4], n E [N] and 0 < 1 < T and the bound (168) holds for 8max (t) where we for convenience recall
mmax (1) S ē min
(2KYT GATM--((2K)37722 8(2K))/2TM' 1)
Wo
We
(183)
We want to show that T does not hold, By (72) we know that for t S T the bound (t)x(() 2 holds and therefore Assumption 4 implies that wii, t 0) So and since moreover Smax (t) S for t < T we find that u(i,, n. t = 0) E 24- + (0) CC Ω. By continuity of u(i,, n, t) in t we conclude that there is 0 such that v(i, n,t) t) Ω for all 0 t T + e, i E [7], and n [N]. Thus, we can apply
We assume that T < and prove by contradiction that this
Lemma D for L = T and then find using Assumption 3 and the bounds 173 174). (188) and (82
56i n, 1) < (->-
WD
2
sup 2 By... EQ
max || D.. fv
On Iw]lmam
10 1)1
dt 2
(2K)
No
+ C1 max(1, Tmaxx) (11/22
In(() 5
Sman (1) + Ci
2a + C1 Cidmax(11)
s
S wo
1,, n, 2) 1 n, 22
max(1, T mmax (t)) + f (t)2 fmmax(()
(2K)
16
Wo
(18((,, n, 8) + 35 maxx (t)
max(1, Fmaxx (t))
(2K)²
16
(184)
If 18(i, n, t) 2 3/40max we find 1ō(i, n, t)2 2 Smax mmax(t)2/2 and we conclude that d 1
(6i n. 121 -
Wo
Smax (1)
(1)(-) Smax (1) + Amax (1) + 16
1
max(1, Fimax (t))8
dt 2
(2K)
2
(185)
1 Wo
8 (2K) mmax (t)(- (t) + 8/2).
30

---

Now, if max (T) S 8/2 we conclude by continuity that there is E > 0 such that mmx (t) < 8 for t E [T, T + =]. On the other hand, if max (t) > we conclude that 4 ( 2) < for all i, such that 16(1, n, t) 2 3/45max. This in particular implies that 5max (f) max (T) for E [ + and some E > 0 (omax is non-increasing at T). This is contradiction, and we conclude that T = co.
dr 2
Finally, we prove the decay of maxx Assume that for t T the bound min m'ax (t)) R holds. Note that the function m (4) is not necessarily differentiable but its left and right derivative exist (as the maximum of finitely many differentiable functions). Then we obtain, similar (14)). for the right derivative the bound
max maxx(t) dl+ 2
1
(2K) mmxx (t) CARK (m(()) S
Wo
Smax(t).
(186)
For
()1-1)/2 C116(2K) G1662K)/K-11/2
N In(s)
max (t)
(187)
wo
the previous bound simplifies to
1 wi
dt 2 max (t)² S
16 (2K)r 8 (t)²
(188)
which implies that for t > 7 by Gronwall's Lemma
8max (1)2 S exp (-16
(2K) [2(]) (t 77) mmax (nax(1122
wo
(189)
Thus in finite time we achieve
mmax (t) < C116(2K)KK K1-
N V [n(S) S CK3I (in(s)
(190)
Wg
This ends the proof. Note that we also get an exponential rate of convergence.
Remark The exponent 9/2 for the lower bound of S in (88) is not tight because we could use that mmax is small, but this provides only small improvement that does not justify the additional technicalities.
Before proving and stating Theorem 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 II be approximately balanced, i.e,, assume that Assumption 7 hollds Let Mi > 0 be such that
21)l S VMi
M 2
sup V Vi (O)C*
max (g(k)) S &
(191)
Let 0 <, X S be a fixed number and assume that víi, n, t) follow the gradient dynamics 8 and are initialized such that
no 00] 1
(192)
and assume that the dynamics exists for all times. Then for all i E [4], k E [K], and all times
4KM1
(w(i,ktt)
+ +IK.
(193)
Proof. Note that by assumption
s ((b(t=)) = = N s
Σ ((((m) f(n))²
((m))) N (2/M)) 2
5
= -NM,
(194)
2
nED
31

---

Then we conclude that
N Mi + / 2 R^(0())) 2 (0(1)) 2
max max
Σ
I0(2, m, 1)1
SE [E(K]
REnL (K)
X min
NA
[I7'(A))
max kE|K]
ww t)l
max
[w(i, t)
iE [E(
4K VE (E((
(195)
Thus we conclude that
max E [K) |w(i,, k1) S 4KMr
21K.
(196)
This ends the proof.
Now we state a refined version of Theorem 5 does not require an a-priori bound but instead proves boundedness of the embeddings v(i,n, n,, for all times using Lemma
Theorem 8. Let II be approximately balanced, i.e., assume that Assumption hoolds. Assume that the function f I.D + R is slot-wise linear. Assume that there is M1 >> O such that
VMI
/MI
sup (0))C
[f(ri...., .1)] S 2
max (g(k)) -
(197)
2
Let 0 < X S be a fixed number and assume that for
4KM
R=
2TK +
(198)
Assumption 2 holds with Ω = (0) and some M > O. Note that M depends R and thus on K and X Let C be the constant from Lemmal) [ Assume that v(i, n, follow the gradient dynamics and are initialized such that
A
wii,n,t= 0)]
(199)
8C²
Then there are constants C2. C3 2 depending on M, I, D, X, and C1 such that for
C2 N2K In?(N/)), NKI-1
In(N/A))
S2 max
(200)
the bound
(0)e 11/8 4C₁ 4C2K(1-1)/2
N
8 (t) S max
[I((s)
A
(201)
8max (0)e CI.DAYy
N ( In(S)
< maxx
holds for all 2
Remark We emphasize again that since M might depend À we obtain explicit rate in X. 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 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 least 1 S g-1 over the randomness of D.
Now let T > 0 be the largest time such that for all 0 < T and all i [7] and E [N] have
v(, n, 1) E Ω and the bound
8 mmxx
(202)
4C₁
32

---

(we assume w.l.o.g. that C1 2 1) holds for 0 st<T By assumption, those relations hold at t t= 0 (note that max (t) < 2 max,,n [u n)D. We can bound for ' using the a-priori estimate (193)
from Lemma 2 and find
4KM 2TK + 1. X
|u(i, n. T)) [w(i,, (k), T) max (T)
(203)
Thus, by continuity, there is E > 0 such that v(i, t) E (R( for some ε > 0 and all t E [o, T Now we can apply Lemmal for the interval [T, + s]. We note that by slot-wise linearity we have Da.f = 0 and w 2 0. Therefore, we get
dt 2 6 ((i. - Ajj((, 12 + +CK V In(s)
Smax
S
(204)
N2
C2K(-1)/ In(s))
max + C2 C,83
S
max
Now (with the same reasoning as in the proof of Theorem 4) there is a constant C2 depending on D, I, and M such that for
(
(12)
S So
In
(205)
X²
the bound
N2
C2 In(S))
(nn)))
5
(206)
S
holds. Clearly, we get for suitable constant C2
N2K -1
Su S C2 X2
In2 (N/A).
(207)
Similarly, we get for
2. 16
16
2 C
S S1
In
(208)
A
the bound
C 11/2 1)/2
N In(S)
X²
<
(209)
5
16C
holds. And we find
NKI-1
S S C3 X1
((/A))).
(210)
Now we can continue to bound 200 for S satisfying (055 and 208) and using (022 as follows
d
8(i n. n) - 115(2, 1) 16C1
Smax +
4 max
4 max.
(211)
dt 2
Now, if mmaxx < A/(8C1) then there is > 0 such that (max(() < a/(4C1) holds for t e
[0. T + E). Thus, we assume that dmax 2 A/(8C1). Let E [4], n [N] be any index such that 18(1 n, 7)] = 6max( Then we conclude that
d
2 82 max +
X2
X X
X²
661 n. T)P
8max S
8max +
6max 0.
(212)
2
16C
SC1
16C1
This implies that
d 1
6
max (T) 50
(213)
d 2
from which we conclude that there E > 0 such that 8max S A/(4C1) holds for < 5 + In either case, get contradiction and thus T co. In particular, apply Lemma I for all times t 2 0. Suppose now that for some
4C, 4C
N In(S) 5
8max 2
(214)
33

---

then we conclude from Lemma] 1 that
X
N
8 max
(t) -
2 mmax (t)² -CIK(--1)/2 [I(S))
omax
dl
S
(215)
max (t) + -0- 4 mmaxx (t)
4 max (t).
Thus we conclude from Gronwall's inequality that
max (0)x A0/4
4CCK(1-1)/2
N In(S)
max (1)2 S
(216)
A
5
This ends the proof.
Proof of Theorem 5 The proof is the same the proof of Theorem above with the only exception that we do not need to show that v(i, n, t) E 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 general and standard matrix concentration bound stated in Lemma
simple consequence of the
We start with the matrices
A NNN with entries
A j.j Ai.ng = [n E D n. nı, ny = n22
(217)
the matrix counting the appearance of a pair of tokens in slot i and (this is closely related the matrices A k which will consider below). Then we have the following result.
Lemma Let A NNN be as defined above. Assume that S 2 12N. Then the following bound holds for > 0
5 N
PD ||A E(AW))) (1 2( (1+
SS-
S
(218)
This implies
PD
[A
S N +1+
S (I())
N S S-
(219)
Proof. As in the proof of Theorem we use Poissonization, i.e., we consider dataset generated by first sample $ ~ Poi(S) and then Ds ~ U(NTT)S. A uniform sample n ~ U([N(') satisfies n, = n1 and n, 12 with probability
1
=
N2
(220)
Therefore the distribution of A is
S
1 A
Poi
(221)
B1.ng
N2
and the entries are independent. We can now apply Lemma where we use In(S) as n in the statement of the lemma and find using > N that
s Po (I B(A 2 (1 7) In(S))
<
4.
(222)
N
34

---

Let us define
5 r(S,N,n) +1) (n+)) N
(223)
Note that P($= = -59 1/(3/5). This implies
PD (1 - E(A 2 r(S,
=P' (IA k.j [[ 2 r(S, N.) - =)
Pb (1 -(AA') '0) }| 2 r, and = S)
(224)
P (B 5)
S 4NS
S S-
3V$
Here we used S 2 12N in the last step. The relation S/N.
219 follows because |EA' = NS/N2
Next we prove a similar but slightly more involved result for the matrices that capture similar statistics as A s. but which in addition only consider n such that II(n) = k. We thus consider any [K]' and any indices i, E which we consider fixed for now. Consider the sets Nii, k) IT," (ks),, k) n, (k;) [N]. For a given dataset D (n y..., n3) we define
the matrices A A E (((G)))N(( by
M Knn D: [I(n) k. n, n1, n2 n22
(225)
for a1 E NNii,)) n2 W k). Thus, the entries are the number of datapoints that are mapped to by Il and whose entries i and j are equal n, and 712. 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 gives the following statement.
Lemma 4. Let A e ((.M)))) be as defined above. Assume that all i [1] and k [K] the bound N/(2K) < (k)) s 2N/K holds and let S 2 24N. Then the following bound holds for 1 >
PD E(( | 2 (1 + 1) In(S) max (1. 1, (2)) (1)
S N
< 3S 7++322 (226)
Note that in particular for K 2 2 we find S 2 24N the bound
S
PD |ARR B(A
2 + m)
5-933/2
(227)
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 5 ~ Poi(S) and then D3 ~ U([N''S.. Note that a uniform sample n ~ 2([[N') satisfies II(n) k. = n21 and
ny n2 for n1 Nii k) and n2 E Nj. k) with probability
[II-((k) (A))
p(k, i.j)
(228)
NO 11 (k.) 11, (k2))
and thus the distribution of Ahii
S
111 (k))
A A1.
Poi
=: =:Poi\(kk, i,))
(229)
N7
|T,
(k:)) ) |1T (k:
35

---

for all nı E (i,k) and n2 E WN k) and the entries are independent. Note that the assumption N/(2K) & in (k) - 2N/K K implies that
|IT- (k))
2 2'2
II (E(K]M65
[K)
Ni-2
(230)
1 (k:) 1 (kj)1
2 11-2
which implies
2
S
(2K) N2 S (k,i,j) -
(231)
K
N2
Now we apply Lemma
where 1y In(S) as 7 in the statement of the lemma and find that PD [144t - ((AA))))) (1 17) In In(
(2())) (1-
(1-1)/2
1)/2
S
SN
S
(232)
N
Let us set
N, (,,,) = (1 + 17) In(S)max
1, (1))
(1
11/2
V N
S
(233)
Again, P S) S 8/88///5 1/(3/S). This implies
Pp (1 - E(A || 2 r(S, K, I, I,,)
Pb (114 B(ARAJ)) KK, =S
Ps (ILAR. [E(AAN))) 2
and S S)
(234)
PD(S = S)
SN S-n
-1
S5 14332
3VS
Here we used 8 2 24N the last step.
We also consider the matrix kki. INNIIN already defined in 1115 which equivalently obtain by embedding Akii, suitable, i.e, we have
{o A if i. = and nı E Nii k), na ((jk))
AJd A
(235)
otherwise.
Since k is submatrix of A k...j
we find
A E.A
A.A
k.i. EAA
(236)
so that Lemma applies to A A...j
also need a simpler concentration statement for the frequency of datapoints n such that n; = n.
We define the vectors B' E R by
B', = Kn E D n, = =n]]-
(237)
We need a simple upper bound on the vectors B¹. Lemma 5. The following bound holds for
S N
PP
B
+ +mxx InS)SNN- (1/3nnnn()))))) < S55
(238)
Proof. Note that Bi (0,SNNN1) The variance of Ber(p) variable is p(1 p) S p. Bernstein's one-sided inequality then reads
P (in(S,,N-) - SN > 2) < exp
(-
t²
(-min min
4SN 3t /4 ))
t2
I exp
2SN 1 + 2t/3
(239)
36

---

Now we apply this bound with = max(2 14/3, (n()) which implies
S N 1 )) SS-7 S 2 max 2 2/n1n(5)
P Bin(S,
(240)
N
As before for the A matrices we also need an extension of the previous result a setting where we in addition require I1(n) k for given &o and [AV]. Recall that Mi, k) 11, (k;). Then
consider the vector B (((.) given by
B = In E D: [[()) = k. n, = n)1 = |D
(241)
Again, this agrees with the definition of given in (14). The following Lemma holds.
Lemma Let Bhi (i())))) be as defined above. Assume that for all [T] and [K] the bound [I, (A+) S 2N/K holds. Then the following bound holds for vj > 0 and n E (i,
| E(B max (2(2
(1 13/2
Vnln(s)
S 4 N' N' (3 =
))
S 25
(242)
K
Note that viewing and as constant we get for S N. II(N) the bound
PD
[ -E(B 2 Cr.vnnln(s N
S 28
(243)
Proof. The proof is along the lines of Lemma 5 but slightly more technical. Note that the entries of Bk. are distributed according to Bin(S, p(k, i)) where
III (k)
p(k,
(244)
III (Aks))
Note that by assumption
2N K
1 2
p(k, i) <
N
(245)
N
K
Since the variance of a Ber(p) variable is p(1 p) S p Bernstein's 's inequality then implies
12
P(BinnS,, i)) S p(k, i)l 2 1) S 2xxp
25p(k,, i) + 2t/3
(246)
12
< 2xxp
min
31
4Sp(k, i)
Now we apply this bound with 12 i),
1), 4/3/ (n()) which implies
PP [Bin(S,p(ks) i)) - S i)) 2
1/33 < 28
(247)
Applying
245 we find
(1-1)/2
S In((s) N 3 (pn((s)) S 225-"
P
[Sin(5, i)) S p(k, 2 max
K
(248)
To conclude that (243 holds, we just need to show that for S> N In(N) the expression S/(N In(S)) bounded which follows by monotonicity of 5/ In(S) in S.
Recall that we introduced kid
[N[ in 114
and then
Bl.i
if e
Bl.
(249)
0 otherwise.
For our analysis we want to summarize all those concentration bounds one result. simplify the statement, only consider main regime of interest.
37

---

Theorem 9. Assume that the bound N/(2K) v II (k)] & 2N/K K holds for all i E [T] and k [K]. Assume that
S max
(250)
Then with probability at least i 1
S-1 the following bounds hold simultaneously for all k E [K]
S | B(A 5ll(5))
(251)
N
|| A M.J ((AA)))) - 65) (x) (1-1)/2
s V N
(252)
In(S))
B - N
VV
(253)
N
2 K
(1-1)/2
S
1hth E(B 1 <
[I(())
(254)
N
Proof. Applying Lemma 3 with we obtain that the first bound holds with probability at least 5 (here we used S 12N). Similarly, we obtain for S 2 (K/2)1-1N
max 1, (x)
12/
S N
(2)
11-11/2
S N
(255)
Setting 7 5 in Lemma we get that the second bound holds (by the union bound) with probability For at least 1 372 5-3 for all i E [T] and & E [K] For S 2 aN In(N) and N 2 a we find
S (S)
aNIn(N) In(ali(N)N) 3ln(N)
aNln(N) N
(256)
which implics for = and $ 2 3N In())
S max 2 2 N. 4/3n S <4yIn(s)
(257)
N
and thus applying Lemma 5 with == 2 implies that the third bound holds with probability least 1 - 5-2 Finallyy we find for $ 2 3(K22) -1N Wln(N) and = 3 the bound
(1-1)/2
S 4 N' 3 In()s)
2
(T 1)/2
S In(S) N
2
Vnln(s)
(258)
K
Lemma 6 implies that the last bound holds with probability at least - --IK'S All bounds simultaneously then hold with probability at least
1
1 2 - S- 2-- s-1 2 S-!
(259)
4
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 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.
38

---

Theorem 10 (Corollary 3.7 in (26)). Let X be sequence of independent random symmetric
matrices with dimension We assume that there is a function g (0, 0o) - [0, oo] and symmetric matrices Ag such that
Re ox* & el(0) for 0 > 0.
(260)
Define
==
Amax ΣAn) An).
(261)
Then for all & E R the following bound holds
IP
Xmax
ΣxH
Xx
2t
t
S d inf e -8t+g(0)pp
(262)
We now apply the previous concentration bound to specific setting of interest.
Lemma 7. Consider a random matrix Idd xd2 whose entries are independent random variables
with Asj Poi(A) for some X 0. Then the following bound holds for all 1) 2 0
PP (114 B(AIl 2 (17 + 1) max (1mxx(d1.d2)x,))
2(d1 + da)e
(263)
Proof. To bound the norm of this non symmetric matrix the usual approach consider
Od, xdy
A A-E -B(A))
Q AT B(A)
(264)
O 0s 0d
Consider the index
I= ((i,)) i [di], {di + 1,...,,., d2 d2}}
(265)
then we can write
Q 2 Σ (Nag - A)(E,s + E,a)
(266)
(i,E
where Fij is the matrix whose entry (i, is 1 and all other entries vanish and Ni.- ~ Poi(A). Let Xi, = (Ne. (E + Eja). Note that by induction one directly finds
EE Egj if is even,
(Fi F.gs)*
(267)
E.s + Ezs
i odd.
In any get ( + E E + Eg in the sense of symmetric matrices and we set Agj Ei + Ejj- Now we obtain for any 8 E E the relation
(Poi()) - 2)26
@2h+1 ((Poi(A - 2)22+1
0x Aij
+ (Ea, + E3,)
(2k)
(2k + 1)!
k=U
R=U
(268)
2 (Poi(A) - 2
g2k+1 (Poi(A) - A)2h+1 +1
S Am
(2k)
(2k: 1)!
Note that the first summand is invariant under 0 0 while the term in absolute values changes its sign. This implies that
(2ME(Poi(A) 2)26
DO++E((Pot( - 2)20+1 )
(2k)
(2k + 1)!
(*E((oi(A) - 2)a
(-0)'E(Poi(X) 2)
max
k!
k!
A=0
k=0
Be ((Po(A))))) A), Ee -((Po((A)))) -X)
max
(1) sle -11)+A0) = N6e -1)-
max
(269)
39

---

Here we used the moment generating function of the Poisson distribution and in the last step that and then for 0 2 0
E 2sinh(0) 2 20.
(270)
Thus we infer from (68 and 269 that
Ee 0x e Aeee -1 -1)- Aij el
car(10 A
(271)
using that A = Aij the second step. Note that here
dp
Σ Ai
E (i.)EE
F.g
Σ
d2 Fa,e
Σ x=d++1
d1 Pag
(272)
(iJ))T
which implies
p A
Σ Asj (Wi)er
max(di,
(273)
Thus apply Theorem 10 and get
P Amax
Σ Xis
2t
S (d1 + d2) inf e (t)+ 020
(Kx))T
(274)
S (di d2) inf e 0t A maxddd e50
If (n, < we apply this bound with t = and 0 = 1 and find
P
Amax
Σ
Xis 2 (1+
s (di + d2) exp (-(1 + +) + maxx(d1, - 2))
(K.))r
< (d1 + d2)e
(275)
If max(d1, d2) 2 we set 0 Vmax (d1, d2) < 1 and t = (17 1) (di, d2)A and find
using that for 0 < f < 1 the bound e" - 1 - < 8 holds
Amax
Σ Xij
2 (n + 1) max (d1, d2)A
S Eexp (-Ot + mmxx(d1,
(i.j)e
(276)
-(1+ V max(d1, d2)
max((d1 d2) + Amaxxddi, d2) max(di, / 12) 7)
(d) + d2) exp
(dj + d
Thus combine both cases and obtain the relation
P Amax (Q) Z (n + 1) max (Vmax(dr, d2)), s (dh + dz)e
(277)
We observe that by 271
EEe ((-xs)) Be -8x < (1014A)
(278)
Thus the same reasoning shows that
P (Amiia (Q) S - (11 - 1) max
S (d1 + d2)e
(279)
and since ||A EAA |||]| max [a(()) the claim follows by applying the union bound over
(27 and (279
40

---

I Kronecker Products
Let here state some basic properties of the Kronecker product of matrices. For two matrices A E IR* and B E Rex the Kronecker product is defined by
011 B
B
A & B
E
Ropxxn
(280)
an1 B
anm
B
We will need the property that the operator norm satisfies
| + B|
||||
|B]|.
(281)
Note that we use the notation V @ w vw to denote the outer product. Formally we have www" w but it is convenient and common to drop the transposed here.
41

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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 Section 2 and Section 4 and Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [Yes]
42

---

Justification:
All assumptions are stated and complete proofs can be found in the sumpple-
mentary 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 cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's 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 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 properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that 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 include experiments. If the paper includes experiments, 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 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 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 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 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
43

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
icy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for 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
/ /pipp/.cc/pullc//uuuii wwwtii/gguie///eeeeeeoocc CoodeSubeiiie ssi ssionnoll onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how 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 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 possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support 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, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
44

---

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 96% CI, if the hypothesis of Normality of enrors 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 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 com- puter 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 compute required for each of the individual experimental runs well 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 into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htap:::///euripss/ http://urrip.cccpplppp/pa httppp/c/ccppp..ttttttccccclleea
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
45

---

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 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.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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 (c.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, wwppppswpptseeee whw.ree.wthhhe....c. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
46

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/codeemoodll as part 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included 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 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 expect authors to adhere 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.
47


Reviews:
[{"review_id": "s6Qmh3bdMv", "reviewer": "Reviewer_Rx9k", "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.", "questions": "Is it possible to (approximately) verify whether a given data set satisfies the assumptions you need to have the GD guarantee?", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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 (?).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review.\n\nRegarding your question about the testability of the assumptions:\nThe assumptions on the model used for learning are (in principle) testable. \nFor the dataset, we can test whether samples follow a uniform distribution.\nWhat is difficult to test is whether a partition induced by $\\Pi$ (such that $f=g\\circ \\Pi$) exists (see Theorem 3). It is also not possible to test whether $\\Pi$ is balanced. \nNote, however, that if any algorithm such as gradient descent indeed reveals the partition induced by $\\Pi$\nthen we can verify post-hoc whether the training set satisfies the structural assumptions.\nFinally, we emphasize that while the criticism of lacking testability is valid, it applies to most theory results.\n\nWe hope that this addresses the reviewers' concerns and that they consider improving their score if they find it satisfactory. We are happy to address any further questions or concerns and welcome additional feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gdh3175fKd", "reviewer": "Reviewer_NSnU", "summary": "The 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)$.", "questions": "1. 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.", "limitations": "The authors have adequately addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "-  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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review. \nLet us first answer the reviewer's questions and then briefly comment on the criticism raised in the review.\n\n> What is $N_0$?\n\n- We denote by $N\\_0(I,K,\\eta)$ a constant depending on $I$, $K$, and $\\eta$, i.e., we will change the statement to 'Then there is a constant $N\\_0(I,K,\\eta)$ such that for all $N\\geq N\\_0$...'. In other words, the statement only holds for a sufficiently large number of tokens. Note that a more general statement of Theorem 1 can be found in Theorem 6 in Appendix B where this restriction is not required, and instead the condition on the sample size $S$ is that it needs to be larger than a maximum of four terms. For simplicity, we only kept the dominating term for large $N$ in Theorem 1 because this is our main setting of interest.\n\n> ... what is $\\boldsymbol{u}$?...\n\n- We apologize, there is a typo (overlooked a missing macro when changing the notation) in Eq. (7) and (8): $\\boldsymbol{u}$ should be $\\bar{\\boldsymbol{v}}$, i.e., the concatenation of all token embeddings (for all $n$ and all $i$). Note that the loss function $\\hat{\\mathcal{L}}$ is defined in Eq. (6) and indeed depends on all token embeddings. \nHowever, for one specific datapoint $\\boldsymbol{n}^s$ the value of $f(\\boldsymbol{n}^s)$\n depends only on\n$v(i,n)$ if $\\boldsymbol{n}^s_i=n$. This happens with probability $1/N$. \nSo we will typically find that $S/N$ terms of the $S$ terms in Eq. (6) depend on the token embedding $v(i,n)$ for any $i$ and $n$. So we expect the gradient with respect to $v(i,n)$\nof $\\hat{\\mathcal{L}}$ to be of order $S/N$. Thus, it is natural to reweight by $N/S$ to make it of order 1. \n\n> Technical novelty\n\n- It is true that the proof itself contains no entirely new technical ingredients, however,\nthis does not necessary mean that it is simple, and the proof required us to \ncombine several different techniques, and it was still quite non-trivial for us to find the right decompositions and control all terms.\n\n> Problem setting is a bit restricted\n\n- As every model, our setting is a simplification of real-world settings, and we feel that the setting is already quite complex.\nNevertheless, the suggested extensions seem very reasonable and interesting.\nNote that it is not directly clear what the 'right' representation would be for gradient descent-based algorithms in the considered extension. We will add your suggestions to the conclusion.\n\nAs short remarks about the Minor Nits:\n\n- We will clarify that we use $|u|$ to denote the norm of a vector and $\\lVert \\cdot\\rVert$ for matrix norms. \n\n- We thank the reviewer for pointing out the typos that we fixed in the current version of the manuscript.\n\nWe hope that this addresses the reviewers' concerns, and we are happy to address any further questions or concerns and welcome additional feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zku5CvsOn3", "reviewer": "Reviewer_8unk", "summary": "This 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).", "questions": "I 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.", "limitations": "The authors have properly addressed the limitations of their work.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review. \nWe would like to briefly comment on the weaknesses and the question raised.\n\n\n> The authors provide no algorithm\n\n- This is a valid point, and we will clarify this in the paper: It does not seem too easy to design a problem-specific algorithm. However, we can reduce the problem to other combinatorial problems for which efficient solvers exist. \nIn our specific case, the problem can be reduced to a constraint satisfaction problem (see details below). Standard solvers can solve this problem, but might be slow.\nHowever, as Theorem 3 shows, this cannot be avoided (unless $P=NP$). \nOur results in Section 4 are in line with the general recent trend to depart from worst-case analysis, and here we show that efficient algorithms exist under the more restrictive assumptions on the problem setting. Indeed, our viewpoint here is that for typical real-world settings the data will not be adversarial, and the sample size large and in this case, the problem much simpler than the hardest problem instances.\nMoreover, we show that then gradient-based training can succeed (in polynomial time). \nWe are not aware of results that would allow us to prove that the approach outlined above (i.e., applying general purpose solvers to a constraint satisfaction problem reduction) is, in fact, efficient on these simpler instances, but empirically this is often true. \n\nWe will add this discussion to the updated paper.\n\n> $N^2$ can still be large if $N$ is large.\n\n- It is true that ideally we want a linear scaling (up to logarithmic terms) in $N$ which would match the information theoretic bound.  We believe that the exponent of the $N^2$ term can be slightly improved, at least for slotwise-linear\n$\\hat{f}$ at the price of weaker guarantees, but this we leave for future work.\nNote that while large language models may use tokenizers with $O(10^5)$ tokens, they are\nalso trained on $O(10^{12})$ unique tokens, so the scaling in our results is not too far away from real-world settings. \n\n> Missing experiments.\n\n- There is no particular reason for not applying this approach to recover clusters in real-world data, we did not try this. This was meant as a theory paper, and we feel its content is already quite dense and a solid contribution without experiments. However, the suggestion to confirm the findings with an experiment on sentiment analysis is intriguing, and we will try to implement this in the future. \n\n\nWe hope that this addresses the reviewers' concerns and that they consider improving their score if they find it satisfactory.  We are happy to address any further questions or concerns and welcome additional feedback.\n\n---\n---\n---\n\nDetails on a general algorithm for the considered setting:\n\nLet us sketch how to reduce our problem setting to a constraint satisfaction problem which can be solved with standard solvers (which of course might have exponential runtime).\nIndeed, we introduce variables $t_{kn}$ for $k\\in [K]$ and $n\\in [N]$ which \nare 1 if token $n$ is in cluster $k$ and 0 otherwise. Then we need the condition\n$t_{1n}\\lor t_{2n}\\lor \\ldots\\lor t_{kn}$ to encode that \nevery variable is assigned to at least one cluster. In addition, we consider variables $s_{\\boldsymbol{k}}=g(\\boldsymbol{k})$\nthat encode whether $g(\\boldsymbol{k})$ is 0 or 1 (extensions to general values are possible). \nWe add for every datapoint $(\\boldsymbol{n}^s,f(\\boldsymbol{n}^s))$ and every $\\boldsymbol{k}$ the constraint \n$$ \n\\lnot t\\_{\\boldsymbol{k}\\_1\\boldsymbol{n}\\^s\\_1}\\lor\n\\ldots\\lor \\lnot t\\_{\\boldsymbol{k}\\_I\\boldsymbol{n}\\^s\\_I}\\lor (\\lnot) s\\_{\\boldsymbol{k}}\n$$\nwhere the last term is negated if $f(\\boldsymbol{n}^s)=0$ and not negated if $f(\\boldsymbol{n}^s)=1$.\nNote that if $\\boldsymbol{n}^s$ is in cluster $\\boldsymbol{k}$ (i.e., $\\Pi(\\boldsymbol{n}^s)=\\boldsymbol{k}$)\n then the first part evaluates to false.\n Thus, the condition ensures that if $\\boldsymbol{n}^s$ is in cluster $\\boldsymbol{k}$\n  the cluster must have the prescribed value\n$s_{\\boldsymbol{k}}=g(\\boldsymbol{k})=f(\\boldsymbol{n}^s)$ for the clause to be true.\n\nAny satisfying assignment of the $\\land$ of all these conditions gives rise to a partition induced by a\n$\\Pi$ (encoded by the $t_{kn}$) and $g$ (encoded by $s_{\\boldsymbol{k}}$) while proof of non-existence shows that no such partition exists.\nSo far, we did not ensure that each token is assigned only to a single cluster\nbut this can be achieved in postprocessing or by adding additional constraints\n(such as $\\lnot t\\_{k\\_1n}\\lor \\lnot t\\_{k\\_1n}$).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "21eTbmw3hD", "reviewer": "Reviewer_w8Q8", "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.", "questions": "Is assuming knowledge of $k$ realistic?\n\nOn 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. \n\nFrom 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.\n\nThe paper does not include experiments on real data.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The authors propose an interesting and novel algorithmic problem.\n\nThe 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.\n\nThe paper is poorly written with many language errors. Especially the paragraph from lines 69 to 81 is incomprehensible. \n\nIn Theorem 1, the $M$ is not defined. Should this be $I$?\n\nThe proposed embedding approach seems somewhat unnatural for this discrete problem.\n\nGradient descent seems like a rather crude approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful review.\n\nBelow, we comment on the weaknesses pointed out in the review and reply to the questions raised.\n\n> The usage of the word \"interact\" in the abstract is rather vague.\n\n- We will make the abstract a bit more verbose and clarify that\n the interaction is through a function which has the property that it is invariant under \n the exchange of tokens from the same class.\n \n > The paper is poorly written with many language errors. Especially the paragraph from lines 69 to 81 is incomprehensible.\n \n - We apologize that the writing of the paper (actually mostly section 1 and 2) could be strengthened, and we have used the time since submission to improve the writing of the paper, and we will make sure that the paper is carefully proofread before publication. In particular, we rewrote the paragraph from line 69-81 based on your feedback as follows:\n \nIn this section, we illustrate our problem with an example and define the setup more formally.\nConsider the set of all animals. Those can be grouped into classes \nsuch as mammals, birds, or reptiles (in fact there is a rich hierarchical structure which we ignore here).\nThose groups were conceived by findings sets of animals that share many properties.\nOnce 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. \nNote that this is a specific instance of a general problem in scientific inference where we want \nto uncover a hidden grouping of similar entities from sparse observations about these entities.\n\nHere our main motivation, however, stems from the analysis of large language models\nwhere a similar problem arises implicitly during training.\nThey are trained by next token prediction, so we do not expect them to learn structure \nby deductive reasoning such as \ncows are mammals, and mammals have lungs, so cows have lungs.\n 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\non token embeddings can uncover a hidden cluster structure of the data.\nNote that if the hidden structure is recovered, then generalization to unseen prompts is possible.\n\nWe now define our formal setup...\n\n---\n---\n> In Theorem 1, the $M$ is not defined. Should this be $I$?\n\n\n - Yes, you are correct that all $M$ should be $I$, thanks for pointing this out.\n \n> The proposed embedding approach seems somewhat unnatural for this discrete problem. Gradient descent seems like a rather crude approach. \n \n - Let us clarify the motivation of this paper. One key motivation is to understand how \n large language models learn complex relations between entities. While our setting is clearly a substantial simplification, this is necessary to make progress on the theoretical side. In this context, it is indeed natural to consider gradient-based training and \n continuous embeddings of the discrete tokens because this is currently the main paradigm in language modelling (and also used in vision systems to some extent). On the other hand, \n it is true that when considering the problem in isolation, it is more natural to \n consider different approaches, such as a reduction to other combinatorial problems.\n In our specific setting, it is possible to reduce the problem to a constraint satisfaction problem, and we refer to our response to R. `8unk` for details. Note that it is not trivial to come up with guarantees on the run-time for direct approaches (indeed, they need to be super-polynomial for general instances).\n\n> Is assuming knowledge of $k$ realistic?\n\n- It is correct that $K$ can be unknown in practice.  To deal with this in practice\nwe can first start with a rather large estimate of $K$ \nand then try to merge clusters to find a minimal representation. Note that the \ngradient-based approach does not require prior knowledge of $K$.\n\n> On line 132, it is written that order  $N^I$ samples are needed to ... \n\n- It is true that we here ignored the term $\\ln(N^I)$ for simplicity because this is an informal argument where the logarithmic term is not relevant. To avoid confusion, we will instead use the correct expression.\n\n> From section 4, it seems that the method needs at least $O(NDI)$ memory. Is this scalable?\n\n- Memory of $DNI$ seems to be scalable to rather large settings. Taking in particular into account that large models are currently trained on huge datasets, i.e., $S$ is enormous, the memory cost of the method is much smaller than the size of the dataset.\n\n> Assumption 1 is quite strong. In practice, one often encounters highly imbalanced clusterings.\n\nThis is true. We consider imbalanced datasets in Theorem 6 where we investigate the effect on the sample complexity. For the gradient-based analysis, we focused on balanced clustering for simplicity. However, an extension to an imbalanced setting would be possible but require additional notational burden.\n\n---\n---\nWe hope that this rebuttal clarifies the motivation of our paper and addresses the concerns raised by the reviewer. If our rebuttal is satisfactory, we would appreciate it if the score can be upgraded\nto reflect this.\nWe are happy to address any further questions or concerns and welcome additional feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their careful and insightful reviews.

All reviewers agree that the problem formulation is interesting and relevant ('simple and clean' R. `Rx9k`, 'novel and clear' R. `NSnU`, 'important insights' R `8unk`, and
'interesting and novel' R. `w8Q8`), the results are acknowledged ('progression [...] natural and well-motivated' R. `Rx9k`, 'solid theoretical study' R. `8unk`) and
the presentation of the results is appreciated ('generally easy for readers to follow' R. `NSnU`, 'the authors interpreted and commented all their results in a clear manner' R. `8unk`).
There are some complaints about the language (R. `w8Q8`) and the lack of a general  algorithm 
and experiments (R. `8unk`) that are addressed in the individual responses below.

Meta-review: The paper proposes the problem of learning partitions of tokens from sequences of tokens. The manuscript establishes sample complexity, NP completeness and sample complexity for gradient descent, under some additional technical assumptions.

Some concerns were raised during the rebuttal phase, including for instance: restrictiveness of the assumption, lack of algorithm and experimentation. After this, the reviewers were still supportive of accepting the paper. In addition, I consider this submission a pure learning theoretic paper, and the reviewers' criticisms should be interpreted from that viewpoint (in addition to the supportive scores of the reviewers.)

For a camera-ready version, please take into the account the comments from all reviewers.
Decision: Accept (poster)