Paper: A Multi-Level Framework for Accelerating Training Transformer Models (ICLR-2024)

Reviewer comment (Reproducibility):
The paper lacks in-depth technical explanations about the proposed operators (Coalescing, De-coalescing, and Interpolation) and their implementation. Why the V-cycle training process is better than the previous width/depth expansion methods like bert2BERT and Network Expansion?

Author response:
**(1) Detailed explanations about the operators and their implementation**

Thanks for pointing it out. We included comprehensive implementation details of the coalescing, de-coalescing and interpolation operations in Appendix I. Please kindly refer to it for technical explanations.

**(4) Why the v-cycle training process is better than previous methods**

Thanks for your insightful question. Considering the performance advantages of our approach over bert2BERT and Network Expansion, we note several key distinctions from the previous methods. First, our method expands the transformer model both in width and depth, but bert2BERT and Network Expansion only expand in one dimension, which may limit the potential for performance improvement. Second, bert2BERT adopts a strategy of initializing the expanded parameters in each layer with parameters from higher layers. According to our observations, however, this approach results in a higher initial loss and necessitates more steps for loss reduction, due to the lesser correlation of these inherited parameters with the original ones. Last, Network Expansion utilizes the EMA parameters to enlarge the model depth, which is orthogonal to our method. We would like to incorporate the EMA method into our framework in the future.

Reviewer comment (Evaluation):
The paper does not extensively discuss the potential trade-offs or limitations introduced by the proposed framework. For example, are there any trade-offs in terms of model accuracy, generalization ability, or robustness? A thorough analysis of these aspects would provide a more comprehensive understanding of the framework's impact on model performance.

Author response:
**(3) Trade-offs of the proposed framework**

In our experimental setup, the interpolated model is trained until its validation loss is comparable to that of the baseline model. As a result, hyper-parameters that negatively impact the model's training outcomes tend to show a reduced acceleration ratio, because they require more training steps to match the baseline's performance (the maximal training steps are still the same as baseline). In Appendix D, we maintained a similar validation loss (with differences less than 0.002) across all models to examine the trade-offs of various hyper-parameters in terms of acceleration outcomes. For example, a larger interpolation ratio $\alpha$ might lower the initial loss of the interpolated model, but could also restrict its learning potential. In line with Network Expansion, we determined the optimal post-trade-off $\alpha$ that offers the best balance between acceleration and baseline-level performance. Another key trade-off involves the number of pre-training steps for the small model, which is a balance between enhancing the small model's performance and computational cost. This specific trade-off is also elaborated upon in Appendix D. Lastly, inspired by Network Expansion, we investigated the trade-off related to the size of small model, with the results updated in Appendix D.

Reviewer comment (Reproducibility):
Algorithm 1 in section 3.4 can help the readers understand the whole flow of the framework but is also kind of sketchy.

Author response:
Thanks for the recommendation. To provide a clearer understanding, we have added more detailed descriptions (Algorithms 2-4) about the operators in Appendix I. This expansion aims to elucidate Algorithm 1 further and make its processes more understandable.

Reviewer comment (Reproducibility):
How different are the original model and the final model merged by the coalesced model and de−coalesced model? Do they have the same dimension? What are the differences between these two models?

Author response:
Our objective is to accelerate the pre-training of large models. The architecture and parameter size of the final model are identical to those of the original model. While the model size would be reduced during the coalescing stage, we ensure its original dimensions are recovered during the de-coalescing phase.

Reviewer comment (Novelty):
I feel the overall novelty of this paper is a bit limited, as compared with LiGO. I find the major differences lie in two aspects: LiGO learns linear mapping matrices via SGD, while this work intuitively defines the coalescing matrix as $[I, I]^T$, seeking to directly coalesce two neighboring neurons and adjacent layers; and as discussed in Appendix B, LiGO gradually learns to increase the model size, whereas this paper introduces V-cycle, a first-coalescing-then-decoalescing learning process equipped with interpolation.

