Question of the equation in Flash Attention 2 Paper

[jeffrey-sunh1]

Hi Tri,
@tridao

I have a question about the equation in chapter 3.1.1 Forward pass of Flash attention 2 paper.

As you mentioned "Only at the every end of the loop do we scale the final $$\tilde{O}^{(last)}$$ by $$diag(l^{(last)} )^{−1}$$ to get the right output."

Suppose 2 is the last one, I can use $$diag(\ell^{(2)} )^{−1}$$ to scale $$\tilde{O}^{(2)}$$, then we get the following formula:

$$diag(\ell^{(2)} )^{−1}\tilde{O}^{(2)}$$ $$=diag(\ell^{(2)} )^{−1}[diag(\ell^{(1)})^{-1}O^{(1)}+e^{S^{(2)}-m^{(2)}}V^{(2)}]$$ $$=diag(\ell^{(2)} )^{−1}diag(\ell^{(1)})^{-1}O^{(1)}+diag(\ell^{(2)} )^{−1}e^{S^{(2)}-m^{(2)}}V^{(2)}\quad \quad equation-1 $$

The expected result should be

$$O^{(2)}=diag(\ell^{(1)}/\ell^{(2)})^{-1}O^{(1)}+diag(\ell^{(2)})^{-1}e^{S^{(2)}-m^{(2)}}V^{(2)}\quad \quad equation-2 $$

The right-hand side of the plus sign in equation 1 is equal to that in equation 2, but the left-hand sides of the plus sign do not match. Could you provide some clarification or guidance on this matter?

[SimpleTheoryOfTypes]

The formula for $O_2$ should be:

$O_2 = \text{diag}\bigg(\frac{l_2}{l_1}\bigg)^{-1} e^{m_1-m_2} O_1 + \text{diag}(l_2)^{-1} e^{S_2-m_2} \times V_2 = \text{diag}(l_2)^{-1} e^{S_1-m} V_1 + \text{diag}(l_2)^{-1} e^{S_2-m} V_2 = O$

[jeffrey-sunh1]

@SimpleTheoryOfTypes So these are typos in the paper?

[SimpleTheoryOfTypes]

Yes, I believe so. It would be best for the original authors to confirm.