Learning Partitions from Context

Authors: simon  buchholz
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:07:31

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.

Main paper: 
 introduction: 
 Modern machine learning systems are able to learn extremely complicated relations from data. They often rely on learned embeddings of discrete tokens in a continuous space. This is notably true for Large Language Models (LLMs) [9, 17, 25, 23] which encode their input by converting text into a sequence of discrete tokens that are embedded in a high dimensional embedding space and those embeddings are fed into, e.g., a transformer architecture [28] which allows to predict the next token. But also in the other domains, e.g., in vision, discrete embeddings are frequently used as a component of deep learning architectures [27, 10] as this enables capturing complex concepts that are often discrete.It was observed that after training these word embeddings exhibit many interesting structures. The most prominent example probably is the observation that the difference of the Word2Vec embedding vectors of the nouns 'king' and 'queen' approximately equals the difference of the embeddings of 'man' and 'woman' [21, 22] . Similarly it was found that using word similarity as an inductive bias to structure latent spaces helps downstream performance [4] . Thus a properly structured latent space seems to be an important ingredient to capture the intricate correlations in complex data.A proper theoretical understanding of such complex models currently remains an elusive goal. However, there have been various attempts to understand various components of deep learning models. Many works investigated the the behaviour of feedforward-networks in particular focusing on shallow networks [2, 8] and asymptotic regimes [15, 29] . More recently several works investigated transformer architectures (often focusing on the one layer case with linearized attention mechanism) [16, 1, 30, 12] . On a technical level [11] is closely related as they study the large depth limit of transformers through the lens of particle systems (however in their case time corresponds to depth, while in our case it corresponds to training time). A feature shared by many of those works is that little structure is assumed on the input, i.e., fixed token embeddings are assumed and often those are even assumed to be isotropic Gaussian.Here we instead focus on the dynamics of the token embeddings and study how these can recover structure present in the data. There is no ground truth target for the embeddings for large scale models used in deep learning and their embeddings need to capture a variety of nuanced correlations and relations that are hard to formalize. Therefore we focus on a simplified problem that nevertheless shares important features with more complex real world settings. One general heuristic is that the embeddings contain information about the similarity of tokens. We focus here on the strongest form of similarity namely equality. Indeed, our central assumption is that tokens can be clustered in a small number of groups such that tokens within a cluster behave exactly the same, i.e., they interact in the same manner with other tokens.Then the central questions is what assumptions allow us to recover a hidden structure. The crucial feature of this setting is that we only get information about a token by its interaction with other tokens about which we also only learn through their interaction behaviour. 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 behaviour where a global structure emerges from local interactions [13, 24, 3] . A related question was investigated in [6] where they study associative memory and also want to identify a hidden structure, however, they learn the class memberships directly through the interaction with a class embedding (and they train the interaction instead of the embedding). In [19] the dynamics of word embeddings and transformer layers was investigated when the data follows a topic model. This work shares the crucial feature that membership of a word in a certain topic is only transmitted through the co-occurrence with other words from the same topic. In contrast to their work we here do not focus on learning class membership from token frequencies and in fact consider uniform token distribution. Instead we view this problem as a logical inference problem: Given a set of facts about a set of tokens can the hidden structure of the tokens discovered.We summarize the main contributions of the paper as follows:• We introduce a learning problem that shares important features with learning through the interaction behaviour that is crucial for LLMs and complex systems.• We analyse this problem from a complexity-theoretic viewpoint where we show that it is in general hard, and from an information-theoretic viewpoint where we show, roughly, that for N different tokens in the alphabet order N ln(N ) samples are sufficient to identify the latent structure.• We then carefully investigate the gradient dynamics of token embeddings finding local recovery of the cluster structure and global recovery for tensor-product functions on the tokens if we have more than N 2 ln(N ) samples for an alphabet with N tokens.Notation. We write [N ] = {1, . . . , N } for the set of the first N integers. The cardinality of a finite set A is denoted by |A| and we also denote the standard Euclidean norm of any vector v ∈ R d by |v|.The expressions λ max (A) and λ min (A) denote the largest and smallest eigenvalue of a symmetric matrix A respectively. We denote the uniform distribution over a set A by U(A). For two subsets A, B ⊂ R d we denote by A + B = {a + b : a ∈ A, b ∈ B} their Minkowski sum. We denote the permutation group on k elements by S k . An overview of the used variable names can be found in Appendix A. 
