Self-attention Does Not Need O(n2) Memory

arxiv.org

xformers.ops.memory_efficient_attentionはこれを実装してると引用してる。

K*Q全部一旦計算する必要なくてある場所のattention計算するためにはその行列のその行だけでいいから各場所について1つずつ計算すれば必要メモリはO(n)って話。

√nってなんだろうと思ったが、並列計算のためにchunkに分ける話だった。