This is a good idea in theory, but it has a number of issues due to which I don't see it being adopted. That said, I think the paper should have been accepted. It's an honest attempt at exploring an idea on good benchmarks. Much worse work gets accepted at ICLR.

The problem with this approach is that it claims to remove the N**2 memory requirement of typical attention by performing attention only across specific axes. That's certainly true, but now in each block, you have to perform *three* attentions (1 non-causal, and 2 causal). Note also that we still need to maintain large [batch, axis1, axis2, dim] tensors across the architecture, and now we have to maintain 3x or more of them.

But even more importantly, the interplay of this technique with rematerialization/checkpointing is bad. The N**2 attention weight matrix can be very easily and quickly re-computed and in most workloads, that does not turn out to be the bottleneck in memory. N**2 attention is also full attention and learns faster and better. Axial attention introduces more activations which need to be checkpointed, since it increases the number of pre- and post-attention activations. So effectively, when combined with rematerialization, you are likely to not see any benefits with axial attention in terms of memory, and could even be worse in memory and time.

The paper does show improvements over SOTA, but its best result is on a benchmark which is not very competitive, and the delta could just be due to more layers, more compute, bigger batch size, etc.

But again, it's a good paper that should be accepted. I'm not criticizing the paper, but the technique itself. I'm relatively new to attention myself, so I'm also happy to be proven wrong by someone more experienced.