SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\min  _ {\boldsymbol {S} _ {B}, \boldsymbol {S} _ {I}, \boldsymbol {S} _ {O}, \boldsymbol {S} _ {L}} \mathbb {E} _ {\mathbf {X} \sim \mathcal {D}} \mathcal {L} (\mathbf {X}, \boldsymbol {M} _ {2}), \quad \text {w . r . t} \quad \boldsymbol {M} _ {2} = \Phi_ {\boldsymbol {S} _ {B}, \boldsymbol {S} _ {I}, \boldsymbol {S} _ {O}, \boldsymbol {S} _ {L}} (\boldsymbol {M} _ {1}), \tag {7}
$$
where  $\mathcal{L}$  is a loss function, and  $\mathcal{D}$  is a data distribution. By replacing the full space of  $\mathcal{S}$  with four small spaces, the spatial requirements of the training process are reduced exponentially.
Procedures of Applying Mango. The procedures of Mango can be divided into four steps: (i) concatenating weights  $\theta^{L_1,D_1}$  of a pretrained model  $\mathbf{M}(L_1,D_1)$  to construct a tensor  $\mathcal{M}_1$ ; (ii) training the growth operator  $\Phi_{\pmb {S}_B,\pmb {S}_I,\pmb {S}_O,\pmb {S}_L}$  items of Eq. ⑦ in a few steps (e.g., 100) to make transferred models maintaining function; (iii) recovering weight tensor  $\mathcal{M}_2$  through the multi-linear operator  $\Phi_{\pmb {S}_B,\pmb {S}_I,\pmb {S}_O,\pmb {S}_L}$ ; (iv) splitting  $\mathcal{M}_2$  to the weight  $\theta^{L_2,D_2}$  of the target model  $\mathbf{M}(L_2,D_2)$  as initialization to continue training.
Remark. Tensor decomposition is often used to compress neural networks. Actually, there are spaces for further reducing the operator size. However, we are not exploring the border of compression ratio, but the possibility of the learning ability of multi-linear operators and the influence on model growth. Therefore, we decompose the huge  $\mathcal{S}$  into four interpretable smaller tensors.  $\mathcal{S}_B$  means the interactions on the parameters in the same layer.  $S_{I}$  and  $S_O$  denote the transformations of input and output dimensions in one parameter, respectively.  $S_L$  indicates the relationship among layers.  $R$  means the low-rank level of  $\mathcal{S}$ . Smaller  $R$  means less correlation of these four tensors.
3.3 Comparison with Recent Advances
Basically, the goal of model growth is to widen and deepen the width and depth of the pretrained models, and operation on width and depth can be easy to implement. Therefore, most of the previous studies grow width and depth separately, although there may exist some correlation that may influence the training efficiency. In this part, we analyze the difference among methods of model growth with the help of the tensor diagrams.
We illustrate a comparison among bert2BERT, LiGO, and Mango in Figure 5. The red circle means a pretrained model, while the blue and gray circles denote trainable and untrainable operator parameters, respectively. The circle with an oblique stroke represents a super-diagonal tensor, where the diagonal elements are all 1. This super-diagonal tensor means that growth operators expand other modes along the mode of the super-diagonal tensor.
Table 1: The comparison among bert2BERT, LiGO and Mango. In a case of  $\mathbf{M}(L_1, D_1) \rightarrow \mathbf{M}(L_2, D_2)$ , there usually are  $I_1 = O_1 = D_1$  and  $I_2 = O_2 = D_2$ .  $R = \max(R_1, R_2, R_3, R_4)$ .
Method	Reference	Operator Parameter	Trainability	Spatial Complexity
bert2BERT	Chen et al. [5]	SI, SO, SK	×	2L1D1D2 + L1L2
LiGO	Wang et al. [53]	SI, SO, SL	✓	2B1D1D2 + L1L2
Mango	-	SB, SI, SO, SL	✓	2RD1D2 + R2(B1B2 + L1L2)
bert2BERT. As shown in Figure 5, the parameters in the growth operator of bert2BERT are all frozen, since they are designed by heuristic inspiration to preserve function. bert2BERT expands modes of  $I_{1}$  and  $O_{1}$  along with the mode of  $L_{1}$ , indicating that one weight is expanded without knowledge from any other weight. To further increase the training ability, bert2BERT applies  $S_{K}$  to construct Advanced Knowledge Initialization (AKI) to utilize knowledge of other layers to help accelerate training.
LiGO. Different from bert2BERT, the LiGO operator can be trained. In addition, LiGO expands  $I_{1}$  and  $O_{1}$  along the mode of  $B_{1}$ . With  $S_{L}$ , LiGO can combine knowledge among layers. Operators of Net2Net and StackBERT grow with or depth separately, which can be regarded as a sub-solution to  $S_{I}, S_{O}$ , and  $S_{L}$  of LiGO. However, one weight in LiGO will not leverage the knowledge of weights in the same layer. Therefore, LiGO only implements partial mapping like bert2BERT.
Mango. We employ Mango on each mode of the weight  $\mathcal{M}_1$ . Mango can approach the full mapping tensor  $\mathcal{S}$  as a low-rank approximation with rank  $R$ , rather than partial mapping operators like LiGO and bert2BERT. Therefore, Mango can obtain  $\mathcal{M}_2$  to capture adequate knowledge from pretrained weights as formulated in Eq.6. Moreover, the diagonal tensor in LiGO and contraction between a diagonal tensor and  $\mathcal{S}_K$  in bert2BERT are subsets of  $\mathcal{S}_L$  and  $\mathcal{S}_B$  in Mango. Therefore, Mango can be a generalization of bert2BERT and LiGO with more expressive power.
In this section, we design a set of experiments to validate the proposed Mango. To begin with, we conduct an ablation study to analyze the influence of Mango on width and depth in Sec. 4.1. Then we conduct image classification with DeiT-B [46] on ImageNet to show the acceleration in computer vision task in Sec. 4.2. Later we conduct a pretraining experiment with BERT-Base [10] to demonstrate the effectiveness in natural language processing tasks in Sec. 4.3. At last, we also employ a pretraining experiment with GPT-Base [39] to show the wide adaption of Mango in Sec. 4.4
Ratio of saving FLOPs. We evaluate the training efficiency in terms of the ratio of saved floating-point operations (FLOPs). Let us assume that training a model from scratch to convergence on metric  $\Psi$  (e.g., MLM loss or accuracy) necessitates  $\xi_{Scratch}$  FLOPs. Then, for a method that reaches the same metric  $\Psi$  with  $\xi_*$  FLOPs, its FLOP saving ratio  $r$  can be computed as
$$
r = \frac {\xi_ {S c r a t c h} - \xi_ {*}}{\xi_ {S c r a t c h}}. \tag {8}
$$
(a) Width: DeiT-T-A  $\rightarrow$  DeiT-S.
(b) Depth: DeiT-T-B  $\rightarrow$  DeiT-S.
Figure 6: Influence from ranks on expanding width, depth, and both of them. The green curve means the accuracy of training operators for 100 steps. The red curve means the final acceleration ratio on training DeiT-S. Mango can achieve training acceleration in every case. Moreover, in each sub-graph, when the accuracy increases with a higher rank, the acceleration ratio keeps almost fixed.
(c) Both: DeiT-T-C  $\rightarrow$  DeiT-S.
In this experiment, we explore the influence of Mango on growing width, depth and both of them in addition to rank setting. We use three tiny vision Transformers (ViTs) [11], i.e., DeiT-T-A, DeiT-T-B, and DeiT-T-C, for growing to DeiT-S [46] on ImageNet [9]. Structures of these DeiTs can be found in Table 4 of the Appendix. For ease of setting, we set all the ranks the same in the range of \{1, 4, 7, 10\}. We train Mango operators for 100 steps, which only requires negligible time. We use Adam with learning rate 1e-3 and weight decay 1e-2 for 300 epoch optimization. The batch size is 1024.
Results are shown in Figure 6. Mango achieves acceleration of at most  $31.0\%$ ,  $46.0\%$ , and  $41.3\%$  on expanding width, depth, and both of them, respectively, which shows Mango can fit various growth scenarios. Interestingly, these acceleration ratios are higher when the operator accuracies are better along the types of pretrained models, e.g., as DeiT-T-B at most achieves  $73.81\%$  accuracy which is higher than  $72.89\%$  of DeiT-T-A, DeiT-T-B derives higher acceleration ratio. This phenomenon suggests that when there are multiple pretrained models can be selected, the better accuracy for one pretrained model through Mango, the faster training speed can be obtained for the target models.
In each expanding case, the accuracies of operator training tend to increase with higher ranks. However, all the cases on the three pretrained models show that better accuracy will not lead to faster training with the same pretrained model. For example, in Figure 6(a), the operator with rank 10 has an accuracy that is  $1.19\%$  higher than the operator with rank 1. Nevertheless, the two ranks reach the same acceleration ratio. This result suggests that rank 1 is enough to use, and also enjoys the advantages of spatial and temporal complexities compared to bigger ranks. And we choose rank 1 for Mango to construct the later experiments.
4.2 Results on Large-Scale Vision Transformer
We conduct the experiment to show the training acceleration on large-scale vision Transformers. We train the Deit-B from Deit-S [46] on ImageNet. Structures of the two models are shown in Table 4 of the Appendix. Ranks of Mango are all set to 1 for complexity benefit without performance loss. The operators of Mango and Ligo are all trained within 100 steps. We use Adam as the optimizer with learning rate 1e-3 and weight decay 1e-2. The batch size is 1024. The training epoch is 300.
Results are shown in Figure 7(a). Mango saves  $76\%$  FLOPs from Scratch which converges to an accuracy of  $80.45\%$  in the end. Compared with schemes of training from Scratch (i.e., Srcatch and StackBERT), methods of training from a smaller model (i.e., Mango, bert2BERT, and LiGO) can attain  $70\%$  accuracy in a short time, even bert2BERT start in low accuracy. Compared with the recent SOTA models, Mango has surpassed bert2BERT for  $+12.0\%$ , and LiGO for  $+20.7\%$ . This experiment shows the ability of Mango to achieve significant improvement in training acceleration. To investigate the influence of Mango on transferring ability, we also conduct an experiment on downstream tasks, including CIFAR10 [26], CIFAR100 [26], Flowers [31], Cars [25], and ChestXRay8 [54]. Results are shown in Table 2. It is easy to see that Mango achieves similar results to the Scratch, which indicates Mango has not influenced the transferring ability.
(a) DeiT-S  $\rightarrow$  DeiT-B.
(b) BERT-Small  $\rightarrow$  BERT-Base.
Figure 7: Results of pretraining DeiT-B, BRET-Base and GPT-Base. Compared to baselines, Mango can achieve the highest savings in FLOPs with  $76.4\%$  for DeiT-B,  $39.2\%$  for BERT-Base, and  $59.9\%$  for GPT-Base from the Scratch model. We also illustrate the corresponding results on wall time in Figure 10 of the appendix.
(c) GPT-Small  $\rightarrow$  GPT-Base.
Table 2: Results of transfer learning performance of DeiT-B. Mango can achieve similar performance to the Scratch model in downstream tasks while saving  $76.4\%$  FLOPs.
Method	FLOPs (×1e18)	Ratio (Saving)	CIFAR10	CIFAR100	Flowers	Cars	ChestXRay8	Average
Training from Scratch
Scratch	12.9	-	99.03	90.22	97.27	91.89	55.66	86.82
StackBERT	11.3	12.6%	99.11	90.10	97.44	91.71	55.63	86.80
Training from the Pretrained Model: M(12,384) → M(12,768)
bert2BERT	4.6	64.4%	98.99	90.47	97.51	91.88	55.34	86.84
LiGO	5.7	55.7%	99.11	90.52	97.18	91.82	55.45	86.82
Mango	3.0	76.4%	99.13	90.23	97.49	91.83	55.46	86.83
4.3 Pretraining on BERT
In this experiment, we conduct the validation to show the training acceleration on BERT [10, 61]. The dataset is the concatenation of English Wikipedia and Toronto Book Corpus [71]. We train the BERT-Base from BERT-Small. The training epoch is 40. The batch size is 768. We list the structures of the two models in Table5 of the Appendix. The ranks of Mango are all 1. Mango and LiGO are both warmly trained for 100 steps. The optimizer is set to AdamW. The learning rate is 1e-4 and the weight decay is 1e-2.
We illustrate the training curves in Figure 7(b). The methods of training from BERT-Small (i.e., Mango, bert2BERT, and LiGO) can all surpass the progressive training StackBERT of acceleration ratio  $29.5\%$ . bert2BERT is over StackBERT by  $+6.1\%$ . Mango achieves the highest acceleration of  $39.2\%$  FLOPs which is  $+3.6\%$  more than bert2BERT. Loss of StackBERT steeply decreases within training for Stacking layers, which can also demonstrate the pretrained weights are helpful for training efficiency. We show the effectiveness of Mango on SQuAD and GLUE benchmark as in Table 3. Mango shows the same transferring ability with faster convergence, which indicates a promising use for practical training.
4.4 Pretraining on GPT
We also implement the experiment to show the training acceleration on GPT [39]. The dataset is the concatenation of English Wikipedia and Toronto Book Corpus [71]. We train the GPT-Base from GPT-Small. The structures of the two models are shown in Table 5 of the Appendix. The ranks of Mango are all 1. Mango and LiGO are trained for 100 steps. We use Adamw with learning rate 1e-4 and weight decay 1e-2. The batch size is 512. The training epoch is 35.
Table 3: Results of downstream tasks of BERT-Base on GLUE [48], SQuADv1.1 [40], and SQuADv2.0 [41] dataset. Mango can also achieve similar performance with the Scratch model while enjoying training efficiency.
Model	FLOPs (×1e19)	Ratio (Saving)	SQuADv1.1 (F1)	SQuADv2.0 (F1)	SST-2 (Acc)	MNLI (Acc)	MRPC (Acc)	COLA (Mcc)	QNLI (Acc)	STS-B (Acc)	QQP (Acc)	GLUE Avg.	SQuAD Avg.
Training from Scratch
Scratch	8.9	-	89.21	77.90	92.18	84.19	87.55	56.35	91.50	89.16	90.25	84.45	83.56
StackBERT	6.3	29.5%	89.82	78.21	92.94	84.63	87.65	61.61	90.95	87.13	90.20	85.01	84.01
Training from the Pretrained Model: M(12,384) → M(12,768)
bert2BERT	5.7	35.6%	90.02	78.99	92.89	84.92	86.91	60.32	91.81	88.11	90.72	85.10	84.50
LiGO	5.9	33.5%	90.09	78.34	92.75	84.99	87.44	61.10	91.33	87.94	90.42	85.14	84.22
Mango	5.4	39.2%	90.17	78.77	92.71	84.86	87.94	62.88	91.49	88.73	90.62	85.60	84.47
We compare the Scratch model, bert2BERT, LiGO, and Mango in Figure 7(c). We observe that the proposed Mango achieves a  $59.9\%$  acceleration ratio. While GPT is different from BERT in different structures, including layer normalization, mask method, and training strategy, Mango can always keep the highest performance. Although LiGO is lower than bert2BERT at the beginning, it converges slower and reaches  $38.1\%$  at last. By contrast, the Mango is almost at the lowest loss in the whole training process and achieves  $+16.7\%$  more than bert2BERT and  $+21.8\%$  more than LiGO, which shows a significant acceleration than the baselines.
Training Transformers can pose a significant demand on computational resources. Reusing pretrained models as an initialization strategy for the target model offers a promising approach to reducing resource costs and accelerating training. However, previous studies only mapped partial weights when growing models, which may fail to consider potential correlations in the whole model and result in inadequate growth mapping. Inspired by this observation, we propose to consider the interaction among all weights in the model to further improve acceleration ability. Specifically, we utilize a full mapping to comprehensively reuse the pretrained model, taking into account all correlations between model weights. As the full mapping tensor is huge and cannot be employed in practice, we propose to use Mango, a multi-linear operator, to reduce computation and space complexity. Experimental results demonstrate that Mango consistently achieves significant acceleration on various large-scale models (e.g., DeiT-B, BERT-Base, and GPT-Base). In the future, we hope that this method can contribute to green AI and significantly reduce the cost of training Transformers.
Limitation. While Mango significantly reduces training costs for large models, it still requires weights of a small pretrained model as the prior knowledge. Additionally, the Mango operator necessitates extra training for initializing. However, obtaining small model weights is relatively simple within the community, and given the substantial cost savings, the resources required for training the operator are minimal.
Societal Impact. The acceleration from Mango significantly impacts energy conservation and environmental protection, promoting the growth of Green AI, which boosts efficiency and reduces computational demand, minimizing energy usage and carbon emissions. Mango can also democratize AI, allowing broader access to technologies of pertaining models without the need for large resources, and promoting sustainable and accessible AI solutions that respect our environmental limits.
[69] Zhuo Zhang, Yuanhang Yang, Yong Dai, Qifan Wang, Yue Yu, Lizhen Qu, and Zenglin Xu. Fedpetuning: When federated learning meets the parameter-efficient tuning methods of pretrained language models. In Findings of the Association for Computational Linguistics: ACL 2023, Toronto, Canada, July 9-14, 2023, pages 9963-9977. Association for Computational Linguistics, 2023.
[70] Jiawei Zhao, Florian Schäfer, and Anima Anandkumar. Zero initialization: Initializing residual networks with only zeros and ones. CoRR, abs/2110.12661, 2021.
[71] Yukun Zhu, Ryan Kiros, Richard S. Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In ICCV, pages 19-27. IEEE Computer Society, 2015.