Author response:
We define the coalescing matrices as a binding scheme, deriving from the observed inter- and intra layer similarities, which are a form of inductive prior. We actually already consulted with the author of LiGO regarding the initialization of the mapping matrices in their work, which were not detailed in their paper and the code has not been released due to certain constraints. The author indicated that in LiGO, matrices B and A are initialized uniformly but normalized with a softmax function to mimic a selection matrix, akin to our de-coalescing matrix $T_{out}=[I,I]$. This suggests that in practice, transformations aren't learned from scratch without any inductive priors in LiGO. Although the concept of LiGO is orthogonal to ours and thus it seems that combining both methods could yield further acceleration, our empirical experiments indicate that a learned transformation for de-coalescing matrices doesn't enhance the performance of our framework. This is evidenced in Appendix J, Figure 7, where we show the continuing training of the interpolated model after training the de-coalescing matrices. The results demonstrate that both models eventually converge to similar performance levels with and without learned transformations. The empirical evidence suggests that our approach and LiGO are originated from different perspectives and a direct combination does not lead to advantages.

Reviewer comment (Novelty):
The interpolation operation, which the authors claim mitigates the low-rank issue encountered in LiGO, is reminiscent of well-known PEFT methods like LoRA. More comparisons with this line of research would enhance the persuasiveness and credibility of the proposed method.

Author response:
We agree that our framework bears a certain degree of resemblance with LoRA. Notably, when we apply width coalescing alone, the interpolation operator can be redefined as $W_{intp} = W_{large} + T_{in} F_{in} W_{large} F_{out} T_{out}$. Our work focus on updating a segment of the low-rank weights, specifically $F_{in} W_{large} F_{out}$, for both forward and backward propagation. This approach is different from LoRA, which involves using both the original and low-rank weights in the forward propagation. As a result, our approach significantly reduces FLOPs and walltime in the pre-training phase of the smaller model. In contrast, LoRA does not alter the forward propagation process of the original model and even requires additional computations for the low-rank parameters. While LoRA do not need gradient calculations for weights in the feedforward, layernorm and embedding layers, as well as biases in all layers, the backpropagation process still needs to process gradients from the top layer down by calculating the product of weights in each layer and the gradients of the layer outputs, which is a major contributor to FLOPs in backpropagation. Therefore, the FLOPs saving for LoRA is marginal compared with the coalesced model. Furthermore, our measurements of training time on models starting from scratch and those trained with LoRA revealed similar rates (around 4.0 iterations per second), which further confirms that the computation reduction of BERT with LoRA is marginal. Finally, in experiments comparing the pre-training of BERT-Base with LoRA against Coalesced BERT (Appendix K, Figure 8), we observed that the coalesced model converges much faster than BERT-Base with LoRA. This underscores the importance of focusing on intensive low-rank updates, i.e. the coalescing operation. We have added the discussion about LoRA in Appendix K.

Reviewer comment (Novelty):
The authors introduce the multigrid algorithm with a detailed description. However, it seems the proposed framework has little to do with this algorithm.

Author response:
Our algorithm is inspired by the multigrid algorithm proposed for solving large linear systems and bears similarity with it in many aspects. The coalescing and de-coalescing operations in our approach are akin to grid mapping in multigrid. Similarly, the interpolation operation plays the role of the residual correction in multigrid. Furthermore, the overall v-cycle training process corresponds to the v-cycle in multigrid. These parallels reflect a strong conceptual link between our method and the multigrid approach.

Reviewer comment (Theory):
What is the significance of the coalescing operation within the overarching framework, and what benefits does it offer compared to initiating the training process with smaller models? In the coalescing step, the compression matrices F_in and F_out in equation 1 and 2, and R in equation 4 are heuristically defined. The recover matrices G in equation 7, and T_in and T_out are also manually defined without further explanation or theoretical basis. I feel the key point of this framework lies in modeling the correlation of parameters between large models and small models. I am not very convinced, from a methodological point of view, why the proposed framework can help to converge faster on the training set D.

