Swing-by Dynamics in Concept Learning and Compositional Generalization
Authors: Yongyi Yang
Core Francisco Park
Ekdeep Singh Lubana
Maya Okawa
Wei Hu
Hidenori Tanaka
Venue: ICLR 2025 Poster (2025)
Keywords: compositional generalization
concept learning
learning dynamics
out of distribution generalization
Primary area: interpretability and explainable AI
TL;DR: We studied the model learning dynamics on a structured identity mapping learning task and demonstrates that the analysis reveals the underlying mechanisms of phenomena observed in concept learning dynamics.

Abstract: Prior work has shown that text-conditioned diffusion models can learn to identify and manipulate primitive concepts underlying a compositional data-generating process, enabling generalization to entirely novel, out-of-distribution compositions. 
Beyond performance evaluations, these studies develop a rich empirical phenomenology of learning dynamics, showing that models generalize sequentially, respecting the compositional hierarchy of the data-generating process. 
Moreover, concept-centric structures within the data significantly influence a model's speed of learning the ability to manipulate a concept.
In this paper, we aim to better characterize these empirical results from a theoretical standpoint.
Specifically, we propose an abstraction of prior work's compositional generalization problem by introducing a structured identity mapping (SIM) task, where a model is trained to learn the identity mapping on a Gaussian mixture with structurally organized centroids. 
We mathematically analyze the learning dynamics of neural networks trained on this SIM task and show that, despite its simplicity, SIM's learning dynamics capture and help explain key empirical observations on compositional generalization with diffusion models identified in prior work.
Our theory also offers several new insights---e.g., we find a novel mechanism for non-monotonic learning dynamics of test loss in early phases of training.
We validate our new predictions by training a text-conditioned diffusion model, bridging our simplified framework and complex generative models.
Overall, this work establishes the SIM task as a meaningful theoretical abstraction of concept learning dynamics in modern generative models.

Published as a conference paper at ICLR 2025
SWING-BY DYNAMICS IN CONCEPT LEARNING AND COMPOSITIONAL GENERALIZATION
Yongyi Yangl.2. 1,2.3, Core Francisco Park Ekdeep Singh Lubana' Maya Okawa!.?, 1.2,
Wei Hu Hidenori Tanaka
1 CBS-NTT Physics of Intelligence Program, Harvard University, Cambridge, MA, USA
2 Physics Informatics Laboratories, NTT Research, Inc., Sunnyvale, CA, USA
3 Computer Science and Engineering, University of Michigan, Ann Arbor, MI, USA
4 Department of Physics, Harvard University, Cambridge, MA, USA
ABSTRACT
Prior work has shown that text-conditioned diffusion models can learn to identify and manipulate primitive concepts underlying a compositional data-generating process, enabling generalization to entirely novel, ditto--distributionn composi- tions. Beyond performance evaluations, these studies develop rich empiri- cal phenomenology of learning dynamics, showing that models generalize se- quentially, respecting the compositional hierarchy of the data-generating pro- cess. Moreover, concept-centric structures within the data significantly influence model's speed of learning the ability to manipulate concept. In this paper, we aim to better characterize these empirical results from a theoretical standpoint. Specificallly, we propose an abstraction of prior work's 's compositional generaliza- tion problem by introducing a structured identity mapping (SIM) task, where a model is trained to learn the identity mapping on Gaussian mixture with struc- turally organized centroids. We mathematically analyze the learning dynamics of neural networks trained on this SIM task and show that, despite its simplicity, SIM's learning dynamics capture and help explain key empirical observations on compositional generalization with diffusion models identified in prior work. Our theory also offers several new insights- e.g., find novel mechanism for non- monotonic learning dynamics of test loss in early phases of training. We validate our new predictions by training text-conditioned diffusion model, bridging our simplified framework and complex generative models. Overall, this work estab- lishes the SIM task as meaningful theoretical abstraction of concept learning dynamics in modern generative models.
INTRODUCTION
Human cognitive abilities have been argued to generalize to unseen scenarios through the identifica- tion and systematic composition of primitive concepts that constitute the natural world (e.g., shape, size, color) (Fodor, 2001; Fodor et al., 1975; Reverberi et al., 2012; Frankland & Greene, 2020; Russin et al., 2024; Franklin & Frank, 2018; Goodman et al., 2008). Motivated by this perspective, the ability to compositionally generalize to entirely unseen, out-of-distribution problems has been deemed desirable property for machine learning systems, leading to decades of research the topic (Smolensky, 1990; Lake & Baroni, 2018; Hupkes et al., 2020; Ramesh et al., 2021; Rombach al., 2022; Lake & Baroni, 2023; Kaur et al., 2024; Du & Kaelbling, 2024).
Recent work has shown that modern neural network training pipelines can lead emergent abilities that allow model to compositionally generalize when it is trained on a data-generating process that itself is compositional in nature (Lake & Baroni, 2023; Ramesh et al., 2023; Okawa et al., 2023; Lepori et al., 2023; Arora & Goyal, 2023; Zhou et al., 2023; Khona et al., 2024; Rosenfeld et al., 2020; Nagarajan et al., 2020; Arjovsky et al., 2019). For example, Okawa et al. (2023); Park et al. (2024) show that text-conditioned diffusion models can learn to identify concepts that constitute the training data and develop abilities to manipulate these concepts flexibly, enabling generations that represent novel compositions entirely unseen during training. These papers also provide spectrum of intriguing empirical results regarding a example, they reveal that abilities to manipulate individual concepts are learned in sequential
model's 's learning dynamics in a compositional task. For

---

Published as a conference paper at ICLR 2025
(a)
Real data
(b) Synthetic data
Our theory
Color
Loss
Test loss
Train loss
Shape
Shape
square apples"
"blue large triangle"
"(1, 1)
10
10
10'
164
(0,0,1) (0,1,1) (1111)
Optimization steps
Figure I: Structured Identity Mapping Task and Swing-by Generalization Dynamics. (a) Given the input "blue square apples on tree with circular yellow leaves," multimodal model learns to generate concepts in the following order: "apple," "blue" (color), and "square" (shape) (example adapted from Li et al. (2024)). (b) A multimodal synthetic task introduced by Okawa et al. (2023); Park et al. (2024). The training set of the task consists of four distinct compositions of concepts, depicted as blue nodes on cubic graph. A diffusion model is trained on this dataset to systematically study the dynamics of concept learning. With the test prompt "small, blue, triangle." the diffusion model sequentially learns the correct size, shape, and finally color. (c) In this work, we introduce a structured identity mapping task as the foundation for a systematical and theoretical studying of the dynamics of concept learning. The model is trained on a Gaussian mixture data, where the centroids are positioned at certain nodes of hyperrectangle (blue dots) and is evaluated on an out- of-distribution test set (red dot). Our theoretical results not only reproduce and explain previously characterized empirical phenomena but also depict comprehensive picture of the non-monotonic learning dynamics in the concept space and predict a "multiple-descent" curve of the test loss (red curve).
order dictated by the -ata-generating process; the speed of learning such abilities is modulated by data-centric measures (e.g., gradient of loss with respect concept values, such as color of an object); and the most similar composition seen during training often controls performance on unseen compositions.
In this work, we aim to demystify the phenomenology of compositional generalization identified in prior work and better ground the problem (or at least specific variant of it called systematicity) via a precise theoretical analysis. To that end, we instantiate a simplified version of the compositional generalization framework introduced by Okawa et al. (2023); Park al. (2024)-calles the "concept space" (see Fig. Ih)-that is amenable to theoretical analysis. In brief, a concept space is a vector space that serves as an abstraction of real concepts. For each concept (e.g., color), a binary number can be used to represent its value (c.g.. 0 for red and 1 for blue). In this way, a binary string can be mapped tuple (c.g.. (1, O, 1) might represent "big blue triangle") and then fed into the diffusion model as conditioning vector. The model output is then passed through a classifier which produces a vector indicating how accurately the corresponding concepts are generated (e.g. a generated image of big blue triangle might be classified as (0.8, 0.1, 0.9)). In this way, the process of generation becomes a vector mapping, and a good generator essentially performs as an identity mapping in the concept space.
We argue that in fact the salient characteristic of concept space its preemptively defined organi- zation of concepts in a systematic manner, not the precise concepts used for instantiating the frame- work itself. Grounded in this argument, we define learning problem called the Structured Identity Mapping (SIM) task wherein a regression model is trained to learn the identity mapping from points sampled from a mixture of Gaussians with structurally organized centroids (see Fig. Ic). Through a detailed analysis of the learning dynamics of MLP models, both empirically and theoretically, we find that SIM, despite its simplicity, can both capture the phenomenology identified by prior work and provide precise explanations for it. Our theoretical findings also lead to novel insights, e.g.. predicting the existence of a novel mechanism for non-monotonic learning curves (similar to epoch- wise double descent (Nakkiran et al., 2021), but for out- out-of-distnnn didistiibbtiion data) in the early phase of training. which we empirically verify to be true by training a text-conditioned diffusion model. Our contributions are summarized below.
Conceptually, we can think of idealized perfect classifier here.
2

---

