Notations and background information#

Throughout this note, we will use the following convention:

1. $x$x,$y$y,$a$a,$b$b,$v$v,$u$u,$ur$ur and$ui$ui are real-valued vectors and$z$z is a complex-valued vector that can be rewritten in terms of two real-valued vectors as$z=a+ib$z=a+ib.

2. $N$N and$M$M are two integers that we will use for the dimension of the input and output space respectively.

3. $f:{\mathcal{R}}^N\to {\mathcal{R}}^M$f:RN→RM is our basic real-to-real function such that$y=f(x)$y=f(x).

4. $g:{\mathcal{C}}^N\to {\mathcal{R}}^M$g:CN→RM is our basic complex-to-real function such that$y=g(z)$y=g(z).

For the simple real-to-real case, we write as$J_f$Jf​ the Jacobian matrix associated with$f$f of size$M\times N$M×N.This matrix contains all the partial derivatives such that the entry at position$(i,j)$(i,j) contains$\frac{\mathrm{\partial }{y}_i}{\mathrm{\partial }{x}_j}$∂xj​∂yi​​.Backward mode AD is then computing, for a given vector$v$v of size$M$M, the quantity$v^T{J}_f$vTJf​.Forward mode AD on the other hand is computing, for a given vector$u$u of size$N$N, the quantity$J_{f}u$Jf​u.

For functions that contain complex values, the story is a lot more complex. We only provide the gist here and the full description can be found atAutograd for Complex Numbers.

The constraints to satisfy complex differentiability (Cauchy-Riemann equations) are too restrictive for all real-valued loss functions, so we instead opted to use Wirtinger calculus.In a basic setting of Wirtinger calculus, the chain rule requires access to both the Wirtinger derivative (called$W$W below) and the Conjugate Wirtinger derivative (called$CW$CW below).Both$W$W and$CW$CW need to be propagated because in general, despite their name, one is not the complex conjugate of the other.

To avoid having to propagate both values, for backward mode AD, we always work under the assumption that the function whose derivative is being calculated is either a real-valued function or is part of a bigger real-valued function. This assumption means that all the intermediary gradients we compute during the backward pass are also associated with real-valued functions.In practice, this assumption is not restrictive when doing optimization as such problem require real-valued objectives (as there is no natural ordering of the complex numbers).

Under this assumption, using$W$W and$CW$CW definitions, we can show that$W=C{W}^{\ast }$W=CW∗ (we use$\ast$∗ to denote complex conjugation here) and so only one of the two values actually need to be “backwarded through the graph” as the other one can easily be recovered.To simplify internal computations, PyTorch uses$2\ast CW$2∗CW as the value it backwards and returns when the user asks for gradients.Similarly to the real case, when the output is actually in${\mathcal{R}}^M$RM, backward mode AD does not compute$2\ast CW$2∗CW but only$v^T(2\ast CW)$vT(2∗CW) for a given vector$v\in {\mathcal{R}}^M$v∈RM.

For forward mode AD, we use a similar logic, in this case, assuming that the function is part of a larger function whose input is in$\mathcal{R}$R. Under this assumption, we can make a similar claim that every intermediary result corresponds to a function whose input is in$\mathcal{R}$R and in this case, using$W$W and$CW$CW definitions, we can show that$W=CW$W=CW for the intermediary functions.To make sure the forward and backward mode compute the same quantities in the elementary case of a one dimensional function, the forward mode also computes$2\ast CW$2∗CW.Similarly to the real case, when the input is actually in${\mathcal{R}}^N$RN, forward mode AD does not compute$2\ast CW$2∗CW but only$(2\ast CW)u$(2∗CW)u for a given vector$u\in {\mathcal{R}}^N$u∈RN.

Default backward mode gradcheck behavior#

# Code from https://github.com/pytorch/pytorch/blob/58eb23378f2a376565a66ac32c93a316c45b6131/torch/autograd/gradcheck.py#L99-L105# Notation changes in this code block:# s here is y above# x, y here are a, b aboveds_dx=compute_gradient(eps)ds_dy=compute_gradient(eps*1j)# conjugate wirtinger derivativeconj_w_d=0.5*(ds_dx+ds_dy*1j)# wirtinger derivativew_d=0.5*(ds_dx-ds_dy*1j)d[d_idx]=grad_out.conjugate()*conj_w_d+grad_out*w_d.conj()# Since grad_out is always 1, and W and CW are complex conjugate of each other, the last line ends up computing exactly `conj_w_d + w_d.conj() = conj_w_d + conj_w_d = 2 * conj_w_d`.

Fast backward mode gradcheck#

While the above formulation of gradcheck is great, both, to ensure correctness and debuggability, it is very slow because it reconstructs the full Jacobian matrices.This section presents a way to perform gradcheck in a faster way without affecting its correctness.The debuggability can be recovered by adding special logic when we detect an error. In that case, we can run the default version that reconstructs the full matrix to give full details to the user.

The high level strategy here is to find a scalar quantity that can be computed efficiently by both the numerical and analytical methods and that represents the full matrix computed by the slow gradcheck well enough to ensure that it will catch any discrepancy in the Jacobians.

Gradgradcheck implementation#

PyTorch also provide a utility to verify second order gradients. The goal here is to make sure that the backward implementation is also properly differentiable and computes the right thing.

This feature is implemented by considering the function$F:x,v\to v^T{J}_f$F:x,v→vTJf​ and use the gradcheck defined above on this function.Note that$v$v in this case is just a random vector with the same type as$f(x)$f(x).

The fast version of gradgradcheck is implemented by using the fast version of gradcheck on that same function$F$F.