Added dropout support to memory efficient variant

[usryokousha]

Hey Phil,

I have been using this repository for a project and I wanted to add dropout for completeness. I checked consistency with perceiver-ar impl.. I hope this is helpful.

-Matt

[lucidrains]

@usryokousha oh sure, thanks Matt! just for your information, the field is slowly starting to realize that the traditional dropout is pretty useless

however, structured dropout, like https://github.com/lucidrains/x-transformers#forgetful-causal-mask or https://arxiv.org/abs/2206.00826 can still be used, but would not need to exist within the attention operation

[lucidrains]

let's merge it for completeness sake though! hope rabe or flash attention is working well for your project! just one more note, you should use the CUDA implementation here for optimal performance

[usryokousha]

Phil, thanks for pointing out the two papers on dropout! I wonder how the Bayesformer paper's proposed dropout stands up in non-causal attention. In my own experiments I have always turned it off because I found it hurt training. The CUDA optimized flash attention package looks very appealing! This is going to help in my future projects for sure!