Explaining Grokking and Information Bottleneck through Neural Collapse Emergence
Authors: Keitaro Sakamoto
Issei Sato
Venue: ICLR 2026 Poster (2026)
Keywords: deep learning
grokking
information bottleneck
neural collapse
training dynamics
Primary area: optimization

Abstract: The training dynamics of deep neural networks often defy expectations, even as these models form the foundation of modern machine learning.
Two prominent examples are grokking, where test performance improves abruptly long after the training loss has plateaued, and the information bottleneck principle, where models progressively discard input information irrelevant to the prediction task as training proceeds.
However, the mechanisms underlying these phenomena and their relations remain poorly understood.
In this work, we present a unified explanation of such late-phase phenomena through the lens of neural collapse, which characterizes the geometry of learned representations.
We show that the contraction of population within-class variance is a key factor underlying both grokking and information bottleneck, and relate this measure to the neural collapse measure defined on the training set.
By analyzing the dynamics of neural collapse, we show that distinct time scales between fitting the training set and the progression of neural collapse account for the behavior of the late-phase phenomena.
Finally, we validate our theoretical findings on multiple datasets and architectures.

Published as a conference paper at ICLR 2026
EXPLAINING GROKKING AND INFORMATION BOTTLE- NECK THROUGH NEURAL COLLAPSE EMERGENCE
Keitaro Sakamoto & Issei Sato
Department of Computer Science
The University of Tokyo
Tokyo, Japan
(sakakei 199a, 1999, sato]0g ecc. ucw..eeoo..uu.ooa....
ABSTRACT
The training dynamics of deep neural networks often defy expectations, even these models form the foundation of modern machine learning. Two prominent examples are grokking, where test performance improves abruptly long after the training loss has plateaued, and the information bottleneck principle, where mod- els progressively discard input information irrelevant to the prediction task as training proceeds. However, the mechanisms underlying these phenomena and their relations remain poorly understood. In this work, we present a unified expla- nation of such late-phase phenomena through the lens of neural collapse, which characterizes the geometry of learned representations. We show that the contrac- tion of population within-class variance and information bottleneck, and relate this measure to the neural collapse mea- sure defined on the training set. By analyzing the dynamics of neural collapse, we show that distinet time scales between fitting the training set and the progression of neural collapse account for the behavior the late-phase phenomena. Finally, we validate our theoretical findings on multiple datasets and architectures.
key factor underlying both grokking
INTRODUCTION
Deep neural networks (DNNs) have demonstrated remarkable success across a variety of tasks, cluding computer vision, natural language processing, and reinforcement learning, yet their training dynamics under gradient descent often reveal unexpected behavior. In particular, when training continues beyond the point where the training loss has been sufficiently reduced, several intriguing late-phase phenomena have been reported. One such phenomenon is grokking (Power et al. 2022), where models initially converge to generalize to unseen data. However, when training continues for a sufficiently long time, the mod- els unexpectedly generalize. Another example the information bottleneck (IB) framework (Tisbyy) eet al. 2000] Tishby & Zaslavsky 2015, Shwartz- -iiartz--iiiva & Tishby] (207)). an informmaiion-hhorree perspective on representation learning DNNs that formalizes the goal of retaining fask-relevant information while compressing the task-irrelevant input information. One particularly intriguing ob- servation here that DNNs do not move directly toward an IB-optimal solution; instead, they first enter fitting phase where they memorize the training data, followed by later compression phase the later training stage during which task-irrelevant input information is discarded.
overfitting solution that perfectly fits the training data but fails
Taken together, these phenomena share the characteristic that the network evolves toward a more desirable state in the late phase of training, suggesting that some internal change occurs within the training dynamics during this transition. However, the mechanisms driving these phenomena, as well as their relationships, remain poorly understood. Bridging this gap is essential for deepening understanding of DNN training and for developing more effective training strategies.
In this work, we focus on the geometric structure of the network's 's representation space and, for the first time, demonstrate that the dynamics of neural collapse (Papyan et al. 2020) provide unified
explanation for these late-phase phenomena, offering new insights into their underlying mecha- nisms. More specifically, the contributions of this work summarized as follows.

---