setting and motivation: 
 In this section, we illustrate our problem with an example and define the setup more formally. Consider the set of all animals. Those can be grouped into classes such as mammals, birds, or reptiles (in fact there is a rich hierarchical structure which we ignore here). Those groups were conceived by findings sets of animals that share many properties. Once these groups are found, we can predict unobserved properties by first identifying the cluster to which an animal belongs and then predict that the property is shared with animals in the same cluster. Note that this is a specific instance of a general problem in scientific inference where we want to uncover a hidden grouping of similar entities from sparse observations about these entities. Figure 1 : Illustration of the setting for I = 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, 6) .Here our main motivation, however, stems from the analysis of large language models where a similar problem arises implicitly during training. They are trained by next token prediction, so we do not expect them to learn structure by deductive reasoning such as cows are mammals, and mammals have lungs, so cows have lungs. Instead, their learning signal is whether a token can be replaced by another token for a given context. Thus, it is a natural question whether gradient descent-based training on token embeddings can uncover a hidden cluster structure of the data. Note that if the hidden structure is recovered, then generalization to unseen prompts is possible.We now introduce our formal setup that captures key aspects of the discussion. We consider I sets of N 1 , N 2 , . . . , N I tokens or entities (such as words). For simplicity we identify these with tokens from the set [N i ]. For each of the sets [N i ] there is a partition P i in K i classes which we can identify with the set [K i ]. Then we can encode the partitions through maps Π i : [N i ] → [K i ] so that the partition is given by P i = (Π -1 i (k i )) ki∈ [Ki] , i.e., Π i encode the class membership. We consider the mapEQUATIONThis structure is illustrated in Figure 1 Now we assume that there is a function g: [K 1 ] × [K 2 ] × • • • [K I ]→ R which depends only on the classes. The map f = g • Π extends this map to the tokens such that it only depends on the group a token belongs to. In the case where g (and thus f ) maps to {0, 1} this can be interpreted as truth assignments, i.e., a statement consisting of a sequence (n 1 , . . . , n I ) is true if f (n 1 , . . . , n I ) = 1 and false otherwise and this case is the main motivation of our work. More generally, f could output the index of a suitable next token where in the 0, 1 case 0 could correspond to a negation while 1 to an end of sentence token. Our goal is to learn the partitions P i or, equivalently, Π up to a permutation and thereby identify the hidden structure.We assume that we are given data in the form of samples (n s , f (n s ) where n s = (n s 1 , . . . , n s I ) andn s i ∈ [N i ].In other words, we try to learn the underlying structure from the interactions of a token with the other tokens which is the same for every element of the partition. To simplify the notation and statements we assume in the following that N i = N and K i = K for some N , K and all 1 ≤ i ≤ I. Our main interest concerns the case where N is large, i.e., there are many entities and K and I are small, i.e., the number of groups is small. Let us summarize several features that this model shares with real world problem such as learning suitable embeddings in latent space.• Hierarchical structures, i.e., groups of objects that share certain features, as discussed here are abundant in language and science. • We only receive an indirect learning signal for the value of Π i (n) through its interaction with other tokens. • Interactions can be very complex, i.e., here the output depends on the interaction of I different tokens and ignoring parts of the context makes learning infeasible.On the other hand, many important features are abstracted away, e.g.:• Here we assume that tokens from the same element of the partition interact in exactly the same way with other tokens while in reality there are many different partitions of the tokens depending on the broader context (e.g., we can group species by habitat, color, or, size each resulting in different partitions) or there are exceptions, e.g., mammals generally do not lay eggs but the platypus does. • Many more complex notions of similarity or further properties of embeddings such as a vector space structure are not covered. Also there can be many uninformative features. • We do not consider noisy data or errors in this work which is crucial for real world applications. 
complexity-theoretic and information-theoretic analysis: 
 We now study this learning problem in different settings. Let us first briefly discuss complexitytheoretic and information-theoretic properties of the learning problem to understand the general boundaries of this learning task. We first study the information-theoretic viewpoint, i.e., the question how many samples are necessary to identify Π (and potentially g). We focus on the case where we sample Π and the data-samples uniformly at random. To learn an unstructured map [N ] I → R we generally need of the order N I ln(N I ) independent samples (N I when sampling without replacement).For the structured setting we show that if Π is drawn uniformly at random then generally order K I N ln(N ) samples are sufficient to learn f and the partition induced by Π. In other words, for every token n i and each of the K I classes Π -1 (k) we need of the order of ln(N ) samples n such that Π(n) = k and n i = n i . In particular, for N ≫ K I any token will interact only with K I ln(N ) ≪ N other tokens, i.e., a very sparse subset of the other tokens.We require the following necessary condition for identifiability: For everyEQUATIONNote that if this condition is indeed necessary because if it is not satisfied then it is not possible to distinguishΠ -1 i (k i ) and Π -1 i (k ′ i ).Clearly, we can generally only identify Π up to a permutation symmetry, i.e., we can only find Π such that there are permutations π i ∈ S K such that Π i = π i Πi . We have the following result. Theorem 1. Assume that g : [K] I → R is a function satisfying the assumption (2). Assume we randomly sample maps Π i such that Π i (n i ) = k i with probability K -1 for all i, n i , and k i and such that (Π i (n i )) i∈I,ni∈[N ] are independent. Assume we are given S samples (n, g • Π(n)) where n ∼ U([N ] I ). Then there is a constant N 0 (I, K, η) such that with probability at least 1 -2e -η for N ≥ N 0 (I, K, η) andEQUATIONwe can recover Π and g up to permutations of[K].This result is a special case of Theorem 6 in Appendix B which shows similar bounds for arbitrary maps Π that are not necessarily random. In the more general setting there are additional dependencies on the size of the preimages Π -1 i (k). Note that this dependency cannot be avoided because if there is a k such that |Π -1 (k)| = 1 and g(k) = 1 and g(k ′ ) = 0 for k ̸ = k ′ then order N I samples are necessary to find n such that Π(n) = k and thus Π. The general proof idea is to bound the probability that any fixed Π ′ ̸ = Π is compatible with the dataset. It turns out that it is possible to bound this probability in terms of the partitions distance of the partitions induced by Π and Π ′ . Then we are left with bounding the number of partitions and we conclude with the union bound. We now show that this this bound is essentially tight. Theorem 2. Let g :[K] I → R be a function such that g(1, k 2 , . . . , k I ) = g(2, k 2 , . . . , k I ) for all k 2 , . . . , k I ∈ [K] except when k 2 = k 3 = . . . = k I = 1. Assume that N is divisible by K and that |Π -1 i (k)| = N/K for all i ∈ [I] and k ∈ [K]. Given 3 ≤ S ≤ N K I-1 ln(N/K)/4 samples (n s , g • Π(n s )) where n ∼ U([N ] I ) i.i.d. Then the function Π is identifiable with probability at most 2e - √ N/K .The proof of this result can be found in Appendix B. Next, we emphasize that while typically a rather small number of samples is sufficient to learn Π it can generally be very hard to do this in practice. More concretely we show that for I ≥ 3 even deciding whether there is a map Π = Π 1 ⊗ . . . ⊗ Π I :[N ] I → [K] I such that f = g • Π given access to samples of the form (n s , t s ) is NP-complete. We show that this is true even if g is known, I = 3 and K = 2. Theorem 3. Consider the map g : {0, 1} 3 → {0, 1} given byEQUATION)Then it is an NP-complete problem to decide given samples of the form(n s 1 , n s 2 , n s 3 , t s ) ∈ [N ] 3 ×{0, 1} whether there is a map Π = Π 1 ⊗ Π 2 ⊗ Π 3 with Π i : [N ] → {0, 1} such that t s = g • Π(n s 1 , n s 2 , n s 3 ) for all samples.The proof of this result can be found in Appendix C. Now that we established under what the conditions Π can in principle be learned, and clarified that this might be hard in general, we next discuss how we can find Π in practice. First, we remark that Theorem 3 rules out the existence of any general fast algorithms to learn Π. Given the combinatorial nature and the hardness of the problem it is natural to reformulate the task as a constraint satisfaction problem which can then be solved using standard SAT solvers (see, e.g., [14] for a review). Indeed, we can introduce Boolean variables t i kn for i∈ [I], k ∈ [K],and n ∈ [N ] which encode whether Π i (n) = k and r kv for every k ∈ [K] I and v ∈ Im(f ) in the (finite) image of f that encode whether g(k) = v. It is then relatively straightforward to then express the conditions for the map Π as a constraint satisfaction problem which is satisfiable if and only if there are maps Π and g such that t s = g(Π(n s )) holds for all samples. We outline the construction in more detail in Appendix C. We leave the task of developing and studying efficient algorithms for the considered problem for future work because the main motivation of this paper is rather to understand how the complex statistical patterns can be extracted using simple gradient based algorithms. This will be investigated in the next section. 
