Adding BatchNorm

[egretholier]

Cons:

• Yet another argument to AttentionLayers
• This is different from FFNBN in the Leveraging Batch Normalization for Vision Transformers paper.

Pros:

• Time series transformer papers are empirically claiming better performance. A Transformer-based Framework for Multivariate Time Series Representation Learning made the original claim and A Time Series is Worth 64 Words: Long-term Forecasting with Transformers made it popular.

[lucidrains]

@egretholier hey, thanks for adding this! could you throw in a test as well as make sure it asserts out if autoregressive or has any masking (masked batch norm is doable, but probably too complex for what you wish to do)

[lucidrains]

@egretholier for the record, i did try https://arxiv.org/abs/2003.07845 back in the day (collaborating with Aran Komat during a time when transformers were relatively unknown) and did not see much of anything.