Repository: ecmwf/anemoi-training (language: python)

Issue:
Losses / Validation Metrics Implementation
### Issues with Loss Implementation in `anemoi-training`

1. **Nomenclature:**
   - The current implementation uses `self.loss_weights` for weighting in the spatial dimension and `self.data_variances` for weighting in the feature dimension. However, this naming can be confusing given that the default strategy for feature dimension weighting is not based on `data_variances`.
   - **Proposal:** Rename `self.loss_weights` to `self.node_weight` (for spatial dimension) and `self.data_variances` to `self.feature_weight` (for feature dimension). Additionally, modify the function responsible for generating feature scaling values to support various feature scaling strategies.

2. **Consistency in Feature Scaling for Loss Calculation:**
   - The loss calculated by `self.metrics` does not currently apply feature scaling, whereas the loss used during training does apply the prescribed feature scaling. Previously, feature scaling was applied consistently across both metrics and training losses.
   - **Proposal:** Either revert to the previous behavior where feature scaling is applied to both, or introduce an option for users to specify whether feature scaling should be applied in `self.metrics`.

3. **Monitoring and Recording Loss Types During Validation:**
   - Currently, only `WeightedMSELoss` is available for monitoring and recording during validation. This limitation restricts flexibility in evaluating model performance.
   - **Proposal:** Extend the configuration options to allow for:
     a) Selection of the type of loss/metric to record.
     b) Choice of whether to record metrics on processed or unprocessed features when evaluating single features (This is crucial for evaluating performance of model's that use different weighting strategies)


---

### What are the steps to reproduce the bug?

NA - Inspect code

### Version

develop

### Platform (OS and architecture)

NA

### Relevant log output

_No response_

### Accompanying data

_No response_

### Organisation

ECMWF
Add ScaleTensor
## ScaleTensor
- Allows dynamic setup of scalars
- Each scalar can be assigned to operate on a dimension or set of dimensions
- Scalar is then resolved at scale time, allowing relative indices

```python
    Examples
    --------
    >>> tensor = torch.randn(3, 4, 5)
    >>> scalars = ScaleTensor((0, torch.randn(3)), (1, torch.randn(4)))
    >>> scaled_tensor = scalars.scale(tensor)
    >>> scalars.get_scalar(tensor.ndim).shape
    torch.Size([3, 4, 1])
    >>> scalars.add_scalar(-1, torch.randn(5))
    >>> scalars.get_scalar(tensor.ndim).shape
    torch.Size([3, 4, 5])
```
split feature weights from val_metrics
Split validation metrics from feature weights functionality, which will allow different groupings of variables in the loss function.

@HCookie @Rilwan-Adewoyin 


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bb550b02..f8d107fe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,9 @@ Keep it human-readable, your future self will thank you!
 
 ## [0.2.2 - Maintenance: pin python <3.13](https://github.com/ecmwf/anemoi-training/compare/0.2.1...0.2.2) - 2024-10-28
 
+### Added
+- Included more loss functions and allowed configuration [#70](https://github.com/ecmwf/anemoi-training/pull/70)
+
 ### Changed
 
 - Lock python version <3.13 [#107](https://github.com/ecmwf/anemoi-training/pull/107)
diff --git a/docs/modules/losses.rst b/docs/modules/losses.rst
index 73045d62..07ce18f2 100644
--- a/docs/modules/losses.rst
+++ b/docs/modules/losses.rst
@@ -3,22 +3,153 @@
 ########
 
 This module is used to define the loss function used to train the model.
+
+Anemoi-training exposes a couple of loss functions by default to be
+used, all of which are subclassed from ``BaseWeightedLoss``. This class
+enables scalar multiplication, and graph node weighting.
+
+.. automodule:: anemoi.training.losses.weightedloss
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+************************
+ Default Loss Functions
+************************
+
 By default anemoi-training trains the model using a latitude-weighted
 mean-squared-error, which is defined in the ``WeightedMSELoss`` class in
-``aifs/losses/mse.py``.
+``anemoi/training/losses/mse.py``. The loss function can be configured
+in the config file at ``config.training.training_loss``, and
+``config.training.validation_metrics``.
+
+The following loss functions are available by default:
+
+-  ``WeightedMSELoss``: Latitude-weighted mean-squared-error.
+-  ``WeightedMAELoss``: Latitude-weighted mean-absolute-error.
+-  ``WeightedHuberLoss``: Latitude-weighted Huber loss.
+-  ``WeightedLogCoshLoss``: Latitude-weighted log-cosh loss.
+-  ``WeightedRMSELoss``: Latitude-weighted root-mean-squared-error.
+-  ``CombinedLoss``: Combined component weighted loss.
+
+These are available in the ``anemoi.training.losses`` module, at
+``anemoi.training.losses.{short_name}.{class_name}``.
+
+So for example, to use the ``WeightedMSELoss`` class, you would
+reference it in the config as follows:
+
+.. code:: yaml
+
+   # loss function for the model
+   training_loss:
+      # loss class to initialise
+      _target_: anemoi.training.losses.mse.WeightedMSELoss
+      # loss function kwargs here
+
+*********
+ Scalars
+*********
+
+In addition to node scaling, the loss function can also be scaled by a
+scalar. These are provided by the ``Forecaster`` class, and a user can
+define whether to include them in the loss function by setting
+``scalars`` in the loss config dictionary.
+
+.. code:: yaml
+
+   # loss function for the model
+   training_loss:
+      # loss class to initialise
+      _target_: anemoi.training.losses.mse.WeightedMSELoss
+      scalars: ['scalar1', 'scalar2']
+
+Currently, the following scalars are available for use:
+
+-  ``variable``: Scale by the feature/variable weights as defined in the
+   config ``config.training.loss_scaling``.
 
-The user can define their own loss function using the same structure as
-the ``WeightedMSELoss`` class.
+********************
+ Validation Metrics
+********************
 
-.. automodule:: anemoi.training.losses.mse
+Validation metrics as defined in the config file at
+``config.training.validation_metrics`` follow the same initialise
+behaviour as the loss function, but can be a list. In this case all
+losses are calculated and logged as a dictionary with the corresponding
+name
+
+***********************
+ Custom Loss Functions
+***********************
+
+Additionally, you can define your own loss function by subclassing
+``BaseWeightedLoss`` and implementing the ``forward`` method, or by
+subclassing ``FunctionalWeightedLoss`` and implementing the
+``calculate_difference`` function. The latter abstracts the scaling, and
+node weighting, and allows you to just specify the difference
+calculation.
+
+.. code:: python
+
+   from anemoi.training.losses.weightedloss import FunctionalWeightedLoss
+
+   class MyLossFunction(FunctionalWeightedLoss):
+      def calculate_difference(self, pred, target):
+         return (pred - target) ** 2
+
+Then in the config, set ``_target_`` to the class name, and any
+additional kwargs to the loss function.
+
+*****************
+ Combined Losses
+*****************
+
+Building on the simple single loss functions, a user can define a
+combined loss, one that weights and combines multiple loss functions.
+
+This can be done by referencing the ``CombinedLoss`` class in the config
+file, and setting the ``losses`` key to a list of loss functions to
+combine. Each of those losses is then initalised just like the other
+losses above.
+
+.. code:: yaml
+
+   training_loss:
+      __target__: anemoi.training.losses.combined.CombinedLoss
+      losses:
+         - __target__: anemoi.training.losses.mse.WeightedMSELoss
+         - __target__: anemoi.training.losses.mae.WeightedMAELoss
+      scalars: ['variable']
+      loss_weights: [1.0,0.5]
+
+All kwargs passed to ``CombinedLoss`` are passed to each of the loss
+functions, and the loss weights are used to scale the individual losses
+before combining them.
+
+.. automodule:: anemoi.training.losses.combined
    :members:
    :no-undoc-members:
    :show-inheritance:
 
+*******************
+ Utility Functions
+*******************
+
 There is also generic functions that are useful for losses in
-``aifs/losses/utils.py``.
+``anemoi/training/losses/utils.py``.
 
 ``grad_scaler`` is used to automatically scale the loss gradients in the
 loss function using the formula in https://arxiv.org/pdf/2306.06079.pdf,
 section 4.3.2. This can be switched on in the config by setting the
 option ``config.training.loss_gradient_scaling=True``.
+
+``ScaleTensor`` is a class that can record and apply arbitrary scaling
+factors to tensors. It supports relative indexing, combining multiple
+scalars over the same dimensions, and is only constructed at
+broadcasting time, so the shape can be resolved to match the tensor
+exactly.
+
+.. automodule:: anemoi.training.losses.utils
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
diff --git a/src/anemoi/training/config/training/default.yaml b/src/anemoi/training/config/training/default.yaml
index 870eeb7a..74070fcd 100644
--- a/src/anemoi/training/config/training/default.yaml
+++ b/src/anemoi/training/config/training/default.yaml
@@ -33,11 +33,35 @@ swa:
 # use ZeroRedundancyOptimizer ; saves memory for larger models
 zero_optimizer: False
 
+# loss functions
+
 # dynamic rescaling of the loss gradient
 # see https://arxiv.org/pdf/2306.06079.pdf, section 4.3.2
 # don't enable this by default until it's been tested and proven beneficial
+
+# loss function for the model
+training_loss:
+  # loss class to initialise
+  _target_: anemoi.training.losses.mse.WeightedMSELoss
+  # Scalars to include in loss calculation
+  # Available scalars include, 'variable'
+  scalars: ['variable']
+  ignore_nans: False
+
 loss_gradient_scaling: False
 
+# Validation metrics calculation,
+# This may be a list, in which case all metrics will be calculated
+# and logged according to their name
+validation_metrics:
+  # loss class to initialise
+  - _target_: anemoi.training.losses.mse.WeightedMSELoss
+    # Scalars to include in loss calculation
+    # Available scalars include, 'variable'
+    scalars: []
+    # other kwargs
+    ignore_nans: True
+
 # length of the "rollout" window (see Keisler's paper)
 rollout:
   start: 1
diff --git a/src/anemoi/training/diagnostics/callbacks/__init__.py b/src/anemoi/training/diagnostics/callbacks/__init__.py
index cf085eab..72d4d242 100644
--- a/src/anemoi/training/diagnostics/callbacks/__init__.py
+++ b/src/anemoi/training/diagnostics/callbacks/__init__.py
@@ -44,6 +44,7 @@
 from anemoi.training.diagnostics.plots import plot_loss
 from anemoi.training.diagnostics.plots import plot_power_spectrum
 from anemoi.training.diagnostics.plots import plot_predicted_multilevel_flat_sample
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
 
 if TYPE_CHECKING:
     import pytorch_lightning as pl
@@ -605,6 +606,11 @@ def _plot(
         # reorder parameter_names by position
         self.parameter_names = [parameter_names[i] for i in np.argsort(parameter_positions)]
 
+        if not isinstance(pl_module.loss, BaseWeightedLoss):
+            logging.warning(
+                "Loss function must be a subclass of BaseWeightedLoss, or provide `squash`.", RuntimeWarning
+            )
+
         batch = pl_module.model.pre_processors(batch, in_place=False)
         for rollout_step in range(pl_module.rollout):
             y_hat = outputs[1][rollout_step]
diff --git a/src/anemoi/training/losses/combined.py b/src/anemoi/training/losses/combined.py
new file mode 100644
index 00000000..11d2b4fe
--- /dev/null
+++ b/src/anemoi/training/losses/combined.py
@@ -0,0 +1,135 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+from typing import Callable
+
+import torch
+
+from anemoi.training.train.forecaster import GraphForecaster
+
+
+class CombinedLoss(torch.nn.Module):
+    """Combined Loss function."""
+
+    def __init__(
+        self,
+        *extra_losses: dict[str, Any] | Callable,
+        losses: tuple[dict[str, Any] | Callable] | None = None,
+        loss_weights: tuple[int, ...],
+        **kwargs,
+    ):
+        """Combined loss function.
+
+        Allows multiple losses to be combined into a single loss function,
+        and the components weighted.
+
+        If a sub loss function requires additional weightings or code created tensors,
+        that must be `included_` for this function, and then controlled by the underlying
+        loss function configuration.
+
+        Parameters
+        ----------
+        losses: tuple[dict[str, Any]| Callable]
+            Tuple of losses to initialise with `GraphForecaster.get_loss_function`.
+            Allows for kwargs to be passed, and weighings controlled.
+        *extra_losses: dict[str, Any] | Callable
+            Additional arg form of losses to include in the combined loss.
+        loss_weights : tuple[int, ...]
+            Weights of each loss function in the combined loss.
+        kwargs: Any
+            Additional arguments to pass to the loss functions
+
+        Examples
+        --------
+        >>> CombinedLoss(
+                {"__target__": "anemoi.training.losses.mse.WeightedMSELoss"},
+                loss_weights=(1.0,),
+                node_weights=node_weights
+            )
+        --------
+        >>> CombinedLoss(
+                losses = [anemoi.training.losses.mse.WeightedMSELoss],
+                loss_weights=(1.0,),
+                node_weights=node_weights
+            )
+        Or from the config,
+
+        ```
+        training_loss:
+            __target__: anemoi.training.losses.combined.CombinedLoss
+            losses:
+                - __target__: anemoi.training.losses.mse.WeightedMSELoss
+                - __target__: anemoi.training.losses.mae.WeightedMAELoss
+            scalars: ['variable']
+            loss_weights: [1.0,0.5]
+        ```
+        """
+        super().__init__()
+
+        losses = (*(losses or []), *extra_losses)
+
+        assert len(losses) == len(loss_weights), "Number of losses and weights must match"
+        assert len(losses) > 0, "At least one loss must be provided"
+
+        self.losses = [
+            GraphForecaster.get_loss_function(loss, **kwargs) if isinstance(loss, dict) else loss(**kwargs)
+            for loss in losses
+        ]
+        self.loss_weights = loss_weights
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Calculates the combined loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        kwargs: Any
+            Additional arguments to pass to the loss functions
+            Will be passed to all loss functions
+
+        Returns
+        -------
+        torch.Tensor
+            Combined loss
+        """
+        loss = None
+        for i, loss_fn in enumerate(self.losses):
+            if loss is not None:
+                loss += self.loss_weights[i] * loss_fn(pred, target, **kwargs)
+            else:
+                loss = self.loss_weights[i] * loss_fn(pred, target, **kwargs)
+        return loss
+
+    @property
+    def name(self) -> str:
+        return "combined_" + "_".join(getattr(loss, "name", loss.__class__.__name__.lower()) for loss in self.losses)
+
+    def __getattr__(self, name: str) -> Callable:
+        """Allow access to underlying attributes of the loss functions."""
+        if not all(hasattr(loss, name) for loss in self.losses):
+            error_msg = f"Attribute {name} not found in all loss functions"
+            raise AttributeError(error_msg)
+
+        @functools.wraps(getattr(self.losses[0], name))
+        def hidden_func(*args, **kwargs) -> list[Any]:
+            return [getattr(loss, name)(*args, **kwargs) for loss in self.losses]
+
+        return hidden_func
diff --git a/src/anemoi/training/losses/huber.py b/src/anemoi/training/losses/huber.py
new file mode 100644
index 00000000..e42105c7
--- /dev/null
+++ b/src/anemoi/training/losses/huber.py
@@ -0,0 +1,103 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
+
+LOGGER = logging.getLogger(__name__)
+
+
+class WeightedHuberLoss(BaseWeightedLoss):
+    """Node-weighted Huber loss."""
+
+    name = "whuber"
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        delta: float = 1.0,
+        ignore_nans: bool = False,
+        **kwargs,
+    ) -> None:
+        """Node- and feature weighted Huber Loss.
+
+        See `Huber loss <https://en.wikipedia.org/wiki/Huber_loss>`_ for more information.
+
+        Parameters
+        ----------
+        node_weights : torch.Tensor of shape (N, )
+            Weight of each node in the loss function
+        delta : float, optional
+            Threshold for Huber loss, by default 1.0
+        ignore_nans : bool, optional
+            Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
+        """
+        super().__init__(
+            node_weights=node_weights,
+            ignore_nans=ignore_nans,
+            **kwargs,
+        )
+        self.delta = delta
+
+    def huber(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """Calculate the Huber loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+
+        Returns
+        -------
+        torch.Tensor
+            Huber loss
+        """
+        diff = torch.abs(pred - target)
+        return torch.where(diff < self.delta, 0.5 * torch.square(diff), self.delta * (diff - 0.5 * self.delta))
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted Huber loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted Huber loss
+        """
+        out = self.huber(pred, target)
+
+        if feature_scale:
+            out = self.scale_by_variable_scaling(out, feature_indices)
+        return self.scale_by_node_weights(out, squash)
diff --git a/src/anemoi/training/losses/logcosh.py b/src/anemoi/training/losses/logcosh.py
new file mode 100644
index 00000000..6112d472
--- /dev/null
+++ b/src/anemoi/training/losses/logcosh.py
@@ -0,0 +1,98 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import torch
+
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
+
+LOGGER = logging.getLogger(__name__)
+
+
+class LogCosh(torch.autograd.Function):
+    """LogCosh custom autograd function."""
+
+    @staticmethod
+    def forward(ctx, inp: torch.Tensor) -> torch.Tensor:  # noqa: ANN001
+        ctx.save_for_backward(inp)
+        abs_input = torch.abs(inp)
+        return abs_input + torch.nn.functional.softplus(-2 * abs_input) - np.log(2)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:  # noqa: ANN001
+        (inp,) = ctx.saved_tensors
+        return grad_output * torch.tanh(inp)
+
+
+class WeightedLogCoshLoss(BaseWeightedLoss):
+    """Node-weighted LogCosh loss."""
+
+    name = "wlogcosh"
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        ignore_nans: bool = False,
+        **kwargs,
+    ) -> None:
+        """Node- and feature weighted LogCosh Loss.
+
+        Parameters
+        ----------
+        node_weights : torch.Tensor of shape (N, )
+            Weight of each node in the loss function
+        ignore_nans : bool, optional
+            Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
+
+        """
+        super().__init__(
+            node_weights=node_weights,
+            ignore_nans=ignore_nans,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted LogCosh loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted LogCosh loss
+
+        """
+        out = LogCosh.apply(pred - target)
+
+        if feature_scale:
+            out = self.scale(out, feature_indices)
+        return self.scale_by_node_weights(out, squash)
diff --git a/src/anemoi/training/losses/mae.py b/src/anemoi/training/losses/mae.py
new file mode 100644
index 00000000..bea16ac2
--- /dev/null
+++ b/src/anemoi/training/losses/mae.py
@@ -0,0 +1,83 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
+
+LOGGER = logging.getLogger(__name__)
+
+
+class WeightedMAELoss(BaseWeightedLoss):
+    """Node-weighted MAE loss."""
+
+    name = "wmae"
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        ignore_nans: bool = False,
+        **kwargs,
+    ) -> None:
+        """Node- and feature weighted MAE Loss.
+
+        Also known as the Weighted L1 loss.
+
+        Parameters
+        ----------
+        node_weights : torch.Tensor of shape (N, )
+            Weight of each node in the loss function
+        ignore_nans : bool, optional
+            Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
+
+        """
+        super().__init__(
+            node_weights=node_weights,
+            ignore_nans=ignore_nans,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted MAE loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted MAE loss
+        """
+        out = torch.abs(pred - target)
+
+        if feature_scale:
+            out = self.scale(out, feature_indices)
+        return self.scale_by_node_weights(out, squash)
diff --git a/src/anemoi/training/losses/mse.py b/src/anemoi/training/losses/mse.py
index 88ad0d0b..87365f8c 100644
--- a/src/anemoi/training/losses/mse.py
+++ b/src/anemoi/training/losses/mse.py
@@ -13,80 +13,69 @@
 import logging
 
 import torch
-from torch import nn
+
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
 
 LOGGER = logging.getLogger(__name__)
 
 
-class WeightedMSELoss(nn.Module):
-    """Latitude-weighted MSE loss."""
+class WeightedMSELoss(BaseWeightedLoss):
+    """Node-weighted MSE loss."""
+
+    name = "wmse"
 
     def __init__(
         self,
         node_weights: torch.Tensor,
-        data_variances: torch.Tensor | None = None,
-        ignore_nans: bool | None = False,
+        ignore_nans: bool = False,
+        **kwargs,
     ) -> None:
-        """Latitude- and (inverse-)variance-weighted MSE Loss.
+        """Node- and feature weighted MSE Loss.
 
         Parameters
         ----------
         node_weights : torch.Tensor of shape (N, )
             Weight of each node in the loss function
-        data_variances : Optional[torch.Tensor], optional
-            precomputed, per-variable stepwise variance estimate, by default None
         ignore_nans : bool, optional
             Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
 
         """
-        super().__init__()
-
-        self.avg_function = torch.nanmean if ignore_nans else torch.mean
-        self.sum_function = torch.nansum if ignore_nans else torch.sum
-
-        self.register_buffer("weights", node_weights, persistent=True)
-        if data_variances is not None:
-            self.register_buffer("ivar", data_variances, persistent=True)
+        super().__init__(
+            node_weights=node_weights,
+            ignore_nans=ignore_nans,
+            **kwargs,
+        )
 
     def forward(
         self,
         pred: torch.Tensor,
         target: torch.Tensor,
         squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
     ) -> torch.Tensor:
         """Calculates the lat-weighted MSE loss.
 
         Parameters
         ----------
         pred : torch.Tensor
-            Prediction tensor, shape (bs, lat*lon, n_outputs)
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
         target : torch.Tensor
-            Target tensor, shape (bs, lat*lon, n_outputs)
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
         squash : bool, optional
             Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
 
         Returns
         -------
         torch.Tensor
             Weighted MSE loss
-
         """
         out = torch.square(pred - target)
 
-        # Use variances if available
-        if hasattr(self, "ivar"):
-            out *= self.ivar
-
-        # Squash by last dimension
-        if squash:
-            out = self.avg_function(out, dim=-1)
-            # Weight by area
-            out *= self.weights.expand_as(out)
-            out /= self.sum_function(self.weights.expand_as(out))
-            return self.sum_function(out)
-
-        # Weight by area
-        out *= self.weights[..., None].expand_as(out)
-        # keep last dimension (variables) when summing weights
-        out /= self.sum_function(self.weights[..., None].expand_as(out), axis=(0, 1, 2))
-        return self.sum_function(out, axis=(0, 1, 2))
+        if feature_scale:
+            out = self.scale(out, feature_indices)
+        return self.scale_by_node_weights(out, squash)
diff --git a/src/anemoi/training/losses/rmse.py b/src/anemoi/training/losses/rmse.py
new file mode 100644
index 00000000..34b913a9
--- /dev/null
+++ b/src/anemoi/training/losses/rmse.py
@@ -0,0 +1,83 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from anemoi.training.losses.mse import BaseWeightedLoss
+
+LOGGER = logging.getLogger(__name__)
+
+
+class WeightedRMSELoss(BaseWeightedLoss):
+    """Node-weighted RMSE loss."""
+
+    name = "wrmse"
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        ignore_nans: bool = False,
+        **kwargs,
+    ) -> None:
+        """Node- and (inverse-)variance-weighted RMSE Loss.
+
+        Parameters
+        ----------
+        node_weights : torch.Tensor of shape (N, )
+            Weight of each node in the loss function
+        ignore_nans : bool, optional
+            Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
+        """
+        super().__init__(
+            node_weights=node_weights,
+            ignore_nans=ignore_nans,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted RMSE loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted RMSE loss
+        """
+        mse = super().forward(
+            pred=pred,
+            target=target,
+            squash=squash,
+            feature_indices=feature_indices,
+            feature_scale=feature_scale,
+        )
+        return torch.sqrt(mse)
diff --git a/src/anemoi/training/losses/utils.py b/src/anemoi/training/losses/utils.py
index 5ddef3d6..5e7c98a8 100644
--- a/src/anemoi/training/losses/utils.py
+++ b/src/anemoi/training/losses/utils.py
@@ -11,10 +11,17 @@
 from __future__ import annotations
 
 import logging
+import uuid
+from typing import TYPE_CHECKING
+from typing import Callable
+from typing import Union
 
 import torch
 from torch import nn
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 LOGGER = logging.getLogger(__name__)
 
 
@@ -27,7 +34,7 @@ def grad_scaler(
 
     Uses the formula in https://arxiv.org/pdf/2306.06079.pdf, section 4.3.2
 
-    Use <module>.register_full_backward_hook(grad_scaler, prepend=False) to register this hook.
+    Use <module>.register_full_backward_hook(grad_scalar, prepend=False) to register this hook.
 
     Parameters
     ----------
@@ -52,3 +59,329 @@ def grad_scaler(
         (channels * channel_weights) / torch.sum(channel_weights, dim=-1, keepdim=True) * grad_in[0]
     )  # rescaled gradient
     return new_grad_in, grad_in[1]
+
+
+TENSOR_SPEC = tuple[Union[int, tuple[int]], torch.Tensor]
+
+
+class Shape:
+    """Shape resolving object."""
+
+    def __init__(self, func: Callable[[int], int]):
+        self.func = func
+
+    def __getitem__(self, dimension: int) -> int:
+        return self.func(dimension)
+
+
+class ScaleTensor:
+    """Dynamically resolved tensor scaling class.
+
+    Allows a user to specify a scalar and the dimensions it should be applied to.
+    The class will then enforce that additional scalars are compatible with the specified dimensions.
+
+    When `get_scalar` or `scale` is called, the class will return the product of all scalars, resolved
+    to the dimensional size of the input tensor.
+
+    Additionally, the class can be subsetted to only return a subset of the scalars, but
+    only from those given names.
+
+    Examples
+    --------
+    >>> tensor = torch.randn(3, 4, 5)
+    >>> scalars = ScaleTensor((0, torch.randn(3)), (1, torch.randn(4)))
+    >>> scaled_tensor = scalars.scale(tensor)
+    >>> scalars.get_scalar(tensor.ndim).shape
+    torch.Size([3, 4, 1])
+    >>> scalars.add_scalar(-1, torch.randn(5))
+    >>> scalars.get_scalar(tensor.ndim).shape
+    torch.Size([3, 4, 5])
+    """
+
+    tensors: dict[str, TENSOR_SPEC]
+    _specified_dimensions: list[tuple[int]]
+
+    def __init__(
+        self,
+        scalars: dict[str, TENSOR_SPEC] | TENSOR_SPEC | None = None,
+        *tensors: TENSOR_SPEC,
+        **named_tensors: dict[str, TENSOR_SPEC],
+    ):
+        """ScaleTensor constructor.
+
+        Parameters
+        ----------
+        scalars : dict[str, TENSOR_SPEC] | TENSOR_SPEC | None, optional
+            Scalars to initalise with, by default None
+        tensors : TENSOR_SPEC
+            Args form of (dimension, tensor) to add to the scalars
+            Will be given a random uuid name
+        named_tensors : dict[str, TENSOR_SPEC]
+            Kwargs form of {name: (dimension, tensor)} to add to the scalars
+        """
+        self.tensors = {}
+        self._specified_dimensions = []
+
+        scalars = scalars or {}
+        scalars.update(named_tensors)
+
+        for name, tensor_spec in scalars.items():
+            self.add_scalar(*tensor_spec, name=name)
+
+        for tensor_spec in tensors:
+            self.add_scalar(*tensor_spec)
+
+    @property
+    def shape(self) -> Shape:
+        """Get the shape of the scale tensor.
+
+        Returns a Shape object to be indexed,
+        Will only resolve those dimensions specified in the tensors.
+        """
+
+        def get_dim_shape(dimension: int) -> int:
+            for dim_assign, tensor in self.tensors.values():
+                if isinstance(dim_assign, tuple) and dimension in dim_assign:
+                    return tensor.shape[list(dim_assign).index(dimension)]
+
+            error_msg = (
+                f"Could not find shape of dimension {dimension} with tensors in dims {list(self.tensors.keys())}"
+            )
+            raise IndexError(error_msg)
+
+        return Shape(get_dim_shape)
+
+    def validate_scalar(self, dimension: int | tuple[int], scalar: torch.Tensor) -> None:
+        """Check if the scalar is compatible with the given dimension.
+
+        Parameters
+        ----------
+        dimension : int | tuple[int]
+            Dimensions to check `scalar` against
+        scalar : torch.Tensor
+            Scalar tensor to check
+
+        Raises
+        ------
+        ValueError
+            If the scalar is not compatible with the given dimension
+        """
+        if isinstance(dimension, int):
+            dimension = [dimension]
+
+        for scalar_dim, dim in enumerate(dimension):
+            if dim not in self or scalar.shape[scalar_dim] == 1 or self.shape[dim] == 1:
+                continue
+
+            if self.shape[dim] != scalar.shape[scalar_dim]:
+                error_msg = (
+                    f"Scalar shape {scalar.shape} at dimension {scalar_dim}"
+                    f"does not match shape of scalar at dimension {dim}. Expected {self.shape[dim]}",
+                )
+                raise ValueError(error_msg)
+
+    def add_scalar(
+        self,
+        dimension: int | tuple[int],
+        scalar: torch.Tensor,
+        *,
+        name: str | None = None,
+    ) -> None:
+        """Add new scalar to be applied along `dimension`.
+
+        Dimension can be a single int even for a multi-dimensional scalar,
+        in this case the dimensions are assigned as a range from the given int.
+        Negative indexes are also valid, and will be resolved against the tensor's ndim.
+
+        Parameters
+        ----------
+        dimension : int | tuple[int]
+            Dimension/s to apply the scalar to
+        scalar : torch.Tensor
+            Scalar tensor to apply
+        name : str | None, optional
+            Name of the scalar, by default None
+        """
+        if not isinstance(scalar, torch.Tensor):
+            scalar = torch.tensor([scalar]) if isinstance(scalar, (int, float)) else torch.tensor(scalar)
+
+        if isinstance(dimension, int):
+            if len(scalar.shape) == 1:
+                dimension = (dimension,)
+            else:
+                dimension = tuple(dimension + i for i in range(len(scalar.shape)))
+
+        try:
+            self.validate_scalar(dimension, scalar)
+        except ValueError as e:
+            error_msg = f"Validating tensor {name!r} raised an error."
+            raise ValueError(error_msg) from e
+
+        if name is None:
+            name = str(uuid.uuid4())
+
+        if name in self.tensors:
+            self._specified_dimensions.remove(self.tensors[name][0])
+            self.tensors[name] = (dimension, self.tensors[name][1] * scalar)
+        else:
+            self.tensors[name] = (dimension, scalar)
+        self._specified_dimensions.append(dimension)
+
+    def subset(self, scalars: str | Sequence[str]) -> ScaleTensor:
+        """Get subset of the scalars, filtering by name.
+
+        See `.subset_by_dim` for subsetting by affected dimensions.
+
+        Parameters
+        ----------
+        scalars : str | Sequence[str]
+            Name/s of the scalars to get
+
+        Returns
+        -------
+        ScaleTensor
+            Subset of self
+        """
+        if isinstance(scalars, str):
+            scalars = [scalars]
+        return ScaleTensor(**{name: self.tensors[name] for name in scalars})
+
+    def subset_by_dim(self, dimensions: int | Sequence[int]) -> ScaleTensor:
+        """Get subset of the scalars, filtering by dimension.
+
+        See `.subset` for subsetting by name.
+
+        Parameters
+        ----------
+        dimensions : int | Sequence[int]
+            Dimensions to get scalars of
+
+        Returns
+        -------
+        ScaleTensor
+            Subset of self
+        """
+        subset_scalars: dict[str, TENSOR_SPEC] = {}
+
+        if isinstance(dimensions, int):
+            dimensions = (dimensions,)
+
+        for name, (dim, scalar) in self.tensors.items():
+            if isinstance(dim, int):
+                dim = (dim,)
+            if len(set(dimensions).intersection(dim)) > 0:
+                subset_scalars[name] = (dim, scalar)
+
+        return ScaleTensor(**subset_scalars)
+
+    def resolve(self, ndim: int) -> ScaleTensor:
+        """Resolve relative indexes in scalars by associating against ndim.
+
+        i.e. if a scalar was given as effecting dimension -1,
+        and `ndim` was provided as 4, the scalar will be fixed
+        to effect dimension 3.
+
+        Parameters
+        ----------
+        ndim : int
+            Number of dimensions to resolve relative indexing against
+
+        Returns
+        -------
+        ScaleTensor
+            ScaleTensor with all relative indexes resolved
+        """
+        resolved_scalars: dict[str, TENSOR_SPEC] = {}
+
+        for name, (dims, scalar) in self.tensors.items():
+            if any(d < 0 for d in dims):
+                dims = [d if d >= 0 else ndim + d for d in dims]
+            resolved_scalars[name] = (dims, scalar)
+
+        return ScaleTensor(**resolved_scalars)
+
+    def scale(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Scale a given tensor by the scalars.
+
+        Parameters
+        ----------
+        tensor : torch.Tensor
+            Input tensor to scale
+
+        Returns
+        -------
+        torch.Tensor
+            Scaled tensor
+        """
+        return tensor * self.get_scalar(tensor.ndim)
+
+    def get_scalar(self, ndim: int, device: str | None = None) -> torch.Tensor:
+        """Get completely resolved scalar tensor.
+
+        Parameters
+        ----------
+        ndim : int
+            Number of dimensions of the tensor to resolve the scalars to
+            Used to resolve relative indices, and add singleton dimensions
+        device: str | None, optional
+            Device to move the scalar to, by default None
+
+        Returns
+        -------
+        torch.Tensor
+            Scalar tensor
+
+        Raises
+        ------
+        ValueError
+            If resolving relative indices is invalid
+        """
+        complete_scalar = None
+
+        tensors = self.resolve(ndim).tensors
+
+        for dims, scalar in tensors.values():
+            missing_dims = [d for d in range(ndim) if d not in dims]
+            reshape = [1] * len(missing_dims)
+            reshape.extend(scalar.shape)
+
+            reshaped_scalar = scalar.reshape(reshape)
+            reshaped_scalar = torch.moveaxis(reshaped_scalar, list(range(ndim)), (*missing_dims, *dims))
+
+            complete_scalar = reshaped_scalar if complete_scalar is None else complete_scalar * reshaped_scalar
+
+        complete_scalar = torch.ones(1) if complete_scalar is None else complete_scalar
+
+        if device is not None:
+            return complete_scalar.to(device)
+        return complete_scalar
+
+    def to(self, *args, **kwargs) -> None:
+        """Move scalars inplace."""
+        for name, (dims, tensor) in self.tensors.items():
+            self.tensors[name] = (dims, tensor.to(*args, **kwargs))
+
+    def __mul__(self, tensor: torch.Tensor) -> torch.Tensor:
+        return self.scale(tensor)
+
+    def __repr__(self):
+        return f"ScalarTensor:\n - With {list(self.tensors.keys())}\n - With dims: {self._specified_dimensions}"
+
+    def __contains__(self, dimension: int | tuple[int] | str) -> bool:
+        """Check if either scalar by name or dimension by int/tuple is being scaled."""
+        if isinstance(dimension, tuple):
+            return dimension in self._specified_dimensions
+        if isinstance(dimension, str):
+            return dimension in self.tensors
+
+        result = False
+        for dim_assign, _ in self.tensors.values():
+            result = dimension in dim_assign or result
+        return result
+
+    def __len__(self):
+        return len(self.tensors)
+
+    def __iter__(self):
+        """Iterate over tensors."""
+        return iter(self.tensors)
diff --git a/src/anemoi/training/losses/weightedloss.py b/src/anemoi/training/losses/weightedloss.py
new file mode 100644
index 00000000..81829b66
--- /dev/null
+++ b/src/anemoi/training/losses/weightedloss.py
@@ -0,0 +1,216 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+from __future__ import annotations
+
+import functools
+import logging
+from abc import ABC
+from abc import abstractmethod
+
+import torch
+from torch import nn
+
+from anemoi.training.losses.utils import ScaleTensor
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BaseWeightedLoss(nn.Module, ABC):
+    """Node-weighted general loss."""
+
+    scalar: ScaleTensor
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        ignore_nans: bool = False,
+    ) -> None:
+        """Node- and feature_weighted Loss.
+
+        Exposes:
+        - self.avg_function: torch.nanmean or torch.mean
+        - self.sum_function: torch.nansum or torch.sum
+        depending on the value of `ignore_nans`
+
+        Registers:
+        - self.node_weights: torch.Tensor of shape (N, )
+
+        Parameters
+        ----------
+        node_weights : torch.Tensor of shape (N, )
+            Weight of each node in the loss function
+        ignore_nans : bool, optional
+            Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
+
+        """
+        super().__init__()
+
+        self.scalar = ScaleTensor()
+
+        self.avg_function = torch.nanmean if ignore_nans else torch.mean
+        self.sum_function = torch.nansum if ignore_nans else torch.sum
+
+        self.register_buffer("node_weights", node_weights, persistent=True)
+
+    @functools.wraps(ScaleTensor.add_scalar, assigned=("__doc__", "__annotations__"))
+    def add_scalar(self, dimension: int | tuple[int], scalar: torch.Tensor, *, name: str | None = None) -> None:
+        self.scalar.add_scalar(dimension=dimension, scalar=scalar, name=name)
+
+    def scale(
+        self,
+        x: torch.Tensor,
+        feature_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Scale a tensor by the variable_scaling.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Tensor to be scaled, shape (bs, ensemble, lat*lon, n_outputs)
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+
+        Returns
+        -------
+        torch.Tensor
+            Scaled error tensor
+        """
+        # Use feature_weights if available
+        if len(self.scalar) == 0:
+            return x
+
+        scalar = self.scalar.get_scalar(x.ndim).to(x)
+
+        if feature_indices is None:
+            return x * scalar
+        return x * scalar[..., feature_indices]
+
+    def scale_by_node_weights(self, x: torch.Tensor, squash: bool = True) -> torch.Tensor:
+        """Scale a tensor by the node_weights.
+
+        Equivalent to reducing and averaging accordingly across all
+        dimensions of the tensor.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Tensor to be scaled, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+            If False, the loss returned of shape (n_outputs)
+
+        Returns
+        -------
+        torch.Tensor
+            Scaled error tensor
+        """
+        # Squash by last dimension
+        if squash:
+            x = self.avg_function(x, dim=-1)
+            # Weight by area
+            x *= self.node_weights.expand_as(x)
+            x /= self.sum_function(self.node_weights.expand_as(x))
+            return self.sum_function(x)
+
+        # Weight by area, due to weighting construction is analagous to a mean
+        x *= self.node_weights[..., None].expand_as(x)
+        # keep last dimension (variables) when summing weights
+        x /= self.sum_function(self.node_weights[..., None].expand_as(x), dim=(0, 1, 2))
+        return self.sum_function(x, dim=(0, 1, 2))
+
+    @abstractmethod
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted scaled loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted loss
+        """
+        out = pred - target
+
+        if feature_scale:
+            out = self.scale(out, feature_indices)
+        return self.scale_by_node_weights(out, squash)
+
+    @property
+    def name(self) -> str:
+        """Used for logging identification purposes."""
+        return self.__class__.__name__.lower()
+
+
+class FunctionalWeightedLoss(BaseWeightedLoss):
+    """WeightedLoss which a user can subclass and provide `calculate_difference`."""
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        ignore_nans: bool = False,
+    ) -> None:
+        super().__init__(node_weights, ignore_nans)
+
+    @abstractmethod
+    def calculate_difference(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """Calculate Difference between prediction and target."""
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+        feature_indices: torch.Tensor | None = None,
+        feature_scale: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted scaled loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, ensemble, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+        feature_indices:
+            feature indices (relative to full model output) of the features passed in pred and target
+        feature_scale:
+            If True, scale the loss by the feature_weights
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted loss
+        """
+        out = self.calculate_difference(pred, target)
+
+        if feature_scale:
+            out = self.scale(out, feature_indices)
+        return self.scale_by_node_weights(out, squash)
diff --git a/src/anemoi/training/train/forecaster.py b/src/anemoi/training/train/forecaster.py
index 36f6fa9a..8fae16b3 100644
--- a/src/anemoi/training/train/forecaster.py
+++ b/src/anemoi/training/train/forecaster.py
@@ -13,6 +13,7 @@
 import os
 from collections import defaultdict
 from collections.abc import Mapping
+from typing import Union
 
 import numpy as np
 import pytorch_lightning as pl
@@ -29,8 +30,8 @@
 from torch.utils.checkpoint import checkpoint
 from torch_geometric.data import HeteroData
 
-from anemoi.training.losses.mse import WeightedMSELoss
 from anemoi.training.losses.utils import grad_scaler
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
 from anemoi.training.utils.jsonify import map_config_to_primitives
 from anemoi.training.utils.masks import Boolean1DMask
 from anemoi.training.utils.masks import NoOutputMask
@@ -83,22 +84,35 @@ def __init__(
         self.save_hyperparameters()
 
         self.latlons_data = graph_data[config.graph.data].x
-        self.loss_weights = graph_data[config.graph.data][config.model.node_loss_weight].squeeze()
+        self.node_weights = graph_data[config.graph.data][config.model.node_loss_weight].squeeze()
 
         if config.model.get("output_mask", None) is not None:
             self.output_mask = Boolean1DMask(graph_data[config.graph.data][config.model.output_mask])
         else:
             self.output_mask = NoOutputMask()
-        self.loss_weights = self.output_mask.apply(self.loss_weights, dim=0, fill_value=0.0)
+        self.node_weights = self.output_mask.apply(self.node_weights, dim=0, fill_value=0.0)
 
         self.logger_enabled = config.diagnostics.log.wandb.enabled or config.diagnostics.log.mlflow.enabled
 
-        self.metric_ranges, self.metric_ranges_validation, loss_scaling = self.metrics_loss_scaling(
-            config,
-            data_indices,
-        )
-        self.loss = WeightedMSELoss(node_weights=self.loss_weights, data_variances=loss_scaling)
-        self.metrics = WeightedMSELoss(node_weights=self.loss_weights, ignore_nans=True)
+        variable_scaling = self.get_feature_weights(config, data_indices)
+
+        self.val_metric_ranges, _ = self.get_val_metric_ranges(config, data_indices)
+
+        # Kwargs to pass to the loss function
+        loss_kwargs = {"node_weights": self.node_weights}
+        # Scalars to include in the loss function, must be of form (dim, scalar)
+        scalars = {"variable": (-1, variable_scaling)}
+
+        self.loss = self.get_loss_function(config.training.training_loss, scalars=scalars, **loss_kwargs)
+
+        assert isinstance(self.loss, torch.nn.Module) and not isinstance(
+            self.loss,
+            torch.nn.ModuleList,
+        ), f"Loss function must be a `torch.nn.Module`, not a {type(self.loss).__name__!r}"
+
+        self.metrics = self.get_loss_function(config.training.validation_metrics, scalars=scalars, **loss_kwargs)
+        if not isinstance(self.metrics, torch.nn.ModuleList):
+            self.metrics = torch.nn.ModuleList([self.metrics])
 
         if config.training.loss_gradient_scaling:
             self.loss.register_full_backward_hook(grad_scaler, prepend=False)
@@ -136,15 +150,119 @@ def __init__(
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.model(x, self.model_comm_group)
 
+    # Future import breaks other type hints TODO Harrison Cook
+    @staticmethod
+    def get_loss_function(
+        config: DictConfig,
+        scalars: Union[dict[str, tuple[Union[int, tuple[int, ...], torch.Tensor]]], None] = None,  # noqa: FA100
+        **kwargs,
+    ) -> Union[torch.nn.Module, torch.nn.ModuleList]:  # noqa: FA100
+        """
+        Get loss functions from config.
+
+        Can be ModuleList if multiple losses are specified.
+
+        Parameters
+        ----------
+        config : DictConfig
+            Loss function configuration, should include `scalars` if scalars are to be added to the loss function.
+        scalars : Union[dict[str, tuple[Union[int, tuple[int, ...], torch.Tensor]]], None], optional
+            Scalars which can be added to the loss function. Defaults to None., by default None
+            If a scalar is to be added to the loss, ensure it is in `scalars` in the loss config
+            E.g.
+                If `scalars: ['variable']` is set in the config, and `variable` in `scalars`
+                `variable` will be added to the scalar of the loss function.
+        kwargs : Any
+            Additional arguments to pass to the loss function
+
+        Returns
+        -------
+        Union[torch.nn.Module, torch.nn.ModuleList]
+            Loss function, or list of metrics
+
+        Raises
+        ------
+        TypeError
+            If not a subclass of `BaseWeightedLoss`
+        ValueError
+            If scalar is not found in valid scalars
+        """
+        config_container = OmegaConf.to_container(config, resolve=False)
+        if isinstance(config_container, list):
+            return torch.nn.ModuleList(
+                [
+                    GraphForecaster.get_loss_function(
+                        OmegaConf.create(loss_config),
+                        scalars=scalars,
+                        **kwargs,
+                    )
+                    for loss_config in config
+                ],
+            )
+
+        loss_config = OmegaConf.to_container(config, resolve=True)
+        scalars_to_include = loss_config.pop("scalars", [])
+
+        # Instantiate the loss function with the loss_init_config
+        loss_function = instantiate(loss_config, **kwargs)
+
+        if not isinstance(loss_function, BaseWeightedLoss):
+            error_msg = f"Loss must be a subclass of 'BaseWeightedLoss', not {type(loss_function)}"
+            raise TypeError(error_msg)
+
+        for key in scalars_to_include:
+            if key not in scalars or []:
+                error_msg = f"Scalar {key!r} not found in valid scalars: {list(scalars.keys())}"
+                raise ValueError(error_msg)
+            loss_function.add_scalar(*scalars[key], name=key)
+
+        return loss_function
+
     @staticmethod
-    def metrics_loss_scaling(config: DictConfig, data_indices: IndexCollection) -> tuple[dict, torch.Tensor]:
+    def get_val_metric_ranges(config: DictConfig, data_indices: IndexCollection) -> tuple[dict, dict]:
+
         metric_ranges = defaultdict(list)
         metric_ranges_validation = defaultdict(list)
+
+        for key, idx in data_indices.internal_model.output.name_to_index.items():
+            split = key.split("_")
+            if len(split) > 1 and split[-1].isdigit():
+                # Group metrics for pressure levels (e.g., Q, T, U, V, etc.)
+                metric_ranges[f"pl_{split[0]}"].append(idx)
+            else:
+                metric_ranges[f"sfc_{key}"].append(idx)
+
+            # Specific metrics from hydra to log in logger
+            if key in config.training.metrics:
+                metric_ranges[key] = [idx]
+
+        # Add the full list of output indices
+        metric_ranges["all"] = data_indices.internal_model.output.full.tolist()
+
+        # metric for validation, after postprocessing
+        for key, idx in data_indices.model.output.name_to_index.items():
+            # Split pressure levels on "_" separator
+            split = key.split("_")
+            if len(split) > 1 and split[1].isdigit():
+                # Create grouped metrics for pressure levels (e.g. Q, T, U, V, etc.) for logger
+                metric_ranges_validation[f"pl_{split[0]}"].append(idx)
+            else:
+                metric_ranges_validation[f"sfc_{key}"].append(idx)
+            # Create specific metrics from hydra to log in logger
+            if key in config.training.metrics:
+                metric_ranges_validation[key] = [idx]
+
+        return metric_ranges, metric_ranges_validation
+
+    @staticmethod
+    def get_feature_weights(
+        config: DictConfig,
+        data_indices: IndexCollection,
+    ) -> torch.Tensor:
         loss_scaling = (
             np.ones((len(data_indices.internal_data.output.full),), dtype=np.float32)
             * config.training.loss_scaling.default
         )
-
         pressure_level = instantiate(config.training.pressure_level_scaler)
 
         LOGGER.info(
@@ -155,12 +273,9 @@ def metrics_loss_scaling(config: DictConfig, data_indices: IndexCollection) -> t
         )
 
         for key, idx in data_indices.internal_model.output.name_to_index.items():
-            # Split pressure levels on "_" separator
             split = key.split("_")
             if len(split) > 1 and split[-1].isdigit():
-                # Create grouped metrics for pressure levels (e.g. Q, T, U, V, etc.) for logger
-                metric_ranges[f"pl_{split[0]}"].append(idx)
-                # Create pressure levels in loss scaling vector
+                # Apply pressure level scaling
                 if split[0] in config.training.loss_scaling.pl:
                     loss_scaling[idx] = config.training.loss_scaling.pl[split[0]] * pressure_level.scaler(
                         int(split[-1]),
@@ -168,29 +283,13 @@ def metrics_loss_scaling(config: DictConfig, data_indices: IndexCollection) -> t
                 else:
                     LOGGER.debug("Parameter %s was not scaled.", key)
             else:
-                metric_ranges[f"sfc_{key}"].append(idx)
-                # Create surface variables in loss scaling vector
+                # Apply surface variable scaling
                 if key in config.training.loss_scaling.sfc:
                     loss_scaling[idx] = config.training.loss_scaling.sfc[key]
                 else:
                     LOGGER.debug("Parameter %s was not scaled.", key)
-            # Create specific metrics from hydra to log in logger
-            if key in config.training.metrics:
-                metric_ranges[key] = [idx]
-        loss_scaling = torch.from_numpy(loss_scaling)
-        # metric for validation, after postprocessing
-        for key, idx in data_indices.model.output.name_to_index.items():
-            # Split pressure levels on "_" separator
-            split = key.split("_")
-            if len(split) > 1 and split[1].isdigit():
-                # Create grouped metrics for pressure levels (e.g. Q, T, U, V, etc.) for logger
-                metric_ranges_validation[f"pl_{split[0]}"].append(idx)
-            else:
-                metric_ranges_validation[f"sfc_{key}"].append(idx)
-            # Create specific metrics from hydra to log in logger
-            if key in config.training.metrics:
-                metric_ranges_validation[key] = [idx]
-        return metric_ranges, metric_ranges_validation, loss_scaling
+
+        return torch.from_numpy(loss_scaling)
 
     def set_model_comm_group(self, model_comm_group: ProcessGroup) -> None:
         LOGGER.debug("set_model_comm_group: %s", model_comm_group)
@@ -274,16 +373,48 @@ def calculate_val_metrics(
         y: torch.Tensor,
         rollout_step: int,
         enable_plot: bool = False,
-    ) -> tuple[dict, list]:
+    ) -> tuple[dict, list[torch.Tensor]]:
+        """Calculate metrics on the validation output.
+
+        Parameters
+        ----------
+            y_pred: torch.Tensor
+                Predicted ensemble
+            y: torch.Tensor
+                Ground truth (target).
+            rollout_step: int
+                Rollout step
+            enable_plot: bool, defaults to False
+                Generate plots
+
+        Returns
+        -------
+            val_metrics, preds:
+                validation metrics and predictions
+        """
         metrics = {}
         y_preds = []
         y_postprocessed = self.model.post_processors(y, in_place=False)
         y_pred_postprocessed = self.model.post_processors(y_pred, in_place=False)
-        for mkey, indices in self.metric_ranges_validation.items():
-            metrics[f"{mkey}_{rollout_step + 1}"] = self.metrics(
-                y_pred_postprocessed[..., indices],
-                y_postprocessed[..., indices],
-            )
+
+        for metric in self.metrics:
+            metric_name = getattr(metric, "name", metric.__class__.__name__.lower())
+
+            if not isinstance(metric, BaseWeightedLoss):
+                # If not a weighted loss, we cannot feature scale, so call normally
+                metrics[f"{metric_name}/{rollout_step + 1}"] = metric(
+                    y_pred_postprocessed,
+                    y_postprocessed,
+                )
+                continue
+
+            for mkey, indices in self.val_metric_ranges.items():
+                metrics[f"{metric_name}/{mkey}/{rollout_step + 1}"] = metric(
+                    y_pred_postprocessed[..., indices],
+                    y_postprocessed[..., indices],
+                    feature_indices=indices,
+                    feature_scale=mkey == "all",
+                )
 
         if enable_plot:
             y_preds.append(y_pred)
@@ -292,7 +423,7 @@ def calculate_val_metrics(
     def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
         train_loss, _, _ = self._step(batch, batch_idx)
         self.log(
-            "train_wmse",
+            f"train_{getattr(self.loss, 'name', self.loss.__class__.__name__.lower())}",
             train_loss,
             on_epoch=True,
             on_step=True,
@@ -335,7 +466,7 @@ def validation_step(self, batch: torch.Tensor, batch_idx: int) -> None:
         with torch.no_grad():
             val_loss, metrics, y_preds = self._step(batch, batch_idx, validation_mode=True)
         self.log(
-            "val_wmse",
+            f"val_{getattr(self.loss, 'name', self.loss.__class__.__name__.lower())}",
             val_loss,
             on_epoch=True,
             on_step=True,


Test patch:
diff --git a/tests/train/test_loss_function.py b/tests/train/test_loss_function.py
new file mode 100644
index 00000000..73d7f246
--- /dev/null
+++ b/tests/train/test_loss_function.py
@@ -0,0 +1,64 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+import torch
+from omegaconf import DictConfig
+
+from anemoi.training.losses.mse import WeightedMSELoss
+from anemoi.training.losses.weightedloss import BaseWeightedLoss
+from anemoi.training.train.forecaster import GraphForecaster
+
+
+def test_manual_init() -> None:
+    loss = WeightedMSELoss(torch.ones(1))
+    assert loss.node_weights == torch.ones(1)
+
+
+def test_dynamic_init_include() -> None:
+    loss = GraphForecaster.get_loss_function(
+        DictConfig({"_target_": "anemoi.training.losses.mse.WeightedMSELoss"}),
+        node_weights=torch.ones(1),
+    )
+    assert isinstance(loss, BaseWeightedLoss)
+    assert loss.node_weights == torch.ones(1)
+
+
+def test_dynamic_init_scalar() -> None:
+    loss = GraphForecaster.get_loss_function(
+        DictConfig(
+            {
+                "_target_": "anemoi.training.losses.mse.WeightedMSELoss",
+                "scalars": ["test"],
+            },
+        ),
+        node_weights=torch.ones(1),
+        scalars={"test": ((0, 1), torch.ones((1, 2)))},
+    )
+    assert isinstance(loss, BaseWeightedLoss)
+
+    torch.testing.assert_close(loss.node_weights, torch.ones(1))
+    assert "test" in loss.scalar
+    torch.testing.assert_close(loss.scalar.get_scalar(2), torch.ones((1, 2)))
+
+
+def test_dynamic_init_scalar_not_add() -> None:
+    loss = GraphForecaster.get_loss_function(
+        DictConfig(
+            {
+                "_target_": "anemoi.training.losses.mse.WeightedMSELoss",
+                "scalars": [],
+            },
+        ),
+        node_weights=torch.ones(1),
+        scalars={"test": (-1, torch.ones(2))},
+    )
+    assert isinstance(loss, BaseWeightedLoss)
+    torch.testing.assert_close(loss.node_weights, torch.ones(1))
+    assert "test" not in loss.scalar
diff --git a/tests/train/test_loss_scaling.py b/tests/train/test_loss_scaling.py
index 2da5ae00..034325e5 100644
--- a/tests/train/test_loss_scaling.py
+++ b/tests/train/test_loss_scaling.py
@@ -131,7 +131,7 @@ def fake_data(request: SubRequest) -> tuple[DictConfig, IndexCollection]:
 def test_loss_scaling_vals(fake_data: tuple[DictConfig, IndexCollection], expected_scaling: torch.Tensor) -> None:
     config, data_indices = fake_data
 
-    _, _, loss_scaling = GraphForecaster.metrics_loss_scaling(config, data_indices)
+    loss_scaling = GraphForecaster.get_feature_weights(config, data_indices)
 
     assert torch.allclose(loss_scaling, expected_scaling)
 
@@ -140,7 +140,9 @@ def test_loss_scaling_vals(fake_data: tuple[DictConfig, IndexCollection], expect
 def test_metric_range(fake_data: tuple[DictConfig, IndexCollection]) -> None:
     config, data_indices = fake_data
 
-    metric_range, metric_ranges_validation, _ = GraphForecaster.metrics_loss_scaling(config, data_indices)
+    metric_range, metric_ranges_validation = GraphForecaster.get_val_metric_ranges(config, data_indices)
+
+    del metric_range["all"]
 
     expected_metric_range_validation = {
         "pl_y": [
@@ -161,5 +163,5 @@ def test_metric_range(fake_data: tuple[DictConfig, IndexCollection]) -> None:
     expected_metric_range["sfc_cos_d"] = [data_indices.internal_model.output.name_to_index["cos_d"]]
     expected_metric_range["sfc_sin_d"] = [data_indices.internal_model.output.name_to_index["sin_d"]]
 
-    assert dict(metric_ranges_validation) == expected_metric_range_validation
-    assert dict(metric_range) == expected_metric_range
+    assert metric_ranges_validation == expected_metric_range_validation
+    assert metric_range == expected_metric_range
diff --git a/tests/train/test_scaler.py b/tests/train/test_scaler.py
new file mode 100644
index 00000000..bc0770cf
--- /dev/null
+++ b/tests/train/test_scaler.py
@@ -0,0 +1,133 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import pytest
+import torch
+
+from anemoi.training.losses.utils import ScaleTensor
+
+
+def test_scale_contains() -> None:
+    scale = ScaleTensor(test=(0, 2))
+    assert "test" in scale
+
+
+def test_scale_contains_indexing() -> None:
+    scale = ScaleTensor(test=(0, 2))
+    assert 0 in scale
+
+
+def test_scale_tuple_contains_indexing() -> None:
+    scale = ScaleTensor(test=((0, 1), 2))
+    assert (0, 1) in scale
+
+
+def test_scale_tuple_not_contains_indexing() -> None:
+    scale = ScaleTensor(test=(0, 2))
+    assert (0, 1) not in scale
+
+
+def test_scale_contains_subset_indexing() -> None:
+    scale = ScaleTensor(test=(0, 2), wow=(0, 2))
+    assert "test" in scale
+    scale = scale.subset("wow")
+    assert "wow" in scale
+    assert "test" not in scale
+
+
+def test_scale_contains_subset_by_dim_indexing() -> None:
+    scale = ScaleTensor(test=(0, 2), wow=(1, 2))
+    assert "test" in scale
+    scale = scale.subset_by_dim(1)
+    assert "wow" in scale
+    assert "test" not in scale
+
+
+@pytest.mark.parametrize(
+    ("scalars", "input_tensor", "output"),
+    [
+        ([[0, torch.Tensor([2])]], torch.tensor([1.0, 2.0, 3.0]), torch.tensor([2.0, 4.0, 6.0])),
+        ([[0, torch.Tensor([0.5])]], torch.tensor([10.0, 20.0, 30.0]), torch.tensor([5.0, 10.0, 15.0])),
+        ([[-1, torch.Tensor([0.5])]], torch.tensor([10.0, 20.0, 30.0]), torch.tensor([5.0, 10.0, 15.0])),
+        ([[0, torch.Tensor([0])]], torch.tensor([5.0, 10.0, 15.0]), torch.tensor([0.0, 0.0, 0.0])),
+        (
+            [[0, torch.Tensor([0.5])], [-1, torch.Tensor([3])]],
+            torch.tensor([10.0, 20.0, 30.0]),
+            torch.tensor([15.0, 30.0, 45.0]),
+        ),
+        (
+            [[0, torch.Tensor([0.5])], [0, torch.Tensor([3])]],
+            torch.tensor([10.0, 20.0, 30.0]),
+            torch.tensor([15.0, 30.0, 45.0]),
+        ),
+    ],
+)
+def test_scale_tensor_one_dim(
+    scalars: list[list[int, torch.Tensor]],
+    input_tensor: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+
+    scale = ScaleTensor()
+    for scalar in scalars:
+        scale.add_scalar(*scalar)
+
+    torch.testing.assert_close(scale.scale(input_tensor), output)
+
+
+def test_invalid_dim_sizes() -> None:
+    scalar = ScaleTensor()
+    scalar.add_scalar(0, torch.ones((5,)))
+
+    with pytest.raises(ValueError, match=r"Validating tensor 'invalid' raised an error."):
+        scalar.add_scalar(0, torch.ones((15,)), name="invalid")
+
+
+def test_invalid_dim_sizes_negative_indexing() -> None:
+    scalar = ScaleTensor()
+    scalar.add_scalar(0, torch.ones((5,)))
+    scalar.add_scalar(-1, torch.ones((15,)), name="invalid")
+
+    with pytest.raises(ValueError, match=r"Validating tensor 'invalid' raised an error."):
+        scalar.resolve(1)
+
+
+def test_valid_dim_sizes_negative_indexing() -> None:
+    scalar = ScaleTensor()
+    scalar.add_scalar(0, torch.ones((5,)))
+    scalar.add_scalar(-1, torch.ones((15,)))
+
+    scalar.resolve(2)
+
+
+@pytest.mark.parametrize(
+    ("scalars", "input_tensor", "output"),
+    [
+        ([[0, 2.0]], torch.ones([4, 6]), torch.ones([4, 6]) * 2),
+        ([[0, [[1.0, 1.0], [1.0, 2.0]]]], torch.ones((2, 2)), [[1, 1], [1, 2]]),
+        ([[(0, 1), [[1.0, 1.0], [1.0, 2.0]]]], torch.ones((2, 2)), [[1, 1], [1, 2]]),
+        ([[(1, 0), [[1.0, 1.0], [1.0, 2.0]]]], torch.ones((2, 2)), [[1, 1], [1, 2]]),
+        ([[(0, 1), [[1.0, 2.0], [1.0, 1.0]]]], torch.ones((2, 2)), [[1, 2], [1, 1]]),
+        ([[(1, 0), [[1.0, 2.0], [1.0, 1.0]]]], torch.ones((2, 2)), [[1, 1], [2, 1]]),
+    ],
+)
+def test_scale_tensor_two_dim(
+    scalars: list[list[int, torch.Tensor]],
+    input_tensor: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+
+    scale = ScaleTensor()
+    for scalar in scalars:
+        scale.add_scalar(*scalar)
+
+    if not isinstance(input_tensor, torch.Tensor):
+        input_tensor = torch.tensor(input_tensor, dtype=torch.float32)
+    if not isinstance(output, torch.Tensor):
+        output = torch.tensor(output, dtype=torch.float32)
+
+    torch.testing.assert_close(scale.scale(input_tensor), output)