Paper: Transferring Learning Trajectories of Neural Networks (ICLR-2024)

Reviewer comment (Theory):
There is a lack of motivation albeit the promising results. It is a lack of the evidence whether having the same trajectory between tasks is always good.

Author response:
One of our motivations is to utilizing existing learning trajectories for subsequent trainings. For example, as we demonstrated in Section 4.3, we believe that learning transfer leads to low-cost updating of the foundation model within fine-tuned models in future. Also, we provided Theorem 1 as the theoretical evidence of similarity of learning trajectories in simple 2-MLP case. Aside from whether it's good or not, the similarity of the trajectory actually exists at least in initial phase of training of 2-MLP, so we investigated how we can utilize the similarity. We believe that future work in real application would reveal the effectiveness of utilizing the similarity of trajectories.

Reviewer comment (Experiments):
To experimentally prove that an initialization or architecture affects the similarity more than the dataset, it is necessary to verify it on more datasets.

Author response:
Thank you for pointing out it. We have added additional empirical analysis of Assumption (P) in Appendix E.8. If we focus on the results on Cars/CUB for Conv8, ResNet18 and ResNet34, we can see that the results largely depends on model architectures (whether it is a ResNet family or not) rather than datasets themselves. Moreover, from a theoretical viewpoint, Theorem 1 guarantees that datasets or distributions do not matter for 2-MLP. This is another example of architecture-dependence rather than dataset-dependence. Of course, we also observed some dataset-dependences such that CIFAR-10 vs CUB/Cars/ImageNet for Conv8, or CUB/Cars vs ImageNet/iNat2017 for ResNets. However, overall trends still seem to depend on architectures rather than datasets.

Reviewer comment (Evaluation):
Does it show the good performance to make the model to have the same trajectory in tasks that are not related to each other?

Author response:
We sincerely note that, in contrast to transfer learning, we did not consider "tasks that are not related to each other" in our learning transfer problem. More precisely, we consider **a single task and multiple initializations** in learning transfer, in contrast to multiple tasks and a single (pre-trained) initialization in transfer learning. Also, as observed in Section 4.3, we believe that updating the foundation model within fine-tuned models is the most promising application. More precisely, when the source trajectory is trained from poorly generalizing initialization (and thus inherits poor generalization), it acquires better generalization after transferred onto more generalizing initialization. Therefore, **even if the trajectory is same, the difference between source/target initializations may lead to better generalization capability**.

Reviewer comment (Evaluation):
Are these transferred models' ensembles better than scratch ensembles? What is it like from an ECE perspective?

Author response:
We have added such evaluations of ensembles in Table 4 in Appendix E.5. In summary, we cannot obtain any informative results from ECE (Empirical Calibration Error) evaluations since ECE metric cannot distinguish between standard-ensemble models and single models. This is possibly because we take an ensemble of only three models. However, we instead introduce another metric "matching probability", which is the probability that a prediction from each member of the ensemble matches the ensembled prediction, i.e., it measures diversity of each member in ensemble. Through this metric, we observed that the diversity in ensemble of transferred models is the same level as that of standard ensemble.

Reviewer comment (Theory):
In Sec. 3.4, there are an explanation why the transfer of the linear trajectory is more stable and has less variance than the transfer of the actual one. The authors explain that it may be because the actual trajectory contains noisy information. I think the theoretical or empirical evidence is necessary to support the explanation.

Author response:
We have added more analysis on this hypothesis in Appendix E.10. Although we cannot prove something that noises in the actual trajectory causes the instability, we provide an empirical evidence of the instability. More precisely, we evaluated the L2 losses of gradient matching (Eq. 6 in the main text) for actual/linear transfer. We observed that the loss for actual transfer is constantly accumulated from the beginning, while the loss for the linear one is suppressed in the early phase. This implies that each gradient direction along the actual trajectory is too diverse from each other to be aligned, due to its stochastic nature.

Reviewer comment (Reproducibility):
I think we need to do the work shown in Fig. 5 to select the optimal parameter. Then shouldn't we put this process into the pseudo-code?

