Have I understood the loss function from the original U-Net paper correctly?

Question

In the original U-Net paper, it is written

The energy function is computed by a pixel-wise soft-max over the final feature map combined with the cross entropy loss function.

...

$$ E=\sum_{\mathbf{x} \in \Omega} w(\mathbf{x}) \log \left(p_{\ell(\mathbf{x})}(\mathbf{x})\right) \tag{1}\label{1} $$

where $w(\mathbf{x})$ is a weight map (I'm not interested in that part right now), and $p_{k}(\mathbf{x})$ is

$$ p_{k}(\mathbf{x})=\exp \left(a_{k}(\mathbf{x})\right) /\left(\sum_{k^{\prime}=1}^{K} \exp \left(a_{k^{\prime}}(\mathbf{x})\right)\right) $$

The pixel-wise softmax with $a_{k}(\mathbf{x})$ being the activation in feature channel $k$ at pixel position $\mathbf{x}$ and $K$ the number of classes. Then $\ell(\mathbf{x})$ from $p_{\ell(\mathbf{x})}$ is the true label of each pixel, i.e. if the pixel at position $\mathbf{x}$ is part of class $1$, then $p_{\ell(\mathbf{x})}$ is equal to $p_1(\mathbf{x})$.

As far as is understand $-E$ should be the cross-entropy function. Right? I've already done the math for the binary case (ignoring $w(\mathbf{x})$) and it seemed to be equal.

Answer 1

Yes, $E$ is the cross-entropy function and a direct generalization of the binary case.

For the binary case, probability to belong to the class $1$ is given by a sigmoid function $\sigma(x)$ of the output $x$, and the probability to belong to the class $0$ is $1 - \sigma(x)$.

Therefore the binary crossentropy will give: $$ -\sum_i (l_i \log \sigma(x) + (1 - l_i) \log (1 - \sigma(x)) $$ Where the sum is over all samples in the dataset. Because $l_i$ is a binary variable one of these two terms would be zero. The gradient descent forces the model to predict true label with more confidence.

For the multiclass case, now the output is a $K$-vector and the softmax function forces its elements to sum to one: $$ \sum_{i = 1}^{K} p_i = 0 $$ The true label is one-hot encoded vector, with $1$ on the position of the true label, and $0$ elsewhere. The generalization of the binary case is: $$ -\sum_i \sum_{j = 1}^{K}(l_{ij} \log \sigma_j(x) + (1 - l_{ij}) \log (1 - \sigma_j(x)) $$ In this case, there will be $K$ non-vanishing contributions.The minization of $E$ forces the classifier to predict the true label more confidently, and all other (to make $\log(1 - \sigma_j(x))$) as small as possible.