# Automatic differentiation (AD) of a loss function which maps unitary matrix onto number

Is it possible to estimate whether automatic differentiation (AD) techniques could enable a more efficient way to repeatedly compute the derivative $$\delta L / \delta u^*_{ij}$$ of a specific loss function $$L$$ than using an analytic expression?

Here, I consider the specific loss function, $$L(U) = \sum_{n=1}^N (U^\dagger P U)_{nn}^2 \;,$$ which maps a unitary matrix $$U$$ onto a number. The analytic expression of the derivative is available as $$\frac{\delta L}{\delta u^*_{ij}} = 2 (U^\dagger P U)_{jj} (P U)_{ij} \;.$$ $$U$$ is a unitary matrix and $$P$$ is a given hermitian matrix, both of dimension $$N\times N$$, where $$N$$ is of the order of $$10^4$$. The derivative, I am interested in, $$\delta L / \delta u^*_{ij}$$ itself is a $$N \times N$$ matrix.

So the computation of both $$L$$ as well as $$\delta L / \delta u^*_{ij}$$ has a complexity of $$\mathcal O(N^3)$$, if I implement the analytic expressions. For example, if I use BLAS level 3 routines.

My question is, if there is a way to estimate whether AD techniques could offer a more efficient way to repeatedly calculate the derivative $$\delta L / \delta u^*_{ij}$$ for all $$i,j$$? With "efficient" I mean "less floating point operations". Are there general rules how to estimate the complexity and the number of floating point operations of AD approaches (such as forward and reverse mode, ...) for such a case?

PS: with "repeatedly" I mean about 100 evaluations of $$\delta L / \delta u^*_{ij}$$ with different unitary matrices $$U$$.

• It's hard to imagine a black box beating an informed implementation of the derivative given that even just evaluating $L(U)$ costs $O(N^3)$. Jun 14, 2023 at 17:12
• I have doubts about the indices in the derivative. You should have the diagonal entries on the original symmetric product, and the mixed indices on the inner derivative. And are you sure about the factor 2, as there are 4 copies of $U$ involved in symmetric roles. Jun 14, 2023 at 18:48
• You are right about the indices. Regarding the factor 2, I thought that I have to consider $u_{ij}$ and $u^*_{ij}$ as independent? But maybe that's not correct. Jun 14, 2023 at 19:33
• You can confirm that there must be a factor 4 there (or more likely a completely different expression) by testing the case $N=1$. Jun 15, 2023 at 7:37
• According to gregs answer the factor 2 seems right. Jun 15, 2023 at 18:23

## 1 Answer

Your analytic derivative expression doesn't look right.

Let's calculate the gradient (in the Wirtinger sense) \eqalign{ \def\l{L} \def\o{{\tt1}} \def\p{\partial} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\LR#1{\left(#1\right)} \def\qiq{\quad\implies\quad} P &= P^H \qquad\qquad\;\:\qiq &P^*=P^T \\ Q &= U^HPU \qquad\quad\qiq &Q^H=Q,\;\;Q^*=Q^T \\ \l &= I:\LR{Q\odot Q^*} &\{{\rm loss\:is\:a\;}{\sf Real}{\rm\:scalar}\} \\ &= I:\LR{Q\odot Q^T} \\ d\l &= I:\LR{Q\odot dQ^T + Q^T\odot dQ} \\ &= 2\LR{I\odot Q^T}:dQ \\ &= 2\LR{I\odot Q^T}:\LR{U^HP\:dU} \\ &= 2\LR{U^HP}^T\LR{I\odot Q^T}:dU \\ &= 2\,P^*U^*\LR{I\odot Q^*}:dU \\ \grad{\l}{U} &= 2\,P^*U^*\LR{I\odot Q}^* \qiq &\grad{\l}{U^*} = 2\,{PU}\LR{I\odot Q} \\ } This gradient is a $$(N\times N)$$ matrix. Its $$(i,j)^{th}$$ component can be found by multiplying on the left and right by the standard $$\{e_i,\,e_j\}$$ basis vectors $$\grad{\l}{U^*_{ij}} = 2\,e_i^T{PU}\LR{I\odot Q}e_j$$ In index notation, this can be reduced to the expression in your post.

But note that this gradient does not incorporate any kind of constraint. Therefore, if it is used in a gradient descent algorithm, it is likely to take steps which will destroy the unitarity of $$U$$.

Notation: In the above, $$\{\odot,\,:\}$$ denote the Hadamard and Frobenius products \eqalign{ F &= A\odot B &\qiq F_{ij} = A_{ij} B_{ij} \\ \phi &= A:B &\qiq \phi = \sum_{i=1}^n\sum_{j=1}^n F_{ij} \;\doteq\; {\rm trace}\LR{A^TB} \\ } and the transpose, complex and hermitian conjugates are denoted by $$\{U^T,\,U^*,\,U^H\}$$

## Update

Let $$\o$$ denote the all-ones vector and $$q={\rm diag}(Q),\,$$ then the gradient can be written as \eqalign{ \grad{\l}{U^*} = 2\LR{PU}\odot\LR{\o q^T} \\ } which has an $${\cal O}\LR{N^2}$$ operation count, assuming that $$(PU)$$ has already been calculated during the evaluation of the cost function (e.g. to monitor convergence of the gradient descent iterations).

In some languages (e.g. Julia) you can access the columns of a matrix at zero cost, in which case using the $$k^{th}$$ component of $$q$$ to scale the $$k^{th}$$ column of $${PU}$$ is even faster than BLAS.

• My bad -- the gradient expression can be reduced to the one in your post.
– greg
Jun 15, 2023 at 17:42
• In the comment above I argued that there should be a factor 4, as one can check from the $N=1$ case. Indeed, if I am not mistaken in dimension 1 the function reduces to $L = PU^4$, but your expression for the gradient reduces to $2P^2U^3$ rather than $4P^2U^3$ as it should. Or am I doing something wrong here? Jun 15, 2023 at 22:23
• @FedericoPoloni In complex scalar $(N=1)$ case, the Wirtinger derivative of the squared modulus $\;|\phi|^2=\phi^*\phi\;$ is $\;\phi^*\;$ not $\;2\phi^*\;$
– greg
Jun 15, 2023 at 22:48
• Ah, thanks, I missed that it was a different definition that the one I am used to for functions $\mathbb{C}^N\to\mathbb{R}$. Jun 16, 2023 at 6:41