Author response:
As we noted in Appendix D.3.3, there are only two hyperparameters: batch size B and trajectory length T. The batch size is determined based on the dataset size or number of classes. The trajectory length is determined based on the number of epochs for the source trajectory. However, our method works almost similarly in other hyperparameter values. (We have added ablation results on T in Appendix E.5.) Thus basically they should be simply chosen based on the available computational budget.

Reviewer comment (Experiments):
One potential usage of the proposed method suggested by the authors is to transfer the update of a foundation model to its fine-tuned versions. However, all experiments are limited to network architectures of relatively smaller scale, and to the cases where fine-tuning task shares exactly the same number of classes as the pre-training task, which differs from the realistic use-case of foundation models, which are of typically larger scale, and are used for various down-stream tasks with task-specific heads.

Author response:
> 1. One potential usage of the proposed method suggested by the authors is to transfer the update of a foundation model to its fine-tuned versions. However, all experiments are limited to network architectures of relatively smaller scale, 

Due to our limited resources, we could not add experiments for foundation models. However, our theoretical result (Theorem 1) provides an evidence that learning transfer works more effectively in the case of networks with more neurons. Intuitively, more neurons give us more degree of freedom to match them by some permutation. Also, we emphasize that our paper should be considerred as the first step to the proposed new problem, where we demonstrated that learning transfer is indeed possible for MLPs/CNNs/ResNets and promising for updating pre-trained models (Section 4.3). Extending our approach to more complex architectures such as Transformers or Diffusion Models needs more work (for e.g., how to exploit the symmetry of self-attention, how to deal with the inverse diffusion process, etc), so it remains for future work.

> and to the cases where fine-tuning task shares exactly the same number of classes as the pre-training task, which differs from the realistic use-case of foundation models, which are of typically larger scale, and are used for various down-stream tasks with task-specific heads.

Sorry for the confusion. We notice that we have forgotten to remove the statement "a 10-classes subset of CIFAR-100" in Section 4.1 from pre-submitted version of our paper, since we are using CIFAR-100 with 100 classes in the submitted version. Thus, we sincerely note that, in all experiments of fine-tuning scenario, fine-tuning task **does not** share the same number of classes as the pre-training task. Indeed, before transferring learning trajectories, we replace the final layer of the given (target) pre-trained initialization by a task-specific head, as well as the standard way. Sorry again for confusing you by the misprint.

Reviewer comment (Experiments):
The authors claim that method can accelerate the training of an ensemble of neural networks. Although the computational cost of the proposed method is briefly described in the appendix, there is no figure or table systematically comparing the cost with traditional training / fine-tuning approaches. More crucially, no experiment compares the performance of the ensemble obtained with GMT / FGMT and traditional approaches. One concern is that, transfering one (or a limited number of) source trajectory, the diversity of the resulted target networks is limited (this is partially endorsed by the landscape visualization in Figure 7), which could hurt the performance of the ensemble.

Author response:
> 2. The authors claim that method can accelerate the training of an ensemble of neural networks. 

We sincerely note that we have made no such a claim on ensemble. Sorry for the confusion. At least, we did not intend to accelerate training of ensemble by the proposed method as is. Indeed, the results of ensemble is not always as good as the standard ones as provided in Appendix E.5. However, we believe that future method may achieve such acceleration as discussed in Introduction due to the lightweight nature of learning transfer.

> Although the computational cost of the proposed method is briefly described in the appendix, there is no figure or table systematically comparing the cost with traditional training / fine-tuning approaches. 

We have added a table on wall-clock time for transfer in Appendix E.6. However, we would like to emphasize that the results are highly unreliable because computing nodes are shared with other research
groups and has non-negligible I/O delay due to its distributed storage. Also our implementation is not so optimized. This is because our main purpose in this paper is to demonstrate the possibility and potential applicability of learning transfer, as we claimed in the last lines of Introduction. Actual speed-up for ensembles or foundation-updates is promising challenges for future research.

> One concern is that, transfering one (or a limited number of) source trajectory, the diversity of the resulted target networks is limited (this is partially endorsed by the landscape visualization in Figure 7), which could hurt the performance of the ensemble.

