Hey I am a little bit confused about the scale.

Inside SineSPE() you deal with the scale (both d^0.25 and num_realizations^0.25) On the other hand when you show the application in pytorch, after applying the filter you divide by sqrt(num_realizations) again, why is that? https://github.com/aliutkus/spe/blob/main/src/pytorch/examples/test_spe.ipynb

optimizes memory usage and speed by sharing the SPE along all layers.

This is done in the following way:

• not redrawing noise each time, so as to share qbar and kbar on all layers. The strategy picked is to keep qbar and kbar untouched as long as their shapes are ok. They must be manually resetted if required.
• remove the use of einsum. It's indeed much nicer, but for some mysterious reason, it apparently did not allow to save RAM when reusing qbar and kbar.

the notebook tries to apply the SPE many times in a row, to simulate many layers.

:warning: note that my_spe.reset() must now be called explicitly each time a new spe must be computed (typically at each batch during training).

Hi there,

I tried to use ConvSPE with images (ndim=2) but spe.py failed. The following snippet replicates the error with the current spe implementation. Notice, I am using a 2D (50x50) input. This PR fixes this error. If this fix is wrong, please suggest a better solution. Thanks

import spe
import torch

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

keys_dim = 64
num_heads = 1
num_realizations = 64
kernel_size = 20
n = 50
batchsize = 8

poscoder = spe.ConvSPE(ndim=2,in_features=keys_dim,kernel_size=kernel_size,num_heads=num_heads,num_realizations=num_realizations)
filter = spe.SPEFilter(gated=True, code_shape=poscoder.code_shape).to(device)

poscoder.to(device=device)

q = torch.rand(batchsize, n, n, num_heads, keys_dim, device=device, requires_grad=True)
k = torch.rand(batchsize, n, n, num_heads, keys_dim, device=device, requires_grad=True)


poscode = poscoder(q.shape[:3])
q, k = filter(q, k, poscode)

I was trying to reproduce your experiments and when I run git submodule init && git submodule update after cloning, but I faced several errors when the [email protected]:cifkao/fast-transformers.git was cloned: [email protected]: Permission denied (publickey) Then I decided to change url [email protected]:cifkao/fast-transformers.git to https://github.com/cifkao/fast-transformers.git and it seems to fix the issue

instead of deciding how many features should be left unchanged or should be modulated by SPE, let's train that !

The idea is to modify the model so that we have $P_d(m,n) \leftarrow \lambda P_d(m-n) + (1-\lambda) 1$

So that depending on this $\lambda$, we turn this dimension into a normal (unchanged) one, or use SPE for it.

In practice, due to the cross-terms thing, this is implementing as drawing noise again, but this time this noise is the same for all time lags (it's hence small). This way, the covariance matrix will be full of one for QdKd, but still zero for QdKd' (d\neq d').

I didn't test heavily, please have a look

this implements a new two step behaviour of first generating the pos encoding, before applying them on demand with a Filter module, that allows gating

this enables the layer-wise gating parameters

For the jax implementation, on line 210 of spe.py, should the axis summed over be -1 instead of -2? When using -2, the size of the last output dimension is num_realizations, rather than the query/key dimension:

return (spe[:, :keys.shape[1]] * keys[..., None]).sum(axis=-1)

Hello,

I implemented the algorithm in the vision transformer architecture the following way:

#inside __init__()
self.spe = SineSPE(num_heads=head_cnt,in_features=in_dim,num_sines=5,num_realizations=64)
self.filter = SPEFilter(gated=False,code_shape=self.spe.code_shape)

#inside forward()
q,k=self.filter(q,k,self.spe(q.shape[:2]))
qk,kp = performer(...)
out=lin_attention(...)

The model I am using has 4 layers 6 heads and embedding dimension 384, patch_size=4.

Training 100 epochs with CIFAR100 converges to 42.3% and without SPE 45.3%. Although this can be expected, with SPE the training time is around 6x longer, is that normal? Performers + ViT takes 39 minutes Perfomers + ViT + SPE takes around 4 hours For both I am using 2 Titan XP GPUs.

This is very problematic to me because I was considering scaling up those experiments with imagenet.

I would also like to know how can I implement the indexing T=N^2 for images (where did you do it in the lra benchmark?), according to section 2 of the paper.

Many thanks!