Author response:
Thanks for your question about the role of the coalescing operation. The coalescing operation establishes a crucial link between the large and small models. To evaluate the effect of the coalescing operation, we remove the coalescing operation in our algorithm, i.e. randomly initialize the small model within the v-cycle. The comparative result, presented in Appendix F, Figure 5(a), suffers from an 8.3\% drop in FLOPs saving when the small model is randomly initialized. The dramatic drop demonstrates the necessity of the coalescing operation.

Additionally, we examined the interpolation loss curve between the large model prior to coalescing and the de-coalesced model (with or without coalescing). By interpolating models across various alpha values and assessing their validation loss, we could chart a linear path in the optimization space[1]. The results, displayed in Appendix F Figure 5(b), reveal lower losses along the interpolation path for the de-coalesced model with coalescing operation. This finding underscores the tighter correlation between the de-coalesced and original models when coalescing is applied. The detailed discussion of these effects is presented in Appendix F.

Reviewer comment (Evaluation):
Aside from conserving computational resources, what benefits do the interpolating model M_{k} and the de-coalesced model M_{k, de-coalesced} offer compared to continuing training the de-coalesced model M_{k, de-coalesced}?

Author response:
Direct training of the de-coalesced model poses practical challenges. The process of de-coalescing involves duplicating attention heads/neurons, resulting in identical gradients for these copied elements, which does not inherently enhance the network's learning capacity. We elucidate this problem mathematically using a simple feedforward neural network example in Appendix G.

We conducted experiments to continuously train the de-coalesced GPT-Base. The training curve, depicted in Appendix G, Figure 6, indicates that the convergence performance of de-coalesced model is markedly inferior compared to the model training from scratch. This empirical result corroborates our theoretical understanding and highlights the limitations of directly training the de-coalesced model.

Reviewer comment (Experiments):
Most efforts of the experiments are on BERT models. Competitive compared methods, say LiGO, are not included in the results of GPT and DeiT in Table 2 & 3.

Author response:
Thanks for your constructive suggestion. Following your advice, we conducted more comparison experiments on GPT and DeiT during rebuttal. Comparative results are presented in RTable 3 and RTable 4. In addition, we compare with Network Expansion[2] on BERT-Base and update the results in Table 1 of the revised version. These results affirm that our multi-level framework maintains its superiority across various model architectures.

**RTable3: Comparative results on GPT-Base.**
| Method            | FLOPs Saving  | Walltime Saving     | LAMBADA | PTB   | WikiText-2 | WikiText103 |
|-------------------|---------|------------|---------|-------|------------|-------------|
| GPT-Base          | 0\%     | 0\%        | 54.5    | 146.3 | 49.8       | 50.2        |
| StackBERT         | 9.5\%   | 8.4\%      | 53.3    | 140.6 | 46.5       | 46.9        |
| bert2BERT         | 11.5\%  | 8.3\%      | 53.9    | 147.1 | 48.8       | 49.4        |
| LiGO              | 14.1\%  | 6.9\%      | 54.0    | 139.7 | 50.1       | 50.5        |
| Network Expansion | 15.2\%  | 12.2\%     | 54.7    | 143.7 | 50.7       | 51.2        |
| Ours              | 24.1\%  | 16.5\%     | 53.2    | 142.5 | 47.2       | 47.5        |