In Table 4 of Appendix E.5, we have added results on diversity of transferred models in ensemble, compared to standard ensemble. Although we evaluated calibration error as Reviewer EPpK suggested, it fails to evaluate the diversity because it cannot distinguish between single models and standard ensembles (possibly because we take an ensemble of only three models). Instead, we introduce another metric "matching probability", which is the probability that a prediction from each member of the ensemble matches the ensembled prediction, i.e., it measures diversity of each member in ensemble. Through this metric, we observed that the diversity in ensemble of transferred models is the same level as that of standard ensemble.

Reviewer comment (Experiments):
Although the theoretical analysis abounds, it would be much more convincing to show these statements — Assumption (P), Theorem 3.1, and Lemma 3.2 — hold in real experiments.

Author response:
> Lack of experiments verifying the Assumption (P), Theorem 3.1, and Lemma 3.2

For Assumption (P) and Theorem 3.1, we have added additional results in Appendix E.8 with more datasets and architectures. Overall trends are consistent with Figure 3 in the main text, which also provides empirical analysis for Assumption (P) and Theorem 3.1.

For Lemma 3.2, we have added empirical analysis in Appendix E.9. We evaluated the cosine similarity between $\theta_{2,\pi_s}^s$ and $\theta_{2,\pi_{s+1}}^s$ during GMT or FGMT, which are expected to be similar to each other by Lemma 3.2. The results show that Lemma 3.2 (with $t=s$ in the statement) actually holds in the early phase of transfer, and also holds in the later phase of FGMT, due to its stability as discussed in Section 3.4.

Reviewer comment (Experiments):
It is not clear how \( T \), the number of time steps, or rather, the number of samples taken from the source trajectory, influences the performance of the transfer result. It seems that \( T \) does not really matter in the Naïve baseline and Oracle baseline, where the parameter-aligning permutation \(\pi\) remains the same across time steps, which 1) would be good to be verified by the authors in the main text and 2) makes it interesting to explore the value of \( T \) that GMT / FGMT requires to have good performance, because the computational cost is proportional to \( T \).

Author response:
> It seems that T does not really matter in the Naïve baseline and Oracle baseline, ..., which 1) would be good to be verified by the authors in the main text 

Thanks for pointing out it. We have added such explanation in Baselines paragraph (Section 4).

> 2) makes it interesting to explore the value of T that GMT / FGMT requires to have good performance, because the computational cost is proportional to T.

Thanks for the suggestion. We have added some experiments with varying T in Appendix E.7. The results indicate that too small T (such as T=5) makes the transfer a bit worse, but it stops improving validation accuracy beyond some T (such as T=15 or 20 in our experiments). In summary, the choice of T slightly influences the results of our algorithm if T is enough size.

Reviewer comment (Experiments):
In the "Generalization ability" part in section 4.3, the experiments and figures should be further clarified. From Figure 7, I assume that the task is to fit the CUB or Cars dataset (for the previously chosen 10-class subset), and the two "Standard" curves are for the target trajectories, which are transferred through FGMT to initialization pretrained with ImageNet-10%, leading to the green curves, but the explanations are really not clear. However, the main text explicitly states that the target initialization \(\theta_2^0\) is "pre-trained on full ImageNet", which implies only *1* (instead of *2*) possible combination for FGMT for each fine-tuning dataset: starting from ImageNet-10%-pretrained, transferring the finetuning trajectory which starts from the ImageNet-Full-pretrained initialization). Another example is, I am not sure what validation the authors refer to when they mention that the ImageNet-Full pretrained initialization has "validation accuracy \(\approx\) 72%" and the ImageNet-10% one has "validation accuracy \(\approx\) 50%".

Author response:
> In the "Generalization ability" part in section 4.3, the experiments and figures should be further clarified.

Sorry for the confusion. We have added more explicit explanations in Section 4.3 in the latest revision. Briefly, we investigated the fine-tuning scenario (i.e., transferring a learning trajectory starting from pre-trained models) where the source initialization $\theta_1^0$ has worse generalization ability (pre-trained on 10% subset of ImageNet) than target initialization $\theta_2^0$ (pre-trained on full ImageNet). The learning trajectory starting from the former initialization should have also worse generalization. However, just by transferring the trajectory onto the latter initialization, the transferred trajectory obtains better generalization inherited from the latter initialization (even though it has never been trained on the latter initialization!).

