Repository: DeepRegNet/DeepReg

Issue:
Refactor losses
## Subject of the issue

The existing loss classes all hardcoded reduction, the corresponding doc has been repeated multiple times.
Also, the multi-scaling class can be a mixin class.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=h1) Report
> Merging [#740](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=desc) (98e3d5d) into [main](https://codecov.io/gh/DeepRegNet/DeepReg/commit/232bc8a1b88cfbb49f83be311de6c5483cd6021c?el=desc) (232bc8a) will **decrease** coverage by `0.56%`.
> The diff coverage is `85.56%`.

> :exclamation: Current head 98e3d5d differs from pull request most recent head 4accc0f. Consider uploading reports for the commit 4accc0f to get more accurate results
[![Impacted file tree graph](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/graphs/tree.svg?width=650&height=150&src=pr&token=QCS45ENG57)](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##              main     #740      +/-   ##
===========================================
- Coverage   100.00%   99.43%   -0.57%     
===========================================
  Files           39       40       +1     
  Lines         2464     2470       +6     
===========================================
- Hits          2464     2456       -8     
- Misses           0       14      +14     
```


| [Impacted Files](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [deepreg/dataset/preprocess.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9kYXRhc2V0L3ByZXByb2Nlc3MucHk=) | `100.00% <ø> (ø)` | |
| [deepreg/loss/\_\_init\_\_.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9sb3NzL19faW5pdF9fLnB5) | `100.00% <ø> (ø)` | |
| [deepreg/model/network.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9tb2RlbC9uZXR3b3JrLnB5) | `100.00% <ø> (ø)` | |
| [deepreg/loss/util.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9sb3NzL3V0aWwucHk=) | `70.00% <57.14%> (-30.00%)` | :arrow_down: |
| [deepreg/registry.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9yZWdpc3RyeS5weQ==) | `97.46% <77.77%> (-2.54%)` | :arrow_down: |
| [deepreg/loss/deform.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9sb3NzL2RlZm9ybS5weQ==) | `100.00% <100.00%> (ø)` | |
| [deepreg/loss/image.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9sb3NzL2ltYWdlLnB5) | `100.00% <100.00%> (ø)` | |
| [deepreg/loss/kernel.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9sb3NzL2tlcm5lbC5weQ==) | `100.00% <100.00%> (ø)` | |
| [deepreg/loss/label.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy9sb3NzL2xhYmVsLnB5) | `100.00% <100.00%> (ø)` | |
| [deepreg/util.py](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree#diff-ZGVlcHJlZy91dGlsLnB5) | `100.00% <100.00%> (ø)` | |
| ... and [1 more](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=footer). Last update [232bc8a...4accc0f](https://codecov.io/gh/DeepRegNet/DeepReg/pull/740?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

mathpluscode: @YipengHu @kate-sann5100 could you review this PR, I did some refactoring:

- Multi-scale is now a mixin class, so that we do not have to define the scales and kernel params when defining each loss.
- The `reduction` arg is removed from loss init, and we define it when building loss.

I've checked on the demo with 1 GPU or 2 GPUs, the losses range are good.

Patch:
diff --git a/deepreg/dataset/preprocess.py b/deepreg/dataset/preprocess.py
index dd178e51..ea83b64e 100644
--- a/deepreg/dataset/preprocess.py
+++ b/deepreg/dataset/preprocess.py
@@ -141,7 +141,7 @@ class RandomAffineTransform3D(RandomTransformation3D):
         :param batch_size: total number of samples consumed per step, over all devices.
         :param scale: a positive float controlling the scale of transformation
         :param name: name of the layer
-        :param kwargs: extra arguments
+        :param kwargs: additional arguments
         """
         super().__init__(
             moving_image_size=moving_image_size,
@@ -214,7 +214,7 @@ class RandomDDFTransform3D(RandomTransformation3D):
         deformation field variance
         :param low_res_size: tuple = (1, 1, 1).
         :param name: name of layer
-        :param kwargs: extra arguments
+        :param kwargs: additional arguments
         """
 
         super().__init__(
diff --git a/deepreg/loss/__init__.py b/deepreg/loss/__init__.py
index 18f67103..450be4e3 100644
--- a/deepreg/loss/__init__.py
+++ b/deepreg/loss/__init__.py
@@ -6,7 +6,6 @@ from deepreg.loss.image import (
     GlobalMutualInformationLoss,
     LocalNormalizedCrossCorrelation,
     LocalNormalizedCrossCorrelationLoss,
-    SumSquaredDifference,
 )
 from deepreg.loss.label import (
     CrossEntropy,
@@ -14,4 +13,5 @@ from deepreg.loss.label import (
     DiceScore,
     JaccardIndex,
     JaccardLoss,
+    SumSquaredDifference,
 )
diff --git a/deepreg/loss/deform.py b/deepreg/loss/deform.py
index 0d97941e..7143a5ef 100644
--- a/deepreg/loss/deform.py
+++ b/deepreg/loss/deform.py
@@ -66,12 +66,13 @@ class GradientNorm(tf.keras.layers.Layer):
     y_true and y_pred have to be at least 5d tensor, including batch axis.
     """
 
-    def __init__(self, l1: bool = False, name: str = "GradientNorm"):
+    def __init__(self, l1: bool = False, name: str = "GradientNorm", **kwargs):
         """
         Init.
 
         :param l1: bool true if calculate L1 norm, otherwise L2 norm
         :param name: name of the loss
+        :param kwargs: additional arguments.
         """
         super().__init__(name=name)
         self.l1 = l1
@@ -112,11 +113,12 @@ class BendingEnergy(tf.keras.layers.Layer):
     y_true and y_pred have to be at least 5d tensor, including batch axis.
     """
 
-    def __init__(self, name: str = "BendingEnergy"):
+    def __init__(self, name: str = "BendingEnergy", **kwargs):
         """
         Init.
 
-        :param name: name of the loss
+        :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
         super().__init__(name=name)
 
diff --git a/deepreg/loss/image.py b/deepreg/loss/image.py
index 4239f52a..f8b8fdd0 100644
--- a/deepreg/loss/image.py
+++ b/deepreg/loss/image.py
@@ -2,54 +2,12 @@
 import tensorflow as tf
 
 from deepreg.constant import EPS
-from deepreg.loss.util import NegativeLossMixin
-from deepreg.loss.util import gaussian_kernel1d_size as gaussian_kernel1d
-from deepreg.loss.util import (
-    rectangular_kernel1d,
-    separable_filter,
-    triangular_kernel1d,
-)
+from deepreg.loss.kernel import gaussian_kernel1d_size as gaussian_kernel1d
+from deepreg.loss.kernel import rectangular_kernel1d, triangular_kernel1d
+from deepreg.loss.util import NegativeLossMixin, separable_filter
 from deepreg.registry import REGISTRY
 
 
-@REGISTRY.register_loss(name="ssd")
-class SumSquaredDifference(tf.keras.losses.Loss):
-    """
-    Sum of squared distance between y_true and y_pred.
-
-    y_true and y_pred have to be at least 1d tensor, including batch axis.
-    """
-
-    def __init__(
-        self,
-        reduction: str = tf.keras.losses.Reduction.NONE,
-        name: str = "SumSquaredDifference",
-    ):
-        """
-        Init.
-
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
-        :param name: name of the loss
-        """
-        super().__init__(reduction=reduction, name=name)
-        self.flatten = tf.keras.layers.Flatten()
-
-    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
-        """
-        Return loss for a batch.
-
-        :param y_true: shape = (batch, ...)
-        :param y_pred: shape = (batch, ...)
-        :return: shape = (batch,)
-        """
-        loss = tf.math.squared_difference(y_true, y_pred)
-        loss = self.flatten(loss)
-        return tf.reduce_mean(loss, axis=1)
-
-
 class GlobalMutualInformation(tf.keras.losses.Loss):
     """
     Differentiable global mutual information via Parzen windowing method.
@@ -64,21 +22,18 @@ class GlobalMutualInformation(tf.keras.losses.Loss):
         self,
         num_bins: int = 23,
         sigma_ratio: float = 0.5,
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "GlobalMutualInformation",
+        **kwargs,
     ):
         """
         Init.
 
         :param num_bins: number of bins for intensity, the default value is empirical.
         :param sigma_ratio: a hyper param for gaussian function
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
         :param name: name of the loss
+        :param kwargs: additional arguments.
         """
-        super().__init__(reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
         self.num_bins = num_bins
         self.sigma_ratio = sigma_ratio
 
@@ -195,8 +150,8 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
         kernel_type: str = "rectangular",
         smooth_nr: float = EPS,
         smooth_dr: float = EPS,
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "LocalNormalizedCrossCorrelation",
+        **kwargs,
     ):
         """
         Init.
@@ -205,13 +160,10 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
         :param kernel_type: str, rectangular, triangular or gaussian
         :param smooth_nr: small constant added to numerator in case of zero covariance.
         :param smooth_dr: small constant added to denominator in case of zero variance.
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
-        :param name: name of the loss
+        :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
-        super().__init__(reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
         if kernel_type not in self.kernel_fn_dict.keys():
             raise ValueError(
                 f"Wrong kernel_type {kernel_type} for LNCC loss type. "
@@ -346,16 +298,16 @@ class GlobalNormalizedCrossCorrelation(tf.keras.losses.Loss):
 
     def __init__(
         self,
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "GlobalNormalizedCrossCorrelation",
+        **kwargs,
     ):
         """
         Init.
-        :param reduction: using AUTO reduction,
-            calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
+
         :param name: name of the loss
+        :param kwargs: additional arguments.
         """
-        super().__init__(reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
 
     def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         """
diff --git a/deepreg/loss/kernel.py b/deepreg/loss/kernel.py
new file mode 100644
index 00000000..05179142
--- /dev/null
+++ b/deepreg/loss/kernel.py
@@ -0,0 +1,94 @@
+import math
+
+import tensorflow as tf
+
+
+def rectangular_kernel1d(kernel_size: int) -> tf.Tensor:
+    """
+    Return a the 1D rectangular kernel for LocalNormalizedCrossCorrelation.
+
+    :param kernel_size: scalar, size of the 1-D kernel
+    :return: kernel_weights, of shape (kernel_size, )
+    """
+
+    kernel = tf.ones(shape=(kernel_size,), dtype=tf.float32)
+    return kernel
+
+
+def triangular_kernel1d(kernel_size: int) -> tf.Tensor:
+    """
+    Return a the 1D triangular kernel for LocalNormalizedCrossCorrelation.
+
+    Assume kernel_size is odd, it will be a smoothed from
+    a kernel which center part is zero.
+    Then length of the ones will be around half kernel_size.
+    The weight scale of the kernel does not matter as LNCC will normalize it.
+
+    :param kernel_size: scalar, size of the 1-D kernel
+    :return: kernel_weights, of shape (kernel_size, )
+    """
+    assert kernel_size >= 3
+    assert kernel_size % 2 != 0
+
+    padding = kernel_size // 2
+    kernel = tf.constant(
+        [0] * math.ceil(padding / 2)
+        + [1] * (kernel_size - padding)
+        + [0] * math.floor(padding / 2),
+        dtype=tf.float32,
+    )
+
+    # (padding*2, )
+    filters = tf.ones(shape=(kernel_size - padding, 1, 1), dtype=tf.float32)
+
+    # (kernel_size, 1, 1)
+    kernel = tf.nn.conv1d(
+        kernel[None, :, None], filters=filters, stride=[1, 1, 1], padding="SAME"
+    )
+
+    return kernel[0, :, 0]
+
+
+def gaussian_kernel1d_size(kernel_size: int) -> tf.Tensor:
+    """
+    Return a the 1D Gaussian kernel for LocalNormalizedCrossCorrelation.
+
+    :param kernel_size: scalar, size of the 1-D kernel
+    :return: filters, of shape (kernel_size, )
+    """
+    mean = (kernel_size - 1) / 2.0
+    sigma = kernel_size / 3
+
+    grid = tf.range(0, kernel_size, dtype=tf.float32)
+    filters = tf.exp(-tf.square(grid - mean) / (2 * sigma ** 2))
+
+    return filters
+
+
+def gaussian_kernel1d_sigma(sigma: int) -> tf.Tensor:
+    """
+    Calculate a gaussian kernel.
+
+    :param sigma: number defining standard deviation for
+                  gaussian kernel.
+    :return: shape = (dim, )
+    """
+    assert sigma > 0
+    tail = int(sigma * 3)
+    kernel = tf.exp([-0.5 * x ** 2 / sigma ** 2 for x in range(-tail, tail + 1)])
+    kernel = kernel / tf.reduce_sum(kernel)
+    return kernel
+
+
+def cauchy_kernel1d(sigma: int) -> tf.Tensor:
+    """
+    Approximating cauchy kernel in 1d.
+
+    :param sigma: int, defining standard deviation of kernel.
+    :return: shape = (dim, )
+    """
+    assert sigma > 0
+    tail = int(sigma * 5)
+    k = tf.math.reciprocal([((x / sigma) ** 2 + 1) for x in range(-tail, tail + 1)])
+    k = k / tf.reduce_sum(k)
+    return k
diff --git a/deepreg/loss/label.py b/deepreg/loss/label.py
index 4da08306..6791ee86 100644
--- a/deepreg/loss/label.py
+++ b/deepreg/loss/label.py
@@ -1,104 +1,54 @@
 """Provide different loss or metrics classes for labels."""
 
-from typing import List, Optional
-
 import tensorflow as tf
 
 from deepreg.constant import EPS
-from deepreg.loss.util import NegativeLossMixin, cauchy_kernel1d
-from deepreg.loss.util import gaussian_kernel1d_sigma as gaussian_kernel1d
-from deepreg.loss.util import separable_filter
+from deepreg.loss.util import MultiScaleMixin, NegativeLossMixin
 from deepreg.registry import REGISTRY
 
 
-class MultiScaleLoss(tf.keras.losses.Loss):
+class SumSquaredDifference(tf.keras.losses.Loss):
     """
-    Base class for multi-scale loss.
+    Actually, mean of squared distance between y_true and y_pred.
 
-    It applies the loss at different scales (gaussian or cauchy smoothing).
-    It is assumed that loss values are between 0 and 1.
-    """
+    The inconsistent name was for convention.
 
-    kernel_fn_dict = dict(gaussian=gaussian_kernel1d, cauchy=cauchy_kernel1d)
+    y_true and y_pred have to be at least 1d tensor, including batch axis.
+    """
 
     def __init__(
         self,
-        scales: Optional[List] = None,
-        kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.NONE,
-        name: str = "MultiScaleLoss",
+        name: str = "SumSquaredDifference",
+        **kwargs,
     ):
         """
         Init.
 
-        :param scales: list of scalars or None, if None, do not apply any scaling.
-        :param kernel: gaussian or cauchy.
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
-        :param name: str, name of the loss.
+        :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
-        super().__init__(reduction=reduction, name=name)
-        assert kernel in ["gaussian", "cauchy"]
-        self.scales = scales
-        self.kernel = kernel
+        super().__init__(name=name, **kwargs)
+        self.flatten = tf.keras.layers.Flatten()
 
     def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         """
-        Use _call to calculate loss at different scales.
+        Return mean squared different for a batch.
 
-        :param y_true: ground-truth tensor, shape = (batch, dim1, dim2, dim3).
-        :param y_pred: predicted tensor, shape = (batch, dim1, dim2, dim3).
-        :return: multi-scale loss, shape = (batch, ).
-        """
-        if self.scales is None:
-            return self._call(y_true=y_true, y_pred=y_pred)
-        kernel_fn = self.kernel_fn_dict[self.kernel]
-        losses = []
-        for s in self.scales:
-            if s == 0:
-                # no smoothing
-                losses.append(
-                    self._call(
-                        y_true=y_true,
-                        y_pred=y_pred,
-                    )
-                )
-            else:
-                losses.append(
-                    self._call(
-                        y_true=separable_filter(
-                            tf.expand_dims(y_true, axis=4), kernel_fn(s)
-                        )[..., 0],
-                        y_pred=separable_filter(
-                            tf.expand_dims(y_pred, axis=4), kernel_fn(s)
-                        )[..., 0],
-                    )
-                )
-        loss = tf.add_n(losses)
-        loss = loss / len(self.scales)
-        return loss
-
-    def _call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+        :param y_true: shape = (batch, ...)
+        :param y_pred: shape = (batch, ...)
+        :return: shape = (batch,)
         """
-        Return loss for a batch.
+        loss = tf.math.squared_difference(y_true, y_pred)
+        loss = self.flatten(loss)
+        return tf.reduce_mean(loss, axis=1)
 
-        :param y_true: ground-truth tensor.
-        :param y_pred: predicted tensor.
-        :return: negated loss.
-        """
-        raise NotImplementedError
 
-    def get_config(self) -> dict:
-        """Return the config dictionary for recreating this class."""
-        config = super().get_config()
-        config["scales"] = self.scales
-        config["kernel"] = self.kernel
-        return config
+@REGISTRY.register_loss(name="ssd")
+class SumSquaredDifferenceLoss(MultiScaleMixin, SumSquaredDifference):
+    """Define loss with multi-scaling options."""
 
 
-class DiceScore(MultiScaleLoss):
+class DiceScore(tf.keras.losses.Loss):
     """
     Define dice score.
 
@@ -129,10 +79,8 @@ class DiceScore(MultiScaleLoss):
         background_weight: float = 0.0,
         smooth_nr: float = EPS,
         smooth_dr: float = EPS,
-        scales: Optional[List] = None,
-        kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "DiceScore",
+        **kwargs,
     ):
         """
         Init.
@@ -141,15 +89,10 @@ class DiceScore(MultiScaleLoss):
         :param background_weight: weight for background, where y == 0.
         :param smooth_nr: small constant added to numerator in case of zero covariance.
         :param smooth_dr: small constant added to denominator in case of zero variance.
-        :param scales: list of scalars or None, if None, do not apply any scaling.
-        :param kernel: gaussian or cauchy.
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
-        :param name: str, name of the loss.
+        :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
-        super().__init__(scales=scales, kernel=kernel, reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
         if background_weight < 0 or background_weight > 1:
             raise ValueError(
                 "The background weight for Dice Score must be "
@@ -162,7 +105,7 @@ class DiceScore(MultiScaleLoss):
         self.smooth_dr = smooth_dr
         self.flatten = tf.keras.layers.Flatten()
 
-    def _call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         """
         Return loss for a batch.
 
@@ -211,12 +154,11 @@ class DiceScore(MultiScaleLoss):
 
 
 @REGISTRY.register_loss(name="dice")
-class DiceLoss(NegativeLossMixin, DiceScore):
-    """Revert the sign of DiceScore."""
+class DiceLoss(NegativeLossMixin, MultiScaleMixin, DiceScore):
+    """Revert the sign of DiceScore and support multi-scaling options."""
 
 
-@REGISTRY.register_loss(name="cross-entropy")
-class CrossEntropy(MultiScaleLoss):
+class CrossEntropy(tf.keras.losses.Loss):
     """
     Define weighted cross-entropy.
 
@@ -229,25 +171,19 @@ class CrossEntropy(MultiScaleLoss):
         binary: bool = False,
         background_weight: float = 0.0,
         smooth: float = EPS,
-        scales: Optional[List] = None,
-        kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "CrossEntropy",
+        **kwargs,
     ):
         """
         Init.
 
         :param binary: if True, project y_true, y_pred to 0 or 1
         :param background_weight: weight for background, where y == 0.
-        :param scales: list of scalars or None, if None, do not apply any scaling.
-        :param kernel: gaussian or cauchy.
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
-        :param name: str, name of the loss.
+        :param smooth: smooth constant for log.
+        :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
-        super().__init__(scales=scales, kernel=kernel, reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
         if background_weight < 0 or background_weight > 1:
             raise ValueError(
                 "The background weight for Cross Entropy must be "
@@ -258,7 +194,7 @@ class CrossEntropy(MultiScaleLoss):
         self.smooth = smooth
         self.flatten = tf.keras.layers.Flatten()
 
-    def _call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         """
         Return loss for a batch.
 
@@ -296,6 +232,11 @@ class CrossEntropy(MultiScaleLoss):
         return config
 
 
+@REGISTRY.register_loss(name="cross-entropy")
+class CrossEntropyLoss(MultiScaleMixin, CrossEntropy):
+    """Define loss with multi-scaling options."""
+
+
 class JaccardIndex(DiceScore):
     """
     Define Jaccard index.
@@ -323,10 +264,8 @@ class JaccardIndex(DiceScore):
         background_weight: float = 0.0,
         smooth_nr: float = EPS,
         smooth_dr: float = EPS,
-        scales: Optional[List] = None,
-        kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "JaccardIndex",
+        **kwargs,
     ):
         """
         Init.
@@ -335,26 +274,19 @@ class JaccardIndex(DiceScore):
         :param background_weight: weight for background, where y == 0.
         :param smooth_nr: small constant added to numerator in case of zero covariance.
         :param smooth_dr: small constant added to denominator in case of zero variance.
-        :param scales: list of scalars or None, if None, do not apply any scaling.
-        :param kernel: gaussian or cauchy.
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
-        :param name: str, name of the loss.
+        :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
         super().__init__(
             binary=binary,
             background_weight=background_weight,
             smooth_nr=smooth_nr,
             smooth_dr=smooth_dr,
-            scales=scales,
-            kernel=kernel,
-            reduction=reduction,
             name=name,
+            **kwargs,
         )
 
-    def _call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         """
         Return loss for a batch.
 
@@ -394,7 +326,7 @@ class JaccardIndex(DiceScore):
 
 
 @REGISTRY.register_loss(name="jaccard")
-class JaccardLoss(NegativeLossMixin, JaccardIndex):
+class JaccardLoss(NegativeLossMixin, MultiScaleMixin, JaccardIndex):
     """Revert the sign of JaccardIndex."""
 
 
diff --git a/deepreg/loss/util.py b/deepreg/loss/util.py
index 05083461..044569cd 100644
--- a/deepreg/loss/util.py
+++ b/deepreg/loss/util.py
@@ -1,8 +1,90 @@
 """Provide helper functions or classes for defining loss or metrics."""
-import math
+
+from typing import List, Optional, Union
 
 import tensorflow as tf
 
+from deepreg.loss.kernel import cauchy_kernel1d
+from deepreg.loss.kernel import gaussian_kernel1d_sigma as gaussian_kernel1d
+
+
+class MultiScaleMixin(tf.keras.losses.Loss):
+    """
+    Mixin class for multi-scale loss.
+
+    It applies the loss at different scales (gaussian or cauchy smoothing).
+    It is assumed that loss values are between 0 and 1.
+    """
+
+    kernel_fn_dict = dict(gaussian=gaussian_kernel1d, cauchy=cauchy_kernel1d)
+
+    def __init__(
+        self,
+        scales: Optional[Union[List, float, int]] = None,
+        kernel: str = "gaussian",
+        **kwargs,
+    ):
+        """
+        Init.
+
+        :param scales: list of scalars or None, if None, do not apply any scaling.
+        :param kernel: gaussian or cauchy.
+        :param kwargs: additional arguments.
+        """
+        super().__init__(**kwargs)
+        if kernel not in self.kernel_fn_dict:
+            raise ValueError(
+                f"Kernel {kernel} is not supported."
+                f"Supported kernels are {list(self.kernel_fn_dict.keys())}"
+            )
+        if scales is not None and not isinstance(scales, list):
+            scales = [scales]
+        self.scales = scales
+        self.kernel = kernel
+
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+        """
+        Use super().call to calculate loss at different scales.
+
+        :param y_true: ground-truth tensor, shape = (batch, dim1, dim2, dim3).
+        :param y_pred: predicted tensor, shape = (batch, dim1, dim2, dim3).
+        :return: multi-scale loss, shape = (batch, ).
+        """
+        if self.scales is None:
+            return super().call(y_true=y_true, y_pred=y_pred)
+        kernel_fn = self.kernel_fn_dict[self.kernel]
+        losses = []
+        for s in self.scales:
+            if s == 0:
+                # no smoothing
+                losses.append(
+                    super().call(
+                        y_true=y_true,
+                        y_pred=y_pred,
+                    )
+                )
+            else:
+                losses.append(
+                    super().call(
+                        y_true=separable_filter(
+                            tf.expand_dims(y_true, axis=4), kernel_fn(s)
+                        )[..., 0],
+                        y_pred=separable_filter(
+                            tf.expand_dims(y_pred, axis=4), kernel_fn(s)
+                        )[..., 0],
+                    )
+                )
+        loss = tf.add_n(losses)
+        loss = loss / len(self.scales)
+        return loss
+
+    def get_config(self) -> dict:
+        """Return the config dictionary for recreating this class."""
+        config = super().get_config()
+        config["scales"] = self.scales
+        config["kernel"] = self.kernel
+        return config
+
 
 class NegativeLossMixin(tf.keras.losses.Loss):
     """Mixin class to revert the sign of the loss value."""
@@ -27,97 +109,6 @@ class NegativeLossMixin(tf.keras.losses.Loss):
         return -super().call(y_true=y_true, y_pred=y_pred)
 
 
-def rectangular_kernel1d(kernel_size: int) -> tf.Tensor:
-    """
-    Return a the 1D rectangular kernel for LocalNormalizedCrossCorrelation.
-
-    :param kernel_size: scalar, size of the 1-D kernel
-    :return: kernel_weights, of shape (kernel_size, )
-    """
-
-    kernel = tf.ones(shape=(kernel_size,), dtype=tf.float32)
-    return kernel
-
-
-def triangular_kernel1d(kernel_size: int) -> tf.Tensor:
-    """
-    Return a the 1D triangular kernel for LocalNormalizedCrossCorrelation.
-
-    Assume kernel_size is odd, it will be a smoothed from
-    a kernel which center part is zero.
-    Then length of the ones will be around half kernel_size.
-    The weight scale of the kernel does not matter as LNCC will normalize it.
-
-    :param kernel_size: scalar, size of the 1-D kernel
-    :return: kernel_weights, of shape (kernel_size, )
-    """
-    assert kernel_size >= 3
-    assert kernel_size % 2 != 0
-
-    padding = kernel_size // 2
-    kernel = tf.constant(
-        [0] * math.ceil(padding / 2)
-        + [1] * (kernel_size - padding)
-        + [0] * math.floor(padding / 2),
-        dtype=tf.float32,
-    )
-
-    # (padding*2, )
-    filters = tf.ones(shape=(kernel_size - padding, 1, 1), dtype=tf.float32)
-
-    # (kernel_size, 1, 1)
-    kernel = tf.nn.conv1d(
-        kernel[None, :, None], filters=filters, stride=[1, 1, 1], padding="SAME"
-    )
-
-    return kernel[0, :, 0]
-
-
-def gaussian_kernel1d_size(kernel_size: int) -> tf.Tensor:
-    """
-    Return a the 1D Gaussian kernel for LocalNormalizedCrossCorrelation.
-
-    :param kernel_size: scalar, size of the 1-D kernel
-    :return: filters, of shape (kernel_size, )
-    """
-    mean = (kernel_size - 1) / 2.0
-    sigma = kernel_size / 3
-
-    grid = tf.range(0, kernel_size, dtype=tf.float32)
-    filters = tf.exp(-tf.square(grid - mean) / (2 * sigma ** 2))
-
-    return filters
-
-
-def gaussian_kernel1d_sigma(sigma: int) -> tf.Tensor:
-    """
-    Calculate a gaussian kernel.
-
-    :param sigma: number defining standard deviation for
-                  gaussian kernel.
-    :return: shape = (dim, )
-    """
-    assert sigma > 0
-    tail = int(sigma * 3)
-    kernel = tf.exp([-0.5 * x ** 2 / sigma ** 2 for x in range(-tail, tail + 1)])
-    kernel = kernel / tf.reduce_sum(kernel)
-    return kernel
-
-
-def cauchy_kernel1d(sigma: int) -> tf.Tensor:
-    """
-    Approximating cauchy kernel in 1d.
-
-    :param sigma: int, defining standard deviation of kernel.
-    :return: shape = (dim, )
-    """
-    assert sigma > 0
-    tail = int(sigma * 5)
-    k = tf.math.reciprocal([((x / sigma) ** 2 + 1) for x in range(-tail, tail + 1)])
-    k = k / tf.reduce_sum(k)
-    return k
-
-
 def separable_filter(tensor: tf.Tensor, kernel: tf.Tensor) -> tf.Tensor:
     """
     Create a 3d separable filter.
diff --git a/deepreg/model/network.py b/deepreg/model/network.py
index 22af33cc..b9b8552a 100644
--- a/deepreg/model/network.py
+++ b/deepreg/model/network.py
@@ -210,8 +210,11 @@ class RegistrationModel(tf.keras.Model):
                 )
                 return
 
+            # do not perform reduction over batch axis for supporting multi-device
+            # training, model.fit() will average over global batch size automatically
             loss_layer: tf.keras.layers.Layer = REGISTRY.build_loss(
-                config=dict_without(d=loss_config, key="weight")
+                config=dict_without(d=loss_config, key="weight"),
+                default_args={"reduction": tf.keras.losses.Reduction.NONE},
             )
             loss_value = loss_layer(**inputs_dict)
             weighted_loss = loss_value * weight
diff --git a/deepreg/registry.py b/deepreg/registry.py
index d85fbe78..f031e8ea 100644
--- a/deepreg/registry.py
+++ b/deepreg/registry.py
@@ -1,12 +1,14 @@
 from copy import deepcopy
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional
 
 BACKBONE_CLASS = "backbone_class"
 LOSS_CLASS = "loss_class"
+METRIC_CLASS = "metric_class"
 MODEL_CLASS = "model_class"
 DATA_AUGMENTATION_CLASS = "da_class"
 DATA_LOADER_CLASS = "data_loader_class"
 FILE_LOADER_CLASS = "file_loader_class"
+
 KNOWN_CATEGORIES = [
     BACKBONE_CLASS,
     LOSS_CLASS,
@@ -114,7 +116,7 @@ class Registry:
         return decorator
 
     def build_from_config(
-        self, category: str, config: dict, default_args: Optional[dict] = None
+        self, category: str, config: Dict, default_args: Optional[dict] = None
     ) -> Any:
         """
         Build a class instance from config dict.
@@ -169,7 +171,7 @@ class Registry:
         """
         return self.register(category=MODEL_CLASS, name=name, cls=cls, force=force)
 
-    def build_model(self, config: dict, default_args: Optional[dict] = None) -> Any:
+    def build_model(self, config: Dict, default_args: Optional[dict] = None) -> Any:
         """
         Instantiate a registered model class.
 
@@ -194,7 +196,7 @@ class Registry:
         """
         return self.register(category=BACKBONE_CLASS, name=name, cls=cls, force=force)
 
-    def build_backbone(self, config: dict, default_args: Optional[dict] = None) -> Any:
+    def build_backbone(self, config: Dict, default_args: Optional[dict] = None) -> Any:
         """
         Instantiate a registered backbone class.
 
@@ -219,7 +221,7 @@ class Registry:
         """
         return self.register(category=LOSS_CLASS, name=name, cls=cls, force=force)
 
-    def build_loss(self, config: dict, default_args: Optional[dict] = None) -> Callable:
+    def build_loss(self, config: Dict, default_args: Optional[dict] = None) -> Callable:
         """
         Instantiate a registered loss class.
 
@@ -247,7 +249,7 @@ class Registry:
         )
 
     def build_data_loader(
-        self, config: dict, default_args: Optional[dict] = None
+        self, config: Dict, default_args: Optional[dict] = None
     ) -> Any:
         """
         Instantiate a registered data loader class.
@@ -291,7 +293,7 @@ class Registry:
         )
 
     def build_data_augmentation(
-        self, config: dict, default_args: Optional[dict] = None
+        self, config: Dict, default_args: Optional[dict] = None
     ) -> Callable:
         """
         Instantiate a registered data augmentation class.
diff --git a/deepreg/util.py b/deepreg/util.py
index 0abb02c9..4b4a2779 100644
--- a/deepreg/util.py
+++ b/deepreg/util.py
@@ -9,7 +9,6 @@ import numpy as np
 import pandas as pd
 import tensorflow as tf
 
-import deepreg.loss.image as image_loss
 import deepreg.loss.label as label_loss
 from deepreg.dataset.load import get_data_loader
 from deepreg.dataset.loader.interface import DataLoader
@@ -163,7 +162,7 @@ def calculate_metrics(
         y_pred = pred_fixed_image[sample_index : (sample_index + 1), :, :, :]
         y_true = tf.expand_dims(y_true, axis=4)
         y_pred = tf.expand_dims(y_pred, axis=4)
-        ssd = image_loss.SumSquaredDifference()(y_true=y_true, y_pred=y_pred).numpy()
+        ssd = label_loss.SumSquaredDifference()(y_true=y_true, y_pred=y_pred).numpy()
     else:
         ssd = None
 
diff --git a/docs/source/docs/registered_classes.md b/docs/source/docs/registered_classes.md
index f3577fc4..2c3aade7 100644
--- a/docs/source/docs/registered_classes.md
+++ b/docs/source/docs/registered_classes.md
@@ -31,14 +31,14 @@ The category is `loss_class`. Registered keys and values are as following.
 | key             | value                                                     |
 | :-------------- | :-------------------------------------------------------- |
 | "bending"       | `deepreg.loss.deform.BendingEnergy`                       |
-| "cross-entropy" | `deepreg.loss.label.CrossEntropy`                         |
+| "cross-entropy" | `deepreg.loss.label.CrossEntropyLoss`                     |
 | "dice"          | `deepreg.loss.label.DiceLoss`                             |
 | "gmi"           | `deepreg.loss.image.GlobalMutualInformationLoss`          |
 | "gncc"          | `deepreg.loss.image.GlobalNormalizedCrossCorrelationLoss` |
 | "gradient"      | `deepreg.loss.deform.GradientNorm`                        |
 | "jaccard"       | `deepreg.loss.label.JaccardLoss`                          |
 | "lncc"          | `deepreg.loss.image.LocalNormalizedCrossCorrelationLoss`  |
-| "ssd"           | `deepreg.loss.image.SumSquaredDifference`                 |
+| "ssd"           | `deepreg.loss.label.SumSquaredDifferenceLoss`             |
 
 ## Data Augmentation
 
diff --git a/examples/custom_image_label_loss.py b/examples/custom_image_label_loss.py
index ed87c3a3..63e7960d 100644
--- a/examples/custom_image_label_loss.py
+++ b/examples/custom_image_label_loss.py
@@ -15,19 +15,16 @@ class RootMeanSquaredDifference(tf.keras.losses.Loss):
 
     def __init__(
         self,
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "RootMeanSquaredDifference",
+        **kwargs,
     ):
         """
         Init.
 
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
         :param name: name of the loss
+        :param kwargs: additional arguments.
         """
-        super().__init__(reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
         self.flatten = tf.keras.layers.Flatten()
 
     def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
diff --git a/examples/custom_parameterized_image_label_loss.py b/examples/custom_parameterized_image_label_loss.py
index b77c0f5f..c107adc2 100644
--- a/examples/custom_parameterized_image_label_loss.py
+++ b/examples/custom_parameterized_image_label_loss.py
@@ -16,20 +16,17 @@ class LPNorm(tf.keras.losses.Loss):
     def __init__(
         self,
         p: int,
-        reduction: str = tf.keras.losses.Reduction.NONE,
         name: str = "LPNorm",
+        **kwargs,
     ):
         """
         Init.
 
         :param p: order of the norm, 1 or 2.
-        :param reduction: do not perform reduction over batch axis.
-            this is for supporting multi-device training,
-            model.fit() will average over global batch size automatically.
-            Loss returns a tensor of shape (batch, ).
         :param name: name of the loss.
+        :param kwargs: additional arguments.
         """
-        super().__init__(reduction=reduction, name=name)
+        super().__init__(name=name, **kwargs)
         if p not in [1, 2]:
             raise ValueError(f"For LPNorm, p must be 0 or 1, got {p}.")
         self.p = p


Test patch:
diff --git a/test/unit/test_loss_image.py b/test/unit/test_loss_image.py
index fb43b932..8355f0b0 100644
--- a/test/unit/test_loss_image.py
+++ b/test/unit/test_loss_image.py
@@ -4,6 +4,7 @@ pytest style.
 Notes: The format of inputs to the function dissimilarity_fn
 in image.py should be better converted into tf tensor type beforehand.
 """
+
 from test.unit.util import is_equal_tf
 from typing import Tuple
 
@@ -15,29 +16,6 @@ import deepreg.loss.image as image
 from deepreg.constant import EPS
 
 
-class TestSumSquaredDistance:
-    @pytest.mark.parametrize(
-        "y_true,y_pred,shape,expected",
-        [
-            (0.6, 0.3, (3,), 0.09),
-            (0.6, 0.3, (3, 3), 0.09),
-            (0.6, 0.3, (3, 3, 3), 0.09),
-            (0.6, 0.3, (3, 3, 3), 0.09),
-            (0.5, 0.5, (3, 3), 0.0),
-            (0.3, 0.6, (3, 3), 0.09),
-        ],
-    )
-    def test_output(self, y_true, y_pred, shape, expected):
-        y_true = y_true * np.ones(shape=shape)
-        y_pred = y_pred * np.ones(shape=shape)
-        expected = expected * np.ones(shape=(shape[0],))
-        got = image.SumSquaredDifference().call(
-            y_true,
-            y_pred,
-        )
-        assert is_equal_tf(got, expected)
-
-
 class TestGlobalMutualInformation:
     @pytest.mark.parametrize(
         "y_true,y_pred,shape,expected",
@@ -62,7 +40,7 @@ class TestGlobalMutualInformation:
         expected = dict(
             num_bins=23,
             sigma_ratio=0.5,
-            reduction=tf.keras.losses.Reduction.NONE,
+            reduction=tf.keras.losses.Reduction.AUTO,
             name="GlobalMutualInformation",
         )
         assert got == expected
@@ -228,7 +206,7 @@ class TestLocalNormalizedCrossCorrelation:
         expected = dict(
             kernel_size=9,
             kernel_type="rectangular",
-            reduction=tf.keras.losses.Reduction.NONE,
+            reduction=tf.keras.losses.Reduction.AUTO,
             name="LocalNormalizedCrossCorrelation",
             smooth_nr=1e-5,
             smooth_dr=1e-5,
diff --git a/test/unit/test_loss_kernel.py b/test/unit/test_loss_kernel.py
new file mode 100644
index 00000000..488f4e36
--- /dev/null
+++ b/test/unit/test_loss_kernel.py
@@ -0,0 +1,87 @@
+from test.unit.util import is_equal_tf
+
+import numpy as np
+import pytest
+import tensorflow as tf
+
+from deepreg.loss.kernel import (
+    cauchy_kernel1d,
+    gaussian_kernel1d_sigma,
+    gaussian_kernel1d_size,
+    rectangular_kernel1d,
+    triangular_kernel1d,
+)
+
+
+@pytest.mark.parametrize("sigma", [1, 3, 2.2])
+def test_cauchy_kernel1d(sigma):
+    """
+    Testing the 1-D cauchy kernel
+    :param sigma: float
+    :return:
+    """
+    tail = int(sigma * 5)
+    expected = [1 / ((x / sigma) ** 2 + 1) for x in range(-tail, tail + 1)]
+    expected = expected / np.sum(expected)
+    got = cauchy_kernel1d(sigma)
+    assert is_equal_tf(got, expected)
+
+
+@pytest.mark.parametrize("sigma", [1, 3, 2.2])
+def test_gaussian_kernel1d_sigma(sigma):
+    """
+    Testing the 1-D gaussian kernel given sigma as input
+    :param sigma: float
+    :return:
+    """
+    tail = int(sigma * 3)
+    expected = [np.exp(-0.5 * x ** 2 / sigma ** 2) for x in range(-tail, tail + 1)]
+    expected = expected / np.sum(expected)
+    got = gaussian_kernel1d_sigma(sigma)
+    assert is_equal_tf(got, expected)
+
+
+@pytest.mark.parametrize("kernel_size", [3, 7, 11])
+def test_gaussian_kernel1d_size(kernel_size):
+    """
+    Testing the 1-D gaussian kernel given size as input
+    :param kernel_size: int
+    :return:
+    """
+    mean = (kernel_size - 1) / 2.0
+    sigma = kernel_size / 3
+
+    grid = tf.range(0, kernel_size, dtype=tf.float32)
+    expected = tf.exp(-tf.square(grid - mean) / (2 * sigma ** 2))
+
+    got = gaussian_kernel1d_size(kernel_size)
+    assert is_equal_tf(got, expected)
+
+
+@pytest.mark.parametrize("kernel_size", [3, 7, 11])
+def test_rectangular_kernel1d(kernel_size):
+    """
+    Testing the 1-D rectangular kernel
+    :param kernel_size: int
+    :return:
+    """
+    expected = tf.ones(shape=(kernel_size,), dtype=tf.float32)
+    got = rectangular_kernel1d(kernel_size)
+    assert is_equal_tf(got, expected)
+
+
+@pytest.mark.parametrize("kernel_size", [3, 5, 7, 9])
+def test_triangular_kernel1d(kernel_size):
+    """
+    Testing the 1-D triangular kernel
+    :param kernel_size: int (odd number)
+    :return:
+    """
+    expected = np.zeros(shape=(kernel_size,), dtype=np.float32)
+    expected[kernel_size // 2] = kernel_size // 2 + 1
+    for it_k in range(kernel_size // 2):
+        expected[it_k] = it_k + 1
+        expected[-it_k - 1] = it_k + 1
+
+    got = triangular_kernel1d(kernel_size)
+    assert is_equal_tf(got, expected)
diff --git a/test/unit/test_loss_label.py b/test/unit/test_loss_label.py
index e116d215..2343dbdb 100644
--- a/test/unit/test_loss_label.py
+++ b/test/unit/test_loss_label.py
@@ -16,22 +16,27 @@ import deepreg.loss.label as label
 from deepreg.constant import EPS
 
 
-class TestMultiScaleLoss:
-    def test_call(self):
-        loss = label.MultiScaleLoss()
-        with pytest.raises(NotImplementedError):
-            loss.call(0, 0)
-
-    def test_get_config(self):
-        loss = label.MultiScaleLoss()
-        got = loss.get_config()
-        expected = dict(
-            scales=None,
-            kernel="gaussian",
-            reduction=tf.keras.losses.Reduction.NONE,
-            name="MultiScaleLoss",
+class TestSumSquaredDistance:
+    @pytest.mark.parametrize(
+        "y_true,y_pred,shape,expected",
+        [
+            (0.6, 0.3, (3,), 0.09),
+            (0.6, 0.3, (3, 3), 0.09),
+            (0.6, 0.3, (3, 3, 3), 0.09),
+            (0.6, 0.3, (3, 3, 3), 0.09),
+            (0.5, 0.5, (3, 3), 0.0),
+            (0.3, 0.6, (3, 3), 0.09),
+        ],
+    )
+    def test_output(self, y_true, y_pred, shape, expected):
+        y_true = y_true * np.ones(shape=shape)
+        y_pred = y_pred * np.ones(shape=shape)
+        expected = expected * np.ones(shape=(shape[0],))
+        got = label.SumSquaredDifference().call(
+            y_true,
+            y_pred,
         )
-        assert got == expected
+        assert is_equal_tf(got, expected)
 
 
 class TestDiceScore:
@@ -69,7 +74,7 @@ class TestDiceScore:
         y_true = tf.ones(shape=shape) * value
         y_pred = tf.ones(shape=shape) * value
 
-        got = label.DiceScore(smooth_nr=smooth_nr, smooth_dr=smooth_dr)._call(
+        got = label.DiceScore(smooth_nr=smooth_nr, smooth_dr=smooth_dr).call(
             y_true,
             y_pred,
         )
@@ -126,9 +131,7 @@ class TestDiceScore:
             background_weight=0.0,
             smooth_nr=1e-5,
             smooth_dr=1e-5,
-            scales=None,
-            kernel="gaussian",
-            reduction=tf.keras.losses.Reduction.NONE,
+            reduction=tf.keras.losses.Reduction.AUTO,
             name="DiceScore",
         )
         assert got == expected
@@ -186,7 +189,7 @@ class TestCrossEntropy:
         y_true = tf.ones(shape=shape) * value
         y_pred = tf.ones(shape=shape) * value
 
-        got = label.CrossEntropy(smooth=smooth)._call(
+        got = label.CrossEntropy(smooth=smooth).call(
             y_true,
             y_pred,
         )
@@ -194,19 +197,18 @@ class TestCrossEntropy:
         assert is_equal_tf(got[0], expected)
 
     @pytest.mark.parametrize(
-        "binary,background_weight,scales,expected",
+        "binary,background_weight,expected",
         [
-            (True, 0.0, None, -np.log(EPS)),
-            (False, 0.0, None, -0.6 * np.log(0.3 + EPS)),
-            (False, 0.2, None, -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
-            (False, 0.2, [0, 0], -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
-            (False, 0.2, [0, 1], 0.5239465),
+            (True, 0.0, -np.log(EPS)),
+            (False, 0.0, -0.6 * np.log(0.3 + EPS)),
+            (False, 0.2, -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
         ],
     )
-    def test_call(self, y_true, y_pred, binary, background_weight, scales, expected):
+    def testcall(self, y_true, y_pred, binary, background_weight, expected):
         expected = np.array([expected] * self.shape[0])  # call returns (batch, )
         got = label.CrossEntropy(
-            binary=binary, background_weight=background_weight, scales=scales
+            binary=binary,
+            background_weight=background_weight,
         ).call(y_true=y_true, y_pred=y_pred)
         assert is_equal_tf(got, expected)
 
@@ -216,9 +218,7 @@ class TestCrossEntropy:
             binary=False,
             background_weight=0.0,
             smooth=1e-5,
-            scales=None,
-            kernel="gaussian",
-            reduction=tf.keras.losses.Reduction.NONE,
+            reduction=tf.keras.losses.Reduction.AUTO,
             name="CrossEntropy",
         )
         assert got == expected
@@ -272,7 +272,7 @@ class TestJaccardIndex:
         y_true = tf.ones(shape=shape) * value
         y_pred = tf.ones(shape=shape) * value
 
-        got = label.JaccardIndex(smooth_nr=smooth_nr, smooth_dr=smooth_dr)._call(
+        got = label.JaccardIndex(smooth_nr=smooth_nr, smooth_dr=smooth_dr).call(
             y_true,
             y_pred,
         )
@@ -329,9 +329,7 @@ class TestJaccardIndex:
             background_weight=0.0,
             smooth_nr=1e-5,
             smooth_dr=1e-5,
-            scales=None,
-            kernel="gaussian",
-            reduction=tf.keras.losses.Reduction.NONE,
+            reduction=tf.keras.losses.Reduction.AUTO,
             name="JaccardIndex",
         )
         assert got == expected
diff --git a/test/unit/test_loss_util.py b/test/unit/test_loss_util.py
index 18776055..76b9e4dc 100644
--- a/test/unit/test_loss_util.py
+++ b/test/unit/test_loss_util.py
@@ -4,149 +4,73 @@
 Tests for deepreg/model/loss/label.py in
 pytest style
 """
-
 from test.unit.util import is_equal_tf
+from typing import List, Optional, Union
 
 import numpy as np
 import pytest
 import tensorflow as tf
 
-from deepreg.loss.util import (
-    NegativeLossMixin,
-    cauchy_kernel1d,
-    gaussian_kernel1d_sigma,
-    gaussian_kernel1d_size,
-    rectangular_kernel1d,
-    separable_filter,
-    triangular_kernel1d,
-)
-
-
-@pytest.mark.parametrize("sigma", [1, 3, 2.2])
-def test_cauchy_kernel1d(sigma):
-    """
-    Testing the 1-D cauchy kernel
-    :param sigma: float
-    :return:
-    """
-    tail = int(sigma * 5)
-    expected = [1 / ((x / sigma) ** 2 + 1) for x in range(-tail, tail + 1)]
-    expected = expected / np.sum(expected)
-    got = cauchy_kernel1d(sigma)
-    assert is_equal_tf(got, expected)
+from deepreg.loss.label import DiceLoss, DiceScore
+from deepreg.loss.util import MultiScaleMixin, separable_filter
 
 
-@pytest.mark.parametrize("sigma", [1, 3, 2.2])
-def test_gaussian_kernel1d_sigma(sigma):
-    """
-    Testing the 1-D gaussian kernel given sigma as input
-    :param sigma: float
-    :return:
-    """
-    tail = int(sigma * 3)
-    expected = [np.exp(-0.5 * x ** 2 / sigma ** 2) for x in range(-tail, tail + 1)]
-    expected = expected / np.sum(expected)
-    got = gaussian_kernel1d_sigma(sigma)
-    assert is_equal_tf(got, expected)
+class TestMultiScaleMixin:
+    def test_err(self):
+        with pytest.raises(ValueError) as err_info:
+            MultiScaleMixin(kernel="unknown")
+        assert "Kernel unknown is not supported." in str(err_info.value)
 
+    def test_get_config(self):
+        loss = MultiScaleMixin()
+        got = loss.get_config()
+        expected = dict(
+            scales=None,
+            kernel="gaussian",
+            reduction=tf.keras.losses.Reduction.AUTO,
+            name=None,
+        )
+        assert got == expected
 
-@pytest.mark.parametrize("kernel_size", [3, 7, 11])
-def test_gaussian_kernel1d_size(kernel_size):
-    """
-    Testing the 1-D gaussian kernel given size as input
-    :param kernel_size: int
-    :return:
-    """
-    mean = (kernel_size - 1) / 2.0
-    sigma = kernel_size / 3
+    @pytest.mark.parametrize("kernel", ["gaussian", "cauchy"])
+    @pytest.mark.parametrize("scales", [None, 0, [0], [0, 1], [1, 2]])
+    def test_call(self, kernel: str, scales: Optional[Union[List, float, int]]):
+        """
+        Test MultiScaleMixin using DiceLoss.
 
-    grid = tf.range(0, kernel_size, dtype=tf.float32)
-    expected = tf.exp(-tf.square(grid - mean) / (2 * sigma ** 2))
+        :param kernel: kernel name.
+        :param scales: scaling parameters.
+        """
+        shape = (2, 3, 4, 5)
+        y_true = tf.random.uniform(shape=shape)
+        y_pred = tf.random.uniform(shape=shape)
 
-    got = gaussian_kernel1d_size(kernel_size)
-    assert is_equal_tf(got, expected)
+        loss = DiceLoss(kernel=kernel, scales=scales)
+        loss.call(y_pred=y_pred, y_true=y_true)
 
 
-@pytest.mark.parametrize("kernel_size", [3, 7, 11])
-def test_rectangular_kernel1d(kernel_size):
-    """
-    Testing the 1-D rectangular kernel
-    :param kernel_size: int
-    :return:
-    """
-    expected = tf.ones(shape=(kernel_size,), dtype=tf.float32)
-    got = rectangular_kernel1d(kernel_size)
-    assert is_equal_tf(got, expected)
+def test_negative_loss_mixin():
+    """Test DiceScore and DiceLoss have reversed sign."""
+    shape = (2, 3, 4, 5)
+    y_true = tf.random.uniform(shape=shape)
+    y_pred = tf.random.uniform(shape=shape)
 
+    dice_score = DiceScore().call(y_pred=y_pred, y_true=y_true)
+    dice_loss = DiceLoss().call(y_pred=y_pred, y_true=y_true)
 
-@pytest.mark.parametrize("kernel_size", [3, 5, 7, 9])
-def test_triangular_kernel1d(kernel_size):
-    """
-    Testing the 1-D triangular kernel
-    :param kernel_size: int (odd number)
-    :return:
-    """
-    expected = np.zeros(shape=(kernel_size,), dtype=np.float32)
-    expected[kernel_size // 2] = kernel_size // 2 + 1
-    for it_k in range(kernel_size // 2):
-        expected[it_k] = it_k + 1
-        expected[-it_k - 1] = it_k + 1
-
-    got = triangular_kernel1d(kernel_size)
-    assert is_equal_tf(got, expected)
+    assert is_equal_tf(dice_score, -dice_loss)
 
 
 def test_separable_filter():
-    """
-    Testing separable filter case where non
-    zero length tensor is passed to the
-    function.
-    """
-    k = np.ones((3, 3, 3, 3, 1), dtype=np.float32)
-    array_eye = np.identity(3, dtype=np.float32)
-    tensor_pred = np.zeros((3, 3, 3, 3, 1), dtype=np.float32)
-    tensor_pred[:, :, 0, 0, 0] = array_eye
-    tensor_pred = tf.convert_to_tensor(tensor_pred, dtype=tf.float32)
-    k = tf.convert_to_tensor(k, dtype=tf.float32)
-
-    expect = np.ones((3, 3, 3, 3, 1), dtype=np.float32)
-    expect = tf.convert_to_tensor(expect, dtype=tf.float32)
-
-    get = separable_filter(tensor_pred, k)
-    assert is_equal_tf(get, expect)
-
-
-class MinusClass(tf.keras.losses.Loss):
-    def __init__(self):
-        super().__init__()
-        self.name = "MinusClass"
-
-    def call(self, y_true, y_pred):
-        return y_true - y_pred
-
-
-class MinusClassLoss(NegativeLossMixin, MinusClass):
-    pass
-
-
-@pytest.mark.parametrize("y_true,y_pred,expected", [(1, 2, 1), (2, 1, -1), (0, 0, 0)])
-def test_negative_loss_mixin(y_true, y_pred, expected):
-    """
-    Testing NegativeLossMixin class that
-    inverts the sign of any value
-    returned by a function
-
-    :param y_true: int
-    :param y_pred: int
-    :param expected: int
-    :return:
-    """
-
-    y_true = tf.constant(y_true, dtype=tf.float32)
-    y_pred = tf.constant(y_pred, dtype=tf.float32)
-
-    got = MinusClassLoss().call(
-        y_true,
-        y_pred,
-    )
+    """Testing separable filter case where diagonal ones are propagated."""
+    k = tf.ones(shape=(3, 3, 3, 3, 1), dtype=tf.float32)
+
+    array_eye = np.identity(3)
+    x = np.zeros((3, 3, 3, 3, 1))
+    x[:, :, 0, 0, 0] = array_eye
+    x = tf.convert_to_tensor(x, dtype=tf.float32)
+
+    expected = tf.ones(shape=(3, 3, 3, 3, 1), dtype=tf.float32)
+    got = separable_filter(x, k)
+
     assert is_equal_tf(got, expected)
diff --git a/test/unit/test_registry.py b/test/unit/test_registry.py
index 57d6de62..d5f96b07 100644
--- a/test/unit/test_registry.py
+++ b/test/unit/test_registry.py
@@ -143,6 +143,7 @@ class TestRegistry:
             expected = re.sub(r":-+", "", lines)
             expected = expected.replace(" ", "")
             expected = expected + "\n"
+
             assert got == expected
 
         # rewrite the file