analysis of gradient descent dynamics: 
 In this section we investigate under what conditions the clustering induced by Π can be learned using gradient descent on token embeddings. Our main finding is that for uniformly random Π and S sufficiently large gradient descent can be used to uncover or at least preserve the cluster structure of the embeddings. This shows that while the general problem is NP hard typical random instances with sufficiently many samples can be solved with straightforward algorithms quickly. This is in spirit similar to the results found in, e.g., [5] . Let us start by introducing the setting.Setting. We assume that we have token embeddings for each of the I sets [N 1 ] to [N I ], i.e., we assume that there are vectors v(i, n) ∈ R D for some D and all i ∈ [I], n ∈ [N ]. Based on these embeddings we assume that we are given a function f : R ID → R that transforms the embeddings into a prediction. We will abuse notation and write for n ∈EQUATIONi.e., we will suppress the map from tokens to embeddings in the notation. Now we consider gradient descent for the embedding vectors using the least square loss on training samples, i.e., the loss of a sample (n,t = f (n)) is ( f (n) -f (n)) 2 .We assume that we are given a dataset D = {n 1 , . . . , n S } ∼ U([N ] I ) S . Then the empirical loss reads (the division by 2 is convenient for the gradient evaluation later)EQUATION)We also define shorthands for certain concatenations of embeddings. For a sample n ∈ [N ] I we denote the collection of embeddings by v(n) = (v(1, n 1 ), . . . , v(I, n I )). Using the convention (5) we can then write f (v(n)) = f (n).Moreover, we define v(i) ∈ R DN as the concatenation of the vectors v(i, 1), . . . , v(i, n), i.e., the combined embedding for the i-th slot and v ∈ R DN I as the concatenation of the vectors v(i) for 1 ≤ i ≤ I, i.e., all token embeddings concatenated. We consider the regularized loss given byEQUATION)Note that the scaling by N/S (instead of usual 1/S) is natural because every token embedding v(i, n) occurs in approximately S/N of the samples and so the scaling ensures that the gradient of Rλ with respect to the token embedding v(i, n) is of order one. Now, we consider the continuous gradient descent of the loss with respect to the embeddings, i.e., we consider (omitting the time variable from the notation)EQUATIONThis introduces a time dynamics on the token embeddings. We indicate the time dependence by v(i, n, t) but we drop t if not necessary. Our main goal is to understand which conditions ensure that the token embeddings v(i, n, t) and v(i, n ′ , t) converge to each other ifΠ i (n) = Π i (n ′ ) as t → ∞.To investigate this we define the center of the class embeddings byEQUATIONand we consider the deviations from the class centres given byEQUATION)Thus the vectors δ(i, n, t) capture whether we recover the cluster structure, in particular if all norms |δ(i, n, t)| are small then we essentially recovered the hidden structure. Therefore we defineEQUATIONSimilarly, to the notation introduced before we consider for k ∈ [K] I the vector w(k) = (w(1, k 1 ), . . . , w(I, k I )). As in (5) we abuse notation and writeEQUATIONSimilar to v(i) and v we introduce w(i) as the concatenation of (w(k, i)) k∈[K] and w as the concatenation of w(i).Assumptions. Our first result for the gradient dynamics states that clusters are stable under the dynamics if the initial loss is sufficiently small. More precisely this means that we assume that v(i, n) and v(i, n ′ ) are close initially whenever Π i (n) = Π i (n ′ ), i.e., δ max (0) is small. In addition, we assume that |ĝ(k) -g(k)| is small. To capture this we define r max (t) = maxEQUATION)Then the result shows that δ max stays small for all times if S ≳ N 2 under mild additional assumptions.In other words if we start from the correctly learned substructures and ĝ(k) ≈ g(k) for all k ∈ [K] I then this remains true for all times. Note that while we phrase smallness as an assumption on the mean embeddings w(i, k) this is generally a consequence of δ max small and a small empirical loss L(v). Let us now state the required assumptions.Assumption 1. We assume that the map Π :EQUATION)This assumption ensures that clusters are of approximately equal size. We have already seen in Section 3 that different cluster sizes increase the sample complexity of learning Π. Assumption 2. We assume that there is a convex set Ω ⊂ R D and a constant M ′ ≥ 1 such that the following bound holdsEQUATION)Here it is convenient to introduce M = 16M ′2 so that later certain errors in a Taylor approximation are bounded by M . We also assume thatEQUATIONThis is a rather mild assumption. For C 3 functions f and Ω bounded this is always true. The next assumption entails a rigidity of approximate minimizers of the loss. Assumption 3. We assume that for all embeddings w(i, kEQUATIONthe boundEQUATIONholds for some positive constant ω 0 . Of course the bound r max ≤ 1 could be replaced by any other constant.Note that this condition can only hold if D ≤ K I-1 , i.e., the latent space dimension cannot be too large. The high level intuition of this assumption is essentially that(at least if k ( f (w(k)) -g(k)) 2 is small) there is no direction v ∈ R D such that v • ∇ w(i,k) f (w(k)) ≈ 0 for all k such that k i = k, i.e., we cannot move one single embedding without changing the output f (w(k)) for at least one k. If this condition does not hold then we cannot guarantee that k ( f (w(k)) -f (k)) 2 is minimized for a unique w(i, k) (for all other embeddings fixed). This generally prevents concentration of v(i, n).Note that this condition does not ensure that there is a unique minimizer w, in particular there could still be a rotationally invariant family of embeddings w(i, k) such that f (w(k)) = g(k) for all k ∈ [K] I . Finally, we need a further mild assumption that ensures that mean token embeddings w(i, k) stay bounded in some set if the loss is small. This can be achieved, e.g., iff → ∞ if |w(k)| → ∞. Assumption 4.We assume that for all collections of mean embeddingsEQUATION)there is a convex set Ω 0 ⊂ R d such that w(i, k) ∈ Ω 0 for all i ∈ [I], k ∈ [K].Again, the right hand side of the bound r max ≤ 1 could be replaced by any other constant.Results. The first stability theorem can then be stated as follows. Theorem 4. Let Π : [N ] I → [K] I be approximately balanced as stated in Assumption 1 Assume that the functions g : [K] I → R and f : R ID → R satisfy Assumption 3 for some ω 0 > 0 and Assumption 4 for some convex set Ω 0 . Assume that Assumption 2 holds for some M and the set Ω = Ω 0 + B 2 (0). Then there are constants c 1 , C 2 , C 3 , C 4 > 0 depending on I, M , D, and ω 0 such that for all initial embeddings v(i, n, t = 0EQUATIONand sample sizeEQUATIONthe following holds with probability at least 1 -S -1 over the randomness of the dataset. When considering the gradient dynamics of the embeddings given by (8) the boundEQUATIONholds and moreoverEQUATIONIn particular δ max (t) → 0 if r max (t) → 0, i.e., all token embeddings for one fixed class converge to the same point.This results shows that for order N 2 ln(N ) samples and initialization sufficiently close to a global minimum the cluster structure remains stable. We do not provide conditions that ensure r max (t) → 0 which guarantees convergence to 0 loss and perfect recovery of the clusters.Next we note that we cannot expect the clustering to be stable in general even if δ max (0) is arbitrarily small if initialization is not close to a minimum. This is true even in the simplest case where I = K = 1, i.e, we consider gradient descent for a single function value. Then gradient descent does not necessarily converge to a global minimum and close by points do not necessarily stay close because gradient descent is not well posed. Let us clarify this by an example. EQUATION)We find that v(1, n, t)→ 2/3 (which is a local minimum of f as t → ∞ if v(1, n, t = 0) > 0. On the other hand v(1, n, t) → -1 if v(1, n, t = 0) < 0. So in this case δ max (0) = O(σ 2 ln(N )) but δ max (t) → 5/3 as t → ∞. Slight modifications show that also δ max → ∞ is possible.The previous example shows that without additional assumption we cannot expect to recover the structure of the data. Therefore we impose additional restrictions on the function f . As apparent from Example 1 and also from the bound in Lemma 1 in Appendix E, it is the curvature of the function f that can push token embeddings of the same class apart. We therefore consider the function class of slow-wise linear functions defined as follows. Definition 1. We call f slot-wise linear if for every v = (v(1), . . . , v(I)) ∈ R D×I and any i ∈EQUATIONLet us denote for v ∈ R D by ṽ ∈ R D+1 the vector v where we append a 1. The most general slot-wise linear function is then of the formEQUATIONwhere T : R (D+1) I → R is a linear map. Note that this class covers linearized attention where the embeddings v(i, n) are split in three separate parts that are used to form key, query, and value vectors.For this function class we can show stronger clustering results.Theorem 5. Let Π be approximately balanced, i.e., assume that Assumption 1 holds. Let f : R ID → R be a slot-wise linear. Assume that v(i, n, t) follow the gradient dynamics (8) and that v(i, n, t) ∈ Ω for all i ∈ [I], n ∈ [N ] and t > 0 for some convex set Ω. Assume that Assumption 2 holds with constant M for the set Ω. Let C 1 be the constant from Lemma 1 (which depends on on I, D, M ). Assume that at initializationEQUATION)Then there are constant C 2 , C 3 ≥ 0 depending on M , I, D, such that forEQUATIONthe boundEQUATIONholds for all t ≥ 0.The high level summary of this result is that for order N 2 ln(N ) samples the clusters can be recovered up to an error of order λ -1 ln(S)N/S. Note that the a-priori assumption that the gradient flow is restricted to the set Ω might appear difficult to guarantee in practice. However, we conjecture that the results extend to gradient dynamics clamped at the boundary. Moreover, in Lemma 2 we prove that the mean embeddings w(i, k, t) stay within a ball of radiusR = O( √ λ -1) for all times. This allows us to prove Theorem 8 which does not require any a-priori bounds on the evolution of v(i, n, t) but comes at the price that the constants C 1 , C 2 , and C 3 depend on λ so we cannot infer the explicit λ dependence. Let us make an important remark. Remark 1. While we state our results for a fixed function f this function could in principle be time dependent, e.g., f could be given by a neural network and we consider gradient descent not only on the token embeddings but also on the network parameters. The only requirement is that the assumptions hold uniformly for all times t. In particular, for Theorem 5 we only need to ensure that the derivatives of f stay uniformly bounded in time. This can, e.g., be guaranteed by clipping the parameters of the slot-wise linear map f . Our results so far show that we can recover the right cluster structure in the sense that the embeddings from the same group cluster. However, this leaves open the question whether there is any non-trivial dynamics at all, i.e., all embeddings could cluster at the same point. This is in general not the case as can be seen from Corollary 1 which states that in the setting of Theorem 5 and for large times the dynamics of the cluster-means follow the equationEQUATIONwhere (2K) -I ≤ α k,i ≤ (2/K) I are positive numbers. This shows that w(i, k) follow generally a non-trivial dynamic (and this also justifies the scaling as this expression is of order 1). So in the generic case the cluster structure will be revealed if the numbers of samples is sufficiently large, however, there is no general guarantee that the clusters are well separated. As an illustration of this result we refer to Figure 2 where the clustering of the embeddings becomes apparent.Proof idea and overview. Let use here give a quick summary of the main steps of the proof and where to find them. The first important ingredient is an expansion of the loss gradient. We Taylor expand the loss of a sample v(n) around the point w(Π(n)) to second order with remainder terms, the relevant calculations can be found in Appendix D (see Proposition 1 for the outcome).A second ingredient are concentration bounds for certain datapoint-statistics random variables and random matrices. Those are derived in Section G with the necessary results collected in Theorem 9.Combining the Taylor expansion with the concentration result we can extract the dominant terms of the expansion (see Appendix E). Moreover, we obtain such an expansion for ẇ(i, k) (see Corollary 1) and thus the displacements δ(i, n) (see Corollary 2). This expansion can then be used to control∂ t |δ(i, n)| 2 (see Lemma 1) which is sufficient to control δ max .Discussion of assumptions Let us contemplate the differences and similarities to training token embeddings in neural networks.• For the first main result Theorem 4 we make minimal assumptions on f so this could in principle be a neural network applied to the token embeddings. The second result, Theorem 5, is more restrictive but covers subclasses of linearized attention. • An important feature of the results is that f itself could be time dependent (see Remark 1).Differences to standard training of neural networks are:• We use continuous time gradient descent instead of stochastic gradient descent. This is a frequently used modification and in suitable limits those converge (see, e.g., [20] ). • We use mean squared error while sequence modelling usually relies on cross entropy loss.This simplification is frequently used in theoretical analysis but it is expected that results generally extend to the non-convex cross entropy loss. • A more crucial difference is that the embedding space dimension in practice is usually chosen large to provide large representation capacity. Here D has to be rather small to allow a unique optimal solution of the token embeddings that allows us to recover the cluster structure. 
