Repository: chainer/chainer

Issue:
Stable SoftmaxCrossEntropy
The probabilities from the softmax function can become numerically zero. Hence,  for the forward method in the softmax_cross_entropy something like

```
numpy.log(numpy.maximum(1.0e-8,p))
```

and

```
log(fmax(1.0E-8, y[_j * n_channel + t]) ) 
```

should be added for the cpu and gpu, respectively.



Discussion:
Thank you for reporting a problem.
How about using `log_sum_exp` to avoid under flow? 

Do you think we should add `logsumexp` to CuPy? `logsumexp` is in different situation from other functions in that it is not in numpy but in spicy (`spicy.misc.logsumexp`). 

I dont see how logsumexp solves the problem


Resolving patch:
diff --git a/chainer/functions/loss/softmax_cross_entropy.py b/chainer/functions/loss/softmax_cross_entropy.py
--- a/chainer/functions/loss/softmax_cross_entropy.py
+++ b/chainer/functions/loss/softmax_cross_entropy.py
@@ -3,10 +3,23 @@
 
 from chainer import cuda
 from chainer import function
-from chainer.functions import softmax
 from chainer.utils import type_check
 
 
+def logsumexp(x):
+    xp = cuda.get_array_module(x)
+    m = x.max(axis=1, keepdims=True)
+    y = x - m
+    xp.exp(y, out=y)
+    return xp.log(y.sum(axis=1, keepdims=True)) + m
+
+
+def softmax_log(x):
+    # TODO(unno): Use cudnn (cudnn v2 doesn't support CUDNN_SOFTMAX_LOG)
+    log_z = logsumexp(x)
+    return x - log_z
+
+
 class SoftmaxCrossEntropy(function.Function):
 
     """Softmax activation followed by a cross entropy loss."""
@@ -32,11 +45,12 @@ def check_type_forward(self, in_types):
 
     def forward_cpu(self, inputs):
         x, t = inputs
-        self.y, = softmax.Softmax().forward((x,))
-        yd = numpy.rollaxis(self.y, 1)
-        yd = yd.reshape(len(yd), -1).T
+        log_y = softmax_log(x)
+        self.y = numpy.exp(log_y)
+        log_yd = numpy.rollaxis(log_y, 1)
+        log_yd = log_yd.reshape(len(log_yd), -1)
 
-        p = yd[six.moves.range(t.size), numpy.maximum(t.flat, 0)]
+        log_p = log_yd[numpy.maximum(t.flat, 0), six.moves.range(t.size)]
         # deal with the case where the SoftmaxCrossEntropy is
         # unpickled from the old version
         if getattr(self, 'normalize', True):
@@ -48,14 +62,15 @@ def forward_cpu(self, inputs):
         if count == 0:
             return numpy.zeros((), dtype=x.dtype),
 
-        y = (numpy.log(p) * (t.flat != self.ignore_label)).sum(keepdims=True) \
+        y = (log_p * (t.flat != self.ignore_label)).sum(keepdims=True) \
             * (-1.0 / count)
         return y.reshape(()),
 
     def forward_gpu(self, inputs):
         cupy = cuda.cupy
         x, t = inputs
-        self.y, = softmax.Softmax(self.use_cudnn).forward((x,))
+        log_y = softmax_log(x)
+        self.y = cupy.exp(log_y)
         if getattr(self, 'normalize', True):
             count = float((t != self.ignore_label).sum())
         else:
@@ -65,12 +80,12 @@ def forward_gpu(self, inputs):
         if count == 0:
             return cupy.zeros((), dtype=x.dtype),
 
-        y = cupy.rollaxis(self.y, 1, self.y.ndim)
+        log_y = cupy.rollaxis(log_y, 1, log_y.ndim)
         ret = cuda.reduce(
-            'S t, raw T y, int32 n_channel, T inv_count', 'T out',
-            't == -1 ? 0 : log(y[_j * n_channel + t])',
+            'S t, raw T log_y, int32 n_channel, T inv_count', 'T out',
+            't == -1 ? 0 : log_y[_j * n_channel + t]',
             'a + b', 'out = a * inv_count', '0', 'crossent_fwd'
-        )(t, y.reduced_view(), y.shape[-1], -1.0 / count)
+        )(t, log_y.reduced_view(), log_y.shape[-1], -1.0 / count)
         return ret,
 
     def backward_cpu(self, inputs, grad_outputs):


