Differentiating a Complex Valued Function

username: bchen
cooked: <p>Hi,</p>
<p>If I wanted to parameterize a function f: R → C via a real valued neural network followed by some operations that make it complex, how are the gradients and higher-order derivatives of f calculated?</p>
<p>As a simple example below, suppose I have a single linear layer and then multiply the output by a complex number. If I perform backpropagation, it seems that the gradients are only with respect to the real part of the output. In the below, the gradient of the weight matrix should be [[(2+3i)1, (2+3i)2], [(2+3i)1, (2+3i)2], but instead it returns [[2x1, 2x2], [2x1, 2x2]].</p>
<pre><code class="lang-auto">import torch

linear_layer = torch.nn.Linear(2, 2)
test = linear_layer.forward(torch.Tensor([[1., 2.]]))
(test*(2+3j)).sum().backward()
for i in linear_layer.parameters(): print(i.grad)
</code></pre>
<p>If this is what is going on, could I just take the imaginary part and call backward on that?</p>
<pre><code class="lang-auto">(test*(2+3j)).imag().sum().backward()
</code></pre>
<p>This seems to work for this toy example, but wanted to get some confirmation. Also, I am curious if this also holds for higher order derivatives calculated using functorch. E.g. if I take a hessian, will it only take the hessian of the real part of the output?</p>

username: KFrank
cooked: <p>Hi bchen!</p>
<aside class="quote no-group" data-username="bchen" data-post="1" data-topic="165209" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/bbe5ce/48.png" class="avatar"> bchen:</div>
<blockquote>
<p>If I wanted to parameterize a function f: R → C via a real valued neural network followed by some operations that make it complex, how are the gradients and higher-order derivatives of f calculated?</p>
</blockquote>
</aside>
<p>The answer is rather nuanced and the best (pytorch) explanation is here<br>
in the <a href="https://pytorch.org/docs/1.13/notes/autograd.html#autograd-for-complex-numbers" rel="noopener nofollow ugc">Autograd-for-Complex-Numbers documentation.</a></p>
<aside class="quote no-group">
<blockquote>
<p>If I perform backpropagation, it seems that the gradients are only with respect to the real part of the output.</p>
</blockquote>
</aside>
<p>There are two things going on here:</p>
<p>First, when you call <code>.backward()</code>, it uses, by default <code>gradient = None</code>.</p>
<p>The pytorch framework is designed to optimize real-valued loss functions.<br>
Because of this, when you call <code>some_complex_loss.backward()</code> pytorch<br>
computes (in effect) <code>some_complex_loss.real.backward()</code>.  (I’m pretty<br>
sure of this, but you might want to scour the above documentation to see<br>
if I’ve missed some nuance.)</p>
<p>Note, in particular, that <code>.backward()</code> computes, in the language of<br>
Wirtinger derivatives, <code>d / dz*</code> (not <code>d / dz</code>).</p>
<p>Second, you are computing the gradient with respect to a real tensor.<br>
(By default, <code>Linear</code> is instantiated with real weights.)  This causes<br>
<code>backward()</code> not to store (or maybe not even compute) the full complex<br>
gradient – presumably because tensors and their <code>.grad</code> properties are<br>
required to be of the same type.</p>
<aside class="quote no-group">
<blockquote>
<p>If this is what is going on, could I just take the imaginary part and call backward on that?</p>
</blockquote>
</aside>
<p>If you need the gradient of the imaginary part, yes, call <code>.backward()</code><br>
on <code>.imag</code>.  (Be sure to take into account that <code>.imag</code> is <em>real</em> so that<br>
<code>t_complex = t_complex.real + 1j * t_complex.imag</code>.)</p>
<p>Here is an illustration of these points:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.13.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; # create a complex Linear that happens to have real coefficients -- experimental
&gt;&gt;&gt; linear_layer = torch.nn.Linear (2, 2).to (dtype = torch.complex64)
&lt;path_to_pytorch_install&gt;\torch\nn\modules\module.py:975: UserWarning: Complex modules are a new feature under active development whose design may change, and some modules might not work as expected when using complex tensors as parameters or buffers. Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml if a complex module does not work as expected.
  warnings.warn(
&gt;&gt;&gt;
&gt;&gt;&gt; # input to Linear must also be complex
&gt;&gt;&gt; test = linear_layer (torch.tensor ([[1. + 1.j, 2. + 2.j]]))
&gt;&gt;&gt;
&gt;&gt;&gt; (test * (2 + 3j)).sum().backward (retain_graph = True)        # gradient = None (by default)
&gt;&gt;&gt; linear_layer.weight.grad   # now you get the full complex grad (of real part of "loss")
tensor([[-1.-5.j, -2.-10.j],
        [-1.-5.j, -2.-10.j]])
&gt;&gt;&gt; linear_layer.bias.grad     # now you get the full complex grad (of real part of "loss")
tensor([2.-3.j, 2.-3.j])
&gt;&gt;&gt;
&gt;&gt;&gt; linear_layer.zero_grad()
&gt;&gt;&gt; (test * (2 + 3j)).sum().real.backward (retain_graph = True)   # same as without .real
&gt;&gt;&gt; linear_layer.weight.grad
tensor([[-1.-5.j, -2.-10.j],
        [-1.-5.j, -2.-10.j]])
&gt;&gt;&gt; linear_layer.bias.grad
tensor([2.-3.j, 2.-3.j])
&gt;&gt;&gt;
&gt;&gt;&gt; linear_layer.zero_grad()
&gt;&gt;&gt; (test * (2 + 3j)).sum().imag.backward (retain_graph = True)   # .imag is (of course) different
&gt;&gt;&gt; linear_layer.weight.grad
tensor([[ 5.-1.j, 10.-2.j],
        [ 5.-1.j, 10.-2.j]])
&gt;&gt;&gt; linear_layer.bias.grad
tensor([3.+2.j, 3.+2.j])
</code></pre>
<p>Note the sign of the imaginary part of <code>linear_layer.bias.grad</code>.  This is<br>
due to the use of <code>d / dz*</code> as the gradient.</p>
<aside class="quote no-group">
<blockquote>
<p>Also, I am curious if this also holds for higher order derivatives calculated using functorch.</p>
</blockquote>
</aside>
<p>I haven’t used <code>functorch</code>, but I believe that it uses autograd under the<br>
hood, with an overlay of “more efficient” loops.  So I expect that these<br>
principles still apply and the autograd documentation still tells you the<br>
(piece-wise) details of what is happening.</p>
<p>Best.</p>
<p>K. Frank</p>