Repository: chainer/chainer

Issue:
Decorrelated batch normalization does not group channels correctly
### Code to reproduce

The following assertion fails.
```python
import chainer.functions as F
import numpy as np
np.random.seed(0)
xs = [np.random.randn(4, 3) for _ in range(2)]
y_foreach = np.concatenate(
    [F.decorrelated_batch_normalization(x, groups=1).array for x in xs],
    axis=1
)
y_groups = F.decorrelated_batch_normalization(
    np.concatenate(xs, axis=1),
    groups=2
).array
np.testing.assert_allclose(y_foreach, y_groups)
```
(batchsize = 4 and channels = 6 = 2 * 3)

Moreover, `L.DecorrelatedBatchNormalization` should maintain an array of the shape `(groups, channels//groups, channels//groups) == (2, 3, 3)` for the average projections.

### Error messages, stack traces, or logs

```
Traceback (most recent call last):
  File "new_issue.py", line 13, in <module>
    np.testing.assert_allclose(y_foreach, y_groups)
  File "/Users/tos/.pyenv/versions/n/lib/python3.7/site-packages/numpy/testing/_private/utils.py", line 1501, in assert_allclose
    verbose=verbose, header=header, equal_nan=equal_nan)
  File "/Users/tos/.pyenv/versions/n/lib/python3.7/site-packages/numpy/testing/_private/utils.py", line 827, in assert_array_compare
    raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=1e-07, atol=0

Mismatch: 100%
Max absolute difference: 2.05411278
Max relative difference: 14.27449499
 x: array([[ 1.354603, -0.458379,  0.976996,  0.244642, -0.951322,  1.426518],
       [ 0.419707,  1.380985, -0.957405, -1.235305,  1.176817,  0.295545],
       [-0.43201 , -1.314752, -1.041418, -0.483977, -1.03046 , -1.305253],
       [-1.3423  ,  0.392146,  1.021826,  1.47464 ,  0.804965, -0.41681 ]])
 y: array([[ 1.386369, -0.478125,  0.997635, -0.169201, -0.451325,  0.587504],
       [ 1.713926,  1.378891, -0.770033, -1.363784,  1.54728 , -0.022264],
       [ 0.265987, -0.853668,  0.12491 , -0.455366, -1.42594 , -1.920292],
       [-0.798459, -0.343282,  1.465721, -0.579473,  0.62617 , -0.463181]])
```


Discussion:

Resolving patch:
diff --git a/chainer/functions/normalization/decorrelated_batch_normalization.py b/chainer/functions/normalization/decorrelated_batch_normalization.py
--- a/chainer/functions/normalization/decorrelated_batch_normalization.py
+++ b/chainer/functions/normalization/decorrelated_batch_normalization.py
@@ -1,11 +1,52 @@
+import numpy
+
 from chainer import backend
 from chainer import function_node
 from chainer.utils import argument
 from chainer.utils import type_check
 
 
+# {numpy: True, cupy: False}
+_xp_supports_batch_eigh = {}
+
+
+# routines for batched matrices
+def _eigh(a, xp):
+    if xp not in _xp_supports_batch_eigh:
+        try:
+            xp.linalg.eigh(xp.ones((2, 2, 2), xp.float32))
+        except ValueError:
+            _xp_supports_batch_eigh[xp] = False
+        else:
+            _xp_supports_batch_eigh[xp] = True
+    if _xp_supports_batch_eigh[xp]:
+        return xp.linalg.eigh(a)
+    ws = []
+    vs = []
+    for ai in a:
+        w, v = xp.linalg.eigh(ai)
+        ws.append(w)
+        vs.append(v)
+    return xp.stack(ws), xp.stack(vs)
+
+
+def _matmul(a, b, xp):
+    if hasattr(xp, 'matmul'):  # numpy.matmul is supported from version 1.10.0
+        return xp.matmul(a, b)
+    else:
+        return xp.einsum('bij,bjk->bik', a, b)
+
+
+def _diag(a, xp):
+    s0, s1 = a.shape
+    ret = xp.zeros((s0, s1, s1), a.dtype)
+    arange_s1 = numpy.arange(s1)
+    ret[:, arange_s1, arange_s1] = a
+    return ret
+
+
 def _calc_axis_and_m(x_shape, batch_size, groups):