Published as a conference paper at ICLR 2025
Structured Identity Mapping (SIM: A faithful abstraction of concept space. We empirically validate our SIM task by training Multi-Layee Perceptrons (MLPs), demonstrating the reproduc- tion of key compositional generalization phenomena characterized in recent diffusion model stud- ies (Okawa et al., 2023; Park et al., 2024). Our findings show: (i) learning dynamics of OOD test loss respect the compositional hierarchical structure of the data generating process; (ii) the rate which a model disentangles a concept and learns the capability to manipulate it is dictated by the sensitivity of the data-generating process changes values of said concept (called "concept signal" in prior work); and (iii) network outputs corresponding to weak concept signals exhibit slowing down in concept space. These results also suggest that the structured nature of the data, rather than specific concepts, drove observations reported in prior work.
Swing-by Dynamics: Theoretical analysis reveals mechanisms underlying learning dy- namies of a compositional task. Building on the successful reproduction of phenomenology with MLPs trained on the SIM task, we further simplify the architecture to enable theoretical analysis. We demonstrate that: (i) analytical solutions with a linear regression model repro- duce the observed phenomenology above, and (ii) the analysis of symmetric 2-layer network (f(:; U) = UU x:) identifies a novel mechanism of non-monotonic learning dynamics gen- eralization loss, which we term Swing-by Dynamics. Strikingly, we show that the learning dy- namics of compositional generalization loss can exhibit multiple descents in its early phase of learning, corresponding to multiple phase transitions in the learning process.
Empirical confirmation of the predicted Swing-by phenomenon in diffusion models. We ver- ify the predicted mechanism of Swing-by Dynamics in text-conditioned diffusion models, observ- ing the non-monotonic evolution of generalization accuracy for unseen combinations of concepts, as predicted by our theory.
In summary, our theoretical analysis of networks trained on SIM tasks provides mechanistic explana- tions for previously observed phenomenology in empirical works and introduces the novel concept of Swing-by Dynamics. This mechanism is subsequently confirmed in text-conditioned diffusion models, bridging theory and practice in compositional generalization dynamics.
PRELIMINARIES AND PROBLEM SETTING
Throughout the paper, we use bold lowercase letters (e.g., X) represent vectors, and use bold uppercase letters (e.g., A) to represent matrices. We use the unbold and lowercase version of corre- sponding letters with subscripts to represent corresponding entries of the vectors or matrices, e.g. T, represent the i-th entry of æ and d, represent the (i,j)-th -th entry of A. For a vector æ and a natural number k, we use F.g to represent the k-dimensional vector that contains the first k entries of x. For natural number k, we use [k] to represent the set whose entries are all except the k-th entry being 1; the dimensionality of this vector is determined by the context if not specified. In the theory part, we frequently consider functions of time, denoted by variable t. If a function g(f) is a function of time i, we denote the derivative of g with respect to t by ((uu)
k). and 1, to represent vector
dg de
Moreover, we sometimes omit the argument t, i.e., g means g(t) for a time
==
{
& is true
determined by the context. For a statement Φ. we define ((6)
to be the indicator
e is false
function of that statement.
2.1 PROBLEM SETTING
Now formally define SIM, which is an abstraction of the concept space. For each concept class, we model them as a Gaussian cluster in the Euclidean space, placed along a unique coordinate di- rection. The distance between the cluster mean and the origin represents the strength of the concept signal, and the covariance of the Gaussian cluster represents the data diversity within the correspond- ing concept class. Additionally, we allow more coordinate directions than the clusters, meaning that some coordinate directions will not be occupied by cluster, which we call non-informative di- rections, and they correspond to the free variables in the generalization task. See Fig. Ic for an illustration of the dataset of SIM.
Training Set. Let d N be the dimensionality of the input space and E [d] be the number of concept classes, i.e., there are s Gaussian clusters, and n E number of samples from each
(x(p)) is generated by the following process: for each cluster. The training set Upe(el { S 1=1
3

---

Published as a conference paper at ICLR 2025
p E [s), each training point of the p-th cluster is sampled i.i.d. from Gaussian distribution ai N
Pp p diag (o)²). where Up 0 is the distance of the p-th cluster center from the origin, and vector with only the first s entries being non-zero, and of describing the data variance on the i-th direction. There is also optionally a cluster centered at O in addition to the & clusters.
Loss function. The training problem is to learn identity mapping on Rd. For model / R x Rd >> Rd and a parameter vector 0 E Th"", we train the model parameters via the mean square error loss.
(())
(0:2
(s)
(2.1)
2nn
Evaluation. We evaluate the model at a Gaussian cluster centered at the point that combines the cluster means of all training clusters. When the variance of the test set is small, the expected loss within the test cluster is approximately equivalent to the loss at its cluster mean. Therefore, for simplicity, in this paper, we focus on the loss at the sum of the test cluster, which a single test point x Ep=1 Mp 1p In App. B, we report further results for the case of various combinations of training clusters, which leads to multiple OOD test points.
OBSERVATIONS ON THE SIM TASK
We first begin by summarizing our key empirical findings on the SIM task. all experiments we MLP models of various configurations, including different number of layers and both linear and non-linear (specifically, ReLU) activations. Throughout this section and the subsequent sections, we frequently consider the model output at the test point 2 over training time, which we call output trajectory of the model.
Due to space constraints, we only present the results for subset of configurations in the main paper and defer other results to App. We note that the findings reported in this section are one-to-one correspondence with results identified using diffusion models in Sec. 5 and prior work (Park et al., 2024).
3.1 GENERALIZATION ORDER CONTROLLED BY SIGNAL STRENGTH AND DIVERSITY
One interesting finding from previous work is that if alter the strength of one concept signal from small to large, the contour of the learning dynamics would dramatically change (Park et al., 2024). Moreover, it is also commonly hypothesised that with more diverse data, the model generalizes better (Gong et al., 2019; Díez-Pastor et al., 2015). Recall that in the SIM task, the distance μa of each cluster represents the corresponding signal strength, and the variance de represents the data diversity. In Fig. 2. we present the output trajectory under the setting of = 2. in which case the trajectory can be visualized in a plane. There are two components be learned in this task and, from the contour of the curve, we can tell the order of different components being learned.
Fig. 2 (a) presents the output trajectory for setting with fixed and balanced ar, and varied μ.. The results show that when A1 < P2, the dynamics exhibit an upward bulging, indicating a preference for the direction of stronger signal. As fu is gradually increased, this contour shifts from upward bulging to a downward concaving, and consistently maintains the stronger signal preference.
In Fig. 2 (b), the μ is fixed to an unbalanced position, with one signal stronger than the other. As we mentioned above, when a is balanced, the model will first move towards the cluster with stronger signal strength. However, when the level of diversity of the cluster with weaker signal is gradually increased, the preference of the model shifts from one cluster to another.
very concrete conclusion can be thus drawn from the results in Fig. 2 (a) and (b): the generalization order is jointly controlled by the signal strength and data diversity, and, generally speaking, the model prefers direction that has a stronger signal and more diverse data. We note that the conclusion here is more qualitative and in Sec. we provide more precise quantitative characterization of how these two values control the generalization order.
3.2 CONVERGENCE RATE SLOW DOWN IN TERMINAL PHASE
In Fig. 2, the arrow-like markers on the line indicate equal training time intervals. In the later phase of training, we observe that the arrows get denser, indicating a slowing down of the learning dynamics: at the terminal phase training. the time required to reduce the number of training steps to reduce the same amount of loss is significantly larger than the beginning.
4

---

Published as a conference paper at ICLR 2025
11,
Trpeenay Tran Text
3.0
12. Trepetey Then/Tee.
Trajectory Thann The
2.0
(........5)
(0.5. 0.05)
Loxid mensbro
Higholinraton
1.0
e212 0.051
Trai
Tret
0.0
0.0
1,0
25
(a)
(b)
(c)
Figure 2: Learning dynamics of MLP SIM task. The figures show the output trajectory of the MLP on a two-dimensional setting (i.e., s = 2). and each marker represents an optimization timepoint. Notice that we only plot the center of the training set can have varied shapes based on the configuration of σ. (a) one-layer linear model with Ø:2 = (.05, .05) and varied μt. Concepts with larger signal (p²) learnt first. (b) one-layer linear model with f4.2 = (1, 2) and varied o. Concepts i with larger diversity (o,) learnt first. (c) 4 layer linear models under f422 (1, and a2 (.05, .05) and different dimensionality. high dim: 64, low dim: d = 2. Notice that (a) and (b) are both in high dim setting. The lower the dimensionality, the stronger Swing-by has.
circle, but the actual training
3.3 SWING-BY DYNAMICS
The results in Fig. (a) and (b) are both performed with one-layer models and under a high dimen- sional setting (d 64). Despite the overall trend being similar in other settings, it is worth explor- ing the change of trajectory as we increase the number of layers, and / or reduce the dimension.
Test loss
2
High dimensional
In Fig. (c), we perform experiments with deeper mod- els, and optionally with lower dimension. Under these changes, we find that the model shows an interesting irregular behavior, where it initially heads towards the OOD test point, but soon turns toward the training set cluster with the strongest signal. This indicates the model, while seems to be generalizing OOD at the beginning, is memorizing the train distribution and unable general- ize OOD at this point. However, with enough training, we find the model start again move towards the test point and thus generalizes OOD. We call this overall dynamic of the output trajectory Swing-by, which we could be suggestive of a non-moneee non-monotonic test loss curve. To assess this further, we track the value of the loss function during training in Fig. 3, demonstrating a double descent- like curvee?. We also note that the Swing-by phenomena seems to be strongest when dimensionality d of the dataset is low, and is rather modest with high dimensional settings. In the high dimensional setting, the OOD loss descent slows down at some point but does not actually exhibit non-monotonic behavior. This low dimensional preference also be explained perfectly by our theory, further described in Sec. 4.
Low dimensional
1
10
Training time
Figure The test loss of multi-layer models.
THEORETICAL EXPLANATION
We next study the training dynamics of a specific class of linear models that are tractable the SIM task and explain the empirical phenomenology of OOD learning dynamics seen in previous section. In Sec. 4.1, first analyze one-layer model whose dynamics be solved analytically. We show that it can explain most phenomena observed in the experiment; however, it fails to reproduce Swing-
We would like note that, it possible understand Swing-by Dynamics as special of so-called epoch-wise double descent (Nakkiran al., 2021; Olmin Lindsten, 2024; Schaeffer et al., 2023); how- ever, epoch-wise wise double descent is generally understood as a consequence of either noisy training or over- parametrization, affecting model's in in-distribution generalization. In contrast, Swing-by by is distributional phe- nomenon where the model fits the training distribution and is hence unable to generalize well OOD, revealing novel mechanism that although leads towards familiar double descent-like curve.
5

---

Published as a conference paper at ICLR 2025
by Dynamics, suggesting that Swing-by Dynamics is intrinsic to deep models, which highlights the fundamental difference between shallow and deep models. In Sec. 4.2, we further analyze the dynamics of symmetric 2-layer linear model, which successfully captures Swing-by Dynamies. Our theoretical results reveal a multi-stage behavior of the model Jacobian during training, which leads the non-monotonic behavior in model output. We show that each stage in the Swing-by Dynamics precisely corresponds to each stage in the Jacobian evolution.
Throughout this section, we assume f(0; æ)) is a linear function of æ. In this case the Jacobian of f (f:n) with respect to æ a matrix that is completely determined by 0, which we denote by We die In this way, the output of the model can be written as f(0; x) Wax. Using the trace trick (with detailed calculations provided in App. C.I), it easy to show that the overall loss function equal to
( = (We - 1) 1/2
(4.1)
where A 1
an Ep= Ek=1 (p) T is the empirical covariance. In this section, we assume is zx
large, in which case A converges to the true covariance of the dataset A = Eeeplaa 1/= which is
diagonal matrix A -dia(a)) a), defined by ap 0
P S 8 for any E [d].
p p>s s
Remark. Notice that in the linear setting we might not directly train We: instead, we train its components. For example, we might have A = (W1, W2) and have Wo = W, W.W, Then, what we actually train is W1 and 22, instead of Wo. As many previous works have emphasized (Arora et al., 2018; & Telgarsky, 2018; Arora et al., 2019; Advani et al., 2020), although the deep linear model has the same capacity the loss landscape of deep linear models can be non-convex.
one-layer linear model, their dynamics can be vastly different and
4.1 A ONE-LAYER MODEL THEORY AND ITS LIMITATIONS
As a warm-up, we first study the dynamics of one-layer linear models, i.c., f(W: Wx. in which case the Jacobian Wo is simply W. As we will show. this setting can already explain most of the observed phenomenology from the previous section including the order of generalization and the terminal phase slowing down, but fails to capture the Swing-by Dynamics, which we will explore in next subsection. Here we present Theorem 4.1, which gives the analytical solution of the one-layer model on the SIM task.
Theorem 4.1. Let (t)) E IRddd be initialized as W = w and updated by W VVCWW, with L be defined by eq. (4.1) with f(W, 2)) Wz, then we have for any z E 1Rd,
f(W(t), 2)x = 1 (RSS)) [1 - cxp ( (axxt)= +
exp (-a,t) wk. (0)
(4.2)
C.(1)
N₂(0)
See App. C.2 for proof of Theorem 4.1. The Theorem shows that the k-th dimension of the output of a one-layer model evaluated on the test point z can be decomposed into two terms: the growth term
(t) 1 (k [1 exp (-a+t) Pk. and the noise term N₂. (t) Σi= cxp (-a,t) Wx.a (0))4- The following properties can be observed for these two terms: (i) the growth term converges to when A: < 8 and when k s, while the noise term converges 0: (ii) both terms converge at an exponential rate; and (iii) the noise term is upper bounded by Σi= (0,s(0))4 If the model initialization small in scale, specifically (0) < maxxc(1) s4s then (t) will always be small, and thus can be omitted. With this assumption in effect, the model output is dominated by the growth term. A closer look the growth term then explains part of the observed phenomenology.
Generalization Order and Terminal Phase Slowing Down. It can be observed that (+) con- verges at an exponential rate, which leads an exponential decay of evolution speed and explains the terminal phase slowing down. Moreover, the exponential convergence rate of (t(t) is controlled by the coefficient a = (oeî + (&)). Therefore, the direction with larger ag. i.c., larger μs and / or ok, converges faster, hence explaining the order of generalization different concepts. The
6

---

Published as a conference paper at ICLR 2025
theorem also reveals the proportional relationship between Mk (concept signal strength) and 0% (data diversity).
The Limitation of the One Layer Model Theory. While we have demonstrated that Theorem 4.1 effectively explains both the generalization order and the terminal phase slowing down, in the solu- tion eq. (4.2), the learning of each direction is independent. This independence omits the possible interaction between the dynamics of different directions in deeper models, and leads to monotonic and rather regular output trajectory (this is verified by the experiment results in Sec. 3.1). However, as the experiments in Sec. 3.3 show, when the number of layers becomes larger, the model actually exhibits a non- non-monotonic trace that can have detours. The theory based on the one-layer model fails in capturing this behavior. In the subsequent subsection, we introduce more comprehensive theory based on a deeper model, and demonstrate that this model explains all the phenomena observed in Sec. 3, especially the Swing-by Dynamics.
4.2 A SYMMETRIC TWO-LAYER LINEAR MODEL THEORY
In this subsection, we analyze a symmetric 2-1ayer linear model, namely =U; æ. where U E Raxdd and d' 2 d. We demonstrate that it accurately captures all the observations presented in Sec. 3, and, more importantly, the theory derived from this model provides a comprehensive understanding of the evolution of the model Jacobian and output, offering a clear and intuitive ex- planation for the underlying mechanism of the model's seemingly irregular behaviors. Due to space constraints, we focus on providing an intuitive explanation of the multi-stage behavior of the model Jacobian and output, and defer the formal proofs the appendix. It is also worth noting that this symmetric 2-layer linear model is a frequently studied model in theoretical analysis (Li et al., 2020; Stöger & Soltanolkotabi, 2021; Jin et al., 2023), and most existing theoretical results for this model focus on the implicit bias of the solution found, instead of on the non-monotonic behavior during training. which is the focus of our analysis.
For convenience, we denote the Jacobian of f at time point by W(1) Wu( The gradient flow update of the j-h entry of W is given by
vis min + ag)) 2
1 thi (3a; az) + 1l (( (3 + a:)]
Σ
Wk. Why (a: + ag 20x)-
2
(1)(())
$1,00
Nty(t)
(4.3)
As noted in eq. (4.3), we decompose the update of waj into three terms. We call Gts(t) (r,(()(a; + ag) the growth term, (t)( _ [3r,;(a;; + az) + 1l (e (3a, + a:)] the suppression term, and (t) = -
Exer kez (t) (t) a, 2ax) the noise term. The name Keds
of these terms suggests their role in the evolution of the Jacobian: the growth term Gi,, always has the same sign as Wi.j. and has a positive contribution to the update, so it always leads to the direction that increases the absolute value of w,,; the suppression term St,5 also has the same sign' as Way but has a negative contribution in the update of Wi.j. so it always leads to the direction that decreases the absolute value of wajj and the effect direction of the noise term is rather arbitrary since de- pends on the sign of wayj and other terms. It is proved in Lemma D.8 that under mild assumptions, the noise term will never be too large; for brevity, we omit in the following discussion and defer the formal treatment of it to the rigorous proofs in App. D.
4.2.1 THE EVOLUTION OF ENTRIES OF JACOBIAN
In order to better present the evolution of the Jacobian, we divide the entries of the Jacobian into three types: the major entries are the first diagonal entries, and the minor entries are the off- diagonal entries who are in the first s rows or first 8 columns, and other entries are irrelevant entries. Notice that the irrelevant entries do not contribute to the output of the test point so we will not discuss them. Moreover, we also divide minor entries into several groups The minor entries in the p-th row or column belongs to the p-th group (thus each entry belongs to two groups). See Fig. for an illustration of the division of the entries.
Notice that since W- -UT is a PSD matrix, the diagonal entries are always non-negative.
7

---

Published as a conference paper at ICLR 2025
G
1.4
Thandient Memotta Memoiittinnn
12
1.5
Value of
0.6
W1.1
1.0
W2,2
0.8
1.0
Loss
W2.1
0.6
Welid Growth
Train
0.4
0.5
0
250 500 Lrtuus
750
1000
1250
1500
1750
2000
2nd
Number opochs
0.0
0.5
1.0
1.5
Figure 5: The learning dynamics symmetrie 2-layer linear model. Left: The change of the test loss and the Jacobian entries with time predicted by the theory; Right: the corresponding model output trajectory. The figures are plotted under s = 2 and all entries of W are initialized positive.
Initial Growth. In this section we assume wi. J are ini-
major entries minor entries
tialized around a very small value such that w < mmax
(This has been shown crucial for generalization (Xu et al., 2019; Liu et al., 2022b), and especially in compositional tasks (Zhang et al.); See App. D.1 for specific assumptions). It is evident that when all Wiyj
group
group
group
close to (we call this period
initial phase), the growth term (ww), while the suppres- sion term and the noise term are (w²²). This suggests that the evolution of w.y is dominated by the growth term. Therefore, in the initial phase, every value in the Jacobian grows towards the direction increasing its absolute value, with the speed assumed that is ordered
d 5
Inrelevant antries
group group group
determined by a; aj. Since in descending order, it is evident that each entry grows faster than those below its right. The Initial Growth stage formally characterized by Lemmas D. to D.3.
Figure 4: An illustration of the en- tries of the Jacobian.
First Suppression. In the Initial Growth stage, the first major entry will be the one that grows exponentially faster than all other entries, making the first one that leaves the initial phase. Once the first major entry becomes significant and non-negligible, it will effect the suppression term of all minor entries the first group. When the difference between aj and ag is large enough, the first major entry is able flip the growth direction of the first group of minor entries and push their values to 0. The suppression stages are characterized by Lemma D.7.
Second Growth and Cycle. Once the suppression of the first group of minor entries takes effect, the second major entry becomes the one that grows fastest. Thus, the second major entry will be the second one that leaves the initial stage. Again, when the second major entry becomes large enough, will suppress the second group of minor entries and push their value to This process continues like this: the growth of major entry is followed by the suppression of the corresponding group of minor entries, which, in turn, leaves space for the growth of the next major entry. The general growth stages are characterized by Lemma D.4 and the fate of off-diagonal entries is characterized by Lemma D.8.
Growth Slow Down and Stop. Notice that the suppression term of a major entry is also influ- enced by its own magnitude. Therefore, when a major entries becomes significantly large, also suppresses itself, leading to the slowing down of its growth. Note that this effect only slows down the growth but will not reverse the direction, since for major entries the suppression term always smaller than the growth term, until wa,, becomes 1 where the growth and suppression terms equal and the evolution stops. The terminal stage the growth of major entries are characterized by Lemma D.5.
8

---

Published as a conference paper at ICLR 2025
4.2.2 EXPLAINING MODEL BEHAVIOR
Recall that we have f (t); 3)2 = ESn (k,p(t)ppp We now explain how the stage- wise evolution of Jacobian described Sec. 4.2.1 determines the evolution of the model output.
Generalization Order and Terminal Phase Slowing Down. From the discussions in Sec. 4.2. 1, by the end of the training, all the major entries converge to 1 and all minor entries converge to O. The major entries grows in the order of corresponding ap. which is determined by Me and op. and slows down when approaching the terminal. This explains our observation that directions with larger Hp and or op is learned first, well as the terminal phase slowing down.
Swing-by Dynamics and Non-monotonie Loss Curve. We argue that the Swing-by Dynamics and the non-monotonic loss curve is caused by the multi-stage major growth vs. minor growth / suppression process. Importantly, in certain configurations, minor entries growing towards larger absolute values (which is the incorrect solution) can lead to the decay of the OOD test loss, and cause "illusion of generalizing" that the output trajectory is moving towards improving OOD generalization. However, this effect is later eliminated by the suppression of the corresponding minor entries, leading to a double (or multiple) descent-lik, like loss curve and reversal in the output trajectory.
More concretely, consider the first (initial) growth stage as an example. In this stage, for each k E [s]. f (U(t); , is dominated by (k.1(t) since W grows fastest among all the entries in the k-th row. If wk,1 happens to be initialized positive, then f (U))); Ø) grows towards 1, which is the correct direction", and loss thus decays. Since in a symmetric initialization, each entry has equal chance of being initialized positive or negative, when s small, it easy to have many minor entries initialized positive, whose growth contributes to the decaying of loss. This causes an illusion that the model is going towards the right direction of OOD generalization. After the minor entries of the first group are suppressed, their contribution to the decaying of the loss is canceled, which leads to the output trajectory turning back to the direction of memorizing a training cluster and a transient loss increase.
Fig. 5 presents the loss curve and the Jacobian entry evolution predicted by the theory with a specific initialization. Notice how, as claimed above, the first and second descending of loss accurately corresponds the initial and second growth of the major entries, and the ascending of the loss corresponds to the suppression of the minor entries. When > 2, there are multiple turns of growth and suppression stages and can possibly leads to a multiple-descent-ll loss curve, which confirm and illustrate in App. E.I.
Remark on the Role of Out-of-Distributtion If one of the training cluster very close to the test point, i.e. there is an E [s] where σp Z Mp, then the setting becomes highly in-distribution, and intuitively there shouldn't be a significant Swing-by since the training loss monotonically decays. We note that this intuition is captured by our theory through Assumption D.5, which requires that the signal strength of different directions to be distinct enough, and this essentially prevents a too large op (compared to pp).
Remark on Failure Modes. We note that our theory also provides an explanation on instances when the model fails to achieve OOD generalization when one more of our assumptions outlined in App. D.I breakdown. A specific case is when a major entry Wi.k is overly suppressed by a corresponding minor entry before it can begin to grow, causing the growth term G becomes nearly zero. Consequently, the model output at & in that direction converges to 0, instead of μa as expected. See App. E.3 for more discussions and illustrations.
Remark on Existing Work. There has been extensive research on the non-monotonic behavior of linear neural networks (in various settings). We note that existing studies either focus on layer networks (Pezeshki et al., 2022; Heckel & Yilmaz, 2020) or diagonally initialized networks (Lampinen & Ganguli, 2018; Saxe et al., 2019; Even et al., 2023; Pesme & Flammarion, 2023; Olmin & Lindsten, 2024), which essentially make the evolution of each direction decoupled. This decoupling simplifies the learning dynamics and can overlook critical aspects thereof (as we dis- cussed in the preceding subsection). In contrast, our analysis. through a careful treatment of each entry of the Jacobian, does not need to make the diagonal initialization assumption, hence allowing
NNotice that this is true even when & a 1, f.c. WK,,1 is minor entry.

---

Published as a conference paper at ICLR 2025
us capture and characterize the rich behaviors that arise from the interaction between different directions.
5 DIFFUSION MODEL RESULTS
Tying back to our original motivation of devising abstraction of concept space first explored in text-to-image generative diffusion models, we now aim verify if our findings and predictions made on SIM task can be reproduced in a more involved practical setup with diffusion models. To this end, we borrow the setup from Okawa et al. (2023); Park et al. (2024) and train conditional image diffusion models two concepts si22 ze and color.
0.7
Color Concept Signal
Color Concept Signal
0.358
0.579
0.137
0.358
0.579
0.6
Color Concept Signal
0.247
0.468
0.689
0.247
0.689
Concept Space MSE
0.4
Size
0.3
0.4
0.2
10-
02
0.0
10-5
Color
Number of steps
160
'1o
Number of gradient steps
5000 10000 15000 20000
(a)
(b)
(c)
Figure 6: Main observations reproduced on prompt-to-image diffusion models. (a) Signal strength controls generalization speed and order; and transient bend towards the concept with stronger signal.
(b) A double descent-like curve for the concept space MSE caused by Swing-by Dynamies. (c) Concept learning gradually slows down, Our theory predicts speed of concept learn- ing slows down at an exponential rate, which broadly matches the experimental results.
In Fig. 6, we repeat the experiments did in Section 4 of Park et al. (2024): we consider a synthetic setup where the model learns to generate the image of a circle of the indicated size and color given by text input. The training set only contains samples from input pairs (red, big), (blue, big) and (red, small 11), and the test set contains samples from an OOD input pair (blue, ue, smal 1 1)) A pretrained classier is applied to map the image generated by the diffusion model back to the concept space so that we can plot the model output curve in the concept space, as shown in Fig. 6 (a). In order to compare model behaviors under different signal strength, we tune the contrast between the color of red samples and blue samples in the training set to control the strength of the color signal (See Fig. 19 for an illustration). See App. G for experiment details.
Fig. 6 (a) shows that the level of concept signal largely alters the generalization dynamics. Specif- ically, we see that the order of compositional generalization is determined by the color concept signal, and can be reversed when we tuning the concept signal. Additionally, Fig. (a) along with the corresponding loss curve plot Fig. 6 (b) also shows Swing-by Dynamics where the generalization dynamics show a bend towards the concept with stronger signal; the bend transient and the gen- eration eventually converges to the correct class (small blue circle), and the corresponding test loss curve shows a double descent-like trend. Fig. 6 (c) confirms that the speed of compositional gener- alization, quantified by the concept space traversal distance per step, decelerates at an exponential rate, as expected from our theoretical findings (Theorem 4.1).
CONCLUSION
In this paper, we propose SIM task as a further abstraction of the "concept space" previously ex- plored by Okawa et al. (2023); Park et al. (2024). We conduct comprehensive investigation into the behaviors of regression model trained on SIM, both empirically and theoretically, demonstrating that the learning dynamics on SIM effectively captures the phenomena observed on image gener- ation task, establishing SIM as a basis for studying compositional generalization. Critically, our theoretical analysis uncovers the underlying causes of several phenomena that previously observed on compositional generalizations, as well as predicting new ones that characterizes the multi-stage and monooooooni learning dynamics, which have been largely overlooked in earlier research. Our diffusion model experiments further verify the validity of our analysis. Additional discussions and potential future work directions can be found in App. E.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
Part of this work was done while YY and WH were visiting the Simons Institute for the Theory of Computing.
REFERENCES
Madhu S Advani, Andrew M Saxe, and Haim Sompolinsky. High-dimensional dynamics of gener- alization enror in neural networks. Neural Networks, 132:428-44 4444284446, 2020.
Jacob Andreas. Measuring compositionality in representation learning. arXiv:1902.07186 2019.
arXiv preprint
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint. arXiv:1907.02893, 2019.
Sanjeev Arora and Anirudh Goyal. A theory for emergence of complex skills in language models. arXiv preprint arXiv:2307.15953 2023.
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis gradient descent for deep linear neural networks. arXiv preprint arXiv:1810.02281, 2018.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Emanuele Bugliarello and Desmond Elliott. The role syntactic planning in compositional image captioning. arXiv preprint arXiv: 2101 2021.
Colin Conwell and Tomer Ullman. Testing relational understanding in text-guided image generation. arXiv preprint arXiv:2208.0002 2022.
Colin Conwell and Tomer Ullman. A comprehensive benchmark human-like relational reasoning for text-to-image foundation models. In ICLR 2023 Workshop Mathematical and Empirical Understanding of Foundation Models, 2023.
José F Díez-Pastor, Juan J Rodríguez, César García-Osorio, and Ludmila Kuncheva. Diver- sity techniques improve the performance of the best imbalance learning ensembles. Information Sciences, 325:98-111 117, 2015.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021. URL https::///// /tta/aiivv.. www.or/aaaaaa...eee 2200 1929.
Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference Machine Learning, pp. 1655-1664. PMLR, 2019.
Yilun Du and Leslie Kaelbling, Compositional generative modeling: A single model is not all you need, arXiv preprint arXiv:2402.0f 103, 2024.
Yilun Du, Shuang Li, Yash Sharma, Josh Tenenbaum, and Igor Mordatch. Unsupervised learning of compositional energy concepts. Advances in Neural Information Processing Systems, 34:15608- 15620, 2021.
Yilun Du, Conor Durkan, Robin Strudel, Joshua Tenenbaum, Sander Dieleman, Rob Fer- gus, Jascha
Arnaud Doucet, and Will Grathwohl. Reduce, reuse. recycle: Compositional generation with energy-based diffusion models and meme. arXiv preprint 1555: 2023.
Mathieu Even. Scott Pesme. Suriya Gunasekar, and Nicolas Flammarion. (s) gd over diagonal linear networks: Implicit bias, large stepsizes and edge of stability. Advances in Neural Information Processing Systems, 36:29406-29448, 2023.

---

Published as a conference paper at ICLR 2025
Weixi Feng, Xuehai He, Tsu-Jui Fu, Varun Jampani, Arjun Akula, Pradyumna Narayana, Sugato Basu, Xin Eric Wang, and William Yang Wang. Training-fre free structured diffusion guidance for compositional text-to-image synthesis. arXiv preprint arXiv:2212.05032, 2022.
Jerry Fodor. Language, thought and compositionality. Royal Institute of Philosophy Supplements, 48:227-242, 2001.
Jerry A Fodor al. The language of thought, volume 5. Harvard university press Cambridge, MA, 1975.
Steven M Frankland and Joshua D Greene. Concepts and compositionality: in search of the brain's language of thought. Annual review of psychology, 71:273-303, 2020.
Nicholas T Franklin and Michael J Frank. Compositional clustering in task structure learning. PLoS
computational biology, 14(():e10061 16, 2018.
Tejas Gokhale, Hamid Palangi. Besmira Nushi, Vibhav Vineet, Eric Horvitz, Ece Kamar, Chitta Baral, and Yezhou Yang. Benchmarking spatial relationships in text-to-image generation. arXiv preprint arXiv:2212.10011 2022.
Zhiqiang Gong, Ping Zhong, and Weidong Hu. Diversity in machine learning. leee Access, 7: 2019.
Noah D Goodman, Joshua Tenenbaum, Thomas L Griffiths, and Jacob Feldman. Compositionality in rational analysis: trammar-based induction for concept learning. The probabilistic mind: Prospects for Bayesian cognitive science, 2008
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition, 2015.
Reinhard Heckel and Fatih Furkan Yilmaz. Early stopping in deep networks: Double descent and how to eliminate it. arXiv preprint arXiv:2007.10099 2020.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023.
Dicuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
Ben Hutchinson, Jason Baldridge, and Vinodkumar Prabhakaran. Underspecification in scene description-tto-dppiiionn tasks, 2022.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv:1810.0200 02032, 2018.
Jikai Jin, Zhiyuan Li, Kaifeng Lyu. Simon Shaolei Du, and Jason D Lee. Understanding incre- mental learning of gradient descent: A fine-grained analysis of matrix sensing. In International Conference Machine Learning, pp. 15200-15220 PMLR, 2023.
Justin Johnson, Bharath Hariharan, Laurens Van Der Maaten, Li Fei-Fei, C Lawrence Zimickk and Ross Girshick. Clevr: A diagnostic dataset for compositional language and elementary visual reasoning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2901-2910, 2017.
Simran Kaur, Simon Park, Anirudh Goyal, and Sanjeev Arora. Instruct-skillmix A powerful pipeline for Ilm instruction tuning. arXiv preprint arXiv:2408. /4774, 2024.
Mikail Khona, Maya Okawa, Jan Hula, Rahul Ramesh, Kento Nishi, Robert Dick, Ekdeep Singh Lubana, and Hidenori Tanaka. Towards an understanding of stepwise inference in transformers: A synthetic graph navigation model. arXiv preprint arXiv:2402.07757, 2024.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances neural information processing systems, 34:21696-21707 2021.
12

---

Published as a conference paper at ICLR 2025
Nupur Kumari, Bingliang Zhang. Sheng-Yu Yu Wang, Eli Shechtman, Richard Zhang, and Jun- Yan Zhu. Ablating concepts in text-to-image diffusion models. arXiv preprint arXiv:2303.1333 2023.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of seguencc-to-sequenee recurrent networks. In International conference on machine learning, pp. 2873-2882. PMLR, 2018.
Brenden M Lake and Marco Baroni. Human-like systematic generalization through a meta-learning neural network. Nature, 15--22, 2023.
Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv:1809.10377 2018.
Evelina Leivada, Elliot Murphy, and Gary Marcus. Dall-e 2 fails to reliably capture common syn- tactic processes. arXiv preprint (tt.::2210/22880 2022.
Michael A Lepori, Thomas Serre, and Ellie Pavlick. Break it down: Evidence for structural compo- sitionality in neural networks. arXiv preprint arXiv:2301.10884 2023.
Martha Lewis, Qinan Yu, Jack Merullo, and Ellie Pavlick. Does clip bind concepts? probing com- positionality in large image models. arXiv preprint 1rXiv:2212.105333 2022.
Hao Li, Yang Zou, Ying Wang, Orchid Majumder, Yusheng Xie, R Manmatha, Ashwin Swami- nathan, Zhuowen Tu, Stefano Ermon, and Stefano Soatto. On the scalability of diffusion-based text-to-image generation. arXiv preprint arXiv:2204. 0228330 2024.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, arXiv preprint wwXiv:2012.09839 2020.
Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum. Compositional visual generation with composable diffusion models. In Computer Vision-ECCV 2022: 17th Euro- pean Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings Part XVII, pp. 439. Springer, 2022a.
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference Learning Representations, 2022b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
Ekdeep Singh Lubana, Kyogo Kawaguchi, Robert P Dick, and Hidenori Tanaka. A percolation model of emergence: Analyzing transformers trained on a formal language. arXiv preprint arXiv:2408. 12578, 2024.
Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient descent quantizes relu network features. arXiv preprint
2018.
Gary Mareus, Ernest Davis, and Scott Aaronson. A very preliminary analysis of dall-e 2. arXiv preprint vrXiv:2204.13807, 2022.
Hancheng Min, Enrique Mallada, and René Vidal. Early neuron alignment in two-layer relu net- works with small initialization. arXiv preprint arXiv:2307. 12851, 2023.
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. arXiv preprint arXhv:2010.15775. 2020.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechanies: Theory and Experiment, 2021(12): 2021.
Maya Okawa, Ekdeep Singh Lubana, Robert P. Dick, and Hidenori Tanaka. Compositional abilities emerge multiplicatively: Exploring diffusion models on a synthetic task. Advances in Neural Information Processing Systems, 2023.
13

---

Published as a conference paper at ICLR 2025
Amanda Olmin and Fredrik Lindsten. Towards understanding spoch-wiss double descent in two- layer linear neural networks. arXiv preprint arXiv:2407.09843 2024.
Core Francisco Park, Maya Okawa, Andrew Lee, Ekdeep Singh Lubana, and Hidenori Tanaka. Emergence hidden capabilities: Exploring learning dynamics in concept space. Advances in Neural Information Processing Systems, 2024.
Scott Pesme and Nicolas Flammarion. Saddle-to-saddee dynamics in diagonal linear networks. Ad- vances Neural Information Processing Systems, 36:7475-7505, 2023.
Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218-29238 2021.
Mohammad Pezeshki, Amartya Mitra, Yoshua Bengio, and Guillaume Lajoie. Multi-scale feature learning dynamics: Insights for double descent. In International Conference on Machine Learn- ing. pp. 17669-17690. PMLR, 2022.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alee Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation, 2021.
Rahul Ramesh, Mikail Khona, Robert P Dick, Hidenori Tanaka, and Ekdeep Singh Lubana. How capable can transformer become? a study on synthetic, interpretable tasks. arXiv preprint arXiv:2311.12997, 2023.
Royi Rassin, Shauli Ravfogel, and Yoav Goldberg. Dalle-2 is seeing double: flaws in word-to- concept mapping in text2image models. arXiv preprint arXiv:2210. 10606, 2022.
Carlo Reverberi, Kai Görgen, and John-Dylan Haynes. Compositionality rule representations in human prefrontal cortex. Cerebral cortex, 22(6):1237-1245 2012.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation, 2015.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. arXiv preprint arXiv:2010.05761, 2020.
Jacob Russin, Sam Whitman McGrath, Ellie Pavlick, and Michael Frank. Is human composition- ality meta-learned? Behavioral and Brain Sciences, 47:e162, 2024.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics learning in deep linear neural networks. arXiv preprint arXiv:1312.6120,, 2013.
Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development deep neural networks. Proceedings of the National Academy of Sciences, 116 (23):11537-11546, 2019.
Rylan Schaeffer, Mikail Khona, Zachary Robertson, Akhilan Boopathy, Kateryna Pistunova, Ja- son Rocks, Ila Rani Fiete, and Oluwasanmi Koyejo. Double descent demystified: Identifying, interpreting & ablating the sources of a deep learning puzzle. arXiv preprint arXiv:2303.1 14151, 2023.
Lukas Schott, Julius Von Kügelgen, Frederik Träuble, Peter Gehler, Chris Russell, Matthias Bethge, Bernhard Schölkopf. Francesco Locatello, and Wieland Brendel. Visual representation learning does not generalize strongly within the same domain. arXiv preprint arXiv:2107.08221. 2021.
Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate dall-e learns to compose. arXiv preprint arXiv:2110.11405 2021.
Paul Smolensky. Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artificial intelligence, 46(1-2):159-22 1990.
14

---

Published as a conference paper at ICLR 2025
Sam Spilsbury and Alexander Ilin. Compositional generalization in grounded language learning via induced model sparsity. arXiv preprint 2200:22077.05118, 2022.
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruc- tion. Advances in Neural Information Processing Systems, 34:23831-23843, 2021.
Tristan Thrush, Ryan Jiang, Max Bartolo, Amanpreet Singh, Adina Williams, Douwe Kiela, and Candace Ross. Winoground: Probing vision and language models for visio-linguistic composi- tionality. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recogni- tion, pp. 9238-5248, 2022.
Vishaal Udandarao, Ameya Prabhu, Adhiraj Ghosh, Yash Sharma, Philip HS Torr, Adel Bibi, Samuel Albanie, and Matthias Bethge. No" zero-shot" without exponential data: Pretraining concept frequency determines multimodal model performance. arXiv preprint arXiv:2404.04115, 2024.
Josef Valvoda, Naomi Saphra, Jonathan Rawski, Adina Williams, and Ryan Cotterell. Benchmark- ing compositionality with formal languages. In Proceedings of the 29th International Conference on Computational Linguistics, pp. 6007-6018, 2022.
Thaddäus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Composi- tional generalization from first principles. Advances in Neural Information Processing Systems, 36, 2024.
Guangyue Xu, Parisa Kordjamshidi, and Joyce Chai. Prompting large pre-trained vision-language models for compositional concept learning, arXiv preprint 1XXx:2211.05077, 2022.
Zhi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. In Neural Information Processing: 26th International Conference, ICONIP 2019, Sydney, NSW, Australia, December 12-15, 2019, Proceedings, Part I 26, pp. 264-274. Springer, 2019.
Mert Yuksekgonul, Federico Bianchi, Pratyusha Kalluri, Dan Jurafsky, and James Zou. When and why vision-language models behave like bag-of-words models, and what to do about it? arXiv preprint arXiv:2210.01936, 2022.
Tian Yun, Usha Bhalla, Ellie Pavlick, and Chen Sun. Do vision- -aiguuge pretrained models learn primitive concepts? arXiv preprint arXiv:2203. 17271, 2022.
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang. Yaoyu Zhang, and Zhi-Qin John Xu. Initialization is critical to whether transformers composite functions by reasoning or memorizing. In The Thirty-eighth Annual Conference Neural Information Processing Systems.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint a Xxx::2310.160222 2023.
15

---

Published as a conference paper at ICLR 2025
A RELATED WORK
In this section, we provide some context for this paper by reviewing some existing work on compo- sitional generalization and the study of deep linear networks.
Compositional Generalization. Prior work compositionality has often focused on benchmark- ing of pretrained models (Thrush et al., 2022; Andreas, 2019; Lewis et al., 2022; Lake & Baroni,
2018; Yun et al., 2022; Lepori et al., 2023; Johnson al., 2017; Conwell & Ullman, 2022; Yuksek- gonul et al., 2022; Schott et al., 2021: Gokhale et al., 2022; Valvoda et al., 2022) or proposition of protocols that allow generation of compositional samples (Du et al., 2021; 2023; Liu et al., 2022a; Xu et al., 2022; Yuksekgonul et al., 2022; Bugliarello & Elliott, 2021; Spilsbury & Ili., 2022; Ku- mari al., 2023; Du & Kaelbling, 2024). While perfect compositionality in natural settings is still lacking (Marcus et al., 2022; Leivada et al., 2022; Conwell & Ullman, 2022; 2023; Gokhale et al., 2022; Du et al., 2023; Liu et al., 2022a; Singh al., 2021; Rassin et al., 2022; Feng et al., 2022; Hutchinson et al., 2022), several works have demonstrated via use of toy settings that this is unlikely to be an expressibility issue, as was hypothesized. e.g., by Fodor et al. (1975), since the model can in fact learn to perfectly compose in said toy settings. The ability to compose is in fact rather dis- tinetly emergent (Okawa et al., 2023; Lubana et al., 2024) and the model learning it often correlates with distinctive patterns in the learning dynamics, as identified by Park et al. (2024). We note that there has in fact been some work on understanding compositional generalization abilities in neural networks (Wiedemer et al., 2024; Udandarao et al., 2024; Ramesh et al., 2023), but, unlike us, the focus of these papers is not on the model's learning dynamics.
Learning Dynamics of Deep Linear Networks. Deep linear networks has been commonly studied model for learning dynamies, and existing works mostly focus on the final solution found by the model, which primarily concerns the stationary point of the dynamics (Arora et al., 2018; Ji & Telgarsky, 2018; Du Hu, 2019; Arora et al., 2019; Advani al., 2020). There have also been works that try to characterize the full learning dynamics; however, they generally require the learning of each direction (ncuron) to be decoupled (Saxe et al., 2013; Olmin & Lindsten, 2024; Lampinen & Ganguli, 2018; Even et al., 2023; Pesme& & Flammarion, 2023; Pesme et al., 2021), which can be realized through a specific initialization choice. The decoupling assumption ignores the interaction between different neurons and highly simplify the dynamics, and
mentioned in Sec. 4.1, make it unable to capture some important phenomena in practice. The symmetric 2-layer linear model is also specific model that is frequently studied, especially in matrix sensing (Li et al., 2020; Stöger & Soltanolkotabi, 2021; Jin et al., 2023), and as we noted in Sec. 4.2, current theoretical results of this model focus on the implicit biases in the solutions learned, while our analysis. on the other hand, aims at characterizing the full learning dynamics and focus on its OOD behavior.
B MODEL COMPOSITIONALLY GENERALIZE IN TOPOLOGICALLY CONSTRAINER ORDER
In this section, we introduce another phenomenon observed on SIM task learning that we do not put in the main paper: the order of compositional generalization happens in topologically constrained order.
In this section, instead of the single test point ã, we introduce a hierarchy of test points. Specifically, let I {0, 1). be the index set of test points. For each v E I, define test point
((v)
1pp
(B.1)
and call (+()) the test point with the index v. Intuitively, the index v describes which training sets are combined into the current test point. If |||||| then ((w) is the center of one of the training clusters.
We assign the component-wise ordering < to the index set I, i.e., for u,v E I, we say 3 if and only Vi E [m], uz S v. It's easy to see that a partial-ordering.
Interestingly, in the SIM experiment, the order of the generalization in different test points strictly follow the component-wise order. This finding can be descrihed formally in the following way: the
16

---

Published as a conference paper at ICLR 2025
loss function is order homomorphism between z the index set, and & on the real number. Let (() the loss function of the test point 2, then we have the following empirical observation:
Yu, u l ( s (z!10)
(B.2)
14
Training Time
Figure 7: The loss each test point different timepoints during training for 2-layer MLP with ReLU activation. Each graph represents timepoint. Each node in the graph represents a test point, with index printed on it, and edges connecting nodes with Hamming distance The color of the graph represents the loss of corresponding test point. Notice that truncate the loss at in order to unify the scale. From lest to right: epoch = 3.5.5.
In Fig. 7 we show the loss of each test point in several timepoints, with o
There is clear trend that the test points that wise order) will only be learned after all of its predecessors are all learned. We call this phenomenon the topological constraint since the constraint is based on the topology of the graph in Fig. 7.
the right of the graph (larger in the component-
C PROOFS AND CALCULATIONS
In the main text we have omitted some critical proofs and calculations due to space limitation. this section we provide the complete derivations. Note that we postpone the proof of related theorems of Sec. 4.2 to App. because of their length.
C.1 THE LOss FUNCTION WITH LINEAR MODEL AND INFINITE DATA LIMIT
In this subsection we derive the transformed loss function eq, (4.1), as well as the expression of data matrix A. For convenience we denote We by W. We have
((0)
W
I) (p)
(C.1)
2ns
1
Tr
(P) (W - T) -1x=
(C.2)
2ns
( I)_ (W- - 1) 2) 11 [p)
Tr
(C.3)
2ns
1
1 (p) (P)
TTr
(W (W (W) T -1)
(C.4)
2
n15
1
Th
1//2(W - 1) TW A 1/2
(C.5)
2
1
(W 1) 1/2 1/2
(C.6)
2
Let & be the data generating process. can be viewed as two components: first assign one of the clusters, and then draw Gaussian vector from Gaussian distribution in that cluster. Specifically,
17

---

Published as a conference paper at ICLR 2025
let be an arbitrary sample from the traning set, then the distribuition of æ is equal to
(1) + diag(o)),
(C.7)
where y is uniform random variable taking values in [8] and & (0, I) is random Gaussian vector that is independent from Here ~ represents having the same distribution.
When n + co, the data matrix A converges to the true covariance, which is is
A E (rrr
(C.8)
(n()) + +diag(o)E) (u(") + diag(o))) ]]
(C.9)
(a()) ( E diE diag(o)
(C.10)
( + diag(o)²
(C.11)
Σ 1 _ diag(o))
(C.12)
diag (1)² + diag(o)²²
(C.13)
C.2 PROOF OF THEOREM 4.1
In this subsection for the notation-wise wise convenience we denote W 0. Since the model one- layer, the loss function eq. (4.1) becomes
<(W) = (W 1)A F
(C.14)
and the gradient is
VL(W) = (W 1JA WA A. We denote the k-th row of W and by ur and A respectively. Then we have
(C.15)
ib Aw, + ax.
(C.16)
The solution of this differential equation is (x(() exp (-At) (0) A-'as + A-'aa ak:
(C.17)
where we use the convention x (0-1) = 0 to avoid the non-invertible case of A.
Thus for any IRd we have
f(W ((W(t);2) 2); (wxx(t),2
(C.18)
((I - c -At) A-1 ak, 2) + (6- (c-A' (wx( 2)
(C.19)
"
1 1- c app
ap2 +
aik wk. (0)
(C.20)
ap
1=1
Zk
(C.21)
and this proves the claim.
D THEORETICAL ANALYSIS OF THE TWO LAYER MODEL
In this section we provide detailed analysis of the symmetric two-layer linear model described in Sec. 4.2.
18

---

Published as a conference paper at ICLR 2025
In this section we assume finite step size, i..., W N - Rdxd is initialized by W (0) and updated by
+(t+1) - W(O) -(V) (C(()))))
(D.1)
"
+(t))+++ [AW + (D.2)
The update of each entry w (t) can be decomposed into three terms, as we described in the main text:
Wa, 1) w),((1) - (1) +a,)
wk Wk, + +az + 20x)
(D.3)
(t) + +a,)
(D.4)
(t)
whi
[us + a₁) + 1(s4)) + a:)]
(D.5)
(1)(1)
Σ (t) + aj + 24x).
(D.6)
knd
(t)
D.1 ASSUMPTIONS
We need make several assumptions to prove the results. Below we make several assumptions that all commonly hold in the practice. The first assumption to make is that both the value of ag and the initialization of W is bounded.
Assumption D.1 (Bounded Initialization and Signal Strength). There exists a 0, 1, such that
vk,a S 0g 70, Vi,j, w & [w)) S Bw.
(D.7)
(D.8)
The second assumption that the step size is small enough. Assumption D.2 (Small Step Size). There exists a constant K 2 20, such that S 9Kyaa
Next, we define a concept called initial phase. The definition of initial phase related a constant P > O.
Definition D.1. Assume there is constant P >> O. For entry (i. and time t. (f/ S P3w, we say this entry is in initial phase.
As the name suggests, the initial phase, the entries shouldn't be too far away from the initializa- tion, and we achieve this by upper bound of P.
Assumption D.3 (Small Initial Phase). PwB < 0.4.
The next assumption to make is that the initialization value (i.e. w) should not be too large. Assumption D.4 (Small Initialization).
min(K - 1,, K- 1/21
1
w S min
(D.9)
PK-d332
V28
and k > 1.1, and < 11 KKC-,,P2
Finally, we also assume that the signal strength difference is significant enough.
Assumption D.5 (Significant Signal Strength Difference). For any i > j. we have
0, + aj
log P
(D.10)
20
10k log PBw + log P3
and there exists a constant C > 1 such that a; 3a, C-1
19

---

Published as a conference paper at ICLR 2025
D.2 THE CHARACTERIZATION OF THE EVOLUTION OF THE JACOBIAN
In this subsection, we provide a series of lemmas that characterize each stage the evolution of the Jacobian matrix W.
The whole proof is based induction, and in order avoid a complicated induction, we make the following assertion, which obviously holds at initialization.
Assertion D.1. For all E N, if * j. then the entry (i:j) stays in the initial phase for all time.
We will use Assertion D.I as an assumption throughout the proves and prove it at the end. This is essentially another way of writing inductions.
We have the following corollary that directly followed by Assertion D.1.
Corollary D.1. For all
and all i.j. (t) s 2Pyad33'2
Now, we are ready to present and prove the major lemmas. The first lemma is to post (rather loose) upper bound of the value of the entries.
Lemma D.1 (Upper Bounded Growth). Consider entry (i.j). We have for all E N, at timepoint the absolute value of the jitt) entry satisfies
[ (t) < [u,, exp (nt + ay)m].
(D.11)
Proof. Since of the Noj term we only its absolute value, the positive case and negative case are symmetric. WLOG we only consider the case where wi, (0) > 0 here.
The claim is obviously satisfied at initialization. We use it as the inductive hypothesis. Suppose at timepoint 2 T the claim satisfied, we consider the time step / +
Since Assertion D.1 guaranteed that every non-diagonal entry is the initial phase, and the S term has different symbol with way (0), we have
(t) + Niy(t)) < 2y
(D.12)
We have
(t -w:,(() S mwi,(t)(a aj) + + 2Pmadd33
(D.13)
naa ay)w,, (0) exp (nt + az)rt + 2Pryad3
(D.14)
2Pmrad332
wh,(0) In + an) n(a + az) + wi,; exp [M a,)
aa)
(D.15)
From Assumption D.4, we have
2P 2Pmadd322
7(a, + a,) + (i.i(0) exp (nt(a; + ag)m]
< naa; + a;) + 2Pyad33
(D.16)
$ n(a; + aaj) + 2(x: - 1)na
(D.17)
s un(a; + az)
(D.18)
s exp(xnla: + +a,]) - 1,
(D.19)
thus have
( wiy (t) + (exp(<n(a; [a; + a,]) - 1] wi., w
(D.20)
s wi, (0) exp In(t + 1)(a; + nx)
(D.21)
Next, we prove that Lemma D.I is tight in the initial stage of the training, up to a constant K in the exponential term.
20

---

Published as a conference paper at ICLR 2025
We have for all E [Ti]. at timepoint every entry (i,j) is in the initial phase, and the absolute value of the (i,j)-th entry satisfies
Lemma D.2 (Lower Bounded Initial Growth). Let Ti
log 2
(wa5()) 2 [00) exp [ntl(a= + uz)xx²
(D.22)
and Wie (t) > O.
Proof. Similar the proof of Lemma D.I, may just assume wi,((0) 0.
Moreover, we also use the claim as an inductive hypothesis and prove it by induction. Since here the inductive hypothesis states that every entry is in the initial phase, we have
(t. + (t) & 42
(D.23)
We have Wi.s (t + (t) 2 naa + w,) , (0) exp |nt(a; + +a,)x - - 2Pmadd32
(D.24)
2 2Pmmad33
(0) exp (tta: + 0z)xx* '1 1/aa + +
w/ exp (nt(a aj 1|
(D.25)
From Assumption D.4, we have
2P'mad32
Why (0) exp Int(a; a,)xx=1] $ 2PPmadd22
(D.26)
(1 x n(a; + aa))
(D.27)
Moreover, notice that when >> 1.1, for any x < 0.1, we have N -1/22 + 1 2 e
Since
Assumption D.2 ensured that 1) S 10(arta,)' we have
Wi. (t (++ 1) 2 Wa, (ay(() 18 n/2 aj)
(D.28)
2 W exp (naa + a1))
(D.29)
( exp [nt( + 1)(a: + a))r
(D.30)
Finally, from Lemma D.I, we have when
((y(() < (0) exp (ni(a + az)ss)
(D.31)
< Bw exp yynT)))
(D.32)
PBw,
(D.33)
which confirms that every entry (i.j) stays in the initial phase before time Tj.
Notice that the time bound Lemma D.2 is a uniform one which applies to all entries. For the major entries, we might want to consider a finer bound of the time that it leaves the initial phase. This can be proved by essentially repeating the same proof idea Lemma D.2.
Lemma D.3 (Lower Bounded Initial Growth for Diagonal Entries). Consider an diagonal entry
( at timepoint the entry (i, is in the initial phase, and the absolute value of the (i, i)-th entry satisfies
(1, Let T
log "5,1
2na. We have for all t
Prew
(t,((t) (0, exp (2ntaaı
(D.34)
We omit the proof of Lemma D.3 since it is almost identical to the proof of Lemma D.2, only with replacing ya by a, and Bw by (0).
Next, characterize the behavior of one diagonal entry after it leaves the initial phase.
21

---

Published as a conference paper at ICLR 2025
Lemma D.4 (Lower Bounded After-Initial Growth for Diagonal Entries). Consider diagonal entry (i, 1). If at time to we have [Wo, 2 Pßw, and for a X E (P,ww, - K K)1,, before time T(A) we have wt + to) < A for all & E [T(A),, then we have
wi + (n) 2 (1, exp [2n]aae] - -)x
(D.35)
Moreover, (0), Wh (Lo), (to +1) 2 0.
Proof. Notice that since W = TUT is PSD matrix, its diagonal entries are always non-negative, this ensures that (O), Wi, (1n), (1, + +1 0.
For the time after and before to T(A), we use induction prove the claim, with the claim itself as the inductive hypothesis. It clearly holds when =
Notice that when (t'(tt) < X, we have
(t)(t)- (t') = (t, wr,) 2 (()(;,((()(1 -)).
(D.36)
Thus we have
(4 + 1 1)-11;i- ())
(D.37)
2 27a,(1 - (to) exp [ntt(a, + a,) ->xx 2P 2Pryadd3 w?
(D.38)
2 2P
(t)) exp (2nta (1 1)1= 2naa(1 - wi exp [2nta - A)x-
(D.39)
Since A < 1 - K-1, and (1,) 2 2Bw 2 from Assumption D.4, we have 2P
< 22Pmaad3'
(D.40)
(i,s(tu) exp 2ntta 1)x-11
(1-x $ 2K -1 (1- x 1122) x1/2) nja
(D.41)
2 2 - 1122) na;(1 - A).
(D.42)
Moreover, since Assumption D.2 ensured that S 2ka((1-)) S (1=; using the fact that
> 1.1 then x -122g + 2 xxxxx for any x < 0.1, we can get
wt, + 2 w;, (f) + (t) 1:= 2/-, 11] -1I (11A1) 2 wi, (t) exp (2na,x '(1
(D.43)
(D.44)
> (1o) exp [2n(1 1)x= (1 2)]
(D.45)
Next, we provide an uniform upper bound (over time) of the diagonal entries. Remember that we mentioned in the gradient flow case, the diagonal term stops evolving when it reaches 1. In the discrete case, since the step size is not infinitesimal, Lemma D.5 shows that it can actually exceed little bit but not much since the step size is small.
Lemma D.5 (Upper Bounded Diagonal Entry). For any diagonal entry (i,1) and any time I, 0 S (t. < 2K-!
Proof. First notice that since W(t) is PSD, its diagonal entry (t)) should always be non-negative, thus (t) 2 is always satisfied. In the following we prove (!) < 1 + 2K -1
We use induction to prove this claim. The inductive hypothesis is the claim it self. It is obviously satisfied at initialization. In the following we assume the claim is satisfied timepoint and prove it for timepoint + 1. Notice that since K < 10, we have 1 + K-1 < 1 2K < 2.
Notice that by Assertion D. and Assumption D.4,
(K 11-
Nt,(t)) < 2 <
K2yd32 a S K-1 a;;
(D.46)
22

---

Published as a conference paper at ICLR 2025
If (c,,(t) 1- we have
Gi,(t) - St,(t) 2a, wa,1) s -2a,K-
(D.47)
Therefore,
Wia(t +1 = (),((1) + (t, (1) No.
(D.48)
S w;,((t) - a,
(D.49)
$ (t,((t
(D.50)
$ 1 11
(D.51)
Moreover, since w;,((t) S 1 2K < 2, we have
1
(t), (t, (t) + [N,, & 4a, + 4a; + K-K a, S 9ya S
(D.52)
Kn
When (1) S K-1, using eq. (D.52), we have
(t + 1) < wi, (t) + 7 (lG,, (f) + (() + (NN..(0)) < =1+2K
(D.53)
The above results together shows that w., (t + 1) S 1K
Corollary D.2 (Upper Bounded Diagonal Update). For any diagonal entry (i,1) and any time (1,,(++1) w.) < K-1
Corollary D.2 is a direct consequence of Lemma D.5 (and actually proved Corollary D.2 in the proof of Lemma D.5).
The next lemma lower bounds the final value of diagonal entries. Together with Lemma D.5 we show that in the terminal stage of training the diagonal entries oscillate around by the amplitude not exceeding -
Lemma D.6. Consider a diagonal entry (i, 1). If at time ln we have we,((n) 2 1 - 2K Ihen for all 2 to we have (t,) 2 2K-1
Proof. We use an induction here. The inductive hypothesis is the claim itself. This obviously holds when L' lo. We assume (t') 2 1 2K at timepoint L' and prove the claim for + 1.
If (l')sll) < 1 K-1, then from Lemma D.4 we know
( + 1) 2 (t') 2 1 - 2K
(D.54)
If (ww,((t)) 11 K-1, then from Corollary D.2 we have
Wi,a(t' + +1) 2 w;, (t') K-1 1-2K11
(D.55)
Now, we are ready to prove Assertion D.1 by considering the suppression. We first prove a lemma that upper bounds the absolute value the minor entries after corresponding major entry becomes significant.
Lemma D.7 (Suppression). Consider an off-diagonal entry (i,j) where > j. If there exists a time to such that (to) 0.8, then for any t' 2 to we have
|Ws,s(t') (t') | S max (1wis (to)),
(D.56)
Proof. Since > 10. from Lemma D.6 and Lemma D.4 we know w.,,(t') > 0.8 for all t' 2 to.
In this proof, we use an induction with the inductive hypothesis being the claim itself, i.e., we assume the claim is true at timepoint t' and prove for t' + 1. The claim obviously holds for t' to. Since in this proof only the absolute value of No., WLOG may assume that (),((')) 0.
23

---

Published as a conference paper at ICLR 2025
If (t'(tt) then we have proved the claim. In the following we may assume wi,((') 2
We have
1
Gi,(t)-Sunn s (t) + +a,) ;
wi (t),(t) + +a,)
(D.57)
(t)(t)(a; + - [0.4(3a; + a,)l
(D.58)
3
5 wh (t')a:
= w (t))aj
(D.59)
wwa
(D.60)
where in (i) we use Assumption D.5.
Thus we have
No,(t) s -
(t))
(D.61)
--C wa + 2Pyod:3
(D.62)
0,
(D.63)
where is is from Assumption D.4 and Assumption D.5. This confirms that wt'(t'+1) < wi,(t')
max { [Wr,((oo), w).
Next, we prove wiy(( (I' +1) 2 max {/or, (4n) w) Notice that Lemma D.5 stated that [wa,]] 2.
Notice that we also have (t') - K thus from Assumption D.2,
[t),(+') (t)) + [t,s < 10 uy (t') + 2
(D.64)
s 10|w, mo (t') + 22dd
(D.65)
9K'
10 2 (t') + 2w
(D.66)
9Kn
|wo + w
(D.67)
2n7
We have
Wi,(t' +1) 2 wi)(t') - ((G + (t)(((t))++++
(D.68)
2 (nGGss ((') + [S,,((') + +NN,s(())
(D.69)
((wn)(()) (t')l + w)
(D.70)
maxf/w: ( (t))),
(D.71)
With all the lemmas proved above, we are now ready to prove Assertion D.1.
Lemma D.S (Assertion D.I). For all E N, i + j, then the entry (i. stays in the initial phase for all time.
Proof. Notice that since W is symmetric, we only need prove the claim for j. Moreover, From Lemma D.7, we only need to prove that there exists a timepoint t", such that wx,((t) 2 0.8, and [Ws, (t*) & PBW..
D.P by Lemma D.3, we have (t.) 2 Pßw. By Lemma D.3 and Lemma D.4, have for any to such that wr,,(t) S X, where 0.85,
log wi, (10
Let to
2maa
(t) 2 (t,)(to) exp [0.3n(t - to)a,k=-1 2 exwexp exp [0.3n(t
(D.72)
to)a,x-1)
(D.73)
24

---

Published as a conference paper at ICLR 2025
Let t' be the first time that (t')(t') arrives above 0.8. Let ta = min
log Piu
+ to.
1
2 to. If
0.3ma:
klg 0.3.
1a 1', we have (1*) 2 0.8. If
0.3ra, lo. have
(t,((t)) 2 wr (0) exp (0.3 a, 1)
(D.74)
0.8
2 PBw exp
log P3w
(D.75)
= 0.8.
(D.76)
Moreover, from Lemma D.I and Assumption D.5, we have
[tt,(tt)| < [0) exp [nt k<(;; + az)]
(D.77)
S BB exp
n log 0.8
P
+ log
PBWw wi,;(0)
x 0i + aj
(D.78)
0.15
20,
a, + aj
S BB exp
10k² log
PBw + log P3
(D.79)
2a;
S Bw cxp [log(P))
(D.80)
PwB.
(D.81)
The claim is thus proved by combining the above bounds on [t")((t)) and wt,((t) with Lemma D.7.
E ADDITIONAL DISCUSSIONS
In this section, we further discuss the findings and theoretical predictions presented this paper.
E. MULTIPLE DESCENTS
In Fig. 5, we verified our theoretical predictions of the Swing-by Dynamics through an experiment of an s 2 example. However, our theory, there can be multiple growth / suppression stages when s > 2, which should give us a multiple descent-like curve. We note here that based on the conditions given in App. D.1, is indeed possible see multiple descent but only with a subtle choice of the signal strengths (44) and under specific initialization conditions.
In Fig. and 9, illustrate two settings where the loss curves exhibit epoch- wise triple and quadru- ple descent. In both settings we use symmetric 2-layer linear model, same as the model used in Sec. 4.2. Note that tuned initialization random seed to generate these results. Moreover, since the time scale of each descent vary, we use a log scale for the number of epochs to make the results more apparent.
worth noting that in Fig. 8 and 9, each major entry starts to grow only after the corresponding minor entry is suppressed (for example in Fig. 8, W2,2 starts to grow after W2,3 is suppressed, W1.2 starts to decay after W2,2 is close to and W1,1 starts to grow after W1,2 is suppressed), and each ascending / descending stage of loss curve aligns well with stage of growth / suppression of the minor and major entries. These correspondence match exactly with our theoretical prediction and shows the correctness and preciseness of our theory.
E.2 THE BREAKDOWN OF THE INITIALIZATION ASSUMPTION
In Sec. 3.3, we mentioned that Swing-by Dynamics seems be more significant when the dimen- sionality of the dataset is low, and in Sec. 4.2.2, we attributed the reason of it to the fact that when the dimensionality of the dataset small, it's easier have more minor entries initialized positive, which lead an illusion of learning in the minor entry growth stage, whose later suppression leads to the non-monotonic output trajectory behavior of Swing-by Dynamics.
In this section, we note that, another reason for the Swing-by by Dynamics to be less significant is that Assumption D.4 breaks down when the dimension is high, if we use standard Gaussian initialization
25

---

Published as a conference paper at ICLR 2025
1.0
10
-7
-3.5
0.8
0.8
Loss
W1.1
-25
W1.1
Value of W
0.6
22.2
0.6
W9.2
-2.
W9.3
Valee
W4.4
21.2
-15
0.4
W1.2
3
3
32.3
-10
W1.3
-2
W
-0.
0.0
10a
10
10a
10a
10*
100
10a
10a
102
10*
epochs
Number epochs
Figure 8: An illustration of epoch-wise triple descent of the symmetric 2-layer linear model. The dataset has dimensionality and number of in- formative directions = 8 = 3, signal strength values f (1.0, 1.5, 2.2), and noise values (0.05, 0.05, 0.05).
Figure 9: An illustration of epoch-wise quadruple descent with the symmetric 2-layer linear model. The dataset has dimensionality and number of informative directions d = 4, signal strength values μ (1.0, 1.5, 2.2, 2.7), and noise values o = (0.05, 0.05, 0.05, 0.5).
initialize the model weights. Specifically, in Assumption D.4, we require that all entries of W are initialized around relatively small value w, which indicates that there is no huge difference between the magnitude of the initialization of major entries and minor entries.
However, notice that W TUT and thus Wa,j (u i, uz). where u, R2d is the i-th row of U. If we use Gaussian distribution to initialize U, i.c. u, N (0, r²)), where is small real number, then we have the expectation of Wi.j be
Eun, ( d
i
(E.1)
i i=
which highlights the different between major entries and minor entries in initialization when d' is large (and notice that d' is lower bounded by d, which is the dataset dimensionality). Moreover. when d' is small, the variance of w.,, will be large, so there is greater chance for them be away from O.
E.3 FAILURE MODES
breakdown the assumptions in App. D.1 can also lead to the model converging to "wrong" solutions that do not fully generalize OOD. For example, a minor entry happens be initialized too large (breaking the Assumption D.4), and / or the corresponding signal strength distinction is not large enough (breaking the Assumption D.5), then suppressed until it grows to significant value, which can, in turn, lead to too strong suppression on the corresponding major entry. In this case, a major entry might be suppressed to 0 (or least, leave the initial phase from below) before it starts to grow, and thus never has chance to grow. state that only learns compositionally
possible that the minor entry is not
This case corresponds to the model being "trapped" generalize to combination of certain (but not all) concepts.
In Fig. 10, we exhibit generalization. Notice how the loss value converges to a non-zero value and the major entry W1,1 is suppressed the very beginning and never grows. Additionally, the output trajectory trapped at point that combines only two directions, missing the third direction.
of such failure mode where the model fails to fully achieve OOD
E.4 FUTURE DIRECTIONS
We note that, current characterization of the model learning dynamics relies on the critical assump- tions in App. D.I. Although those assumptions reasonable and common in practice, the model
26

---

Published as a conference paper at ICLR 2025
Loss
value of w
W1.3
Loss
21.2
W1.3
2
W5.3
10a
10a
10²
10
10
Number of epachs
Figure 10: An illustration of failure case where the model doesn't successfully generalize OOD. The dataset has = = 3. μ (0.7,1.7,3 1.7,3) and o = (0.05, 0.05, 0.05). Left: The evolution of values of Jacobian and the OOD loss evaluated at test point ã; Right: The output trajectory (orange curve).
behavior still shows some regularity when those assumptions breakdown. We have discussed some of the possible consequences when one of those assumptions breakdown above, but more an intu- important future direction is to systematically
itive way, instead of systematic way. Therefore, characterize what will happen beyond the assumptions given in App. D.I. Among which, one spe- cific and very important topic is the failure modes, i.e. under what conditions the model fails to generalize OOD.
Another important direction is to generalize current analysis to more complex models, such as deep linear networks, two-layer ReLU networks, models in the NTK regime. The key point of current analysis of the 2-layer symmetric model is to correctly slice the learning dynamics of each entry of the Jacobian into multiple stages, such that each stage, the learning dynamics dominated by rather simple dynamics.
Currently, since the model is 2-layer and without bias terms, there only first-order and second- order terms in the learning dynamics. However, if we consider deeper models, there might be higher-order terms in the dynamics, and it is important identify and simplify the effect these higher-order interactions in order to make the problem tractable.
For ReLU networks, it is known that there will be "early-alignment" stage when trained on linear-separable data (Maennel et al., 2018; Min et al., 2023), where each neuron converge to a fixed direction, and make the model behaves like a linear model. We claim that investigating the early-alignment of 2-layer ReLU networks on the SIM task can be the starting point of theoretically characterizing the behavior of ReLU networks the SIM task.
F ADDITIONAL SIM EXPERIMENT DETAILS AND RESULTS
In this section, we present the results of SIM experiments under different settings, including linear and non-linear models. The consistent behavior observed across these settings confirms the univer- sality of our findings and explanations.
F.1 EXPERIMENT DETAILS
In all SIM experiments, including those presented in main paper and in appendix, the number of training samples in each Gaussian cluster is 5000. We use MLP models with either linear activations ReLU activations, and all the models size of 128 and a learning rate of 0.1 for 40 epochs. Unless otherwise specified, the dimensionality of all data points is d 64, and the hidden layer dimensionality of the models is also 64 by default.
trained using stochastic gradient descent with batch
27

---

Published as a conference paper at ICLR 2025
It is important to note that in theory, assumed that all training clusters and the test point are aligned with the standard coordinate. However, our experiments, in order to make the results more universal and general, add random rotation to all the train / test points.
F.2 ADDITIONAL EXPERIMENT RESULTS
Fig. 11 and Fig. 12 repeat the learning order experiments described in Sec. 3.1, using 2-layer model with and without ReLU activation, respectively. It is easy to see that despite showing more non-regular curves, in multi-layer models the overall trends described in Sec. 3.1 and Sec. 3.2 preserved.
== (1, Trajectory Train Test
2.0
3.0
(2, Trajectory
Peontory/TTaiin Train Test
(3, Trajectory Train Test
(0.05, 0.05)
2.0
1.0
(0.1, 0,05)
0.05)
σ (1, 0.05)
1.0
(2, 0.05)
Train
Tast
0.0
0.0
0.0
2.0
0.0
1.0
2.0
3.0
Figure 11: Output trajectory of 2-layer models with linear activations. The number of informa- tive directions in the dataset is s = 2. Left: μ12 (1,2) with varied o's; Right: σ.2 = (0.05, 0.05) with varied pts.
μl (1,, 2) Trajectory Train Test
Tapectory/Train///
2.0
3.0
μ (2, Trajectory/ Train Test
(3, 2) Trajectory Train Test
0.05)
2.0
1.0
(0.1. 0.05)
(0.5, 0.05)
1.0
g (2, 0.05)
Train
Test
0.0
0.0
0.0
1.0
2.0
2.0
3.0
Figure 12: Output trajectory of 2-layer models with ReLU activations. The number of in- formative directions the dataset = Left: μ12 (1,2) with varied o values; Right: 2 (0.05, 0.05) with varied a values.
Fig. 14, we present the output trajectory for two settings that exhibit significant Swing-by and Fig. 14 the corresponding loss curve. Specifically, the dataset has dimensionality of d = 3, and is not randomly rotated. The models used have 3 layers, 3 hidden dimensions and linear activations. Comparing these results with the curve presented in Fig. 11. Fig. 12 and in Fig. is evident that models with more layers and fewer input dimensions are easier to have Swing-by, which confirms theoretical prediction in Sec. 4.2.2.
28

---

Published as a conference paper at ICLR 2025
4.0
2.0
1.0
2.0
Train
Train
Test
Test
0.0
0.0
0.0
1.0
2.0
0.0
2.0
4,0
Figure 13: Output trajectory of 3-layer models with linear activations and 3 hidden dimensions. The dataset has dimensionality 3, number of informative directions and variance 0:2 (0.05, 0.05). Left: μ12 (1, 2); Right: μ42 = (2, 4).
Test loss
Small μ
5
Large μ
0
0
10
20
30
40
Training time
Figure 14: The loss curve of corresponding models in Fig. 13. Small ic μ12 (1,2); Large μc:
μ. = (2,
F.3 FURTHER EXPERIMENTS THE ROLE OF OUT-OF-DISTTIBUTION
In Sec. 4.2.2, briefly discussed that Swing-by must be OOD phenomenon: the test point should be sufficiently away from the training clusters. In this subsection, we verify that this condition is satisfied in our SIM experiments.
In Fig. 15, we repeat the experiment of Fig. (c), but addition the output curve, we also plot each training point. It is clear that there is no training point that is close to the test point.
In order to further ensure the separation of the training clusters with the testing point, in Fig. 16, we again repeat the experiment of Fig. (c), but we make a truncation of the training distribution: force that no training point can lie within the distance of min s Mp of the testing point It evident that only few training samples are discarded and this truncation has little impact on the dynamics.
F.4 FURTHER EXPERIMENTS ON THE TRAJECTORY TRAINING SET
In order make it clearer to compare the model behavior training and test point, here in Fig. 17 and Fig. 18 we repeat the setting of Fig. 5, with additional information provided: 1) in Fig. 17, we added the training loss curve In Fig. 18 we added trajectories, corresponding to the model output given the training cluster means input. We omit the training cluster centered at the origin since the 2-layer linear model always output for this point.
TThis repeatedly sampling training points and discarding those within the specified distance of the testing point, until the training reaches the desired size.
29

---

Published as a conference paper at ICLR 2025
2.0
2.0
1.0
1.0
Train
Train Test
Test
0.0
0.0
0.0
1.0
2.0
0.0
1.0
2.0
Figure 15: Reproducing Fig. 2 (c), with train- ing points plotted.
Figure 16: Reproducing Fig. (c), with train- ing distribution truncated.
1.5
1.0
2.4
1.2
W1,1
1.0
1.0
Train
W1.1
0.8
Test
Loss
best
train
0.4
0.2
0.0
0.0
250
500
750
10000 1250 1500 1750 2000
0.0
Mumber of spoths
Figure Reproducing Fig. 5 left, with training loss plotted.
Figure 18: Reproducing Fig. 5 right, with the output trajectory of training set.
As shown in Fig. 17, the training loss monotonically decreases, exhibiting different behavior com- pared to the non-monotonic test loss curve. This is expected because we used a small learning rate, and it is well-known that, under such conditions, the training loss must decay monotonically. In Fig. 18, the two curves of training trajectories both exhibit non-regular behaviors, but at different stages of the training. This observation aligns with our analysis of the multi-stage behavior of the learning.
G DIFFUSION MODEL EXPERIMENTS
We describe experimental details for the diffusion model experiments. We largely follow Park al. (2024) in these experiments.
G.1 SYNTHETIC DATA
Fig. 19 illustrates the DGP. We borrow part of the compositional data generating process (DGP) introduced by in Park et al. (2024). The DGP generates set images circles based on the concept variables color=(red, blue) and size=(big, small]. Each concept variable can be selected as composed to yield four classes 00, 01, 10, 11 respectively corresponding to (zed, b g), (red, small), (blue, big), (blue, small). 1)). Here, class average pixels values of red
30

---

Published as a conference paper at ICLR 2025
01
01
Sire
Size
Color
Color
0.
10
Data Generation process with different concept signals. Figure from (Park et al., 2024). The data generating processed used to train the diffusion model, where we can control the strength of the two concept signals independently. Left: data distribution with stronger concept signal in the color dimension. Right: A data distribution with stronger concept signal in size.
Figure 19:
and blue colors will control the concept signal for color and the difference between small 1 and big. In Fig. 6, we fix the big circle's iss  diameter to 70% of the image and the small circle's diameter to 30% of the image. We then adjust the absolute difference between the blue color and red color from 0.2 (very similar colors) to 0.7 (very different colors). The DGP randomizes the location of the circle, the background color and adds some noise to avoid having a very narrow data distribution. Please refer to Park al. (2024) for further detail.
In Fig. 19, show different data distributions, one with big color concept signal and one with a big size concept signal.
G.2 MODEL TRAINING
We train a conditional diffusion model on the synthetic data defined above. specific, we train variational diffusion model (Kingma et al., 2021) to generate 3 x 32 x 32 images conditioned on 4-dimensional vector where the first element of the vector specifies the size of the circle and the others specifies the RGB colors.
Model Architecture We use a conditional U-Net (Ronneberger al., 2015) with hidden dimen- sions (64, 128, 256] before each downsampling layer and two ResNet (He et al., 2015) layers in each level. The conditioning vector is first transformed into the same dimensions as the hidden dimen- sions using 2-layer MLP and are added to the representation after every downsampling layer. The U-Net has a self attention layer (Dosovitskiy et al., 2021) I) its bottleneck. used LayerNorm (Ba et al., 2016) for normalization layers and GELU (Hendrycks & Gimpel, 2023) activations.
Diffusion We use learned linear noise schedule for the diffusion process as defined in Kingma et al. (2021), initialized with "max 10, "min -5. We assume data noise of 10-3 Variational diffusion models do not require fixing the number of diffusion steps at training time, but use 100 steps for generation at inference time.
Training We train our model with the Adam optimizer (Loshchilov Hutter, 2019) with learn- ing rate 1 x 10-3 and weight decay 0.01. We use batch size of 128 and train for 20k steps.
G.3 EVALUATION
We evaluate the concept space representation of the generated output image using a trained classifier. Since we have the ground truth DGP, we used large amount of data train a perfect classifier. We used a U-Net backbone followed by a max pooling layer and a MLP classifier to classify each concept variable color and size. We train this classifier for 10k steps and achieve 100% accuracy on held out test set. We average over 32 generated images and model run seeds to get the ensemble average concept space representation.
31

---

Published as a conference paper at ICLR 2025
The concept space MSE in Fig. 6 (b) is simply calculated the MSE distance in the concept space defined in Park et al. (2024). The concept learning speed dC/dr) is quantified by estimating the movement speed in the same concept space by finite difference method.
32


Reviews:
[{"review_id": "wKdyl5BbLC", "reviewer": "Reviewer_DJMW", "summary": "The paper sets out to explain previously reported characteristics of the learning dynamics of text-conditioned diffusion models with respect to  their ability to compose concepts to generalize to unseen combinations. Specifically, the paper aims to study why models learn concepts in a specific order and how the structure of the data influence a model's learning speed. To this end, the authors propose a simple reconstruction task (which they term structured identity mapping, SIM) in which an MLP should reconstruct its inputs drawn from a mixture of Gausssians. The authors argue that this task allows them to study and explain the key empirical observations on compositionally generalizing diffusion models and lead to the novel characterization of a non-monotonic learning dynamic.", "questions": "# Questions\n- LL86: \"The model output is then passed through a classifier which produces a vector indicating how accurately the corresponding concepts are generated (e.g. a generated image of big blue triangle might be classified as (0.8, 0.1, 0.9)). In this way, the process of generation becomes a vector mapping, and a good generator essentially performs as an identity mapping in the concept space.\" In this setup, the vector mapping is $c \\circ g$, comprising the generator $g$ _and_ classifier $c$. While it is clear that a good generator should be the identity, the role of the classifier also has to be analyzed. Can we be sure a priori, that $c$ is an identity mapping?\n- How is training done? Is a fixed number of points sampled in the beginning to be used as a training set, or are points instead drawn for each batch?\n- LL193: Why is $\\mu_k$ equated with signal strength? Intuitively, instead of the absolute value of $\\mu_k$, the distance between clusters should be more meaningful, which might be high for an individual cluster even if $\\mu_k$ is low.\n    - it is also not clear why this should matter to the model. The model could simply normalize each input dimension such that inputs are always balanced, e.g., in the task of Fig. 2b\n- App. B: This observation mainly seems to imply that test loss _increases with the distance from the training set_, which, again does not say anything about OOD generalization or compositional generalization, and instead is a clear effect of the increasing probability density of the training set.\n    - Also, why is the loss truncated instead of normalized?\n\n# Minor Suggestions\n- Fig. 1 and LL73 show that the data is clustered around nodes of the hypercube, yet LL150 and LL160 explain that each cluster has a different $\\mu_p$ (distance from the origin). How do these statements fit together?\n- L163 “There is also optionally a cluster centered at 0.” Is this in addition to the $s$ clusters, or is this just specifying that $\\mu_p$ might be 0?\n- LL170: \"We evaluate the model at a Gaussian cluster centered at the point that combines the cluster means of all training clusters.\" What does this mean? Say the test cluster is $\\boldsymbol x_k^\\text{test} \\sim \\mathcal N(\\boldsymbol \\mu_\\text{test}, \\boldsymbol \\sigma_\\text{test}^2)$, does this just mean that $\\boldsymbol \\mu_\\text{test}$ is the mean of all $\\mu_p \\boldsymbol 1_p$? \n- §3.1 and Fig. 2: What are the markers at each optimization timepoint? The model output for the given input $\\boldsymbol x = \\boldsymbol \\mu$?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This work studies a relevant problem: The learning dynamics of compositional generalization in diffusion models are important to understand how models can learn in a sample-efficient manner, how generalization can be achieved, or how training data should be curated, to name a few ways insights could be impactful.\n\nThe paper is easy to follow for the most parts and builds on a prior line of work in this area.\n\nThe non-monotonic training dynamics of a symmetric two-layer linear model are well explained by the theory and offer a novel (as far as I can tell) characterization of the training behavior of basic models. The theoretical training dynamics outlined in §4 and especially §4.2 and Fig. 5 model the empirical observations from §3 and Fig. 2 well.", "weaknesses": "In summary, I find the paper misses the mark, as the SIM task is, as far as I understand, a poor setting to study compositional generalization, and the insights on this simple task translate poorly to the training of a diffusion model, even for the simple toy setting that is used (which itself is approximating the compositional generalization of text-conditioned diffusion models that this paper aims to study). I will elucidate the issues I see with the setting and results below.\n\nAs it stands, I cannot recommend acceptance of this work; however, I could see how the results on the learning dynamics of symmetric linear models from §4 might be interesting in their own right, without considering any (as I understand, faulty) connection to compositional generalization. Maybe a reconsideration of the scope of the paper and a reframing of the results without overclaiming meaningful insights regarding the training dynamics of text-conditioned diffusion models could be an interseting contribution on its own.\n\n# The SIM Setting\nAs far as I understand, previous works took the \"concept space\" as an abstracted view of the training and test data in order to study training trajectories of diffusion models. In the SIM task, this concept space is instead interpreted directly as the input and output space of a model. While I can see that this choice is motivated by the intuition in LL89, \"a good generator essentially performs as an identity mapping in the concept space\", this simplification of the setting is not explicitly mentioned, e.g., in the introduction, and the limitations this simplification imposes on the applicability and transferability of the results are not discussed.\n\nFurther, the training distribution in the SIM task is a mixture of Gaussians. In general, I find that this setup is not explained very clearly, see minor issues below. The central issue I see is that this setup entirely misses the point of (compositional) generalization. First, by introducing the training clusters as Gaussians, any point in $\\mathbb R^d$ has probability $> 0$. So it is simply not true that any \"[test] point is outside of the training distribution—not just the training data, necessitating out-of-distribution generalization.\" (LL173). Second, since test loss is only tracked on an individual point, not a test distribution, the experiments cannot even consider distribution shifts. Overall, describing this setting as requiring \"generalization\" to an \"OOD test point\" is misleading.\n\nIf we understand that in this setting, any test point has non-zero probablity in the training set, the conclusions from §3 are mostly unsurprising. Consider the probability density of a given point in the grids in Fig. 2 belonging to the training set (in fact, I encourage the authors to include this information in the figure). I expect that for Fig. 2a, we will see that the trajectories are skewed towards higher-density regions, which roughly speaking can be interpreted as the model reducing the risk of a certain output. Similarly, we can understand the takeaway from Fig. 2b, the \"terminal slow down\" differing for different values of $\\boldsymbol \\sigma$ as an effect of the probability density. For $\\boldsymbol \\sigma = (2, 0.05)$, the overall density at point $(1, 2)$ is much higher than for $\\boldsymbol \\sigma = (0.05, 0.05)$, making it much less likely that this point (or points close to it) are sampled, so that the model takes much longer to learn it.\n\nWhile the \"transient memorization\" is an intersting behavior of even simple models, and the author's explanation of this phenomenon seems interesting, I find it hard to justify translating any insights from this setting to a compositional generalization regime.\n\n# Transferring Results to Diffusion Models\n§5 is very bare-bones, to the point that it is unclear how well the observations from the SIM setting actually translate here. E.g., there doesn't seem to be as much of a \"terminal slow down\" in Fig. 6 as was shown in Fig. 2a/b. The \"transient memorization\" also seems much less pronounced.\n\nAdditionally, it is unclear _why_ results form the SIM setting should translate here, as in this case the test point is truly out of distribution, and, if I understand the training setting correctly, the trianing set only contains discrete values for each factor.", "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": "UjsCXYzckz", "reviewer": "Reviewer_caEk", "summary": "This paper investigates the dynamics of neural networks in achieving compositional generalization. The authors propose a Structured Identity Mapping (SIM) task, where models are trained on a Gaussian mixture dataset organized around concept-centric clusters. They find that concept strength and diversity strongly influence the speed of convergence and the ability to generalize to compositional out-of-distribution (OOD) test points. They also observe a phenomenon called \"transient memorization,\" where models initially memorize the training distribution but eventually reorient towards OOD generalization with extended training. Theoretical analyses on simple linear models substantiate these findings, highlighting the role of signal strength and diversity in learning order and the occurrence of non-monotonic loss behavior. Finally, the authors validate these findings by training diffusion models, observing similar patterns in generalization dynamics.", "questions": "The paper effectively accomplishes its aims, though a few points could be noted for improvement:\n1. **Limited Theoretical Scope**: While the theoretical framework successfully supports the empirical observations in more complex models, it is based on simple linear models.  I acknowledge the difficulty in analyzing more complex neural networks theoretically. Nonetheless, for future work, extending the analysis to objects like kernel methods, particularly the Neural Tangent Kernel (NTK) [1], could provide a deeper understanding as NTK has been shown to capture certain behaviors of neural networks during training [2]. \n2. **Contextualization of Findings**: Some insights, such as the speed and order of generalization, have been observed in different forms in previous OOD works [3, 4, 5]. Further contextualizing the findings within these works could enhance the paper’s relevance and depth.\n\n[1] [Jacot 2018] https://arxiv.org/abs/1806.07572\n\n[2] [Lee 2020] https://arxiv.org/abs/2007.15801\n\n[3] [Nagarajan 2020] https://arxiv.org/abs/2010.15775v3\n\n[4] [Arjovsky 2019] https://arxiv.org/abs/1907.02893\n\n[5] [Rosenfeld 2020] https://arxiv.org/abs/2010.05761", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is clearly written, with a logical flow that makes each insight and conclusion easy to follow.\n2. The simplicity of the problem setup enhances the clarity and robustness of both the empirical observations and the theoretical contributions.\n3. The diffusion model results are compelling, mirroring the behavior observed in simpler settings and offering explanations for phenomena noted in prior work.", "weaknesses": "See the questions section for more information.", "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": "qa1dabYrKM", "reviewer": "Reviewer_pvo7", "summary": "The paper introduces the structured identity mapping (SIM) task, where a network is trained to learn the identity mapping from inputs which are sampled from largely disjoint Gaussian distributions, and evaluated on a test point defined as the average of their means. By analyzing learning dynamics of one-layer and symmetric two-layer linear networks on this task, the paper shows that several existing empirical observations can be modeled - higher SNR features are learnt earlier and faster, an epoch-wise double-descent phenomenon (termed Transient Memorization) occurs on the test point, and that learning slows down exponentially with time.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Paper formalizes a proxy task (SIM) to model the learning dynamics of compositional generalization, and shows that real-world text-conditioned diffusion models exhibit similar behavior on a specific task. \n- Theoretical setup is clearly explained and theoretical conclusions are also well-elaborated. Limitations of theoretical results on the one-layer and symmetric two-layer linear models are also clearly discussed.", "weaknesses": "- It is not at all apparent that $\\hat{x}$ is \"outside of the training distribution\", since $p(x_k^{(p)} = \\hat{x}) > 0$ for training data $x_k^{(p)}$, despite what is claimed on L174. Why not choose sigma such that $\\sigma_p \\geq 0$ while keeping $\\sigma_{q \\neq p} = 0$? Furthermore components of $\\sigma$ seems to be as large as $2$ in Fig 1., with $\\mu$ ranging from $0 - 2$. In such cases, the converse seems to hold -- that $x_k^{(p)}$ is very much in-distribution of the training data.\n- As a result, the introduced terminology \"Transient Memorization\" does not seem to be different from epoch-wise double descent. The paper stresses that the main difference is OOD vs in-distribution test loss (L263-265), but (a) as mentioned above the test sample considered are in fact in-distribution, and (b) even for double descent, it is actually essential for testing samples to differ significantly from the training ones as pointed out by [1]. \n-  The proposed SIM task is claimed to be a \"further abstraction of the 'concept space' previously explored\" (L531), however it appears instead to be an (over)-simplification of previous investigations. The motivation of the SIM model and how it relates with more general real-world settings are also unclear (for instance, why the proposed auto-encoding loss?).\n- The paper's main empirical results section (Page 10) appear very underdeveloped and all important experiment details are left to the Appendix. Figure 6 is also not clearly explained, and different plots are distinguished based on \"difference of class mean pixel values\". Without reading App G, it is almost impossible to interpret what these results and experiments are doing. $\\Delta$ Color is not defined in App G either.\n- Paper concludes from Fig. 2 that \"deceleration is not determined by ... the loss value ... but more depends on the data and training time\" (L215). This is not self-evident, since (1) the training loss is not even plotted in Fig 2, and (2) regions with denser arrows are clearly also regions of lower test loss.\n- The paper also claims that \"there is a timescale determined by the training data such that if the model does not achieve OOD generalization within that period, significantly more computation will be required for the model to achieve it.\" (L238-240), which purportedly can be observed in Fig 2(b). However, the figure simply shows that configurations with lower $\\sigma$ values take longer to correctly predict the test point. This conclusion seems self-fulfilling, since when $\\sigma$ is larger, the density of training samples around $\\hat{x}$ will be much higher.\n- Minor: L534-L535 incomplete sentence - \"We make a comprehensive.\"\n- The work's primary area is stated as interpretability / explainability in AI, but the contributions to these areas are unclear. The bulk of the theoretical contributions seem to be instead exploring learning dynamics for one-layer/deep-linear models, which I have insufficient experience to evaluate with regards to recent literature. \n\n[1]  Double Descent Demystified: Identifying, Interpreting & Ablating the Sources of a Deep Learning Puzzle, 2023.", "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": "j13pcjRGj3", "reviewer": "Reviewer_c6mH", "summary": "This paper presents a theoretical framework for understanding the dynamics of compositional generalization in neural networks. Building on prior work that examines how models generalize by manipulating underlying primitive concepts, the authors introduce a Structured Identity Mapping (SIM) task, where a model learns the identity mapping on a Gaussian mixture with structured centroids. Through analysis of this SIM task, the paper reveals mechanisms behind non-monotonic learning dynamics, which is validated through experiments on text-conditioned diffusion models.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper’s explanation of learning dynamics, including the multi-stage Jacobian evolution and its impact on concept learning order and speed, sheds light on underlying model behaviors. I think the interesting point is the high similarity between the diffusion experiment and the theoretical model behavior, even though the correlation between the two is not clearly demonstrated.\n\n2. Experiments on SIM and diffusion models are well-designed, with relevant results that support the theoretical predictions, enhancing the practical relevance of the results.", "weaknesses": "1. The theoretical model studied in this paper differs significantly from practical diffusion models. Although the authors attempt to demonstrate a relationship between the theoretical analysis and practical model behavior through empirical results, this connection remains unclear, making it challenging to confirm whether the behavior observed in diffusion models is directly related to the theoretical insights presented. While I acknowledge the value of the analysis for single-layer and two-layer linear networks, it would be beneficial if the authors conducted further analysis on simplified diffusion settings to bridge the gap between theory and practice more effectively.\n\n2. In reference [1], the authors also discuss learning primitives in compositional tasks. Additional discussion on this work could be beneficial. Based on this, I have the following question: In [1], small initialization is an important factor for the model to learn primitives. In this paper’s theoretical part, small initialization is also assumed, but in practical diffusion models, is small initialization a crucial factor for learning primitives, or is it merely a theoretical necessity?\n\n3. Could Figure 5 display the training data’s trajectory evolution? I am curious whether there is a similarity between the training and test trajectories in the early stages of training.\n\n[1] Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing. arXiv:2405.05409.", "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)