LEMON: Lossless model expansion

We thank the reviewers for acknowledging our method is versatile and interesting. We now address the reviewer's concerns:

1. The paper could benefit from experiments on the generalizability of LEMON to other architectures beyond Vision Transformers and BERT models. For example, CNN models are completely unexplored in terms of experiments. I suggest adding model expansion experiments for ResNet18 --> ResNet50 and EfficientNetB0 --> EfficientNetB4 (or other variants).

Thank your for your constructive feedback. In response to your suggestion, we have included experiments with expansion from ResNet-50 to WideResNet-110 in Appendix F.2. We compare LEMON with the best baseline method, i.e., bert2BERT-AKI. Our results indicate that LEMON is able to complete the training process within 60 epochs, which is 66.7% of the time required for training from scratch. Moreover, we observed that the bert2BERT-AKI baseline did not offer a similar acceleration, underscoring LEMON's advantages.

Regarding the choice of models, we did not opt for the expansion from ResNet-18 to ResNet-50 because they utilize different architectural blocks (Basic block vs. BottleNeck block). Instead, we expand ResNet-50 to WideResNet-110 to showcase LEMON’s capabilities involving both depth and width expansion.

2. Sections 5.1 and 5.2 experimentally study the effect of learning rates and schedules. However, the paper lacks a theoretical analysis of why this happens. For example, why does a small learning rate lead to lower final performance? I think only experimental verification is not enough and this requires some theoretical analysis.

We thank the reviewer for highlighting the importance of theoretical analysis in understanding the effects of learning rates and schedules on model performance. Indeed, several theoretical works have attempted to shed light on the relationship between learning rate, training schedules, and model generalization. For instance, He et al. [1] establishes a generalization bound correlated with the batch size to learning rate ratio, and Li et al. [2] demonstrates that networks trained with a large initial learning rate and subsequent annealing can generalize better than networks trained with a consistently small learning rate, based on the learning order of the model.

We believe that our study, primarily empirical in nature, can provide valuable, practical insights into this intricate relationship and can serve as a useful reference for theoreticians in the field.

Recognizing the significance of a combined theoretical and empirical approach, we are committed to incorporating a deeper theoretical analysis into our future research, which may involve collaborations with theoreticians.

3. Can "incremental" model expansion help achieve even better performance? For example, instead of expanding from "Model(small) pretrained --> Model (huge)", would it be better to expand in steps as "Model (small) pretrained --> Model (middle) -->Model (big) --> Model (huge)"

We are grateful for the reviewer’s suggestion regarding the potential of 'incremental' model expansion. Indeed, this concept is an interesting aspect of model scaling and could potentially enhance performance. However, due to its complexity and the constraints of the rebuttal period, we have not yet explored this incremental approach. But we recognize the potential value of this method and are keen to investigate it in our future work.

[1] He, Fengxiang, Tongliang Liu, and Dacheng Tao. "Control batch size and learning rate to generalize well: Theoretical and empirical evidence." Advances in neural information processing systems 32 (2019).

[2] Li, Yuanzhi, Colin Wei, and Tengyu Ma. "Towards explaining the regularization effect of initial large learning rate in training neural networks." Advances in Neural Information Processing Systems 32 (2019).

We are grateful for the reviewer’s appreciation of our work’s simplicity and effectiveness, as well as our study on training configurations. We now address the questions raised:

1. In Figure 7 $c$ and (d), the loss curves for BERT language modeling are still decreasing when training is stopped. It might be better to train the model till convergence to evaluate whether or not the proposed method can have the same performance as the model trained from scratch.

We appreciate the reviewer's observation concerning the training duration in these figures. The decision to stop training when the logarithmic validation MLM loss reaches 1.64 was intentional. This specific loss value was selected because, beyond this point, the benefits of continued training diminish due to the reduced learning rate. Additionally, the choice of a logarithmic scale for loss visualization can sometimes exaggerate the appearance of a decreasing trend, which is important to consider when interpreting these curves.

This choice allows for a more stable comparison across methods, as the validation loss is subject to fluctuations due to the stochastic nature of MLM masks during training/validation.

Regarding the comparison of LEMON's performance with the model trained from scratch, we refer to the results presented in Table 2. In this table, we compare the downstream performance of BERT(12,768) models expanded using LEMON from BERT(6,384)/BERT(6,512) over the complete 165k/132k training steps. The results demonstrate that the models expanded via LEMON outperform the performance of models trained from scratch for the complete 220k steps. This evidence strongly supports the effectiveness of LEMON.

