Support Taylor Series Linear Attention

[hoangmit]

🚀 Feature

Can we support this in xformers?

First proposed in

• https://arxiv.org/abs/2312.04927
• Mentioned in https://hazyresearch.stanford.edu/blog/2023-12-11-zoology2-based
• and in https://hazyresearch.stanford.edu/blog/2024-03-03-based

Size trade off version: https://arxiv.org/abs/2402.04347

Another variant: https://arxiv.org/abs/2402.10644

Existing implementations:

• https://github.com/HazyResearch/based
• https://github.com/HazyResearch/zoology
• https://github.com/lucidrains/taylor-series-linear-attention