Lottery Ticket Hypothesis

How to find winning tickets with fastai

The Lottery Ticket Hypothesis

The Lottery Ticket Hypothesis is a really intriguing discovery made in 2019 by Frankle & Carbin. It states that:

A randomly-initialized, dense neural network contains a subnetwork that is initialised such that — when trained in isolation — it can match the test accuracy of the original network after training for at most the same number of iterations.

Meaning that, once we find that subnetwork. Every other parameter in the network becomes useless.

The way authors propose to find those subnetwork is as follows:

Lottery Ticket Hypothesis

Initialize the neural network
Train it to convergence
Prune the smallest magnitude weights by creating a mask \(m\)
Reinitialize the weights to their original value; i.e at iteration \(0\).
Repeat from step 2 until reaching the desired level of sparsity.

from fasterai.sparse.all import *

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64), device=device)

What we are trying to prove is that: in a neural network A, there exists a subnetwork B able to get an accuracy \(a_B > a_A\), in a training time \(t_B < t_A\).

Let’s get the baseline for network A:

learn = Learner(dls, resnet18(num_classes=2), metrics=accuracy)

Let’s save original weights

initial_weights = deepcopy(learn.model.state_dict())

learn.fit(5, 1e-3)

epoch	train_loss	valid_loss	accuracy	time
0	0.595053	0.543628	0.717185	00:07
1	0.546125	0.531834	0.728687	00:07
2	0.521566	0.521680	0.746279	00:07
3	0.478598	0.478660	0.764547	00:07
4	0.421515	0.719835	0.633288	00:07

We now have our accuracy \(a_A\) of \(79\%\) and our training time \(t_A\) of \(5\) epochs

To find the lottery ticket, we will perform iterative pruning but, at each pruning step we will re-initialize the remaining weights to their original values (i.e. before training).

We will restart from the same initialization to be sure to not get lucky.

learn = Learner(dls, resnet18(num_classes=2), metrics=accuracy)
learn.model.load_state_dict(initial_weights)

<All keys matched successfully>

We can pass the parameters lth=True to make the weights of the network reset to their original value after each pruning step, i.e. step 4) of the LTH. To empirically validate the LTH, we need to retrain the found “lottery ticket” after the pruning phase. Lottery tickets are usually found following an iterative pruning schedule. We set the start_epoch parameter to \(5\) to begin the pruning process after \(5\) epochs.

schedule = Schedule(sched_iterative, start_pct=0.25)

sp_cb = SparsifyCallback(50, 'weight', 'local', large_final, schedule, lth=True)

As our iterative schedule makes \(3\) pruning steps by default, it means that we have to train our network for start_epoch + \(3*t_B\), so \(20\) epochs in order to get our LTH. After each step, the remaining weights will be reinitialized to their original value

learn.fit(20, 1e-3, cbs=sp_cb)

Pruning of weight until a sparsity of [50]%
Saving Weights at epoch 0
Sparsity at the end of epoch 0: [0.0]%
Sparsity at the end of epoch 1: [0.0]%
Sparsity at the end of epoch 2: [0.0]%
Sparsity at the end of epoch 3: [0.0]%
Sparsity at the end of epoch 4: [0.0]%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 5: [16.67]%
Sparsity at the end of epoch 6: [16.67]%
Sparsity at the end of epoch 7: [16.67]%
Sparsity at the end of epoch 8: [16.67]%
Sparsity at the end of epoch 9: [16.67]%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 10: [33.33]%
Sparsity at the end of epoch 11: [33.33]%
Sparsity at the end of epoch 12: [33.33]%
Sparsity at the end of epoch 13: [33.33]%
Sparsity at the end of epoch 14: [33.33]%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 15: [50.0]%
Sparsity at the end of epoch 16: [50.0]%
Sparsity at the end of epoch 17: [50.0]%
Sparsity at the end of epoch 18: [50.0]%
Sparsity at the end of epoch 19: [50.0]%
Final Sparsity: [50.0]%
Sparsity in Conv2d 1: 50.00%
Sparsity in Conv2d 7: 50.00%
Sparsity in Conv2d 10: 50.00%
Sparsity in Conv2d 13: 50.00%
Sparsity in Conv2d 16: 50.00%
Sparsity in Conv2d 20: 50.00%
Sparsity in Conv2d 23: 50.00%
Sparsity in Conv2d 26: 50.00%
Sparsity in Conv2d 29: 50.00%
Sparsity in Conv2d 32: 50.00%
Sparsity in Conv2d 36: 50.00%
Sparsity in Conv2d 39: 50.00%
Sparsity in Conv2d 42: 50.00%
Sparsity in Conv2d 45: 50.00%
Sparsity in Conv2d 48: 50.00%
Sparsity in Conv2d 52: 50.00%
Sparsity in Conv2d 55: 50.00%
Sparsity in Conv2d 58: 50.00%
Sparsity in Conv2d 61: 50.00%
Sparsity in Conv2d 64: 50.00%

