Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping (2010.13369v1)

Abstract: Recently, Transformer-based LLMs have demonstrated remarkable performance across many NLP domains. However, the unsupervised pre-training step of these models suffers from unbearable overall computational expenses. Current methods for accelerating the pre-training either rely on massive parallelism with advanced hardware or are not applicable to LLMing. In this work, we propose a method based on progressive layer dropping that speeds the training of Transformer-based LLMs, not at the cost of excessive hardware resources but from model architecture change and training technique boosted efficiency. Extensive experiments on BERT show that the proposed method achieves a 24% time reduction on average per sample and allows the pre-training to be 2.5 times faster than the baseline to get a similar accuracy on downstream tasks. While being faster, our pre-trained models are equipped with strong knowledge transferability, achieving comparable and sometimes higher GLUE score than the baseline when pre-trained with the same number of samples.

Accelerating the Training of Transformer-Based LLMs through Progressive Layer Dropping

The paper "Accelerating Training of Transformer-Based LLMs with Progressive Layer Dropping" addresses the computational challenges associated with pre-training Transformer-based LLMs, notably BERT. As the unsupervised pre-training of these models is computationally expensive, the authors propose a novel method utilizing progressive layer dropping to enhance training efficiency without heavy reliance on expensive hardware setups.

Overview

LLMs based on Transformers have significantly advanced NLP tasks, achieving superior performance benchmarks. However, the computational demand during the pre-training phase remains a considerable challenge. Traditional methods to reduce this load often involve large-scale parallelism requiring substantial computational resources, making them impractical for most research facilities.

This paper introduces an architectural adjustment and training method called Progressive Layer Dropping. The focus is on reducing the computational cost of training Transformer networks by dynamically adapting the network architecture during training. The essence of this technique lies in strategically omitting certain layers within Transformer models, thus reducing computational load while maintaining model performance.

Key Contributions

1. Switchable-Transformer (ST) Blocks: The introduction of ST blocks is core to the proposed method, enabling the activation or deactivation of specific Transformer layers during training through a gating mechanism. These blocks utilize identity mappings to stabilize the training process, informed by a comprehensive analysis of training dynamics highlighting challenges associated with stochastic depth approaches.
2. Progressive Layer Dropping Schedule: The authors propose a progressive schedule for layer dropping that dynamically adjusts the dropout rate over time, allowing lower layers to be more reliably present during early training stages. This approach adapts to the natural phases of Transformer model training—beginning with full-depth models during initial training stages (when variance is high) and progressively increasing drop rates as training stabilizes.
3. Implementation and Efficiency Gains: Through extensive experimentation with BERT, the authors demonstrate that their method achieves a 24% reduction in time per training sample, and enables pre-training to be completed 2.5 times faster compared to traditional methods. This speedup is obtained while retaining comparable or superior accuracy on downstream tasks such as those in the GLUE benchmark.

Implications and Future Work

The proposed method holds substantial theoretical and practical implications. By reducing the computational overhead associated with large-scale models, it broadens accessibility, enabling more institutions to engage in high-caliber NLP research. In practice, this approach offers a potential pathway to faster deployment of LLMs in real-world applications where latency and cost constraints are paramount.

Potential avenues for future exploration include broader application of this method across different Transformer architectures and tasks beyond NLP, such as vision models. Integrating this approach with other acceleration techniques, such as mixed-precision training or distributed computing, could lead to further improvements in training efficiency without sacrificing model accuracy.

In summary, the paper advances understanding in model optimization by providing an empirically validated method that reduces training time significantly while maintaining model performance, thus offering a more viable route to training large-scale Transformer models.