Third/Fourth derivative of cross-entropy loss

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];
);

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]]
]

• It might help a bit to observe that xent = Log[Total[Exp[z]]] Total[q] - z . q. Now Total[Exp[z]] is just a scalar funtion (and not in the denominator). Jan 14, 2022 at 8:11
• good point, that simplifies derivations, updated Jan 14, 2022 at 14:13
• actually, Total[q]=1, so this turns out to be equivalent to derivatives of the log-partition function $J(z)=\log \sum_i \exp z_i$ Jan 15, 2022 at 20:31
• I'm not sure I'm correctly interpreting the question, but here I propose an alternative method: wolframcloud.com/obj/b50c9dd5-bfec-4651-940d-ccf640238be3 . Is this computation reproducing what you expect?
– jose
Jan 25, 2022 at 1:39
• That seems correct, but this expression is not human readable. Just compare your approach for obtaining the second derivative....it's just diagonal - rank1 matrix Jan 25, 2022 at 5:42

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))
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

• This does not appear to be Mathematica code. Jan 16, 2022 at 15:31
• Thanks, this is very useful! Do you have a reference or did you derive this yourself? How does this extend to higher derivatives? Jan 16, 2022 at 18:25
• @bbgodfrey technically it isn't, but it's 80% of the way towards an answer, the other 20% can leverage Mathematica einsum implementation mentioned here mathematica.stackexchange.com/questions/261720/… Jan 16, 2022 at 18:27