[Paper Review] The Lottery Ticket Hypothesis for Pre-trained BERT Networks

Tianlong Chen, Jonathan Frankle+5 more|arXiv (Cornell University)|Jul 23, 2020

Topic Modeling47 references36 citations

TL;DR

The paper shows that sparse, trainable subnetworks exist in pre-trained BERT at non-trivial sparsities (40–90%), some subnetworks found at pre-training initialization can train to full accuracy on downstream tasks, and a universal subnetwork from MLM transfers across tasks.

ABSTRACT

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.

Motivation & Objective

• Assess whether matching subnetworks ( lottery tickets ) exist in pre-trained BERT models for downstream NLP tasks.
• Determine whether these subnetworks can be found at pre-trained initialization and/or during early training.
• Evaluate transferability of found subnetworks across different downstream tasks.
• Identify whether universal subnetworks exist that transfer to many tasks without accuracy loss.

Proposed method

• Use iterative magnitude pruning (IMP) to identify sparse subnetworks from a pre-trained BERT base initialized at theta_0.
• Prune weights by magnitude globally to a target sparsity and rewind to a specified training step i (including i=0).
• Evaluate subnetworks by training them on downstream tasks and comparing to full unpruned BERT performance.
• Compare IMP subnetworks to randomly pruned and randomly reinitialized subnetworks to establish the importance of pruning mask and initialization.
• Test both transferability across tasks and universality by pruning MLM-derived subnetworks and evaluating on multiple tasks.

Experimental results

Research questions

• RQ1 Do matching subnetworks exist in BERT at non-trivial sparsities when initialized at pre-trained theta_0?
• RQ2 Do these subnetworks transfer to other downstream tasks, and are there universal subnetworks that transfer broadly?
• RQ3 Does rewinding to intermediate training states improve subnetwork performance or transferability?
• RQ4 How does transfer performance compare when subnetworks are derived from MLM pre-training versus other downstream tasks?

Key findings

• Matching subnetworks exist at 40%–90% sparsity across GLUE tasks and SQuAD when pruned from the pre-trained initialization.
• Subnetworks found on MLM transfer universally to other tasks; subnetworks from other tasks transfer only in a limited fashion.
• Rewinding to intermediate training steps does not notably improve performance and can even hurt some tasks.
• Subnetworks prune to 70% sparsity from MLM yield universal transfer across tasks studied.
• Subnetworks from MLM achieve best transfer performance, achieving near-same-task performance on target tasks in many cases.