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.602693	0.557460	0.710419	00:03
1	0.538421	0.569738	0.694858	00:02
2	0.515489	0.555442	0.728687	00:02
3	0.470217	0.560604	0.691475	00:03
4	0.444224	0.617595	0.652233	00:03

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

epoch	train_loss	valid_loss	accuracy	time
0	0.591790	0.597947	0.722598	00:04
1	0.539337	0.593441	0.705683	00:05
2	0.508178	0.581475	0.708390	00:04
3	0.476695	0.577352	0.718539	00:04
4	0.433314	0.499315	0.753721	00:04
5	0.555141	0.545766	0.713126	00:04
6	0.520976	0.575154	0.718539	00:04
7	0.488574	0.636114	0.631258	00:04
8	0.457624	0.484251	0.767253	00:04
9	0.417182	0.444555	0.788904	00:04
10	0.494244	0.641547	0.703654	00:04
11	0.438013	0.482871	0.760487	00:04
12	0.399567	0.506693	0.748985	00:04
13	0.375169	0.413867	0.803789	00:04
14	0.343471	0.522576	0.732747	00:04
15	0.407545	0.486564	0.761164	00:04
16	0.359361	0.480845	0.783491	00:04
17	0.321976	0.413148	0.805142	00:04
18	0.283537	0.439845	0.814614	00:04
19	0.252996	0.424548	0.835589	00:04

Sparsity at the end of epoch 0: 0.00%
Sparsity at the end of epoch 1: 0.00%
Sparsity at the end of epoch 2: 0.00%
Sparsity at the end of epoch 3: 0.00%
Sparsity at the end of epoch 4: 0.00%
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.00%
Sparsity at the end of epoch 16: 50.00%
Sparsity at the end of epoch 17: 50.00%
Sparsity at the end of epoch 18: 50.00%
Sparsity at the end of epoch 19: 50.00%
Final Sparsity: 50.00%

Sparsity Report:
--------------------------------------------------------------------------------
Layer                          Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
conv1                          Conv2d          9,408      4,704         50.00%
layer1.0.conv1                 Conv2d          36,864     18,432        50.00%
layer1.0.conv2                 Conv2d          36,864     18,432        50.00%
layer1.1.conv1                 Conv2d          36,864     18,432        50.00%
layer1.1.conv2                 Conv2d          36,864     18,432        50.00%
layer2.0.conv1                 Conv2d          73,728     36,864        50.00%
layer2.0.conv2                 Conv2d          147,456    73,728        50.00%
layer2.0.downsample.0          Conv2d          8,192      4,096         50.00%
layer2.1.conv1                 Conv2d          147,456    73,728        50.00%
layer2.1.conv2                 Conv2d          147,456    73,728        50.00%
layer3.0.conv1                 Conv2d          294,912    147,456       50.00%
layer3.0.conv2                 Conv2d          589,824    294,912       50.00%
layer3.0.downsample.0          Conv2d          32,768     16,384        50.00%
layer3.1.conv1                 Conv2d          589,824    294,912       50.00%
layer3.1.conv2                 Conv2d          589,824    294,912       50.00%
layer4.0.conv1                 Conv2d          1,179,648  589,824       50.00%
layer4.0.conv2                 Conv2d          2,359,296  1,179,647     50.00%
layer4.0.downsample.0          Conv2d          131,072    65,536        50.00%
layer4.1.conv1                 Conv2d          2,359,296  1,179,648     50.00%
layer4.1.conv2                 Conv2d          2,359,296  1,179,648     50.00%
--------------------------------------------------------------------------------
Overall                        all             11,166,912 5,583,455     50.00%

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%

epoch	train_loss	valid_loss	accuracy	time
0	0.585928	0.599651	0.695535	00:03
1	0.567199	0.571735	0.725304	00:03
2	0.520258	0.527225	0.738160	00:04
3	0.487749	0.526812	0.746279	00:03
4	0.455229	0.509876	0.732747	00:03
5	0.502195	0.805385	0.704330	00:03
6	0.465073	0.489304	0.769959	00:03
7	0.424250	1.209978	0.462111	00:03
8	0.393979	0.475942	0.756428	00:03
9	0.361146	0.519064	0.751015	00:03
10	0.433051	0.483362	0.756428	00:03
11	0.400122	0.484240	0.786198	00:03
12	0.360677	0.424378	0.807848	00:03
13	0.323188	0.448461	0.805819	00:03
14	0.291294	0.436022	0.820027	00:03
15	0.360917	0.482217	0.803112	00:03
16	0.320696	0.397914	0.816644	00:03
17	0.269400	0.516504	0.771989	00:03
18	0.256572	0.471974	0.799053	00:02
19	0.222436	0.480946	0.779432	00:03

Sparsity at the end of epoch 0: 0.00%
Saving Weights at epoch 1
Sparsity at the end of epoch 1: 0.00%
Sparsity at the end of epoch 2: 0.00%
Sparsity at the end of epoch 3: 0.00%
Sparsity at the end of epoch 4: 0.00%
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.00%
Sparsity at the end of epoch 16: 50.00%
Sparsity at the end of epoch 17: 50.00%
Sparsity at the end of epoch 18: 50.00%
Sparsity at the end of epoch 19: 50.00%
Final Sparsity: 50.00%

