PaperlayerView exact revision on arXiv

Source-linked AI summary

The Lottery Ticket Hypothesis for Pre-trained BERT Networks

Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, Michael Carbin

arXiv:2007.12223v2cs.LGcs.CLcs.NEstat.ML

TL;DR

Pre-trained BERT had not been fully understood under the lottery ticket hypothesis, especially regarding initialization and cross-task transfer. The paper uses pruning to identify sparse subnetworks and evaluates them across NLP tasks, finding matching tickets at substantial sparsity and universal transfer from masked-language-modeling subnetworks.

  • Problem

    The paper investigates whether pre-trained BERT contains trainable, transferable sparse subnetworks and whether universal subnetworks exist across downstream tasks.

  • Method

    The study uses unstructured magnitude pruning with rewinding to identify subnetworks in pre-trained BERT and evaluates their task and transfer performance.

  • Results

    Matching subnetworks appear at 40% to 90% sparsity, while 70%-sparse subnetworks found through masked language modeling transfer universally and maintain accuracy.

  • Takeaways & Limitations

    Lottery ticket observations extend to pre-trained BERT, including universal subnetworks that can replace full BERT while retaining transfer ability.

  • Takeaways & Limitations

    Transfer comparisons are confounded because winning tickets from different tasks have different sizes.

Abstract

from arXiv · show

In natural language processing (NLP), enormous pre-trained models like BERT have become the standard starting point for training on a range of downstream tasks, and similar trends are emerging in other areas of deep learning. In parallel, work on the lottery ticket hypothesis has shown that models for NLP and computer vision contain smaller matching subnetworks capable of training in isolation to full accuracy and transferring to other tasks. In this work, we combine these observations to assess whether such trainable, transferrable subnetworks exist in pre-trained BERT models. For a range of downstream tasks, we indeed find matching subnetworks at 40% to 90% sparsity. We find these subnetworks at (pre-trained) initialization, a deviation from prior NLP research where they emerge only after some amount of training. Subnetworks found on the masked language modeling task (the same task used to pre-train the model) transfer universally; those found on other tasks transfer in a limited fashion if at all. As large-scale pre-training becomes an increasingly central paradigm in deep learning, our results demonstrate that the main lottery ticket observations remain relevant in this context. Codes available at https://github.com/VITA-Group/BERT-Tickets.

1 Introduction

This paper asks whether pre-trained BERT contains sparse subnetworks that can train independently and transfer across downstream tasks. It finds such subnetworks at initialization, including universal subnetworks identified through masked language modeling.

  • Motivation: BERT models are a useful setting for testing trainable, transferable subnetworks because pre-training is amortized across downstream tasks.The paper explicitly asks whether pre-trained weights contain task-specific or universal matching subnetworks.
  • Research gap: The study addresses limited understanding of lottery ticket behavior in pre-trained BERT, focusing on initialization and transfer across tasks.It compares the transformer architecture and BERT’s lengthy pre-training initialization with prior lottery ticket results.
  • Findings: 40% to 90% sparsity: unstructured magnitude pruning finds matching subnetworks on standard GLUE and SQuAD tasks.These subnetworks are evaluated on a range of downstream NLP tasks.
  • Findings: Pre-trained initialization: matching subnetworks are found before downstream training rather than only after early training.This differs from prior NLP findings, where matching subnetworks emerged after some training.
  • Findings: 70% sparsity: subnetworks found with masked language modeling transfer universally to other tasks while maintaining accuracy.Subnetworks found on most downstream tasks instead have task-specific sparsity patterns and limited transfer.
  • Implications: The results extend lottery ticket observations to BERT and indicate that universal subnetworks can replace full BERT without inhibiting transfer.The paper frames smaller-from-the-start training as a possible consequence of these findings.

2 Related Work

Related work applies lottery ticket methods and compression techniques to NLP models, but prior results generally find matching subnetworks only after training. This paper instead studies universally trainable compressed BERT subnetworks across downstream tasks.

  • BERT compression: Compression research for BERT includes pruning, quantization, parameter-sharing, and distillation.The paper focuses specifically on pruning because it underlies the lottery ticket hypothesis.
  • BERT compression: BERT pruning typically targets post-training inference compression, whereas this work seeks compressed models that remain trainable across downstream tasks.The contrast is between task-specific inference efficiency and universal trainability.
  • Lottery tickets in NLP: Prior NLP studies found matching subnetworks early in training on Transformers and LSTMs, but not at initialization.A concurrent BERT study examines related questions using different pruning structures and methods.

3 Preliminaries

