PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping

Minjia Zhang, Yuxiong He

arXiv:2010.13369v1cs.LG

TL;DR

Transformer pre-training is computationally expensive, and some existing acceleration approaches rely on specialized hardware or large clusters. The paper uses progressive layer dropping with Switchable-Transformer blocks to accelerate BERT pre-training. It reports reduced training time while maintaining similar downstream accuracy and strong GLUE performance.

  • Problem

    Transformer pre-training remains computationally expensive, while some acceleration methods require specialized hardware or large clusters and direct layer dropping can destabilize training.

    Source: §1 · pp. 1–2

  • Method

    Progressive layer dropping uses Switchable-Transformer blocks and a schedule that increases layer dropping over training, bypassing blocks during selected updates.

    Source: §1 · p. 2, §4.2 · pp. 5–6

  • Results

    24% time reduction when training the same number of samples; PLD reaches similar downstream accuracy in 15.56h versus the baseline's 39.15h, a 2.5× speedup.

    Source: §5.1 · pp. 7–8

  • Takeaways & Limitations

    At the same number of samples, the method achieves a 1.1% higher GLUE score than the baseline, indicating strong knowledge transferability.

    Source: §1 · p. 2

Abstract

from arXiv · show

Recently, Transformer-based language models 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 language modeling. In this work, we propose a method based on progressive layer dropping that speeds the training of Transformer-based language models, 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.

1 Introduction

Transformer pre-training has become central to NLP, but its computational cost remains high and existing acceleration can require specialized hardware or large clusters. The paper addresses unstable layer dropping with a Switchable-Transformer block and a progressive schedule, reporting faster BERT pre-training and strong downstream performance.

1 Introduction

  • Existing acceleration approaches can depend on Tensor Cores or large distributed clusters, placing them beyond the resources of many research labs and individual practitioners.Large-batch distributed training can also converge to sharp local minima with poor generalizability.Source: §1 · pp. 1–2
  • Directly dropping Transformer layers destabilizes pre-training and can cause divergence or suspicious local optima, motivating the paper's focus on stable layer dropping.Reducing model depth outright can also lower downstream accuracy relative to full-model pre-training.Source: §1 · p. 2
  • BERT pre-training with the proposed method is 2.5× faster than the baseline to reach similar downstream accuracy, while same-sample training yields a 1.1% higher GLUE score.The method also reduces training time by 24% for the same number of samples.Source: §1 · p. 2
  • Switchable-Transformer blocks bypass layers during selected training intervals, and a progressive schedule increases layer dropping as training evolves.The schedule distributes a global drop rate across blocks to favor different layers.Source: §1 · p. 2

2 Background and Related Work

The paper situates BERT pre-training among Transformer-based representation-learning methods and prior efforts to reduce its cost. It reviews depth reduction and layer dropping as opportunities that also carry accuracy or stability challenges.

2 Background and Related Work

  • BERT stacks Transformer layers, each applying attention and feed-forward sublayers with residual connections and layer normalization.Pre-training on unlabeled corpora provides initialization for downstream fine-tuning.Source: §2 · p. 2
  • Because BERT's forward and backward FLOPS scale linearly with Transformer depth, removing layers can reduce per-iteration computation, while ELECTRA instead trades higher per-step cost for fewer steps.The cited schedules reduce most FLOPS in later steps as keep probability approaches a fixed value.Source: §2 · pp. 2–3
  • Training a shallower BERT can significantly reduce accuracy, while prior progressive network growth requires choosing in advance how many steps to train at each depth.Direct stochastic dropping of attention and feed-forward sublayers can converge to bad local optima or diverge at higher learning rates.Source: §2 · p. 3

3 Preliminary Analysis

The preliminary analyses examine normalization placement, representation changes during training, and the effects of removing layers. Together, they identify conditions under which layer dropping is less disruptive.

