Transformers in Flax

[conceptofmind]

Hi,

I previously worked with lucidrains to implement 18 different Vision Transformers and PaLM in Google's Flax/JAX. I am going to release versions in DeepMind's Haiku soon as well. If this is of any interest I can open up a PR and add them to the corresponding places on the list.

Vision Transformers: https://github.com/conceptofmind/vit-flax

PaLM: https://github.com/conceptofmind/PaLM-flax

Thank you,

Enrico

[cmhungsteve]

Hello,

I added your Vision Transformer repo to the list. PaLM is a pure Transformer, which is out of this repo's scope.

Thanks.

[conceptofmind]

Hi @cmhungsteve,

Thank you for adding the Vision Transformer to the list.

The ViT-flax repository I linked above currently includes Flax transformer architecture implementations for papers such as:

• Going deeper with Image Transformers
• CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention
• CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification
• CvT: Introducing Convolutions to Vision Transformers
• DeepViT: Towards Deeper Vision Transformer
• LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference
• Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding
• Three things everyone should know about Vision Transformers
• Learning to Merge Tokens in Vision Transformers
• Rethinking Spatial Dimensions of Vision Transformers
• RegionViT: Regional-to-Local Attention for Vision Transformers
• ScalableViT: Rethinking the Context-oriented Generalization of Vision Transformer
• Better plain ViT baselines for ImageNet-1k
• Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
• Twins: Revisiting the Design of Spatial Attention in Vision Transformers
• Vision Transformer for Small-Size Datasets

I will open up a PR for the matching sections. And provide an update when all the Haiku versions are complete.

Thank you,

Enrico