**RTable4: Comparative results on DeiT-B.**
| Method            | FLOPs Saving  | Walltime Saving     | Imagenet  (Top 1 Acc)  | CIFAR10 | CIFAR100 | Flowers | Cars   |
|-------------------|---------|------------|-------------|---------|----------|---------|--------|
| DeiT-B            | 0\%     | 0\%        | 81.1\%      | 99.1\%  | 90.8\%   | 97.8\%  | 92.1\% |
| StackBERT         | 23.8\%  | 15.1\%     | 81.2\%      | 99.1\%  | 90.8\%   | 97.6\%  | 92.1\% |
| bert2BERT         | -0.1\%  | -0.13\%    | 81.6\%      | 99.1\%  | 90.7\%   | 97.7\%  | 92.2\% |
| LiGO              | 25.4\%  | 12.0\%     | 81.7\%      | 99.1\%  | 90.7\%   | 97.8\%  | 92.1\% |
| Network Expansion | 25.0\%  | 22.5\%     | 81.5\%      | 99.1\%  | 90.7\%   | 97.8\%  | 92.1\% |
| Ours              | 27.1\%  | 24.3\%     | 81.5\%      | 99.1\%  | 90.8\%   | 97.6\%  | 92.1\% |

Reviewer comment (Evaluation):
In Table 4, as the number of Levels increases, the author's method not only saves more computational resources but also further improves the performance of the final model. What is the rationale for this observed enhancement in performance?

Author response:
Regarding the observed enhancement in Table 4, we find that the enhancement is predominantly the improvements on the MRPC and CoLA datasets. It should be noted that Level-3 does not consistently surpass the baseline across all tests. For instance, its performance is marginally lower than the baseline on the STS-B dataset. Compared with other downstream datasets, the MRPC and CoLA datasets have limited sample sizes, leading to wide fluctuations in finetuning performance. Similar fluctuations are also observed in Table 1 of LiGO. The calculation of the average score does not account for the number of samples or the datasets' reliability and thus result in the average score of level-3 seeming significantly better than baseline actually is. Based on this observation, we could not conclude that level-3 generalizes much better than the baseline.

Lastly, we have implemented learnable mapping matrices in our algorithm. However, our findings suggests that the effect of such an approach is indistinguishable. This is further supported by our experiments on the effects of various coalescing matrices, which indicate that the choice of different coalescing matrices may not be particularly crucial in our set-up.

Reviewer comment (Evaluation):
Though significant speedup in the NLP transformer, the proposed method has limited improvement in FLOPs and time reduction in the large vision model.

Author response:
Thank you for raising an interesting question regarding the performance of our framework on vision transformers. In fact, due to limited computing resources before submission, we chose DeiT-S to conduct the vision transformer experiment. DeiT-S has 12 layers but only 6 heads, which means its complexity and redundancy is less than other larger models. Therefore, the savings of FLOPs and walltime are limited on DeiT-S. 

After the first submission, we managed to get more computing resources and conducted the vision transformer experiment on DeiT-B, which has 12 layers and 12 heads, i.e. equivalent to the scale of BERT-Base and GPT-Base. Experimental results are presented in the RTable 2. The FLOPs and walltime saving on DeiT-B is 27.1\% and 24.3\%. The results show that our method performs well on vision transformers as long as the parameter size is close to that of other architectures.

Reviewer comment (Reproducibility):
It's unclear how to initialize the matrix F. It seems the F can be arbitrary and the initialization of F is not discussed sufficiently.

Author response:
Thanks for your question about the initialization of the width coalescing matrix $F$. Here we illustrate the initialization process with a transformer example with 12 layers and 12 heads. Considering to the redundancy in the attention heads as observed in Figure 1, we hope to bind similar attention heads and simultaneously optimize them in a more efficient manner. The width coalescing matrix $F^{k+1, l}_{w, out}=[\mathbf{I}/2, \mathbf{I}/2]^T$ used in our experiments is created by merging the $i$ and $i+6$ ($i=1,2,...,6$) attention heads pair by pair. For the feedforward layer with $n$ neurons, the $i$ and $i+n/2$ neurons is also merged in a paired fashion.