The paper defines subnetworks, matching tickets, and universal subnetworks, then identifies them with iterative magnitude pruning and rewinding. Experiments use pre-trained BERT across masked language modeling and downstream NLP tasks.

  • Experimental settings: The experiments use BERTBASE with 12 transformer blocks, hidden size 768, 12 attention heads, and 110M parameters.Each downstream task adds a task-specific classification layer containing less than 3% of the network’s parameters.
  • Experimental settings: Evaluation covers masked language modeling, nine GLUE tasks, and SQuAD v1.1 using standard hyperparameters and validation or development datasets.The downstream tasks include sentence-pair classification and other task categories described in the experimental setup.
  • Preliminaries: A subnetwork is BERT with a binary pruning mask that fixes selected weights to zero.The masked model is written as f(x; m ⊙θ, ·).
  • Preliminaries: A matching subnetwork is one whose trained evaluation metric on task T is no lower than the unpruned pre-trained BERT baseline.The definition compares training from the selected initialization against the full model.
  • Preliminaries: A universal subnetwork uses one shared pruning mask while allowing task-specific configurations for multiple tasks.The paper denotes these configurations as γ_Ti for tasks {T_i}^N.
  • Identifying subnetworks: Iterative magnitude pruning trains the unpruned network, repeatedly removes 10% of remaining lowest-magnitude weights, and rewinds surviving weights to step i.Returning the rewound masked network yields the subnetwork at target sparsity s.
  • Evaluation: Transfer is assessed against both unpruned-task performance and same-task subnetwork performance.Dark cells indicate that transfer matches or exceeds the relevant comparison.

4 The Existence of Matching Subnetworks in BERT

Pre-trained BERT contains matching subnetworks at substantial sparsities, and their specific masks and initialization matter. Unlike prior larger-scale NLP and transformer settings, these winning tickets can be found at pre-trained initialization, while rewinding offers no consistent benefit.

  • Winning tickets: 40% to 90% sparsity: IMP finds winning tickets for MLM and every downstream task evaluated.A winning ticket is within one standard deviation of full BERT; the highest sparsities range from 40% on SQuAD to 90% on QQP and WNLI.
  • Winning-ticket ingredients: Random pruning and random reinitialization perform substantially worse than IMP subnetworks, showing that both the pruning mask and pre-trained initialization matter.On MNLI, performance drops by 15 percentage points with random pruning and 21 percentage points with reinitialization.
  • Winning-ticket ingredients: At high sparsities, random pruning approaches IMP only at 10% to 20% sparsity, confirming the importance of mask structure for IMP performance.The comparison covers CoLA, SST-2, and SQuAD across sparsities.
  • Rewinding: Rewinding does not notably improve downstream performance and can make STS-B and RTE subnetworks non-matching.The authors note that small training sets may cause overfitting in these tasks.
  • Standard pruning: Standard pruning varies by task: it improves some IMP results by up to two percentage points but reduces others by up to three percentage points.Improvements occur for QQP, QNLI, MRPC, and MLM; drops occur for STS-B, WNLI, RTE, and SST-2.
  • Summary: Downstream subnetworks share many pruned weights, whereas MLM subnetworks show larger differences; overall, BERT extends lottery-ticket observations to pre-trained initialization.This initialization result differs from prior NLP settings where matching subnetworks emerged only after training.

5 Transfer Learning for BERT Winning Tickets

The section evaluates whether BERT winning-ticket subnetworks transfer across downstream tasks and whether transfer depends on task, sparsity, initialization, or training data.

  • Transfer methodology: Winning tickets are identified on source task S, trained on target tasks T, and evaluated using TRANSFER(S, T), with a new target-specific classification layer.Transfer comparisons use subnetworks found through the IMP routine.
  • Transfer methodology: Transfer is assessed against same-task performance at 70% sparsity to control for source subnetworks that are too sparse to remain winning tickets.Figure 3 reports TRANSFER(S, T) − TRANSFER(T, T), averaged over three runs.
  • Transfer results: Only three of eleven source tasks transfer to more than two other tasks, although seven transfer to at least half of the other tasks when a 2.5-percentage-point performance drop is allowed.Transferability is therefore usually limited under the exact same-task criterion.
  • Transfer results: MLM produces the best-transfer subnetwork and matches same-task performance across all target tasks, making it universal under the reported criterion.Directly pruning θ0 instead matches same-task performance in only four cases, though it still outperforms every non-MLM source task.
  • Transfer results: Transferability appears related to training-set size: MRPC and WNLI transfer weakly, whereas MNLI and SQuAD transfer to four and three other tasks.MLM has by far the largest training set and the strongest transfer.
  • Initialization: Initializing from θ0 is not noticeably better for transfer than initializing from weights fine-tuned on SQuAD by varying amounts.The comparison uses rewinding and standard-pruning variants focused on SQuAD.

6 Conclusions and Implications

