Paper: Enhancing Neural Training via a Correlated Dynamics Model (ICLR-2024)

Reviewer comment (Writing):
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, there was confusion due to 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 on the actual training are coming in Sec. 3.4 would help remove this momentary confusion.

Author response:
**Weakness 1** - To improve the flow of the paper, we added the following opening paragraph to Sec. 3 (it includes the previous opening of subsection 3.2 and further details).

"We present several approaches to apply our method, starting by discussing the underlying correlated dynamics (Sec. 3.1). Followingly we introduce the basic CMD version, called post-hoc CMD, which analyzes dynamics after they occur (Sec. 3.2). We then present the efficient online CMD (Sec. 3.3) and embedded CMD (Sec. 3.4) which embeds modeled dynamics into the training process. In the last subsection we demonstrate accuracy landscape visualization using CMD."

Reviewer comment (Writing):
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) was not understood. It is unclear if this is pre-empting the point made in Sec. 3.4 that, once embedded, a weight's $a_i, b_i$ are frozen.

Author response:
**Weakness 2** - Below we provide a detailed explanation about this point. Since it is not very important, yet somewhat confusing - we decided to remove it from the manuscript to avoid the misunderstanding of future readers.

Detailed explanation:

The message here is that the coefficients are always used as static scalars for modeling. These coefficients, though subject to iterative updates during the training process as detailed in Section 3.3, are consistently employed as static scalars. 

**Scalar Nature of CMD Coefficients:** In all CMD modelings, the coefficients \(a_i, b_i\) are used as scalars throughout the process. Even when we model the dynamics at different epoch ranges (e.g., epochs 1-100 and 1-200), which results in different coefficient sets, each set is a set of scalar coefficients.

**Iterative Updates in Online CMD:** In the online CMD approach, these coefficients are updated iteratively throughout the training process. However, it's crucial to understand that at each iteration, the coefficients \(a_i, b_i\) are re-evaluated and assigned new scalar values. They do not evolve into or represent time-dependent trajectories such as \(w_i(t)\). Each update is a discrete change, resulting in a new set of scalar values.

**Potential Misconceptions:** The iterative update process in online CMD could suggest to readers a trajectory-like behavior for these coefficients, as they are updated through time (and trajectories are time-dependent). However, this is not the case. Despite the dynamic nature of the updates, the coefficients \(a_i, b_i\) do not transition into trajectories; they are updated as distinct, individual scalar values at each time step. 

For instance, if at time t, we would like to approximate the dynamics at time interval [t1, t2]:t2<t, then we would use the coefficients obtained at time t (the latest available time).

Reviewer comment (Presentation):
Sec. 3.5 (CMD Accuracy and Loss Landscape Visualization) feels too quickly presented. It is unclear what the main takeaway is — whether that CMD can be used to visualize the landscape or that CMD training effectively finds the optimal parameters. This section might be better suited for the Appendix.

Author response:
**Weakness 3** – The visualization options available when using CMD introduced in Section 3.5 are quickly and succinctly presented.  The main takeaway is that CMD could be used for visualization of the loss/accuracy landscape, as another interesting application that CMD offers. Appendix G is an extended section presenting different visualization results for different scenarios. Although most of the information is available only in the appendix, we decided to shortly present this idea in the main paper, so that interested readers will know to look for the expansion in the Appendix.

Reviewer comment (Writing):
The work on federated learning is very interesting and a great application of the approach; however, the explanation of the approach feels a little rushed (possibly due to lack of space) and Eq. 16 is a bit confusing. Improving the discussion around Sec. 4 would be helpful and could be aided by moving Sec. 3.5 to the Appendix to free up space.