2. Vision Transformer is a pre-norm Transformer, and in BERT language modeling, the authors also used the pre-norm variant. Since the authors claimed compatibility of the algorithm with different variants, it would be better to see the experiments on different variants (at least a post-norm variant) to verify the claim. While in Appendix, the authors showlossless expansion for other variants, it is also important to evaluate the performance metrics.

We are grateful for the reviewer's constructive feedback regarding the evaluation of LEMON with different Transformer variants. To address this point, we have expanded our experiments to include Post-LN BERT models, specifically testing the expansion from BERT(6,384) to BERT(12,768) as follow. We repoort number of training steps needed to achieve a log validation MLM loss of 1.67. (See details in Appendix F.3)

The results demonstrate that LEMON not only outperforms these baselines but also achieves 36.57% less computational savings than training from scratch.

3. Since this work also studies the initialization of model parameters, it might be interesting to compare the proposed idea with other initialization approach, such as Mimetic initialization (https://arxiv.org/abs/2305.09828), also mentioned in the related work.

We appreciate the reviewer’s suggestion to compare LEMON with Mimetic initialization. We recognize that Mimetic initialization addresses a different scenario, primarily focusing on scenarios where no pre-trained (small) models are available. This distinction makes our work and Mimetic initialization orthogonal yet potentially complementary. For example, Mimetic initialization might be used to effectively train the initial small model, which could then be efficiently expanded using LEMON.

At present, the lack of open-sourced code for Mimetic initialization limits our ability to conduct related experiments. However, we are keenly interested in this possibility and plan to explore such comparisons as soon as the code becomes available.

We are grateful for the reviewer’s recognition of our work's comprehensiveness, versatility, and significance. Now we answer the reviewer's questions:

1. After reading this paper, I would like to know more about the practical application of the model expansion method in real-world scenarios. Please provide more examples to illustrate the effectiveness and applicability of the proposed approach.

We appreciate the reviewer's interest in the practical applications of our model expansion method. Our approach can be beneficial in various real-world scenarios:

1. Streamlining model training across scales: Companies often require models of different scales to meet diverse user needs or optimize inference speed. For instance, training models like Llama2 7B, 11B, and 70B typically requires substantial resources when done independently. Our method accelerates this process by allowing the sequential expansion of a smaller model (e.g., 7B) into larger versions (11B, then 70B). This approach significantly reduces computational resources and time.

2. Dataset testing and adaptation: When a new or modified dataset is introduced, it's common practice to test it with smaller-scale models before deploying larger ones. Our model expansion method enables an efficient transition from testing with a small model to full-scale deployment with a larger model.

3. Efficient research prototyping: Researchers often aim to build upon existing open-sourced models by scaling them up. Our method allows for the direct expansion of these pre-existing models into larger versions, facilitating prototyping and experimentation without the need for training large models from scratch.

2. How does the performance change in Figure 6b when using a learning rate larger than the default value?

We appreciate the reviewer’s question regarding the effect of using a maximum learning rate larger than the default value in Figure 6b. In our experiments, increasing the learning rate above the default value led to training instability and divergence. We believe that the default learning rate is often tuned to be close to the 'critical value' – a threshold beyond which the training becomes unstable.

Indeed, this is an interesting topic, and we plan to delve deeper into this topic in our future work.

3. Regarding the significant drop in accuracy during the early stages of training, it may be attributed to the transfer of pre-trained parameters from one local optimum to another. Considering the requirements of certain real-time systems, such performance drop in model accuracy is unacceptable. I think maybe we can smooth out this process by, for example, setting a mask to control the number of trainable parameters in each epoch, gradually transitioning them to avoid a drastic drop in accuracy. Overall, solving this problem could be of significant importance for many real-world applications.

We thank the reviewer for raising an important concern regarding the initial accuracy drop in our method, particularly in the context of real-time systems where stable performance is crucial. Indeed, our method in its current form might not be ideally suited for such applications due to this challenge. And we are actively investigating solutions to this issue.

Regarding the learning rate setting:

(1) Constant or minimally changing learning rate: Our investigations indicate that the initial performance drop is primarily attributed to a sudden increase in the learning rate. In scenarios where the learning rate is constant or undergoes minimal changes, we have observed that this drop is significantly less pronounced. This indicates that our expansion method may be effectively utilized in such settings.

(2) Decreasing learning rate: In cases where the learning rate decreases over time, we propose an interim solution of continuing to use the small/source model at the point of expansion. Once the expanded model gains its expected performance, the system can seamlessly transition to the expanded model, ensuring both stability and the benefits of expansion.

Thank you for highlighting this critical aspect. It forms an important part of our ongoing research, and we are committed to developing more effective solutions to reduce the initial performance drop without compromising on training efficiency.

We thank the reviewer for acknowledging the clarity, organization, and feasibility of our work. We now address the concerns raised:

1. The challenge arising with the ‘symmetry breaking’ is described in the third paragraph of section Introduction: “the expanded model will never gain more capacity than the source model.”This statement raises confusion as training a model with smaller capacity but larger size appears to be of limited value, which incurs greater overhead but achieves limited performance.

We appreciate the reviewer's feedback and the opportunity to clarify the distinction between 'weight symmetry' and 'break weight symmetry.' We acknowledge that our initial explanation might have caused some confusion.

In our paper, 'weight symmetry' refers to a situation where, after neuron replication, the weights associated with these replicated neurons are initially identical and continue to remain identical throughout the training process. Hence, 'weight symmetry' limits the neural network's representation power. To avoid 'weight symmetry,' 'breaking weight symmetry' becomes crucial, where we set the fan-out weights of replicated neurons to be unequal.

This will be more clearly articulated in the revised manuscript, particularly in the Section 1 and 3. The terms 'weight symmetry' and 'break weight symmetry' are highlighted in blue.

2. Have the considered baselines for expansion in Section 6 been confirmed to be lossless? If not, it is necessary to present the gap with the original model. If they are indeed lossless, an analysis should be provided to explain why the proposed method achieves higher validation accuracy compared to AKI, which also breaks symmetry, as shown in Figure 8.

Thank you for your question. The baseline models, namely bert2BERT-AKI and bert2BERT-FPI, are NOT lossless. This is reflected in the MLM loss and accuracy for BERT and ViT, as shown below.

It's important to note that the MLM validation loss may slightly vary from the source model due to the different MLM masks applied in each batch. However, we have carefully verified LEMON's losslessness by comparing the outputs of small and large models across multiple data samples.

To provide a clearer understanding of the effectiveness our model expansion approach, we have added data points in our plots to show performance right after expansion, before any additional training.

3. The results in Table 2 are a bit confusing. It is unclear whether model expansion or longer training duration indeed contributes to the improved performance.

We thank the reviewer for highlighting the need for clarity regarding the results in Table 2. Our aim is to compare the performance of models expanded using LEMON (trained for 132/165k steps), with models trained from scratch for a longer duration of 220k steps. The improved performance is primarily due to the model expansion process, as the training steps for the expanded models are significantly fewer (132/165k) compared to the 220k steps for the models trained from scratch.

We realize that our initial presentation might have led to some confusion. Specifically, the mention of 'A potential reason for this may be its longer training duration' refers to a different comparison – between LEMON expanded models using BERT(6,384) and BERT(6,512).

To avoid this confusion, we have revised Table 2 by including the total training steps for each case and modifying the caption.

4. The novelty of this paper seems quite limited. The key idea of model expansion seems a simple extension of net2net. Are there any essential technical differences? In my opinion, a simple extension of an existing approach is insufficient for a top-tier conference.

We thank the reviewer for their comments on the novelty of our paper. While we recognize the foundational role of Net2Net in the field of model expansion, we respectfully disagree that our work is a mere extension of it, particularly when considering Transformers with LayerNorm.

Width and depth expansion:

• Net2Net was initially developed for CNNs using ReLU activation and BatchNorm. While groundbreaking, it had its limitations. For instance, its depth expansion is only lossless under specific conditions (e.g., $\sigma(I\sigma(x))=\sigma(x), \forall x$). Moreover, Net2Net struggles with 'weight symmetry' issues in the absence of noise addition, losing its lossless characteristic when noise is introduced.
• As neural network architectures have evolved, particularly with the emergence of Transformers using LayerNorm, the constraints of Net2Net have become more pronounced. These architectures present challenges that Net2Net was not designed to address. Recognizing this gap, we introduce a suite of lossless method with 'weight symmetry breaking', applicable to both divisible and indivisible width, depth expansion, and across various architectures (See Table 1). The ability to maintain losslessness in these scenarios marks a substantial advancement.

Training recipe:

• Another significant advancement is our exploration of the impact of training recipes on the performance of expanded models. Net2Net does not explore this aspect, whereas our work provides novel insights into optimizing training strategies post-expansion, contributing to the enhanced performance of the expanded models.

These novel contributions significantly address the critical limitations of existing model expansion techniques and may pave the way for new research and practical applications. The technical contribution and the practical applicability of our work align well with the standards of a top-tier conference.