Shivam PandeyAI Research EngineerIM
Writing / KERNELIZED

AI · Compute · Autograd · MATH

Max Kernel: Forward and Backward Pass (MATH)

The forward and backward mathematics of a max-reduction kernel using indexed tensor notation.

Shivam Pandey2 min read

  • A∈Rs2s1A \in \mathbb{R}^{s_2 s_1} where s2s_2 and s1s_1 are the index set, for example in a 5D tensor A∈RijklmA \in \mathbb{R}^{ijklm} a possible index set could be s2={i,j}s_2 = \{ i, j \} and s1={k,l,m}s_1 = \{k, l, m \}.

  • Bs2=maxs1(As2s1)B_{s_2} = max_{s_1}(A_{s_2 s_1}) is a max operation over s1 index set of A, and the resulting index set is the index set of AA reduced over s1s_1 s.t. B∈Rs2B \in \mathbb{R}^{s_2}

Max (Reduction) operation

Bs2=max⁡s1(As2s1)=As2s1m∣s1m=argmaxs1(As2s1)\begin{aligned}B_{s_2} = \max_{s_1}(A_{s_2 s_1}) = A_{s_2 s^m_1} \big|_{s^m_1 = argmax_{s_1}(A_{s_2 s_1})} \end{aligned}

Backward Pass:

Here we will deduce the pullback of AA under the max⁡\max operation, w.r.t. BB.

∂Bs2∂As2′s1′=∂As2s1m∂As2′s1′=1s2=s2′⋅1s1m=s1′\begin{aligned}\frac{\partial B_{s_2}}{\partial A_{s'_2 s'_1}} = \frac{\partial A_{s_2 s^m_1}}{\partial A_{s'_2 s'_1}} = \mathbb{1}_{s_2 = s'_2} \cdot \mathbb{1}_{s^m_1 = s'_1} \end{aligned}

Full reduction: With Loss Drivative

Let’s assume the output Bs2B_{s_2} being used by some frisky function to generate loss Os3O_{s_3} and somehow we have the pullback of BB w.r.t. OO as dBs3s2=∂Os3∂Bs2dB_{s_3 s_2} = \frac{\partial O_{s_3}}{\partial B_{s_2}}, and now we are interested in find out what’s the pullback of As2′s1′A_{s'_2 s'_1} w.r.t. Os3O_{s_3} i.e. ∂Os3∂As2′s1′\frac{\partial O_{s_3}}{\partial A_{s'_2 s'_1}}.

∂Os3∂As2′s1′=∑s2∂Os3∂Bs2∂Bs2∂As2′s1′=∑s2dBs3s21s2=s2′⋅1s1m=s1′=dBs3s2′1s1m=s1′∣s1m=argmaxs1(As2′s1)\begin{aligned}\frac{\partial O_{s_3}}{\partial A_{s'_2 s'_1}} = \sum_{s_2} \frac{\partial O_{s_3}}{\partial B_{s_2}} \frac{\partial B_{s_2}}{\partial A_{s'_2 s'_1}} = \sum_{s_2} dB_{s_3 s_2} \mathbb{1}_{s_2 = s'_2} \cdot \mathbb{1}_{s^m_1 = s'_1} = dB_{s_3 s'_2} \mathbb{1}_{s^m_1 = s'_1} \big|_{s^m_1 = argmax_{s_1}(A_{s'_2 s_1})} \end{aligned} ∂Os3∂As2′s1′=dBs3s2′1s1m=s1′∣s1m=argmaxs1(As2′s1)\begin{aligned} \frac{\partial O_{s_3}}{\partial A_{s'_2 s'_1}} = dB_{s_3 s'_2} \mathbb{1}_{s^m_1 = s'_1} \big|_{s^m_1 = argmax_{s_1}(A_{s'_2 s_1})} \end{aligned}

Application in Attention: Specialize the Expressions

Setup:

Sij=Sijp−max⁡j(Sijp)Pij=softmax(Sij)Sij∈RM×NPij∈RM×NOϕ∈RdPϕij=∂Oϕ∂Pij∈Rϕ×M×N  ⟹  KnowndSϕij=∂Oϕ∂Sij∈Rϕ×M×N=Pij[dPij−dPi:T∘Pi:]\begin{aligned}S_{ij} = S^p_{ij} - \max_j(S^p_{ij}) \\ P_{ij} = softmax(S_{ij}) \\ S_{ij} \in \mathbb{R}^{M \times N} \\ P_{ij} \in \mathbb{R}^{M \times N} \\ O_{\phi} \in \mathbb{R} \\ dP_{\phi ij} = \frac{\partial O_{\phi}}{\partial P_{ij}} \in \mathbb{R}^{\phi \times M \times N} \implies \text{Known} \\ dS_{\phi ij} = \frac{\partial O_{\phi}}{\partial S_{ij}} \in \mathbb{R}^{\phi \times M \times N} = P_{ij} \left[dP_{ij} - dP_{i:}^T \circ P_{i:} \right] \end{aligned}

Differentiation

∂Oϕ∂Si′j′p=∑ij∂Oϕ∂Sij∂Sij∂Si′j′p=∑ijdSϕij(1ij=i′j′−1i=i′1jm=j′∣argmaxj(Si′jp))=dSi′j′−∑jdSi′j1jm=j′∣jm=argmaxj(Si′jp)=dSi′j′−1jm=j′∣jm=argmaxj(Si′jp)∑jdSi′j\begin{aligned}\frac{\partial O_{\phi}}{\partial S^p_{i'j'}} = \sum_{ij} \frac{\partial O_{\phi}}{\partial S_{ij}} \frac{\partial S_{ij}}{\partial S^p_{i'j'}} = \sum_{ij} dS_{\phi ij} (\mathbb{1}_{ij = i'j'} - \mathbb{1}_{i = i'} \mathbb{1}_{j^m = j'} \big|_{argmax_{j}(S^p_{i'j})}) \\ = dS_{i'j'} - \sum_{j} dS_{i'j} \mathbb{1}_{j^m = j'} \big|_{j^m={argmax_{j}(S^p_{i'j})}} = dS_{i'j'} - \mathbb{1}_{j^m = j'} \big|_{j^m={argmax_{j}(S^p_{i'j})}} \sum_{j} dS_{i'j} \end{aligned}

Replacing dSijdS_{ij}:

∂Oϕ∂Si′j′p=dSi′j′−1jm=j′∣jm=argmaxj(Si′jp)∑jdSi′j\begin{aligned}\frac{\partial O_{\phi}}{\partial S^p_{i'j'}} = dS_{i'j'} - \mathbb{1}_{j^m = j'} \big|_{j^m={argmax_{j}(S^p_{i'j})}} \sum_{j} dS_{i'j} \\ \end{aligned}