Enhancing Neural Training via a Correlated Dynamics Model
Authors: Jonathan Brokman
Roy Betser
Rotem Turjeman
Tom Berkov
Ido Cohen
Guy Gilboa
Venue: ICLR 2024 poster (2024)
Keywords: dynamics modeling
training analysis
low dimensional model for training
Primary area: learning theory

Abstract: As neural networks grow in scale, their training becomes both computationally demanding and rich in dynamics. Amidst the flourishing interest in these training dynamics, we present a novel observation: Parameters during training exhibit intrinsic correlations over time. Capitalizing on this, we introduce  \emph{correlation mode decomposition} (CMD). This algorithm clusters the parameter space into groups, termed modes, that display synchronized behavior across epochs. This enables CMD to efficiently represent the training dynamics of complex networks, like ResNets and Transformers, using only a few modes. Moreover, test set generalization is enhanced.
    
We introduce an efficient CMD variant, designed to run concurrently with training. Our experiments indicate that CMD surpasses the state-of-the-art method for compactly modeled dynamics on image classification. Our modeling can improve training efficiency and lower communication overhead, as shown by our preliminary experiments in the context of federated learning.

Published as a conference paper at ICLR 2024
ENHANCING NEURAL TRAINING VIA A CORRELATED
DYNAMICS MODEL
Jonathan Brokman"
Roy Betser"
Rotem Turjeman" rot ao0 com
jonathanrrok@gmal.. com
roybelcampus technion. ac
Tom Berkov"
Ido Cohen
Guy Gilboa' guy www.gglbbaae...ommm technion,
@ com
wdocch@ariel.acco
ABSTRACT
As neural networks grow in scale, their training becomes both computationally demanding and rich in dynamics. Amidst the flourishing interest in these training dynamics, we present novel observation: Parameters during training exhibit in- trinsic correlations over time. Capitalizing on this, we introduce correlation mode decomposition (CMD). This algorithm clusters the parameter space into groups, termed modes, that display synchronized behavior across epochs. This enables CMD to efficiently represent the training dynamics of complex networks, like ResNets and Transformers, using only a few modes. Moreover, test set general- ization is enhanced.
We introduce efficient CMD variant, designed ing. Our experiments indicate that CMD surpasses the state-of-the-art method for compactly modeled dynamics image classification. Our modeling can improve training efficiency and lower communication overhead, as shown by our prelimi- nary experiments in the context of federated learning.
concurrently with train-
INTRODUCTION
The research of training dynamics in neural networks aims to comprehend the evolution of model weights during the training process. Understanding stochastic gradient descent (SGD) is of paramount importance as it is the preferred training method used in practice. By studying these dy- namics, researchers strive to achieve predictability, enhance training efficiency, and improve overall performance. In this work we show, for instance, how a good model can improve distributive and federated learning.
Outside the scope of deep learning, time profiles and correlation methods have been valuable for studying complex systems since the 1990s, as demonstrated by the foundational work of Aubry al. (1991; 1992). Studies in variational image-proce proccssioocssnnggggo such as Burger et al. (2016) and Gilboa (2018) also highlight the emergence of specific time profiles during gradient descent. However, these methods often rely on properties like convexity and homogeneity, which neural networks training dynamics typically lack.
Research into deep learning training dynamics dates back to the early days of neural networks, with initial studies focusing on two-layer models (Saad & Solla, 1995b;a; Mace & Coolen, 1998; Chizat & Bach, 2018: Aubin et al., 2018: Goldt et al., 2019; Oostwal et al., 2021). These simpler models continue to offer valuable insights, but our interest lies in the dynamics of contemporary, complex models that feature significantly larger numbers of weights and layers. These complexities introduce higher degrees of nonlinearity. Another popular simplification is the first-order Taylor approximation of the dynamics, particularly in the realm of infinitely wide neural networks, as in Jacot et al. (2018); Lee et al. (2019). However, this often results in performance shortcomings
different approach examines the low-dimensional sub-space hypothesis, positing that effective training occur in limited set of directions within the weight space, see Li et al. (2018a); Gur- Ari et al. (2018); Gressmann et al. (2020). Recently, Li et al. (2023) employed the eigenvectors
The contribution of these researchers to this work is equal. Technion Israel Institute of Technology, Ariel University

---