Author response:
**Weakness 4** – We have added an explanation regarding Eq. (16) (and omitted details on the visualization section). Additionally, in Appendix F.3 we have added a new discussion that outlines the advantages of using dynamics modeling for efficient communication in Federated Learning, beyond the scope of CMD.  It highlights the potential for improved model interpretability and adaptability to future dynamics modeling techniques, enhancing efficiency, scalability, and understanding in Federated Learning. See also Appendices F.1,F.2 for technical details and communication overhead complexity of our FL experiments.

Reviewer comment (Experiments):
Consider whether DMD/Koopman methods could be used to predict how the individual modes are going to evolve, thereby reducing the computational cost even more.

Author response:
**Minor Comment 2** - The use of DMD/Koopman methods in this context is interesting. If it works - it can open the possibility to extrapolate the reference trajectories forward in time, and in turn forward the full set of trajectories.

Reviewer comment (Evaluation):
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]. It would be interesting to know whether the correlated modes discovered in this work might explain this improved success (e.g., the GNN is able to better learn the correlated modes).

Author response:
**Minor Comment 3** - Thank you for bringing this research to our attention. We were not acquainted with this approach for accommodating highly nonlinear dynamics with DMD. In fact, it may alleviate concerns we raise in Appendix H.2 (DMD in Neural Networks). We added it to the Previous work section. The question raised is interesting. We think there might be a connection, we note there are still some differences between this work and our method. Mainly, in that paper the model dynamics are approximated using a set of  eigenfunctions (as in other DMD approaches) and in CMD the dynamics are approximated by reference dynamics from the original trajectories. However, We do think that the use of a GNN might assist learning eigenfunctions that are more aligned to the system (similarly to our method). We anticipate additional fruitful research in this direction.

Reviewer comment (Writing):
There are several places where the first quotation marks in a quote are backwards. This can be remedied by using `` instead of '' in LaTeX.

Author response:
**Minor Comments 1, 4, 5** – In Figure 2, 'GD' in the legends of the plots stands for SGD. A clarification has been added to the figure caption. Additionally, we added an explanation for the acronyms in Table 2 - 'Mom' stands for momentum and 'SCH' stands for leaning rate scheduler. All the backward quotation marks have been fixed as well. Thank you for pointing these issues out.

Reviewer comment (Presentation):
In Fig. 2, CMD is compared to GD. It should be clarified whether this should be SGD instead.

Author response:
**Minor Comments 1, 4, 5** – In Figure 2, 'GD' in the legends of the plots stands for SGD. A clarification has been added to the figure caption. Additionally, we added an explanation for the acronyms in Table 2 - 'Mom' stands for momentum and 'SCH' stands for leaning rate scheduler. All the backward quotation marks have been fixed as well. Thank you for pointing these issues out.