3.1 Training Stability: PostLN or PreLN?

  • PreLN keeps gradient norms more balanced across depth than PostLN, whose gradients can vanish toward lower layer IDs; PreLN also has input-gradient norm-preserving ratios near 1.PostLN is more sensitive to hyperparameters and often diverges with aggressive learning rates, whereas PreLN supports more stable optimization.Source: §3.1 · p. 3

3.2 Corroboration of Training Dynamics

  • As training proceeds, PreLN's upper layers produce increasingly similar representations, consistent with successive layers iteratively refining estimates rather than creating entirely new representations.PostLN's dissimilarity remains relatively high and bumpy; the analysis compares embedding L2 distance and inverse cosine similarity across layers and steps.Source: §3.2 · p. 4, Figure 4 · p. 4

3.3 Effect of Lesioning

  • With a keep ratio of 0.5, layer removal substantially harms PostLN and can cause divergence at higher learning rates, while its effect on PreLN is modest.Higher PreLN layers are less dependent on individual layers, while lower layers that derive new features should be dropped less often.Source: §3.3 · p. 4

4 Our Approach: Progressive Layer Dropping

Progressive layer dropping (PLD) accelerates Transformer pre-training using Switchable-Transformer blocks and a schedule that progressively reduces the retained layers. The approach combines identity-mapped blocks, stochastic sublayer gates, and depth-aware retention.

4.1 Switchable-Transformer Blocks

  • Switchable-Transformer blocks use PreLN identity mapping and gates that can bypass each sublayer during training.The proposed block differs from the original Transformer block through these two changes.Source: Figure 5 · p. 5, §4.1 · p. 5
  • Each sublayer gate is sampled from a Bernoulli distribution, choosing whether its transformation is active during training.The gate takes values 0 or 1, with the probability of choosing 1 denoted by p_i.Source: §4.1 · p. 5

4.2 A Progressive Layer Dropping Schedule

  • Retention starts at 1 and decreases toward a floor, while lower layers receive lower drop probabilities and both sublayers share a schedule.The schedule distributes the global retention target across layers according to depth.Source: §4.2 · pp. 5–6
  • Because bypassed blocks retain identity connections, they can skip forward-backward computation and gradient updates during training.The resulting stochastic training algorithm is PLD.Source: §4.2 · p. 6
  • With a 12-layer BERT and a 0.5 keep-probability floor, PLD averages 9 active layers and saves approximately 25% of FLOPS; full-depth blocks return for fine-tuning and testing.The reported expected depth is approximately 3L/4 when training is long enough.Source: §4.2 · p. 6

5 Evaluation

The evaluation reports faster BERT pre-training with competitive GLUE performance, then examines convergence, fine-tuning sensitivity, and design choices behind progressive layer dropping.

5 Evaluation

  • The evaluation uses a 12-layer BERT-base baseline pretrained on Wikipedia and BookCorpus, with GLUE fine-tuning results reported as medians over five initializations.The pretraining corpus contains roughly 2.8B word tokens; GLUE covers nine sentence or sentence-pair understanding tasks.Source: §5 · p. 7

5.1 Experimental Results

  • At lrmax=1e−4, PLD and baseline converge at nearly the same rate, while at 1e−3 the baseline diverges and PLD converges healthily.The authors report that the architectural changes stabilize training and allow more aggressive learning rates.Source: §5.1 · p. 7
  • 24% fewer training seconds per sample and 2.8× end-to-end wall-clock speedup accompany PLD’s lower expected depth per step.The per-sample reduction is slightly below the 25% GFLOPS reduction because the output layer still requires some computation.Source: §5.1 · p. 7, Figure 7 · p. 7
  • 2.5× speedup (15.56h vs. 39.15h) yields similar downstream accuracy; PLD scores 82.3 versus the baseline’s 82.1 GLUE score at checkpoint 100.At checkpoint 186, PLD scores 83.2 versus 82.1 for the baseline; at checkpoint 87, PLD scores 81.6 versus 82.1.Source: Table 1 · p. 7, §5.1 · p. 8

