SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yu Pan\*, Ye Yuan, Yichun Yin, Zenglin Xu\*, Lifeng Shang, Xin Jiang, Qun Liu
$\spadesuit$ Harbin Institute of Technology Shenzhen, Shenzhen, Guangdong, China
\* Pengcheng Laboratory, Shenzhen, China
$\diamond$  Peking University, Beijing, China
$\hat{\mathcal{O}}$  Huawei Noah's Ark Lab, Shenzhen, Guangdong, China
Training large models from scratch usually costs a substantial amount of resources. Towards this problem, recent studies such as bert2BERT and LiGO have reused small pretrained models to initialize a large model (termed the "target model"), leading to a considerable acceleration in training. Despite the successes of these previous studies, they grew pretrained models by mapping partial weights only, ignoring potential correlations across the entire model. As we show in this paper, there are inter- and intra-interactions among the weights of both the pretrained and the target models. As a result, the partial mapping may not capture the complete information and lead to inadequate growth. In this paper, we propose a method that linearly correlates each weight of the target model to all the weights of the pretrained model to further enhance acceleration ability. We utilize multi-linear operators to reduce computational and spacial complexity, enabling acceptable resource requirements. Experiments demonstrate that our method can save  $76\%$  computational costs on DeiT-base transferred from DeiT-small, which outperforms bert2BERT by  $+12.0\%$  and LiGO by  $+20.7\%$ , respectively.
Transformers [47] have recently achieved great successes in various scenarios [11, 4, 20]. Generally, Transformers tend to be larger for more expressive power and better performance (e.g., ViTG [6] and GPT-3 [3]). As the size of models continues to grow, training Transformers takes longer and can result in higher  $\mathrm{CO}_{2}$  emissions, which conflicts with the principles of Green AI [43]. Thus, training Transformers efficiently is crucial not only
Width Expanding Operator
Depth Expanding Operator
Figure 1: Expanding operators.  $W_{s}$  and  $W_{t}$  mean small pretrained and target weights, respectively.  $\Phi_{\theta_w}$  and  $\Phi_{\theta_d}$  denote width and depth expanding operators with parameters  $\theta_w$  and  $\theta_d$ .
for financial gain but also for environmental sustainability [5]. To achieve efficient training, it is a wise option to grow a pretrained small model into a larger one, since the pretrained small model has already learned knowledge from the data, which allows for faster training compared to starting from scratch. [13]. Moreover, there are numerous pretrained models that are easily accessible [53], reducing the cost of utilizing a smaller pretrained model. Furthermore, empirical evidence shows that Transformers have inductive biases that facilitate scaling fitting [42, 22]. This demonstrates the feasibility of learning from pretrained models [53].
*This work is done when Yu Pan is an intern at Huawei Noah's Ark Lab.
Figure 2: Attention pattern maps between BERT-Small (shorted as "S") and BERT-Base (shorted as "B") in four examples. "L" represents the layer index, while "A" denotes the index of the attention head. For example, "S-(L10, A6)" denotes the attention map from layer 10, head 6 of the BERT-Small model. The examples are derived from different sentences. In each example, we can observe the similarities of attention maps in both inter-layer and intra-layer connections, even across different scaled models.
The concept of reusing pretrained models essentially involves exploring the efficient mapping from the pretrained model to the target model. This mapping process can be represented as linear growth operators, as depicted in Figure 1. One viable approach to mapping is to leverage knowledge from other weights. For instance, StackBERT [13] utilizes low-layer information to construct new higher layers, by duplicating layers to increase depth during training epochs. bert2BERT [5] expands width through the expanding operator of Net2Net [7]. Differently, bert2BERT utilizes weights of neighbor layers for enhancing the ability, which also shows the benefit of taking knowledge from other weights. Moreover, distinct from the aforementioned research directly employing a fixed operator, LiGO [53] trains expanding operators which are tied with layers to achieve superior transfer accuracy by considering knowledge from other layers, which is advantageous for training efficiency. However, they neglect the potential connectivity between models.
Potential Connectivity between Models. The attention maps of BERT, as depicted in Figure 2, indicate that there are similarities not only in the weights of the same or different layers, but also across different models. These similarities suggest that we can leverage a full mapping transformation to capture this correlation in both inter-layer and intra-layer connections, even across different scaled models. By doing so, we can reuse the parameters of pretrained small models and accelerate the training of larger models. However, previous studies have overlooked the overall network connectivity and have instead expanded each weight through partial mapping transformation. For instance, Net2Net [7] only considers preserving functionality by transforming weights individually, while bert2BERT [5] increases model width head by head in Transformers. On the other hand, LiGO [53] primarily focuses on extending the weights of the same type (e.g., query, key, and value in Transformers), which heavily affects the performance of the transferred models.
Multi-linear Operators. Based on the above observations, we propose mapping the entire pretrained model to each weight of the target model, instead of employing a partial transformation. This approach can ensure high expressive power, promoting the potential connectivity in models for better training efficiency. However, the completed mapping has a huge parameter tensor, occupying an enormous space (see Section 3.2). To address this issue, we propose Mango, a multi-linear structure (i.e., tensor ring [33,52,51]), which decomposes the large mapping tensor into four smaller tensors that are bonded by ranks to construct multi-linear operators. Mango allows for efficient training by reducing the space required for the mapping tensor. Formally, Mango can be considered as a generalization of LiGO and bert2BERT (see Section 3.3). By using the full mapping approach, Mango can save  $76\%$  computational costs on DeiT-base transferred from DeiT-small, which outperforms bert2BERT by  $+12.0\%$  and LiGO by  $+20.7\%$ , respectively.
Efficient Training from Scratches. Model scratches can be roughly regarded as models without knowledge priors. Some training strategies for scratches are universal and orthogonal to our method.
For example, Adam [23] and large-batch size training [66] accelerate the training process from angles of optimizer and magnitude of input data, respectively. Wang et al. [50] enable the training of very deep neural networks with limited computational resources via a technique called active memory. Shoeybi et al. [45] use mixed precision training to assist training. Low-rank methods benefit training for less memory and time [21]. Wu et al. [59] take notes of rare words for better data efficiency. Dropping layers [68], knowledge inheritance [38], and merging tokens [2] are also efficient methods for training. Another line of work [13, 63, 14, 44] is termed progressive training, which gradually increases the model size within the training process for training efficiency. Li et al. [27] employ a neural architecture search (NAS) method to search optimal sub-networks for progressive training on ViTs. Xia et al. [60] suggest a three-phase progressive training regime to achieve a good trade-off between training budget and performance. Wang et al. [55] introduces a novel curriculum learning approach for training efficiency through firstly learning simpler patterns, then progressively introducing more complex patterns.
Efficient Training from Pretrained Models. Pretrained models usually contain abundant data knowledge, which is helpful for training [32]. By preserving the function of a pretrained model while expanding the model size, it is feasible to give the corresponding larger model an initial state with high performance. Net2Net [7] is the first work to propose the concept of function-preserving transformations by expanding width by splitting neurons and growing depth with identity layers. However, Net2Net splits neurons randomly. Towards this problem, a series of studies [56, 58, 49, 57] propose to select the optimal subset of neurons to be split by utilizing functional steepest descent. bert2BERT [5] expands small transformers by following the function preserving idea. Recently, LiGO [53] utilizes a trainable linear operator to learn a good expanding formula. Different from these prior studies, our method tries to implement a full mapping that achieves comprehensive utilization of the whole smaller model.
Neural Network Initialization. Our method is related to neural network initialization techniques. Xavier [12] and Kaiming [16] initialization aim to control input variance equal to that of output. Generalizing Xavier and Kaiming methods, a universal weight initialization paradigm proposed by Pan et al. [34] can be widely applicable to arbitrary Tensorial CNNs. In addition, Hardt and Ma [15] has shown theoretically that network training benefits from maintaining identity, particularly for improving the efficiency of residual networks. Fixup [67] and ZerO [70] both set residual stem to 0 (not residual connections) to ensure the identity of signals, thereby successfully initializing ResNets. SkipInit [8] replaces Batch Normalization with a multiplier whose value is 0. ReZero [1], on the other hand, adds extra parameters of value 0 to maintain identity, resulting in faster convergence. IDInit is an initialization approach to keep the identity matrix for stable training of networks [36]. In comparison, our work explores fully reusing smaller pretrained models as efficient initialization.
Low-rank Techniques in Neural Networks. Low-rank methods are feasible for reducing spatial and temporal complexities in neural networks [52, 35, 62]. For example, Idelbayev and Carreira-Perpinañan [19] uses matrix decomposition to compress convolutional neural networks (CNNs) for faster inference. LoRA [17] applies low-rank matrices for fine-tuning large language models (LLMs) in affordable resources. As a parameter-efficient tuning (PETuning) method [69] for federated learning, FedPara [18] utilizes the Hadamard product on low-rank parameters for reducing communication time in federated learning. Pan et al. [33] and Li et al. [28] use tensor ring decomposition for reducing the size of neural networks. Tucker decomposition and block-term Tucker decomposition have been used in T-Net and BT-layers [24, 64] for improving the model performance, respectively. Ma et al. [30] apply block-term Tucker to compress Transformers. Yin et al. [65] propose to use ADMM to optimize tensor training to achieve better compression and performance. These techniques have built a solid foundation for our work to implement multi-linear transformation for transferring knowledge from pretrained models to target models.
3 Mango Operator
This section presents the proposed Mango operator. First, we provide an overview of the tensor diagram, the concept of the tensor ring matrix product operator (TR-MPO), and the formulation of Transformer architecture in Sec. 3.1. Next, we delve into the details of the proposed multi-linear mapping operator (i.e., Mango) in Sec. 3.2. Finally, we compare Mango with recent advances in terms of tensor diagrams in Sec. 3.3
Tensor Diagram. Tensor decomposition is a technique that involves splitting a tensor into multiple smaller tensors, typically more than the two matrices that result from matrix factorization. To better illustrate the interactions among multiple tensors, tensor diagrams are often used. A tensor diagram is com
Figure 3: Tensor diagram instances.
posed of two primary elements: a tensor vertex and a tensor contraction. A tensor vertex represents a tensor, and its order is determined by the number of edges connected to it. Each edge is assigned an integer that indicates the dimension of the corresponding mode. As shown in Figure 3 a 3rd-order tensor  $\mathcal{T} \in \mathbb{R}^{\mathbf{i}_0 \times \mathbf{i}_1 \times \mathbf{i}_2}$  can be drawn as a circle with three edges. The process of taking the inner product of two tensors on their matching modes is called tensor contraction. For example, Two tensors,  $\mathcal{M} \in \mathbb{R}^{\mathbf{i}_0 \times \mathbf{i}_1 \times \mathbf{i}_2}$  and  $\mathcal{N} \in \mathbb{R}^{\mathbf{j}_0 \times \mathbf{j}_1 \times \mathbf{j}_2}$ , can be contracted in corresponding positions to form a new tensor of  $\mathbb{R}^{\mathbf{i}_0 \times \mathbf{i}_1 \times \mathbf{j}_2 \times \mathbf{j}_3}$ , when they have equal dimensions:  $\mathbf{i}_2 = \mathbf{j}_0 \triangleq e_0$ . The contraction operation can be formulated as
$$
\left(\boldsymbol {\mathcal {M}} \times_ {2} ^ {0} \boldsymbol {\mathcal {N}}\right) _ {i _ {0}, i _ {1}, j _ {2}, j _ {3}} = \sum_ {m = 0} ^ {e _ {0} - 1} \boldsymbol {\mathcal {M}} _ {i _ {0}, i _ {1}, m} \boldsymbol {\mathcal {N}} _ {m, j _ {2}, j _ {3}}. \tag {1}
$$
The tensor diagram is an elegant tool for succinctly visualizing and manipulating multi-dimensional arrays or tensors. By representing tensor operations as networks of nodes and edges, these diagrams provide an intuitive way to understand the underlying mathematical operations and their interactions. As a result, we have chosen to use tensor diagrams to illustrate our method and to analyze the connections between our approach and prior studies.
Tensor Ring Matrix Product Operator (TR-MPO). TR-MPO means tensor ring (TR) of an MPO [37] format. Given a  $2N$ -order tensor  $\mathcal{X} \in \mathbb{R}^{I_1 \times J_1 \times I_2 \times J_2 \dots I_N \times J_N}$ , its TR-MPO decomposition can be mathematically expressed as
$$
\boldsymbol {\mathcal {X}} _ {i _ {1}, j _ {1}, i _ {2}, j _ {2} \dots , i _ {N}, j _ {N}} \approx \sum_ {r _ {1}, r _ {2}, \dots , r _ {N} = 1} ^ {R _ {1}, R _ {2}, \dots , R _ {N}} \boldsymbol {\mathcal {G}} _ {r _ {1}, i _ {1}, j _ {1}, r _ {2}} ^ {(1)} \boldsymbol {\mathcal {G}} _ {r _ {2}, i _ {2}, j _ {2}, r _ {3}} ^ {(2)} \boldsymbol {\mathcal {G}} _ {r _ {3}, i _ {3}, j _ {3}, r _ {4}} ^ {(3)} \dots \boldsymbol {\mathcal {G}} _ {r _ {N}, i _ {N}, j _ {N}, r _ {1}} ^ {(N)}, \tag {2}
$$
where  $\{R_1, R_2, \ldots, R_N\}$  denote the ranks,  $\mathcal{G}^{(n)} \in \mathbb{R}^{R_n \times I_n \times I_n \times R_{n+1}}$  denotes a 4th-order core tensor and  $R_1 = R_{N+1}$ , which indicates ring-like structure.
Transformer Architecture. The Transformer [47] is a deep learning architecture that has revolutionized the artificial intelligence field including computer vision (CV) and natural language processing (NLP). As shown in Figure 4, a Transformer block consists of two main sub-layers: the multi-head self-attention (MHSA) layer and the feed-forward neural network (FFN) layer.
(1) MHSA Layer. The MHSA layer in the Transformer block computes the attention scores between each input element and every other element, allowing the model to attend to different parts of the input sequence during processing. Inputs of MHSA are query matrix  $\mathbf{Q} \in \mathbb{R}^I$ , key matrix  $\mathbf{K} \in \mathbb{R}^I$  and value matrix  $\mathbf{V} \in \mathbb{R}^I$  with parameters  $\mathbf{W}^Q \in \mathbb{R}^{I \times O}$ ,  $\mathbf{W}^K \in \mathbb{R}^{I \times O}$ ,  $\mathbf{W}^V \in \mathbb{R}^{I \times O}$ ,  $\mathbf{W}^O \in \mathbb{R}^{I \times O}$ . Usually, it constrains  $I = O$ . Moreover, MHSA separates these parameters into  $n$  heads:  $\{\mathbf{W}^{Q,i}\}^n$ ,  $\{\mathbf{W}^{K,i}\}^n$ ,  $\{\mathbf{W}^{V,i}\}^n$ ,  $\{\mathbf{W}^{O,i}\}^n$ . The MHSA mechanism can be formulated as follows:
$$
\operatorname {A t t} _ {i} (\mathbf {Q}, \mathbf {K}, \mathbf {V}) = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q W} ^ {Q , i} (\mathbf {K W} ^ {K , i}) ^ {T}}{\sqrt {d _ {k}}}\right) \mathbf {V W} ^ {V, i} \mathbf {W} ^ {O, i ^ {T}},
$$
$$
\operatorname {M H S A} \left(\left(\mathbf {Q}, \mathbf {K}, \mathbf {V}\right)\right) = \sum_ {i = 1} ^ {n} \operatorname {A t t} _ {i} (\mathbf {Q}, \mathbf {K}, \mathbf {V}), \tag {3}
$$
where  $d_{k}$  is the dimensionality of the key vectors. The self-attention mechanism is performed multiple times in parallel, each time using a different set of learned parameters  $\mathbf{W}^{Q,i}$ ,  $\mathbf{W}^{K,i}$ , and  $\mathbf{W}^{V,i}$  to compute multiple "heads" of attention. The resulting attention heads are linearly transformed by a learned weight matrix  $\mathbf{W}^{O,i}$  to produce the output. At last, MHSA concatenates the output as the final result of the MHSA layer.
(2) FFN Layer. The FFN layer in the Transformer block is responsible for applying a non-linear transformation to the output of the self-attention layer.  $\mathbf{X} \in \mathbb{R}^I$  is the input. The weights are
Figure 4: The full mapping of the Mango operator. The left sub-figure shows that the parameters of a Transformer layer are  $\mathbf{W}^Q, \mathbf{W}^K, \mathbf{W}^V, \mathbf{W}^O, \mathbf{W}^{IN}$  and  $\mathbf{W}^{OUT}$ .  $I$  means an input dimension size.  $O$  denotes an output dimension size.  $L$  is the layer number. We concatenate all the parameters into a tensor  $\mathcal{M}$  and then consider a full mapping operator  $\mathcal{S}$  to transform this tensor. However,  $\mathcal{S}$  is huge, thereby, we use a multi-linear method TR-MPO to decompose it into four smaller tensors  $\{\mathcal{S}_B, \mathcal{S}_I, \mathcal{S}_O, \mathcal{S}_L\}$  from the Mango operator  $\Phi_{\mathcal{S}_B, \mathcal{S}_I, \mathcal{S}_O, \mathcal{S}_L}$ .
$\mathbf{W}^{IN}\in \mathbb{R}^{I\times kO}$  and  $\mathbf{W}^{OUT}\in \mathbb{R}^{kI\times O}$ , where usually  $I = O$  and  $k$  is a ratio that is often set to 4. The FFN layer can be formulated as
$$
\operatorname {F F N} (\mathbf {X}) = \operatorname {G e L U} \left(\mathbf {X} \mathbf {W} ^ {I N}\right) \mathbf {W} ^ {\text {O U T}}. \tag {4}
$$
We neglect biases in the formulation as it is usually set to 0 at initialization. The output of the FFN layer is obtained by applying two linear transformations to the output of the MHSA layer.
Finally, both the self-attention output and the FFN output are processed through a residual connection and layer normalization to prevent the model from collapsing or overfitting to the training data. Apparently, the parameters in a Transformer are mainly based on its linear transformation matrices, i.e.,  $\mathbf{W}^Q$ ,  $\mathbf{W}^K$ ,  $\mathbf{W}^V$ ,  $\mathbf{W}^O$ ,  $\mathbf{W}^{IN}$  and  $\mathbf{W}^{OUT}$ .
3.2 Multi-linear Operator
The essential of model growth is to transfer the knowledge from the small model to the bigger counterpart. Prior study [5] finds that taking weights from the neighbor layer as initial parameters can further improve the convergence speed, which gives the insights that knowledge from other layers can help training as there are many similar attention maps among layer. Nevertheless, the fact is that this similarity exists all over the pretrained model as shown in Figure 2, which motivates us to consider the possibility of whether we can utilize the knowledge from all weights. Based on this consideration, we construct a mapping of the whole model. As the correlation is hard to formulate heuristically, we learn the mapping parameter to implicitly transfer the pretrained model.
Notation. Here, we give the notations for clearly formulating our method. A model with  $L$  layers and a  $D$  hidden size can be denoted as  $\mathbf{M}(L,D)$ . Following Section 3.1 parameters of  $j$ -th layer are a set  $\pmb{\theta}_j = \{\mathbf{W}_j^Q \in \mathbb{R}^{I \times O}, \mathbf{W}_j^K \in \mathbb{R}^{I \times O}, \mathbf{W}_j^V \in \mathbb{R}^{I \times O}, \mathbf{W}_j^{IN} \in \mathbb{R}^{I \times kO}, \mathbf{W}_j^{OUT} \in \mathbb{R}^{kI \times O}\}$ ,  $j \in [L]$ , usually satisfying  $I = O = D$ . The weight of  $\mathbf{M}(L,D)$  is  $\pmb{\theta}^{L,D} = \{\theta_j\}_{j=1}^L$ . A growth operator with parameter  $\pmb{S}$  is denoted as  $\Phi_{\pmb{S}}$ . A mapping from  $\mathbf{M}(L_1, D_1) \to \mathbf{M}(L_2, D_2)$  can be denoted as  $\pmb{\theta}^{L_2, D_2} = \Phi_{\pmb{S}}(\pmb{\theta}^{L_1, D_1})$ . In the case of  $L_1 < L_2$ , this mapping represents a growth mapping. After growing,  $\pmb{\theta}^{L_2, D_2}$  will be used as initial weights for the target model.
Full Mapping Operator. To utilize all weights for knowledge transferring, we first concatenate weights across layers. As shown in Figure 4, we concatenate  $j$ -th layer to form a tensor of shape  $B \times I \times O$  along with order  $I$  and  $O$  where  $B = 2k + 4$ . Then the final weight tensor  $\mathcal{M} \in \mathbb{R}^{B \times I \times O \times L}$  can be derived by combining the concatenated tensors of  $L$  layers. Giving
bert2BERT
LiGO
Figure 5: Growth processes with the tensor diagram. The circle with an oblique stroke means a super-diagonal tensor. The blue color means a trainable operator, while the gray color denotes an untrainable operator. The red color means the smaller model. Each tensor diagram means the growth process that growing  $M_{1}$  to a bigger model through growth the operator  $S_{*}$ .
a small model  $\mathcal{M}_1\in \mathbb{R}^{B_2\times I_1\times O_1\times L_1}$ , a big model  $\mathcal{M}_2\in \mathbb{R}^{B_1\times I_2\times O_2\times L_2}$ , and a parameter  $\pmb {S}\in \mathbb{R}^{B_1\times I_1\times O_1\times L_1\times B_2\times I_2\times O_2\times L_2}$  of a full mapping operator  $\Phi_{\mathcal{S}}$ , we can transfer  $\mathcal{M}_1$  to  $\mathcal{M}_2$  as
$$
\left(\mathcal {M} _ {2}\right) _ {b _ {2}, i _ {2}, o _ {2}, l _ {2}} = \sum_ {b _ {1}, i _ {1}, o _ {1}, l _ {1}} \left(\mathcal {M} _ {1}\right) _ {b _ {1}, i _ {1}, o _ {1}, l _ {1}} \mathcal {S} _ {b _ {1}, i _ {1}, o _ {1}, l _ {1}, b _ {2}, i _ {2}, o _ {2}, l _ {2}}, \tag {5}
$$
where  $b_{1}, i_{1}, o_{1}, l_{1}, b_{2}, i_{2}, o_{2}$ , and  $l_{2}$  are entries of corresponding tensors.
Multi-linear Mapping Operator. Note that  $\mathcal{S}$  is extremely big, which makes it infeasible to display this mapping by applying a whole transformation tensor. Therefore, we propose a multi-linear operator named Mango which decomposes  $\mathcal{S}$  through a tensor ring matrix product operator (TR-MPO) in four small tensors  $\{\pmb {S}_B\in \mathbb{R}^{R_1\times B_1\times B_2\times R_2},\pmb {S}_O\in \mathbb{R}^{R_2\times O_1\times O_2\times R_3},\pmb {S}_L\in \mathbb{R}^{R_3\times L_1\times L_2\times R_4},\pmb {S}_I\in$ $\mathbb{R}^{R_4\times I_1\times I_2\times R_1}\}$ .  $R = \{R_{1},R_{2},R_{3},R_{4}\}$  is the rank of TR-MPO. Then, we can update  $\Phi_S$  to  $\Phi_{\pmb {S}_B,\pmb {S}_I,\pmb {S}_O,\pmb {S}_L}$ , and Eq. (5) can be reformulated with a multi-linear form as
$$
\begin{array}{l} \left(\mathcal {M} _ {2}\right) _ {b _ {2}, i _ {2}, o _ {2}, l _ {2}} = \sum_ {b _ {1}, i _ {1}, o _ {1}, l _ {1}, r _ {1}, r _ {2}, r _ {3}, r _ {4}} \\ \left(\mathcal {M} _ {1}\right) _ {b _ {1}, i _ {1}, o _ {1}, l _ {1}} \left(\boldsymbol {S} _ {B}\right) _ {r _ {1}, b _ {1}, b _ {2}, r _ {2}} \left(\boldsymbol {S} _ {O}\right) _ {r _ {2}, o _ {1}, o _ {2}, r _ {3}} \left(\boldsymbol {S} _ {L}\right) _ {r _ {3}, l _ {1}, l _ {2}, r _ {4}} \left(\boldsymbol {S} _ {I}\right) _ {r _ {4}, i _ {1}, i _ {2}, r _ {1}}. \tag {6} \\ \end{array}
$$
The total size of  $S_B, S_I, S_O$ , and  $S_L$  is exponentially less than  $S$ , which makes Mango viable for practical implementation while maintaining the full correlation between the small and big models.
Training Target. After designing the multi-linear operators, we train these operators to obtain the function preserving  $\mathcal{M}_2$ . The training target can be denoted as