The conclusions find that pre-trained BERT contains sparse, trainable subnetworks and universal subnetworks that preserve transfer across downstream tasks. They propose IMP with MLM as a possible second stage of BERT pre-training, while noting that hardware speedups remain future work.

  • Conclusions: Pre-trained BERT contains sparse subnetworks at non-trivial sparsities that train in isolation to full performance across downstream tasks.The findings extend core lottery-ticket observations to BERT with pre-trained initialization.
  • Implications: Universal subnetworks transfer to all evaluated downstream tasks, allowing a smaller subnetwork to replace full BERT while maintaining its transfer ability.The conclusion frames transfer preservation as the key practical property.
  • Implications: IMP using MLM could serve as a second BERT pre-training stage that produces an equally capable subnetwork with far fewer parameters.The authors identify hardware-platform speedup measurements as future work.

Broader Impact

The broader-impact discussion identifies potential reductions in BERT fine-tuning costs but does not claim new learning capabilities or significant societal harm.

  • Potential impact: Universal lottery-ticket subnetworks may reduce the energy and financial costs of fine-tuning BERT on suitable hardware platforms.The potential benefit depends on appropriate hardware support.
  • Scope and risk: The subnetworks appear to have the same expressive power as their full networks, so the work does not enable learning paradigms unavailable to full BERT.This is presented as the basis for assessing societal risk as limited.

A Further Results on the Existence of Matching Subnetworks in BERT

Additional results examine stricter winning-ticket criteria, performance across sparsities, run variability, metric choice, and MLM transfer comparisons.

  • Stricter winning-ticket criterion: Under the stricter criterion of matching or exceeding unpruned BERT, winning-ticket sparsities decrease for several tasks, including SQuAD from 40% to 20% and MLM from 70% to 50%.MNLI and WNLI are unchanged under this criterion.
  • Sparsity and accuracy: Figure 11 reports IMP-subnetwork performance across all sparsities for each task, providing broader context for the sparsity–accuracy relationship.The figure complements the highest-sparsity summaries.
  • Performance summaries: Table 4 reports task performance at the highest winning-ticket sparsity, averaging five runs and reporting standard deviations as errors.Winning tickets match or exceed full BERT under the table’s criterion.
  • Performance summaries: Table 5 reports mean, median, and best performance across five runs for Full BERTBASE and the identified sparse subnetworks.The best sparse-subnetwork numbers are reported as consistent with HuggingFace results.
  • Metric sensitivity: Except for STS-B, winning-ticket sparsities remain the same across the reported alternative evaluation metrics.STS-B differs between 50% Pearson and 40% Spearman criteria.
  • MLM transfer comparison: Table 7 compares target-task subnetworks with MLM subnetworks at the same sparsity level to assess transfer performance.The comparison directly evaluates MLM-based transfer against corresponding target-task subnetworks.

B Further Results on Transfer Learning for BERT Winning Tickets

MLM-found IMP subnetworks are broadly universal across downstream tasks, including at the most extreme task-specific winning-ticket sparsities. Additional analyses show transfer remains strong under lower sparsity and dataset-size constraints.

  • Universality: For five of ten downstream tasks, MLM subnetworks remain winning tickets at each task’s highest winning-ticket sparsity.Three additional tasks have gaps of at most 0.5 percentage points; MRPC and CoLA have gaps of 1.6 and 2.6 percentage points, respectively.
  • Additional sparsities: At 50% sparsity, weights trained on the source task improve transfer for MNLI but degrade transfer for MLM.

C Finding Subnetworks with Multi-task Pruning

Multi-task pruning provides only a marginal transfer gain over pruning subnetworks using MLM alone. Adding downstream-task information before pruning therefore does not significantly improve transfer.

  • Multi-task-trained BERT subnetworks at 70% sparsity provide only a marginal transfer gain over subnetworks found using MLM alone.The multi-task objective combines MLM with downstream tasks before pruning.

D Similarity between Sparsity Patterns

Downstream-task subnetworks have highly overlapping sparsity patterns, whereas MLM subnetworks differ more. This structural similarity contrasts with the observed transfer behavior.

  • Downstream-task subnetworks share more than 90% of pruned weights in common.The overlap ratio is computed between masks from pairs of tasks.
  • Despite their similarity, downstream-task subnetworks transfer poorly between tasks, while the less-similar MLM subnetwork transfers well.

E Influence of Training Dataset Size on Transfer

The dataset-size analysis tests whether MLM subnetworks’ strong transfer depends on MLM’s larger training set. More MLM examples yield a small, consistent improvement, but transfer remains at least as good as same-task performance under downstream-sized datasets.

  • The MLM transfer analysis constrains training examples to match SST-2, CoLA, and SQuAD, and also evaluates 160,000 examples.The corresponding dataset sizes are 67,360, 8,576, and 88,656 training examples.
  • More training samples produce a small but consistent improvement in MLM subnetwork transfer performance at 70% sparsity.
  • The transfer figures compare source-task subnetworks with target-task performance across sparsities and pruning or rewinding settings.Figures 5–9 cover 50% or task-specific sparsities, rewinding versus standard pruning, and multi-task pruning.
  • Transfer performance matches or outperforms same-task performance even when MLM training examples equal the downstream training-set size.
Loading 2007.12223v2…