5.2 Ablation Studies

  • PLD is more robust than the baseline to fine-tuning learning-rate choice, often performing better with large rates.The grid search tests rates from 1e-5 through 1e-4.Source: §5.2 · p. 8, Figure 8 · p. 9
  • Keep ratios from 0.5 to 0.9 are identified as a good range, while a ratio as small as 0.3 may cause divergence.The keep-ratio convergence curves are shown in Figure 10.Source: §5.2 · p. 8, Figure 10 · p. 10
  • With large learning rates, PreLN scores 82.6 versus PLD’s 83.2 despite using 24% more compute, while shallow BERT scores 1.4 points below PLD.The 9-layer shallow model has the same training GFLOPS as PLD; PreLN with PostLN hyperparameters scores 80.2 versus PostLN’s 82.1.Source: §5.2 · p. 9, Table 3 · p. 10
  • Fixed-ratio random dropping scores 0.5 points below PLD, and the authors report that the depth schedule accounts for most accuracy gains while the temporal schedule supports half-precision stability.Without the temporal schedule, half-precision training diverges and must switch to slower full precision; temporal-only scheduling matches random dropping in GLUE score.Source: §5.2 · pp. 9–11

6 Conclusion

The conclusion presents progressive layer-wise dropping with a Switchable-Transformer block as a way to train BERT faster while maintaining competitive performance.

6 Conclusion

  • The proposed training strategy achieves competitive performance to training a deep model from scratch at a faster rate.The approach uses a Switchable-Transformer block and a progressive layer-wise drop schedule, informed by analysis of architecture and stochastic-depth training.Source: §6 · p. 11

A Pre-training Hyperparameters

The appendix identifies a table of pre-training hyperparameters for the baseline and PLD.

A Pre-training Hyperparameters

B Establishing Identity Mapping with PreLN

PreLN expresses each Transformer layer as its input plus a residual transformation, preserving an identity path. The paper links this path to more balanced gradients and contrasts it with PostLN’s gradient and training difficulties.

B Establishing Identity Mapping with PreLN

  • PreLN writes each layer output as its input plus a residual transformation, preserving an identity mapping path.The residual transformation combines self-attention and feed-forward functions applied after layer normalization.Source: B Establishing Identity Mapping with… · p. 15
  • The gradient decomposition includes a direct term that propagates information to shallower blocks without depending on the complexity of intervening transformations.A second term propagates through the Transformer blocks.Source: B Establishing Identity Mapping with… · p. 15
  • Unlike PreLN, PostLN repeatedly normalizes the skip-connection signal, impeding information propagation and causing vanishing gradients and training instability.The passage attributes avoiding vanishing or exploding gradients, stable optimization, and performance gain to PreLN.Source: B Establishing Identity Mapping with… · p. 15

C PreLN From the View of Unrolled Iterative Refinement

The iterative-refinement view treats successive PreLN layer outputs as estimates of a shared latent representation, with zero-mean residuals supporting an unbiased expectation. In fine-tuning, PLD is more robust to learning-rate choice than the baseline, particularly at large learning rates.

C PreLN From the View of Unrolled Iterative Refinement

  • Under the iterative-refinement view, successive layer outputs estimate the same unknown latent representation, and PreLN makes the expected difference between consecutive outputs zero.The residual is therefore zero-mean, which the authors report empirically verifying; the residual-function sum determines estimate variance.Source: C PreLN From the View of Unrolled… · pp. 15–16
  • The learning-rate evaluation excludes very small datasets because their validation scores vary substantially across random seeds.Fine-tuning uses batch size 32, five epochs, and median development-set scores across five initializations without ensembling.Source: C PreLN From the View of Unrolled… · p. 16
  • PLD is more robust to learning-rate choice than the baseline and often performs better with large learning rates.The baseline’s fine-tuning results are often much worse at large learning rates.Source: C PreLN From the View of Unrolled… · p. 16

Related papers

Loading 2010.13369v1…