Skip to content

7. Losses

7.1 elbo

This module contains the code for the Losses.

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

Computes the Evidence Lower Bound (ELBO) loss, combining a likelihood loss and KL divergence.

Initializes the ELBO loss with specified likelihood loss function, sample count, and KL weight.

Parameters:

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

Loss function for computing likelihood loss.

 required
num_samples int

Number of samples for Monte Carlo approximation.

 1
kl_weight float

Scaling factor for the KL divergence component.

 0.001
**kwargs Any

Additional keyword arguments.

 {}
Source code in illia/losses/tf/elbo.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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 specified likelihood loss
    function, sample count, and KL weight.

    Args:
        loss_function: Loss function for computing likelihood loss.
        num_samples: Number of samples for Monte Carlo approximation.
        kl_weight: Scaling factor for the KL divergence component.
        **kwargs: Additional keyword arguments.
    """

    # 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)

7.1.1.1 __call__(y_true, y_pred, model)

Computes the ELBO loss, averaging over multiple samples.

Parameters:

Name Type Description Default
y_true Tensor

True labels.

 required
y_pred Tensor

Predictions from the model.

 required
model Model

TensorFlow model containing Bayesian layers.

 required

Returns:

Type Description
Tensor

Average ELBO loss across samples.
Source code in illia/losses/tf/elbo.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor, model: Model) -> tf.Tensor:
    """
    Computes the ELBO loss, averaging over multiple samples.

    Args:
        y_true: True labels.
        y_pred: Predictions from the model.
        model: TensorFlow model containing Bayesian layers.

    Returns:
        Average ELBO loss across samples.
    """

    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

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

Computes the KL divergence loss for Bayesian modules within a model.

Initializes the KL divergence loss with specified reduction method and weight.

Parameters:

Name Type Description Default
reduction Literal['mean']

Method to reduce the loss, currently only "mean" is supported.

 'mean'
weight float

Scaling factor for the KL divergence loss.

 1.0
**kwargs Any

Additional keyword arguments.

 {}
Source code in illia/losses/tf/elbo.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def __init__(
    self,
    reduction: Literal["mean"] = "mean",
    weight: float = 1.0,
    **kwargs: Any,
) -> None:
    """
    Initializes the KL divergence loss with specified reduction
    method and weight.

    Args:
        reduction: Method to reduce the loss, currently only "mean"
            is supported.
        weight: Scaling factor for the KL divergence loss.
        **kwargs: Additional keyword arguments.
    """

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

    # Set parameters
    self.reduction = reduction
    self.weight = weight

7.1.2.1 __call__(model)

Computes the KL divergence loss across all Bayesian layers in the model.

Parameters:

Name Type Description Default
model Model

TensorFlow model containing Bayesian layers.

 required

Returns:

Type Description
Tensor

KL divergence cost scaled by the specified weight.
Source code in illia/losses/tf/elbo.py 
63
64
65
66
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
def __call__(self, model: Model) -> tf.Tensor:
    """
    Computes the KL divergence loss across all Bayesian layers in
    the model.

    Args:
        model: TensorFlow model containing Bayesian layers.

    Returns:
        KL divergence cost scaled by the specified weight.
    """

    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