Skip to content

8. Losses

This module implements the Kullback-Leibler (KL) divergence loss for Bayesian neural networks in Tensorflow.

8.1 KLDivergenceLoss(reduction='mean', weight=1.0, **kwargs)

Computes the Kullback-Leibler divergence loss across all Bayesian modules.

Initializes the Kullback-Leibler divergence loss computation.

Parameters:

Name Type Description Default
reduction Literal['mean']

Reduction method for the loss.

 'mean'
weight float

Scaling factor applied to the total KL loss.

 1.0
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
None

None.
Source code in illia/losses/tf/kl.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    reduction: Literal["mean"] = "mean",
    weight: float = 1.0,
    **kwargs: Any,
) -> None:
    """
    Initializes the Kullback-Leibler divergence loss computation.

    Args:
        reduction: Reduction method for the loss.
        weight: Scaling factor applied to the total KL loss.
        **kwargs: Additional keyword arguments.

    Returns:
        None.
    """

    # Call super class constructor
    super().__init__(**kwargs)

    # Set attributes
    self.reduction = reduction
    self.weight = weight

8.1.1 __call__(*args, **kwargs)

Computes Kullback-Leibler divergence for all Bayesian modules in the model.

Parameters:

Name Type Description Default
*args Any

Unused positional arguments.

 ()
**kwargs Any

Must include 'model' as a keyword argument.

 {}

Returns:

Type Description
Tensor

Scaled Kullback-Leibler divergence loss as a scalar array.
Source code in illia/losses/tf/kl.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def __call__(self, *args: Any, **kwargs: Any) -> tf.Tensor:
    """
    Computes Kullback-Leibler divergence for all Bayesian
    modules in the model.

    Args:
        *args: Unused positional arguments.
        **kwargs: Must include 'model' as a keyword argument.

    Returns:
        Scaled Kullback-Leibler divergence loss as a scalar array.
    """

    model = kwargs.get("model")
    if model is None:
        raise ValueError("Model must be provided as a keyword argument")

    kl_global_cost: tf.Tensor = tf.constant(0.0, dtype=tf.float32)
    num_params_global: int = 0

    # Iterate through the model's layers
    for layer in model.layers:
        if isinstance(layer, BayesianModule):
            kl_cost, num_params = layer.kl_cost()
            kl_global_cost += kl_cost
            num_params_global += num_params

    # Compute mean KL cost and scale by weight
    kl_global_cost = tf.divide(
        kl_global_cost, tf.cast(num_params_global, tf.float32)
    )
    kl_global_cost = tf.multiply(kl_global_cost, self.weight)

    return kl_global_cost

This module implements the Evidence Lower Bound (ELBO) loss for Bayesian neural networks in TensorFlow.

8.2 ELBOLoss(loss_function, num_samples=1, kl_weight=0.001, **kwargs)

Computes the Evidence Lower Bound (ELBO) loss function for Bayesian neural networks.

This combines a reconstruction loss and a KL divergence term, estimated using Monte Carlo sampling.

Initializes the ELBO loss with sampling and KL scaling.

Parameters:

Name Type Description Default
loss_function Callable[[Tensor, Tensor], Tensor]

Module for computing reconstruction loss.

 required
num_samples int

Number of MC samples for estimation.

 1
kl_weight float

Weight applied to the KL loss.

 0.001

Returns:

Type Description
None

None.
Source code in illia/losses/tf/elbo.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def __init__(
    self,
    loss_function: Callable[[tf.Tensor, tf.Tensor], tf.Tensor],
    num_samples: int = 1,
    kl_weight: float = 1e-3,
    **kwargs: Any,
) -> None:
    """
    Initializes the ELBO loss with sampling and KL scaling.

    Args:
        loss_function: Module for computing reconstruction loss.
        num_samples: Number of MC samples for estimation.
        kl_weight: Weight applied to the KL loss.

    Returns:
        None.
    """

    # Call super class constructor
    super().__init__(**kwargs)

    # Set attributes
    self.loss_function = loss_function
    self.num_samples = num_samples
    self.kl_weight = kl_weight
    self.kl_loss = KLDivergenceLoss(weight=kl_weight)

8.2.1 __call__(y_true, y_pred, *args, **kwargs)

Computes the ELBO loss using KL regularization and reconstruction error.

Parameters:

Name Type Description Default
y_true Tensor

Ground truth targets.

 required
y_pred Tensor

Model predictions.

 required
*args Any

Unused positional arguments.

 ()
**kwargs Any

Must include 'model' containing Bayesian layers.

 {}

Returns:

Type Description
Tensor

Scalar tensor representing the total ELBO loss.
Source code in illia/losses/tf/elbo.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def __call__(
    self, y_true: tf.Tensor, y_pred: tf.Tensor, *args: Any, **kwargs: Any
) -> tf.Tensor:
    """
    Computes the ELBO loss using KL regularization and reconstruction error.

    Args:
        y_true: Ground truth targets.
        y_pred: Model predictions.
        *args: Unused positional arguments.
        **kwargs: Must include 'model' containing Bayesian layers.

    Returns:
        Scalar tensor representing the total ELBO loss.
    """

    model = kwargs.get("model")
    if model is None:
        raise ValueError("Model must be provided as a keyword argument")

    loss_value: tf.Tensor = tf.constant(0.0, dtype=tf.float32)

    for _ in range(self.num_samples):
        current_loss = self.loss_function(y_true, y_pred) + self.kl_loss(model)
        loss_value += current_loss

    # Average the loss across samples
    loss_value = tf.divide(loss_value, tf.cast(self.num_samples, tf.float32))

    return loss_value