sharenoise

[aliutkus]

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.

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

[cifkao]

Maybe I would prefer to require calling reset manually with the (maximum) query length (shape) instead of calling it automatically if the shape doesn't match. Then in forward we can just check that we have at least the required length and truncate it if needed. That would allow to store the PEs for multiple iterations in case we find the redrawing is slow or want deterministic behavior.

[aliutkus]

Maybe I would prefer to require calling reset manually with the (maximum) query length (shape) instead of calling it automatically if the shape doesn't match. Then in forward we can just check that we have at least the required length and truncate it if needed. That would allow to store the PEs for multiple iterations in case we find the redrawing is slow or want deterministic behavior.

Yes, reset should be called manually, at the start of the forward of the hosting module.
but your solution does not work: we really do need to call it again at each forward pass, during training stage, because if we don't do it, we don't update the parameters of the SPE

[aliutkus]

Maybe I would prefer to require calling reset manually with the (maximum) query length (shape) instead of calling it automatically if the shape doesn't match. Then in forward we can just check that we have at least the required length and truncate it if needed. That would allow to store the PEs for multiple iterations in case we find the redrawing is slow or want deterministic behavior.

Yes, reset should be called manually, at the start of the forward of the hosting module.
but your solution does not work: we really do need to call it again at each forward pass, during training stage, because if we don't do it, we don't update the parameters of the SPE

Still, we could store the pure noise z for reusability, yes

[aliutkus]

looks a bit redundant to me, what about a max_length ?

[cifkao]

Fixed

[aliutkus]

I don't really like it that we need to provide the key_shape each time, we can't do otherwise ?

[cifkao]

Fixed