Sparsity Report:
--------------------------------------------------------------------------------
Layer                          Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
conv1                          Conv2d          9,408      4,704         50.00%
layer1.0.conv1                 Conv2d          36,864     18,432        50.00%
layer1.0.conv2                 Conv2d          36,864     18,432        50.00%
layer1.1.conv1                 Conv2d          36,864     18,432        50.00%
layer1.1.conv2                 Conv2d          36,864     18,432        50.00%
layer2.0.conv1                 Conv2d          73,728     36,864        50.00%
layer2.0.conv2                 Conv2d          147,456    73,728        50.00%
layer2.0.downsample.0          Conv2d          8,192      4,096         50.00%
layer2.1.conv1                 Conv2d          147,456    73,728        50.00%
layer2.1.conv2                 Conv2d          147,456    73,728        50.00%
layer3.0.conv1                 Conv2d          294,912    147,456       50.00%
layer3.0.conv2                 Conv2d          589,824    294,912       50.00%
layer3.0.downsample.0          Conv2d          32,768     16,384        50.00%
layer3.1.conv1                 Conv2d          589,824    294,912       50.00%
layer3.1.conv2                 Conv2d          589,824    294,912       50.00%
layer4.0.conv1                 Conv2d          1,179,648  589,824       50.00%
layer4.0.conv2                 Conv2d          2,359,296  1,179,648     50.00%
layer4.0.downsample.0          Conv2d          131,072    65,536        50.00%
layer4.1.conv1                 Conv2d          2,359,296  1,179,648     50.00%
layer4.1.conv2                 Conv2d          2,359,296  1,179,648     50.00%
--------------------------------------------------------------------------------
Overall                        all             11,166,912 5,583,456     50.00%

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(10, 1e-3, cbs=sp_cb)

Pruning of weight until a sparsity of 50%
Saving Weights at epoch 0

epoch	train_loss	valid_loss	accuracy	time
0	0.594171	0.610432	0.649526	00:04
1	0.561133	0.920322	0.684709	00:03
2	0.574408	0.606864	0.659675	00:03
3	0.567595	0.552262	0.714479	00:02
4	0.516829	0.581788	0.694858	00:03
5	0.553904	0.525427	0.740866	00:02
6	0.490872	0.523433	0.715156	00:02
7	0.494422	0.939405	0.379567	00:02
8	0.438296	0.471724	0.783491	00:02
9	0.402010	0.540082	0.765900	00:02

Sparsity at the end of epoch 0: 0.00%
Sparsity at the end of epoch 1: 0.00%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 2: 16.67%
Sparsity at the end of epoch 3: 16.67%
Sparsity at the end of epoch 4: 16.67%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 5: 33.33%
Sparsity at the end of epoch 6: 33.33%
Resetting Weights to their epoch 0 values
Sparsity at the end of epoch 7: 50.00%
Sparsity at the end of epoch 8: 50.00%
Sparsity at the end of epoch 9: 50.00%
Final Sparsity: 50.00%

Sparsity Report:
--------------------------------------------------------------------------------
Layer                          Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
conv1                          Conv2d          9,408      4,704         50.00%
layer1.0.conv1                 Conv2d          36,864     18,432        50.00%
layer1.0.conv2                 Conv2d          36,864     18,432        50.00%
layer1.1.conv1                 Conv2d          36,864     18,432        50.00%
layer1.1.conv2                 Conv2d          36,864     18,432        50.00%
layer2.0.conv1                 Conv2d          73,728     36,864        50.00%
layer2.0.conv2                 Conv2d          147,456    73,728        50.00%
layer2.0.downsample.0          Conv2d          8,192      4,096         50.00%
layer2.1.conv1                 Conv2d          147,456    73,728        50.00%
layer2.1.conv2                 Conv2d          147,456    73,728        50.00%
layer3.0.conv1                 Conv2d          294,912    147,456       50.00%
layer3.0.conv2                 Conv2d          589,824    294,912       50.00%
layer3.0.downsample.0          Conv2d          32,768     16,384        50.00%
layer3.1.conv1                 Conv2d          589,824    294,912       50.00%
layer3.1.conv2                 Conv2d          589,824    294,911       50.00%
layer4.0.conv1                 Conv2d          1,179,648  589,824       50.00%
layer4.0.conv2                 Conv2d          2,359,296  1,179,648     50.00%
layer4.0.downsample.0          Conv2d          131,072    65,536        50.00%
layer4.1.conv1                 Conv2d          2,359,296  1,179,648     50.00%
layer4.1.conv2                 Conv2d          2,359,296  1,179,648     50.00%
--------------------------------------------------------------------------------
Overall                        all             11,166,912 5,583,455     50.00%

learn.validate()

[0.6421075463294983, 0.6617050170898438]