Reviewer comment (Theory):
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`.

Author response:
**Weakness 1.1** - Eq. (2)'s approximation of mode dynamics varies in preciseness. Consequently, correlation varies as well, as indicated by Fig. 1 as you noted. This raises questions about Eq. (2)’s representation of the original dynamics upon less-than-perfect correlation. To answer these questions we added an experiment (Fig. 6 in Appendix A.1) that visualizes the original dynamics alongside their representation while displaying average mode correlation. It shows good approximation of the general dynamics even upon less-than-perfect-correlation.

Reviewer comment (Presentation):
The y-axis in Figure 1 (Right) is unclear.

Author response:
**Weakness 1.2** - These are the weight values, added the label “weight” to the y-axis, and the label “epoch” to the x-axis.

Reviewer comment (Theory):
Any theoretical evidence of Hypothesis 1 is not provided.

Author response:
**Weakness 1.4** - We agree additional investigations are required to substantiate our findings with a corresponding theory. From a theoretical point of view - it is known that NNs exhibit large redundancies, hence low-dimensional representations should exist. While our correlation-based model is novel, it aligns with previous research, see for instance [3], which strengthens its utility. From a practical point of view - acting upon our hypothesis indeed results with reduced redundancy while exhibiting high performance.

Reviewer comment (Evaluation):
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.

Author response:
**Weakness 1.6** - 

**a.** Fig. 2 and Table 1 present distinct experiments, each fine-tuning a different pre-trained Vit-b-16 model: Fig. 2 uses post-hoc CMD with a model pre-trained on JFT-300M (dataset of 300M samples), while Table 1 employs online CMD with a model pre-trained on ImageNet1K (1.2M samples). We acknowledge that this should have been clearly stated, and revised the captions of Fig. 2 and Table 1 to include this distinction. We additionally added Fig. 18 in Appendix D, an experiment that aligns with the experiment setting of Table 1.	

**b.** The title should read “Baseline SGD” and not just “training” (as in “regular training”). Thank you for noticing, it is fixed now. As Table 1’s caption suggests, this whole table solely consists of test accuracy results. For further details on these experiments see Tables 5, 6 in the Appendix D.

Reviewer comment (Experiments):
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.

Author response:
**Weakness 1.7** - In order to support our claim we added Fig. 16 to Appendix C.5. This figure illustrates online CMD's performance with varying warm-up durations, indicating that extended warm-up periods do not impact the outcomes. Table 4 in Appendix C.5 implicitly shows this claim as well: This table compares post-hoc CMD with online CMD. The former selects the reference trajectories based on the full training trajectories and the latter selects reference trajectories using a short 20-epoch warm-up period. The table demonstrates comparable performance between the two methods over various numbers of modes. 
To conclude, Both the added experiment in Fig. 16 and Table 4 in the Appendix support that there is no need to wait until the end of training before selecting the reference trajectories. References to Table 4 and Fig. 16 have been added in the main body of the paper to support the claim.

Reviewer comment (Experiments):
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.

Author response:
**Strength 3 and Weakness 1.8** - The inherent boost in performance was initially surprising to us as well. See the Results section, under “CMD as a regularizer”, where we attribute this (beneficial) phenomenon to smoothing of the trajectories, i.e. reduction of parameter fluctuations - which are typical upon SGD. The boost in train and test accuracy upon such smoothing was shown and explained in previous work, see for instance [1], [2], which introduced Exponential Mean Averaging (EMA) and Stochastic Weight Averaging (SWA) respectively. We compare ourselves to [1], [2] in Table 2.

Reviewer comment (Experiments):
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.

Author response:
**Weakness 2** - Regarding the concern about limited datasets and transferability to large-scale difficult tasks:


**a. Data Scope:** Besides Cifar10, our research also incorporates CelebA-HQ for detailed image generation and PASCAL VOC 2012 for image segmentation, offering a variety of datasets.

**b. Architectural Scale:** Our work encompasses a range of architectures. This includes smaller models like LeNet, mid-scale ones such as ResNets, and larger frameworks like WideReNet (36.5M parameters),  ViT-b-16 (86M parameters). We also experiment with the StarGAN-v2 (50M parameters, composing a complex of Generator, Style Encoder, a Latent-Noise Mapping networks), demonstrating our engagement with large-scale and challenging architectures.

**c. Learning Tasks:** A concern about carry-over to other difficult learning tasks was raised. While our results section does focus on image classification (as is the case in previous work and benchmarks of this domain), we conduct further experiments, showing the effectiveness of CMD in image segmentation and image generation, particularly exploring the dynamics of GAN training. Moreover, we believe that we are the first to test such methods for image generation.

Reviewer comment (Theory):
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.

Author response:
**Weakness 3** - See our answer to **Weakness 2** regarding complex learning scenarios \ difficult tasks. 

Investigating a generalized version, as you suggest (as in Eq. (1)), is a proper direction for future research. We do remark that DMD, which uses a linear combination approach as you proposed, has limitations in expressiveness. This is detailed in our analysis in Appendix H.2, titled “DMD in Neural Networks,” where we discuss the constraints of DMD in representing complex learning dynamics. 

Though simple, our approach, which is data-driven and nonlinear, shows modeling capabilities of complex dynamics in various scenarios. This underscores the ongoing challenge of balancing expressiveness and simplicity in model reduction techniques.

Reviewer comment (Writing):
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.

Author response:
**Weakness 1** – We replaced citations with the preferable \citep, when appropriate.

Reviewer comment (Presentation):
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.

Author response:
**Weakness 2** – To demonstrate results across an extended range of epochs, we added Fig. 18 in Appendix D. This figure illustrates the test accuracy of a Vit-b-16 model over 90 epochs of CIFAR10 training using regular SGD training and Online CMD, and includes examples with other models (GoogleNet and WideRes28-10). These experiments align with the experimental setting of Table 1, namely Vit-b-16 is initially pre-trained on ImageNet1K [1], which contains a little more than 1.2 million images. Figure 2's left plot showcases Vit-b-16 pre-trained on the JFT-300M dataset [2], which encompasses over 300 million images. Given the advanced starting point of this pre-trained model, the training duration is significantly reduced (compared to Fig. 18 and Table 1), rendering it less noteworthy. Therefore, we only applied post-hoc CMD in this scenario. Details regarding the pre-trained models have been added to the captions of Fig.2 and Table 1, as well as Appendices B,D, where the experimental details are described.

Regarding the right part of the middle plot in Figure 2, please refer to Fig. 15 in Appendix B, which features more examples of generated images and quantitative comparisons as well. This will let you get a better grasp of the difference between the methods, as there are qualitative differences between the generated images. StarGAN-v2 in its original setting requires EMA (a training regularization technique) to obtain its good results. In Figs. 2, 15 we demonstrate that CMD produces regularization that can be used in this context as well. A more thorough analysis of CMD regularization is provided in the context of image classification, in the Results Sec. 5, under “CMD as a regularizer”.

Reviewer comment (Presentation):
In the middle plot of Figure 2, there appears to be a visual discrepancy or blur on the right side that might need clarification or correction.

Author response:
Regarding the right part of the middle plot in Figure 2, please refer to Fig. 15 in Appendix B, which features more examples of generated images and quantitative comparisons as well. This will let you get a better grasp of the difference between the methods, as there are qualitative differences between the generated images. StarGAN-v2 in its original setting requires EMA (a training regularization technique) to obtain its good results. In Figs. 2, 15 we demonstrate that CMD produces regularization that can be used in this context as well. A more thorough analysis of CMD regularization is provided in the context of image classification, in the Results Sec. 5, under “CMD as a regularizer”.

Reviewer comment (Experiments):
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.

Author response:
**CMD Algorithm's Design and Implementation:**
Additional details and information regarding the CMD design is provided in Appendix C due to space limitations. We present full psuedo-code algorithms (Appendix C.1), explain how the sampled weights are chosen (Appendix C.2), discuss different methods to divide the model weights to modes (Appendix C.3), etc. Additional references to Appendix C have been added to the methods section to ensure that this information is not overlooked.

**Comparison with Other Methods:**
   Our comparison primarily focuses on P-BFGS, which is the SOTA in low-dimensional training dynamics. This focus is due to the similarities between CMD and P-BFGS in managing training dynamics using a limited number of modes, a departure from earlier methods that utilized thousands. CMD, like P-BFGS, demonstrates exceptional performance with only a few to several dozen modes. 
   
To the best of our knowledge, P-BFGS, published in IEEE Transactions on Pattern Analysis and Machine Intelligence (March 2023), was the first to showcase such performance. Prior to P-BFGS, there were no direct competitors in such extremely  low-dimensional framework. Additionally, to the best of our knowledge, we are the first to propose such modeling for applications in segmentation and generative models, which our work has successfully demonstrated. Other comparisons: We also compare CMD with known regularization techniques like Exponential Moving Average (EMA) and Stochastic Weight Averaging (SWA). We see limited connection with pruning methods and would appreciate further insight into their suggested relevance.

Reviewer comment (Writing):
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.

Author response:
**Weakness 4** – Our work is motivated by previous works related to dynamics analysis (specifically DMD, as mentioned in the Introduction section). Therefore, the nomenclature used in our work is taken from the dynamics analysis area. Particularly, the term 'modes' in our work refers to highly correlated groups of model parameters, as described in Hypothesis 1 in Section 3.1.

Reviewer comment (Experiments):
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.

Author response:
**Weakness 5** – Our work encompasses a range of architectures, demonstrating CMD on larger and complex models as well. This includes smaller models like LeNet-5 (less than 1M parameters), mid-scale ones such as ResNets (up to 36M parameters), and larger frameworks like ViT-b-16 (86M parameters), demonstrating consistency across scales. The CMD parameters used for each experiment (Table 6, Appendix D) show that three model scales (LeNet-5, ResNet18 and ViT-b-16) use the exact same hyper-parameters. We also experiment with the complex StarGAN-v2 architecture (which is a composition of a Generator, a Style Encoder, and a Latent-Noise Mapping network), demonstrating our engagement with large-scale and challenging architectures. These support CMD as a practical method relevant for complex models at scale.

Reviewer comment (Writing):
The motivation behind applying CMD in federated learning seems a bit unclear and could benefit from a more explicit demonstration or explanation.

Author response:
**Weakness 6** – It is true that the motivation should be stated more explicitly, thus we added to Sec. 4’s opening paragraph: “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.”

Reviewer comment (Presentation):
Could you clarify the visual discrepancy observed in the middle plot of Figure 2? What does the blur on the right side represent?

Author response:
Regarding the right part of the middle plot in Figure 2, please refer to Fig. 15 in Appendix B, which features more examples of generated images and quantitative comparisons as well. This will let you get a better grasp of the difference between the methods, as there are qualitative differences between the generated images. StarGAN-v2 in its original setting requires EMA (a training regularization technique) to obtain its good results. In Figs. 2, 15 we demonstrate that CMD produces regularization that can be used in this context as well. A more thorough analysis of CMD regularization is provided in the context of image classification, in the Results Sec. 5, under “CMD as a regularizer”.

Reviewer comment (Novelty):
Could you elaborate on the motivation and rationale behind applying the CMD method in federated learning?

Author response:
**Weakness 6** – It is true that the motivation should be stated more explicitly, thus we added to Sec. 4’s opening paragraph: “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.”

Reviewer comment (Writing):
Authors are encouraged to improve the writing quality of the current manuscript.

Author response:
**Regarding the writing quality:** We made several adjustments throughout the manuscript following the comments here. We are opened for further suggestions. 

For improved readability, an opening paragraph was added to Sec. 3 (Method) so that the reader anticipates the flow mentioned in Strength 3. The opening paragraph of Sec. 4 (CMD for Efficient Communication in Distributed Learning) was added motivation. In Appendix H (which is referenced in the Introduction section 1) we added discussion on: commonalities and differences of DMD and CMD. To better navigate the appendix we added references to Appendix C for details on CMD design, Table 6 for hyper parameters used in the classification models, and Table 4 in the Appendix for performance comparison between online and post-hoc performance.
Small fixes were introduced: added axis labels to the right plots of Fig. 2, improved Table column titles from “(regular) training” to “baseline SGD”, citations have been improved using \citep, additionally to slight rephrasing and editing.

Reviewer comment (Experiments):
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.

Author response:
**Regarding Federated Learning (FL):** We acknowledge these are preliminary results, as emphasized in the abstract. Nonetheless, the domain (low-dim. representation of training dynamics) lacks practical examples, thus it was important for us to show this idea, which is also timely. While we plan to further research this direction with CMD, the efficiency of saving communicated parameters via low-dimensional dynamics modeling is not restricted to CMD. We hope that as this domain will progress, other approaches may take inspiration from our proposal, and use their techniques for FL as well.