information:highlighting, linguistic:voice_change
Numerical algorithms

The attention mechanism in Transformers takes three arguments: a "query vector" q {\displaystyle q}, a list of "key vectors" k 1, …, k N {\displaystyle k_{1},\dots,k_{N}}, and a list of "value vectors" v 1, …, v N {\displaystyle v_{1},\dots,v_{N}}, and outputs a softmax-weighted sum over value vectors: o = ∑ i = 1 N e q T k i − m ∑ j = 1 N e q T k j − m v i {\displaystyle o=\sum _{i=1}^{N}{\frac {e^{q^{T}k_{i}-m}}{\sum _{j=1}^{N}e^{q^{T}k_{j}-m}}}v_{i}} The standard softmax method involves several loops over the inputs, which would be bottlenecked by memory bandwidth. The FlashAttention method is a communication-avoiding algorithm that fuses these operations into a single loop, increasing the arithmetic intensity. It is an online algorithm that computes the following quantities: z i = q T k i m i = max (z 1, …, z i) = max (m i − 1, z i) l i = e z 1 − m i + ⋯ + e z i − m i = e m i − 1 − m i l i − 1 + e z i − m i o i = e z 1 − m i v 1 + ⋯ + e z i − m i v i = e m i − 1 − m i o i − 1 + e z i − m i v i {\displaystyle {\begin{aligned}z_{i}&=q^{T}k_{i}&\\m_{i}&=\max(z_{1},\dots,z_{i})&=&\max(m_{i-1},z_{i})\\l_{i}&=e^{z_{1}-m_{i}}+\dots +e^{z_{i}-m_{i}}&=&e^{m_{i-1}-m_{i}}l_{i-1}+e^{z_{i}-m_{i}}\\o_{i}&=e^{z_{1}-m_{i}}v_{1}+\dots +e^{z_{i}-m_{i}}v_{i}&=&e^{m_{i-1}-m_{i}}o_{i-1}+e^{z_{i}-m_{i}}v_{i}\end{aligned}}} and returns o N / l N {\displaystyle o_{N}/l_{N}}. In practice, FlashAttention operates over multiple queries and keys per loop iteration, in a similar way as blocked matrix multiplication. If backpropagation is needed, then the output vectors and the intermediate arrays, {\displaystyle,} are cached, and during the backward pass, attention matrices are rematerialized from these, making it a form of gradient checkpointing.