Published as a conference paper at ICLR 2024
of the covariance between weight trajectories define 40-dimensional training sub-space. They achieve strong performance in this constrained sub-space, setting the SOTA for compactly modeled dynamics in image classification. Other studies proposed to represent the training dynamics of complex models through a concise set of representative trajectories win (1):
my (t) ~ by + Σ
am Wm (4),
where bi, ai. are scalars and wt( the trajectory the i'th tunable weight (parameter) of the network. Common solution approaches employ dynamical analysis techniques, relying Koopman theory and Dynamic Mode Decomposition (DMD), see Schmid (2010); Dogra & Redman (2020); Naiman & Azencot (2021). Our experiments and analysis indicate that DMD and its variants do not adequately capture the complex dynamics inherent in neural network training (see Fig. 1 and Appendix H).
Our main observation is that network parameters are highly correlated and can be grouped into "Modes" characterized by their alignment with a common evolutionary profile. This captures the complex dynamics, as illustrated in Fig. 1. Our findings are substantiated through comprehensive experiments; a selection of key graphs presented in the main body, while the Appendix offers rich array of additional results, ablation studies, stability evaluations and details regarding various architectures and settings. Our contributions are three-fold:
The introduction of Correlation Mode Decomposition (CMD), data-driven approach that effi ciently models training dynamics.
Devising efficient CMD which can be performed online during training, out out-performing the state- of-the-art method for training dynamics dimensionality reduction of Li al. (2023) in both effi ciency and performance.
We pioneer the use modeled training dynamics to reduce communication overhcad. Our pre- liminary experiments offer promising avenue for future research.
DMO
weight
OLD
Figure 1: Correlated dynamics. Left: Correlation matrix of weight trajectories, clustered to modes, for simple network (See Fig. 10a in the Appendix) trained on MNIST (LeCun et al. (1998)). Middle: Clustered correlation matrix of weight trajectories for ResNet18 (He et al. (2016a)) trained CIFAR10 (Krizhevsky al. (2009)). The block diagonal structure indicates high correlations of the parameter dynamics within each mode. Right: Three weight trajectories from SGD training are reconstructed via DMD, which uses exponentials (Top) and CMD, which uses reference trajectories (Bottom). These highlight the effectiveness of CMD in capturing complex training dynamics. See Appendix A. Figs. 6, for extended findings.
PREVIOUS WORK
Dimensionality reduction techniques, particularly those using correlation matrices, are key to our approach and have roots in statistics and information visualization, see Nonato & Aupetit (2088); Sorzano et al. (2014); Van Der Maaten al. (2009); Van der Maaten Hinton (2008). Various stud- ies explore low-dimensional sub-space training using similar methods. For example, random projec- tions can capture substantial amount of the gradient's information eLi al., 20 20188 8a). Hessian- based methods focus its top cigenvectors, (Gur-Ari et al., 2018). Dimensionality has been further duced via smart weight sampling, (Gressmann et al., 2020). The work by Li et al. (2023) bears the most resemblance research. The authors affirm low-dimensional training dynamics, using top
2

---

Published as a conference paper at ICLR 2024
eigenvectors of the covariance matrix to span the training space. The gradient is fully computed for all of the weights before being projected onto this basis, making their method more involved than regular (stochastic) gradient descent. They also introduce efficient Quasi-Newton training scheme that benefits from the low-dimensional representation. Our approach distinguishes itself from Li et al. (2023) by focusing on dynamics mode decomposition, rather than training in low- dimensional sub-space. Both methods use dimensionality reduction to compactly model the dynam- ics, but operate differently. Our analysis is not based on eigenvectors but on nonlinear clustering, yielding improved performance. DMD and Koopman theory. (Schmid, 2010; Tu, 2013; Williams al., 2015; Kutz et al., 2016), decompose dynamics to modes using well-studied operators. They extend to neural training dynamics analysis (Manojlović et al., 2020; Dogra & Redman, 2020; Red- man et al., 2021), including advanced methods for highly non-linear cases Ych et al. (2023). Simánek et al. (2022) also incorporates DMD for training optimization. Contrary to these, we extract modes from the observed correlation in the dynamics without using predefined operator, implying more adaptive, data-driven approach. The geometry loss landscapes has been a topic of interest, see Shalev-Shwartz & Ben-David (2014); Keskar et al. (2016); Li et al. (2018b); Draxler et al. (2018). We showcase our dimensionality-redui
technique in this regard as well.
We demonstrate the relevance of our approach for communication overhead in distributed learning, a widely-studine challenge addressed by various techniques, including order scheduling and data compression (Ho et al., 2013; Cui et al., 2014; Zhang et al., 2017; Chen et al., 2019; Shi et al., 2019; Hashemi et al., 2019; Peng et al., 2019; Dryden et al., 2016; Alistarh et al., 2017). In Federated Learning, communication is particularly problematic (Li et al., 2020). For instance, Chen et al. (2021) finds communication takes up more than 50% of the training time in their experiments. Solutions involve adaptive weight updates to reduce communication load (Seide et al., 2014; Alistarh et al., 2017: Hsich et al., 2017: Luping et al., 2019; Konečný et al., 2016; McMahan et al., 2017; Ström, 2015; Chen al., 2021). In our work we propose for the first time to use dynamics modeling to reduce the number of communicated parameters.
METHOD
We present several approaches apply our method, starting by discussing the underlying correlated dynamics (Sec. 3.1). Following this discussion we introduce the basic CMD version, called post- hoc CMD, which analyzes dynamics after they occur (Sec. 3.2). We then present the more efficient online CMD (Sec. 3.3) and embedded CMD (Sec. 3.4), which embeds CMD dynamics to the training process. In Sec. 3.5 we demonstrate accuracy landscape visualization via CMD.
3.1 THE UNDERLYING CORRELATED DYNAMICS
Hypothesis The dynamics of neural network parameters can be clustered into very few highly correlated groups, referred to as "Modes".
Let W (w) denote the set of time-profiles of the neural network trainable weights during the training process, where is the weight index, t is the training time, and N is the total number of such weights. We refer W as the 'tajeectories' where each wt(() is single 'trajectory', The trajectories are clustered to M modes, where M < N. Let W m C W be the subset of trajec- mode m within W. Using minor abuse of notation, when is discrete we use W, Wm to denote the matrices where the (discrete) trajectories are arranged as rows. De- tails of the discrete case will be provided Sec. 3.2. The correlation of two trajectories u, v is denoted corr(u, (u, v). Any two time trajectories u, u that are perfectly correlated (or anti-correlated ie. |corr(u, v) = 1, can be expressed as an affine transformation of each other, u = av + b. Thus, under the assumption of correlated dynamics, the following approximation is used to describe the dynamics of a weight trajectory
tories corresponding
wi (t) s a.w (t) + b,, (ty(t) E Wm:
(2)
where a; b, are scalar affine coefficients associated with wi, and wm (t) is a single common weight trajectory chosen to represent mode termed 'reference weight'. Note that in terms of the usual dynamics decomposition approach of Eq. (1), the case of Eq. (2) is a specific case where a single (m)(t) models each trajectory. Figs. 1, 7 demonstrate that Eq. (2) approximates the dynamics using
3

---

Published as a conference paper at ICLR 2024
notably small set of modes (10 or less in most cases), supporting hypothesis of correlated dynamics. Surprisingly, parameters of the same layers are typically distributed between the modes, see Fig. 10b in the Appendix.
While neural networks are notoriously known for their complex dynamics, the data-driven approach of Eq. (2) is simple yet effective capturing these training intricacies. The reference weights hold the advantage not being restricted to common assumptions typically made in classical dynamics analysis. For instance, we show DMD classical analysis approach) to be too restrictive even for simple cases in Appendix H.2.
Our key novelty lies in leveraging the correlation between parameter trajectories compactly model the training dynamics. We will show that this approach outperforms current modeling methods.
3.2 POST-HoC CMD
In this section we introduce the basic CMD version, called post-hoc CMD, which analyzes dynamics after they occur. This algorithm models dynamics as per Eq. (2): It splits the trajectories to modes, associates a reference trajectory each mode and approximates the affine coefficients.
Consider a discrete time setting, E [I, where E is the number of epochs, and denote Nm the number of trajectories in mode m. Then Wm: W; E R. W RN*E Wm m
E RR- XE We consider ID vectors such as u E ER:E to be row vectors, and define ü == u E Di= u(t), (u, v) unf.
(a.0) ||u]| VuuT, corr(u, v) |
A,t B s.t. A(i) ag, B(i b, and let Ams Bmn
R Nm be the parts of A, B corresponding to the weights in mode m. Denote um
E R
where I = (1, 1) E B and d.. [A BTl ER In this context, Eq. (2) reads
W m Am Wm-
Step 1. Find reference trajectories. Given the trajectories w/ E W, use correlation select Wm. To avoid a large N x N matrix, we sample K trajectories to construct the absolute trajectory correlation matrix C₁, i.c.
(4) where Wkis Wk2 take any pair of the K sampled trajectories and M < K K C1 is used to cluster the K trajectories to M modes. In each mode, Wm is chosen as the weight trajectory with the highest average correlation to the other weights in the mode. Details the sampling and clustering procedures are available in Appendix C.2, C.3.
C1 12] = corr(wk, wk2)l. C1 E R
Step 2. Associate weights to modes. Given (wm) J.M and W, partition W to ( W Let
C2 E R.N Cali, m] = corr Wm)]-
(5)
To get [W m
J.M. mm=1, we obtain the mode of each w; W as arg max of the 7th row C2li,
Step 3. Obtain the affine parameters. Given Wiss W m. get Åm by pseudo-inverse of Eq. (3)
Åm W "n w ( (
w arg min W m Am Ammumle
Am
where ||- || denotes the Frobenius norm. Reconstruction of the modeled dynamics W performed by re-plugging Ån to the right-hand-side of Eq. (3)
Recon is
m
W Recon
A Wm + BTT,
m
where, e.g., mode weights at end of training appear at the last column, denoted W Recon
1==
We tested post-hoc CMD on image classification, segmentation and generative image style transfer. The performance upon replacing the weights with the CMD modelled weights produces similar, and even enhanced results, compared to the original network weights, see Fig. 2. Further experimental details and additional post-hoe results, and ablation studies are in Appendix B. As far as we know, are the first to assess dynamics modeling in image segmentation and generative style transfer. Post-hoc CMD is resource- -ntunccenteesvveeeeeee due to full training trajectory usage. Thus we introduce efficient CMD variants for real-time computation during training.
4

---

Published as a conference paper at ICLR 2024
Source
Reference
Original
CMD
0.30
sosan
M 1
cr an
030
Figure 2: Post-hoc CMD examples. CMD representation highly general and models well diverse architectures operating on various tasks. Here, the successful modeling evident through its ability to follow the performance dynamics while even providing a performance boost. Left: CIFARIO performance on ViT-b-16 (Dosovitskiy et al., 2020), pre-trained on JFT-300M. Middle: StarGAN- v2 (Choi et al., 2020) style transfer, qualitative result of the original network compared to the result of CMD modeling of its dynamics. Right: Segmentation results, using PSPNet Architecture (Zhao et al., 7). See more details and results in Appendix B and Figs. 15, 18. 'GD' in the legend stands for SGD.
3.3 ONLINE CMD
We introduce an iterative CMD version that mitigates the computational drawbacks of post-hoc CMD. Designed to run alongside the regular training process, with negligible overhead we term it online. Transforming post-hoc CMD into iterative algorithm, we eliminate the need for full training trajectories, reducing memory and computation.
Warm-up phase. We observed that the reference trajectories can be selected once during training. and do not require further updates (see Fig. 16 and Table in Appendix C.5). Thus we perform Step 1 of Sec. 3.2 once at a pre-defined epoch F, and the selected reference weights continue to model the dynamics for the remainder of the training process. After the warm-up phase, we perform efficient iterative updates, variants of Step 2 and Step of Sec. 3.2.
Step 2 (iterative). This step requires computing C2, the trajectory correlation matrix of Eq. (5). C2 calculated using inner products and norms. Denote w, R² as the discrete trajectory of a weight w at times 1 through t. Given (Wm,t-1 W 1), ||w,,,-111 112,, we may obtain (Wm,t. Wr,)), [[ as:
(wm, wr ) (Wm.t-1, Wi,c-1) 1) + Wm. (t), ||w,, ||w
w/ (t).
(8)
Step 3 (iterative). Let (m be Am evaluated on trajectories up to time t. Given Am - 1), we have by Eq. (6)
(m) = (Am(t- - (0000111)mmmmmm0mm (mm + Wm (t)oo
(t) (1) (m..)
Wm, TT m. Wm,t- 1T 1mm
w Wm (t) 1
Wmm( (t) )
(10)
where W (t), Wm:(t are the thh column of W m. wm respectively, and um, are columns 1 through t of wm- Note that we only used the weights at the current time t, along with the 2 matrices Wm,e-1 W" m.t-1 2m "n m.t which are also iteratively maintained via Eq. (10), In all of our experi- ments invertability of this x was never an issue (generally, pseudo-inverse can be used). Thus we use Eq. (9), Eq. (10) for iterative updates, see derivation of Eq. (9) in Appendix C.4.
Performance Our experiments show that after the F warm-up iterations we may fix not only Steps 1, but also Step 2, and continue with iterative Step 3 updates, while maintaining on-par performance with post-hoe CMD, as shown Table in Appendix C.5. Online CMD outperforms the current SOTA fow-dimensional model of Li et al. (2023) various image classification architectures, as highlighted in Table 1, while requiring a substantially shorter warm-up phase.
Efficiency Unlike Eq. (6), which is performed once at the end of the warm-up phase, the iterative updates, Eq. (9), Eq. (10), do not require trajectory storage at all. Moreover, their overhead is computationally negligible compared to training step, see Appendix C.6 for the caleulations. Thus we have an online dynamics modeling procedure with negligible overhead setting the ground for training enhancement. The full Online CMD algorithm is available in Al Algorithm 3 in Appendix C.1.
5

---

Published as a conference paper at ICLR 2024
3.4 EMBEDDING THE MODELED DYNAMICS IN TRAINING
This method gradually integrates CMD dynamics into training, aimed at leveraging modeled dynam- ies for enhanced learning. This improves accuracy and reduces the number of trained parameters. Experiments show that gradual CMD embedding stabilizes training. whereas instantaneous CMD use in SGD is unstable (Fig. 3). Two naive approaches failed: 1) periodic CMD weight assignment reverted to regular SGD paths; 2) training only the reference weights after the warm-up phase.
Let us describe the gradual embedding approach of choice. We rely upon the online algorithm of Sec. 3.3, where after a short warm-up phase the affine parameters ag, b, are iteratively updated at each epoch. Every L epochs we further embed additional parameters. A parameter wi for which the model is stable (change of a; by in time is negligible) is CMD embedded from then on, i.e. it gets its value according to the CMD model and not by gradient updates. Additionally, from this point on, a, b, are fixed.
0.92
0.91
Taired Parameters (%6]
80
0.90
0.89
60
0.4
0.88
5GD
40
0.3
0.2
CHD
0.86
0.87
Emterced CMD
Reference Welcchss The ng
Online CND
20
0.1
0.85
40
120
100
150
Epoch
Epoch
Spach
Figure 3: Embedding CMD in training. Left: Naively assigning CMD-modeled weights to the model every 20 epochs, and performing SGD in between. Dots represent the CMD model perfor- mance. Performance is repeatedly drawn back to the CMD-less (regular SGD) case. Middle: Online CMD and Embedded CMD, compared to regular SGD. An additional naive approach which trains only the reference weights post warm-up is also presented (green). Performance is portrayed along the training epochs. CMD runs use M 10 modes, F 20 warm-up epochs. We use an embed- ding rate of P = 10%, and L = 10 epochs (Right), which lead to = 50% reduction in the total number trained parameters in the whole training cycle. Both online and embedded CMD variants considerably improve performance, compared to regular SGD. Both naive approaches do not.
As shown in Appendix C.8, the affine parameters a,, by tend to exhibit minimal or no changes after certain epoch, denoted T;. Thus, after 7; we may a,, b, i.e. use a; < az(Tz), by +- b(T:)Vl 7. Denote I(t)= {i: f > T. To embed the modeled dynamics to w;, we perform
(t(() - SSG update
if $ I(t)
(11)
aaww (t) + by if E Z(t)
After we say that w; is an embedded weight, and a,, b, its embedded coefficents. Note: Embedded coefficients are frozen in time, embedded weights are not. Let e; be criterion of change in a,, b;, for instance (t())= / Io (t) (t 12 + (t) (t 2) One way to obtain T; is by
74 = (t) is in the bottom P% of the unembedded weights.
(12)
Other reasonable ways include т, = (t() < e, for un-embedded w;, using a small threshold >> 0. The full online embedding algorithm is available in Algorithm 4 in the Appendix. Note that we do not embed the reference weights { {wm) M mm=1- Sec. 4 demonstrates the benefit of embedded locally store the embedded coefficients
dynamics in the context of federated learning, where to significantly reduce the communication overhead.
3.5 CMD ACCURACY AND Loss LANDSCAPE VISUALIZATION
Reducing the search space to 1-2 dimensions allows for loss landscape visualization. To this end we use CMD with 2 modes, as seen in Fig. 4, where different parts of the dataset are used for accuracy visualization. Expectedly, training and validation sets differ in landscape smoothness. Smoothness well minima points also vary across the CIFAR10 10 classes. See further details in Appendix G.

---

Published as a conference paper at ICLR 2024
4,003
harny Sct
0.55
-0.020
-0.20
BAAAT SHE AUTOTOO
0500
0.300
sbM3 Sel 20000
E20
-3.012
0.900
-0.022
-0.012
0.003
-0.722
024
0.005
0.10
003
-4014
0.900
-4.11
0.000
-0018
-6.316
-1822
-2003
-8024
-0.16
-5.020
-5.032
-0:24
-0 01G
500
-0.20
-E,032
-0.224
-0.220
-0.972
-O 524
-0 0115
FAL
Visualization of accuracy landscape. A grid is created based the reference weights, and colored by accuracy values. The accuracy landscape of the training set (left) is very smooth, compared that of the validation set (second left). The landscape of the Automobile class (3rd left) is rather regular whereas the Dog class (right) is much more irregular and suboptimal. The original CMD model is marked with a black dot, the optimal model is marked with red dot.
Figure 4:
4 CMD FOR EFFICIENT COMMUNICATION IN DISTRIBUTED LEARNING
We propose for the first time to use training dynamics modeling for efficient communication in distributed learning. In conventional (non-distributed) learning, discussed above, a reduction in trained parameters leads to decreased computation. Nonetheless, in distributed scenarios, modeling trajectories may reduce the No. of communicated parameters as well, offering further significant benefits. Specifically in Federated Learning (FL), communication efficiency is crucial (Li et al., 2020). For example, Chen et al. (2021) find in their FL. experiments that communication accounts for over 50% of the total training time. Following their experimental setting, we concentrate on FL, where client models handle unique per-client data. Although developed for FL, our method is adaptable for distributed learning.
Let C be the set of client models, trained locally on distinct data. Periodical 'synchronization rounds' occur at time to follows: A subset C(t.) C is randomly selected and aggregated in a Central Server using standard averaging (MeMahan et al., 2017). We consider the (less standard) practice that aggregates only subset of the weights
Central Server: w! (ughiin (t.))
Σ
welient (1,); Vi & I(t.):
(13)
iC(t.) cleeneC(t.))
where Itt,) is the set of indices of the un-aggregated weights at time t.. Un-aggregated weights do not require communication with the central server, hence increasing the size of I while maintaining harness the embedded CMD modeling of Sec. 3.4, where performance is beneficial. To this end, (t.)) (E=)),ez(e.) as follows: The coefficients {as, bi) are the analyzed set of weight trajectories is {w: maintained in the Central Server as in Sec. 3.4, and calculated w.r.t. {wimain (,,)) Let T. be the round in which o, by are embedded. At tg Ts.is the pair {a,, b,) is distributed to the client models. Each client maintains local copy of all the embedded coefficients fa;, bihez The client synchronization is performed, similarly to Eq. (11),
mmai (tx) aaw (1,) +
if E I(t.) if E I(I.),
Local Client: w.
client (l.) +
(14)
where I(t,) fi:tt: > Te,i), and T. is updated the Central Server via Eq. (12). Note that only the weights corresponding to indices i & I(t,)) and the set {w, J.M mm= require communication. Ad- ditionally, each a,, b; pair is sent once during the training from the main model to the clients (when they are selected as embedded coefficients). Note: CMD is performed solely on {w:main (t,)),
In basic FL, N weights from each the |C| selected clients are sent to a central server during each synchronization step. These aggregated into main model, which is then distributed back to all (C] clients. Thus, before employing any compression, the total data volume transmitted for each synchronization, is given by:
Sync. Volume (Baseline) N[CC + [C))
(15)
where usually this is measured floating point numbers (depending on the variable-type of the N weights).
7

---

Published as a conference paper at ICLR 2024
In our scenario, due to the progressive change in I, the communicated volume is not constant. Thus, we consider the following average:
2N E, [C],
Sync. Volume (CMD) N noo eenedded (ICI+ + [C]) +
(16)
where Nroot-ente emneeddedd the number of un-embedded parameters, averaged over E, synchronization rounds. Note that upon full embedding. the M reference weights are still required for communica- tion. The term E, [C] accounts for the following: Each a,, pair is embedded and communicated once during training. The N pairs form total of 2N communicated parameters. Communication is performed to all clients, i.c. 2N|C| total volume. Average round communication is thus 2N E, [C].
In Fig. 5 our approach is compared to Chen et al. (2021), showing comparable accuracy with significant reduction of communicated weights.
5 RESULTS
We provide focused comparison between our approach and the most closely related contempo- rary works: Li et al. (2023) that uses the SOTA P-BFGS for compactly modeled dynamics im- age classification, and Chen al. (2021)'s APF and aggressive APF (A-APF) aimed at reducing communicated parameters in FL. While we are first to leverage dynamics modeling for efficient communication, APF targets parameters that nearly constant time, which is infact plateau dynamics.
Table 1: CIFARIO test accuracy results for different modeling methods. P-BFGS stands for the algorithm proposed in Li et al. (2023). Warm-up phase length (equal to CMD or longer) is noted in each P-BFGS column title. Best result in each row is in bold, second best is underlined. Vit-b-16 model is pre-trained on ImageNetIK. For the full experimental details see Appendix D, Tables 5, specifically.
Model
Baseline SGD
Online CMD
P-BFGS (equal)
P-BFGS (long)
18
90.91% + 0.52
92.43% = 0.34
89.34% = 0.54
91.03% = 0.38
WideRes
92.70% = 0.48
93.21 % = 0.50
88.65% = 0.77
92.78% = =0.25
PreRes 164
90.64% = 0.98
91.34% + 1.20
90.76% = 1.114
90.91% = =0.77
LeNet-5
75.36% = =0.57
76.77% = 0.60
74.67% 0.35
76.44%
= =0.37
GoogleNet
92.05% + 0.48
93.07% % + 0.49
92.01%
1.05
92.41%
= 0.38
ViT-b-16
97.86% 0.11
97.99% % + 0.28
95.05% 0.69
Dimensionality reduction. Table we compare our online CMD method to P-BFGS (Li et al., 2023) various architectures, ranging from very simple networks to vision transformers. We add results of CMD embedding for some the models Table 10 in the Appendix. For implemen- tation details up, and with an 80-epoch warm-up as they originally propose. Dimensionality 40 unless the warm-up is <40 epochs, then matches the warm-up length due method constraints. Our on- line CMD method outperforms P-BFGS in all experiments. With just handful of modes (10 or less), we achieve superior performance, regardless of the warm-up phase duration. P-BFGS requires substantially more dimensions and longer warm-up up phase than CMD, yielding models with low performance for equal warm-up phases CMD and less than 40 dimensions. Furthermore, On- line CMD has constant memory consumption and negligible computational overhead (see Appendix C.6, C.7). On the other hand, P-BFGS requires a matrix linearly dependent on the number of di- mensions, lcading to memory issues when training large models such Wide-ResNe (Zagoruyko & Komodakis, 2016) or vision transformer ViT-b-16 (Dosovitskiy et al., 2020). In fact, we could not train P-BFGS on ViT-b-16 using single GPU, like CMD. Moreover, the P-BFGS parameters for the long warm-up case using our computational device (125GB RAM memory), therefore only the equal warm-up phase option is available in the results.
Appendix D. P-BFGS appears twice in the table: With CMD-eeqivalent warm-
could not calculate
CMD as a regularizer by-product of applying CMD for dynamics modeling is smoother trajec- tories, demonstrated in Figs. 7. Smoothed trajectories have been shown improve performance,
8

---

Published as a conference paper at ICLR 2024
for instance in using EMA and SWA, (Tarvainen & Valpola, 2017: Izmailov al., 2018). We com- pare this performance boost with CMD in Table 2. using several SGD variants, It can be observed that CMD even slightly out-performs EMA and SWA. See additional information in Appendix D. Notably, EMA and SWA are specifically designed for smoothing and relate to CMD solely through this by-product; they do not involve dynamics modeling.
Table 2: CIFARIO test accuracy results on ResNet18 for different regularizers. Different train- ing optimization techniques used. Best result in each row is in bold, second best is underlined. CMD entails excellent implicit regularization, competing well with dedicated training regularization algorithms. 'Mom' stands for momentum and 'SCH' stands for learning rate scheduler.
Optimizer / Method
Bascline SGD
EMA
SWA
Online CMD
SGD SGD Mom
90.68%640.29 90.51%:10.52
90.16%
92.09%=00100
92.26% 0.12
92.42%+10.43 91.29%++0.23
92.28%:10.44
92.43% +: 0.34
SGD + Mom + SCH
91.28%+10.23
91.29%++0.24
91.68% 0.21
Adam
89.59%+0.44
91.43%
91.65%620.4
91.70% = 0.51
Federated Learning In Table 3 contrast CMD with APF and A-APF. As baseline, per- form FL training without any communication boosting. The Volume criterion quantifies the average number of parameters communicated per round, expressed as ratio to the baseline number of parameters. We note that no additional compression is performed here. We closely followed the experimental setting provided in Chen al. (2021), as detailed in Appendix F.I. CMD preferable in terms of Volume to APF and A-APF seen in Table 3.
Table 3: Federated Learning, ResNet18 on CIFARIO. CMD significantly reduces volume (48.12%) while boosting test accuracy compared the baseline. In contrast, APF marginally re- duces communication but attains a slightly higher aceuracy boost, and A-APF reduces communica- tion (not as much as CMD) with a negligible ac- curacy boost. Overall, CMD offers a favorable trade-off between communication and accuracy.
50
cmd (cliene)
regular
20
(aggresctve)
70
0.3
600
90
500
Round
Round
Figure 5: Federated Learning, ResNet18 on CIFARIO. Test loss and accuracy throughout training. CMD with two modes harnessed for FL, compared to regular SGD (Baseline) and Aggressive APF (A-APF). Each trend-line is the average of 10 runs.
Volume
Test Accuracy
Baseline
100%
91.4% +0.3
Ours
48.12%
91.7% #0.2
APF
90.22%
92.0% +0.2
A-APF
60.75%
91.5% +.0.4
6 CONCLUSION
In this paper, we have presented a new approach model neural network parameter dynamics, leveraging the underlying correlated behavior of the training process. We show the model applies for wide variety of architectures and learning tasks. By incorporating low-dimensional approxi- mation and using iterative update techniques, we introduce Online CMD, an efficient online method that models the training dynamics with just a handful of modes, while enhancing performance, We also incorporated the modeled dynamics into the training process and demonstrated the efficiency in Federated I Learning, improving the communication overhead. The idea to use modeled dynamics for this task is novel in itself.
Our experiments showcase a range benefits: We improve accuracy relative to regular SGD train- ing and also surpass the state-of-the-art in low-dimensional dynamics modeling. The model can also be used effectively for visualization of accuracy and error landscapes. Additionally, our feder- ated learning method significantly reduces communication overhead without compromising perfor- mance. Our approach could be effectively combined with other techniques, opening up new avenues for optimizing complex neural network training.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY
Detailed pseudo-code of our CMD algorithm is described in Algorithms 1, 2. Our other CMD variants, Online CMD and Embedded CMD are descibed in Algorithms 3, 4. All the algorithms are available in the Appendix. The parameters used for our experiments in Sec. 5 are available in the Appendix D. specifically Table 5 (training parameters) and Table (CMD parameters). All of our algorithms: Post-hoc, Online, and Embedded CMD, have an implementation provided here, alongside script which executes them the setting of the results Sec. 5.
ACKNOWLEDGEMENTS
GG would like to acknowledge support by the Israel Science Foundation (Grants No. 534/19 and 1472/23), by the Ministry of Science and Technology (Grant No. 5074/22) and by the Ollendorff Minerva Center.
REFERENCES
Dan Alistarh, Demjan Grubic, Jemyy Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-effiienn sgd via gradient quantization and encoding. Advances in neural in- formation processing systems, 30, 2017.
Travis Askham and Nathan Kutz. Variable projection methods for optimized dynamic mode decomposition. SIAM Journal on Applied Dynamical Systems, 17(1):380-416, 2018
Benjamin Aubin, Antoine Maillard, Florent Krzakala, Nicolas Macris, Lenka Zdeborová, et al. The committee machine: Computational to statistical gaps in learning two-layers neural network. Advances in Neural Information Processing Systems, 31, 2018.
N Aubry, R Guyonnet, and R Lima. Spatio-temporal symmetries and bifurcations via bi-orthogonal decompositions. Journal of Nonlinear Science, 2(2):183-215, 1992.
Nadine Aubry, Régis Guyonnet, and Ricardo Lima. Spatiotemporal analysis of complex signals: theory and applications. Journal of Statistical Physics, 64(3):683-739, 1991.
Martin Burger, Guy Gilboa, Michael Moeller, Lina Eckardt, and Daniel Cremers. Spectral decompo- sitions using homooeoggee functionals. SIAM Journal on Imaging Sciences, 330011488,48 2016.
Chen Chen, Wei Wang, and Bo Li. Round-robin synchronization: Mitigating communication bot- tlenecks in parameter servers. In IEEE INFOCOM 2019-IEEE Conference on Computer Commu- nications. pp. 532-540. IEEE, 2019.
Chen Chen, Hong Xu, Wei Wang, Baochun Li, Bo Li, Li Chen, and Gong Zhang. Communication- efficient federated learning with adaptive parameter freezing. In 2021 IEEE 41st International Conference on Distributed Computing Systems (ICDCS), pp. 1-11. IEEE, 2021.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018.
Yunjey Choi, Youngjung Uh. Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8188-8197, 2020.
Ido Cohen and Guy Gilboa. Latent modes of nonlinear flows: A koopman theory analysis. Elements in Non-local Data Inteructions: Foundations and Applications, 2023.
Henggang Cui, Alexey Tumanov, Jinliang Wei, Lianghong Xu, Wei Dai, Jesse HaaeerKuuharrye Qirong Ho, Gregory R Ganger, Phillip B Gibbons, Garth A Gibson, et al. Exploiting iterative- ness for parallel ml computations. In Proceedings of the ACM Symposium on Cloud Computing, pp. 1-14, 2014.
10

---

Published as a conference paper at ICLR 2024
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprins arXiv:1810.04805 2018.
Felix Dietrich, Thomas N Thiem, and Ioannis G Kevrekidis. On the koopman operator of algorithms.
SIAM Journal on Applied Dynamical Systems, 19(2):860 885, 2020.
Akshunna $ Dogra and William Redman. Optimizing neural networks via koopman operator theory. Advances in Neural Information Processing Systems, 33:2087-2097, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.I 1929, 2020.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309- 1318. PMLR, 2018.
Nikoli Dryden, Tim Moon, Sam Ade Jacobs, and Brian Van Essen. Communication quantization for data-parallel training of deep neural networks. In 2016 2nd Workshop on Machine Learning in HPC Enviromments (MLHPC), pp. 1-8. IEEE, 2016.
Mark Everingham and John Winn. The pascal visual object classes challenge 2012 (voc2012) de- velopment kit. Pattern Anal. Stat. Model. Comput. Learn., Tech. Rep, 2007(1-45):5, 2012.
Guy Gilboa. Nonlinear Eigenproblems in Image Processing and Computer Vision. Springer, 2018.
Sebastian Goldt, Madhu Advani, Andrew Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
Frithjof Gressmann, Zach Eaton-Rosen, and Carlo Luschi. Improving neural network training in low dimensional random bases. Advances in Neural Information Processing Systems, 33:12140- 12150, 2020.
Sylvain Gugger, Zachary Mueller, Sourab Mangrulkar, Thomas Wolf, Philipp Schmid, Lysandre Debut, Dhiraj D Kalamkar, Iker García-Ferrero, and Vasudev Gupta. accelerate. https:// github. http:/ccccchhhg/ae hwwp///hhgginggagccecc wwtp:///acccccccceeceeeeaaeaeeee 2021.
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.0477 2018
Sayed Hadi Hashemi, Sangeetha Abdu Jyothi, and Roy Campbell. Tictac: Accelerating distributed deep learning with communication scheduling. Proceedings of Machine Learning and Systems, 1:418 0300-430.220 200.20301999 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual net- works. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Nether- lands, October 11-14, 2016, Proceedings. Part IV 14, pp. 630-645. Springer, 2016b.
Shiqi He, Qifan Yan, Feijie Wu, Lanjun Wang, Mathias Lécuyer, and Ivan Beschastnikh. Gluefl: Reconciling client sampling and model masking for bandwidth efficient federated learning. Pro- ceedings of Machine Learning and Systems, 5, 2023.
Qirong Ho, James Cipar, Henggang Cui, Seunghak Lee, Jin Kyu Kim, Phillip B Gibbons, Garth A Gibson, Greg Ganger, and Eric P Xing. More effective distributed ml via a stale synchronous parallel parameter server. Advances in neural information processing systems, 26, 2013.

---

Published as a conference paper at ICLR 2024
Kevin Hsieh, Aaron Harlap, Nandita Vijaykumar, Dimitris Konomis, Gregory R Ganger, Phillip B Gibbons, and Onur Mutlu. Gaia:Geeo (Ge-Distrrbbted) machine learning approaching {LAN] speeds. In 14th USENIX Symposium Networked Systems Design and Implementation (NSDI 17), pp. 629-647, 2017.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407, 2018.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for im- proved quality, stability, and variation. arXiv preprint arXiv:1710.1000 2017.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: arXiv:/609.0e. 04836,, 2016.
Jakub Koneěný, Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05422 2016.
Bernar O Koopman.
Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 17(5):315-318 1931.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
J Nathan Kutz, Steven L Brunton, Bingni W Brunton, and Joshua L Proctor. Dynamic mode decom- position: data-driven modeling of complex systems. SIAM, 2016.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-233) 1998.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. arXiv preprint arXiv: 08838, 2018a.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances neural information processing systems, 31, 2018b.
Tao Li, Lei Tan, Zhehao Huang, Qinghua Tao, Yipeng Liu, and Xiaolin Huang. Low dimensional trajectory hypothesis is true: Dnns can be trained in tiny subspaces. IEEE Transactions on Pattern Analysis and Machine Intelligence, 15(3):3411-3420 2023. doi: 10. 10.1099/PAMM.222..1111110
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE signal processing magazine, 3703):50-60, 2020.
WANG Luping, WANG Wei, and Bo. Cmfl: Mitigating communication overhead for federated learning. In 2019 IEEE 39th international conference distributed computing systems (ICDCS), pp. 954-964. IEEE, 2019.
CWH Mace and ACC Coolen. Statistical mechanical analysis of the dynamics of learning in per- ceptrons. Statistics and Computing, 8(1):55-88, 1998.
Iva Manojlović, Maria Fonoberova, Ryan Mohr, Aleksandr Andrejčuk, Zlatko Drmač, Yannis Kevrekidis, and Igor Mezić. Applications of koopman mode analysis to neural networks. arXiv preprint arXiv:2006.1 /765, 2020.
12

---

Published as a conference paper at ICLR 2024
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication- efficient learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Igor Mezić. Spectral properties of dynamical systems, model reduction and decompositions, Non- linear Dynamics, 41:309-325, 2005.
Ilan Naiman and Omri Azencot. A koopman approach to understanding sequence neural models. arXiv preprint arXhv:2102.07824. 2021.
Luis Gustavo Nonato and Michael Aupetit. Multidimensional projection for visual analytics: Link- ing techniques with distortions, tasks, and layout enrichment. IEEE Transactions on Visualization and Computer Graphics, (5(8):2650-255) 2018.
Elisa Oostwal, Michiel Straat, and Michael Biehl. Hidden unit specialization in layered neural networks: Relu vs. sigmoidal activation. Physica A: Statistical Mechanics and its Applications, 564: 125:125 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Yanghua Peng, Yibo Zhu, Yangrui Chen, Yixin Bao, Bairen Yi, Chang Lan, Chuan Wu, and Chuanx- iong Guo. A generic communication scheduler for distributed dnn training acceleration. Pro- ceedings of the 27th ACM Symposium on Operating Systems Principles, pp. 16-29, 2019.
William T Redman, Maria Fonoberova, Ryan Mohr, loannis G Kevrekidis, and Igor Mezic. An operator theoretic view on pruning deep neural networks. arXiv preprint arXiv:2110.14856 2021.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recogmition challenge. International journal of computer vision, 115:211-22 2015.
David Saad and Sara Solla. Dynamics on-line gradient descent learning for multilayer neural networks. Advances in neural information processing systems, 8, 1995a.
David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995b.
Peter Schmid. Dynamic mode decomposition of numerical and experimental data. Journal fluid mechanics, 656:5-28, 2010.
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns. In Fifteenth annual conference of the international speech communication association, 2014.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algo- rithms. Cambridge university press, 2014.
Shaohuai Shi, Xiaowen Chu, and Bo Li. Mg-wfbp: Efficient data communication for distributed synchronous sgd algorithms. In IEEE INFOCOM 2019-IEEE Conference on Computer Commu- nications, pp. 172-180 IEEE, 2019.
Petr Šimánek, Daniel Vašata, and Pavel Kordík. Learning to optimize with dynamic mode decom- position. In 2022 International Joint Conference Neural Networks (IJCNN), pp. 1-8. IEBE, 2022.
Carlos Oscar Sánchez Sorzano, Javier Vargas, and A Pascual Montano, A survey of dimensionality reduction techniques. arXiv preprint arXiv: 1403.2877. 2014.
Nikko Ström. Scalable distributed dnn training using commodity gpu cloud computing. 2015.
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinay Gupta. Revisiting unreasonable ef- fectiveness of data in deep learning era. In Proceedings of the IEEE international conference on computer vision, pp. 843 852, 2017.
13

---

Published as a conference paper at ICLR 2024
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 1-9, 2015.
Antti Tarvainen and Harri Valpola. Mean teachers better role models: sistency targets improve semi-supervised deep learning results. Advances in neural information processing systems, 30, 2017.
Weight-averaged
Jonathan H Tu. Dynamic mode decomposition: Theory and applications. PhD thesis, Princeton University, 2013.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Laurens Van Der Maaten, Eric O Postma, H Jaap van den Herik, et al. Dimensionality reduction: A comparative review. Journal of Machine Learning Research, 10(66-71):13, 2009.
Matthew O Williams, loannis Kevrekidis, and Clarence W Rowley. A data-driven approximation of the koopman operator: Extending dynamic mode decomposition Journal of Nonlinear Science, 25:1307-1346. 2015.
King Fai Yeh, Paris Flood, William Redman, and Pietro Liò. Learning linear embeddings for non- linear network dynamics with koopman message passing. arXiv preprint w3x.h::3300000600 2023.
Sergey Zagoruyko and Nikos Komodakis.
Wide residual networks.
arXiv preprint
arXiv:1605.07146, 2016.
Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang
Wei, Pengtao Xie, and Eric P Xing. Poseidon: An efficient communication architecture for distributed deep learning on {GPU} clusters. In 2017 USENIX Annual Technical Conference (USENIX ATC 17), pp. 181-193, 2017.
Hengshuang Zhao, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia. Pyramid scene parsing network. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 2881-2890, 2017.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
A CORRELATION RELATED EXPERIMENTS
A.1 MANUAL INSPECTION OF MODELED TRAJECTORIES
Mode Size -4909914, Cor11 -0.786, Cor12 0.998 Mode 1. Stze 1498665, Corll 0.062. -09944
Mode 2. Size 425322. Corll 0.551 0092-0.947
Nonna bod Weigre gros
Norms brod ptess
Barch
Mode Sizc 051803, Corr1 Corl2
402582.
Cor12 -0.94
Modo Size 513343, 0.407,
mommo wogitt
Mar Mormateee wengstt
Tpech
Epech
Epech
Mode 6, Size 620700, Corri
0.527. Com2 0,970
Mode
428173, Corll
0.460, Con 0.888
Mode
Size 775012, Corrl 0.612, 0.972
homoeaad Weggo:
hurmalted Wergee
1 1.
Mode 9. Size 648467, Corr1 - 0.442, Corr2 0.954
Spach
Figure 6: For each mode, weight trajectories are portrayed as mean and variance of the normalized trajectories (shaded area) alongside the reference trajectory (solid line). Each subplot represents mode characterized by its parameter count (size), inter-mode correlation (corrl), and correlation of the mean trajectory with the reference trajectory (cor12). Formally, corrl for the cluster mes clus- ("2"") Eieco.. Σjecm.j> corr(w;, WT).. where Nm, is the size of the cluster, and cor2 is ter Cm is
corr Nm. Eiec,, Wi, wm)) Notably, corr2 consistently indicates strong alignment of the mean trajectory with the reference trajectory, demonstrating its efficacy in capturing the general dynamics, even upon less-than-perfee values in corrl. Additionally, larger modes exhibit higher corr2 as well as corrl values, suggesting more consistent behavior within the collective dynamics. The mode- average dynamics display varied patterns, including sharp decreases and subsequent stabilization, reflecting the learning process's 's adaptation phases. This diversity is effectively encapsulated by the reference trajectories, affirming our dynamics modeling approach despite the variation in correla- tion.
15

---

Published as a conference paper at ICLR 2024
(a) mode O
(b) mode
01100
(c) mode 2
(d) mode
(c) mode 4
(I) mode 5
mode 6
(h) mode
mode 8
i mode
Figure 7: CMD Weight Trajectory Reconstruction. randomly sampled weights from each mode, and their CMD reconstruction. CIFAR1O image classification, ResNet18. The CMD modeling provides more stable, less oscillatory solutions.
In Fig. 1 we compare CMD weight trajectory approximation to DMD weight trajectory approx- imation. Here we offer additional examples of CMD weight trajectory approximations. CMD is performed with M 10 modes and 3 weights tion. Manual observation is indispensable, therefore in Fig. 7 we show CMD smoothens the output, implying a regularization effect that may boost performance similarly smoothing methods like EMA and SWA (Tarvainen & Valpola, 2017; Izmailov et al., 2018). Note however that EMA and SWA do not offer modelling (which the core interest in our work).
randomly sampled from each mode for visualiza-
A.2 REFERENCE WEIGHT AND MODE WEIGHTS RELATIONS
In Fig. 8, statistical metrics validate CMD's weight trajectory modeling. The reference weight of a mode, together with the mean and variance of 10,000 randomly sampled weights per mode are presented. Different approaches are used Post-hoc CMD row), Naive reference weight training
16

---

Published as a conference paper at ICLR 2024
Vomaatoe Nanght
Vormatond Mesgghts
Vormalbers
16
90
130
24 Epoch
190
(a) Post-Hoc CMD, Test Acc- 91.24%%
M
Normasten Narghts
Normatoed
Ezoch
Epacc
Eosch
(b) Training the Ref- Weight Space, Test Ace 89.44%
Nomelrre 2nd Whight
kammlzzd Weight
Kormalzss
(c) Gradual CMD Embedding. Test Acc 91.50%%
Figure 8: Reference Weight Trajectories with the Mean and Variance of Weights per Mode. Top: post-hoe CMD; Middle: naive CMD embedding. Only reference weights trained post warm-up phase; Bottom: Gradual CMD Embedding. Test Accuracy of the final model is available in the caption of each row. Blue color represents SGD trained weights and orange represents CMD embedded weights. The gradual embedding case mitigates high performance with a converging mean and decreasing variance.
(middle row) and gradual CMD embedding (bottom row). In each row 3 modes are presented. In the post-hoc CMD case the mean of the weights in the mode do not follow the exact dynamic of the reference weight, it is somewhat similar and more smooth. In addition, the variance increases throughout the epochs. In the naive reference weight training case all of the weights in the mode directly related to the reference weight and change only according to the reference weight, post warm-up (orange). Therefore, there is no variance and the mean of the sampled weights is exactly equal to the reference weight. However, in this case the final test accuracy is not comparable to the other cases (more details about the accuracy are available in Fig. 3). The warm-up phase in this case was 50 epochs. In the gradual CMD embedding case the weights are gradually linked to the reference weights which results with converging mean and decreasing variance. The high performance maintained, terms of test accuracy. The warm-up phase is 20 epochs.
Note: In order to display weights with different ranges of values in the same graph we stretch the w. m.) values to the range of [0, 1] w; (wor(w;)- (w:)
17

---

Published as a conference paper at ICLR 2024
A.3 CORRELATION MATRIX EXAMPLES
Our main observation is that many network parameters, although non soontmooth in their evolution, are highly correlated and can be grouped into "Modes" characterized by their alignment with common evolutionary profile. In Fig. we displayed two examples of a clustered correlation matrix. Bellow we present two additional examples. The first example is of SimpleNet (see Fig. 10b) trained on binary classification task. In this case the parameters are divided into 3 modes. The main mode contains the majority of the network parameters. The second example is of a vision transformer, ViT-b-16, fine tuned CIFARIO. The model parameters are divided into 30 modes. The modes contain a decreasing number of parameters, and some redundancy is visible. Redundancy in the number of modes is implied when all of the weights from mode have high correlation all of weights in a different mode (mode O, mode 2, mode 4, etc. in this case, and mode 1, mode 3, etc.).
1.00
0.50
0.25
-0.25
0.75
Figure 9: Clustered Correlation Matrix. Left SimpleNet trained on MNIST, M = 3 modes. Right ViT-b-16 fine tuned on CIFARIO, M 30 modes. A single main mode visible for SimpleNet, while for ViT-b-16 it appears there is a certain redundancy in the high number of modes.
B POST-HoC CMD EXPERIMENTS
Image Classification First we consider the CIFAR10 (He et al., 2016a) classification problem, using a simple CNN architecture, used in Manojlović et (2020), we refer to it as SimpleNer (Fig- 10a). Cross-Entropy Loss and augmented data (horizontal flip, random crop, etc.) are used for training. The CNN was trained for 100 epochs, using Adam optimizer with momentum, and initial learning rate of 7 1 x 10-3 The CMD Analysis training. In this case we obtained 12 modes for the post-hoc CMD. The results from the CMD and the original SGD training are presented in Fig. 10c. We clearly see the model follows well the SGD dynamic, where both train and test accuracy are even slightly improved. In addition, the distribution of the modes through the entire network is presented in Fig. 10b. Each mode contains weights from each layer of the network. Therefor, trying to simplify the mode decomposition to full layers of weights is expected work as well. After validating our method on a simple scenario we tested CMD on larger architectures, such as ResNet18 (He et al., 2016a), see Fig. 12.
applied on the full 100 checkpoints of this
In recent years, after becoming key tool in NLP (Devlin et 2018) transformers are making a large impact in the field of computer vision. Following Dosovitskiy et al. (2020), we apply our method on a pre-trained vision transformer. Our model was applied on the fine tuning process of ViT-b-16 (Dosovitskiy et al., 2020; Gugger et al., 2021), pre-trained on the JFT-300M dataset (Sun al., 2017), on CIFARIO 10 with 15% validation/training split. We used Adam optimizer with start- ing learning rate of 5 x 10-5 with linear scheduling. The network contains 86 million parameters. Negative Log Likelihood Loss was used. One can see Fig. 2 that our method models the dynamics well and the induced regularization of CMD yields stable, non-oscillatory evolution.
Image Segmentation using PSPNet Our method generalizes well to other vision tasks. We present CMD modeling on the semantic segmentation task for PASCAL VOC 2012 dataset (Ev- eringham & Wimn, 2012). For this task we used PSPNet Architecture (Zhao et al., 2017). (13 108 trainable parameters). The model was trained for 100 epochs, using SGD optimizer with momen- tum of 0.9, and weight decay of 10-1. In addition, we used "poly" learning policy with initial
18

---

Published as a conference paper at ICLR 2024
fully- connected layer with 100 neurons + Relu
fully- connocted layer with
convolation 16 5x5 keruels max pooling
convolution 32 5x5 kermels max pooling
Softmax classification
impul
DisiBaaton roces trough lowwen
Na CMO Aecuncy
POBL
633
PCB
FCBL
(b)
Figure 10: SimpleNet on CIFARIO. (a) architecture (b) CMD Modes distribution through the
Network. (c) CMD modeling performance
learning rate of fj = 1 x 10-2, Post-hoc CMD was executed using M = 10 modes. Pixel accuracy and mloU results from the CMD modeling and the original SGD training presented in Fig. 11. Max CMD performance on the validation set: Pixel accuracy: 88.8% with mloU of 61.6%. Max SGD performance the validation set: Pixel accuracy: 88.1% 1% with mloU 58.9%. We observe CMD follows SGD well in the training, until it reaches the overfitting regime. For the validation set, CMD more stable and surpasses SGD for both quality criteria.
0.94
GO Vs. CMD Pael Accuracy
0.75
Train GO Trai CMD
0.70
0.92
Val CMD
0.65
0.90
Accuracy
mloU
0.60
0.88
Many M Mamo/ in
0.86
0.55
0.84
Train CMD
0.50
Val GO
0.45
20
40
60
80
100
20
60 Epach
100
Epoch
Figure 11: Post-hoc CMD on PSPNet. PASCAL VOC 2012 (Everingham & Winn, 2012) seg- mentation results, PSPNet Architecture. Original SGD training vs. CMD modeling with M = 10. modes. Left: pixel accuracy. Right: mloU. CMD follows SGD well in the training, until it reaches the overfitting regime. For the validation set, CMD is more stable and surpasses SGD for both qual- ity criteria.
Image synthesis using StarGan-v2 We applied our CMD modeling on image synthesis task using StarGan- (Choi et al., 2020). This framework consists of four modules, each one contains millions of parameters. It was modeled successfully by only a few modes, demonstrated in Fig. 15.
Ablaioon Study Our method introduces additional pareemterseeeers such as the number of modes (M) and the number of sampled weights to compute the correlation matrix (K). We performed
19

---

Published as a conference paper at ICLR 2024
Accurecy
90
Man Trulhhaeess
88
(U06)
86
(K00)
624
84
40
60
80
100
Epochs
Nam Modes
(a) Accuracy per Epoch
(b) Best Accuracy
DCCO1TA
5000100
Num Madit.
(c) Weights MSE
Figure 12: ResNet18 CIFARIO Results for Different Number of Modes (a) Test Accuracy per epoch. (b) Maximal Test Accuracy CMD model per number modes used. (c) MSE of Weights reconstruction per number of modes used for CMD, While the MSE decreases as the number of modes increases, this not reflected performance, which is relatively constant.
several robustness experiments, validating our proposed method and examining the sensitivity of the CMD algorithm to these factors.
Robustness to number of modes. Fig. 12 presents the test accuracy results of 10 different CMD) models, with different number of modes for the same SGD training. Surprisingly, the test accuracy close to invariant to the number modes. The best test accuracy per mode presented Fig, 12b. We reach the highest score for only two modes. appears ResNet18 is especially well suited for our proposed model. Weights MSE does in fact reduce as more modes generated (Fig. 12c), but this does effect the model performance. It is important to note that the weights MSE is very low even for modes.
2. Robustness to different random subsets. Here We examine empirically our assump- tion that in the K sampled weights we obtain all the essential modes of the network. In Fig. 13 we show the results of 10 CMD experiments executed on the same SGD training (ResNet18, CIFAR10), where each experiment different random subset the weights is sampled. We fixed the number of modes to M = 10 and the subset size to K = 1000. The mean value of the test-accuracy each epoch shown, as well as the maximal and minimal values (over 10 trials, vertical bars). As can be expected, there some changes in the CMD results. However, is highly robust and outperforms the original GD training- In addition, conducted several experiments, examining different values of K, the values were 50, 125, 250, 500, 1000, 2000. The number of modes was fixed (M 10). For any fixed K, 5 experiments were carried out, Fig. 13. We observe that the algorithm is robust to this meta meta-paranan parameter as well. The results tend to be with somewhat less variance for larger K, but generally very small, compared to the size of the net N.
the sampling size for the initial clustering phase can be
20

---

Published as a conference paper at ICLR 2024
828
80
70
60
CMD
20
Epocha
80
80 100
918
282
1000
1250
1500
1790
2000
(a) Fixed K
(b) Different K Values
Figure 13: ResNet18 CIFARIO Results for Different Sampled Weights. (a) Mean Test Accuracy epoch for different experiments with fixed K 1000. (b) Mean Test Accuracy per epoch for different experiments with different values. M = 10 in all the experiments. Despite different weights leading slightly different models (a), CMD consistently outperforms SGD. Different sizes of yield very similar results (b).
120
Figure 14: ResNet18 CIEARIO Results for Different Random Initializations. Mean test accuracy (and variance) results for 10 different experiments is presented for SGD and CMD. The variance of the CMD results similar to the variance of the SGD results, with a smoother mean.
3. Robustness to random initialization In Fig. 14 we show that the effects of random initialization of the weights is marginal. The results of 10 training experiments are shown (ResNet18, CIFARIO). In each experiment a different random initialization the weights is used. The same weights sampled and the number of modes is fixed (M 10, K 1000). The variance of the CMD results similar to the variance the SGD results, with smoother mean.
21

---

Published as a conference paper at ICLR 2024
Original + EMA
Original + EMA
Original
CMD
Original
CMD
Origiall
Source
Reference
Original
CMD
oa
reference guided synthesis
35-
22
500
0.40
0.00
0.40
0.44
0.31
Figure 15: StarGAN-v2 Style Transfer. Post-hoc CMD modeling of the dynamics of StarGAN- v2 on Celeb-A-HQ dataset Karras et al. (2017). It can be seen that the CMD regularization effect similar to that of the EMA in this style-transfer framework, both qualitatively, and in the image generation measures LPIPS and FID which measure discrepancy from real Celeb-A-HQ images. CMD performed for both reference-image and latent-noise guided generation.
22

---

Published as a conference paper at ICLR 2024
C METHOD DETAILS
C.I ALGORITHMS
In Sec. 3.2 we introduce the CMD algorithm, consisting of three steps. Steps 1 and describe clustering the model weights modes. The clustering use is correlation based clustering with linear complexity and is described in Algorithm The full CMD algorithm is described, including the third step, in Algorithm
Algorithm 1 Correlation- based clustering with linear complexity
Input:
W W ERNNT Matrix of all weights.
K Number sampled weights.
.M/t t Number of wanted modes desired in-cluster distance threshold.
Output:
{Cm ( m=11 Clusters of the network parameters.
procedure
1. Initialize Cn O. L..M.
2. Sample a subset of random weights and compute their correlations.
3. Cluster this set to M modes based on correlation values, update {Cm) accordingly.
4. Choose profiles W/m = arg maxujeC.,, Σ WC,ECm [corr(w;, Wy)]:
for i in (N K) weights do
m = arg maxm corr(w,, Wy.))l
Cm +- C. [w;}
end for return (Cm) )M m=1 end procedure
Algorithm 2 CMD Correlation Mode Decomposition
Input:
W E RRN Matrix of all weights.
K Number sampled weights.
/M Number of wanted modes desired in-cluster distance threshold.
Output:
Recon E NRNTT. CMD reconstruction of network parameters.
procedure CMD
Perform correlation clustering to M modes using Alg. 1. (M is fixed if given, or deduced according the input). Get mode mapping (Cm J
for m + 1 to M do Follow Eq. (6) to compute A, B, the affine parameters. Recon per mode, the CMD reconstructed weights.
Follow Eq. (7) to compute W m end for return W Recon end procedure
Algorithm can return the CMD parameters (A, B vectors, Wm. weights, etc.) instead the re- constructed weights Recon, if wanted. The Online CMD algorithm is presented in Algorithm 3.
C.2 SAMPLING DETAILS
We would like to make sure that every layer is represented in the subset that randomly sampled at the first step of our method (correlation and clustering). The choice of reference weights depends
23

---

Published as a conference paper at ICLR 2024
Algorithm 3 Online CMD Iterative AB Update Algorithm
Inputs:
F Number of warm-up phase epochs.
E Number of total epochs.
Any other input needed for the SGD training (learning rate, momentum, etc.) or CMD algorithm (CMD sampled weights number K, CMD number of modes M, etc.)
Outputs:
Regular SGD final model
CMD final model
Procedure
1. Start SGD training and save model checkpoint each epoch.
2. When epoch F' perform the full CMD algorithm. All CMD parameters are saved
(A, B, Related mode vectors and reference weights time profiles). No checkpoints (additional or previous) are needed.
3. For epoch from F + 1 E: update A,B vectors according to Eq. (9), using the current model weights. Update the reference weights time profiles.
4. After last epoch calculate final CMD model weights according to Eq. (7).
K this set. We therefore allocate a budget of 2-Jlagers] for each layer and randomly sample it uniformly (using a total budget of K 2 parameters). The other K parameters are sampled uniformly throughout 2
the network parmeter space. When K <= |layers| we simply sample all of the K parameters uniformly, In the context of image classification, we observed reduction in test accuracy vari- ance when using this procedure (compared to naively sampling with uniform distribution accross all weights).
C.3 CLUSTERING METHOD
To cluster the correlation map into modes employed full linkage algorithm the Farthest Point hierarchical clustering method. After creating a linkage graph we broke it down into groups either by thresholding the distance between vertices (c), by specifying number for modes (M). Note -In Step 1, we associate weights; in Step 2, only the remaining K weights are considered, reducing C2 from R to (NN-K)MM
The Clustering method could be substituted by any clustering algorithm such the simple k- k-means and others. We do not recommend methods such as PCA to avoid assumptions, such as orthogonal- ity, which may not apply weight trajectory space.
C.4 ITERATIVE STEP 3 DERIVATION
Given Am,t-1, we can re-express Eq. (6) as
Am W m 2i (wm wz )
(17)
[ m.t- rw T
m.t-1 + W (1)mm (1) (wmwin) )
On the other hand, by Eq. (6) at time / 1-1 1 we have W Wm,-1 mmt-1 Plugging this we indeed get Eq. (9), which is:
Am,(!) (Am,t- (Am,t-(@@m,t-tttt
wit wm, W Wm Wm.()) (4)))))))))) (wm,
(18)
C.5 POST-HOC vS ONLINE CMD COMPARISON
Post-Hot Hoc CMD is the basic approach, while Online CMD is highly efficient version of it. While it is often the case that efficiency comes the expense of performance here it not the case,
24

---

Published as a conference paper at ICLR 2024
validated by experiments, Table. 4. The comparison is done ResNet18 trained CIFARIO with different number of modes. The post-hoc experiments are done at the end of the full training (150 epochs). The online approach is performed with short warm-up phase (20 epochs) for steps and in the online CMD algorithm (see Sec. 3.3). The iterative updates are then performed until the end of training (from epoch #21 to the last epoch epoch # 150). Each experiment is performed 5 times, on the same SGD training process. The same weights are sampled for both approaches in each experiment. In Fig. 16 we present an experiment comparing online CMD runs with different warm-up phases. We examine short warm-up phases (20,30,40 epochs) and longer warm-up phases (80, 100, 120 epochs). We also present the regular SGD training case and the post- hoc CMD case. All CMD runs, online CMD with different warm-up phases and post-hoc CMD, achieve similar results. CMD is performed with 10 modes all the experiments in this figure, using the same training process.
Table 4: Comparing Online CMD vs the basic Post-hoc CMD. Accuracy performance is com- parable, where the online version has slightly higher variance. The number the model column represents the number of modes.
Model
Test Accuracy
CMD1 (online)
92.06% 0.05
CMD1 (post-hoc)
92.06% = 0.02
CMD2 (online)
92.07% 0.07
CMD2 (post-hoc)
92.01% 0.04
CMD5 (online)
91.94% = 0.12
CMD5 (post-hoc)
91.96% 0.05
CMD10 (online)
91.98% = 0.17
CMD10 (post-hoc)
92.05% 0.25
0.90
0.85
Test Accuracy
0.80
Online CMD W20
Online CMD W30
0.75
Online CMD W40
0.70
Online CMD W80
0.65
Online CMD W100
Online CMD W120
0.60
Post-hoc CMD
0.55
SGD
0
50
100
150
Epoch
Figure 16: Online CMD with different warm-up phases. Test accuracy for online CMD with different warm-up phases. The value after W in the legend descries the length of the warm-up phase. SGD and post-hoc CMD results are presented as well.
C.6 ONLINE CMD FLOPS COMPUTATION
Here we show rigorously that the Online CMD algorithm computation can be completely neglected compared to the SGD which it accompanies. Upon the online CMD Eq. (9), Eq. (7) are performed on each epoch that occurs after the F' warm-up epochs (note that Eq. (7) is performed only in order to validate the model and is not necessary for the training process itself. However we will take the worst when calculating the computational overhead.) We consider t2 t1 + 1, which also use in practice. Let us calculate the number of FLOPS required for each epoch in this case. Reminder: p matrix multiplied by a matrix takes nm(2p-1 FLOPS.
25

---

Thus in total have 16.V (1)) FLOPS, that be done on the CPU (do not require GPU). This is negligible considering typical actual FLOPS, for instance ResNet18 has 20.07-10" N N00N FLOPS just for forward pass, according to the model zoo provided here. Published as a conference paper at ICLR 2024
Let us start with Eq. (9), which is given by:
Am,12 = (Am.e1 e (W,,t, w m,,)) + W Wm,tit wi ( uy m We begin with the complexity of the first term:
(19)
Am,t (1m,,t m,,t))
(20)
Computational complexity: First multiplication a Nm 2 by a 2 6.Wm FLOPS. The second multiplication (inside the brackets), may be updated iteratively on its own, similarly to Eq. (8), with FLOPS,, i.e. we treat it as O(1). Thus we have 6N,,. O(1) FLOPS.
The last term:
(Wm,tz W m,
(21)
The multipication is similarly ((1) (via another iterative update), and the inverse of a i 0(1) well. The middle term:
W m.t Wn
TT
(22)
Multiplying N,, x -1 by at - t, 2, considering 12 t, yields 2N,,, FLOPS.
The first addition takes 2Nn,. FLOPS, and the remaining multipication by the inverse matrix is Nm 2 miltiplied by the inverse which
2 resulting with 6N,,, FT.OPS.
In total summing all FLOPS and summing over the modes, where the number of modes is much smaller than N, have
14N + O(1) FLOPS.
(23)
Now let move to the easier Eq. (7). Here have element-wise multipication and one elementwise addition to the weights vector, i.e. 2N FLOPS.
1.82 10
C.7 ONLINE CMD MEMORY CONSUMPTION
The final memory consumption of the algorithm, for a random number of modes, after the F warm- up epochs, is approximately 2.25 times the memory of single model. The CMD parameters, and there relative sizes are detailed below.
1. Vectors and B each have size equal the size of full model.
2. M vector which stores the related mode of each weight has the size of 0.25 times the model. This vector has the same length A and B, but the values stored have smaller range positive integers that represent the mode number. This vector can be compressed even more if necessary.
3. The full time profile for each reference weight. A 2 x 2 matrix can be stored for each reference weight instead, as explained in Step 3 in Sec. 3.3. Both cases have negligible size compared to the full model size.
is well established (see for instance Izmailov et al. (2018)) that the space occupied by storing the model is insignificant during the training process compared to the sustained gradients. Therefor, after the warm-up epochs this online method has negligible memory consumption.
C.8 COFFFICIENT STABILIZATION
Analyzing the changes in the affine parameters {a,, by J. among them, with some experiencing more significant changes than others. Additionally, as depicted in Fig. 17, the portion of parameters that remain relatively unchanged increases the training procedure progresses. This motivates us to gradually embed the CMD model into the SGD training process.
throughout training reveals variations
26

---

Published as a conference paper at ICLR 2024
1e0
3.5
Last Epoch 140
Epoch-
3.0
Last Epoch 120
Last Epoch-
Last Epoch
Last Epoch 90
Epoch
Last Epoch 70
Epoch
Last Epoch-50
Epoch
Last Epoch
10-
Last Epoch
0.5
0.0 0.0
0.4
0.0
Figure 17: Relative Change in the Affine Parameters. Histogram of the relative change in the [br) parameters calculated in different epochs. The relative change is compared w.r.t. {6, J at the end of the process, at epoch 150. As we can see, there are weights that do not change their value at all between epochs 20 and 150. Additionally, the portion of B parameters that do not change, or change a little, increases throughout the training procedure.
D RESULTS SECTION EXPERIMENTAL DETAILS
Models In order to generalize our method and verify that it works on various models and ar- chitectures we examined different ResNet models and networks with different architectures. We strive to check models with different properties in order to examine a diverse set of examples. In the ResNet family of models tested ResNet18, Preactivation ResNet- 164 (He et al., 2016b) and Komodakis, 2016). PreRes164 is a very deep network com-
Wide ResNet28-10 (Zagoruyko posed of more than 800 layers. The number of parameters in this model is 2.7 million, less than in ResNet18 (11.2 million). In Wide ResNet28- 10 the feature channels per layer are multiplied by 10 and the depth, the number of residual blocks, is 28. This model has 36.5 million parameters, more than three times the number of parameters in ResNet18. A different architecture, which is very ba- sic, is LeNer-5 (LeCun et al., 1998). This model, which was presented very basic architecture, composed of several convolutional layers. We also tested in GoogleNet (Szegedy al., 2015), another model which is not relatively new. addition these well known models we also demonstrate the CMD methods presented in this work on vision transformers ar- chitecture named ViT-b-16 introduced in Dosovitskiy et (2020). Transformers have become fundamental tool in NLP (Devlin et al., 2018) and many computer vision tasks (Dosovitskiy et al., 2020). The Vit-b-16 model we used was pre-trained on ImageNet (Russakovsky et al., 2015). The code we used for each model is linked to the following model names. ResNet18, Wide ResNet28-10. Preactivation- -estiont--eeee----eee- 164, LeNet-5, GoogleNet, ViT-b-16.
than 20 years ago, has
Table SGD Training Implementation Details. ViT-b- a cosine scheduler used, see more details in text. For generalization different optimizers are used.
Model
Training Epochs
Optimizer
Learning Rate
ResNer18 WideRes
150
SGD
0.05
150
SGD
0.1
PreRes164
150
SGD
0.1
LeNet-5
150
Adam
0.001
GoogleNet
120
Adam
0.001
ViT-b-16
90
SGD
0.03*
Each experiment is performed times and the mean test accuracy and standard deviation presented in the results tables (Tables 2). The training parameters (number of epochs, optimizer, learning rate) and the CMD parameters (warm-up phase, number modes (M), number of sampled weights (K)) for each model are presented in Tables 5, respectively.
Implementation Details
27

---

Published as a conference paper at ICLR 2024
Table CMD Implementation Details. Models with relatively high number of layers (PreRes 164, GoogleNet) require more warm-up epochs and more sampled weights for stable performance. For generalization different number of modes are used.
Model
Warm-up Epochs
Modes
Sampled Weights
ResNet18
20
10
1000
WideRes
20
1000
PreRes164
40
2
10000
LeNet-5
20
10
1000
GoogleNet
30
7
5000
ViT-b-16
20
10
1000
In all the experiments the weight decay was 0 and when SGD is used momentum is 0.9. In the ViT-b-16 training we used a cosine scheduler starting with a short (5 epochs) warm-up phase where the learning rate linearly increases from 0 to 0.03. Then cosine decreasing scheduler is applied for the remaining 85 epochs. The code for this scheduler is from here. The training parameters used training parameters used in other works, specifically in the linked projects from above. Additionally, we wanted to show CMD works with any optimizer, not specifically SGD, so we used Adam in several cases, Both SGD and ADAM optimizers are implemented using the Pytorch library (Paszke et al., 2019). The CMD warm-up phase was enlarged for relatively deep models, with many layers. The number of modes used in each experiment was not optimized. Rather, various number of modes from 2 to 10 modes were chosen for generalization purposes. P-BFGS code is provided by the authors of the original paper here. In Fig. 18 single run comparisons between the regular SGD presented. Each plot describes single run,
test accuracy and the Online CMD test accuracy a specific model (see figure caption), following the experimental details described above. The test accuracy of the Online CMD method is much more smooth compared to the regular SGD training result.
0.983
0.94
0.880
0.92
0.975
0.88
0.970
0.96
0.965
0.84
0.960
0.82
20D
0.955
902
0.82
0.80
40
80
120
0.950
50
75
0.80
30
60
90
120
Epoch
Epoch
Epoch
Figure 18: SGD Online CMD Comparison. Results of single are presented for each model. The Online CMD is superior to the regular SGD in all cases. Left: WideRes28-10 Middle: ViT-b-16, Right: GoogleNet.
For Table used the same learning rate in all experiments (LR 0.05), when momentum was applied the value was 0.9 and the learning rate scheduler (when used) is the Cosine scheduler with no warm-up epochs. All the training sessions are 150 epochs long, with warm-up phase of 20 epochs for CMD. EMA is performed with smoothing factor 3 = 0.9 and SWA is performed on the last 25% of epochs. CMD was performed with M 10 modes.
In order to demonstrate CMD on different dataset we present a small experiment, performed on CIFAR100 (Krizhevsky et al., 2009), in Table 7. The implementation details of this experiment on ResNet18 8 are 150 epochs, learning rate 0.05 and momentum 0.9. CMD is performed with 10 modes, the number of sampled weights (K) is 1000 and 20 warm-up epochs were used. For Embedded CMD used L 10 and P 10. The experiment is repeated 5 times, the mean and standard deviation of the results are presented.
In Table 2 results of CMD are compared to EMA and SWA, which all achieve similar results. To further analyze the relations between these methods we experimented with preforming one on top of the other (Table 9, Using ResNet18 CIFAR10, per epoch for EMA and CMD (post-hoc). Then, we performed SWA, EMA and CMD (post-hoc) on each set of weights SGD weights, EMA weights and CMD weights. The implementation details
reconstructed the full model weight values
28

---

Published as a conference paper at ICLR 2024
Table 7: CMD on CIEARI00. Both CMD methods (online and embedded) improve the baseline SGD training.
Model
Baseline SGD
Online CMD
Embedded CMD
ResNet18
66.81% 0.62
68. 74% 0.44 69.27% 0.57
Table 8: CMD TinyImageNet.
Model
Baseline SGD
Post-hoc CMD 28.4% + 0.72
ResNet18
26.01% 0.31
of this experiment are 150 epochs, learning rate 0.05 and momentum 0.9. CMD (post-hoc) is performed with 10 modes and the number of sampled weights (K) 1000. The results indicate that similar regularization introduced by EMA and SWA is attained by CMD as well.
Table 9: CMD relations with EMA and SWA: The test accuracy the end of the baseline SGD training is 91 24%. Using only one of the methods (SWA, EMA or CMD) leads to similar results as combining methods together. This experiment performed once, therefore there is standard deviation.
Weights/Method
SWA
EMA
CMD
SGD
92.1%
92.07%
92.06%
EMA
92.04
92.07%
CMD
92.09%
92.07%
E EMBEDDING CMD DETAILS
E. ALGORITHM
The full gradual CMD embedding and weight freeing algorithm is presented in Algorithm 4.
Algorithm 4 Gradual CMD Embedding Algorithm
syper-parameters:
F (Number of warm-up epochs);
L (Number of epochs between each embedding epoch);
P (Percentage of embedded coefficients in each embedding epoch);
Any inputs required for the CMD algorithm training process (learning rate, etc.)
Procedure
WiFF& +- weight trajectories of F regular SGD epochs.
A, B. related modes <- Post-hoc CMD on Wi.F-
4++0
b index set of embedded weights
for t
Fo
Regular SGD step Iterative update of A, B via Eq. (9)
SGD update if $ I
w/: (aw,((t)
in accordance with Eq. (11), where if E I then a,, by
if FieI
fixed in time
(1- (.-F%) = 0 then
Diff +- ((Aoaa A)2 + (Bols B)²
& excluding reference weights and embedded
weights
[A, B B) BJ
I - I Uindices of bottom percentile Diff}
end if
end for
29

---

Published as a conference paper at ICLR 2024
E.2 DIFFERENT EMBEDDING SCHEDULING
We examine different values for P. as well several approaches to define P, in our experiments. P can be defined as the percentage from the weights in the model, this noted as the basic case. In this case using a constant value P will lead to the same number of weights embedded in each embedding epoch. A different way to define P is the percentage of the not embedded weights. This relative percentage approach is noted as relative P. Applying constant value I' will lead to reducing number of weights embedded on each embedding epoch, as the number of not embedded weights reduces. Reducing the value of P is another approach. One decrease the value of P linearly, exponentially, etc. In addition a scheduled decrease can be used as well. In Sec. 5, in the federated learning experiments, combined approach involving both a constant P and exponentially decreasing P is used (see Appendix F.1). We demonstrate these different approaches bellow in Fig. 19. We examined different P values {10,20). for the basic case and the relative P case. In addition we used a pre defined scheduled decreasing P. starting from P = 20. We used ResNet18 trained and tested on CIFARIO. The CMD warm-up phase is 20 epochs and I. all cases. The mean test accuracy of 5 different training cycles is presented, with the final accuracy and standard deviation available in the legend. Overall, from the five different options presented, embedding the A.B vectors iteratively, with P 20% and using the relative percentage approach yields the best results. However, results are very similar for most cases.
in
0.92
Test Accuracy
0.75
Test Test Accuracy
0.90
SGD 90.91 0.52
0.50
Online CMD - 92.34 0.42
P=10 99.67 0.57
0.88
P-20 90.54 0.47
P-10 91.64 = 0.35
0.25
0.86
Schedu P 91.74 0.24
0
50
100
150
40
80
120
Epoch
Epoch
Figure 19: Gradual Weight Embedding with Different P Scheduling. Results are improved com- pared to SGD training, while there is some drop in performance compared to the regular CMD. The results very close, most options achieving nearly the same results. The only drop in performance occurs in the basic approach with = 20, where after only 70 epochs (five embedding epochs) all of the weights are embedded.
E.3 EMBEDDED CMD VS ONLINE CMD
We compare the Online CMD algorithm (Algorithm 3) to the gradual CMD embedding algorithm (Algorithm 4) different models, see Table 10. In these experiments the implementation details are the same as the experiments in Table 1. The Embedding algorithm parameters are fixed for all experiments
10, P 10, using the basic constant P embedding approach. The gradually embedded CMD method introduces weight freezing which enhances efficiency. Therefore, a limited drop in performance reasonable. For some models the chosen hyper- parameeers achieve better results than the full online CMD method.
Table 10: Comparing Online CMD vs Embedded CMD. In some cases the Embedded CMD outperforms the Online method, in other cases the online method with no embedding superior.
Model ResNet1 8 WideRes
Baseline SGD
Online CMD 92.34% + 0.42
Embedded CMD
90.91% % + 0.52
91.80% + 0.30
92.70% + 0.48
93.21 + 0.50
92.36% 0.77
PreRes GoogleNet
90.64% + 0.98
91.34% + 1.20
91.71% 0.58
92.05% + 0.48
93.07% + 0.49
93.36% 0.14
30

---

Published as a conference paper at ICLR 2024
F FEDERATED LEARNING DETAILS
F.1 FEDERATED LEARNING EXPERIMENT DETAILS
We closely followed the experimental setting provided in the original paper: 50 Clients with disjoint datasets of non-IlI class distribution that is drawn independently for each client from a Dirichlet dis- tribution with a 1,500 rounds. Each synchronization round, each client completes 10 epochs. The APF thresh- old set to T 0.05, and whenever the number of embedded coefficints crosses 80%, perform thresh &- threah To conform with this scheduling, also perform for our method P t- 5/2 under the same circumstances. For CMD we set P 5%. The "aggressiveness' parameter schedule of A-APF is set as min round 0.5).
10% of the clients are sampled for synchronization i.e. |C't,,] O.1
2000
F.2
APF CODE
Weight freezing can be combined with other strategies to enhance efficiency in Federated Learning. GlueFL He al. (2023) stands out as a framework that adeptly merges weight freezing with client sampling, ensuring optimized communication overhead, especially crucial in cross-device settings like mobile or IoT devices. They have open-sorce code, including comparison to Chen et (2021) in here. We used their project for our APF implementation.
F.3 CALCULATION OF THE NEGLIGIBLE ADDED COMMUNICATION
2x
Let us show that N Nereontedtt
10-4,
In our setup of ResNet18, have N 10 Following
[||++
Chen et al. (2021) have |C|
50, CC 5. As shown in Fig. 5 these result with convergence for N
E ~ 103. Thus we have Nus-cnialta (IC) +[]
2N
2N
2 N
In
F
5-10 Noseennitae
our experiments
2C 0.5, resulting with 5-10
10-4
N
F.4 POTENTIAL OF GENERALIZED DYNAMICS MODELING IN FEDERATED LEARNING
Our work with CMD in FL has demonstrated promising preliminary results. A key contribution here is leveraging dynamics modeling for efficient communication in distributed learning, and is not restricted to CMD. This section discusses the potential of extending the framework to include various dynamics modeling approaches, emphasizing the value it may serve in FL.
The core of this generalized approach remains the use of dimensionality reduction. By communicat- ing oow-dimeniional representations of model parameters between clients and the central server, the volume of transmitted data can be significantly reduced. Alongside the transmission of compressed data, is crucial for client models effectively reconstruct the original dynamics from the reduced information. Maintaining local copies of necessary data and mappings enables clients to accurately translate the low-dimensional representations back the full parameter space.
With dimensionality reduction at its core, we believe this approach may allow for a more inter- pretable understanding the model's behavior. For instance, by focusing on reduced set of pa- rameters, it becomes easier to trace and understand how these parameters influence the learning outcome. This particularly beneficial in FL environments, where understanding the contribution of individual clients to the collective model is often complex and opaque. The ability to map back and forth between the full and reduced parameter spaces not only ensures efficient communication but also provides a clearer insight into the dynamics of the learning process. This transparency has the potential to aid diagnosing the learning process for various issues.
As the field of dynamics modeling continues to evolve, new dynamics modeling approaches will emerge. A generalized framework, not limited to CMD, can seamlessly integrate these techniques, ensuring continued effectiveness and relevance in rapidly evolving domain. Moreover, natural extension to generalize the framework to accommodate a range of dynamics modeling techniques. This generalization will utilize the most suitable dynamics model for specific tasks.
In summary, the generalization of the CMD framework to include a variety of dynamics modeling approaches presents an intriguing prospect. We believe such an approach remain relevant and
31

---

Published as a conference paper at ICLR 2024
adapt for future advancements in FL well as dynamics modeling. The goal is to enhance the adaptability, efficiency, and scalability of FL. providing a foundation for further research and de- velopment in this area. The addition of explainability to this framework not only aids in effective communication but may also enhance our understanding of distributed learning processes. Thus we believe this could be valuable tool for future advancements in the field.
G CMD LANDSCAPE VISUALIZATION: EXTENSION
64.4
Figure 20: Loss and Accuracy Landscapes for Post-hoc CMD. Left- CMD with a single mode, ID curves; Right CMD with two mode, 2D landscapes. The optimal models from each dataset, and the original model are marked on each landscape and described in the legend.
A different feature available due to the reduced number of parameters of the model is visualization feature. Shifting the reference weight values of CMD with a single mode will lead to a 1D loss curve. In the 2 modes case 2D loss landscape is created. Starting from post-hoc CMD with single mode, in Fig. 20 we present the ID curves and 2D landscapes representing the train loss and validation accuracy metrics for different values of the reference weights. The values of the reference weights determine the value of all the other reconstructed weights in the model, as shown in Eq. (7). Changing the reference weights values changes the values of all the other weights as well, creating a new model, easily calculated. An array of values are tested in the area of the original reference weights values (Wp-
wo
2 we + too with wo increments (ID case), or 50 1 wi increments (2D case).
[wo wo We + wo [wi - 1 1 w1 + 2 with 50 50
The main visual difference between the landscapes of the training loss and the validation accuracy in both cases is that the training set loss is smooth and the validation set accuracy is not. For the training set case there is a clear global minimum. For the validation set case there are many local maximum points, around the global maximum. We will note that the training set accuracy behaves similarly to the training set loss function and that the validation set loss behaves similarly to the validation set accuracy function. The smooth landscape in the training set case indicates that the CMD model is much more stable on the training set, compared to the validation set. Generated from the training data, this is an expected feature of the CMD model.
Fine tuning the model on a specific class (or classes) of data could also be done via this visualization method. An example is presented in Fig. 21, where each 2D landscape presents the validation ac- curacy landscape for images of a specific class. For each class different values of (wo, wh) optimize the validation set accuracy. The 2D landscapes produced per class show different behaviour from the full validation set accuracy landscapes presented earlier in this section. These 2D landscapes are divided into regions with the same exact performance. These different landscapes show how each model calculated with different reference weights values affects each class differently. The caption of each sub-figure contains the name of the class and the test accuracy of the optimized model in brackets. The original test accuracy, before any optimization, is 91. 1999%. Some classes lead to better optimized models (bird, frog, etc.) and others lead to lower overall test accuracy (airplane, deer, etc.).
A different approach is embedding a subset of the CMD model parameters into the model trained by SGD. The embedded parameters are dependant on the reference weights and the other parameters have fixed values. In Fig. 22 we illustrate an example (ResNet18 on CIFARIO) where 10% of all the model weights are embedded, 30% are embedded and 100% of the model weights are embedded. As expected, when more weights are embedded, and linked to the reference weights, a more significant change is visible in the 2D loss landscape.
32

---

Published as a conference paper at ICLR 2024
<00
-0
(a) Airplane (91.7%)
(b) Automobile (92.07%)
(c) Bird (92.13%)
-000
<00
(d) Cat (92.01%) %)
(c) Deer, (91.74)%
(f) Dog (91.96%)
-00
(g) Frog (92.1%)
(h) Horse (91.65%)
(i) Ship (92.13%)
4.00
(i) Truck (92.1%)
Figure 21: 2D Validation Accuracy per Class Landscapes for CMD with Two Modes. The caption of each sub-figure contains the name of the class and the test accuracy of the optimized model in brackets. The original test accuracy, before any optimization, 91.99%. Some classes lead to better optimized models (bird, frog, etc.) and others lead to lower overall test accuracy (airplane, deer, etc.).
To conclude this section, using the reduced number of parameters for visualization intents shows interesting options for different subsets of the data or different subsets of the model parameters. The visualization also leads to the conclusion that there are many similar models, that achieve similar results, in the area of the original CMD calculated model. By changing the values of the reference weights we can easily calculate a new model, with different weight values, that achieves similar results.
H DMD ANALYSIS FOR NEURAL NETWORKS
H.1 DYNAMIC MODE DECOMPOSITION
DMD is one of the common and robust tools for system analysis. It is widely applied in dynami- cal systems in applications ranging from fluid dynamics through finance control systems (Kutz et al., 2016). DMD reveals the main modes of the dynamics in various nonlinear cases and yields constructive results in various applications (Dietrich al., 2020). Thus, it is natural to apply DMD-
33

---

Published as a conference paper at ICLR 2024
100%
0.125
0.125
0.125
0.100
0.100
0.075
0.075
0.075
0.050
0,050
0.050
0.025
0.025
6.08
40. MM5 0.3
-000--2-0 ode
Figure 22: Loss Landscapes for Post-hoc CMD with Different Portions of the Embedded Model. Left 10% of the CMD model weights are embedded; Middle 30% of the CMD model weights are embedded: Bottom 100% of the CMD model weights are embedded; A stable be- haviour is maintained when different portions the CMD model are embedded in the SGD trained model. When a larger percentage of the CMD model weights are embedded, the changes are more visible the landscape. The maintained minima as well as smoothness throughout the subplots suggests that incorporating partial CMD modeling in training, as in embedded CMD, is viable
type algorithms model the process of neural network training in order analyze the emerging modes and to reduce its dimensions. For instance, in Naiman & Azencot (2021) it was proposed use DMD as a method analyze already trained sequential networks such Recurrent Neural Networks (RNNs) to asses the quality of their trained hypothesis operator in terms of stability.
DMD is originally formulated as an approximation of Koopman decomposition (Mezić, 2005), where it is argued that any Hamiltonian system has measurements which behave linearly under the governing equations (Koopman, 1931). These measurements, referred to as Koopman eigenfunc- tions, evolve exponentially in time. This fact theoretically justifies the use of DMD, where it can be considered as an exponential data fitting algorithm (Askham & Kutz, 2018). The approximation of Koopman eigenfunctions by DMD can generally not be obtained for systems with one more of the following characteristics (Cohen & Gilboa (2023)):
1. Finite time support.
2. Non-smooth dynamics- not differential with respect to time.
3. Non-exponential decaying profiles.
Below we demonstrate simple example of applying DMD single-layer linear network, show- ing that already this very simple case we reach the limits stated above when using augmentation.
H.2 DMD IN NEURAL NETWORKS
The training process of neural networks is highly involved. Some common practices, such aug- mentation, induce exceptional nonlinear behaviors, which are very hard to model. Here we examine the capacity of DMD to characterize gradient descent process with augmentation. We illustrate this through simple Toy example of linear regression.
Toy example- linear regression with augmentation. Let formulate linear regression prob- lem as follows,
L(x, wi v) = lll wrl/
(24)
where ||-| denotes the Frobenius norm, E mmxxn ạ gdx= and _x In order to optimize for w, the gradient descent process, initialized with some random weights wo.
*
w*x) T..
(25)
where the superscript & denotes epoch number As T' denotes transpose and n is (possibly adaptive) learning rate. We note that introducing an adaptive learning rate by itself already makes the dynamic nonlinear, 7 denotes time interval. Let us further assume augmentation process is performed, such that at each epoch we use different sets of x and corresponding y matrices, that is, (a
w n/ (y**
(26)
34

---

Published as a conference paper at ICLR 2024
Then, the smoothness in Eq. (25) becomes non-smooth in the augmented process, Eq. (26). Actu- ally, we obtain a highly nonlinear dynamic, where at each epoch the operator driving the discrete flow is different. We generally cannot obtain Koopman cigenfunctions and cannot expect DMD to perform well. Our experiments show that this is indeed what happens in practice. In the modeling of simple networks, DMD model the dynamics of classical gradient descent but fails as soon as augmentation is introduced. Fig. 23 presents an experiment of CIFARIO classification using simple CNN (Fig. 10a) with standard augmentations, usually performed on this dataset (horizontal flips and random cropping). In this experiment in order to reconstruct the dynamics, we perform DMD with different dimensionality reduction rates (r = 10, 50, 90) using Koopman node operator, as described in Dogra & Redman (2020). The results presented in Fig. 23 show that when the dimensionality reduction high (r = 10) DMD fails to reconstruct the network's dynamics. How- ever, when the dimensionality reduction is mild (r 90) DMD reconstruction is also unstable and oscillatory. An example of DMD performed on a more complex network (ResNet18), with no data augmentation, is also available in Fig. 23.
80
GD
DMD (r=10)
60
Accuracy
Accuracy
60
DMD (r=
DMD (r=90)
40
GD
40
DMD (r=10)
DMD (r=50)
20
(r=90)
20
20
40
60
80
100
0
25
50 Epochs
75
100
Epochs
Figure 23: DMD Experiments. CIFARIO classification results, Original Gradient Descent (GD) training Vs. DMD modeling. Top: SimpleNet architecture, with augmentations. Bottom: more complex model- ResNet18, no augmentations. 3 values of dimensionality reduction parameter r per experiment. DMD fails to achieve comparable results in all cases.
H.3 DMD AND CMD
The common ground between DMD and CMD is the intention to approximate a complex dynamic with simpler one. In both cases, the original dynamic is decomposed into time profiles and spatial structures. The main difference is that DMD uses for its modes pre-known eigenfunctions, in a linear approach, which is limited exponential profile representation. CMD is nonlinear, using a single representative with improved adaptability. DMD employs Eq. (1), while we employ the special case of Eq. (2).
Several key differences exist - first, DMD is computed by finding a linear operator that approximates the dynamic with squsre problem. The operator is later decomposed into spatial structures (modes) and cigenvalues yielding exponential profiles in time. CMD, conversely, is not based on pre- known temporal funetions, but rather assumes correlation between parameters and uses time profiles present in the dynamic itself. They are not computed using an eigen-decomposition but rather by clustering a correlation matrix and finding a single representative profile in each cluster. Secondly, there is no operator approximation in CMD, thus it does not have, at least directly, the option to extrapolate the dynamic by continuing applying the operator to the end the training data. Thirdly, decomposition in CMD is "transposed" in relation to DMD, in the sense that the CMD "modes" are groups with similar temporal behavior (a cluster the correlation matrix), while in DMD each mode spatial function that has be multiplied by an exponent to form its contribution to the total. To match this logic, the CMD modes should have been a vector including all coefficient vectors (a, b) from all clusters. But then, the whole vector does not have the same temporal behavior, as is divided to clusters. For CMD with a single mode the coefficient vector is similar to the DMD notion of a spatial mode.
35


Reviews:
[{"review_id": "x8iXCSZDOP", "reviewer": "Reviewer_t9P8", "summary": "This manuscript introduces an interesting idea, namely, *correlation mode decomposition (CMD)*, to cluster the parameter into groups. Instead of considering the top eigenvectors of the Hessian, the idea of CMD efficiently models training dynamics. The insights of CMD can be applied to communication-efficient distributed training.", "questions": "NA", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* In general this manuscript is well-structured. \n* This manuscript considers an interesting aspect of modeling the training dynamics of complex networks. The idea of using clustered parameters looks novel to the reviewer.\n* The manuscript has a good logic flow, from the definition of the post-hoc CMD to online CMD and embedded CMD.\n* Sufficient numerical results also justify the effectiveness of the CMD. An extension to FL is also provided in the manuscript.", "weaknesses": "* Authors are encouraged to improve the writing quality of the current manuscript.\n* Regarding the experiments on FL, it remains unclear to the reviewer why only two communication-efficient FL baselines, namely APF and A-APF, are considered for the evaluation. More recent SOTA methods need to be taken into account.", "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": "2lDWgxsihE", "reviewer": "Reviewer_8rcX", "summary": "The paper presents a novel observation and methodology related to the training dynamics of large-scale neural networks. The authors observe that the parameters during the training of neural networks exhibit intrinsic correlations over time. Capitalizing on this observation, they introduce an algorithm called Correlation Mode Decomposition (CMD). CMD clusters the parameter space into groups, termed modes, that display synchronized behavior across epochs. This representation allows CMD to efficiently encapsulate the training dynamics of complex networks like ResNets and Transformers using only a few modes, enhancing test set generalization as a result. An efficient CMD variant is also introduced in the paper, designed to run concurrently with training, and the experiments indicate that CMD surpasses the performance of existing methods in capturing the neural training dynamics.", "questions": "* Could you clarify the visual discrepancy observed in the middle plot of Figure 2? What does the blur on the right side represent?\n\n* Would it be possible to extend the range of epochs shown in the left plot to provide a more comprehensive view of the CMD method's performance over time?\n\n* Could you elaborate on the motivation and rationale behind applying the CMD method in federated learning? \n\n* What considerations were made in choosing the comparative methods and evaluation criteria in this work?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The paper introduces a novel observation regarding the intrinsic correlations over time of parameters during the training of neural networks. This insight is leveraged to develop a new algorithm, Correlation Mode Decomposition (CMD), which is a creative contribution to the field.\n\n* Despite the complexity of the topic, the paper seems to be structured and articulated in a manner that allows the reader to follow the authors' logic and methodologies.", "weaknesses": "* The citation format within the text could be improved for consistency and adherence to academic conventions. Utilizing citation commands like \\citet or \\citep would enhance the readability and professionalism of the references within the text.\n\n* Figure 2 Analysis: The benefits of the CMD method as depicted in Figure 2 are not evidently clear. In the left plot, it would be helpful to see the results over a more extended range of epochs to ascertain the method's effectiveness over a longer training period.\nIn the middle plot, there appears to be a visual discrepancy or blur on the right side that might need clarification or correction.\n\n* Algorithm Explanation and Comparison: A more detailed explanation and justification of the CMD algorithm's design and implementation are necessary for a thorough understanding. The comparison seems limited, primarily focusing on P-BFGS. It would be enriching to see comparisons with other relevant methods, such as pruning methods, to gauge the CMD algorithm's relative performance and novelty.\n\n* The presentation of some concepts and terminologies, such as the \"modes\" in CMD, might be unclear or lack sufficient explanation, making it challenging for readers unfamiliar with the topic.\n\n*  Discussions or demonstrations regarding the scalability of the CMD algorithm when applied to larger or more complex neural network models would be a valuable addition to assess the method's practical applicability.\n\n* The motivation behind applying CMD in federated learning seems a bit unclear and could benefit from a more explicit demonstration or explanation.", "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": "LnUfA81iPA", "reviewer": "Reviewer_yr3M", "summary": "This paper proposes a novel method (Correlation Mode Decomposition, CMD) to reduce the dimension of learning dynamics of neural networks, by leveraging correlations between learning dynamics. Using correlations between (the entire histories of) learning dynamics of each parameters, the proposed method first divides the set of parameters into several clusters (called modes),  identifies one representative parameter from each modes, and then represents the other parameters in each mode by scaling the representative and adding a bias. The paper also provides an online version of the dimensionality reduction (Online CMD) that can be used without memorizing the history of parameters during training, which still requires training of all parameters, and a parameter-efficient version (Embedded CMD) that enables us to reduce the number of trainable parameters gradually during training. The paper also propose to use the dimensionality reduction of learning dynamics for distributed learning. The paper empirically shows the superiority of CMD against standard training and the state-of-the-art methods of dimensionality reduction.", "questions": "See Weaknesses.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. The procedure of CMD seems reasonable and also novel in dimensionality reduction of learning dynamics.\n2. It may be also novel that their proposal to use the dimensionality reduction for distributed training, but less confident since I'm not an expert in this area.\n3. Experimental results (Figure 3) shows a surprising result that Online/Embedded CMD outperforms full SGD on CIFAR-10, which seems somewhat contradictory because Online/Embedded CMD was designed to approximate the full SGD.", "weaknesses": "1. There are many unclear points in experimental results/figures.\n    1. In each mode block in Figure 1 (Left & Middle), correlation between most of parameters tends to be less than 1.0, which does not satisfy the hypothesis behind the proposed method: `Any two time trajectories u, v that are perfectly correlated can be expressed as an affine transformation of each other`.\n    2. The y-axis in Figure 1 (Right) is unclear.\n    3. What are the different/common points between CMD and DMD? The paper highlights CMD vs DMD in Figure 1, but any description for DMD is not provided.\n    4. Any theoretical evidence of Hypothesis 1 is not provided.\n    5. `Figs. 1, 6 demonstrate that Eq. (2) accurately represents the dynamics using a notably small set of modes` in Section 3.1 is overclaimed due to W1-1 and W3.\n    6. Accuracies plotted in Figure 2 (Left) seem inconsistent with the corresponding accuracy in Table 1. It is also weird that test accuracy of CMD is consistently higher than training accuracy of CMD in Figure 2.\n    7. The author claimed that `we observed that the reference trajectories can be selected once during training, and do not require further updates`. However I could not find such results in the paper.\n    8. In Table 1, it is weird that CMD significantly outperforms full training (GD) although CMD is designed to approximate GD. At least, CMD should be worse than GD in training accuracy if CMD behaved as a regularizer as the authors claimed. If CMD outperforms GD even in training accuracy, I'm concerning that there should be some leakage or bug in experiments.\n2. Experiments are limited on a single toy dataset CIFAR-10. I'm concerning whether the proposed method still works well on more large-scale, difficult learning tasks.\n3. Since CMD identifies only single representative for each cluster, the method only leverage proportional relationships between parameter dynamics, which may lead to very limited expression power of the reduction method, especially in more complex learning scenarios.", "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": "KXeKwqfeVf", "reviewer": "Reviewer_JHvF", "summary": "The authors make a novel hypothesis about the way in which deep neural network weights dynamically evolve during training: \"weights can be clustered into a very few highly correlated groups\" (Hypothesis 1). The authors illustrate support for this hypothesis by examining the training of several different types networks on several different tasks, using post-hoc correlation based clustering. The success they found on these examples motivated them to develop an online method, which reduced computational demands and led to better accuracy than state-of-the-art low dimensional training methods. Lastly, the authors provided initial results on improving federated learning (and lowering computational costs) via their method.", "questions": "All my questions are posed in the section above.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. This paper was well motivated and written. It was (for the most part) very easy to follow.\n\n2. The hypothesis of highly correlated modes fits well within existing work (which the authors cite well), but is a highly novel discovery. This makes it impactful and interesting. \n\n3. The results the authors achieved are impressive. They performed better than standard methods of training, and achieved state-of-the-art results among low-dimensional training methods. They tested on a variety of tasks (including federated learning, which - the authors say - makes it the first use of dynamical modeling to reduce communication overhead). \n\n4. The Appendices are full of examples, pseudo-code, clarifying remarks, and extra details. This is a very packed paper.", "weaknesses": "I found no major weaknesses of this paper. \n\nHere are a few points that whose clarification will enhance the quality of the paper: \n\n1. The notation used in Sec. 3.3 (Online CMD) is a little difficult to follow. Additionally, because the use of online CMD to train DNNs had been foreshadowed in earlier parts of the paper, I was confused by the lack of details on how the individual modes were trained - which was answered in Sec. 3.4. Explicitly mentioning, in Sec. 3.3., that more details are coming on the actual training are coming in Sec. 3.4, would help remove this momentary confusion. \n\n2. I did not understand the comment \"Even though $\\tilde{A}_m$ is re-evaluated at each time-step $t$, each $a_i, b_i \\in \\tilde{A}_m(t)$ are fixed scalars.\" (Sec. 3.3). Is this pre-empting the point made in Sec. 3.4 that, once embedded, a weight's $a_i, b_i$ are frozen? \n\n3. I felt like Sec. 3.5 (CMD Accuracy and Loss Landscape Visualization) is too quickly presented. What is the main takeaway? That you can use CMD to visualize the landscape, or that CMD training does a good job of finding the optimal parameters? Maybe this would be better in the Appendix?\n\n4. I think the work on federated learning is very interesting and a great application of the approach. That being said, I did feel like the explanation of the approach was a little rushed (possibly due to lack of space) and Eq. 16 was a bit confusing. Improving the discussion around Sec. 4 would be helpful (and could be aided by putting Sec. 3.5 in the Appendix to get more room). \n\nMinor Comments: \n\n1. The acronyms in Table 2 should be clarified (I assume Mom is momentum but what is SCH?). \n\n2. Could you use DMD/Koopman methods to predict how the individual modes are going to evolve, thereby reducing the computational cost even more? \n\n3. The comparison of DMD with CMD in Fig. 1 is insightful. Recent work has found that using Graph Neural Networks can lead to better approximations of the Koopman operator, specifically for the training of DNNs [1]. Do you think the correlated modes you discover might explain this improved success (e.g., the GNN is able to better learn the correlated modes)?\n\n4. There are several places where the first quotation marks in a quote are backwards. This can be remedied by using `` instead of '' in Latex.  \n\n5. In Fig. 2, CMD is compared to GD. Should this be SGD?\n\n[1] https://arxiv.org/abs/2305.09060", "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)