How does Tensorflow compute SparseCategoricalCrossentropy?

Question

The expression for categorical cross-entropy loss can be obtained via the negative log likelihood.

In particular, let $C$ be the number of mutually exclusive classes in a multi-class classification problem, $\mathcal{D}=\left\{ (\mathbf{x}_{i},\mathbf{y}_{i})\right\} _{i=1}^{n}$ be the data, where $\mathbf{y}_{i}\in\left\{ 0,1\right\} ^{C}$ is the $i$-th label represented using the “one-hot” encoding, i.e. $y_{c}=0$ for all $c\in1:C$, except $y_{c^{*}}=1$ where $c^{*}$ is the number of the correct class. Let also $p_{c}(\mathbf{x}_{i})$ be the probability of the $c$-th class on the $i$-th data vector; then $$\mathsf{P}((\mathbf{x},\mathbf{y}))=p_{c^{*}}(\mathbf{x})=\prod_{c=1}^{C}p_{c}^{y_{c}}(\mathbf{x})$$ and the likelihood of the parameter vector $\boldsymbol{\theta}$ is $$\mathsf{L}(\boldsymbol{\theta}\mid\mathcal{D})=\prod_{(\mathbf{x},\mathbf{y})\in\mathcal{D}}\prod_{c=1}^{C}p_{c}^{y_{c}}(\mathbf{x}).$$ Hence the negative log likelihood is $$-\log\mathsf{L}(\boldsymbol{\theta}\mid\mathcal{D})=-\sum_{(\mathbf{x},\mathbf{y})\in\mathcal{D}}\sum_{c=1}^{C}y_{c}\log p_{c}(\mathbf{x}).$$

However, is this really the expression computed by Tensorflow?

Answer 1

Apparently, Tensorflow computes the average of the negative log likelihood terms rather than their sum:

import tensorflow as tf

def categorical_ce(y, logit, reduce_mean=True):
    cce = \
        -tf.reduce_sum(
            tf.math.log(
                tf.nn.softmax(logit)
            ) *
            tf.one_hot(
                tf.cast(y, tf.int32), 
                logit.shape[1]
            ),
            axis=-1
        )
    if reduce_mean:
        cce = tf.reduce_mean(cce)
    return cce

A quick check bellow reveals that the result matches the one computed by SparseCategoricalCrossentropy:

>>> print(logit)
tf.Tensor(
[[-0.07976019  1.4097637   0.31711328 -0.56984025 -0.2769775   0.14477172
   0.6587809   0.03351004  0.24701062  0.55834836]
 [ 0.07477544  0.7773567  -0.09413219  0.00184567  0.0491263   0.19461593
   0.47267246  0.02392993  0.56105083  0.22566506]], shape=(2, 10), dtype=float32)
>>> print(y)
[9 0]
>>> tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)(y, logit)
<tf.Tensor: shape=(), dtype=float32, numpy=2.3179154>
>>> categorical_ce(y, logit)
<tf.Tensor: shape=(), dtype=float32, numpy=2.3179154>

As jkpate points out, this average is

an estimate of the cross-entropy of the model probability and the empirical probability in the data, which is the expected negative log probability according to the model averaged across the data.

Moreover, computing the average effectively de-couples mini-batch size and learning rate, see Mean or sum of gradients for weight updates in SGD.