T encent P retrain: A Scalable and Flexible Toolkit for Pre-training Models of Different Modalities (2023)

Abstract: Recently, the success of pre-training in text domain has been fully extended to vision, audio, and cross-modal scenarios. The proposed pre-training models of different modalities are showing a rising trend of homogeneity in their model structures, which brings the opportunity to implement different pre-training models within a uniform framework. In this paper, we present TencentPretrain, a toolkit supporting pre-training models of different modalities. The core feature of TencentPretrain is the modular design. The toolkit uniformly divides pre-training models into 5 components: embedding, encoder, target embedding, decoder, and target. As almost all of common modules are provided in each component, users can choose the desired modules from different components to build a complete pre-training model. The modular design enables users to efficiently reproduce existing pre-training models or build brand-new one. We test the toolkit on text, vision, and audio benchmarks and show that it can match the performance of the original implementations.

## introduction pre-training on large-scale data and then finetuning on downstream tasks has become a paradigm for text, vision, and audio tasks @xcite @xcite . in addition to the similarity in the pipeline paradigm, these pre-training models as well have close model structures: on one hand, most of them consist of the following components, embedding, encoder, target embedding, decoder, and target, on the other hand, many modules in above components are shared among models of different modalities. for example, the transformer module (in encoder component) @xcite , which is successful in the field of text, is increasingly being applied to the vision and audio modalities. @xcite . the trend towards homogeneity in pre-training models is becoming more apparent, which makes it possible to integrate them into a uniform framework. a representative work in this direction is huggingface transformers @xcite , which exploits a non-modular design mode. for each pre-training model in huggingface transformers, several separate classes are created, and the code is not refactored with additional abstractions. users can develop their pre-training models independently which is useful to collaborative development in the community. however, in this design mode, users need to implement the model from scratch when adding a new pre-training model, requiring considerable code work. in addition, with the increased number of pre-training models, the number of classes and lines of code also increases linearly. codes with the same function may be written many times, which degrades the readability and maintainability of the project. in response to shortcomings of non-modular design mode, we introduce tencentpretrain, a modular toolkit specially designed for pre-training models of different modalities. as shown in figure 1 , tencentpretrain has five components, namely embedding, encoder, target embedding, decoder, and target. among them, target embedding and decoder components are optional, since the targets of many pre-training models do not involve sequence decoding @xcite . tencentpretrain is hierarchical modular designed with two degrees of freedom. at component level, users are free to combine modules within a component, for example, combining multiple modules in target component to perform multi-task pretraining @xcite . at pre-training model modality embedding encoder target embedding decoder target elmo (peters et al., 2018) text word bi-lstm --bilm infersent @xcite text word gru --cls cove @xcite text word lstm word lstm lm bert @xcite text word, pos, seg transformer --mlm, sp gpt-2 @xcite text word, pos transformer --lm t5 @xcite text word transformer word transformer lm vit @xcite vision patch, pos transformer --cls beit @xcite vision patch, pos transformer --mlm s2t @xcite audio speech, pos transformer word, pos transformer lm vilt @xcite text-vision word_patch, pos, seg transformer --mlm, cls the model level, users can combine modules from different components to constitute a complete pretraining model. modularity in design makes tencentpretrain scalable with the increasing number of newly proposed pre-training models. users are allowed to reuse existing modules with little efforts, avoiding repeated implementation of core functions. at the same time, tencentpretrain provides a robust and clear interface among different components. it brings flexibility, allowing users to build custom model structures through a configuration file without any code work. tencentpretrain is implemented with pytorch (paszke et al., 2019), and it supports distributed training and deepspeed optimization library @xcite . tencentpretrain is fully connected with huggingface transformers, providing comprehensive conversion scripts of pre-training models between the two frameworks. users can switch between the two frameworks at low cost. tencentpretrain is tested on text, vision, and audio benchmarks and is able to reproduce the results of sota pre-training models. the ten-centpretrain toolkit is publicly available at https: //github.com/tencent/tencentpretrain . ## framework the current mainstream pre-training models are basically similar in structure. in the embedding component, the data is mapped into an embedding matrix. and then the matrix is passed through the encoder. finally the target layer performs pretraining tasks according to the output of the encoder layer. if the pre-training task requires sequence generation, the decoder is inserted between the encoder and the target. figure 1 demonstrates the overall framework of tencentpretrain. it divides a pre-training model into five components, and various modules are provided in each component. in practice, a user firstly selects one or multiple modules from each component (modularization within component), and then combine modules from different components to build a pre-training model (modularization cross components). in the rest of this section, we respectively introduce the above five components and modules included in them. ## usage this section provides examples of building pretraining models with tencentpretrain. the modular design enables the users to quickly build the pretraining model through the combination of modules. modules used in pre-training models are specified in configuration files and the examples are shown as follows foot_0 : # bert implementation { "embedding" : [ " word " , " p o s " , " s e g " ] , "encoder" : " t r a n s f o r m e r " , "target" : [ " mlm " , " s p " ] } if the desired pre-training model cannot be built by the combination of existing modules, tencent-pretrain encourages users to develop a new module, and combine it with existing modules. we take the implementation of asr model s2t @xcite as an example. most modules required by s2t are available and we only need to implement a new module, speech embedding, which greatly speeds up the implementation process. tencentpretrain and huggingface transformers are interoperable. the conversion scripts are publicly available foot_1 , and the weights of different pretraining models can be converted between the two frameworks. in practical use, users are free to switch between these two frameworks. with tencentpretrain, we build a pre-trained weight model zoo. each pre-trained weight has two versions which can be loaded by either tencentpretrain or huggingface transformers. currently, the tencentpretrain model zoo includes over 50 pretrained weights. we provide pre-training data as well as training details, allowing users to reproduce results with less effort. the weights (pre-trained by tencentpretrain) are currently downloaded over 500 thousand times per month foot_2 . model hf uer tp transformer 1135 749 795 +bert +822 +130 (+word_pos_seg, bert) +149 (+pos,seg, mlm,sp) +roberta +696 +92 (+word_pos,mlm) +0 +gpt-2 +688 +0 +0 +t5 +1008 +17(+word) +0 +vit +493 -+59(+patch) +s2t +824 -+51(+speech) +vilt +618 -+15 (+word_patch) ## evaluation this section evaluates tencentpretrain framework quantitatively. firstly, we compare tencentpretrain with non-modular framework in terms of implementation cost. then we show that tencentpretrain can reproduce the results of sota models on a range of benchmarks. ## conclusion this paper presents tencentpretrain, a pre-training toolkit characterized by modular design and multimodal support. in tencentpretrain, pre-training models of different modalities are regarded as the combination of multiple modules, which is easy to configure and extensible. furthermore, we quantitatively demonstrate that tencentpretrain facilitates the users to reuse existing modules and decreases the cost of model development. at last, we test tencentpretrain on a range of datasets and show that it can reproduce the sota results.