Test patch:
diff --git a/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py b/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py
--- a/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py
+++ b/tests/chainer_tests/functions_tests/activation_tests/test_softmax.py
@@ -23,9 +23,12 @@ def check_forward(self, x_data, use_cudnn=True):
         y = functions.softmax(x, use_cudnn)
         self.assertEqual(y.data.dtype, numpy.float32)
 
-        y_expect = numpy.exp(self.x)
+        y_expect = numpy.empty_like(self.x)
         for i in six.moves.range(y_expect.shape[0]):
-            y_expect[i] /= y_expect[i].sum()
+            x = self.x[i]
+            log_z = numpy.ufunc.reduce(numpy.logaddexp, x)
+            x -= log_z
+            y_expect[i] = numpy.exp(x)
 
         gradient_check.assert_allclose(y_expect, y.data)
 
@@ -70,6 +73,13 @@ def test_backward_gpu_no_cudnn(self):
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.gy), False)
 
 
+class TestSoftmaxUnstable(TestSoftmax):
+
+    def setUp(self):
+        self.x = numpy.array([[-1000, 1]], dtype=numpy.float32)
+        self.gy = numpy.random.uniform(-1, 1, (1, 2)).astype(numpy.float32)
+
+
 class TestReplicatedSoftmax1(TestSoftmax):
 
     def setUp(self):
diff --git a/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py b/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py
--- a/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py
+++ b/tests/chainer_tests/functions_tests/loss_tests/test_softmax_cross_entropy.py
@@ -15,6 +15,8 @@
 
 class TestSoftmaxCrossEntropy(unittest.TestCase):
 
+    backward_atol = 1e-4
+
     def setUp(self):
         self.x = numpy.random.uniform(-1, 1, (4, 3)).astype(numpy.float32)
         self.t = numpy.random.randint(0, 3, (4,)).astype(numpy.int32)
@@ -28,13 +30,13 @@ def check_forward(self, x_data, t_data, use_cudnn=True):
         loss_value = float(cuda.to_cpu(loss.data))
 
         # Compute expected value
-        y = numpy.exp(self.x)
         loss_expect = 0.0
         count = 0
-        for i in six.moves.range(y.shape[0]):
+        for i in six.moves.range(self.x.shape[0]):
             if self.t[i] == -1:
                 continue
-            loss_expect -= math.log(y[i, self.t[i]] / y[i].sum())
+            log_z = numpy.ufunc.reduce(numpy.logaddexp, self.x[i])
+            loss_expect -= (self.x[i] - log_z)[self.t[i]]
             count += 1
 
         if count == 0:
@@ -69,7 +71,7 @@ def check_backward(self, x_data, t_data, use_cudnn=True):
         f = lambda: func.forward((x.data, t.data))
         gx, = gradient_check.numerical_grad(f, (x.data,), (1,), eps=0.02)
 
-        gradient_check.assert_allclose(gx, x.grad, atol=1e-4)
+        gradient_check.assert_allclose(gx, x.grad, atol=self.backward_atol)
 
     @condition.retry(3)
     def test_backward_cpu(self):
@@ -86,6 +88,15 @@ def test_backward_gpu_no_cudnn(self):
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.t), False)
 
 
+class TestSoftmaxCrossEntropyUnstable(TestSoftmaxCrossEntropy):
+
+    backward_atol = 1e-3
+
+    def setUp(self):
+        self.x = numpy.array([[-1000, 1]], dtype=numpy.float32)
+        self.t = numpy.array([0], dtype=numpy.int32)
+
+
 class TestReplicatedSoftmaxCrossEntropy1(TestSoftmaxCrossEntropy):
 
     def setUp(self):