Backward

For a function f defined on the domain \(f : X \to Y\), we define the gradient-adjoint-dot operator as

It is implied that \(v \in Y\) and the domain of \(\Psi[f, x]\) is \(\Psi[f, x] : Y \to X\).

Notice how the summation eliminates the indexing of the function, while the indexing for the gradient remains.

Using \(\Psi^i = \Psi[f^i, r^i]\), the chain rule above can be re-organized as a sequence of function evaluations of \(\Psi^i\)

The process is illustrated in the back-propagation`section of the figure. We see that for each evaluation of :math:Psi^i`, we obtain the gradient of \(F\) relative to the intermiedate variable \(r^i\), \(\nabla_{r^i} F\). Because we apply \(\Psi^i\) in the decreasing order of \(i\), this method is called backward propagation or reverse accumulation.

This method is also called the adjoint method in the analysis of dynamical systems because the summation is along the adjoint index of the jacobian \(\frac{\partial f_i}{\partial x_j}\). The main drawback of back propagation is that it requires one to store the intermediate results along with the function evaluation in order to compute the gradient-adjoint-dot operators \(\Psi^i\) depends on \(r^i\), which needs to be evaluated before the back propagation. However, the method gives the full gradient against the free variables x_j after one full accumulation, making it at advantageous for certain problems compared to forward accumulation, which we describe next.

In all three AD software packages we checked (TensorFlow, Theono, or autograd), a method to look up the the gradient-adjoint-dot operator is provided, either as a member of the operator entity or as an external dictionary.

Forward

In contrast, we can define a gradient-dot-operator as

It is implied that \(u \in X\) and the domain of \(\Gamma[f, x]\) is \(\Gamma[f, x] : X \to Y\).

Notice the summation is over the indexing of the free variable, \(x_j\). Hence, the name does not have adjoint like the previous operator. One way to think of \(\Gamma[f]\) is that it rotates \(u\) by the jacobian.

With the gradient-dot operator of \(\Gamma^i = \Gamma[f^i, r^i]\), we can write down the forward accumulation rule of AD:

This process is illustrated in the section on forward accumulation in the figure. We see that for each evaluation of \(\Gamma^i\), we obtain the directional derivative of \(r^i\) along \(u\), \(\sum \frac{\partial r^i}{\partial x_j} u_j\). The accumulation goes along the increasing order of \(i\), making the name forward accumulation a suitable one.

The advantage of forward accumulation is that one can evaluate the gradient as the function \(F\) is evaluated, and no intemediate results need to be saved: we see that when \(\Gamma^i\) is requested, \(r_i\) is already evaluated. This is clearly a useful feature when nesting (layers of neural networks or number of time steps) is high. However, the cost is we can only obtain a directional derivative. For some applications, this is useful (e.g., computing Hession for Newton-CG or trust-region Newton-CG methods). When the full gradient is desired, one needd to run the forward accumulation many times - as many times as the number of free parameters, which could be prohibatively high.

We shall note that this method is also called forward senstivity in the analysis of dynamical systems.

Discrete Fourier Transform

Discrete Fourier Transform is the discretized version of Fourier Transform. It is a commonly used density matrix operator in the modelling of physical process. This is mostly because spatial differentiation can be written as multiplication in the spectrum space.

The gradients involve complex numbers, which are tuples of two real numbes. The gradient that is conveniently used is

for \(z = x + \imath y\). It is related to the Wirtinger derivatives (Fourier transform is a harmonic function).

The gradient-adjoint-dot operator of a discrete fourier transform is its dual transform. Specifically,

where \(\Psi\) is the gradient-adjoint-dot operator. Notably, the free variables \(X\) and \(Y\) do not show up in the final expressions. This is because Fourier transforms are linear operators.

We do not include a formal proof in this blog. (The proof is relatively simple)

Gradients of the real fourier transforms (ifft and irfft) are slightly more complicated. Of the complex vector in a real fourier transform, only about half of the complex numbers are indepedent. The other half is the hermitian conjugates. (See What FFTW really computes)

Due to this hermitian property of the complex vector, irfft has two types of gradients:

• gradient over the full complex vector, which follows the same rules as the complex fourier transform

  \begin{equation*} \Psi[\mathrm{rfft}, X](V) = \mathrm{irfft}(V) , \end{equation*}
  \begin{equation*} \Psi^{\mathrm{full}}[\mathrm{irfft}, Y](V)_j = \mathrm{rfft}(V)_j, \end{equation*}

• gradient over the compressed complex vector

The full complex vector and the compressed complex vector are related by a decompress operation, which introduces a factor of 2 to the fourier modes with a hermitian conjugate mode.

\begin{equation*} \mathrm{decompress}(Y) = Y \end{equation*}

\begin{equation*} \Psi[\mathrm{decompress}, Y](V)_j = V_j \left\{ \begin{matrix} 1 & \mathrm{ if } j = N - j, \\ 2 & \mathrm{ if } j \neq N - j \end{matrix} \right. = V_j \mathsf{D}(j, N) . \end{equation*}

\begin{equation*} \Psi^{\mathrm{compressed}}[\mathrm{irfft}, Y](V)_j = \Psi[\mathrm{decompress}, Y](V)_j \Psi^{\mathrm{full}}[\mathrm{irfft}, Y](V)_j = \mathsf{D}(j, N) \mathrm{rfft}(V)_j, \end{equation*}

The complex version of Discrete Fourier Transform is implemented in TensorFlow (GPU only), Theono, and autograd, though it appears the version in autograd is incorrect. The real-complex transforms (rfft and irfft) are not implemented in any of the packages, neither is the decompress operation defined in any of these packages.

Resampling Windows

The resampling window converts a field representation between particles and meshes. It is written as

where \(p^i\) is the position of i-th particle/mesh point and \(q^j\) is the position of j-th mesh/particle point; both are usually vectors themselves (the universe has three spatial dimensions). \(p^i\) and \(q^i\) are themselves vectors with a spatial dimension (we will use the integer symbol \(a\) to index this dimension).

• paint: When \(p^i\) is the position of particles and \(q^j\) is the position of the mesh points, the operation is called a paint.
• readout: When \(p^i\) is the position of the mesh points and \(q^j\) is the position of mesh points, the operation is called a readout.

\(W\) is the resampling window function. A popular form is the cloud-in-cell window, which represents a linear interpolation:

for a given size of the window \(h\). We have used \(a\) as the index of the spatial dimensions.

Most windows are seperatable, which means they can be written as a product of a scalar function \(W_1\),

For these windows,

We can then write down the gradient-adjoint-dot operator of the window

• The first gradient corresponds to the displacement of the source.
• The second gradient corresponds to the displacment of the destination.
• The third gradient corresponds to the evolution of the field.

This looks complicated, and it is. However, in a particle mesh simulation, more often than not either one of the sources or the destination is a fixed grid which does not move. In this case \(v = 0\), and we do not need to compute the corresponding gradient term.

We also note that it is possible to extend these expressions to Smoothed Particle Hydrodynamics (SPH). In SPH, \(h\) is another free variable. We leave this to future work, but note that the symmetric property of the hydro-dynamical force kernel may introduce additional complication.

None of the three packages we surveyed (TensorFlow, Theono and autograd) recognizes these resampling window operators. The resampling window operators cannot be easily expressed as diagonal matrix operators.