Repository: chainer/chainer

Issue:
backward of F.normalize is not stable
`NormalizeL2.backward` computes 0/0 if the input contains a zero vector.  PR #4190, I wrote, caused this.  Sorry.

To begin with, x/(||x|| + eps) is C^1 but not C^2 (at x=0).  The correct backward might not be a good choice.


Discussion:

Resolving patch:
diff --git a/chainer/functions/normalization/l2_normalization.py b/chainer/functions/normalization/l2_normalization.py
--- a/chainer/functions/normalization/l2_normalization.py
+++ b/chainer/functions/normalization/l2_normalization.py
@@ -7,6 +7,25 @@
 from chainer.utils import type_check
 
 
+class _SetItemZero(function_node.FunctionNode):
+
+    """Write values to mask of zero-initialized array"""
+
+    def __init__(self, mask):
+        self.mask = mask
+
+    def forward(self, inputs):
+        x, = inputs
+        xp = cuda.get_array_module(x)
+        y = xp.zeros(self.mask.shape, x.dtype)
+        y[self.mask] = x
+        return y,
+
+    def backward(self, indices, grad_outputs):
+        g, = grad_outputs
+        return g[self.mask],
+
+
 class NormalizeL2(function_node.FunctionNode):
 
     """L2 normalization"""
@@ -43,7 +62,14 @@ def backward(self, indexes, grad_outputs):
         norm = F.broadcast_to(norm, gy.shape)
 
         x_gy_reduced = F.sum((x * gy), axis=self.axis, keepdims=True)
-        x_gy_reduced /= norm_noeps
+
+        # L2 normalize with eps has continuous backward. However,
+        # the backward is not differentiable for the indices of zero vectors.
+        # To avoid nan in double backward, do not compute outside of mask.
+        mask = norm_noeps.array != 0
+        x_gy_reduced, = _SetItemZero(mask).apply((
+            x_gy_reduced[mask] / norm_noeps[mask],))
+
         x_gy_reduced = F.broadcast_to(x_gy_reduced, gy.shape)
         gx = gy * norm - x_gy_reduced * x
         gx = gx / norm ** 2


Test patch:
diff --git a/tests/chainer_tests/functions_tests/normalization_tests/test_l2_normalization.py b/tests/chainer_tests/functions_tests/normalization_tests/test_l2_normalization.py
--- a/tests/chainer_tests/functions_tests/normalization_tests/test_l2_normalization.py
+++ b/tests/chainer_tests/functions_tests/normalization_tests/test_l2_normalization.py
@@ -1,3 +1,4 @@
+import functools
 import unittest
 
 import itertools
@@ -12,7 +13,27 @@
 from chainer.testing import attr
 
 
-@testing.parameterize(*testing.product([
+def _skip_if(cond, reason):
+    """Skip test if cond(self) is True"""
+    def decorator(impl):
+        @functools.wraps(impl)
+        def wrapper(self, *args, **kwargs):
+            if cond(self):
+                raise unittest.SkipTest(reason)
+            else:
+                impl(self, *args, **kwargs)
+        return wrapper
+    return decorator
+
+
+def _is_good_param(param):
+    # Check if 'nonzero' param is valid and meaningful. On the latter point,
+    # x should contain at least a zero if 'nonzeros' param is given.
+    return param['nonzeros'] is None \
+        or param['nonzeros'] < numpy.prod(param['shape'])
+
+
+@testing.parameterize(*filter(_is_good_param, testing.product([
     [
         {'shape': (4, 15), 'axis': 1},
         {'shape': (4,), 'axis': 0},
@@ -26,14 +47,39 @@
         {'shape': (), 'axis': ()},
     ],
     [
-        {'eps': 1e-5},
-        {'eps': 1e-1},
+        # nonzeros (optional int): max number of nonzero elems in input
+        # truezero (bool): flag whether zero elems are exactly zero. If false,
+        #     randomly-chosen small values are used.
+        {'eps': 1e-5, 'nonzeros': None},
+        {'eps': 1e-1, 'nonzeros': None},
+        {'eps': 1e-1, 'nonzeros': 0, 'truezero': True},
+        {'eps': 1e-1, 'nonzeros': 0, 'truezero': False},
+        {'eps': 1e-1, 'nonzeros': 2, 'truezero': True},
+        {'eps': 1e-1, 'nonzeros': 2, 'truezero': False},
     ],
-]))
+])))
 class TestL2Normalization(unittest.TestCase):
 
     def setUp(self):
         self.x = numpy.random.uniform(-1, 1, self.shape).astype(numpy.float32)
+        if self.nonzeros is not None:
+            # Make self.x have limited number of large values
+
+            # get mask of indices to modify at
+            zeros = self.x.size - self.nonzeros
+            while True:
+                rand = numpy.random.uniform(0, 1, self.shape)
+                mask = rand <= numpy.sort(rand.ravel())[zeros - 1]
+                if self.x[mask].shape == (zeros,):
+                    break
+
+            # set zeros or small values to a part of the input
+            if self.truezero:
+                self.x[mask] = 0
+            else:
+                zero_scale = 10. ** numpy.random.randint(-40, -3)
+                self.x[mask] = numpy.random.uniform(
+                    -zero_scale, zero_scale, zeros)
         self.gy = numpy.random.uniform(-1, 1, self.shape).astype(numpy.float32)
         self.ggx = numpy.random.uniform(
             -1, 1, self.shape).astype(numpy.float32)
@@ -83,6 +129,9 @@ def test_backward_gpu(self):
         self.check_backward(
             cuda.to_gpu(self.x), self.axis, cuda.to_gpu(self.gy))
 
+    @_skip_if(
+        lambda self: self.nonzeros is not None,
+        'backward of L2Normalize is non-differentiable at zero vector')
     def check_double_backward(self, x_data, axis, y_grad, x_grad_grad):
         def f(x):
             return functions.normalize(x, eps=self.eps, axis=axis)