> Another example is, I am not sure ... "validation accuracy \approx 72%" and the ImageNet-10% one has "validation accuracy \approx 50%".

These are validation accuracies of the pre-trained models on ImageNet. We have added this explanation in the main text.

Reviewer comment (Theory):
Data augmentation is widely used for the training of DNNs, which essentially (often randomly) modifies the dataset. Does this violate the assumption that the dataset is the same in the problem statement?

Author response:
> Q1. Data augmentation is widely used for the training of DNNs, which essentially (often randomly) modifies the dataset. Does this violate the assumption that the dataset is the same in the problem statement?

We treat data augmentation as just noise in gradients as well as mini-batch order. Basically it draws samples from the same distribution, so it should not matter.

Reviewer comment (Experiments):
Even gradient-based, many optimization methods for DL can have updates very different from SGD (for example, Adam). Considering that the parameter alignment is done with gradient matching in GMT and FGMT, does the choice of optimization method changes the effect of GMT / FGMT? Ablation study regarding this could be worthwhile to add.

Author response:
> Q2. Even gradient-based, many optimization methods for DL can have updates very different from SGD (for example, Adam). Considering that the parameter alignment is done with gradient matching in GMT and FGMT, does the choice of optimization method changes the effect of GMT / FGMT? Ablation study regarding this could be worthwhile to add.

Thank you for the interesting suggestion. We have added some results with Adam in Appendix E.4. The results show that (1) during transfer it fails to achieve high accuracy in contrast to SGD cases, (2) but it still accelerates a bit the subsequent training. We hypothesize that this is related to the phenomenon that MLI (Monotonic Linear Interpolation) fails in Adam cases as previous work [1] have observed, which indicates the trajectory of Adam is more complex than that of SGD. Addressing this difficulty with Adam is also an interesting direction in future work.

