Super Tickets in Pre-Trained Language Models: From Model Compression to Improving Generalization

TL;DR

This work extends the Lottery Ticket Hypothesis to extremely over-parameterized pre-trained language models, revealing a phase transition where pruning yields super tickets that can surpass full-model performance at certain compression levels. It introduces a structured pruning approach that targets attention heads and FFN units, coupled with rewinding to pre-trained weights to identify super tickets. The study demonstrates consistent single-task gains on GLUE, significant improvements for small data tasks, and a phase-transition phenomenon that motivates a tickets-sharing strategy for multi-task learning and domain adaptation. The proposed approach balances model bias and variance, enabling improved generalization and robust cross-task/domain transfer while remaining hardware-friendly due to structured sparsity.

Abstract

The Lottery Ticket Hypothesis suggests that an over-parametrized network consists of ``lottery tickets'', and training a certain collection of them (i.e., a subnetwork) can match the performance of the full model. In this paper, we study such a collection of tickets, which is referred to as ``winning tickets'', in extremely over-parametrized models, e.g., pre-trained language models. We observe that at certain compression ratios, the generalization performance of the winning tickets can not only match but also exceed that of the full model. In particular, we observe a phase transition phenomenon: As the compression ratio increases, generalization performance of the winning tickets first improves then deteriorates after a certain threshold. We refer to the tickets on the threshold as ``super tickets''. We further show that the phase transition is task and model dependent -- as the model size becomes larger and the training data set becomes smaller, the transition becomes more pronounced. Our experiments on the GLUE benchmark show that the super tickets improve single task fine-tuning by $0.9$ points on BERT-base and $1.0$ points on BERT-large, in terms of task-average score. We also demonstrate that adaptively sharing the super tickets across tasks benefits multi-task learning.

Figures (7)

• Figure 1: Illustrations of the phase transition phenomenon. Left: Generalization performance of the fine-tuned subnetworks (the same as Figure \ref{['exp:phase_transtion']} in Section \ref{['sec:st_exp']}). Middle and Right: An interpretation of bias-variance trade-off.
• Figure 2: Illustration of tickets sharing.
• Figure 3: Single task fine-tuning validation results in different GLUE tasks. Upper: Performance Gain. Lower: Percent of weight remaining.
• Figure 4: Single task fine-tuning evaluation results of the winning (blue), the random (orange), and the losing (green) tickets on the GLUE development set under various sparsity levels.
• Figure 5: Phase transition under different randomly sampled training subsets. Note that the settings are the same as Figure \ref{['exp:phase_transtion']} (bottom left), except the data size.