3-pass
\(\mathsf{NO}\) TATIONS
\(\{m_i\}{:}\max_{j=1}^i\left\{x_j\right\}\),
with initial value \(m_0=-\infty.\)
\(\{d_i\}{:}\sum_{j=1}^ie^{x_j-m_N}\),
with initial value \(d_0=0,d_N\) is the
denominator of safe softmax. \(\{a_i\}{:\text{
the final softmax value}}.\)
BODY \(\textbf{for }i\leftarrow 1,
N\textbf{ do}\) \[m_i\leftarrow\max\left(m_{i-1},x_i\right)\]
\(\mathbf{end}\)
\(\textbf{for }i\leftarrow 1, N\textbf{
do}\) \[d_i\leftarrow
d_{i-1}+e^{x_i-m_N}\] \(\mathbf{end}\)
\(\textbf{for }i\leftarrow 1, N\textbf{
do}\) \[a_i\leftarrow\frac{e^{x_i-m_N}}{d_N}\]
\(\mathbf{end}\)
这是 3 step 计算 attention
的方法,每一步都需要上一步的结果才可以继续计算。这样的话由于 sram
中没有足够的存储空间,因此需要多次访存。 ### Online attention \[\begin{aligned}
d_i^{\prime}& =\sum_{j=1}^ie^{x_j-m_i} \\
&= \left(\sum_{j=1}^{i-1} e^{x_j-m_i}\right)+e^{x_i-m_i} \\
&= \left(\sum_{j=1}^{i-1}
e^{x_j-m_{i-1}}\right)e^{m_{i-1}-m_i}+e^{x_i-m_i} \\
&= d_{i-1}' e^{m_{i-1}-m_i}+e^{x_i-m_i}
\end{aligned}\] 找到迭代式之后就可以从 3 step 降到 2 step \[\begin{aligned}&\mathbf{for~}i\leftarrow1,N\textbf{
do}\\&&&m_i&&\leftarrow&\max\left(m_{i-1},x_i\right)\\&&&d_i^{\prime}&&\leftarrow&d_{i-1}^{\prime}e^{m_{i-1}-m_i}+e^{x_i-m_i}\\&\mathbf{end}\\&\mathbf{for~}i\leftarrow1,N\textbf{
do}\\&&&a_i\leftarrow&&\frac{e^{x_i-m_N}}{d_N^{\prime}}\\&\mathbf{end}\end{aligned}\]
好像 FLOPs
计算量并没有减少,甚至还略有增加,因为现在每次都需要计算额外的 scale