Does replacing some layers with linear attention (GQA or KDA) from FLA improves the model training?

[vukrosic]

1. measure baseline for 8M tokens and 20M tokens (follow SETUP_INSTRUCTIONS)
2. replace some layers with linear attention, measure the new training fairly (val loss, time) - you may tell AI to measure it fairly

• From previous experience, last layer should probably be classic (softmax) attention and the ratio should be 20-25% softmax attention
• Make sure to confirm that your results actually train better or faster and the tradeoff is worthed

[bigwolfeman]

I had to do some deep research on this a while back when I was working on a Retnet implementation. I found several papers mentioning it is best practice to interleave standard attention at a 1:4-1:7 ratio to other forms of attention.

MIRAS is one of the best write ups right now on the behavior of various attention-like implementations. https://research.google/blog/titans-miras-helping-ai-have-long-term-memory/ Their basic findings are: Use many different kinds at once and it just keeps getting better. They were using Neural Memory, Attention, and Retention all at once.

Maintaining some amount of full quadratic attention reduces susceptibility to the chunk boundary problem.

My TitanMAC repo is a near full implementation of MIRAS. I am not maintaining retention in the architecture. My experience with Retnet was very poor. I had to patch pytorch to just get it to work. It was never performant and the neural memory does about 90% of the work of attention while being non quadratic and extreme in context length.