conclusion: 
 In this paper we considered a learning problem where we try to recover a partition of tokens from their interaction with other tokens. This can be seen as a toy problem for next token prediction in LLMs but also more broadly as a problem in scientific inference. We studied this problem from different perspectives, namely from an information-theoretic, complexity-theoretic, and gradient descent based viewpoint. We found that order N ln(N ) samples are sufficient to recover the partition for N tokens while we showed that N 2 ln(N ) samples are sufficient for gradient based methods. There are several natural open follow up questions. First, there are some open questions regarding the tightness of our analysis of the gradient descent. In particular it is a natural question whether already Ω(N ln(N )) samples are sufficient to control δ max which is the information-theoretic threshold and would be similar to the optimal results for matrix completion [18, 7] ). Another interesting question for future research is whether Theorem 5 holds for standard initialization schemes for the token embeddings. Secondly, it is of interest to relax the notion of clustering of embeddings to more general notions that still allow to recover some structure but are also applicable to high dimensional latent spaces and potentially to multiple partitions and relations on the same tokens (e.g., tokens belonging to different clusters). Thirdly, it is a natural question whether this work can be connected more closely to empirical findings. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? 
 
 Appendix: In this section we provide the proofs and additional results for Section 3. Let us first start to state the general information-theoretic bound that handles the case where |Π -1 i (k)| might be of arbitrary size. Theorem 6. Let g : [K] I → R satisfy (2) and there is a projectionΠ = Π 1 ⊗. . .⊗Π I : [N ] I → [K] I . Assume there is L > 0 such that for every k ∈ [K] I we have Π -1 (k) ≥ N I /L. Moreover assume for all i ∈ [I] and k ∈ [K] the bound |Π -1 i (k)| ≥ M , i.e., every class has at least M members. Assume we are given S samples {(n s , gEQUATIONfor some η ≥ 2 then with probability at least 1 -6e -η we can recover Π and g up to permutations of[K].In other words, for any mapΠ ′ = Π ′ 1 ⊗ . . . ⊗ Π ′ I and g ′ such that g • Π(n s ) = g ′ • Π ′ (n s ). There are permutations π i ∈ S K such that Π i = π i • Π ′i and the corresponding relation holds for g, g ′ .A few remarks to explain this result are in order. Remark 2. The scaling of S might appear slightly complicated so let us comment on this. We are mostly interested in the case where N large and K stays bounded. In addition, we are primarily interested in the regime where L ≤ CK I bounded and M ≥ cN/K (which holds for random Π), i.e., the sampling probability of each class k ∈ [K] I is of similar size. In this case the first term in the condition (31) for S stays constant as N → ∞. The second term dominates and we see that we need only O(N ln(N )) samples to identify Π and g. This is the setting studied in Theorem 1 in the main text. The result is essentially tight in this limit (up to the dependence on I) as stated in Theorem 2. Note that as remarked in the main text the dependence on L cannot be avoided. Indeed, consider the extreme case where Π i (n) = 1 for n = 1 and Π i (n) = 2 for n > 1 and g(k 1 , . . . , k I ) = 1 iff k 1 = . . . = k I = 0 and g(k 1 , . . . , k I ) = 1 otherwise. Then we have L = N I and we also need order N I samples to sample the point (0, 0, . . . , 0) ∈ [N ] I which is necessary to identify Π.Proof. Let us first introduce some notation. We denote the set of samplesEQUATION)Consider two partition P, Q of a set [N ]. We denote by P A for A ⊂ [N ] the restriction of a partition to a subset (i.e., {P ∩ A : P ∈ P}). The partition distance is defined byEQUATIONi.e., the minimal number of elements that need to be removed such that the partitions agree.We call Π ′ compatible with the datapoints D if there is a g ′ such that g ′ • Π ′ = g • Π on the the data. Now the general strategy is to consider any other candidate map Π ′ and upper bound the probability that is compatible, i.e., for some function g ′ the functions f ′ = g ′ Π and f agree on the data. We will then conclude by applying the union bound over maps Π ′ . Thus we need to prove an upper bound on the number of partitions with partition distance at most ∆ and a lower bound on the error probability for a given ∆. Let us start with the latter. Denote byEQUATIONthe partition generated by Π i . Consider any other mapΠ ′ = Π ′ 1 ⊗ . . . ⊗ Π ′ I : [N ] I → [K] I with corresponding partitions P ′ i .We define for any such Π ′ the quantityEQUATIONAssume now thatEQUATIONand in particular D(P i , P ′ i ) ≤ ∆. Let P 1 = {P 1 , . . . , P K }. Let A be a set of maximal size such that P A 1 = P ′A 1 . After relabeling we can assume that P k ∩ A = P ′ k ∩ A where P ′ k ∈ P ′ 1 and by composing a permutation we can also assume thatEQUATION)The same applies to all partitions P i and P ′ i . Next we claim that if f = g • Π and f ′ = g ′ • Π ′ agree on all samples then with probability K I e -S 2 I L over the randomness of the samples g = g ′ . Define E k := Π ′-1 (k) ∩ Π -1 (k). Using (37) and the assumption on Π we have that for k ∈EQUATION)The probability that none of the S samples is in E k can then be bounded as followsEQUATION)Applying the union bound over k we find thatEQUATION)Now we bound the probability that g ′ = g is compatible, i.e., gEQUATIONWe now lower bound the size of E n1 under the assumption that there is k ∈ [K] such that n 1 ∈ P ′ k ∩A c (where A is as above a set such that P A 1 = P ′A 1 ). Then Π ′ 1 (n 1 ) = k and (by minimality of A c ) we find n 1 / ∈ P k and thus Π 1 n 1 = k ̸ = k for some k. By assumption, we can find k 2 , . . . , kEQUATIONwe conclude that all such n are in E n1 By assumption we haveEQUATION)This implies thatEQUATIONand thereforeEQUATION)The sets E n1 are disjoint so we conclude thatEQUATION)Now we can upper bound the probability that D is compatible with Π ′ and g byEQUATION)Clearly the same reasoning applies to any other index instead of 1.The next step is to upper bound the number of such candidate partitions. The number of partitions such that D(P 1 , P ′ 1 ) ≤ ∆ can be bounded by (N K) ∆ , i.e., we ∆ times select one of N tokens and assign it to another (or the same) class. This bound can be applied to all indices 1 ≤ i ≤ I and by the union bound, we find thatEQUATION)Thus we can bound the the probability that any of the maps Π ′ (there are more such maps than partitions because there is a label assigned to each class but consistency with the data depends only on the underlying permutation) with 0 < ∆(Π ′ ) ≤ M/2 is consistent with the data is bounded byEQUATIONHere the first term corresponds to the upper bound for the probability that (after permutation) g = g ′ is compatible with the data and the second term bounds the probability that any other function is compatible with the data and I accounts for the fact that any i ∈ [I] can be arg max i D(P i , P ′ i ). We now consider the remaining case that there is an index i ∈ [I] such that D(P i , P ′ i ) ≥ M/2 where we can assume w.l.o.g. that i = 1. As above let P 1 = {P 1 , . . . , P K } andP ′ 1 = {P ′ 1 , . . . , P ′ K } such that P k ∩ A = P ′ k ∩ A andA is a set of maximal size. We now claim that we can find an indices k ̸ = k such thatEQUATION)Indeed, first assume that there is an index k such that|P ′ k ∩ P k | ≤ M/K. Since |P ′ k | ≥ M there is another index k such that |P ′ k ∩P k | ≥ M/K. By maximality of A we have P k ∩A = P ′ k ∩A = P k ∩P ′ kand moreoverEQUATIONbecause otherwise exchanging P ′ k and P ′ k would allow to pick a larger A. By our assumption we conclude from here that EQUATIONEQUATIONEQUATION)Clearly Π ′ is constant (and equal to (k, k ′ 2 , . . . , k ′ I )) on A ∩ Ā and therefore also g ′ • Π ′ is constant for any g ′ . On the other hand Π(A) = {k} is constant andEQUATIONfor n ∈ A and n ∈ Ā. Now Π ′ can only be consistent with the data (for any g ′ ) if there is no sample in A or no sample in Ā, i.e., and thusEQUATIONwhere we used the union bound in the last step. Then we find using (54) and (56)EQUATIONFinally we can bound the number of partitions by the number of maps Π i : [N ] → [K] which implies that |{P 1 , . . . , P I :EQUATION)So we get the upper boundEQUATIONCombining ( 49) and (59) and using the union bound we findEQUATIONNote that clearly the assumptions imply M ≤ N/K < N and thus e -S/(2 I L) ≤ e -S∆/(2 I LN ) for ∆ ≤ M/2. We then find (using the simple boundEQUATIONforEQUATIONand η > 2 where we bound the geometric sum by twice its largest term. The first summand in (60) can be bounded byEQUATIONforEQUATION)The proof of Theorem 1 is now a direct consequence of the previous result because when assuming that a uniformly random map Π is chosen we can estimate the quantities M and L in the previous theorem with high probability.Proof of Theorem 1. We observe that |π-1 i (k i )| ∼ Bin(N, K -1). Applying a Chernoff bound on the tail of the binomial variable we obtainEQUATION)By the union bound we getEQUATIONif N ≥ N 0 (η, K, I). Note that M ≥ N/(2K) implies that L ≤ (N/M ) I ≤ (2K) I . Assuming that Π is such that the bounds for M and L hold we can apply Theorem 6 and find that for N ≥ N 0 sufficiently large (depending on I, K, η) andEQUATION(we bounded ln(N K) ≤ 2 ln(N )) the maps Π and g are identifiable with probability at least 1 -e -η .Here we used that as N → ∞ the term indicated above is dominating in (31). Now the union bound over the bad events ends the proof.Proof of Theorem 2. As before, we note D = {n 1 , . . . , n S }. Assume there is n 1 such that Π 1 (n 1 ) = 1 and such thatD ∩ {n 1 } × Π -1 2 (1) × . . . × Π -1 I (1) = ∅. (68) Then Π ′ ̸ = Π is compatible with D where Π ′ i = Π i for i ≥ 2 and Π 1 (n) = Π ′ 1 (n) for n ̸ = n 1 and Π ′ 1 (n 1 ) = 2 ̸ = 1 = Π 1 (n 1 )(by assumption on g). Let us denoteEQUATIONThenEQUATION)To estimate the probability of the event n1 {A n1 ∩ D = ∅} 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 a dataset D as before, i.e., n s ∼ U([N ] I ) for s ∈ [ S]. Then we find thatEQUATIONand those events are independent for n 1 ̸ = n ′ 1 . Thus we find thatEQUATION)By independence we now find thatEQUATIONEQUATIONEQUATION)where we used the upper bound for S and (1 -x) ≤ e -x . Assume D s ∼ U([N ] I ) s and defineEQUATIONNote that p S is an upper bound on the probability that Π is identifiable as explained above. We have shown thatEQUATION)Clearly p s is decreasing in s. This implies thatEQUATIONA Chernoff bound for the Poisson distribution readsEQUATIONwhich implies with λ = 2S and x = S thatEQUATIONWe thus conclude thatEQUATIONfor S ≥ 3 (implying 1 -(2/e) S ≥ 1/2).In this section, we provide the proof of Theorem 3 and a general reduction to a constraint satisfaction problem. We start with the proof of Theorem 3. An important property we will use frequently in the proof is that T has the property that for and values x 1 and x 2 there is x 3 (x 1 , x 2 ) such that T (x 1 , x 2 , x 3 (x 1 , x 2 )) = 0. Indeed x 3 (x 1 , x 2 ) = 0 except for x 1 = x 2 = 1 where x 3 (1, 1) = 0 has this property.Proof. We reduce the problem to 3SAT. We consider an arbitrary formulaEQUATIONon ℓ variables with m clauses C i . We denote f = g • Π.To improve readability we consider a symbol set S instead of [N ] which is used for all 3 slots. We proceed in three steps. First we show that we can write down a set of equations that ensures that two symbols s 0 , s 1 ∈ S satisfy Π i (s j ) = j for all i (if f = g • Π). Suppose the following relations holdEQUATION)Then it is easy to conclude that Π i (s j ) = j holds. Indeed suppose that Π 1 (s 1 ) = 0. Then the lower part of the previous display combined with the definition of g imply Π 2 (s 0) = Π 3 (s 1 ) = Π 2 (s 1 ) = Π 3 (s 0 ) = 1.But this is a contradiction to f (s 1 , s 0 , s 0 ) = 0. Thus Π 1 (s 1 ) = 1 and the same reasoning implies Π i (s 1 ) = 1. Then the second equation in (81) directly implies Π i (s 0 ) = 0.We now consider a symbol setEQUATIONwhereEQUATIONwill encode the variables in the formula ϕ andEQUATIONare auxiliary variables for the clauses. In addition we need further auxiliary variables X = { X1 , . . . , Xℓ },EQUATIONEQUATIONWe now add a set of equations that will ensure that for some valueEQUATIONThis can be achieved by adding the following relations for all iEQUATIONNote that the these relations indeed imply that Π a (X i ) ̸ = Π b ( Xi ) for a ̸ = b, a, b ∈ {1, 2, 3}. This implies then that (86) holds. We add the similar relations for C i,k which again ensure thatEQUATION)Now we encode the clauses of the formula. Consider a clause C l of the form x i1 ∨ x i2 ∨ x i3 . Then we add the relations f ( Xi1 , Xi2 , C l,1 ) = 0EQUATIONFor a given choice of Π we set x i = Π 1 (X i ). We now claim that the relations in the last display can hold if and only if x i1 ∨ x i2 ∨ x i3 evaluates to true. Suppose the formula evaluates to false. ThenΠ j (X i1 ) = Π j (X i2 ) = Π j (X i3 ) = 0 (90)and thereforeEQUATIONThen relation (89) impliesEQUATIONfor i = 1, 2. Using (88) we find that Π 1 ( Cl,1 ) = Π 2 ( Cl,2 ) = 0 which impliesEQUATIONfor any Π 3 (o l ). Therefore the equations cannot all hold.Suppose to the contrary that the formula evaluates to true. Let us first assume that Π 1 (X i1 ) = 1.Then Π 1 ( Xi1 ) = 0 and we can set Π 3 (C l,1 ) = 0 which ensures that the first equation is satisfied (for any Π 2 ( Xi2 )). By definition of g we can also find a value c l,2 such that for c l,2 = Π 3 (C l,2 ) the relation f ( Xi2 , Xi3 , C l,2 ) = f (x i2 , x i3 , c l,2 ) = 0 holds. Now Π 1 ( Cl,1 ) = 1 -Π 3 (C l,1 ) = 1 ensures that for any value of cl,2 = Π 2 ( Cl,2 ) we can choose Π 3 (o l ) such that the relation f ( Cl,1 , Cl,2 , o l ) = f (1, cl,2 , o l ) = 1 holds.The same reasoning applies for Π 1 (X i3 ) = 1 and a similar argument applies if Π 1 (X i2 ) = 1. For clauses containing negations the same construction works except that Xij has to be replaced by X ij for the negated variables in equation ( 89).Putting everything together we have shown that for a given formula ϕ there is a set of relations of the form f (n s ) = t s where n s ∈ S 3 given by ( 81), (87) (plus similar equations for C l,i ), and (89) which has the following properties. For an assignment x = (x 1 , . . . , x n ) ∈ {0, 1} n such that ϕ(x) = 1 evaluates to true there is a map Π such that for all s the relations t s = f (n s ) hold and where Π 1 (X i ) = x i . On the other hand, if for some Π all the relations t s = f (n s ) hold, then the Boolean variables x i = Π(X i ) will satisfy the formula ϕ. This ends the proof.We now show how the problem of finding Π and g can be generally expressed as a constraint satisfaction problem. Recall that we want to find (if they exist) for a given K a projection map Π : [N ] I → [K] I and g : K I → R such that for all given samples (n s , t s ) the relation t s = g •Π(n s ) holds. The general strategy is to introduce Boolean variables encoding the maps Π and g and then express all conditions as suitable constraints for these variables. We first introduce Boolean variables t i kn for i ∈ [I], k ∈ [K], and n ∈ [N ] which are 1 if token n is in cluster k, i.e., Π i (n) = k and 0 otherwise. Then the expressionEQUATIONfor i ∈ [I] and n ∈ [N ] is true if and only if n is assigned to to at least one cluster. In addition, we consider variables r kv for every k ∈ [K] I and v ∈ Im(f ) in the (finite) image of f . These variables shall encode whether g(k) = v is is true or not. Then the constraintsEQUATIONfor v ̸ = v ′ and all k ensure that each cluster is assigned at most one value v. Finally, we add for every datapoint (n s , f (n s )) and every k the constraintEQUATIONThis ensures that if Π(n s ) = k then this cluster must be assigned value f (n s ).We then consider the constraint satisfaction problem consisting of the of all the conditions in (94), (95), and (96). Then any satisfying assignment gives rise to a map Π and g such that g • Π(n 2 ) = t s . Indeed, we set g(k) = v if s kv = 1 if such a v exists and arbitrarily otherwise. Note that for every satisfying assignment and every k there is at most one such v so g is well defined. Moreover, we set Π i (n) = k for any k such that t i nk = 1 and at least one such k exists. On the other hand, we can easily construct a satisfying assignment given Π and g so that there are no solutions Π and g if the constraint satisfaction problem has no solution.Note that we did not ensure that each token is assigned only to a single cluster but this can be achieved in postprocessing or by adding additional constraints (such as ¬t i k1n ∨ ¬t i k1n ).The goal of this and the following section is to lay the groundwork for the proof of the main results.The general strategy of our proofs is to Taylor expand the loss of each term around the mean token embedding w(i, Π i (n i )) to first order plus a remainder term. We can then extract the dominating terms of these expansions using concentration results for the datapoint statistics. It turns out that the linearized dynamics has favourable properties while the remainder terms can be bounded. In this section we derive the expansion of the loss gradient while the required bounds, in particular Theorem 7 and Lemma 1, are derived in the next Section, Appendix E (they rely on concentration results which are deferred to Appendix G).As pointed out above, the goal of this section is to Taylor expand the sample loss for one sample where we expand the loss around the point w(Π(n)). Let us first introduce some notation. Recall that we denoted by v(n) = (v(1, n 1 ), . . . , v(I, n I )) and by v(i) the concatenation of the embeddings(v(i, n)) n∈[N ]. We define similarlyEQUATIONand we denote by δ(i) ∈ R DN the concatenation of δ(i, 1), . . . , δ(i, n), i.e.,EQUATION)and δ ∈ R IDN as the concatenation of δ(1), . . . , δ(I).Consider a data-point n with Π(n) = k. We now consider the derivative of the mean squared error of this term, i.e.,EQUATION)We use greek-indices for the latent space dimension. Fix an i ∈ [I] as the slot with respect to which we take the derivative. We then find for α ∈[D] d dv α (i, n i ) 1 2 f (v(n)) -f (k) 2 = d dv α (i, n i ) 1 2 f (v(n)) -f (k) 2 = d dv α (i, n i ) f (v(n)) • f (v(n)) -f (k) =: h(v(n)).(100)Here we introduced the shorthand h(v(n)) for this function which also depends on i, k, and α. Now we estimate this function by Taylor expanding it around the pointEQUATIONThen we find the following expansion to second orderEQUATIONwhere R (j1,β1),(j2,β2) denotes the remainder terms. We remark that if we assume that v(i, n) ∈ Ω for some convex set Ω and all i ∈ [I], n ∈ [N ] then (by convexity of Ω) also w(i, k i ) ∈ Ω andEQUATION)Let us now (abusing notation again, it will be clear from context which function we refer to) writeEQUATION)We introduce some more quantities to get a concise representation of the Taylor expansion. We define the matrices D i,j L(k) ∈ R D×D containing the derivatives of a function L with respect to w(j 1 , n j1 ) and w(j 2 , n j2 ), i.e., we considerEQUATION)Similarly we define the vectorEQUATION)For the diagonal entries D j,j we need a more fine grained decomposition. Note that by the product rule we haveEQUATION)Using the notations we introduced we can therefore write (recall ĝ(k) = f (w(k)))EQUATION)Collecting finally all remainder terms asEQUATIONwe can thus summarize the Taylor expansion result as followsEQUATIONBased on the expansion (110) we now want to get an expression for the gradient of the total loss on the entire dataset. We decompose the dataset as followsEQUATION) Note that if Π i (n) ̸ = k i then D k,n,i = ∅.We also define similarlyEQUATION)We find the following expansionEQUATIONWe now rewrite or bound the four summands. First we relate those expressions to the datapoint statistics matrices. We defineEQUATION)andEQUATION)Then the first term equalsEQUATION)The second term is similarly given byEQUATIONThe third term is given byEQUATION)where δ was introduced in (98). Let us summarize those findings as a proposition. Proposition 1. Assume that the bound (15) holds for some Ω and v(i, n) ∈ Ω for all i ∈ [I], n ∈ [N ].Then the following expansion holdsEQUATIONProof. The result follows from the calculations above.The goal of this section is to apply the concentration bounds that will be derived in Appendix G to the expansion from the previous section which allows us to extract the (with high probability) dominating terms. Recall that we defined (11) as the maximal deviation norm and 13) as the maximal residual. Finally we introduce the notationδ max = max i∈[I],n∈[N ] |δ(i, n)| inr max = max k∈[K] I |ĝ(k) -g(k)| in (EQUATIONEQUATION)This impliesEQUATION)Theorem 7. Assume that the bounds in Theorem 9 hold and assume f satisfies Assumption 2 for some set Ω and v(i, n) ∈ Ω for i ∈ [I] and [N ]. Then we obtain the following expansion for the loss gradientEQUATION)Here the error terms are bounded byEQUATIONEQUATION)EQUATION)EQUATIONEQUATIONProof. The proof is a bit technical and essentially combines the assumptions and the concentration results in a straigtforward fashion. First we remark that for n ∈ Π -1 i (k) we have as stated before thatEQUATIONand B k,i n = 0 otherwise. Then Proposition 1 implies that (123) holds with the definitions of the error terms given below. We now define and bound the error terms in the decomposition. Note that we haveEQUATION)where we either bound both terms using Assumption 2 or the second term by r max and the first one by Assumption 2 together with M ′ ≤ M . This then implies that the error term capturing the fluctuations of the occupation statisticsEQUATIONcan be bounded byEQUATIONHere we used that there are K I-1 non-vanishing terms in the sum (ifΠ i (n) ̸ = k i we have B k,i n = EB k,i n = 0). Next, we consider the self interaction error termsEQUATIONEQUATIONFor E S,1 we get (similar to above) and using that the operator norm of a D × D matrix A is bounded byEQUATIONThe error term E S,2 could also be absorbed in other terms later on but since it is not dominant we just bound it. We can obtain similarly to our treatment of E S,1 the boundEQUATIONNext, we consider the interaction error termEQUATIONWe note that E(A k,i,j n1,n2 = c for some constant c if Π i (n 1 ) = k i and Π j (n 2 ) = k j and 0 otherwise. This implies that ifEQUATIONhere we usedEQUATIONwhich follows from the definition (10) in the last step. If Π i (n) ̸ = k i then this expression is clearly zero as the corresponding row of A k,i,j is zero. Thus we find thatEQUATIONHere we used the submultiplicativity of the operator norm of Kronecker-products stated in (280). Now we use | δ(j)| ≤ √ N δ max , the concentration bound (251), and bound ∥D i,j L(k)∥ ≤ DM similar as before to findEQUATIONFinally, we bound the Taylor expansion remainder termEQUATIONEQUATIONandEQUATIONwhere h = h n,i,α was introduced in (100).Assuming the bounds ( 15) and ( 16) we findEQUATIONIndeed, 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 followsEQUATION)Thus we obtain using (252)EQUATIONIf we consider the gradient descent dynamics we obtain the following expansion for the time derivative ẇ(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 getEQUATIONwhere E w can be bounded byEQUATION)for S ≥ N ln(N ).Proof. The result follows from Theorem 7 and the definition of w(i, n). We also use the relationn∈|Π -1 i (k)| δ(i, n) = 0 (already stated in (139)) to conclude that terms involving δ(i, n) cancel. The condition S ≥ N ln(N ) allows us to bound ln(S)N/S ≤ 2 to control E T . Note that here we use that E S,1 and E S,2 can be bounded by E I . We therefore get the following bound for δ(i, n). Corollary 2. Under the same assumptions as in Theorem 7 we getEQUATION)where E δ can be bounded byEQUATION)for S ≥ N ln(N ).Proof. This follows from Theorem 7, Corollary 2, and the relationEQUATIONNow we are in the position to control the time evolution of |δ(i, n)| 2 . We obtain the following relation Lemma 1. Under the same assumptions as in Theorem 7 the following bound holdsd dt 1 2 |δ(i, n)| 2 ≤ -λ - ω (2K) I + 2 I sup v1,...,v I ∈Ω max i ∥D 2 i,i f (v 1 , . . . , v n )∥r max |δ(i, n)| 2 + C 1 min(1, r max )K (I-1)/2 ln(S) N S δ max + C 1 K (I-1)/2 ln(S) N 2 S δ 2 max + C 1 δ 3 max (153)where C 1 is a constant depending on D, I, and M , and ω is defined byEQUATIONProof. We haveEQUATIONRecall that by (122) we haveEQUATIONWe consider ω as defined in the statement of the lemma where ω ≥ 0 follows because we sum positive semi-definite rank one matrices. Then we find using the lower bound on the spectrumEQUATIONThus we find d dtEQUATIONwhere we used for the second to last contribution that we sum over K I-1 terms which is cancelled by the (2/K) I factor. We finally bound the error term byEQUATIONThis ends the proof.In this Section we prove our main results, Theorem 4 and Theorem 5. In addition we state and prove Theorem 8 which is a variant of Theorem 5 without a-priori bound on the embeddings. Essentially the proof strategy for both results is to rely on the groundwork from the previous two sections, in particular on Lemma 1. Indeed, we first verify that the conditions of Lemma 1 hold for all times and then application of this Lemma allows us control the evolution of δ max for all times.Proof of Theorem 4. The proof proceeds in three steps. First we use the monotonicity of the loss to deduce that r max (t) can be bounded in terms of r max (0) and δ max (t) and δ max (0). Then we show that choosing the variables as in the statement of the theorem allows us to bound all the terms in Lemma 1 in Appendix E. Then we apply Lemma 1 and deduce the decay of the maximum of the displacements |δ(i, n, t)|. Here we need to check carefully that the assumptions of the lemma are satisfied for all t.First we assume the conclusions on the concentration properties from Theorem 9 in Appendix G hold which occurs with probability at least 1 -S -1 over the randomness of the training data.Let us consider any time T > 0 and we assume that v(i, n, t) ∈ Ω for all i ∈ [I], n ∈ [N ] and 0 ≤ τ ≤ T . This allows us in particular to apply the Assumption 2 up to time T for embeddings v(i, n).We now implement the first step where we bound r max (t) for all 0 ≤ t ≤ T in terms of r max (0), δ max (t), and δ max (t). The idea is to upper bound the initial sample loss L(t = 0) and lower bound the loss L(t > 0). First we note that by a first order Taylor expansion and using Assumption 2EQUATIONThis implies thatEQUATIONNext, we derive a lower bound on the loss. Let k be such that | f (k) -f (k)| = r max . By assumptionEQUATIONand using concentration bounds (e.g., by summing (253) over n and using the lower bound for S) we find thatEQUATION)Thus we can lower bound the loss using the same Taylor expansion as before byEQUATION)Since gradient descent does not increase the loss the bounds (161) and (164) together imply for t ≥ 0EQUATION)This implies r max (t) ≤ 2(2K) I/2 (r max (0) + IM δ max (0)) + IM δ max (t).(166)Then we get the bound r max (t) ≤ 4(2K) I/2 (r max (0) + IM max(δ max (0), δ max (t))).(167)Now we choose C 2 and C 3 such that the initialization condition (20) becomesEQUATIONandEQUATION)Assume in addition that for all times 0 ≤ t ≤ T the bound δ max (t) ≤ δ holds. Using (167) we then find that for 0EQUATION)and similarlyEQUATIONand thereforeEQUATIONThe next step is to bound the various error terms appearing in Lemma 1. Using the last display we can bound for all i ∈ [I] and n ∈ [N ] using Assumption 2 for t ∈ [0, T ]EQUATION)Moreover, we can estimate for 0 ≤ t ≤ TEQUATIONNow we observe that for x = a ln 2 (a) and a ≥ 1 we find (using ln(a) ≤ a) the boundEQUATIONHence we considerEQUATION)Note that tracking only the dependence on K and N (and using K ≤ N ) this condition readsEQUATIONThen we find using (175) combinded with the observation that ln(S)/ √ S is decreasing for S ≥ e 2 for S ≥ S 0 the boundEQUATION)We note that similarly for x = a ln(a) we getEQUATIONThus we consider similarlyEQUATION)where δ was introduced in (168). Note that tracking only N and K this bound becomesEQUATION)Then, similar to before, (using monotonicity of ln(S)/S for S ≥ e) we find for S ≥ S 1C 1 K (I-1)/2 max(1, r max ) ln(S) N S ≤ C 1 K I/2 max(1, r max ) ln(S 1 ) N S 1 ≤ 1 16ω 0 (2K) I δ max(1, r max (t)).(182) After all these preliminary estimates we can now finish the proof.Let us define T ∈ R + 0 ∪ {∞} to be the maximal time such that v(i, n, t) ∈ Ω for all i ∈ [I], n ∈ [N ] and 0 ≤ t ≤ T and the bound (168) holds for δ max (t) where we for convenience recallEQUATIONWe want to show that T = ∞. We assume that T < ∞ and prove by contradiction that this does not hold. By (172) we know that for t ≤ T the bound r max (t) ≤ 1 holds and therefore Assumption 4 implies that w(i, k, t = 0) ∈ Ω 0 and since moreover δ max (t) ≤ 1 for t ≤ T we find that v(i, n, t = 0) ∈ Ω 0 + B 1 (0) ⊂⊂ Ω. By continuity of v(i, n, t) in t we conclude that there is ε > 0 such that v(i, n, t) ∈ Ω for all 0 ≤ t ≤ T + ε, i ω 0 (2K) I δ max (t)(-δ max (t) + δ/2).(185) Now, if δ max (T ) ≤ δ/2 we conclude by continuity that there is ε > 0 such that δ max (t) < δ for t ∈ [T, T + ε]. On the other hand if δ max (t) > δ we conclude that d dt 1 2 |δ(i, n, T )| 2 < 0 for all i, n such that |δ(i, n, t)| ≥ 3/4δ max . This in particular implies that δ max (t) < δ max (T ) for t ∈ [T, T + ε] and some ε > 0 (δ max is non-increasing at T ). This is a contradiction and we conclude that T = ∞.Finally we prove the decay of δ max . Assume that for t ≥ T the bound min(1, r max (t)) ≤ R holds. Note that the function δ max (t) is not necessarily differentiable but its left and right derivative exist (as the maximum of finitely many differentiable functions). Then we obtain, similar to (184), for the right derivative the boundEQUATIONForEQUATIONthe previous bound simplifies toEQUATIONEQUATIONwhich implies that for t ≥ T by Gronwall's LemmaEQUATIONEQUATION)Thus in finite time we achieveEQUATIONThis ends the proof. Note that we also get an exponential rate of convergence.Remark 3. The exponent 9/2 for the lower bound of S in (181) is not tight because we could use that r max is small but this provides only a small improvement that does not justify the additional technicalities.Before proving and stating Theorem 8 let us first prove that the weight decay λ > 0 allows us to derive a-priori bounds on the time evolution as stated in the following lemma. Lemma 2. Let Π be approximately balanced, i.e., assume that Assumption 1 holds. Let M 1 > 0 be such that sup v(1),...,v(I)∈B1(0EQUATION)Let 0 < λ ≤ 1 be a fixed number and assume that v(i, n, t) follow the gradient dynamics (8) and are initialized such that|v(i, n, t = 0)| ≤ 1 (192)and assume that the dynamics exists for all times. Then for all i ∈ [I], k ∈ [K], and all timesEQUATIONProof. Note that by assumptionEQUATIONThen we conclude thatEQUATION)Thus we conclude thatEQUATIONThis ends the proof.Now we state a refined version of Theorem 5 which does not require an a-priori bound but instead proves boundedness of the embeddings v(i, n, t) for all times using Lemma 2. Theorem 8. Let Π be approximately balanced, i.e., assume that Assumption 1 holds. Assume that the function f : R ID → R is slot-wise linear. Assume that there is M 1 > 0 such that sup v(1),...,v(I)∈B1(0EQUATION)Let 0 < λ ≤ 1 be a fixed number and assume that forEQUATIONAssumption 2 holds with Ω = B R (0) and some M > 0. Note that M depends on R and thus on K and λ. Let C 1 be the constant from Lemma 1. Assume that v(i, n, t) follow the gradient dynamics (8) and are initialized such thatEQUATION)Then there are constants C 2 , C 3 ≥ 0 depending on M , I, D, λ, and C 1 such that forEQUATIONthe boundEQUATIONholds for all t ≥ 0. Remark 4. We emphasize again that since M might depend on λ we obtain no explicit rate in λ. However, the S and N dependence is the same as in Theorem 5.Proof. The general strategy of the proof is similar to the proof of Theorem 4 but the proof is slightly simpler as it is easier to obtain a-priori estimates on the evolution of v(i, n, t). We assume that the conclusion of Theorem 9 holds which occurs with probability at least 1 -S -1 over the randomness of D. Now let T ≥ 0 be the largest time such that for all 0 ≤ t ≤ T and all i ∈ [I] and n ∈ [N ] we have v(i, n, t) ∈ Ω and the boundEQUATION(we assume w.l.o.g. that C 1 ≥ 1) holds for 0 ≤ t ≤ T . By assumption those relations hold at t = 0 (note that δ max (t) ≤ 2 max i,n |v(i, n)|). We can bound for t = T using the a-priori estimate (193) from Lemma 2 and findEQUATIONThus, by continuity, there is ε > 0 such that v(i, n, t) ∈ B R (0) for some ε > 0 and all t ∈ [0, T + ε]. Now we can apply Lemma 1 for the interval [T, T + ε]. We note that by slot-wise linearity we have D i,i f = 0 and ω ≥ 0. Therefore we getEQUATIONNow (with the same reasoning as in the proof of Theorem 4) there is a constant C 2 depending on D, I, and M such that forEQUATIONholds. Clearly we get for a suitable constant C 2EQUATIONSimilarly, we get forEQUATIONthe boundEQUATIONholds. And we findEQUATIONNow we can continue to bound (204) for S satisfying (205) and (208) and using (202) as followsEQUATIONNow, if δ max (T ) ≤ λ/(8C 1 ) then there is ε > 0 such that δ max (t) ≤ λ/(4C 1 ) holds for t ∈ [0, T + ε]. Thus we assume that δ max (T ) ≥ λ/(8C 1 ). Let i ∈ [I], n ∈ [N ] be any index such that |δ(i, n, T )| = δ max (T ). Then we conclude thatEQUATIONThis implies thatEQUATIONfrom which we conclude that there is ε > 0 such that δ max (t) ≤ λ/(4C 1 ) holds for 0 ≤ t ≤ T + ε.In either case we get a contradiction and thus T = ∞. In particular we can apply Lemma 1 for all times t ≥ 0. Suppose now that for some tEQUATIONthen we conclude from Lemma 1 that d dtEQUATIONThus we conclude from Gronwall's inequality thatEQUATIONThis ends the proof.Proof of Theorem 5. The proof is the same as the proof of Theorem 8 above with the only exception that we do not need to show that v(i, n, t) ∈ Ω but this holds by assumption which makes the proof strictly simpler.EQUATIONfor n 1 ∈ N (i, k), n 2 ∈ N (j, k). Thus, the entries are the number of datapoints that are mapped to k by Π and whose entries i and j are equal to n 1 and n 2 . Note that this agrees with the definition in (115) given in Section D. Applying concentration bounds for random matrices derived in Appendix H below we obtain the following result. Lemma 4. Let A k,i,j ∈ R N (i,k)×N (j,k) be as defined above. Assume that for all i ∈ [I] and k ∈ [K] the bound N/(2K) ≤ |Π -1 i (k)| ≤ 2N/K holds and let S ≥ N . Then the following bound holds for η > 0EQUATION)Note that for K ≥ 2 we get for S ≥ N the boundEQUATION)Proof. The proof is essentially the same as the proof of Lemma 3 with some additional notational complications. Again we consider a dataset D generated by first sampling S ∼ Poi(S) and thenEQUATIONand thus the distribution of A k,i,j isEQUATIONfor all n 1 ∈ N (i, k) and n 2 ∈ N (j, k) and the entries are independent. Note that the assumptionEQUATIONNow we apply Lemma 7 where we use η ln(S) as η in the statement of the lemma and find thatEQUATIONNote that P( S = S) = S S e -S /S! ≥ 1/(3 √ S). This impliesEQUATION)then we can writeEQUATIONwhere E i,j is the matrix whose entry (i, j) is 1 and all other entries vanish and N i,j ∼ Poi(λ). Let X ij = (N i,j -λ)(E i,j + E j,i ). Note that by induction one directly findsEQUATIONIn any case we get ±(E i,j + E j,i ) k ≤ E i,i + E j,j in the sense of symmetric matrices and we set A ij = E i,i + E j,j . Now we obtain for any θ ∈ R the relationEQUATIONNote that the first summand is invariant under θ → -θ while the term in absolute values changes its sign. This implies∞ k=0 θ 2k E(Poi(λ) -λ) 2k (2k)! + ∞ k=0 θ 2k+1 E(Poi(λ) -λ) 2k+1 (2k + 1)! = max ∞ k=0 θ k E(Poi(λ) -λ) k k! , ∞ k=0 θ k E(Poi(λ) -λ) k k! = max Ee θ(Poi(λ)-λ) , Ee -θ(Poi(λ)-λ)= max e λ(e θ -1)-λθ , e λ(e -θ -1)+λθ = e λ(e |θ| -1)-λ|θ| .(268)Here we used in the last step that λ ≥ 0 and then for θ ≥ 0 e θ -e -θ = 2 sinh(θ) ≥ 2θ.Thus we infer from (267) and (268) thatEe θXij ≤ e λ(e |θ| -1)-λ|θ| A ij = e (λ(e |θ| -1)-λ|θ|)Aij := e g λ (|θ|)Aij .using that A k ij = A ij in the second step. Note that hereEQUATIONwhich implies EQUATIONLet us here state some basic properties of the Kronecker product of matrices. For two matrices A ∈ R n×m and B ∈ R p×q the Kronecker product is defined byEQUATION)We will need the property that the operator norm satisfiesEQUATION)Note that we use the notation v ⊗ w = vw ⊤ to denote the outer product. Formally we have vw ⊤ = v ⊗ w ⊤ but it is convenient and common to drop the transposed here.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: The assumptions required for proving GD are many and some appear to be strong and unverifiable (?).