epoch	train_loss	valid_loss	accuracy	time
0	0.587271	0.612850	0.668471	00:07
1	0.547513	0.573875	0.701624	00:07
2	0.501447	0.514319	0.738160	00:07
3	0.467517	0.518835	0.740189	00:07
4	0.431308	0.621117	0.654263	00:07
5	0.561822	0.546851	0.728011	00:07
6	0.514754	0.570669	0.765223	00:07
7	0.469772	0.576346	0.675913	00:07
8	0.427512	0.454539	0.799053	00:07
9	0.385058	0.537472	0.728011	00:07
10	0.482251	0.442526	0.778078	00:07
11	0.430115	0.510421	0.777402	00:07
12	0.383115	0.424600	0.811908	00:07
13	0.341061	0.663934	0.681326	00:07
14	0.322339	0.383250	0.830853	00:07
15	0.402276	0.441383	0.799053	00:07
16	0.348916	0.384334	0.814614	00:07
17	0.311075	0.409829	0.817321	00:07
18	0.280943	0.405584	0.824763	00:07
19	0.244700	0.385061	0.827470	00:07

We indeed have a network B, whose accuracy \(a_B > a_A\) in the same training time.

Lottery Ticket Hypothesis with Rewinding

In some case, LTH fails for deeper networks, author then propose a solution, which is to rewind the weights to a more advanced iteration instead of the initialization value.

learn = Learner(dls, resnet18(num_classes=2), metrics=accuracy)
learn.model.load_state_dict(initial_weights)

<All keys matched successfully>

This can be done in fasterai by passing the rewind_epoch parameter, that will save the weights at that epoch, then resetting the weights accordingly.

sp_cb = SparsifyCallback(50, 'weight', 'local', large_final, schedule, lth=True, rewind_epoch=1)

learn.fit(20, 1e-3, cbs=sp_cb)

Pruning of weight until a sparsity of [50]%
Sparsity at the end of epoch 0: [0.0]%
Saving Weights at epoch 1
Sparsity at the end of epoch 1: [0.0]%
Sparsity at the end of epoch 2: [0.0]%
Sparsity at the end of epoch 3: [0.0]%
Sparsity at the end of epoch 4: [0.0]%
Resetting Weights to their epoch 1 values
Sparsity at the end of epoch 5: [16.67]%
Sparsity at the end of epoch 6: [16.67]%
Sparsity at the end of epoch 7: [16.67]%
Sparsity at the end of epoch 8: [16.67]%
Sparsity at the end of epoch 9: [16.67]%
Resetting Weights to their epoch 1 values
Sparsity at the end of epoch 10: [33.33]%
Sparsity at the end of epoch 11: [33.33]%
Sparsity at the end of epoch 12: [33.33]%
Sparsity at the end of epoch 13: [33.33]%
Sparsity at the end of epoch 14: [33.33]%
Resetting Weights to their epoch 1 values
Sparsity at the end of epoch 15: [50.0]%
Sparsity at the end of epoch 16: [50.0]%
Sparsity at the end of epoch 17: [50.0]%
Sparsity at the end of epoch 18: [50.0]%
Sparsity at the end of epoch 19: [50.0]%
Final Sparsity: [50.0]%
Sparsity in Conv2d 1: 50.00%
Sparsity in Conv2d 7: 50.00%
Sparsity in Conv2d 10: 50.00%
Sparsity in Conv2d 13: 50.00%
Sparsity in Conv2d 16: 50.00%
Sparsity in Conv2d 20: 50.00%
Sparsity in Conv2d 23: 50.00%
Sparsity in Conv2d 26: 50.00%
Sparsity in Conv2d 29: 50.00%
Sparsity in Conv2d 32: 50.00%
Sparsity in Conv2d 36: 50.00%
Sparsity in Conv2d 39: 50.00%
Sparsity in Conv2d 42: 50.00%
Sparsity in Conv2d 45: 50.00%
Sparsity in Conv2d 48: 50.00%
Sparsity in Conv2d 52: 50.00%
Sparsity in Conv2d 55: 50.00%
Sparsity in Conv2d 58: 50.00%
Sparsity in Conv2d 61: 50.00%
Sparsity in Conv2d 64: 50.00%