[1] Lucas et al. "On monotonic linear interpolation of neural network parameters" (ICML'21)

Reviewer comment (Evaluation):
Why does the performance deteriorate in Figure 5d-5f? It is simply acknowledged without analysis and discussion.

Author response:
> Q3. Why does the performance deteriorate in Figure 5d-5f? It is simply acknowledged without analysis and discussion.

In the case of pre-trained initializations, we can see that Assumption (P) starts to fail in very early iterations in Figure 3. This causes the performance degradation in Figure 5d-5f, which are transfer between pre-trained initializations. Addressing this challenge also remains as future work.

Reviewer comment (Presentation):
In Figure 7, to better understand the effect of FT and FGMT, where would the FGMT-only parameters be in the landscape? What about permuted-source + FT?

Author response:
> Q4. In Figure 7, to better understand the effect of FT and FGMT, where would the FGMT-only parameters be in the landscape? What about permuted-source + FT?

For the first question, there are two reasons: (1) since the accuracy of FGMT-only is less than source or target parameters, plotting FGMT-only parameter gives us no insightful information about their relationships in loss landscape, (2) in Section 4.3, we aimed to investigate the nature of FGMT+FT parameter since it is more important than FGMT-only parameter in practice. For the second question, since permuting a NN does not change its output or accuracy, we do not need to train the permuted-source further. If we did so, it would not change the loss landscape property of the permuted-source parameter.

Reviewer comment (Experiments):
The empirical evaluation of the algorithm is conducted on a limited set of benchmark datasets, and it is unclear how well the algorithm would perform on other types of datasets or in real-world scenarios.

Author response:
We have added additional experiments on more large-scale and complex datasets (SUN397 and iNaturalist2017) than Cars and CUB in Appendix E2. The results show that (1) transfer-only achieves less accuracy compared to Cars/CUB cases but (2) it still accelerates the subsequent training after transfer a bit, which indicates the future possibility of improving our method for complex datasets. Also, we would like to emphasize that Theorem 1 does not restrict the type of datasets or distributions at least in 2-MLP case. Thus we believe that learning transfer will be effective for real-world datasets if future research continues to improve methods.

Reviewer comment (Experiments):
The paper assumes that the source and target tasks are related, and it is unclear how well the algorithm would perform when the tasks are not such related.

Author response:
We sincerely note that, in contrast to transfer learning, there are no "source and target tasks" in our learning transfer problem. More precisely, we consider **a single task and multiple initializations** in learning transfer, in contrast to multiple tasks and a single (pre-trained) initialization in transfer learning. The difference of the problem itself is one of our contributions, as other reviewers acknowledge it. Nevertheless, it would be an interesting future direction to investigate how to extend our problem with **multiple tasks** and multiple initializations, which may integrate transfer learning and learning transfer.

Reviewer comment (Evaluation):
The paper does not provide a detailed analysis of the computational cost of the algorithm, which may be a concern for large-scale neural networks.

Author response:
In Appendix D.3.4, we briefly discussed the computational cost of our algorithms. Also we have added wall-clock time for learning transfer in Appendix E.6. The results also include the comparison between ResNet18 and ResNet34, from which we can see how computational cost grows for large networks in real experiments. For more theoretical aspects, see also answer to Q5. However, it is worth noting that **our main contributions are (1) problem formulation, (2) derivation of algorithm (3) empirical demonstration of learning transfer (4) analysis of transferred solutions**, and thus not an actual speed-up in real world applications, as we explicitly claimed in Introduction.

Reviewer comment (Experiments):
How sensitive is the algorithm's performance to the assumption that the source and target tasks are related, and how well does it perform when the tasks are unrelated?

Author response:
We sincerely note that, in contrast to transfer learning, there are no "source and target tasks" in our learning transfer problem. More precisely, we consider **a single task and multiple initializations** in learning transfer, in contrast to multiple tasks and a single (pre-trained) initialization in transfer learning. The difference of the problem itself is one of our contributions, as other reviewers acknowledge it. Nevertheless, it would be an interesting future direction to investigate how to extend our problem with **multiple tasks** and multiple initializations, which may integrate transfer learning and learning transfer.

Reviewer comment (Experiments):
How does the proposed algorithm compare to other methods for transferring learning trajectories, such as fine-tuning or transfer learning?

Author response:
We sincerely note that, in contrast to transfer learning, there are no "source and target tasks" in our learning transfer problem. More precisely, we consider **a single task and multiple initializations** in learning transfer, in contrast to multiple tasks and a single (pre-trained) initialization in transfer learning. The difference of the problem itself is one of our contributions, as other reviewers acknowledge it. Nevertheless, it would be an interesting future direction to investigate how to extend our problem with **multiple tasks** and multiple initializations, which may integrate transfer learning and learning transfer.

Fine-tuning and transfer learning are not designed for solving our problem, and it is also not obvious how to exploit them in our problem.

Reviewer comment (Experiments):
How well does the algorithm perform on datasets that are not included in the empirical evaluation, and how does its performance compare to other methods on these datasets?

Author response:
We have added additional experiments on more large-scale and complex datasets (SUN397 and iNaturalist2017) than Cars and CUB in Appendix E2. The results show that (1) transfer-only achieves less accuracy compared to Cars/CUB cases but (2) it still accelerates the subsequent training after transfer a bit, which indicates the future possibility of improving our method for complex datasets. Also, we would like to emphasize that Theorem 1 does not restrict the type of datasets or distributions at least in 2-MLP case. Thus we believe that learning transfer will be effective for real-world datasets if future research continues to improve methods.

Reviewer comment (Experiments):
Can the algorithm be extended to handle more complex neural network architectures?

Author response:
The algorithm only exploits the permutation symmetry of intermediate neurons in feed-forward MLPs, so in principle our method can be applied to many network architectures that can be seen as a restricted version of MLPs like CNNs. However, it remains as an important future work how to deal with network architectures beyond feed-forward MLPs, such as Transformers, which involve additional symmetery from their attention structure.

Reviewer comment (Evaluation):
How does the computational cost of the algorithm compare to other methods for transferring learning trajectories, and how does it scale with the size of the neural network and the number of training epochs?

Author response:
The algorithm mainly constitutes of gradient computation (GC) and gradient matching (GM). Roughly speaking, for the neural network with $N$ neurons, GC requires $O(N^2)$ computation steps as well-known with backpropagration, and GM requires $O(N^3)$ steps with Hungarian algorithm. Also the algorithm does not depend on the number of epochs, but rather on the length of trajectory T and batch size B, which should be chosen by computational budgets. See also answer to W3.