questions: questions: Is it possible to (approximately) verify whether a given data set satisfies the assumptions you need to have the GD guarantee?
limitations: limitations: Yes.
paper_summary: paper_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)$.
strength_weakness: strength_weakness: - 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.

-  The problem setting is a bit restricted. It could be interesting to explore extensions such as: 

1) 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$.  

2) Introducing some dependency between tokens and potentially allowing one item to belong to several clusters.


Minor Nits:
- Page 2 line 74. “such that”?
- Page 3, Line 117. Duplicates of “such as” and “e.g.,”
- Page 4, line 135 “need of the order”
- Page 5, equation (7), it should be $||u||$? u should be a vector.
questions: questions: 1. Page 4, Line 145, what is $N_0$?
2. 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: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_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).
strength_weakness: strength_weakness: 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.
2. 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.
3. 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.
questions: 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: limitations: The authors have properly addressed the limitations of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: The usage of the word "interact" in the abstract is rather vague.

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

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

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

Gradient descent seems like a rather crude approach.
questions: questions: Is assuming knowledge of $k$ realistic?

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

From section 4, it seems that the method needs at least $O(DMI)$ memory. Is this scalable?
limitations: limitations: Assumption 1 is quite strong. In practice, one often encounters highly imbalanced clusterings.

The paper does not include experiments on real data.

Comments: title: After Rebuttal
comment: I acknowledge the rebuttal. The authors might consider proving that their required conditions hold (unconditionally) for interesting families of inputs defined by other means. There is no need to respond to this comment.
I thank the authors for their rebuttal, which addresses most of my concerns. I will increase my score. The paper still lacks experiments on real data. Adding experiments on real data would really help make the proposed problem setting more interpretable.
Thanks for the response. I will keep my original score.
Thank you for the reply. I will keep my score.