# Factor of two in Zygote complex gradient

**URL:** <https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545>\
**Category:** Optimization (Mathematical)\
**Tags:** zygote\
**Created:** [March 7, 2022, 11:51pm UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545 "2022-03-07T23:51:17Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![goerz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerz/32/3269_2.png) [@goerz](https://discourse.julialang.org/u/goerz)\
**Post date:** [March 7, 2022, 11:51pm UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/1 "2022-03-07T23:51:17Z")

</div>

Either Zygote’s gradients for a scalar function with respect to a complex vector are off by a factor of two, I’ve been doing matrix algebra wrong for the last 10 years, or I’m just being stupid here. Consider the following example:

```julia
using LinearAlgebra
using Zygote

# Two random normalized complex state vectors |Ψ⟩, |Ψtgt⟩
N = 10
Ψ = rand(N) .* exp.(2π .* 1im .* rand(N)); Ψ ./= norm(Ψ)
Ψtgt = rand(N) .* exp.(2π .* 1im .* rand(N)); Ψtgt ./= norm(Ψtgt)

# Functional J = 1 - |⟨Ψtgt|Ψ⟩|² = 1 - ⟨Ψtgt|Ψ⟩⟨Ψ|Ψtgt⟩
J(Ψ, Ψtgt) = 1 - abs(Ψtgt ⋅ Ψ)^2

# Zygote gradient
∇1 = Zygote.gradient(ϕ -> J(ϕ, Ψtgt), Ψ)[1]

# Manual gradient: ∂/∂⟨Ψ| (1 - ⟨Ψtgt|Ψ⟩⟨Ψ|Ψtgt⟩) = - ⟨Ψtgt|Ψ⟩ |Ψtgt⟩
grad(Ψ, Ψtgt) = -1 .* (Ψtgt ⋅ Ψ) .* Ψtgt
∇2 = grad(Ψ, Ψtgt)

norm(abs.(∇2 - ∇1)) # should be zero, but is not
norm(abs.(∇2 - ∇1/2)) # zero

```

I’m using the notation `⟨a|b⟩ ≡ a ⋅ b` for the inner product here (borrowed from quantum mechanics, `|b⟩` is a (column) vector here, and `⟨b|` would be the conjugate-transpose (row) vector, the “co-state”).

I was always taught ([and have taught others](https://qucontrol.github.io/krotov/master/07_krotovs_method.html#equation-chi-boundary)) that if you have a functional `J(|Ψ⟩)` you can do the derivatives by simply treating states and co-states as independent variables. I actually never bothered to check entrywise that this derivative w.r.t. the co-state matches the [Wirtinger derivative](https://qucontrol.github.io/krotov/master/07_krotovs_method.html#equation-wirtinger-deriv), but that was always the assumption.

So, am I doing something wrong here, or is there something wrong with Zygote?

---

<div class="post-metadata">

**Author:** ![ericphanson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ericphanson/32/215186_2.png) [@ericphanson](https://discourse.julialang.org/u/ericphanson)\
**Post date:** [March 8, 2022, 1:28am UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/2 "2022-03-08T01:28:55Z")

</div>

I think your manual gradient is missing a 2. The variable your differentiating shows up twice so you need the product rule (or alternatively when looking at the absolute value squared you need to chain rule through the squared function which gives a 2).

[http://www.matrixcalculus.org/](http://www.matrixcalculus.org/) seems to agree (although they assume everything is real):

 ![image](https://global.discourse-cdn.com/julialang/original/3X/e/2/e2edbeffd9eb994600c7f5d0706b5b1c55fa1428.jpeg)  
(Sorry for the screenshot, I’m on my phone).

---

<div class="post-metadata">

**Author:** ![goerz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerz/32/3269_2.png) [@goerz](https://discourse.julialang.org/u/goerz)\
**Post date:** [March 8, 2022, 2:41am UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/3 "2022-03-08T02:41:58Z")

</div>

In the real case, I’d agree, but for a complex vector, I don’t think that’s still true. In fact, it’s the same for a complex scalar: For `J(z) = 1 - |z|² = 1 - z z̄`, the Wirtinger derivative `∂J/∂z̄` would be defined as `(∂J/∂Re[z] + 𝕚 ∂J/∂Im[z])/2`, see e.g. [Complex Derivatives, Wirtinger View and the Chain Rule | Ekin Akyürek](https://www.ekinakyurek.me/complex-derivatives-wirtinger/). So, with `z = x + 𝕚y` that would be `-(2x + 𝕚 2y) / 2 = -z`, which also matches the intuition that `z` and `z̄` can simply be treated as independent variables. However, Zygote gives

```julia
J(z) = 1 - abs(z)^2
Zygote.gradient(z -> J(z), 1+1im)
(-2.0 - 2.0im,)

```

I then found that the Zygote manual uses _exactly_ this example in the [chapter on Complex Differentiation](https://fluxml.ai/Zygote.jl/latest/complex/), and that I just didn’t read that closely enough. It also gives the solution of manually defining a Wirtinger derivative. This indeed also works for the vector case. In my example script, I can add

```julia
# Wirtinger derivative
function wirtinger(func, Ψ)
    y, back = Zygote.pullback(func, Ψ)
    du = back(1.0)[1]
    dv = back(1.0im)[1]
    return (du' + im*dv')/2, (du + im*dv)/2
end

∇3 = wirtinger(Ψ -> J(Ψ, Ψtgt), Ψ)[2]
norm(abs.(∇3 - ∇2)) # zero, as it should be

```

which does the right thing.

I’m still a little confused about what just `Zygote.gradient` does (I might have to re-read that chapter a few more times). In the scalar case, I’m pretty sure it does

```julia
∂f(z)/∂z -> ∂f(x, y)/∂x + 𝕚 ∂f(x, y)/∂y

```

which seems somewhat trivially connected to the correct Wirtinger derivative (up to my missing factor of 2).

On the other hand, it would seem like `Zygote.gradient` does a single backward-propagation, whereas `wirtinger` does two. So probably, I can’t use the simple `Zygote.gradient` in general.

---

<div class="post-metadata">

**Author:** ![ericphanson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ericphanson/32/215186_2.png) [@ericphanson](https://discourse.julialang.org/u/ericphanson)\
**Post date:** [March 8, 2022, 3:10am UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/4 "2022-03-08T03:10:33Z")

</div>

Ah I see, I did mentally skip over the Wirtinger part. To me it’s less natural to define two derivatives like that, but googling it, it seems pretty common and probably has its uses.

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [March 8, 2022, 3:28am UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/5 "2022-03-08T03:28:42Z")

</div>

See also [Taking Complex Autodiff Seriously in ChainRules - #62 by Mason](https://discourse.julialang.org/t/taking-complex-autodiff-seriously-in-chainrules/39317/62) and [Complex numbers · ChainRules](https://juliadiff.org/ChainRulesCore.jl/stable/maths/complex.html) (since Zygote is built on top of ChainRules).

---

<div class="post-metadata">

**Author:** ![goerz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerz/32/3269_2.png) [@goerz](https://discourse.julialang.org/u/goerz)\
**Post date:** [March 8, 2022, 8:14pm UTC](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/6 "2022-03-08T20:14:59Z")

</div>

Ok, I think I got the connection between `Zygote.gradient` and the Wirtinger derivative now. It seems Iike I can in fact use a single call to `Zygote.gradient`, because my function `J` is \in \mathbb{R} (whereas [the above `wirtinger` function](https://discourse.julialang.org/t/factor-of-two-in-zygote-complex-gradient/77545/3) also handles the more general case of `J` \in \mathbb{C}). Specifically:

1. `Zygote.gradient(func, Ψ)` is equivalent to `y, back = Zygote.pullback(func, Ψ); back(1.0)[1]`, according to the [Zygote documentation](https://fluxml.ai/Zygote.jl/latest/adjoints/#Pullbacks-1).

2. `Zygote.pullback` falls back to `ChainRules.rrule` (or, at least, is compatible with `ChainRules.rrule`).

3. Given a function J: \mathbb{C}^N \rightarrow \mathbb{R}, and `z⃗, back = Zygote.pullback(J, z⃗₀)`, the pullback for the imaginary unity is zero, that is, `dv = back(1im)[1] = 0` in the `wirtinger` function:

Thus, in the `wirtinger` function, `du = Zygote.gradient(func, Ψ)` and `dv = 0`, and the two outputs are simply `du'/2` and `du/2`. That is, the gradient up to a factor of two, as I was observing.