Backpropagation

Last but not least, KeOps fully supports automatic differentiation. Most of the magic required is implemented by the F::DiffT attributes of KeOps formulas and reductions, as discussed in previous pages.

Backprop through a Sum reduction

Then, to implement the PyTorch backward of the KeOps Genred operator, we simply have to remember that if $(g_i)\in {\mathbb{R}}^{\mathrm{M}\times \mathrm{E}}$ is a “gradient to backpropagate” with respect to the output $(a_i)\in {\mathbb{R}}^{\mathrm{M}\times \mathrm{E}}$ of a Genred call with a Sum reduction, we can write that for all variations $(\delta p,\delta x_i,\delta y_j)$ of the parameters, $i$- and $j$-variables, at order 1:

$$\begin{array}{r}\begin{array}{rl}⟨& \sum _{j=1}^{\mathrm{N}}F(p+\delta p,x_i+\delta x_i,y_j+\delta y_j)-F(p,x_i,y_j),g_i{⟩}_{{\mathbb{R}}^{\mathrm{M}\times E}}\\ =& \sum _{i=1}^{\mathrm{M}}\sum _{j=1}^{\mathrm{N}}(⟨{\partial }_{p}F(\dots )\cdot g_i,\delta p⟩+⟨{\partial }_{x_i}F(\dots )\cdot g_i,\delta x_i⟩+⟨{\partial }_{y_j}F(\dots )\cdot g_i,\delta y_j⟩).\end{array}\end{array}$$

Consequently, performing the appropriate permutations of sums:

$$\begin{array}{r}\begin{array}{rl}{\partial }_{x_i}[\sum _{j=1}^{\mathrm{N}}F(p,x_i,y_j)]\cdot (g_i)& =\phantom{\sum _{i=1}^{\mathrm{M}}}\sum _{j=1}^{\mathrm{N}}({\partial }_{x_i}[F(p,x_i,y_j)]\cdot g_i),\\ {\partial }_{y_j}[\sum _{j=1}^{\mathrm{N}}F(p,x_i,y_j)]\cdot (g_i)& =\phantom{\sum _{i=1}^{\mathrm{M}}}\sum _{i=1}^{\mathrm{M}}({\partial }_{y_j}[F(p,x_i,y_j)]\cdot g_i),\\ {\partial }_p[\sum _{j=1}^{\mathrm{N}}F(p,x_i,y_j)]\cdot (g_i)& =\sum _{i=1}^{\mathrm{M}}\sum _{j=1}^{\mathrm{N}}({\partial }_p[F(p,x_i,y_j)]\cdot g_i).\end{array}\end{array}$$

Backprop through a Log-Sum-Exp reduction

Similarly, when $(a_i)$ is given through a Log-Sum-Exp reduction:

$$\begin{array}{r}{a}_i=\log \sum _{j=1}^{\mathrm{N}}\exp F(p,x_i,y_j),\end{array}$$

straightforward computations show that:

$$\begin{array}{r}\begin{array}{rl}{\partial }_{x_i}[\log \sum _{j=1}^{\mathrm{N}}\exp F(p,x_i,y_j)]\cdot (g_i)& =\phantom{\sum _{i=1}^{\mathrm{M}}}\sum _{j=1}^{\mathrm{N}}{e}^{F(p,x_i,y_j)-a_i}\cdot ({\partial }_{x_i}[F(p,x_i,y_j)]\cdot g_i),\\ {\partial }_{y_j}[\log \sum _{j=1}^{\mathrm{N}}\exp F(p,x_i,y_j)]\cdot (g_i)& =\phantom{\sum _{i=1}^{\mathrm{M}}}\sum _{i=1}^{\mathrm{M}}{e}^{F(p,x_i,y_j)-a_i}\cdot ({\partial }_{y_j}[F(p,x_i,y_j)]\cdot g_i),\\ {\partial }_p[\log \sum _{j=1}^{\mathrm{N}}\exp F(p,x_i,y_j)]\cdot (g_i)& =\sum _{i=1}^{\mathrm{M}}\sum _{j=1}^{\mathrm{N}}{e}^{F(p,x_i,y_j)-a_i}\cdot ({\partial }_p[F(p,x_i,y_j)]\cdot g_i).\end{array}\end{array}$$

In other words, a backward pass through a Genred call that involves a Sum or a Log-Sum-Exp reduction can always be written as a symbolic Map-Reduce computation.

Bootstrapping derivatives of arbitrary order

Applying these commutation rules between the differential operator ${\partial }_{\mathtt{\text{V}}}$ and the Sum or Log-Sum-Exp reductions, the pykeops/torch/generic/generic_red.py module provides full compatibility between KeOps LazyTensors and the torch.autograd package. Thanks to recursive calls to the Genred operator and to our symbolic math engine, everything works just fine – even high-order derivatives.