Published as a conference paper at ICLR 2026
Theorem 4,3
Train
Test
Phase 1 Fitting Training Set
Phase 2 Neural Collapse
Accuracy
Theorem 3.2
Grokking
Time
Theorem 4.1
Time
Phese 2
Compression
Theorem 3.4
(02:0)
Phase
Population
Fitting (Proposition 3.3)
Within-Class Variance
1(2: X)
Neural Collapse
Information Bottleneck
Figure Conceptual relationships in the late-phase training discussed in this work.
First, show that the contraction of within-class variance in representations plays crucial role both grokking and IB dynamics. Specifically, for grokking, we derive an upper bound on the generalization error in terms of the population within-class variance of the learned representa- tions (Theorem (.2).. Similarly, for the IB principle, we show that the redundant information in the representations, which is discarded in the IB compression phase, bounded by the popula- tion within-class variance (Theorem 3.4). These results motivate the analysis of neural collapse, whose properties include the collapse of empirical within-ch representations in the training data, known as NC1 (Section B).
Second, we provide a quantitative analysis of the discrepancy between the population within- class variance and its empirical counterpart, using an approach analogous to generalization error analysis (Theorem I.T.. This allows us to evaluate to what extent the reduction of empirical within-class variance, that is, the progression of neural collapse, implies corresponding reduc- tion in the population within-class variance, a key quantity we identified. In this way, we relate the behaviors of grokking and IB dynamics to the development of neural collapse (Section
Finally, building the preceding results, analyze the development of neural collapse by explicitly tracking the dynamics gradient descent. Leveraging the results Jaoot et al.) (2025). we establish that the empirical within-class variance decreases during training and characterize its time scale (Theorem 4.3) In particular, depending on the strength of weight decay, the time scale on which neural collapse is sufficiently realized can lag far behind the time scale of fitting the training data. This result suggests that the timing of neural collapse emergence underlies the delayed generalization in grokking as well as the compression phase IB dynamics (Section
These relationships illustrated in Figure I along with the corresponding theorems and propo- sitions presented the main text. In addition to the thcoretical analysis, we validate our findings through extensive experiments various datasets and architectures. Taken together, our work deep- ens the understanding of late-phase training phenomena from the perspective of neural collapse.
Notation. We consider K-elass classification problem with a dataset S {10), yr ) con- sisting of N examples, where 2; Rd is the input and y [K] {1,2. K)) is the class label. The input domain is X' C Rdd, and yi denotes the one-hot encoding label in {0, 1]K Let X E RdxN and Y {0, 1] KxN X denote the training inputs and one-hot labels, respectively. We use S. = {( mi) S y/ e} to denote the subset of class-c examples, and, by abuse of nota- tion. also the corresponding index set E [N] y: = c}. We denote its cardinality as n. = [Sc]- Let X and Y denote the random variables representing the input and label, respectively. We write px.y for the joint distribution over (X, Y), and px and py for their marginal distributions. When the variables are clear from the context, we simply write p to refer to the corresponding distribution. Let T(X; Y) be the mutual information (MI) between random variables X and Y. We denote the multivariate Gaussian distribution with mean μ and covariance matrix Σ as N N E). We use log(-)
Code available at http///aoooa//////a//o/aa////aaa
2

---

Published as a conference paper at ICLR 2026
to denote the natural logarithm. For linear-algebraie notation, for matrix A, we use || to denote the spectral norm, and || 12.1 to denote the (2,1) matrix norm, defined as || = E, | 1|,
where a; is the i-th column of A. Finally, we use (a, b) = a. to denote the standard Euclidean inner product.
RELATED WORK
Grokking, On the empirical side, several studies attempted to explain the cause of grokking in terms of the parameter compression (Liu et al., [2023a] Varma et al. 2023) and some complexity measures (Nanda et al., 2023 Liu et al. 2023b Humayun et al. 2024| DeMoss et al. 2025). Other studies sought to relate grokking to other training- rellteed concepts. such as double descent (Davies) jet al., 2023] Huang al. 2024) and optimization stability (Thilak al. 2022). On the theoretical
side, -igh--imensional limit of the linear model was analyzed in the setting of regression Leevi ct al. 2024) and binary classification (Beck et al. 2025). There are also studies analyzing two-layer networks with XOR data (Xu et al.] 2024) and mean-field analysis (Rubin et al. 2024). The most closely related line of work focuses on the transition from the kernel regime to the rich regime (Lyu al. 2024 Kumar al. 2024]. Our study provides novel perspective based on the emergence of
neural collapse, offering a new connection to the representation-leaming view of Liu et al. (0022b).
Information Bottleneck. The IB principle (Tishby et al. 2000) has been analyzed in the context of deep learning by modeling the successive neural network layers a Markoy chain (Tishby & Za- slavsky| 2015 Shwartz-Ziv & Tishby 2017; 1 Saxe et al.) 2018; Geiger 2021 Lorenzen et al. [022]] Adilova et al. 2023 Butakov et al.) 2024a b). In addition these studies analyzing DNN training from the IB perspective, minimizing the IB objective has also been shown to benefit generalization bounds (Kawaguchi et al., 2023 Sefidgaran et al., 2023). Several studies have attempted to explain the IB dynamics. For example, Shwartz-Ziv et al. (2019) attributed the compression phase to the diffusion component of stochastic gradient descent (SGD), while Goldfeld et al. (2019) and Koch & Ghosh (2025) discussed its relation to geometric compression and grokking, respectively. How- ever, these studies lack the rigorous theoretical analysis of the IB dynamics. A related line of work analyzes late- thteephase behavior through reconstruction loss rather than mutual information (Schneider & Prabhushankar 2024, [Schneider] 2025), though this perspective provides only an indirect proxy for the IB dynamics. In contrast, our work offers new explanation of IB dynamics based on neural collapse, a geometric form of compression.
Neural Collapse. Neural collapse (Papyan al. 2020 a late-stage training phenomenon where the representations of the training data converge to a simplex equiangular tight frame (ETF) formed with the class mean representations.
major line of research analyzes the optimality of such so- lutions under an unconstrained feature model (UFM) and its variants (Fang et al. 2021 1 Mixon
al. 2022 Lu & Steinerherger, 2022, Tirer & Bruna 2022] Thrampoulidis et al. 2022 Dang iet al. 2023 Tirer et al.) [023] Sukenni et al. [023 2024| Jiang al. (2024), where the features
are treated directly as optimization variables. Several works have examined the learning dynam-
ics toward neural collapse under UFM (Zhu et al.) 202T Mixon et al. 2022 Ji et al., 2022 Zhou et al. 2022a b). but this setting cannot reflect the input data properties and diverges from the actual parameter-basee training dynamics. Beyond UFM setting, recent studies instead analyze parameter updates under weight regularization, opening a promising new direction (Jacot er al., 2025) Wu & Mondelli 2025). Building on this, our study broadens the understanding of training dynamics by revealing the connection between late-phase phenomena and neural collapse.
3 EMERGENT BEHAVIOR IN LATE-PHASE TRAINING
In this section, we examine two intriguing phenomena that arise in the late-phase training and remain active topics of research: grokking and IB dynamics. Through separate analysis in the following subsections, we demonstrate that the population within-class variance in the representation space plays a critical role and constitutes unifying factor. To this end, in the following, let g Rd + d. denote the feature extractor, which may take various model architectures, W E R xxx the last layer classifier, and f 2d + R with (xx = W (() the full model. We use f(x); to denote the i-th component of the prediction vector f(x).
3

---

Published as a conference paper at ICLR 2026
We now introduce the population within-class variance, which will play central role in this sec- tion. We remark that this metric is not introduced as an ad-hoc measure; rather, it naturally arises from our analysis of the second-phase dynamics presented below. In particular, the scale-invare invariaat formulation in Definition 3.1 ensures that this quantity reflects genuine geometric concentration, disentangled from changes in output scale. This motivates the following definition.
Definition 3.1 (Population within-class variance), To evaluate variance in a scale-invariant manner, we consider the expected within-class variance of the rescaled feature extractor, defined as
g(x)) Bg
Ex|y= [|1((X) [x where g(x)
Bg sup (1g(2))l2. xE
This quantity, as well as its expectation over the label distribution py, serves as a key component in the results of this section. The results established here using the population within-class vari- ance then motivate the subsequent analysis of neural collapse, which characterizes the reduction of empirical within-class variance in the training data.
3.1 GROKKING: PHASE GENERALIZATION DYNAMICS
In the context of grokking, Liu et al. (2022b) empirically shows that the acquisition of structured representations leads to a transition from memorization to generalization, which we further analyze here. As a simple observation, since the network output is given by f(x) = W Wg( achieving good generalization facilitated when the representations of each class are linearly separable in distribution. When the training loss is sufficiently reduced some time step T1, the existence of weight matrix (Ti) indicates that the representations ble; however, this does guarantee that the representations in distribution are linearly separable. Intuitively, when the representations of each class are more tightly clustered, achieving linear sep- arability becomes easier, and better generalization performance can be expected. Formalizing this intuition yields the following theorem.
the training set can be linearly separa-
Theorem 3.2 (Generalization via population within-class variance). For a fixed feature extractor and the last layer W (wi...., w we have
-1
W.-Wk max (EXI (x()) 1| -
Pr arg max [f(x) # y
S (c)
Σ
1
[E[K]
Appte
Exiy=c [|x(x) -EIIx=
This theorem follows directly from applying the union bound and variance- aarsed tail probability bound. The proof is given in Appendix [..]] Theorem 3.22 states that improving test accuracy is facilitated by two conditions: i) the class mean representations Exjy== [x()] become more aligned with the corresponding last-layer weights w., and ii) the population within-class variance decreases. From a grokking perspective, when the training loss is sufficiently reduced, it remains unclear how well the classifier W aligns with the class means of the representations. Nevertheless, independent of the properties of W, reducing the within-class variance of the representations tightens the upper bound on the generalization error. To uncover the mechanism of grokking, remains to show that the within-class variance decreases even after the training accuracy has been sufficiently improved.
3.2 IB DYNAMICS
The IB principle (iihbyy et al. 2000) formulates a constrained optimization problem: it seeks a compact representation Z of the input X that retains as much information as possible about the target Y while compressing X. This formulation is built on the idea that a concise short code can extract the features of X essential for predicting Y, and the IB principle serves as one approach to explaining the representations acquired by DNNs. Under the Markov chain Y + X Z, this can be formulated using MI as finding the conditional distribution [(Z that minimizes
min (Z: -1(21))
p x(
where 3 > 0 is a trade-off parameter controlling the balance between compression and information preservation. When analyzing DNNs within the IB framework, we encounter several difficulties:
4

---

Published as a conference paper at ICLR 2026
for a deterministic network and continuous representation. [2; X) can be infinite, making the analysis ill ill-defined; furthermore. the network parameters are not reflected in the IB analysis (Saxe et al.] 2018: Amjad & Geiger 2019; Goldfeld et al., 2019). To address these issues, analyze the representation after independently adding an arbitrarily small Gaussian noise E ~ ((, x²
1. which is standard approach in IB analysis of DNNs (Saxe et al. 2018 Goldfeld et al., 2019 Buttaov et al. (2024b). It should be noted that such a small amount of noise has a negligible
effect on the network's output, making it good proxy for the actual network. Since we are currently
focusing on the output of the feature extractor g. we denote the representation as Z = ((X)(() B
where B, = (Ig()) I| denotes the output scale of g introduced in Definition 3.1 We use
to denote the realization of Z.
Representation dynamics in the two-dimensional (T(Z; X), plane, which called infor- mation plane, is useful tool for analyzing the training dynamics of DNNs from the IB perspective a preliminary observation, to perform classification accurately
(Shwartz-Ziv & Tishby| 2017). with neural network, both [[2; X)) and (Z; Y) need to be sufficiently large; this follows from the following proposition (see Appendix B.2 for the proof):
Proposition 3.3 (Phase of IB dynamics), For any last-layer classifier W, let lcEly.2, z) denote the cross-entropy loss between the target [K] and the predicted logits Wz z E IR defined by
(CE(y.2) z) -oog log (exp((W 2)u)/ Emn exp((Wz)oo) Then, we have
I(Z 2 I(2;Y) -Br,z (Y Z)] + const.
Proposition 3.3 implies that if the network at initialization discards information about both X and Y. the training process must appropriately increase I(Z;Y) and consequently I(Z; X). over time. Please note that if the initial state of the network sufficiently preserves information about X and Y without collapsing the outputs, then this MI increase phase is unnecessary.
An intriguing observation in the existing information plane work is that, in the late stage of training, DNNs tend to compress T T(2; X) while preserving T(Z; Y), thereby moving toward more optimal solution with respect to the IB objective in Equation ( We explain this behavior through the following theorem, which is based on the degree collapse of the population within-class variance, quantity introduced in Definition 3.I
Theorem 3.4 (Phase 2 of IB dynamics via population within-class variance). Let Z x(X)+ Byg, where E ~ NO, o² I). Here, the variance o² > O is chosen to be small enough to ensure negligible effect on the network output. Then, the superftuous information is bounded as follows:
1 T(Z:X) - - < 202 Ex.Y [|((x) Exiy ((X)]12]
=
We further show the tightness and behavior of this upper bound in Proposition B33 Theorem B.4 the role of within-clat classs variance in reducing superfluous information, corresponding to the evolution toward the upper-left in the information plane. Together with Theorem 3.2, the analysis this section establishes population within-class variance as a key measure. In the next section, examine how this measure decreases during training and, for grokking, how it proceeds on a different time scale from fitting training set.
The proof is given in Appendix B.2
4 EVOLUTION OF WITHIN-CLASS VARIANCE
In this section, motivated by the previous results, we analyze the training dynamics of within-class variance in the learned representation space. We first reduce the discussion of within-class variance to neural collapse, namely the geometric arrangement of class-wise representations in the training set. We then examine how neural collapse progresses under gradient descent, independently of the training loss convergence. These results together shed light on the evolution of within-class variance.
4.1 POPULATION WITHIN-CLASS VARIANCE AND NEURAL COLLAPSE
we examine how the upper bounds based on within- variance can be approximated using training data, Since this requires delving into the specifics of the trained
Building on Theorems 3.2 and 3.4
5

---

Published as a conference paper at ICLR 2026
model, for the remainder of the paper consider the following standard DNN:
f(x) W W.g(2)) oot (Wi-10 (...a (Wiz)...))
(2)
where We E _dexde1,, with dg d denoting the input dimension. Note that Wt and d respectively correspond to W and dzep used in the previous section. Here denotes the element- wise activation function with (0) and 1-Lipschitz, while oout is the activation at the output of the feature extractor g. Throughout Section [4]]] we consider the standard setting dout o, and in Section 4.2
we set out. id, i.e., the identity map, for analytical convenience.
We first provide formal bound on the difference between the population within-class variance, which served as an important measure in the previous section, and its empirical counterpart. The following result is based standard approach of uniform convergence commonly used in gener- alization error analysis. The proof is provided Appendix B33
Theorem 4.1 (Concentration of within-class variance). Suppose the input domain X is bounded, i.., | la I B2 for all x E X. Let ò E (0, 1) and dmax maxogig de, and recall Bg = SuP (1g(x)]l2- Define the complexity measures of 9 as [T(g) = max [IL Well,, } and A(g) = max { {EE'' (l]Wella.1/ Welle)) 3/2 .1). Then, with probability at least 1 - 8, for all
e E [K], we have
ExIY=0 [10(x) - Ex 10 [1)]
nr Σ
g(ar))
1 - nh E jES,
g(0r))
iES.
B.)
V/lo(K//)) + log log (T()))((0) )
log(ne)/logddmmaal log(ne) (og(dmmx
/Tc
|/Tcc Bg
Note that in l[g) both the numerator and the denominator represent the output scale of gi the numerator spectral-norm-basse upper bound, while the denominator reflects the actual output scale. This theorem establishes O(1/ bound on the gap between the population within-class variance studied carlier and the empirical within-class variance the training data. This guarantee justifies focusing on the empirical variance in the subsequent analysis, where investigate its dynamics as a proxy for the population behavior.
4.2 NEURAL COLLAPSE DYNAMICS
Analyzing the variance of class-wise representations in the training data naturally motivates the study of neural collapse. We therefore begin by defining the neural collapse metrics. Neural collapse refers to several characteristic properties the late stage of training: (NCI) the representations collapse to their respective class means; (NC2) the class means form an ETF; and (NC3) each row of WL aligns with the corresponding class mean up a positive scaling.
1 Eies ((x;), and denote their counterparts for ğ as jte Let Me 1 Eies, g(x;) and μG = N
and jag. Among neural collapse metrics, we focus on NC1, defined as NC1 Tr(Ew) where Tr(Ea)
the within-clas covariance Ew N/2 E (g(x.) me) (g(x,, me) and the between- class covariance Σg >K1 =c= (+t. ma) (ML pa) Since we consider the within-class vari- ance rescaled instead of dividing by the between-class variance, as appears in Theorem 4.1 we define the following measure as the rescaled NC1 (RNC1):
K
RNC1 = Tr (Tw)
Σ Σ 10((;)) fic 12..1.
B3
=
N c=1 iES.
Remark 4.2 (Difference between RNCI and NCI). While both metrics are invariant to the scale of g. NCI, which is normalized by the between-class variance, does not reduce to zero when all features collapse to a single point. This reflects class-center separation, a property already captured by NC2, whereas RNCI isolates the within-class variance aspect more directly than NC1.
We next specify the training setup for analyzing the dynamics of neural collapse. The network f is trained by gradient descent with step size 17 >> 0. The loss function is the squared loss with weight

---

Published as a conference paper at ICLR 2026
decay, controlled by a hyperparameter 0, and is defined as follows:
@ 1)=+1)=0011 ((+=))ee1r- E (0(T)) 1 (x) -
YI
where θ is the concatenation of all network parameters (W)) and f- the network defined in Equation (2) at time step Accordingly, we denote the value of RNCI at time step as (T)
We now analyze the dynamics of neural collapse measured in terms of RNC1, together with the convergence of the training loss. Our results build on the recent results of Jacot et al. (2025), which assume (i) a pyramidal network architecture, (ii) a smooth activation function, and (ii) a specific initialization condition; their formal statements are given in Appendix B.4 The next theorem estab- lishes that the training loss and RNC1 converge on different time scales. Here we use the standard Oig-O notations (-) and Q2-)) to describe the dependence on X, 17, C1, and 62-
Theorem 4.3 (Time scales of neural collapse dynamics), Suppose that the network f satisfies As- sumptions B.3 to B.5 and that the input domain X' is bounded. Fix < E1 < $ minee[[] ne and E2 For weight decay X O(1)), learning rate L = ((e2), and time steps Ti < Th satisfying
TI Ω ( 10g - 2)
log -
T2 =
Am
62
the regularized training loss and RNCI are bounded as
Ex (0( 61, RNCI 0 +22
This theorem not only shows that RNCI 1 indeed decreases under gradient descent training, but also clarifies the time scales that govern the convergence of the training loss and RNCI. For the target thresholds é1 and 62, both require a logarithmic order of training steps in their reciprocals. On the other hand, while the convergence of the training loss is independent of the weight decay parameter X, the time scale for the convergence of RNC1 grows inversely with smaller X. This implies that 72 can be much larger than тı depending on the value of X, indicating that the convergence of RNCI may occur substantially later than that of the training loss.
Remark 4.4 (Summary of Theoretical Results). Up to this point, as summarized in Figure I we have developed our analysis starting from the late-phase phenomena. We now provide an overall summary of the insights obtained from our theoretical results in Sections and
Grokking. By combining Theorems 3.2 and 4.1 we showed that the decrease of the empirical within-class variance, namely RNCI, leads to improved test accuracy. Theorem 4.3 further es- tablished the time scale governing this decrease, jointly with the convergence of the training loss. In particular, when 71 < T2 in Theorem 4.3 for example with a small weight decay X, neural collapse occurs later than the convergence of the training loss, and generalization improvement lags behind fitting training set; that is the grokking behavior.
IB Dynamics. For the first fitting phase, Proposition 3.3 demonstrated that this phase is neces- sary whenever the network discards input information. Unlike the first phase of grokking, this fitting phase is not necessarily tied to training loss convergence. For the second compression phase, (he4oemm33.4jand444 ssooweed that it proceeds together with the decrease of RNC1, whose convergence and time scale were established in Theorem 4.3
5
EXPERIMENTS
In this section, we conduct experiments to validate the theoretical results in Sections and
5.1 GROKKING
We first analyze the relationship between grokking, within-class variance, and neural collapse. Fol- lowing Liu et al. (2023a), train an MLP the MNIST dataset (LeCun et al., 2010). The model has four layers with architecture [784, 200, 200, 200, 10) and ReLU activation. The initialization scale is increased by factor of 8 in Liu et al. 2023a), and use the AdamW optimizer

---

Published as a conference paper at ICLR 2026
Train Examples
Unseen Examples
Density
Dersity
-1.5
-1.0
-0.5
0.5
-20
-1.5
-0.5
10
155
Signed distance decision bourdary (class
Signed
boundary
(a) Overlitting phase (time step 71 16,000). Train accuracy = 100%, test accuracy = 15%.
Train Examples
Unseen Examples
Dersity
Dersity
-1.5
-1.0
-0.5
0.5
-2.0
-15
-1.0
-0.5
10
155
Signed distance decision bourdary (class
Signzd
boundary
(b) Convergence phase (time step = 100,000). Train accuracy = 100%, test accuracy = 88%.
time steps during grokking. The margin of each ample is defined as the signed distance from its representation to the decision boundary determined by the last-layer classifier, calculated ((wo wi.g(ar (())) + bo (o-bb)) b1) bb1)/
Figure 2: Margins of individual examples
wille. where 6. de- notes the bias term for class We trained a 4-layer MLP on the MNIST dataset. These results reveal
the link between representation variance and generalization, supporting Theorems 3.2 and 4.1 We
additionally provide similar visualizations for several other class pairs in Appendix D.1
[Loshchllov & Hutter 2019) with learning rate of 1c-3 and weight decay of 0.01. Additional results on other datasets and architectures, including convolutional neural networks (CNNs) and Transformers (Vaswani al. 2017), presented Appendix D.I In examining the dynam- ies of neural collapse, we primarily track the RNCI score, which is the focus of our theoretical analysis. As an additional geometric metric, we also report the NC2 score, another geometric as- pect of neural collapse. We define it as the condition number of the matrix of class mean vectors, NC2 n
=( ((H).... .K)) We include NC2 to provide supplementary geometric insight into the arrangement of the class mean vectors.
We begin by motivating our approach of analyzing grokking through the lens of representation learning. Figure 2 shows how the representations of training and unseen examples evolve in a set- ting where grokking occurs. the overfitting phase (top), the training examples are separated, but the unseen examples exhibit large within-class variance despite their mean shifting toward the cor- rect class, resulting in many misclassifications. As training proceeds, the training examples become further separated and collapse into single points (bottom). At this stage, discussed in Theo- rem 4.1. the collapse of the training representations is to some extent inherited by the underlying distribution, leading to the reduction of the population within-class variance, illustrated in the right panel. Consequently, test accuracy improves, and grokking emerges in a manner consistent with the generalization bound established in Theorem 3.2
Next, Figure 3 presents the grokking dynamics as well as those of RNCI and NC2 scores under different weight decays. which further supports our analysis in two major respects. First: the de- crease of RNCI is synchronized not with fitting the training set but with the emergence of grokking, i.e., the improvement of test accuracy. This phenomenon consistently appears across all weight- decay settings, reinforcing our theoretical result that explains grokking through the progression of neural collapse. Although NC2 is not directly treated in our analysis, it exhibits almost the same behavior as RNCI and converges toward 1, indicating the emergence of neural collapse. Second: stronger weight decay À accelerates the timing of grokking and narrows the gap between training ac- curacy saturation and generalization. This observation aligns with main result linking grokking to RNC1 dynamics and also supports Theorem 4.3) which shows that the convergence of the RNCI score becomes faster the weight decay increases.
8

---

Published as a conference paper at ICLR 2026
Test Accuracy
RNC1 Score
NC2 [condition-number
1.0
0.1.1
11.
30
0.8
0.01
25
0.6
0.001
Aw0.3
15-
0.4
A 0.03
10
A=0.03
0.2
0.008
0.0003
Time stap
Time step
Time stop
Figure 3: Dynamies of test accuracy, RNC1, and NC2 scores throughout training for different weight decay coefficients A. In the test accuracy panel (left), the training accuracy is additionally shown in dashed lines of the same color to visualize grokking behavior. Results are averaged over five dif- ferent seeds with MLP trained on the MNIST dataset. These results demonstrate the connection between neural collapse and grokking, and their time scales, supporting Theorems
and
11z.xi-ia -
nHSICIZ, X) nHSIC;Z;
RNC1 Score
70
A=0.3
1-0.1
0.2
1-011
A= 6.03
0.0.3
A=0.03
20
A- 0.003
1.000
+0.00
0-
101
19
3C²
100
301
201
10a
10a
302
201
Time step
Time stop
Time stop
Figure 4: Dynamics of redundant information in IB framework (estimated via MI and nHSIC) and RNC1 scores throughout training for different weight decay X. Results averaged five differ- ent seeds with an MLP trained on the MNIST dataset. These results show the connection between neural collapse and IB dynamics, as well as their time scales, supporting Theorems 3.4
and
5.2 IB DYNAMICS
In this section, we conduct experiments on IB dynamics. Using the same setup as in the previous experiments, measure how the MI between the learned representation Z and the input X, as well as between Z and the target Y, evolves during training. A fundamental difficulty IB experiments with DNNs is that, beyond toy settings, conventional estimators based binning or kernel density estimation fail to provide accurate estimates in nigh-dimensional input or representation due to the curse dimensionality. This challenge remains an active research area, and in this work. we adopt the recent dimensionaliit--educiii www.:aiittt-rduutto--bbbbeeeee MI estimator of Butakov et al. (2024b). Specifically, we first compress the variables into four dimensions and then estimate each entropy term of MI using the k-NN-based Kozachenko-Leonente method (Kozachenko 1987 Berrett et al. 2019). We denote this estimate by i and provide the details in Appendix C In addition, to support the reliability of our experimental findings and address the inherent difficulty of MI estimation, we use the normalized Schheitt-Schiit independence criterion (nHSIC) (Gretton et al. 2005). proxy widely adopted in crommation-theoretii analysis DNNs. Please see Appendix C for the background.
Figure 4 shows, under the same setting as in Figure the behavior of redundant information in the IB principle discussed in Section 3.2 together with the corresponding RNCI score. As shown in the grokking experiments, the decrease of the RNCI score in the later training stage occurs earlier when the weight decay is stronger (right), which consistent with Theorem 4.3 Theorems 3.4 and [4. establish that this decrease in the RNCI score contributes the reduction of the redundant information, and the figure demonstrates this result. The left figure shows MI estimates, indicating that the stronger weight decay accelerates the decrease of redundant information. Although the

---

Published as a conference paper at ICLR 2026
information reduction is slightly delayed for large weight decay values (A = 0.1), the decrease of RNCI scores actually leads to the reduction of redundant information. Since MI is estimated by decomposing it into differential entropy terms that estimated separately, the resulting MI estimates can take negative values despite the non-negativity of MI, while still capturing the overall decreasing trend. To further support our findings, also include the results using nHSIC to measure superfluous information (middle). This result exhibits the same qualitative behavior as the RNCI score and corroborates our theoretical analysis of the IB dynamics, Additional results for other model architectures and datasets are provided in Appendix D.2
CONCLUSION
In this work, we focus on grokking and IB dynamics as two representative late-phase phenomena of DNNs, whose mechanisms have been elusive. We show that both phenomena be explained in terms of the population within-class variance of the learned representations, and more specifically, by the progression of neural collapse and its associated time scale. These theoretical findings supported by our experiments, Beyond the theoretical perspective, our results also provide practical implications: tracking quantities such the rescaled within-class variance can help determine when continued training will be beneficial, and weight decay can accelerate the transition to this late- phase regime. natural next step to extend the time-scale analysis of neural collapse other architectures beyond MIP or to different initialization methods. Another interesting direction is to analyze the possibility of neural collapse that implicitly arises without weight decay.
7
ACKNOWLEDGGMENTS
This work supported by JSPS KAKENHI Grant Number JP24H00709. We thank the reviewers for their constructive comments and the members of our laboratory for helpful discussions.
REFERENCES
Linara Adilova, Bernhard C Geiger, and Asja Fischer. Information plane analysis for dropout neural networks. In The Eleventh International Conference on Learning Representations, 2023. URL
net wttp://wwwwwweeetteeeeeeeeeee.e. forum?1d=bQQBccoo
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International conference machine learning, pp. 242-252. PMLR, 2019.
Rana Ali Amjad and Bernhard C Geiger. Learning representations for neural network-base based clas- sification using the information bottleneck principle. IEEE transactions on pattern analysis and machine intelligence, 2209):222-22233 2019.
Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Spectrally-normaliieedd margin bounds for neural networks. Advances in neural information processing systems, 30, 2017.
Alon Beck, Noam Itzhak Levi, and Yohai Bar-Sinai. Grokking at the edge linear separabil- ity. In Forty-second International Conference on Machine Learning, 2025. URL https:
Thomas B Berrett, Richard J Samworth, and Ming Yuan. Efficient multivariate entropy estimation
via k-nearest neighbour distances. 2019.
Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli. Memorization and optimiza- tion in deep neural networks with minimum over pree-paramttrizztionn Advances in Neural Informa- fion Processing Systems,
2022.
Ivan Butakov, Aleksandr Tolmachev, Sofia Malanchuk, Anna Neopryatnaya, and Alexey Frolov. Mutual information estimation via normalizing flows. Advances in Neural Information Processing Systems, 37:3027-3057 2024a.
10

---

Published as a conference paper at ICLR 2026
Ivan Butakov, Alexander Tolmachev, Sofia Malanchuk, Anna Neopryatnaya, Alexey Frolov, and Kirill Andreev. Information bottleneck analysis of deep neural networks via lossy compression. In The Twelfth International Conference Learning Representations, 2024b. URL https:
//pperevieww net. net. formm? d=huGECz8dPp
Thomas M. Cover and Joy A. Thomas, Elements of Information Theory (Wiley Series in Telecom- munications and Signal Processing). Wiley-Interscience, USA, 2006. ISBN 0471241954.
Hien Dang, Tho Tran Huu, Stanley Osher, Nhat Ho. Tan Minh Nguyen, et al. Neural collapse in deep linear networks: From balanced to imbalanced data. In International Conference on Machine Learning. pp. 6873-6947 PMLR, 2023.
Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. arXiv preprint 2rXiv:2303.06/73, 2023.
Branton DeMoss, Silvia Sapora, Jakob Foerster, Nick Hawes, and Ingmar Posner. The complexity dynamics of grokking. Physica D: Nonlinear Phenomena, pp. 134859, 2025.
Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguisties: human language technologies, volume (long and short papers), pp. 4171-4186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675- 1685. PMLR. 2019.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer- peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118(43)xe1000911 2021.
Spencer Frei, Niladri S Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pp. 2668-2703, PMLR, 2022.
Kenji Fukumizu, Arthur Gretton, Xiaohai Sun, and Bernhard Schölkopf. Kernel measures of condi- tional dependence. Advances neural information processing syslems, 20, 2007.
Tomer Galanti, András György, and Marcus Hutter. On the role of neural collapse in trans- fer learning. In International Conference on Learning Representations, 2022. URL https: //openzeview nwt. nett/oorrmmddd-wIpd...
Bernhard C Geiger. On information plane analyses of neural network classifiers review. IEEE
Transactions on Neural Networks and Learning Systems, 23(11))7039-7051 2021.
Alison L Gibbs and Francis Edward Su. On choosing and bounding probability metrics. Interna- lional statistical review, 70(3):419-433 2002.
Ziv Goldfeld. Ewout Van Den Berg, Kristjan Greenewald, Igor Meinyk, Nam Nguyen, Brian Kings- bury, and Yury Polyanskiy. Estimating information flow in deep neural networks. In International Conference on Machine Learning, pp. 2308- PMLR, 2019.
Arthur Gretton, Olivier Bousquet, Alex Smola, and Bernhard Schölkopf. Measuring statistical de- pendence with hilbert-schmidt norms. In International conference on algorithmic learning theory, pp. 63-76. Springer, 2005.
Song Guo, Lei Zhang, Xiawu Zheng, Yan Wang. Yuchao Li, Fei Chao, Chenglin Wu, Shengchuan Zhang, and Rongrong Ji. Automatic network pruning via schmett-schmidt independence criterion lasso under information bottleneck principle. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 17458- 17469, October 2023.

---

Published as a conference paper at ICLR 2026
Ting Han, Linara Adilova, Henning Petzka, Jens Kleesick, and Michael Kamp. Flatness is neces- sary, neural collapse is not: Rethinking generalization via grokking. In The Thirty-ninth Anmual Conference Neural Information Processing Systems, 2025.
petttp:://ppenreeeee
i d= Oat.cttDDg
X.Y. Han, Vardan Papyan, and David L. Donoho. Neural collapse under MSE loss: Proximity to and dynamics on the central path. In International Conference on Learning Representations, 2022.
URL wtt.thttt///ppeeee reviirevieee iew. net. Forum?idaw 1 R3
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Eduard Hovy, Laurie Gerber, Ulf Hermjakob, Yew Lin, and Deepak Ravichandran. To- ward semantics-based answer pinpointing. In Proceedings of the First International Confer- ence on Human Language Technology Research, 2001. URL httph:tt///// /wtp///wwwww/a aww.eb.wwwwbogg..
aanhhoogyy/ HO01 1069
Yufei Huang, Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. Unified view of grokking, double descent and emergent abilities: A comprehensive study on algorithm task. In First Conference on Language Modeling, 2024. URL
www. fant/forummicc
[GG1bbmisss
Ahmed Imtiaz Humayun, Randall Balestriero, and Richard Baraniuk. Deep networks always grok and here is why. In Forty-first International Conference Machine Learning, 2024. URL
Cettps://openreviee net 2d= 90KMr
Arthur Jacot, Peter Súkeník, Zihan Wang, and Marco Mondelli. Wide neural networks trained with weight decay provably exhibit neural collapse. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL 1HCN4p]Tb4]
www.errwwiwweeeeeeeeeeeeeee net/forumiid
Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie J Su. An unconstrained layer-peeled perspective on neural collapse. In International Conference on Learning Representutions, 2022.
net URL httph: http:///ppeeeiiiieee http:/ewweeeetttteeeeeeeee forum? Wd--21dyy
Tong Jian, Zifeng Wang, Yanzhi Wang, Jennifer Dy, and Stratis Ioannidis. Pruning adversarially robust neural networks without adversarial examples. In 2022 IEEE International Conference on Data Mining (ICDM), pp. 993-998. IEEE, 2022.
Jiachen Jiang, Jinxin Zhou, Peng Wang, Qing Qu, Dustin G. Mixon, Chong You, and Zhihui Zhu. Generalized neural collapse for a large number of classes. In Forty-first International Conference on Machine Learning, 2024. URL httph:::///// Kkttss:/oppereevce. net forum u ?id-
D487kkB89m
Kedar Karhadkar, Michael Murray, and Guido F Montufar. Bounds for the smallest eigenvalue of the ntk for arbitrary spherical data of arbitrary dimension. Advances in Neural Information 2024.
Processing Systems, 37:138197
Kenji Kawaguchi, Zhun Deng, Xu Ji, and Jiaoyang Huang. How does information bottleneck help deep learning? In International Conference on Machine Learning, pp. 16049-16096. PMLR, 2023.
Robert de Mello Koch and Animik Ghosh. A two-phase perspective on deep learning dynamics. arXiv preprint arXiv:250.112700, 2025.
Leonenko Kozachenko. Sample estimate of the entropy of a random vector. Probl. Pered. Inform., 23:9, 1987.
12

---

Published as a conference paper at ICLR 2026
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking the tran- sition from lazy to rich training dynamics. In The Tweifth International Conference on Learning Representations, 2024. URL https::
htttps::/operreviiiie/ /ttp///pppee...eeeeeeeeeeeee net (otoorumiiiitttmmme
Yann LeCun, Léon Bottou, Genevieve Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-50. Springer, 2002.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online). Available: http://yann.eecur /ew.u/..co//xx///miss... 2, 2010.
Noam Itzhak Levi, Alon Beck, and Yohai Bar- Sinai. Grokking in linear estimators - a solvable model that groks without understanding. In The Twelfth International Conference on Learning Representations, 2024. URL
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th Interna- fional Conference on Computational Linguistics, 2002. URL https: / www.acllebb.o. org/ Catthlooyy/C2----
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 1966 2022a.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. Advances in Neu- ral Information Processing Systems, 35:34651-34663 2022b.
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference on Learning Representations, 2023a. URL https: //openzeview.o wwt. wwtp:/eteetttttmmmdddddioioi
Ziming Liu, Ziqian Zhong, and Max Tegmark. Grokking as compression: A nonlinear complexity perspective. arXiv preprint arXiv:2310.059118 2023b.
Stephan Sloth Lorenzen, Christian Igel, and Mads Nielsen. Information bottleneck: Exact analysis of (quantized) neural networks. In International Conference on Learning Representations, 2022.
URL https: htp//ppeevvviieeeeeeeeeeeeeeeeee net wttp://wwwwweeeeeeeeneeeeeeeeee. 2i 9DZCCrUOW
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- net./ tene/ffoumm ? id=
ence on Learning Representations, 2019. URL
Bkg6RiCqY7
Jianfeng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. Applied and Com- putational Harmonic Analysis, 59:224-244 241, 2022.
Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon Shaolei Du, Jason D. Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provably induce grokking. In The Twelfth International Conference Learning Representations, 2024. URLhttps: /tLltttss//openeevveww wpw.revvee.eeeeeeee/e.ee........ net. forum?
1d=XsHqr9dEGH
Wan-Duo Kurt Ma, JP Lewis, and W Bastiaan Kleijn. The hsic bottleneck: Deep learning without back-propagation. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 5085-5092, 2020.
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. Sampling Theory, Signal Processing, and Data Analysis, 20(2):11, 2022.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning- MIT press, 2018.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretabiliity. In The Eleventh International Conference on Learn ing Representations, 2023.
/ /
13

---

Published as a conference paper at ICLR 2026
Quynh N Nguyen and Marco Mondelli. Global convergence of deep networks with wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:1 1961- 11972. 2020.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40): 24652-24663, 2020.
Roman Pogodin and Peter Latham. Kernelized information bottleneck leads to biologically plausible 3-factor hebbian learning in deep networks. Advances in Neural Information Processing Systems, 33:7296-7307, 2020.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.0217. 2022.
Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking first order phase transition in two layer networks, In The Twelfth International Conference on Learning Representations, 2024. URL
net /tt///oooommdd-----T0
Keitaro Sakamoto and Issei Sato. end training induces information bottleneck through layer- role differentiation: A comparative analysis with layer-wise training. Transactions Machine Learming Research, 2024. ISSN 2835-8856. URL httts:://openrer//// net/ www.nnettfoommomeo forum?
1 1d=03wmRh2SFT
Andrew Michael Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Bren- dan Daniel Tracey, and David Daniel Cox. On the information bottleneck theory of deep learning. In International Conference on Learning Representations, 2018. URL https:
//openzeview. www.nnee///iddd//eeeeeo.. WPG-A-
Johannes Schneider. Learning in NN: from fitting most to fitting few. Neural Computing and Applications, 37728):23423-23444 2025. doi: 10. 100100//000521-055111555
Johannes Schneider and Mohit Prabhushankar. Understanding and leveraging the learning phases of neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 14886 18886-14893, 2024.
Milad Sefidgaran, Abdellatif Zaidi, and Piotr Krasnowski. Minimum description length and gen- eralization guarantees for representation learning. Advances in Neural Information Processing Systems, -85--1255 2023.
Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via informa- tion. arXiv preprint arXiv: htt:::7733000//0 2017.
Ravid Shwartz-Ziv. Amichai Painsky, and Naftali Tishby, REPRESENTATION COMPRES-
SION AND GENERALLZZTION IN DEEP NEURAL NETWORKS, 2019. URI. https:
//open www.weevieew... wwt.:/et./oomaa d=
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1631-1642, Seattle, Washington, USA, October 2013, Association for Computa- tional Linguistics. URL httpt:t////// /wwww ac web. org wttp:o//gg/gaoooooo////333/ 1170
and Christopher Potts. Recursive deep models for semantic compositionality
Peter Súkeník, Marco Mondelli, and Christoph H Lampert. Deep neural collapse is provably optimal for the deep unconstrained features model. In Thirt-seventh Conference on Neural Information Processing Systems, 2023. URL /ttt:s://opemreeiiee..ettt
Peter Súkeník, Christoph Lampert, and Marco Mondelli. Neural collapse low-rank bias: Is deep neural collapse really optimal? Advances in Neural Information Processing Systems, 37: 188250-138288, 2024.
14

---

Published as a conference paper at ICLR 2026
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206.048: 2022.
Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collaps. geometry. Advances in Neural Information Processing Sys- tems, 35:27225-2777 27238, 2022.
Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural col- lapse. In International Conference on Machine Learning, pp. 21478-211111 21505. PMLR, 2022.
Tom Tirer, Haoxiang Huang, and Jonathan Niles-Weed. Perturbation analysis of neural collapse. In International Conference on Machine Learning, pp. 34301-34329, PMLR, 2023.
Naftali Tishby and Noga 7aslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pp. 1-5. leee, 2015.
Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/00040057, 2000.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.023390 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Zifeng Wang, Tong Jian, Aria Masoomi, Stratis loannidis, and Jennifer Dy. Revisiting hilbert- schmidt information bottleneck for adversarial robustness. Advances Neural Information Pro- cessing Systems, 34:586-597, 2021.
Zifeng Wang, Zheng Zhan, Yifan Gong, Yucai Shao, Stratis Ioannidis, Yanzhi Wang, and Jennifer Dy. Dualhsic: Hsic-bottleneck and alignment for continual learning. In International Conference on Machine Learning. pp. (6578-36592. PMLR, 2023.
Diyuan Wu and Marco Mondelli. Neural collapse beyond the unconstrained features model: Land- scape, dynamics, and generalization in the mean-field regime. In Forty-second International Conference on Machine Learning, 2025, URL https: /ttpss:///pppeeeeeeeeee net net/forunurttt ? ZrhGq6644
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmark- ing machine learning algorithms, 2017.
Zhiwei Xu, Yutong Wang. Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in reLU networks for XOR cluster data. In The Twelfth International Conference on Learning Representations, 2024. URL https:::... : /
net/ fww.forrmmmooo
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference on Machine Learning, pp. 27179-27202. PMLR, 2022a.
Jinxin Zhou, Chong You, Xiao Li. Kangning Liu. Sheng Liu, Qing Qu, and Zhihui Zhu. Are all losses created equal: A neural collapse perspective. Advances in Neural Information Processing Systems, 35:31697-31710 2022b.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A ge- ometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34:29820-29 -9933444., 2021.
Difan Zou and Quanquan Gu. An improved analysis training over paramaaamttei deep neural networks. Advances in neural information processing systems, 32, 2019.
15

---

Published as a conference paper at ICLR 2026
A FURTHER COMPARISON WITH RELATED WORK
As concurrent work to ours, we compare our results with Han et al. (2025) in this section. Since the title may appear contradict our claims at first glance, we clarify that their findings are fully con- sistent with ours and highlight the novelty of our contribution. HHan et al. (2025) discusses grokking in relation to neural collapse, which makes their setting similar to ours. Although the title suggests that neural collapse might not be relevant for generalization, this is not what the paper actually argues. Rather, they state that neural collapse appears earlier than the acquisition of generaliza- tion. The emergence of neural collapse does not align with that of generalization, and therefore it does not fully explain generalization ability. In contrast, they observe that flatness correlates more closely with the decrease in test loss and thus appears to offer better explanation in their experi- ments. However, this empirical observation is entirely consistent with our analysis. The apparent contradiction arises because Han clustering (NCC) metric, which mixes multiple properties associated with neural collapse. In partic- ular, NCC also captures the separation of class mean representations that naturally occurs through fitting the training set, and therefore it decreases earlier than the test loss. As we emphasize in Re- mark our analysis isolates the contribution of within-class variance or RNCI score, allowing us to separate out the cause of the generalization and correctly focus on the variance decrease that is actually responsible for the emergence of generalization in grokking. Furthermore, the theoretical analysis in Han al. (2025) is limited to showing that flatness is partially guaranteed under neural collapse. Their work addresses neither the connection to generalization nor the associated training dynamics. In contrast, illustrated in Figure [ our paper clarifies the relationships among sev- eral concepts, including information bottleneck, and further characterizes their training dynamics through a neural-collapse-basee analysis. This leads unified understanding of the mechanisms underlying grokking
(2025) measures neural collapse using the neural collapse
B PROOF OF MAIN RESULTS
B. GROKKING RESULTS
We first provide the following lemma, which is useful for bounding the tail probability with respect to the variance of the random variable.
Lemma B.1 (Cantelli's s inequality). Let X be a real-valued random variable with mean [[ and variance o². Then, for any À > O, we have
o² r2 X2
Pr 2 B
Using this lemma, we can prove Theorem 3.2 as follows.
Proof of Theorem 3.2
The test error of the model prediction is bounded follows:
Pr
arg max (f(z): 4y)
(3)
:E[K]
Eo BaaPxiy
Larg (xx
(4)
Pr(c)-Prr3kE (K (c} s.t. (g(0), wa) > (g(x), wo) IY =c)
(5)
K
< Σ (c Σ Pr((g()), - we) >0 0 Y =
(6)
kec
16

---

Published as a conference paper at ICLR 2026
where we used union bound in the last argument. From Lemma B.1 the term in the last line can be further bounded as follows:
Pr ((g(2)) wa we) > 0
(7)
Pr (gg(z) - =x [9(2)) 20 we) > (Ex [g(x)), w.
w. - -ws) I = c)
(8)
(w. (w.-ww) ] E. (w. w))
if (ExIr=e w. -wx- 0,
(w--wx) Ee(w.- we w (Bxx oo w.
(9)
otherwise,
where Ee is given by E. Ex [o(x) Ex1x=x (o(X) (xx)) 7. In
both cases, Equation ( can be rewritten as
max (ExxIx=(o(XX)) we wx). 03 -1
Pr ((9(#), wa - we) > 0 Y c) 1+
(10)
(we - w³) Eg (w. - wx)
Here, applying the Cauchy-Schwarz inequality to the variance term yields
(we (Ww) De (we wp) = Ex (9(X)) Ex we -ex,) mnx)
(11)
Ex (g((X) (ox)]11 | wxl1
(12)
By combining Equations 6 10 and 12
obtain the desired result:
max (ExIr= ----0
1 w
Pr
arg max (f(x)) *E[K]
S Σ
(c)
1+
Exir= [l9(x) ExIr )))
kipt
K
max (ExII=el0( (0(()) 1w we 11
Σ
py (c)
1+
keff
Exjy=
[lx(x) 1X ox (x)
where the last line follows from Definition 3.1 and dividing both the numerator and denominator by B3 supecx |l2
B.2 IB DYNAMICS
B.2.1 PROOF PROPOSITION 3.3
Proof of Proposition 3.3 By definition, obtained by adding noise to g(X), so the Markov chain
Y + X g(X) Z holds. The first inequality I I(2; X) 2 I(2; Y) follows from the Markov
chain + X Z and the data processing inequality (DPI). For the second inequality, we use
variational approach. From a definition of MI, we have
I(Z;Y) = / dy dz ply. x) log plU. 2)
p(y)p(z) =
f
dy dz ply, z) log p(ulz)
(13)
p(y)
Here, we introduce variational approximation q(v(2)) for the conditional distribution p(y 2). From the non-negativity of the KL divergence, we have
I(Z;Y) 2 / dy dz p(v, z) log
(u)lz)
J dy p(y, z) log ((ul=) HYY..
(14)
p(u)
We model the variational approximation as a softmax function with respect to the last layer output W i.c., q(y) x) exp (Wx),) EK exp ((Wz),))) leading
I(Z:Y) 2 ((. (R 2)] HYY).
(15)
Since the entropy of the target Y constant, we conclude the desired inequality.
n
17

---

Published as a conference paper at ICLR 2026
B.2.2 PROOF OF THEOREM 3.4
Before moving on to the proof of Theorem 3.4| we provide the following lemma, which states that the Gaussian distribution maximizes the entropy among all distributions with the same covariance. Lemma B.2 Coverr & Thomas 2006). Theorem 8.6.5). Let the random vector R have zero mean and covariance matrix Σ = ]] Then, we have h(X) S log (2TT)) det(z)) } with equality if and only if X ~ (O0,))
With this lemma, we show the proof of Theorem 3.4
Proof of Theorem 3.4
Using the Markov chain - - Z and the chain rule of MI, we have
I(Z;X) (Z; X, Y) 1(2; Y) + [Z;
(16)
leading to the first equality. Rewriting the conditional MI with the differential entropies, we have
X(I) h(Zl M( hZZ h(Z[X),
(17)
which again follows from the Markov chain. For the second term, we use the differential entropy of Gaussian distribution (Cover Thomas 2006 Theorem 8,4.1), leading to
X( h (g(X) B,E (X) drep (1 log(2TB303))
(18)
The conditional covariance matrix of Z given Y y given by
Σzlr= Cov x|y [a(X)] O Idrep
(19)
Then, from Lemma B.2 the first term in Equation 17 can be computed as follows:
Eyny [h(Z|Y- = y)]
(20)
2 Eywr [log (2TT))umP det (E
(21)
2 drep (1 + log lor 2T B²0²) + Eyny log |CovxIx=y
1
Ideop
(22)
B30²
CoVx [g(x)
2 drep (1 + log
2T
Byg B²) + Eynr
TT
(23)
B200
where the last inequality follows from the fact that log det (A I) TH(A) for any positive semi-
definite matrix A. By putting Equation (18) and Equation (33 into Equation (17) we have
((2:X|T) S 2Bz02 Ey-y [Ir [
(24)
202 E (x,)))
[1|ā(ar) Ex- (or [1
(25)
which concludes the proof.
B.2.3 REMARK ON UPPER-BOUND TIGHTNESS
In Theorem 3.4 the redundant information I I(Z; X Y) is upper-bounded with population within- class variance. this section, we discuss how informative this upper bound is.
From the proof of Theorem 3.4 obtain the following tight bound on I [Z; X YY):
log det (CovxIy=yl((x) + 7 1) σ²
T(Z; X Y) S 2 Eyyy
(26)
This bound is derived from Lemma B.2, and equality holds when Y = follows the Gaussian distribution. Therefore, this upper bound is tight. In the subsequent step of the proof, we used the inequality log det(A + I) < Tr(A). Since the equality does not hold when is positive defi- nite, the resulting bound in Theorem is not tight. However, regarding the within-class variance term appearing in the theorem, we can show that reducing the variance in each coordinate always decreases both the tight bound above and the upper bound in Theorem 3.4
18

---

Published as a conference paper at ICLR 2026
Proposition B.3. Reducing the population within-class variance in any single coordinate, i.e.,
ExIy=v (o(X) Σx =x)x=g(((X))) for any i E [drepl and y E [K], strictly decreases both the upper bound in Theorem 3.4 and the tight upper bound on I (Z; X Y) in Equation (26
Proof of Proposition B.3
For the upper bound in Theorem 3.4
since it can be written as
Σ
By
[(((X) xx (x)))
(27)
XE xE(4rep]
it is obvious that decreasing any summand decreases the upper bound. We now show that the tight
upper bound in Equation (6) also decreases under such a perturbation. Let the amount of the
variance reduction along coordinate i be 5 >> O. For notational simplicity, we denote the original
covariance matrix by K xI [XX] and the perturbed covariance by K' K še,e,
Then, the value of the upper bound is bounded below as
K do,ee
+1)]
Ey
log det
+
(28)
σ²
1 By
8
K'
K' 02 I
2
log
1
02
02
+1)
ei det
(29)
K"
2 By
log det
02 +1)
(30)
where the equality holds from the matrix determinant lemma, and the inequality follows from the fact that the diagonal element of the inverse of a positive definite matrix is always positive. Since the right-hand side is exactly the upper bound in Equation (26) evaluated the perturbed covariance, this completes the proof.
B.3 CONCENTRATION OF WITHIN-CLASS VARIANCE
The analysis in Section is carried out using Rademacher complexity, which is standard tool for establishing uniform convergence and deriving generalization bounds. For real-valued func- tion class F and fixed training set S, the empirical Rademacher complexity 2t, (F) is defined as E, [Sup)EF _ Σm= <=](xr:)) where e; are independently sampled from Unif ((+1)).
Theorem 4.1 is inspired by Galanti al. 2022 Proposition 3). In contrast their result, the following theorem addresses the three novel aspects; (i) a uniform convergence result for within- class variance, (ii) a refined design the failure probabilities and the union bound, and (iii) a precise and tight upper bound on Rademacher complexity.
Proof of Theorem 4.1 We first define the function classes of DNNs follows:
g = [9: Rd Re g(x) (W-100 (.- We
(Wiz))),
1. [L
1]}
D-1
1-1
Welle.1
2/3
3/2
G2l = geg l1(g) = II Well 2", A(g) =
Σ
< 2
Wellz
t=1
for s,t E N. For each of these classes, define corresponding rescaled classes as follows:
sup (lx(x))122 9 E {). G.tt { (((())) supnex (gg(2))12 : E 9x.t
g(x)
g(x)
Then, have g U..+ Gz.t and 9 = U &
19

---

Published as a conference paper at ICLR 2026
We first fix c E [K] and s,t E N. For any Go,t, have
Ex [lx(X) [xIY=o10(X)]IB
=Ex
O(X) 1
Σ
8(x,))
Σ g(x,)
E. Yx [0(
(31)
ne jESe
The jESc
S Ex|Y=e
x(X)
Σ The
1
(32)
jESc
In the following, we will analyze the gap between Equation 332 and its empirical counterpart. For Ejes, 3(x;)) Since the output of 9 E Ge.t is rescaled, have ll0(x)]l2 for all X, and the output of is bounded with 4. By normalizing the output with this value to apply MMhrr et al. 2018 3.3), with probability at least ōx we have
any fixed I E G..t. we define function RRdr - as h(z) =
Theorem
[((x)
Ex|y=
T(X)
1
Σ jES.
g(æ;)
g(x:))
The Σ g(æ,)) JeS.
The
The iE
12
log(2/5ac))
< 22 (h Gos) 4. 33
(33)
2n.
By choosing the failure probabilities as de,t 8/ (Kst(s +((t+1) and applying the union bound, since have Ex 8,,1 8/K, the above inequality holds with probability at least 1 8 for all s N and g E G...-
Next, we analyze the Rademacher complexity term in Equation (33) using a covering number argu- ment. For set U, we define its covering number NUU, [II)) as the minimal cardinality of a subset V S U such that, for every u E U, there exists v E V satisfying ||u w|l I e. Here, for the real-valued function class F, define its restriction to the training data points
Fis. = { x (f(2))....
f(x..)) E R" [fEF]
I
(34)
As discussed earlier, both the domain of the function and the average Ejes, I g(x,)/n, are re- stricted to the unit ball. Thus, the Lipschitz constant of h given by hLiP 4. By Bartlett et al. 2017 Theorem 3.3) and the definition of Go.e. we have
.,, l12) < V llog ((ho ss.
Bz ne log (2dmax))
2
(35)
B
Using Bartlett et al. and applying Equation (35)
2017
Lemma A.5), modified so that the range is extended from [0, 1] to [0, 4],
we have
4a
12
AVTT
V |188N ((ho
Rr (h o (2.) < inf a>0
,, l12 de
(36)
n
nr.
Ja
4a
(4vTic
Br 20+6+2
< inf
+
log
Vnc (ogg(2maxx)
(37)
a>0
//c
Bg
16
12: 22+1+2 B
log (nc) Vlog(2dmax).
(38)
ne
By
20

---

Published as a conference paper at ICLR 2026
where the last line follows by taking
4/ n/c Substituting this into Equation 33
we have
27
Ex
[(x)
1
Σ ((0,))
g(x;)
1
Σ g(x,)
nhe JESc
ne
iESe
JES.
32 ne
12. 20+e+3 B₁
log(2K sl(s + 1)(1 + 1)/0)
log (ne) Vlog( /log(2mmax 12
(39)
Bn,B.
2nc
For any g E G. let s = [log2 (211(g)) and & = [log, (2A(g))J, so that 9 E Gs with 2 I
l1() $ 2 and 2 & A(g) I 2'. From Equation (39). it follows that with probability at least
1 6, we have
E Ex|y=c
g(X)
1
Σ f(tz)
1
g(x,))
Σ g(zz)
The jES.
The E.
The jES.
log log(ne) (()) Vlog( (//5)//)//) + log log (TT) )
0
VTc
Bg
which completes the proof.
B.4
NEURAL COLLAPSE DYNAMICS
B.4.1 PRELIMINARIES
We first provide the formal assumptions used in Theorem 4.3 adopted from Jacot et al. (2025). A detailed discussion of their validity and typicality is given in Remark B.6
Assumption B.3 (Pyramidal network). Let d1 2 N and de 2 de+1 for all f (
Assumption B.4 (Smooth activation). Let y E (0, 1) and 2 Suppose a satisfies: d'(x) (7, 1] for all E IR; (ii) [a(x)]] S [ar| for every E IR; and (iii) o' is B-Lipschitz continuous.
Assumption B.5 (Initialization). Let X = min (We(O) and X = Jmin (a (W)) where muin denotes the smallest singular value of matrix. Suppose that have
(2) L
AF II do min (Ar, minee de) 2 881 E=3
Eo (0(0))
In the following, for notational convenience, we denote by Z, E RdexN the output of the f-th layer for the entire training set. Specifically, Zo = Xo. Ze (W,Z--1 for l [L 2], and Ze = WiZe- for & E {L 1, L}. Additionally, we denote λ = [O)lle + (3.... A for [L]. We denote by (.) the i-th largest singular value given matrix.
Remark B.6 (Validity of the Assumptions). The derivation of neural collapse from gradient de- scent, rather than from the unconstrained feature model (UFM), was first established in Jacot et al.) (2025), and Theorem 4.3 adopts the same assumptions. To analyze the convergence of the DNN training loss, one typically imposes conditions ensuring a positive lower bound on the Jacobian with respect to the parameters. There multiple well-established ways guarantee this, including width requirements across all layers (Du et al.] 2019 Allen-Zhu et al., 2019 Zou & Gu 2019 and other 2020 Bombbri et al.] 2022; Karhadkar et al.) (024)). Our Assumptions B. to B.5 can be replaced
conditions with mild width assumptions INguyen& & Mondelli
by any of these alternatives. For example,|Alil Allen-Zhu -hu et al. (2019): Zou & Gui 2019) analyze ReLU networks instead of smooth activations but require all layers to be sufficiently wide. In contrast, the pyramidal topology assumption (Assumption B.3) removes the need for width assumptions on every layer and replaces them with a more realistic setting: a wide first layer followed by narrow- ing architecture. The smooth-activation assumption (Assumption 5) is widely used convergence analysis and appears in independent work such as Nguyen & Mondelli (2000) Bombari et al. (2022); Liu et al. (2022a); Frei et al. (2022). Smooth leaky ReLU satisfies this assumption and can approxi- mate ReL.U arbitrarily well for suitable choices of y and B. Assumption B.5 concerns initialization,
21

---

Published as a conference paper at ICLR 2026
and it can be satisfied by choosing the second-layer scale sufficiently small. Moreover, by relaxing Assumption B.3 from dj 2 N to d1 = (((N). Assumption can be shown to hold under standard He/LeCun initialization (LeCun al. 2002 He et al., 2013). which follows directly from Appendix of Nguyen & Mondelli (020). Finally, the use of squared loss is standard this line of work:
it is used both in the gradient-descent NTK-st analyses (Du al. 2019 Allen-Zhu et al., 2019]
Zou & Gu, 2019, Nguyen & Mondelli 2020 Jacot et al., 2025) and in theoretical neural collapse
literature (Han al. 2022; Zhou al. 2022a Sūkeník et al. 2023 2024). As noted in Jaacot et al.
(025), once the training set is interpolated, the second part of the analysis (the time step in The- orem 4.3) showing the emergence of neural collapse does not rely on these assumptions. We note
that none of our results except Theorem 4.3 depend on Assumptions B.3 to B.5
B.4.2 PROOF OF THEOREM
We now provide a proof of Theorem 4.3 The argument follows that of Jacot et al. (2025 Theorem B.2). Specifically, for the convergence of the training loss, we employ the standard Dojsiw-Eoossiewicz (PL) condition to establish linear convergence. On the other hand, the pro- gression of NC1 and RNCI is explained through the development of weight balancedness, i.e., WI Wr = W The key idea is that once the network output interpolates the target suffi- ciently well, which results small training loss, weight balancedness ensures that the degree of interpolation is inherited by the output of the preceding layer, i.e., the representation space.
Since our primary focus on the convergence of RNCI metric, we begin by establishing the necessary result.
Proposition B.7 (Theorem B.1 of Jacot et al. (2025)). If the network satisfies (i) approximate interpolation, i.e., [I Y|r S E1, (ii) approximate balancedness, ie., [[ Wt
WL-1 W,, lll $ Ez. and (iii) bounded representations and weights, i.e. 1-z-lllz S r, L S r, and [eell S for l E [L., then if S min (xx(Y),
V
[K (K-DN 4K
we have
2
E
Tr (r(Ew) S
+ v Vd.
(Y) - E1
Corollary B.8. Under the conditions of Proposition
i S mince[K] v we have B.7
x
E1 mince[K] vhe
RNC1 <
+ VdL-16>
N
E1 VN
Proof of Corollary B.8 By the definition of RNCI, it suffices to bound Tr(Ew) from above and By from below. Substituting Ø (Y) = mince[K] M which follows from YYT
diag (n1,...,, nk.) into Proposition B.7 we have
E1
Tr Tr(Ew)
+
/ dr 12
(40)
N
min vne
For the lower bound on By. we show the existence of training example whose norm can be bounded
from below. From the condition (xx - | p I {1, there exists E [N] such that [)f();;)
allla & CI/VN, which implies |3: - & N < []f(2))]22. Since yı is one-hot vector and by WWLll S r, we have
1
S 1gg(x;))12 S Bg.
(41)
VNV
Combining Equations (40) and (41) yields the desired upper bound. Finally, regarding the condition on ε we used that o (Y) min.e]]] Vn S /N/K and K > 2.
Proposition B.9. Suppose that the network I satisfies Assumptions B.3 to B.5 and that the input
domain X is bounded, i.e., ||0|||2 & for all X. Fix 0 < 61 < 2/2 [[] /ne and eg > O.
22

---

Published as a conference paper at ICLR 2026
Let weight decay parameter X and step size n satisfy the following:
(l)()()))
ei
X S min
(L- 2L(1-3) 1- AF
11 A
100I
18 (8 + )F/2)
1
62
71 S min
231" 5N BB max { (1, (4mx) 34/21 L 2x 4 (Amx)
|
where
31 = 5N3B" ( max 11 {13) L 5/ and m2 = (1+ /41/a)
(1)0 + ro) with
1 min {Ar, minne(s.....)) de} and a 2 3-11-3) dr I1 de Then, there exist time
ro steps
2
log (0(0)) A
Ama log Ex10(0))- Am Amy
log (1mA
TIS
T2 T2 +
log 77
log (1 - 17
log (1 - 11,)
such that for any time step T 2 TH we have
(x((()))
for any time step 72, we have
p
V261
RNC1(7)
+ v dr_1162
2 N (1- \E1)
mince|K] v /2E1
where
== max 22/mxx (2 L-2 | (2 L-1 1|x
Proof of Proposition B.9 The proof follows the same argument as in the proof of Jacot et al. (2025 Theorem B.2). Under Assumptions B.3 to B.5 the PL property of the unregularized loss L (0) is established. While the presence of a weight-decay term slightly shifts the convergence point, the PL condition still holds for the regularized loss (x),). From the PL condition (x) around the initialization, there exists time step
Ama log (x(( Ama
T S
(42)
log (1 n
such that (T1) S 2Amx < ef. Here, note that if the goal is only to ensure that (x(T) < é the choice of oi in the statement suffices.
the regularized loss smooth, meaning that (Z) is 32- 32-Lipschitz continuous and if the learning rate is chosen larger than 1/22. then the loss remains non-increasing for all тΙ To evaluate 32, it is necessary to bound the parameter norm from above. To obtain the norm bound that is independent of X, we 22m, instead of of in the original proof and evaluate
X (0( < E (0(T)) 2Amx,
(43)
which gives |10(T)]l2 < 22/mx- Therefore, Jacot et al. 2025 Lemma C.I) provides ,32 5N3B3 max {1, (4m,))2/22 1/22. Under the learning rate specified in the assumption, the in- equality (x()()) < e? holds for all time steps beyond Tı (or 7().
The remainder of the proof, including weight balancedness, proceeds as in the original argument by combining the above weight norm bound with the interpolation bound [| (X) - Y < /2e, for all 2 Ti- Specifically, the weight balancedness Wz(F) Wz(r) Wt-1 (7) (7) is shown
23

---

Published as a conference paper at ICLR 2026
to converge under sufficiently small learning rate, with the argument based on the analysis of its one-step update. In summary, there exists a time step 72 such that
F2 log 8im
T2 S T1
+
(44)
log (1 - 177
and for any time step T 2 72, we have [T)(T) Wx(r)-
- (T) (T..(()) S c2. Applying Corollary B.8 with { = v 12, & = <2, and the upper bound yields the desired result.
B.9 specifically, by replacing ci with and evaluating the order of ri and in the proposition. By the condition of X S in Proposition B.9 and the definition of my, we can bound m, as constant order depending only on the initialization. Consequenly, the two terms in the upper bound of RNCI by scales independent À and μú; thus, we have RNC1(T) = O ((vEi + V=2)²) = O (Ei + 62). Finally, we evaluate the order of the upper bounds on ri and T2. Since we have log(1 x) for small v, the desired order expression of time steps is obtained.
Proof of Theorem 4.3 The conclusion follows from Proposition
each multiplied
C EXPERIMENTAL DETAILS
The experiments on grokking in the main text follow the classification setup introduced in Liu al. (2023a), where an MLP trained on the MNIST dataset. The MLP has hidden dimensions 200, 200. 200. 10) with ReLU activations, and does not include normalization or dropout lay- ers. The initialization scale is enlarged by a factor of eight across the entire network, and training is performed with the AdamW optimizer at a learning rate 1e-3. For Figure we set the weight decay to 0.01, while Figure 3 shows the results across multiple values of weight decay. The training set size is 1000, the batch size is 100, and the model is trained for 300,000 iterations. Compared to Liu eet al. (2023a), we increased the number of layers in the MLP from three to four to examine behaviors deeper architectures. As consequence, we observe slight instability in training accuracy just before fitting the training set in Figure 3 Nevertheless, the grokking behavior still clearly occurs.
For our experiments the IB dynamics, theoretic quantities: (i) MI estimation based on autoencoders and (ii) information estimation based on HSIC. As we describe below, each of these methods has been widely used in the literature, but they exhibit different characteristics. By employing multiple estimation methods, we aim to enhance the plausibility of the experimental results presented in our work.
adopted estimation methods for information-
MI Estimation via Autoencoder. MI estimation via kernel density estimation (KDE) is widely used in IB studies, but it suffers from poor sample complexity in high-dimensional settings. To address this issue, we adopted the compression- approach proposed by BButakov et al. (2024b), which performs MI estimation in a lower-dimensional space by applying dimensionality reduction via autoencoders.
Following Butakov et al. 2024b), we first train autoencoders on the input X and the representation Z and estimate differential entropies in now-dimensional latent spaces with dimensions dx and dz = 4, respectively. The estimation is performed using the Kozachenko-Leeonnlle estimator Kozachenko, 1987), which is based the density estimation via k-nearest neighbors (k-NN). In practice, we follow Butakov et al. (2024b) and use weighted variant of the estimator, as developed by Berrett et al. (2019). Although training a new autoeneoder for each latent representation can be computationally demanding, Butakov et al. (2024b) showed that using linear autoencoder, that is, principal component analysis (PCA), for compressing Z yields competitive results. Therefore, we use PCA for compressing Z in our experiments. For compressing X, we apply toy CNN autoencoder in which both the encoder and decoder consist of five layers each.
HSIC Estimation. Even approaches based autoencoders require assumptions such invert- ibility for accurate estimation and demand additional computational resources. To further support the findings of our study, also conducted experiments based normalized HSIC. et al. (007 Theorem 4) shows that a normalized version of HSIC (to be defined below) coincides with the chi-square divergence between the joint distribution Px,y y and the product of marginals Px Py.
24

---

Published as a conference paper at ICLR 2026
Given that the MI is defined using the KL divergence instead of chi-square divergence, this suggests that the HSIC-based quantity can be interpreted as a variant of MI. More specifically, Gibbs & Su
(2002) Theorem 5) shows that the chi- -quare divergence upper-bounds the KL divergence. Using
HSIC as an information-theoretii quantity is a common practice in the literature on IB and learning algorithms (Ma et al.] [2020] Pogodin & Latham
2020 Wang et al. 202T Tian et al. 2022 Guo
ct al. 2023 Wang et al. 2023 SSkamoto & Sato, 2024).
We now describe the estimation procedure along with the necessary preliminaries. The HSIC is measure of statistical dependence between two random variables that can capture non-linear rela- tionships. Suppose that we have two random variables X and Y on probability spaces (XY, Bx Px) and (y, Py) respectively, where Bx and By are the Borel o-algebras on A' and V. We consider functions that map elements each sample space to real values. Let F and be reproducing kernel Hilbert spaces (RKHS) with corresponding kernels Ky : X x X -> R and Kig y x y + IR. The mean y4x defined as an element of F such that (Mx, f)= Ex [X)] for all f E F. Similarly, let juy and yxx denote the mean elements of G and F G, respectively. The cross-covariance oper- ator CXY G + F is defined as linear operator such that Cxy == Pxx - pi Please note that an element of F F G, but ing (f.Cxy و = (fg. pXY μu Py) Fog for any f E F and g E G. The Hilbert Schmidt norm of the linear operator C : G - F is defined as ||C]|
regard it as linear operator from G to F by defin-
Ex, (dee Ciby) where {de} and {15) are orthonormal bases of F and G. respectively. We define the HSIC as HSIC Y) = |Cxy 11 which is calculated using kernel functions follows:
HSIC(X, Y) = EX,Y,' (X, X' (",gg
2 [Ex] [xx(X, X XJ Ey (Y, [Y] + Ex,x |N X')] Ey (Y)(
(45)
where (X', (X) independent copy of (X, Y). Given a dataset {(T; Vi ) following Px,r, we can estimate the HSIC as Tr(Kx HNV - 1)) Kx = *x(2,, x,), KY..j = (g M;), and the centering matrix H = IN - 1/N. In our experiments, motivated by the connection to the chi-square divergence described above, we use normalized version of HSIC (nHSIC). Specifically, the nHSIC is defined as the Hilbert- Schmidt norm of the normalized cross tross-covarianc operator, which is given by nHSIC(X, = | Cxx-Cxxypp Maxxyyyy/ lirss 1/1/2137s As in previous studies employing nHSIC for DNN analysis, we define the estimator as Tr (Kx H (NIN) Ky (Ky H + cNIN) "1 and set c 1c-5.
where we denote Kx, H R NxN
D ADDITIONAL EXPERIMENTS
D.1 GROKKING
Decision Boundaries for Additional Class Pairs. Figure showed how representations evolve during training on MNIST, but due to space limitations in the main text, we presented only the decision boundary between class 0 and 1. To demonstrate that the same phenomenon occurs for other class pairs well, 5 results for two additional pairs: class 4 vs. 5 and class vs. 9, For these pairs, we observe the same trend: during the overfitting phase, training examples are already almost separable but still exhibit large within-class variance. As training proceeds, the training examples become more separated and more tightly clustered. Consequently, test examples are better separated and more concentrated in representation space, leading to improved generalization.
Fashion-MNIST. We conduct experiments on Fashion-MNIST (Xiao et al. 2017) as additional experiments different dataset. Following the main text, we use a four-layer MLP with the same training configurations. Tigure[6]show the results, indicating that grokking also occurs on Fashion- MNIST. What is important here not merely the occurrence of grokking, but rather that the decrease in the RNC1 score coincides with the increase in test accuracy, and that stronger weight decay accelerates this timing. These observations reinforce our theoretical results. As supplementary information, we also provide results when increasing the training set size from 1000 to 3000. In this case, the overall trend remains unchanged, but the test accuracy increases slightly earlier.
25

---

Published as a conference paper at ICLR 2026
Train Examoles
Unseen Examples
0255 deas
dids
Denuty
Dencity
-2.0
-1.5
-1.0
0.0
0.5
1.0
1.5
2.0
Signad distance decision coundary (class
Signed decision boundary
(a.1) Overfitting phase (time step 71
16,000): Class vs.
Examoles
Unseen Examples
dbss diss
dots
dass
Denety
Dencit
-2.0
-1.5
-1.0
-0.5
0.0
0.5
1.0
1.5
2.0
-20
-15
-1.0
-0.5
0.0
0.5
14
1.5
Signad distance to decision toundary (clazz VG. class
Sigred oistance decision boundary (dass class 5)
(a.2) Convergence phase (time step 72 100,000): Class vs. 5.
Train Examoles
Unseen Examples
Denety
Denaty
0.5
1.0
Signed distance to dectsion boundary Iclass vs. class
Sigred discance to decision boundary (dass v5. class 91
(b.1) Overfitting phase (time step 16,000). Class vs. 9.
Irain Examoles
Unseen Examples
Dersty
clecs
Dencby
-1.0
-0.5
0.0
0.5
155
2.0
Signad
boundary Iclass
Signed
boundary
(b.2) Convergence phase (time step 72 100,000): Class vs.
Figure 5: Margins of individual examples at two time steps during grokking. The model is 4-layer MLP on the MNIST dataset. Here, we provide the corresponding plots for several class pairs other than the 0-1 pair shown in Figure
Modular Arithmetic. Previous studies on grokking have primarily focused on modular arithmetic tasks. For example, the addition task takes two integers and b as input and outputs their sum modulo a prime number p. i.e., (a + b) mod p. However, this setup is closer to a regression problem than classification, the outputs have ordinal structure rather than being categorical labels. As our analysis is on classification, we consider this problem outside the scope of our study.
CNN. In the main text, we examined grokking classification setting using MLP models, fol- lowing Liu et al. (2023a). This choice is consistent with the theoretical analysis setting in Section where neural collapse grokking occurs in other architectures, we conducted additional experiments with CNNs. We trained CNN consisting of two convolutional layers with max-poolle mox-pooling. followed by two fully connected layers. As in the MLP experiments, we attempted to scale the initialization of all layers by factor of eight, but training was unstable. We therefore considered the modification scaling only the fully connected layers.
studied with an MLP feature extractor. To further investigate whether
Figure shows the results of changing the weight decay parameter X. For all X, test accuracy im- proves at the same time training accuracy, and no grokking behavior is observed. Nevertheless, the results are consistent with our analysis: test accuracy improves in parallel with the progression of neural collapse. As for why there is time lag between fitting the training and the reduction of within-class variance, a possible explanation is that CNNs, due to their inductive bias of local invariance, already extract useful features during the fitting phase. Theorem 4.3 suggests that the
26

---

Published as a conference paper at ICLR 2026
Test Accuracy
RNC1 Score
NC2 Score
0.40
A=01
000
68-
1=0.00
0.6
120
30
0.0
20
5.10
0.2
640
A=0.01
10
A- 0.00
A 0.0C1
2- 0,001
Time stcp
Time Step
Time step
(a) 1000 training samples.
Test Accuracy
RNC1 Score
NC2 Score
1=0.1
1-001
0.
0.6
0.5
220
0.2
0.4
1-000
2-0.11
2- 0.003
1=0000
107 337
357
10*
307
301
101
331
301
100
331
10
Time stop
Time step
Time step
(b) 3000 training samples.
Figure 6: MLP trained on the Fashion-MNIST dataset with different weight decay coefficients X. Test accuracy. RNC1, and NC2 scores are reported. In the test accuracy panel (left), the training accuracy is additionally shown dashed lines of the same color to visualize grokking behavior. Results are averaged over five different seeds.
Accuracy
1.
2.1.1
2- 0.00
0.001
0300
0.172
0.6
0.150
0.2
0.72
8.00
0.0000
0.001
322
Time step
stop
stop
Figure 7: CNN trained the MNIST dataset with different weight decay A. Test accuracy, RNCI, and NC2 scores are reported. In the test accuracy panel (left), the training accuracy is additionally shown in dashed lines of the same color. Results averaged over five different seeds.
fit the output f(X) to the labels Y is propagated to the preceding layer, i.c., the representation space, through the progression of weight balancedness. In contrast, if the representations are al- ready well concentrated from the fitting phase, then no such delay arises after f(X) fits the labels. This interpretation is supported by Figure & which shows the CNN results for the experiment in Figure 2 of the main text. Compared with the grokking scenario with MLPs in Figure 2 the CNN representations already well collapsed at the point when the training set is first fitted (Figure Sa Consequently, the margin distribution for test examples is relatively concentrated from the fitting phase, and training progresses, its center becomes increasingly separated. Finally, Figure also shows that when À= 0.1, the RNC1 score continues to decrease after the training accuracy has reached 1.0, and the test accuracy further improves. Thus, although this case is not grokking be- havior, it demonstrates that continuing training after fitting the training set can further reduce the within-class variance and thereby improve test accuracy.
27

---

Published as a conference paper at ICLR 2026
Trein Examples
Unseen Examples
Dersity
-1.5
-0.5
0.
1.5
2.0
-1.0
Signed distance decision boundary (class class
signed
boundary (class
(a) Fitting phase (time step 71 = 2,700). Train accuracy = 100%, test accuracy = 86%.
Trein Examples
Unseen Examples
Dersity
Dersity
-0.5
05
15
2.0
-20
Signed decision boundary (class va. dlass
signed distance boundary (class 1)
(b) Convergence phase (time step 72 = 100,000). Train accuracy = 100%, test accuracy = 94%.
Figure 8: Margins two time steps for CNN trained the MNIST dataset.
Test Accuracy
FAC1 Scorz
NC1 Score
VC2 Score
1-22
2-3-3
13
2-32
2=02
3-023
3-233
35003
68
30003
0.0
5053
010
Time 200
me
Time
Time
(a) Different weight decay A.
Accuracy
NC1
VC2
2020
3000
5.00 5383
5h0- 5920
50t- 5000
6300
5.00 5950
Slbc 13020
600= 30286
-22- 10000
0.04
1300
(b) Different training sample sizes.
Figure 9: One-layer ViT trained on the MNIST dataset. Test accuracy, RNCI, NC1, and NC2 scores reported. In the test accuracy panel (left), the training accuracy is additionally shown in dashed lines of the same color. Results are averaged over five different seeds.
Transformer MNIST. We also conducted experiments with transformer architecture Vaswani al. 2017) for MNIST classification. Specifically, we use a one-layer vision transformer( ViT) (Dosovitskiy et al. 2020) with hidden dimension 128, four heads, and feedforward dimension 256, without dropout. The input images are divided into patches of size and embedded with convolu- tional layer, followed by learnable positional encoding and a class token. At the position of the class token, a linear head attached for classification. Following prior work of grokking, we adopt an initialization scale of eight, but for training stability, the scaling is applied to the feedforward layers and the linear head.
Figure 9 shows the results when changing weight decay and training sample sizes, where no grokking behavior is observed. In this case, the RNCI score and test accuracy show different trends, whereas the NC1 score decreases in accordance with the improvement in test accuracy. This flects that, to account for test accuracy improvements, not only the reduction in the RNCI score but also class mean separation must be considered. Theorem 3.2 captures both of these elements in
28

---

Published as a conference paper at ICLR 2026
Test Accuracy
Score
NC2 Score
Teal Tann
350
0.020
0.3
250
0.015
0.5
0.7
0.8
0.003
0.5
6033
Time stop
Tiree stop
Time stop
(a) Standard initialization.
Tes: Accuracy
RNC1 Scare
NC2 Score
300
0.5
0.01
0.8
230
100
0.0
0.005
0.5
30
10'
32
107
10
127
30
107
10a
101
327
50'
107
35*
Time steg
Tir step
Time step
(b) Scaled initialization (x8).
Figure 10: One-layer transformer encoder trained on SST-2. Results are averaged over five different seeds, and shaded areas correspond to one standard deviation.
the generalization bound. As shown in the grokking experiments in the main text, when the class means separated, reducing the RNC1 score leads improved test accuracy. In contrast, when the class mean separation is insufficient in the early stages of training. small RNCI score alone does not guarantee good classification performance. This aspect is reflected in the NC1 score; as discussed in Remark 4.2 unlike RNC1, the NC1 score incorporates the information on class-mean separation through its ratio with between-class variance. explains why its decrease coincides with improvements in test accuracy in Figure
Transformer + Text Datasets. To validate our findings across diverse datasets, we conducted experiments on three text classification benchmarks: SST-2, a binary sentiment classification task (Socher et al. 2213)): TREC-6, question classification task with six classes (Hovy et al. 22001] Li & Roth 2002);; and AG-news, topic classification task with four news categories (Zhang al. 2015).
For preprocessing, we use the Hugging Face Bert WordPiece tokenizer (Devlin et al.) (2199) solely for tokenization, restricting the vocabulary to tokens present in the training set. All embeddings, including unknown and padding tokens, are randomly initialized, and the maximum sequence length set to 128. The model configuration is the same as in the MNIST ViT experiment. Training is performed with weight decay of 1e-2 and training set size of 3000. We show results for both standard initialization and a variant where the initialization scale of the feedforward and linear layers is scaled up by a factor of building the prior grokking studies.
The results are shown in Figure 10 (SST-2), Figure II (TREC-6), and Figure 12 (AG-news). In all cases, we observe little difference between different initialization scales, and grokking does occur. As a consistent trend, the figures show that as the model fits the training set, the RNCI score first peaks and then decreases. Notably, for SST-2 and AG-news, this decrease in RNCI score coincides with a gradual improvement in test accuracy. While this behavior is not as abrupt as grokking, supports our theoretical result that continuing training beyond the training accuracy plateau, driving further neural collapse, can benefit generalization.
D.2 IB DYNAMICS
29

---

Published as a conference paper at ICLR 2026
Accuracy
Score
NC2 Score
235
Taan
202
0.5
0.04
150
0.0
0.01
0.00
100
50
0.2
0.00
Time step
Tiree stop
Time stop
(a) Standard initialization.
Tes: Accuracy
RNC1 Scare
NC2 Score
0.5
200
0.0
0.4
200
100
0.2
0.00
307
10'
324
307
10
127
30
107
10
107
327
50'
107
35*
Time steg
Tir step
Time step
(b) Scaled initialization (x8).
Figure 11: One-layer transformer encoder trained on TREC-6. Results are averaged over five differ- ent seeds, and shaded areas correspond to one standard deviation.
Test Accuracy
Scare
NC2 Score
0.31
350
0.9
300
0.0
0.25
200
0015
0.5
200
0.4
0.03
307
107
227
105
2.7
107
107
207
307
325
Time
Tire
Time
(a) Standard initialization.
Test Accuracy
RNC1 Score
NC2 Score
1.0
250
0.9
0.000
250
0.8
0.001
230
0.5
0.02
step
step
Time
(b) Scaled initialization (x8).
One-layer transformer encoder trained on AG-news. Results are averaged over five different seeds, and shaded areas correspond to one standard deviation.
Figure 12:
Fashion-MNIST. As an additional experiment on a different dataset, we conducted experiments on Fashion-MNIST. The experimental setup is the same that of Figure 6al containing four- layer MLP with scaled initialization. Figure B shows the results. As noted before, increasing the weight decay accelerates the decrease in the RNCI score, and the same behavior is observed for
30

---

Published as a conference paper at ICLR 2026
izzxx) -iz,
0.00
0.15
-5
0.10
A-0.0
2. 0.21
0.02
000000
7 0.001
A200.001
Time stop
Time step
Time stop
Figure 13: MLP trained the Fashion-MNIST dataset with different weight decay coefficients A. Dynamics of the redundant information (estimated via MI and nHSIC) and RNCI scores are reported. Results averaged over five different seeds.
Accuracy
1iz -k.
Tran
0.5
2.3
187
307
357
me stop
stop
Figure 14: ResNer18 trained on the CIFARIO dataset. Dynamics of test accuracy, redundant infor- mation (estimated via MI and nHSIC), and RNCI scores are reported. In the test accuracy panel (left), the training accuracy is shown in dashed lines. Results are averaged over five different seeds, and shaded areas correspond one standard deviation.
redundant information measured by MI and nHSIC. Similar to the MNIST experiments in Figure 4 the redundant information estimated via MI decreases slightly later than that estimated via nHSIC. In either case, the results indicate that the decrease in the RNCI score leads to the reduction of these information measures.
ResNet CIFAR1O. We also conducted IB experiments in a more standard setting, training ResNet18 (He et al. 2016) on the CIFAR10 dataset (Krizhevsky) (209).. In the grokking experi- ments, was important delay the decrease of the RNCI score relative to the fit the training set, which was achieved by adopting a large initialization scale and a small sample size. In contrast, the reduction of redundant information in the IB principle does not necessarily require such a timing dis- crepancy. Therefore, in this experiment, we used the standard initialization scale and the full training set 50,000 examples. Figure 14 shows the results, with test accuracy shown on the left for the reference. For both MI and nHSIC. redundant information decreases training progresses. When compared with the behavior the RNCI score, the trends are similar particularly the later phase of training, supporting our theoretical result that links the reduction of IB superfluous information the decrease of the RNCI score. The left panel shows that this later compression phase corresponds the period after the training has already been fit. This suggests that continuing training to promote neural collapse is beneficial not only for generalization but also from the perspective of the IB principle.
31


Reviews:
[{"review_id": "n88IgsSN9M", "reviewer": "Reviewer_gZ1w", "summary": "This paper demonstrates theoretically that the two training late-phase phenomena of grokking and information bottleneck (IB) dynamics are both due mainly to reduced population within-class variance.  It relates the measure of such variance reduction to the neural collapse degree present in the training dataset.  The paper shows that the model training duration ends sooner than the onset of the neural collapse progress, with the former dependent on the time taken to fit the training set and hence resulting in late-phase phenomena because neural collapse is responsible for such phenomena.  Experiments are conducted to validate its theoretical results.", "questions": "Experiments on more general scenarios for models with more classes are desirable for result validation, since their results are not clear to hold.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper broadens the understanding of training dynamics by revealing the connection between late-phase phenomena and neural collapse.  It analyzes the variance of class-wise representations in the training data and shows that the decrease of empirical within-class variance leads to improved test accuracy, which takes place in the second compression phase to exhibit model grokking behavior after the first fitting phase. \n\nBy considering the grokking and IB dynamics as two representative late-phase phenomena of DNNs, the paper is the very first to demonstrate that both phenomena can be explained in terms of the population within-class variance, offering new insights into their underlying mechanisms. \n\nThe paper conducts a quantitative analysis on the discrepancy between the population within-class variance, allowing to evaluate the progression of neural collapse and implying a corresponding reduction in the population within-class variance to relate\nthe behaviors of grokking and IB dynamics to the development of neural collapse.", "weaknesses": "The experimental study for validating theoretical results is limited to DNN models with two classes, like the assumption of theoretical treatment.  It should be evaluated for the models targeting datasets with more classes, for generalization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iBMhHIy2ZN", "reviewer": "Reviewer_ud5o", "summary": "This paper attempts to provide geometric formulation for grokking’s delayed generalization. They do so by showing that grokking is the combination of a deep neural network following IB dynamics and achieving within-class variance collapse. It derives bounds linking this variance collapse to both generalization error and redundant mutual information, and shows empirically that variance collapse (RNC1) lags training-loss convergence under weight decay. The result is a heuristic unification of IB, grokking, and NC1.", "questions": "1) Could flatness or implicit regularization, rather than NC1 dynamics, explain the delayed generalization you observe?\n\n2) Is thm 3.4 a tight bound or a loose ceiling?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper reframes grokking and IB as natural outcomes of a slowly emerging regularity principle (neural collapse). The paper is well-written, offers clear theorems that formalize variance–generalization connections, and provides empirical breadth.\nThe time-scale separation between fitting and collapse is a genuine insight for grokking dynamics, and the presentation is careful and reproducible.", "weaknesses": "However, the analysis treats neural collapse only through NC1, uses an one-sided IB inequality (Thm 3.4), and replaces Papyan 2020’s ETF-based NC2 metric with a condition-number surrogate that cannot ensure ETF geometry. The paper’s main claim is a unification of grokking and information bottleneck (IB) dynamics through neural collapse (NC). This is conceptually interesting but not rigorously established.\n\nThis unification is heuristic not proven: The entire IB connection rests on a one-sided variance-based upper bound (Thm 3.4) that does not prevent degenerate or trivial solutions (when compression is more important than capturing any relevant information). \n\nTheoretical results rely on strong simplifications (smooth activations, squared loss, pyramidal architecture, Assumption A.3, A.4, A.5). The theoretical guarantees in Theorem 4.3 depend on highly restrictive conditions (pyramidal architecture, smooth activations, squared-loss training) that differ from the experimental setups (ReLU + cross-entropy) . Theorem 4.3 and its proof explicitly assume optimization of the squared loss with weight decay, as given in Section 4.2.  All the convergence results refer to that MSE-based objective, not to the cross-entropy loss used in the experiments.\n\nI also think replacing the NC2 metric with a condition-number surrogate without justification is sloppy. The co-occurrence of NC1 and NC2 is critical for the normative statement of neural collapse to hold (theorem 3 in Papyan 2020). NC2 is strictly defined as equal angles and equal norms both satisfied. While when condition number = 1, the perfectly isotropic mean configuration coincides with ETF. Many non-ETF configurations can also yield small condition numbers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PZZIjSvttK", "reviewer": "Reviewer_oBWN", "summary": "The paper investigates learning phaes of neural nets focusing on the later training stages using wihtin-class variance as a key investigative lense.", "questions": "* How is this behavior if you include different kinds of normalization?\n* Could you intuitively state the impact of assumption A.5 (in appendix)? In what sense, does it restrict initialization (e.g. compared to He or other initialization methods)?\n\n* A variance decrease is also possible by simple scaling, e.g., say X are inputs to a layer then we have var[WX] for some weight matrix W, now if we scale X, i.e., we have cX with c<1, we have c^2*var[WX]. \nIf you use regularization on weights you naturally scale weights and push them towards 0. So from this angle the metric appears not very suitable.\n\n\n* You state \"An intriguing observation in the existing information plane work is that, in the late stage of training, DNNs tend to compress I(Z; X) while preserving I(Z; Y ), thereby moving toward a more optimal solution with respect to the IB objective in Equation (1). \n\nThis observation has already been made outside the IB framework, by looking at the reconstruction loss L(X,Z) over time (see [1] below). That is the paper argues that reconstructing the input X given intermediate representation Z gets better (L(X,Z) decreases) before increasing again (L(X,Z) increases), at the same time cross entropy CL(X,Y) loss only reduces, which is qualitatively identical to your Figure 1 showing I(Z,Y), I(Z,X), except that they also show a third phase, i.e., initially L(Z,X) remains constant.  Do you agree?\n\n[1] Schneider, J., & Prabhushankar, M. (2024). Understanding and leveraging the learning phases of neural networks. In Proceedings of the Thirty-Eighth AAAI Conference on Artificial Intelligence and Thirty-Sixth Conference on Innovative Applications of Artificial Intelligence and Fourteenth Symposium on Educational Advances in Artificial Intelligence (pp. 14886-14893).\n\n\n* Why should in-class variance be a good measure? In particular, you show that it decreases, but it a better measure might something that changes across the phases ? Won't it? \n\n*  [2] (see below) argues that first the mean of the samples are fit, followed by those not well-aligned which might actually strongly impact representations though only minimally impacting the loss, which is fairly intuitive. Does not the initial convergence to the mean (per class) also apply to your formulation?   Essentially you minimize the population within class variance (Definition 3.1) if g(X) ~ Mean_(X,Y=c) g(X|c) ~ E_{X|Y=c}[g(X)].  In particular, the RNC1 score seems to measure just that and directly corresponds to the L2-norm used in [2]. Could you discuss the differences and similarities? \nThe paper [2] argues that after the majority has been fit, learning slows down as but there can still be substanital changes -- the paper says a few (strongly different samples from the mean) can be repsonsible for this using the dot product. This seems to be very different from your approach as you are only concerned with the majority, are you? Or is it also relevant in your case ?\n\n[2] Learning in NN: from fitting most to fitting a few. Neural Computing and Applications, 37(28), 23423-23446.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Important problem: Theory is generally lacking in this area\n* Good alignment of theory and practice.\n* Well written.", "weaknesses": "* There are some concerns regarding the suitability of in-class variance, which appear critical   (see questions).\n* As for any theory, there are a number of assumptions.  and lack of common architectural elements\n* There are concerns regarding novelty: Some of the findings appear known, some approaches seem to overlap with prior work, some key statements appear to be known, though not in this context and under a different analytical lense -  but the paper still of interest despite potential overlaps, but clarification is needed in this regard (see questions).\n* Lack of practical implications (though not a must for a theory paper)\n\n(Note, the initial judgement is more on the negative side due to a number of concerns - in particular the in-class variance as metric, but given questions are clarifyable an increase in score is appropriate as this is overall a strong paper)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)