epoch	train_loss	valid_loss	accuracy	time
0	0.594580	0.625727	0.685386	00:07
1	0.543424	0.667005	0.675237	00:07
2	0.509345	0.483691	0.748309	00:07
3	0.470524	0.478201	0.763870	00:07
4	0.451323	0.550612	0.725981	00:07
5	0.518282	0.483968	0.762517	00:07
6	0.478035	0.729290	0.709743	00:07
7	0.425959	0.479093	0.782815	00:07
8	0.393114	0.407376	0.815968	00:07
9	0.350740	0.755570	0.654939	00:07
10	0.442296	0.497535	0.759134	00:07
11	0.390260	0.617039	0.677267	00:07
12	0.346842	0.365099	0.836942	00:07
13	0.311359	0.431787	0.811908	00:07
14	0.276918	0.389922	0.835589	00:07
15	0.342124	0.495095	0.748985	00:07
16	0.308753	0.553223	0.734777	00:07
17	0.278255	0.491012	0.803112	00:07
18	0.254106	0.564554	0.782815	00:07
19	0.224321	0.341632	0.855210	00:07

Super-Masks

Researchers from Uber AI investigated the LTH and found the existence of what they call “Super-Masks”, i.e. masks that, applied on a untrained neural network, allows to reach better-than-random results.

learn = Learner(dls, resnet18(num_classes=2), metrics=accuracy)
learn.model.load_state_dict(initial_weights)

<All keys matched successfully>

To find supermasks, authors perform the LTH method then apply the mask on the original, untrained network. In fasterai, you can pass the parameter reset_end=True, which will reset the weights to their original value at the end of the training, but keeping the pruned weights (i.e. the mask) unchanged.

sp_cb = SparsifyCallback(50, 'weight', 'local', large_final, schedule, lth=True, reset_end=True)

learn.fit(20, 1e-3, cbs=sp_cb)

Pruning of weight until a sparsity of [50]%
Saving Weights at epoch 0
Sparsity at the end of epoch 0: [0.0]%
Sparsity at the end of epoch 1: [0.0]%
Sparsity at the end of epoch 2: [0.0]%
Sparsity at the end of epoch 3: [0.0]%
Sparsity at the end of epoch 4: [0.0]%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 5: [16.67]%
Sparsity at the end of epoch 6: [16.67]%
Sparsity at the end of epoch 7: [16.67]%
Sparsity at the end of epoch 8: [16.67]%
Sparsity at the end of epoch 9: [16.67]%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 10: [33.33]%
Sparsity at the end of epoch 11: [33.33]%
Sparsity at the end of epoch 12: [33.33]%
Sparsity at the end of epoch 13: [33.33]%
Sparsity at the end of epoch 14: [33.33]%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 15: [50.0]%
Sparsity at the end of epoch 16: [50.0]%
Sparsity at the end of epoch 17: [50.0]%
Sparsity at the end of epoch 18: [50.0]%
Sparsity at the end of epoch 19: [50.0]%
Final Sparsity: [50.0]%
Sparsity in Conv2d 1: 50.00%
Sparsity in Conv2d 7: 50.00%
Sparsity in Conv2d 10: 50.00%
Sparsity in Conv2d 13: 50.00%
Sparsity in Conv2d 16: 50.00%
Sparsity in Conv2d 20: 50.00%
Sparsity in Conv2d 23: 50.00%
Sparsity in Conv2d 26: 50.00%
Sparsity in Conv2d 29: 50.00%
Sparsity in Conv2d 32: 50.00%
Sparsity in Conv2d 36: 50.00%
Sparsity in Conv2d 39: 50.00%
Sparsity in Conv2d 42: 50.00%
Sparsity in Conv2d 45: 50.00%
Sparsity in Conv2d 48: 50.00%
Sparsity in Conv2d 52: 50.00%
Sparsity in Conv2d 55: 50.00%
Sparsity in Conv2d 58: 50.00%
Sparsity in Conv2d 61: 50.00%
Sparsity in Conv2d 64: 50.00%

epoch	train_loss	valid_loss	accuracy	time
0	0.611311	0.796523	0.674560	00:07
1	0.560894	0.608468	0.685386	00:07
2	0.529301	0.537641	0.736130	00:08
3	0.476619	0.542131	0.740866	00:07
4	0.441745	0.456982	0.776725	00:07
5	0.538183	0.558080	0.737483	00:07
6	0.497710	0.502065	0.749662	00:07
7	0.457478	0.542794	0.731394	00:07
8	0.425738	0.548435	0.725981	00:07
9	0.393300	0.562673	0.680650	00:07
10	0.475655	0.520430	0.756428	00:08
11	0.433569	0.582280	0.740866	00:08
12	0.408267	0.755556	0.547361	00:08
13	0.364271	0.421615	0.809202	00:07
14	0.324698	0.368889	0.832882	00:07
15	0.394710	0.604650	0.738836	00:07
16	0.350399	0.412777	0.814614	00:07
17	0.318693	0.443487	0.803789	00:07
18	0.280467	0.395299	0.826793	00:07
19	0.251164	0.380674	0.831529	00:07

learn.validate()

(#2) [0.8659006357192993,0.665764570236206]