Source-linked AI summary
Playing the lottery with rewards and multiple languages: lottery tickets in RL and NLP
Haonan Yu, Sergey Edunov, Yuandong Tian, Ari S. Morcos
TL;DR
The paper investigates whether lottery-ticket initializations extend beyond supervised natural-image learning to NLP and reinforcement learning. Using sparse subnetworks from LSTM, Transformer, classic-control, and Atari models, it finds that winning tickets generally outperform matched random initializations, including a Transformer retaining near-equivalent translation performance at one-third the size.
Problem
It was unclear whether lottery-ticket behavior generalized beyond supervised image classification to other domains, architectures, and learning regimes.
Method
The authors evaluate winning versus random sparse-subnetwork initializations in LSTM language modeling, Transformer translation, classic-control RL, and Atari RL.
Results
Winning tickets generally outperform random tickets across NLP and RL, with a Transformer Big reaching 28.9 BLEU at 67% pruning versus 29.2 BLEU unpruned.
Takeaways & Limitations
The findings suggest that lottery-ticket behavior is not restricted to supervised image classification but is a broader feature of deep neural network training.
Abstract
from arXiv · showhide
The lottery ticket hypothesis proposes that over-parameterization of deep neural networks (DNNs) aids training by increasing the probability of a "lucky" sub-network initialization being present rather than by helping the optimization process (Frankle & Carbin, 2019). Intriguingly, this phenomenon suggests that initialization strategies for DNNs can be improved substantially, but the lottery ticket hypothesis has only previously been tested in the context of supervised learning for natural image tasks. Here, we evaluate whether "winning ticket" initializations exist in two different domains: natural language processing (NLP) and reinforcement learning (RL).For NLP, we examined both recurrent LSTM models and large-scale Transformer models (Vaswani et al., 2017). For RL, we analyzed a number of discrete-action space tasks, including both classic control and pixel control. Consistent with workin supervised image classification, we confirm that winning ticket initializations generally outperform parameter-matched random initializations, even at extreme pruning rates for both NLP and RL. Notably, we are able to find winning ticket initializations for Transformers which enable models one-third the size to achieve nearly equivalent performance. Together, these results suggest that the lottery ticket hypothesis is not restricted to supervised learning of natural images, but rather represents a broader phenomenon in DNNs.
1 INTRODUCTION
The paper asks whether lottery-ticket behavior extends beyond supervised image classification to NLP and reinforcement learning. Across these domains, it finds winning-ticket initializations, including a Transformer that retains near-baseline translation performance with one-third of the parameters.
- Motivation: Prior evidence left unclear whether lottery-ticket behavior was intrinsic to DNNs or dependent on supervised learning, architecture, task, dataset bias, or optimization artifacts.Large learning rates could damage the effect, while larger models and datasets sometimes required warmup or late rewinding.
- Research question: The study tests whether lottery tickets generalize to NLP and RL, where recurrent dynamics, gating, shifting data distributions, and reward signals alter optimization.The evaluation includes recurrent LSTMs, Transformers, classic control, and Atari games.
- Findings: Winning-ticket initializations outperform random tickets in LSTM language modeling, Transformer machine translation, classic control, and many Atari games, though Atari results show high variance.The comparison focuses on whether initialization affects sparse-subnetwork training rather than requiring subnetworks to match the full network.
- Conclusion: The results suggest that lottery-ticket behavior is a general DNN phenomenon rather than a property restricted to supervised natural-image learning.The conclusion spans both NLP and RL evidence.
2 RELATED WORK
The related work establishes lottery tickets as sparse subnetworks found through pruning, while prior studies questioned their generality and reported random subnetworks matching winning tickets in some settings.
- Lottery-ticket hypothesis: The lottery ticket hypothesis proposes that small, sparse subnetworks can achieve high performance when trained independently within over-parameterized networks.The cited work reports subnetworks with as few as 0.1% of the original parameters.
- Lottery-ticket hypothesis: Late rewinding was introduced as a refinement that improves winning-ticket performance for large models and datasets.It resets winning tickets to weights from an early training point rather than their initial values.
- Challenges: Liu et al. found that structured-pruning random subnetworks could match winning-ticket performance, challenging the hypothesis.This result raised questions about whether the initialization advantage is general across pruning settings.
- Challenges: Gale et al. reported similar performance between random and winning-ticket subnetworks in ResNets and Transformers, without iterative pruning or late rewinding.The present paper uses these methodological differences to revisit the comparison.
- Pruning methods: This work uses magnitude pruning, removing the smallest-magnitude weights first, while prior pruning research largely evaluated supervised image-classification settings.Other cited approaches include greedy weight ranking, channel redundancy ranking, and variational pruning.
3 APPROACH
The approach generates tickets through one-shot or iterative magnitude pruning, optionally uses late rewinding, and evaluates them across NLP and RL tasks using task-specific performance measures.
- Ticket generation: One-shot pruning trains the full network, removes low-magnitude parameters, resets the remaining weights, and retrains the sparse model.The method evaluates whether the reset sparse subnetwork performs better than a random initialization.
- Ticket generation: Iterative pruning alternates training, pruning a small fraction of remaining weights, re-initialization to a winning ticket, and retraining.Subsequent pruning uses the trained weights of the winning ticket from the preceding iteration.
- Pruning settings: All experiments generally use global iterative pruning at p = 0.2, while RL runs use 20 pruning iterations and pool weights across layers.Global pruning orders weights by magnitude across layers, so layerwise sparsity can differ.
- Rewinding: Late rewinding resets winning-ticket weights to values from an early training point instead of their initialization values.The method tests late rewinding because earlier work found it beneficial for large-scale models and datasets.
- NLP experiments: The NLP evaluation covers two-layer LSTMs for Wikitext-2 language modeling and Transformer Base and Big models for WMT’14 English-to-German translation.The LSTM experiments also compare settings with and without late rewinding and with one-shot pruning.
- RL experiments: The RL evaluation covers three classic-control environments and nine Atari pixel-control games with discrete actions.Classic control tests fully connected networks, while pixel control uses Atari environments.
- RL evaluation: RL ticket performance is measured by averaged episodic reward over the final L game episodes, with three pruning-process repetitions reported as mean ± standard deviation.Winning and random ticket reward curves are compared as pruning increases.
4 RESULTS
Winning tickets outperform random tickets across LSTM language modeling, Transformer translation, classic control, and Atari tasks, although their benefits vary by task. Iterative pruning is generally important, while late rewinding has smaller or setting-dependent effects.
- Language modeling with LSTMs: Winning tickets with late rewinding outperformed random tickets at every pruning level in LSTM language modeling, with up to 90% of parameters removed without noticeably increasing log perplexity.The task used a two-layer LSTM trained on Wikitext-2.
- Language modeling with LSTMs: Iterative pruning was essential for LSTM performance, whereas removing late rewinding caused only slight damage concentrated at intermediate pruning levels.With 80% of parameters pruned, performance fell below that of random tickets when iterative pruning was removed.
- Machine translation with Transformers: Winning tickets outperformed random tickets in Transformer machine translation, and iterative pruning had a larger impact than late rewinding.The Transformer Base baseline achieved 27.6 BLEU on Newstest’14; pruning was performed globally after 1000 updates, with and without late rewinding.
- Machine translation with Transformers: Transformer winning-ticket performance remained robust when only transformer-layer weights were pruned, while random-ticket performance declined much faster.In this setting, winning-ticket BLEU decreased roughly linearly, whereas embedding weights appeared more safely prunable for random initializations.
- Machine translation with Transformers: 28.9 BLEU at 67% pruning versus 29.2 BLEU unpruned showed that a Transformer Big winning ticket could retain 99% of unpruned performance with one-third of the weights.The model was trained from scratch using iterative pruning and late rewinding.
- Reinforcement learning: Winning tickets outperformed random tickets in classic control, while Atari results varied substantially: some games benefited strongly, others slightly or not at all.Kangaroo winning-ticket initializations enabled some learning despite unpruned models failing to reach high rewards, but remained far from optimal performance.
5 CONCLUSION
The study finds winning-ticket initializations across NLP and RL, extending evidence beyond supervised image classification. In RL, their advantage is substantial on classic control and present for many, but not all, Atari games.
- Winning-ticket initializations beat random tickets in recurrent LSTMs, Transformers, classic-control tasks, and many Atari games.
- High-performing Transformer Big models remain possible at pruning rates of at least 67%.
- The results suggest that the lottery-ticket phenomenon is not restricted to supervised image classification but is a general feature of deep neural network training.
- Figure 7 examines how iterative pruning affects several pixel-control games, with variability shown across runs.
A APPENDIX
Table A1 summarizes the games and network architectures used in the reinforcement-learning experiments. It also defines the convolution notation and activation convention used in the table.
- Table A1 summarizes the games and corresponding networks used in the RL experiments.
- Conv(w, x,y,z) denotes convolution filter size, channel number, stride, and padding.
- All layer activations in the listed networks are ReLUs.