Compact Transformers

[dmm]

For those of us (like myself) who'd like to experiment with changing Transformer architecture on a home personal computer.

Links are in the comments.

Comments

[dmm]

Escaping the Big Data Paradigm with Compact Transformers: https://arxiv.org/abs/2104.05704

PyTorch implementation: https://github.com/SHI-Labs/Compact-Transformers

[dmm]

https://www.cs.toronto.edu/~kriz/cifar.html - for CIFAR, but you don't need it, just tell the script a directory, and it will download and unpack it.

[dmm]

I just reproduced it on the latest PyTorch (a later version of PyTorch than their README suggests; I used Windows 10, but it should work elsewhere too).

The result was a bit lower than in their table:

[Epoch 200] Top-1 88.34 Time: 64.41
Script finished in 64.41 minutes, best top-1: 88.34, final top-1: 88.34

They had more luck and got 89.17 (the training is a bit jittery, but generally OK).

[dmm]

So, if I want to change the Transformer attention per my request-for-plot https://github.com/anhinga/JuliaCon2021-poster/tree/main/RFP-draft

The file involved is

https://github.com/SHI-Labs/Compact-Transformers/blob/main/src/utils/transformers.py

and it has lines 207 and 313 which both say

            x = torch.matmul(F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), x).squeeze(-2)

and these lines seem to be the candidates for the change described in my RFP

[dmm]

So, here is my first attempt to test my conjecture that softmax cross-normalization of values is fruitful.

With the modified line 207

            x = torch.matmul(F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), F.softmax(x, dim=0)).squeeze(-2)

the first run is worse than the baseline:

[Epoch 200] Top-1 84.56 Time: 58.91
Script finished in 58.91 minutes, best top-1: 84.61, final top-1: 84.56

versus baseline

[Epoch 200] Top-1 88.34 Time: 64.41
Script finished in 64.41 minutes, best top-1: 88.34, final top-1: 88.34

Now it's time to ponder this initial negative result.

[dmm]

I am pretty sure that this has bugs related to my lack of understanding of the source code in

https://github.com/SHI-Labs/Compact-Transformers/blob/main/src/utils/transformers.py

Time to understand this better.

***

Yes, these are 3D tensors and "batched matrix multiplication". So there is no reason for my formula to be applicable; I need to change it.

[dmm]

So, the fix is "F.softmax(x, dim=1)" instead of "F.softmax(x, dim=0)", and it works much better, but still trails the unmodified version somewhat...

[Epoch 200] Top-1 86.99 Time: 60.12
Script finished in 60.12 minutes, best top-1: 87.10, final top-1: 86.99

versus baseline

[Epoch 200] Top-1 88.34 Time: 64.41
Script finished in 64.41 minutes, best top-1: 88.34, final top-1: 88.34

[dmm]

Let's now scale it to bring the max magnitude closer (so that we don't have to figure out the best learning rates again):

Use "10*F.softmax(x, dim=1)" instead of "F.softmax(x, dim=1)". I am not sure this changes much, but we'll see...

(Perhaps, one does need to play with learning rate schedule; at the beginning it looked like this is better, but then it started to look like this "10*" is counter-productive.)

[Epoch 200] Top-1 85.64 Time: 60.00
Script finished in 60.00 minutes, best top-1: 85.77, final top-1: 85.64

[dmm]

Trying a mixed of original Transformer and the new formula, "F.softmax(x, dim=1)+x":

[Epoch 200] Top-1 88.47 Time: 70.49
Script finished in 70.49 minutes, best top-1: 88.49, final top-1: 88.47

We got some very mild improvement

[dmm]

I decided to take "signed normalization" considerations into account ( https://github.com/anhinga/JuliaCon2021-poster/tree/main/signed-normalization ) and try "x-F.softmax(x, dim=1)" instead.

Got a much less stable training curve (I think) getting ahead and falling behind the previous one all the time, but it ended somewhat behind:


[Epoch 200] Top-1 88.12 Time: 58.26
Script finished in 58.26 minutes, best top-1: 88.20, final top-1: 88.12

I am going to rerun (I'd like to check whether the current setup is deterministic, and if not, what are the variations).

[dmm]

Not deterministic, that's for sure... So one run might not mean all that much...

[dmm]

Now this is way better:

[Epoch 200] Top-1 88.67 Time: 58.07
Script finished in 58.07 minutes, best top-1: 88.67, final top-1: 88.67

But comparing in the presence of this much jitter is a nightmare, unless one configuration is overwhelmingly better.

One might need to do tons of reruns (in parallel, perhaps) to get statistics...