Taylor series of a matrix exponential

Question

I am looking to minimize the value of: $$g(t)=\mathrm{Tr}\left[\exp(X+tY)\right]$$ where both $X$ and $Y$ are symmetrical matrices with real coefficients. In general, $X$ and $Y$ do not commute so $\exp(X+tY)\neq\exp(X)\exp(tY)$. We can further assume that $tY$ is small when compared to $X$ at the minimum.

I assume that one of the the simplest approach is to attempt to write $g(t)$ as a Taylor expansion, like: $$g(t)=g_0+tg_1+\frac12t^2g_2+\dots$$ In this case, a good approximation for the minimum is easily obtained with $t\sim-\frac{g_1}{g_2}$ (Newton's method) and the process can be iterated until we meet a convergence criterion.

The first two coefficients are quite trivial to find. For instance, $g_0=\mathrm{Tr}\left[\exp(X)\right]$ and $g_1=\mathrm{Tr}\left[Y\exp(X)\right]$, as explained here. However, I spent some hours on this but I can't find an easy expression for $g_2$ yet.

Is there a proper way to express $g_2$ so it can be computed numerically?

Answer 1

We have $$\exp(X+tY) = \sum_{n=0}^\infty (X+tY)^n/n! $$ Now each $(X+tY)^n$ can be expanded as a sum of products of $n$ terms, where each term is either $X$ or $tY$. The coefficient of $t^2$ in $\exp(X+tY)$ is thus $$ \sum_{n=2}^\infty \frac{1}{n!} \sum_{a,b,c} X^{a} Y X^{b} Y X^{c}$$ where the second sum is over all ordered triples $(a,b,c)$ of nonnegative integers with $a+b+c = n-2$. Taking the trace, we have $\text{Tr}(X^a Y X^b Y X^c) = \text{Tr}(X^{a+c} Y X^b Y)$ so we can write $$ 2 g_2 = \sum_{n=2}^\infty \frac{1}{n!} \sum_{k=0}^{n-2} \text{Tr}(X^k Y X^{n-2-k} Y)$$ or, with $m = n-2-k$, $$ \sum_{m=0}^\infty \sum_{k=0}^\infty \frac{1}{(m+k+2)!} \text{Tr}(X^k Y X^m Y) $$ Unfortunately, I don't think that's going to help much for numerical evaluation: it's simpler just to use the Taylor series of $\exp(X+tY)$ directly.

Answer 2

The simple way to get it numerically is to use an automatic differentiation library such as JAX.

from jax.config import config

# optional: use 64-bit computation only works on startup!
config.update("jax_enable_x64", True)

from jax import grad
from jax import numpy as jnp
from jax.scipy.linalg import expm
from numpy import random as rng


def f(t, X, Y):
    return jnp.trace(expm(X + t * Y))


eps = 10**-8
t = 0.1
X = rng.randn(3, 3)
Y = rng.randn(3, 3)

Df = grad(f)  # derivative with respect to first input
D2f = grad(Df)  # second derivate with respect to first input

numerical_grad = (f(t + eps, X, Y) - f(t - eps, X, Y)) / (2 * eps)
automatic_grad = Df(t, X, Y)

print(automatic_grad - numerical_grad)

Note: the first time you evaluate Df might be much slower than subsequent calls. You can optionally use jax.jit to JIT-compile f. JAX also supports GPU-compute, which can speed up things considerably.