Thanks for the suggestion regarding more discussion on the initialization of F. To providing further comparison, we adopt another width coalescing matrix $F_{adj}$, which merges $i$ and $i+1$ ($i=1, 3,...,11$) attention heads. We give the definition of $F_{adj}$ in Appendix E due to the limitation of OpenReview' to support latex equations. We evaluate the effect of different width and depth coalescing matrices on BERT-Base. Comparative analysis reveals minimal variation in terms of FLOPs savings among various matrices (<0.3\%). We have updated Appendix E with the elaboration on the initialization of width coalescing and the discussion of different width coalescing matrices.

Reviewer comment (Experiments):
The evaluation result on GPT and DeiT-S lacks a comparison with other works.

Author response:
Thanks for pointing out this point. Following your advice, we conducted more comparison experiments on GPT and DeiT during rebuttal. Comparative results are presented in RTable 1 and RTable 2. In addition, we compare with Network Expansion[1] on BERT-Base and update the results in Table 1 of the revised version. These results affirm that our multi-level framework maintains its superiority across various model architectures.

**RTable1: Comparative results on GPT-Base.**
| Method            | FLOPs Saving | Walltime Saving     | LAMBADA | PTB   | WikiText-2 | WikiText103 |
|-------------------|---------|------------|---------|-------|------------|-------------|
| GPT-Base          | 0\%     | 0\%        | 54.5    | 146.3 | 49.8       | 50.2        |
| StackBERT         | 9.5\%   | 8.4\%      | 53.3    | 140.6 | 46.5       | 46.9        |
| bert2BERT         | 11.5\%  | 8.3\%      | 53.9    | 147.1 | 48.8       | 49.4        |
| LiGO              | 14.1\%  | 6.9\%      | 54.0    | 139.7 | 50.1       | 50.5        |
| Network Expansion | 15.2\%  | 12.2\%     | 54.7    | 143.7 | 50.7       | 51.2        |
| Ours              | 24.1\%  | 16.5\%     | 53.2    | 142.5 | 47.2       | 47.5        |

**RTable2: Comparative results on DeiT-B.**
| Method            | FLOPs Saving  | Walltime Saving     | Imagenet  (Top 1 Acc)  | CIFAR10 | CIFAR100 | Flowers | Cars   |
|-------------------|---------|------------|-------------|---------|----------|---------|--------|
| DeiT-B            | 0\%     | 0\%        | 81.1\%      | 99.1\%  | 90.8\%   | 97.8\%  | 92.1\% |
| StackBERT         | 23.8\%  | 15.1\%     | 81.2\%      | 99.1\%  | 90.8\%   | 97.6\%  | 92.1\% |
| bert2BERT         | -0.1\%  | -0.13\%    | 81.6\%      | 99.1\%  | 90.7\%   | 97.7\%  | 92.2\% |
| LiGO              | 25.4\%  | 12.0\%     | 81.7\%      | 99.1\%  | 90.7\%   | 97.8\%  | 92.1\% |
| Network Expansion | 25.0\%  | 22.5\%     | 81.5\%      | 99.1\%  | 90.7\%   | 97.8\%  | 92.1\% |
| Ours              | 27.1\%  | 24.3\%     | 81.5\%      | 99.1\%  | 90.8\%   | 97.6\%  | 92.1\% |

Reviewer comment (Writing):
Can you explain how the intra- and inter-layer similarity is utilized in the coalescing and de-coalescing procedure?

Author response:
In general, the intra- and inter-layer similarity suggests that it's unnecessary to train similar heads/layers separately. Therefore, we propose the coalescing operation, which binds heads/layers and optimizes them simultaneously. For example, layer 1 and layer 2 would be bound into layer 1' in a smaller model. With the de-coalescing operation, each bound head/layer will then be divided into two identical heads/layers, i.e. layer 1 of the small model is mapped to layer 1 and layer 2 in the de-coalesced model. When interpolated from the de-coalesced, parameters of layer 1 and layer 2 in the original model will be updated in the same way. As a result, the whole procedure can be viewed as updating adjacent heads/layers simultaneously. In this way, we implicitly utilize the intra- and inter-layer similarity to accelerate the training procedure.