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.609106	0.557732	0.705007	00:02
1	0.549768	0.543677	0.712449	00:02
2	0.509367	0.609181	0.719892	00:02
3	0.465070	0.490555	0.757781	00:02
4	0.433047	0.477749	0.753721	00:02

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.583488	0.844444	0.684032	00:02
1	0.559688	0.547341	0.705683	00:02
2	0.512511	0.520819	0.734100	00:02
3	0.465449	0.725977	0.716509	00:02
4	0.440095	0.753626	0.602165	00:02
5	0.536660	0.757445	0.465494	00:02
6	0.488848	0.655531	0.701624	00:02
7	0.447604	0.480934	0.780108	00:02
8	0.407234	0.511213	0.757781	00:02
9	0.392847	0.464021	0.784844	00:02
10	0.483271	0.467506	0.769283	00:02
11	0.433441	0.867711	0.711773	00:02
12	0.378504	0.397733	0.831529	00:02
13	0.353482	0.420315	0.801759	00:02
14	0.335506	0.490255	0.746955	00:02
15	0.405062	0.506038	0.736130	00:02
16	0.347061	0.436380	0.804465	00:02
17	0.314622	0.390379	0.826793	00:02
18	0.277594	0.395576	0.834912	00:02
19	0.253820	0.432128	0.810555	00:02

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 Report:
--------------------------------------------------------------------------------
Layer                Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
Layer 0              Conv2d          9,408      4,704         50.00%
Layer 1              Conv2d          36,864     18,432        50.00%
Layer 2              Conv2d          36,864     18,432        50.00%
Layer 3              Conv2d          36,864     18,432        50.00%
Layer 4              Conv2d          36,864     18,432        50.00%
Layer 5              Conv2d          73,728     36,864        50.00%
Layer 6              Conv2d          147,456    73,728        50.00%
Layer 7              Conv2d          8,192      4,096         50.00%
Layer 8              Conv2d          147,456    73,728        50.00%
Layer 9              Conv2d          147,456    73,728        50.00%
Layer 10             Conv2d          294,912    147,456       50.00%
Layer 11             Conv2d          589,824    294,912       50.00%
Layer 12             Conv2d          32,768     16,384        50.00%
Layer 13             Conv2d          589,824    294,912       50.00%
Layer 14             Conv2d          589,824    294,912       50.00%
Layer 15             Conv2d          1,179,648  589,824       50.00%
Layer 16             Conv2d          2,359,296  1,179,648     50.00%
Layer 17             Conv2d          131,072    65,536        50.00%
Layer 18             Conv2d          2,359,296  1,179,648     50.00%
Layer 19             Conv2d          2,359,296  1,179,648     50.00%
--------------------------------------------------------------------------------
Overall              all             11,166,912 5,583,456     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.618825	0.631097	0.696211	00:02
1	0.564544	0.688483	0.696888	00:02
2	0.513954	0.614130	0.737483	00:02
3	0.482632	0.602386	0.706360	00:02
4	0.442137	0.530012	0.743572	00:02
5	0.512772	0.626846	0.631258	00:02
6	0.478054	0.526319	0.730041	00:02
7	0.429781	0.566376	0.675913	00:02
8	0.395110	0.496819	0.776049	00:02
9	0.359115	0.542083	0.709066	00:02
10	0.444215	0.624481	0.656292	00:02
11	0.400898	0.435920	0.805819	00:02
12	0.364652	0.504759	0.737483	00:02
13	0.323991	0.459131	0.790934	00:02
14	0.287191	0.413124	0.823410	00:02
15	0.371384	0.557266	0.690122	00:02
16	0.322098	0.480477	0.794993	00:02
17	0.269774	0.459394	0.791610	00:02
18	0.244662	0.452999	0.820704	00:02
19	0.224821	0.435102	0.828823	00:02

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 Report:
--------------------------------------------------------------------------------
Layer                Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
Layer 0              Conv2d          9,408      4,704         50.00%
Layer 1              Conv2d          36,864     18,432        50.00%
Layer 2              Conv2d          36,864     18,432        50.00%
Layer 3              Conv2d          36,864     18,432        50.00%
Layer 4              Conv2d          36,864     18,432        50.00%
Layer 5              Conv2d          73,728     36,864        50.00%
Layer 6              Conv2d          147,456    73,728        50.00%
Layer 7              Conv2d          8,192      4,096         50.00%
Layer 8              Conv2d          147,456    73,728        50.00%
Layer 9              Conv2d          147,456    73,728        50.00%
Layer 10             Conv2d          294,912    147,456       50.00%
Layer 11             Conv2d          589,824    294,912       50.00%
Layer 12             Conv2d          32,768     16,384        50.00%
Layer 13             Conv2d          589,824    294,912       50.00%
Layer 14             Conv2d          589,824    294,912       50.00%
Layer 15             Conv2d          1,179,648  589,824       50.00%
Layer 16             Conv2d          2,359,296  1,179,648     50.00%
Layer 17             Conv2d          131,072    65,536        50.00%
Layer 18             Conv2d          2,359,296  1,179,648     50.00%
Layer 19             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.573335	0.585538	0.715156	00:02
1	0.546942	0.641137	0.615697	00:02
2	0.577234	0.629645	0.684709	00:02
3	0.539981	0.616901	0.720568	00:02
4	0.511118	0.567626	0.684032	00:02
5	0.523180	0.540820	0.736807	00:02
6	0.489160	0.508543	0.762517	00:02
7	0.479372	0.565328	0.717862	00:02
8	0.448294	0.574637	0.666441	00:02
9	0.399796	0.437068	0.794317	00:02

Sparsity at the end of epoch 0: [0.0]%
Sparsity at the end of epoch 1: [0.0]%
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.0]%
Sparsity at the end of epoch 8: [50.0]%
Sparsity at the end of epoch 9: [50.0]%
Final Sparsity: [50.0]%

Sparsity Report:
--------------------------------------------------------------------------------
Layer                Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
Layer 0              Conv2d          9,408      4,704         50.00%
Layer 1              Conv2d          36,864     18,432        50.00%
Layer 2              Conv2d          36,864     18,432        50.00%
Layer 3              Conv2d          36,864     18,432        50.00%
Layer 4              Conv2d          36,864     18,432        50.00%
Layer 5              Conv2d          73,728     36,864        50.00%
Layer 6              Conv2d          147,456    73,728        50.00%
Layer 7              Conv2d          8,192      4,096         50.00%
Layer 8              Conv2d          147,456    73,728        50.00%
Layer 9              Conv2d          147,456    73,728        50.00%
Layer 10             Conv2d          294,912    147,456       50.00%
Layer 11             Conv2d          589,824    294,912       50.00%
Layer 12             Conv2d          32,768     16,384        50.00%
Layer 13             Conv2d          589,824    294,912       50.00%
Layer 14             Conv2d          589,824    294,912       50.00%
Layer 15             Conv2d          1,179,648  589,824       50.00%
Layer 16             Conv2d          2,359,296  1,179,648     50.00%
Layer 17             Conv2d          131,072    65,536        50.00%
Layer 18             Conv2d          2,359,296  1,179,648     50.00%
Layer 19             Conv2d          2,359,296  1,179,648     50.00%
--------------------------------------------------------------------------------
Overall              all             11,166,912 5,583,456     50.00%

learn.validate()

[0.6600900292396545, 0.675913393497467]