Third/Fourth derivative of cross-entropy loss

Question

I need a nice formula for the third (or fourth derivative if it's easier) of cross-entropy loss $\frac{\partial^3 J}{\partial z^3}$ where

$$J(p(z)) = -\sum_i q_i\log p(z)_i$$ $$p(z)_i=\frac{\exp z_i}{\sum_i \exp z_i}$$ Can anyone suggest any Mathematica magic that can help me find it?

• First derivative is $H_1=p-q$

• Second derivative is $H_2=\text{diag}(p)-pp'$

• There's symmetric factorization $H_2=Q^TQ$ with $Q=\text{diag}(\sqrt{p})-\sqrt{p}p$

• $H_3=\text{diag}(p)-p\otimes p\otimes p$, at least when $p=1/3,1/3,1/3$

• What is the formula for higher derivatives?

I'm suspecting there's a nice formula by looking at concrete values. For instance, code below computes derivatives for $z_i=1$, counts number of unique values and shows a single matrix slice

$$ H_2=\left( \begin{array}{ccc} \frac{2}{9} & -\frac{1}{9} & -\frac{1}{9} \\ -\frac{1}{9} & \frac{2}{9} & -\frac{1}{9} \\ -\frac{1}{9} & -\frac{1}{9} & \frac{2}{9} \\ \end{array} \right)$$

$$H_3=\left( \begin{array}{ccc} \frac{2}{27} & -\frac{1}{27} & -\frac{1}{27} \\ -\frac{1}{27} & -\frac{1}{27} & \frac{2}{27} \\ -\frac{1}{27} & \frac{2}{27} & -\frac{1}{27} \\ \end{array} \right)$$

$$H_4=\left( \begin{array}{ccc} -\frac{2}{27} & \frac{1}{27} & \frac{1}{27} \\ \frac{1}{27} & -\frac{1}{27} & 0 \\ \frac{1}{27} & 0 & -\frac{1}{27} \\ \end{array} \right)$$

notebook

(* approximate equality testing *)

DotEqual[a_, b_] := 
  Norm[Flatten[{a}] - Flatten[{b}], \[Infinity]] < 1*^-9;
On[Assert];

softmax[z_] := 
  Exp[z]/Total[Exp@z]; (* make entries positive and add up to 1 *)

d = 3; (* number of dimensions *)

z = Array[z00, d]; (* vector of potentials *)

p = softmax[z];  (* vector of  probabilities *)

q = Array[q00, d];  (* target probabilities *)

(* substitution rules to replace q,z with numeric values *)
num := (
   qvals = softmax[Array[1 &, d]];
   zvals = Array[1 &, d];
   Thread[q -> qvals]~Join~Thread[z -> zvals]
   );

xent = Log[Total[Exp[z]]] Total[q] - z . q;
first = D[xent, {z, 1}] /. num;
second = D[xent, {z, 2}] /. num;
third = D[xent, {z, 3}] /. num;
fourth = D[xent, {z, 4}] /. num;
fifth = D[xent, {z, 5}] /. num;

myFirst = (p - q) /. num;

mySecond = DiagonalMatrix[p] - Outer[Times, p, p] /. num;
secondSqrt = DiagonalMatrix[Sqrt[p]] - Outer[Times, Sqrt[p], p] /. num;

Assert[first \[DotEqual] myFirst]
Assert[second \[DotEqual] mySecond]
Assert[Transpose[secondSqrt] . secondSqrt \[DotEqual] mySecond]

myThird = 
  "TODO"; (* figure out formula for third derivative and its \
factorization *)

For[order = 2, order <= 10, order += 1,
 deriv = D[xent, {z, order}] /. num;
 slice = (Composition @@ Table[First, order - 2])@deriv;
 unique = DeleteDuplicates@Sort[Flatten@deriv];
 Print[StringForm["order=``  num unique=``  `` ", order, 
   Length@unique, slice // MatrixForm]]
 ]

Answer 1

You can compute it using generalized Einstein notation. Since the formulas are not easy to read, I will instead post some code using NumPy and the einsum-function that computes the third-order derivative. It is basically a sum of diagonal tensors and outer products.

# let q, z be a numpy vectors of length n
p = np.exp(z) / np.sum(np.exp(z))
functionValue = np.dot(-q, np.log(p))
gradient = p - q
Hessian = np.diag(p) - np.einsum('i, j -> ij', p, p)

# three-dimensional identity tensor, does not exist in numpy
eye_3 = np.einsum('ij, jk -> ijk', np.eye(n), np.eye(n))
# two dimensional diagonal tensor
diag_2_p = np.diag(p)
# three-dimensional diagonal tensor
diag_3_p = np.einsum('ijk, k -> ijk', eye_3, p)

outer_p = np.einsum('i, j, k -> ijk', p, p, p)

third_order_derivative = diag_3_p \
    - np.einsum('ij, k -> ijk', diag_2_p, p) \
    - np.einsum('ik, j -> ijk', diag_2_p, p) \
    - np.einsum('jk, i -> ijk', diag_2_p, p) \
    + 2 * outer_p