Optimizing Deeper Transformers on Small Datasets (2021)

Abstract: It is a common belief that training deep transformers from scratch requires large datasets. Consequently, for small datasets, people usually use shallow and simple additional layers on top of pre-trained models during fine-tuning. This work shows that this does not always need to be the case: with proper initialization and optimization, the benefits of very deep transformers can carry over to challenging tasks with small datasets, including Text-to-SQL semantic parsing and logical reading comprehension. In particular, we successfully train 48 layers of transformers, comprising 24 fine-tuned layers from pre-trained RoBERTa and 24 relation-aware layers trained from scratch. With fewer training steps and no task-specific pre-training, we obtain the state of the art performance on the challenging cross-domain Text-to-SQL parsing benchmark Spider. We achieve this by deriving a novel Data dependent Transformer Fixed-update initialization scheme (DT-Fixup), inspired by the prior T-Fixup work. Further error analysis shows that increasing depth can help improve generalization on small datasets for hard cases that require reasoning and structural understanding.

## introduction in recent years, large-scale pre-trained language models @xcite @xcite trained with transformers @xcite have become standard building blocks of modern nlp systems to help improve generalization when task-specific annotations are limited. in practice, it has been found that deeper transformers generally yield better results with sufficient training data @xcite , especially on tasks involving reasoning and structural understanding. this suggests that additional transformer layers should be employed in conjunction with pre-trained models, instead of simple and shallow neural components, such as a classifier head, currently used by models of many nlp tasks. however, the common belief in the literature is that training deep transformers from scratch requires large datasets, and few attempts have been made on small datasets, to the best of our knowledge. one implication is that although extra transformer layers on top of pre-trained models should help with more challenging problems in principle, it does not work in practice due to limited training data. we show that after resolving several optimization issues with the method proposed in this work, it is possible to train very deep transformers with improved generalization even on small datasets. one advantage of pre-trained models is the reduced computational resources needed when finetuning on small datasets. for instance, it allows practitioners to finetune on a single gpu and obtain strong performance on a downstream task. however, the large size of pre-trained models limits the batch size that can be used in training new transformer layers on a small computational budget. despite their broad applications, training transformer models is known to be difficult @xcite . the standard transformer training approach leverages learning rate warm-up, layer normalization @xcite and a large batch size, and models typically fail to learn when missing any one of these components. the restricted batch size aggravates the training difficulties. even if a large batch size can be feasibly employed, poorer generalization results are often observed @xcite , especially when the dataset size is only several times larger than the batch size. furthermore, many recent works noticed a performance gap in this training approach due to layer normalization @xcite @xcite @xcite . inspired by the recent t-fixup by huang et al. @xcite , which eliminates the need for learning rate warm-up and layer normalization to train vanilla transformers, we derive a data-dependent initialization strategy by applying different analyses to address several key limitations of t-fixup. we call our method the data-dependent transformer fixed-update initialization scheme, dt-fixup. in the mixed setup of additional yet-to-be-trained transformers on top of pre-trained models, dt-fixup enables the training of significantly deeper transformers, and is generally applicable to different neural architectures. our derivation also extends beyond vanilla transformers to transformers with relational encodings @xcite , allowing us to apply the results to one variant called relation-aware transformer @xcite . by applying dt-fixup on different tasks, we show that the impression that deep transformers do not work on small datasets stems from the optimization procedure rather than the architecture. with proper initialization and optimization, training extra transformer layers is shown to facilitate the learning of complex relations and structures in the data. we verify the effectiveness of dt-fixup on spider @xcite , a complex and cross-domain text-to-sql semantic parsing benchmark, and recolr @xcite , a reading comprehension dataset requiring logical reasoning. while text-to-sql semantic parsing is inherently different from reading comprehension, they share similar characteristics which require certain levels of reasoning and structural understanding ability. meanwhile, the sizes of both datasets are less than 10k training samples, which is tiny by deep learning standards and renders large-batch training undesirable due to poor generalization foot_0 . on both datasets, dt-fixup consistently outperforms the standard approach with better generalization and allows the training of significantly deeper transformer models. for spider, we successfully apply dt-fixup to train a text-to-sql parser containing 48 transformer layers, with 24 relation-aware layers trained from scratch on top of 24 pre-trained layers from pre-trained roberta @xcite . our parser achieves 70.9% exact match accuracy on the spider test set, which is the state of the art at the time of writing. at the same time, it requires less training steps and no task-specific pre-training as compared to the prior art @xcite . for reclor, we rank the second on the public leaderboard by simply adding 4 transformer layers on top of roberta. further error analysis shows that the performance improvements by increasing the depth mainly come from better generalization on the harder cases requiring reasoning and structural understanding. even the failed predictions from the deep models are more reasonable than from the shallow ones. ## background in this section, we present the necessary background by first introducing the relation-aware transformer layer, which outperforms the vanilla transformer layer with limited data by injecting additional inductive bias @xcite . then, we introduce the t-fixup technique @xcite for optimizing deeper vanilla transformers and discuss why it does not directly apply in the mixed transformer optimization setup. ## our approach we now follow the analysis framework of t-fixup (huang et al., 2020), but derive the conditions to bound the gradient updates of the self-attention block in the presence of a pre-trained model. based on the derivation, we propose a data-dependent initialization strategy for the mixed setup of the new transformers on pre-trained encodings. ## experiments all the experiments in this paper are conducted with a signle 16gb nvidia p100 gpu. ## 35 42 column err. only ## 60 53 sketch err. only ## 105 88 ## conclusion despite the broad applications of the transformer model, it struggles to perform well for some nlp tasks with limited training data. in this work, we propose a theoretically justified optimization strategy dt-fixup to train deeper transformer model with improved generalization and faster convergence speed on small datasets, which is generally applicable to different neural architectures. on two important tasks, text-to-sql semantic parsing and logical reading comprehension that require reasoning and structural understanding, applying dt-fixup achieves sota or near-sota results by simplying using extra transformer layers on top of the pre-trained models. such observations suggest even boarder applicability of deeper transformers. proof. first, let's inspect the feedforward pass through the transformer blocks, which have nonlinear layers g l 's and skip connections: for vanilla self-attention and mlp @xcite this is due to the fact that the probability from softmax sums to one, so does not alter the overall norm; at initialization, values are at the linear identity range of the nonlinearities. therefore, for all three types of layers: for vanilla self-attention or mlp layers: and for relational self-attention: at initialization, we want v l , w l , r v l of all layers to have the same norm, i.e. v l similarly if f g is vanilla transformer instead of a relational one, we have: the only variable that still depends on l is x x x l , which by expanding the recursion in eq. 9-10, gives: now let x x x θ = µ , and we have assumed that µ 1, which is very common for output of pre-trained encoders, and due to the high dimensionality. and let ξ = n (4µ 2 + 2µ + 2) 21, eq. 17 becomes : this concludes the proof for relational transformers. for vanilla transformers, with ξ = n (4µ 2 ) -1 2 , and following the same steps, but plugging into eq. 18, we have