-    m = batch_size * groups
+    m = batch_size
     spatial_ndim = len(x_shape) - 2
     spatial_axis = tuple(range(2, 2 + spatial_ndim))
     for i in spatial_axis:
@@ -47,34 +88,38 @@ def forward(self, inputs):
         C = c // g
         spatial_axis, m = _calc_axis_and_m(x_shape, b, g)
 
-        if g > 1:
-            x = x.reshape((b * g, C) + x.shape[2:])
-        x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
-
-        mean = x_hat.mean(axis=1)
-        x_hat = x_hat - mean[:, None]
+        # (g, C, m)
+        x_hat = x.transpose((1, 0) + spatial_axis).reshape(g, C, m)
+        mean = x_hat.mean(axis=2, keepdims=True)
+        x_hat = x_hat - mean
         self.eps = x.dtype.type(self.eps)
 
         eps_matrix = self.eps * xp.eye(C, dtype=x.dtype)
-        cov = x_hat.dot(x_hat.T) / x.dtype.type(m) + eps_matrix
-        self.eigvals, self.eigvectors = xp.linalg.eigh(cov)
-        U = xp.diag(self.eigvals ** -0.5).dot(self.eigvectors.T)
-        self.y_hat_pca = U.dot(x_hat)  # PCA whitening
-        y_hat = self.eigvectors.dot(self.y_hat_pca)  # ZCA whitening
-
-        y = y_hat.reshape((C, b * g,) + x.shape[2:]).transpose(
+        cov = _matmul(
+            x_hat, x_hat.transpose(0, 2, 1),
+            xp) / x.dtype.type(m) + eps_matrix
+        # (g, C), (g, C, C)
+        self.eigvals, self.eigvectors = _eigh(cov, xp)
+        U = _matmul(
+            _diag(self.eigvals ** -0.5, xp),
+            self.eigvectors.transpose(0, 2, 1),
+            xp)
+        self.y_hat_pca = _matmul(U, x_hat, xp)  # PCA whitening
+        # ZCA whitening
+        y_hat = _matmul(self.eigvectors, self.y_hat_pca, xp)
+
+        y = y_hat.reshape((c, b) + x_shape[2:]).transpose(
             (1, 0) + spatial_axis)
-        if self.groups > 1:
-            y = y.reshape((-1, c) + x.shape[2:])
 
         # Update running statistics
         if self.running_mean is not None:
+            mean = mean.squeeze(axis=2)
             self.running_mean *= self.decay
             self.running_mean += (1 - self.decay) * mean
         if self.running_projection is not None:
             adjust = m / max(m - 1., 1.)  # unbiased estimation
             self.running_projection *= self.decay
-            projection = self.eigvectors.dot(U)
+            projection = _matmul(self.eigvectors, U, xp)
             self.running_projection += (1 - self.decay) * adjust * projection
 
         return y,
@@ -104,38 +149,45 @@ def forward(self, inputs):
         g = self.groups
         C = c // g
         spatial_axis, m = _calc_axis_and_m(gy_shape, b, g)
+        arange_C = numpy.arange(C)
+        diag_indices = slice(None), arange_C, arange_C
 
-        if g > 1:
-            gy = gy.reshape((b * g, C) + gy.shape[2:])
-        gy_hat = gy.transpose((1, 0) + spatial_axis).reshape(C, -1)
+        gy_hat = gy.transpose((1, 0) + spatial_axis).reshape(g, C, m)
 
         eigvectors = self.eigvectors
         eigvals = self.eigvals
         y_hat_pca = self.y_hat_pca
-        gy_hat_pca = eigvectors.T.dot(gy_hat)
-        f = gy_hat_pca.mean(axis=1)
-
-        K = eigvals[:, None] - eigvals[None, :]
-        valid = K != 0
-        K[valid] = 1 / K[valid]
-        xp.fill_diagonal(K, 0)
-
-        V = xp.diag(eigvals)
-        V_sqrt = xp.diag(eigvals ** 0.5)
-        V_invsqrt = xp.diag(eigvals ** -0.5)
-
-        F_c = gy_hat_pca.dot(y_hat_pca.T) / gy.dtype.type(m)
-        M = xp.diag(xp.diag(F_c))
-
-        mat = K.T * (V.dot(F_c.T) + V_sqrt.dot(F_c).dot(V_sqrt))
-        S = mat + mat.T
-        R = gy_hat_pca - f[:, None] + (S - M).T.dot(y_hat_pca)
-        gx_hat = R.T.dot(V_invsqrt).dot(eigvectors.T).T
-
-        gx = gx_hat.reshape((C, b * g,) + gy.shape[2:]).transpose(
+        gy_hat_pca = _matmul(eigvectors.transpose(0, 2, 1), gy_hat, xp)
+        f = gy_hat_pca.mean(axis=2, keepdims=True)
+
+        K = eigvals[:, :, None] - eigvals[:, None, :]
+        valid = K != 0  # to avoid nan, use eig_i != eig_j instead of i != j
+        K[valid] = xp.reciprocal(K[valid])
+
+        V = _diag(eigvals, xp)
+        V_sqrt = _diag(eigvals ** 0.5, xp)
+        V_invsqrt = _diag(eigvals ** -0.5, xp)
+
+        F_c = _matmul(
+            gy_hat_pca, y_hat_pca.transpose(0, 2, 1),
+            xp) / gy.dtype.type(m)
+        M = xp.zeros_like(F_c)
+        M[diag_indices] = F_c[diag_indices]
+
+        mat = K.transpose(0, 2, 1) * (
+            _matmul(V, F_c.transpose(0, 2, 1), xp)
+            + _matmul(_matmul(V_sqrt, F_c, xp), V_sqrt, xp)
+        )
+        S = mat + mat.transpose(0, 2, 1)
+        R = gy_hat_pca - f + _matmul(
+            (S - M).transpose(0, 2, 1), y_hat_pca, xp)
+        gx_hat = _matmul(
+            _matmul(R.transpose(0, 2, 1), V_invsqrt, xp),
+            eigvectors.transpose(0, 2, 1), xp
+        ).transpose(0, 2, 1)
+
+        gx = gx_hat.reshape((c, b) + gy_shape[2:]).transpose(
             (1, 0) + spatial_axis)
-        if g > 1:
-            gx = gx.reshape((-1, c, ) + gy.shape[2:])
 
         self.retain_outputs(())
         return gx,
@@ -166,22 +218,20 @@ def check_type_forward(self, in_types):
     def forward(self, inputs):
         self.retain_inputs((0, 1, 2))
         x, mean, projection = inputs
+        xp = backend.get_array_module(x)
         x_shape = x.shape
         b, c = x_shape[:2]
         g = self.groups
         C = c // g
         spatial_axis, m = _calc_axis_and_m(x_shape, b, g)
 
-        if g > 1:
-            x = x.reshape((b * g, C) + x.shape[2:])
-        x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
+        x_hat = x.transpose((1, 0) + spatial_axis).reshape(g, C, m)
+        x_hat = x_hat - xp.expand_dims(mean, axis=2)
 
-        y_hat = projection.dot(x_hat - mean[:, None])
+        y_hat = _matmul(projection, x_hat, xp)
 
-        y = y_hat.reshape((C, b * g) + x.shape[2:]).transpose(
+        y = y_hat.reshape((c, b) + x_shape[2:]).transpose(
             (1, 0) + spatial_axis)
-        if g > 1:
-            y = y.reshape((-1, c) + x.shape[2:])
 
         return y,
 
@@ -200,25 +250,21 @@ def __init__(self, groups):
     def forward(self, inputs):
         self.retain_inputs(())
         x, mean, projection, gy = inputs
+        xp = backend.get_array_module(x)
         gy_shape = gy.shape
         b, c = gy_shape[:2]
         g = self.groups
         C = c // g
         spatial_axis, m = _calc_axis_and_m(gy_shape, b, g)
 
-        if g > 1:
-            gy = gy.reshape((b * g, C) + gy.shape[2:])
-            x = x.reshape((b * g, C) + x.shape[2:])
-        x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
-        gy_hat = gy.transpose((1, 0) + spatial_axis).reshape(C, -1)
-        gy_hat_pca = projection.T.dot(gy_hat)
-        gx = gy_hat_pca.reshape(
-            (C, b * g) + gy.shape[2:]).transpose((1, 0) + spatial_axis)
-        if g > 1:
-            gx = gx.reshape((-1, c) + gy.shape[2:])
-        rhs = x_hat - mean[Ellipsis, None]
-        gprojection = (x_hat - rhs).T.dot(gy_hat)
-        gmean = -gx[:, 0]
+        gy_hat = gy.transpose((1, 0) + spatial_axis).reshape(g, C, m)
+        x_hat = x.transpose((1, 0) + spatial_axis).reshape(g, C, m)
+        gy_hat_pca = _matmul(projection.transpose(0, 2, 1), gy_hat, xp)
+        gx = gy_hat_pca.reshape((c, b) + gy_shape[2:]).transpose(
+            (1, 0) + spatial_axis)
+        rhs = x_hat - xp.expand_dims(mean, axis=2)
+        gprojection = _matmul((x_hat - rhs).transpose(0, 2, 1), gy_hat, xp)
+        gmean = -gy_hat_pca[..., 0]
         self.retain_outputs(())
         return gx, gmean, gprojection
 
diff --git a/chainer/links/normalization/decorrelated_batch_normalization.py b/chainer/links/normalization/decorrelated_batch_normalization.py
--- a/chainer/links/normalization/decorrelated_batch_normalization.py
+++ b/chainer/links/normalization/decorrelated_batch_normalization.py
@@ -1,8 +1,13 @@
+import functools
+import warnings
+
 import numpy
 
+import chainer
 from chainer import configuration
 from chainer import functions
 from chainer import link
+import chainer.serializer as serializer_mod
 from chainer.utils import argument
 
 
@@ -62,9 +67,13 @@ class DecorrelatedBatchNormalization(link.Link):
     def __init__(self, size, groups=16, decay=0.9, eps=2e-5,
                  dtype=numpy.float32):
         super(DecorrelatedBatchNormalization, self).__init__()
-        self.avg_mean = numpy.zeros(size // groups, dtype=dtype)
+        C = size // groups
+        self.avg_mean = numpy.zeros((groups, C), dtype=dtype)
         self.register_persistent('avg_mean')
-        self.avg_projection = numpy.eye(size // groups, dtype=dtype)
+        avg_projection = numpy.zeros((groups, C, C), dtype=dtype)
+        arange_C = numpy.arange(C)
+        avg_projection[:, arange_C, arange_C] = 1
+        self.avg_projection = avg_projection
         self.register_persistent('avg_projection')
         self.N = 0
         self.register_persistent('N')
@@ -72,6 +81,16 @@ def __init__(self, size, groups=16, decay=0.9, eps=2e-5,
         self.eps = eps
         self.groups = groups
 
+    def serialize(self, serializer):
+        if isinstance(serializer, serializer_mod.Deserializer):
+            serializer = _PatchedDeserializer(serializer, {
+                'avg_mean': functools.partial(
+                    fix_avg_mean, groups=self.groups),
+                'avg_projection': functools.partial(
+                    fix_avg_projection, groups=self.groups),
+            })
+        super(DecorrelatedBatchNormalization, self).serialize(serializer)
+
     def forward(self, x, **kwargs):
         """forward(self, x, *, finetune=False)
 
@@ -131,3 +150,62 @@ def start_finetuning(self):
 
         """
         self.N = 0
+
+
+class _PatchedDeserializer(serializer_mod.Deserializer):
+
+    def __init__(self, base, patches):
+        self.base = base
+        self.patches = patches
+
+    def __repr__(self):
+        return '_PatchedDeserializer({}, {})'.format(
+            repr(self.base), repr(self.patches))
+
+    def __call__(self, key, value):
+        if key not in self.patches:
+            return self.base(key, value)
+        arr = self.base(key, None)
+        arr = self.patches[key](arr)
+        if value is None:
+            return arr
+        chainer.backend.copyto(value, arr)
+        return value
+
+
+def _warn_old_model():
+    msg = (
+        'Found moving statistics of old DecorrelatedBatchNormalization, whose '
+        'algorithm was different from the paper.')
+    warnings.warn(msg)
+
+
+def fix_avg_mean(avg_mean, groups):
+    if avg_mean.ndim == 2:  # OK
+        return avg_mean
+    elif avg_mean.ndim == 1:  # Issue #7706
+        if groups != 1:
+            _warn_old_model()
+        return _broadcast_to(avg_mean, (groups,) + avg_mean.shape)
+    raise ValueError('unexpected shape of avg_mean')
+
+
+def fix_avg_projection(avg_projection, groups):
+    if avg_projection.ndim == 3:  # OK
+        return avg_projection
+    elif avg_projection.ndim == 2:  # Issue #7706
+        if groups != 1:
+            _warn_old_model()
+        return _broadcast_to(
+            avg_projection, (groups,) + avg_projection.shape)
+    raise ValueError('unexpected shape of avg_projection')
+
+
+def _broadcast_to(array, shape):
+    if hasattr(numpy, 'broadcast_to'):
+        return numpy.broadcast_to(array, shape)
+    else:
+        # numpy 1.9 doesn't support broadcast_to method
+        dummy = numpy.empty(shape)
+        bx, _ = numpy.broadcast_arrays(array, dummy)
+        return bx


Test patch:
diff --git a/tests/chainer_tests/functions_tests/normalization_tests/test_decorrelated_batch_normalization.py b/tests/chainer_tests/functions_tests/normalization_tests/test_decorrelated_batch_normalization.py
--- a/tests/chainer_tests/functions_tests/normalization_tests/test_decorrelated_batch_normalization.py
+++ b/tests/chainer_tests/functions_tests/normalization_tests/test_decorrelated_batch_normalization.py
@@ -5,31 +5,43 @@
 
 
 def _decorrelated_batch_normalization(x, mean, projection, groups):
+    xs = numpy.split(x, groups, axis=1)
+    assert mean.shape[0] == groups
+    assert projection.shape[0] == groups
+    ys = [
+        _decorrelated_batch_normalization_1group(xi, m, p)
+        for (xi, m, p) in zip(xs, mean, projection)]
+    return numpy.concatenate(ys, axis=1)
+
+
+def _decorrelated_batch_normalization_1group(x, mean, projection):
     spatial_ndim = len(x.shape[2:])
     spatial_axis = tuple(range(2, 2 + spatial_ndim))
-    b, c = x.shape[:2]
-    g = groups
-    C = c // g
-    x = x.reshape((b * g, C) + x.shape[2:])
+    b, C = x.shape[:2]
     x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
     y_hat = projection.dot(x_hat - mean[:, None])
-    y = y_hat.reshape((C, b * g) + x.shape[2:]).transpose(
+    y = y_hat.reshape((C, b) + x.shape[2:]).transpose(
         (1, 0) + spatial_axis)
-    y = y.reshape((-1, c) + x.shape[2:])
     return y
 
 
 def _calc_projection(x, mean, eps, groups):
+    xs = numpy.split(x, groups, axis=1)
+    assert mean.shape[0] == groups
+    projections = [
+        _calc_projection_1group(xi, m, eps)
+        for (xi, m) in zip(xs, mean)]
+    return numpy.concatenate([p[None] for p in projections])
+
+
+def _calc_projection_1group(x, mean, eps):
     spatial_ndim = len(x.shape[2:])
     spatial_axis = tuple(range(2, 2 + spatial_ndim))
-    b, c = x.shape[:2]
-    g = groups
-    C = c // g
-    m = b * g
+    b, C = x.shape[:2]
+    m = b
     for i in spatial_axis:
         m *= x.shape[i]
 
-    x = x.reshape((b * g, C) + x.shape[2:])
     x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
     mean = x_hat.mean(axis=1)
     x_hat = x_hat - mean[:, None]
@@ -39,6 +51,11 @@ def _calc_projection(x, mean, eps, groups):
     return projection
 
 
+def _calc_mean(x, groups):
+    axis = (0,) + tuple(range(2, x.ndim))
+    return x.mean(axis=axis).reshape(groups, -1)
+
+
 @testing.parameterize(*(testing.product({
     'n_channels': [8],
     'ndim': [0, 2],
@@ -74,9 +91,8 @@ class TestDecorrelatedBatchNormalization(testing.FunctionTestCase):
     skip_double_backward_test = True
 
     def setUp(self):
-        self.decay = 0.9
-        check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
-        check_backward_options = {'atol': 1e-4, 'rtol': 1e-3}
+        check_forward_options = {'atol': 1e-3, 'rtol': 1e-3}
+        check_backward_options = {'atol': 1e-3, 'rtol': 1e-3}
         if self.dtype == numpy.float16:
             check_forward_options = {'atol': 1e-2, 'rtol': 1e-2}
             check_backward_options = {'atol': 1e-2, 'rtol': 1e-2}
@@ -99,16 +115,11 @@ def forward(self, inputs, device):
 
     def forward_expected(self, inputs):
         x, = inputs
-        C = self.n_channels // self.groups
-        head_ndim = 2
-        spatial_axis = tuple(range(head_ndim, x.ndim))
-        x_hat = x.reshape((5 * self.groups, C) + x.shape[2:])
-        x_hat = x_hat.transpose((1, 0) + spatial_axis).reshape(C, -1)
-        mean = x_hat.mean(axis=1)
-        projection = _calc_projection(x, mean, self.eps, self.groups)
-
+        groups = self.groups
+        mean = _calc_mean(x, groups)
+        projection = _calc_projection(x, mean, self.eps, groups)
         return _decorrelated_batch_normalization(
-            x, mean, projection, self.groups),
+            x, mean, projection, groups),
 
 
 @testing.parameterize(*(testing.product({
@@ -147,8 +158,10 @@ class TestFixedDecorrelatedBatchNormalization(testing.FunctionTestCase):
     def setUp(self):
         C = self.n_channels // self.groups
         dtype = self.dtype
-        self.mean = numpy.random.uniform(-1, 1, C).astype(dtype)
-        self.projection = numpy.random.uniform(0.5, 1, (C, C)).astype(dtype)
+        self.mean = numpy.random.uniform(
+            -1, 1, (self.groups, C)).astype(dtype)
+        self.projection = numpy.random.uniform(
+            0.5, 1, (self.groups, C, C)).astype(dtype)
 
         check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
         check_backward_options = {'atol': 1e-4, 'rtol': 1e-3}
diff --git a/tests/chainer_tests/links_tests/normalization_tests/test_decorrelated_batch_normalization.py b/tests/chainer_tests/links_tests/normalization_tests/test_decorrelated_batch_normalization.py
--- a/tests/chainer_tests/links_tests/normalization_tests/test_decorrelated_batch_normalization.py
+++ b/tests/chainer_tests/links_tests/normalization_tests/test_decorrelated_batch_normalization.py
@@ -1,3 +1,6 @@
+import contextlib
+import os
+import tempfile
 import unittest
 
 import numpy
@@ -11,45 +14,57 @@
 
 
 def _decorrelated_batch_normalization(x, mean, projection, groups):
-    spatial_ndim = len(x.shape[2:])
-    spatial_axis = tuple(range(2, 2 + spatial_ndim))
-    b, c = x.shape[:2]
-    g = groups
-    C = c // g
+    xs = numpy.split(x, groups, axis=1)
+    assert mean.shape[0] == groups
+    assert projection.shape[0] == groups
+    ys = [
+        _decorrelated_batch_normalization_1group(xi, m, p)
+        for (xi, m, p) in zip(xs, mean, projection)]
+    return numpy.concatenate(ys, axis=1)
 
-    x = x.reshape((b * g, C) + x.shape[2:])
-    x_hat = x.transpose((1, 0) + spatial_axis).reshape((C, -1))
 
+def _decorrelated_batch_normalization_1group(x, mean, projection):
+    spatial_ndim = len(x.shape[2:])
+    spatial_axis = tuple(range(2, 2 + spatial_ndim))
+    b, C = x.shape[:2]
+    x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
     y_hat = projection.dot(x_hat - mean[:, None])
-
-    y = y_hat.reshape((C, b * g) + x.shape[2:]).transpose(
+    y = y_hat.reshape((C, b) + x.shape[2:]).transpose(
         (1, 0) + spatial_axis)
-    y = y.reshape((-1, c) + x.shape[2:])
     return y
 
 
 def _calc_projection(x, mean, eps, groups):
+    xs = numpy.split(x, groups, axis=1)
+    assert mean.shape[0] == groups
+    projections = [
+        _calc_projection_1group(xi, m, eps)
+        for (xi, m) in zip(xs, mean)]
+    return numpy.concatenate([p[None] for p in projections])
+
+
+def _calc_projection_1group(x, mean, eps):
     spatial_ndim = len(x.shape[2:])
     spatial_axis = tuple(range(2, 2 + spatial_ndim))
-    b, c = x.shape[:2]
-    g = groups
-    C = c // g
-    m = b * g
+    b, C = x.shape[:2]
+    m = b
     for i in spatial_axis:
         m *= x.shape[i]
 
-    x = x.reshape((b * g, C) + x.shape[2:])
-    x_hat = x.transpose((1, 0) + spatial_axis).reshape((C, -1))
-
+    x_hat = x.transpose((1, 0) + spatial_axis).reshape(C, -1)
     mean = x_hat.mean(axis=1)
     x_hat = x_hat - mean[:, None]
-
     cov = x_hat.dot(x_hat.T) / m + eps * numpy.eye(C, dtype=x.dtype)
     eigvals, eigvectors = numpy.linalg.eigh(cov)
     projection = eigvectors.dot(numpy.diag(eigvals ** -0.5)).dot(eigvectors.T)
     return projection
 
 
+def _calc_mean(x, groups):
+    axis = (0,) + tuple(range(2, x.ndim))
+    return x.mean(axis=axis).reshape(groups, -1)
+
+
 @testing.parameterize(*(testing.product({
     'n_channels': [8],
     'groups': [1, 2],
@@ -62,7 +77,6 @@ class DecorrelatedBatchNormalizationTest(unittest.TestCase):
 
     def setUp(self):
         C = self.n_channels // self.groups
-        head_ndim = 2
 
         self.link = links.DecorrelatedBatchNormalization(
             self.n_channels, groups=self.groups, dtype=self.dtype)
@@ -73,21 +87,19 @@ def setUp(self):
         self.gy = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
 
         if self.test:
-            self.mean = numpy.random.uniform(-1, 1, (C,)).astype(self.dtype)
-            self.projection = numpy.random.uniform(0.5, 1, (C, C)).astype(
+            self.mean = numpy.random.uniform(
+                -1, 1, (self.groups, C)).astype(self.dtype)
+            self.projection = numpy.random.uniform(
+                0.5, 1, (self.groups, C, C)).astype(
                 self.dtype)
             self.link.avg_mean[...] = self.mean
             self.link.avg_projection[...] = self.projection
         else:
-            spatial_axis = tuple(range(head_ndim, self.x.ndim))
-            x_hat = self.x.reshape((5 * self.groups, C) + self.x.shape[2:])
-            x_hat = x_hat.transpose((1, 0) + spatial_axis).reshape((C, -1))
-            self.mean = x_hat.mean(axis=1)
+            self.mean = _calc_mean(self.x, self.groups)
             self.projection = _calc_projection(self.x, self.mean,
                                                self.link.eps, self.groups)
         self.check_forward_options = {'atol': 1e-4, 'rtol': 1e-3}
-        self.check_backward_options = {
-            'atol': 5e-3, 'rtol': 1e-3, 'dtype': numpy.float64}
+        self.check_backward_options = {'atol': 5e-3, 'rtol': 1e-3}
         if self.dtype == numpy.float32:
             self.check_backward_options = {'atol': 5e-2, 'rtol': 5e-2}
 
@@ -114,7 +126,7 @@ def test_forward_gpu(self):
     def check_backward(self, x_data, y_grad):
         gradient_check.check_backward(
             self.link, x_data, y_grad, (),
-            eps=1e-2, **self.check_backward_options)
+            dtype=numpy.float64, **self.check_backward_options)
 
     def test_backward_cpu(self):
         self.check_backward(self.x, self.gy)
@@ -125,4 +137,76 @@ def test_backward_gpu(self):
         self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.gy))
 
 
+# TODO(kataoka) Use `contextlib.nullcontext` if Python 3.7 or higher is assumed
+@contextlib.contextmanager
+def nullcontext():
+    yield
+
+
+@testing.parameterize(*(testing.product({
+    'n_channels': [8],
+    'groups': [1, 2],
+    'dtype': [numpy.float32, numpy.float64],
+})))
+@testing.inject_backend_tests([
+    'test_model_compatibility_npz', 'test_model_compatibility_hdf5',
+], [
+    {},
+    {'use_cuda': True},
+    {'use_cuda': True, 'cuda_device': 1},
+    {'use_chainerx': True, 'chainerx_device': 'native:0'},
+    {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
+    {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
+])
+class TestDecorrelatedBatchNormalizationCompat(unittest.TestCase):
+
+    def setUp(self):
+        fd, path = tempfile.mkstemp()
+        os.close(fd)
+        self.temp_file_path = path
+
+    def tearDown(self):
+        if hasattr(self, 'temp_file_path'):
+            os.remove(self.temp_file_path)
+
+    def check_model_compatibility(self, backend_config, save, load):
+        C = self.n_channels // self.groups
+        old_model = {
+            'avg_mean': numpy.random.uniform(
+                -1, 1, (C,)).astype(self.dtype),
+            'avg_projection': numpy.random.uniform(
+                0.5, 1, (C, C)).astype(self.dtype),
+            'N': numpy.array(0)
+        }
+        save(self.temp_file_path, old_model)
+
+        model = links.DecorrelatedBatchNormalization(
+            self.n_channels, groups=self.groups, dtype=self.dtype)
+        model.to_device(backend_config.device)
+        with (
+                testing.assert_warns(UserWarning) if self.groups != 1
+                else nullcontext()):
+            load(self.temp_file_path, model)
+        x = numpy.random.rand(5, self.n_channels, 2).astype(self.dtype)
+        x = backend_config.get_array(x)
+        with chainer.using_config('train', False):
+            model(x)
+        model(x)
+
+    def test_model_compatibility_npz(self, backend_config):
+        self.check_model_compatibility(
+            backend_config,
+            chainer.serializers.save_npz,
+            chainer.serializers.load_npz,
+        )
+
+    @testing.with_requires('h5py')
+    def test_model_compatibility_hdf5(self, backend_config):
+        self.check_model_compatibility(
+            backend_config,
+            chainer.serializers.save_hdf5,
+            chainer.serializers.load_hdf5,
+        )
+
+
 testing.run_module(__name__, __file__)