Skip to content

Losses

Objectives for training variational Bayesian models.

Variational Loss

bensemble.losses.VariationalLoss 

VariationalLoss(likelihood_model: Module, alpha: float = 1.0, num_batches: int = 1, kl_weight: float = 1.0)

Bases: Module

Variational loss supporting ELBO and Rényi alpha-divergence.

Initializes the VariationalLoss module.

Parameters:

Name Type Description Default
likelihood_model Module

Module computing negative log-likelihood p(y|x).

 required
alpha float

Rényi alpha parameter. When alpha=1.0, standard ELBO is used. Defaults to 1.0.

 1.0
num_batches int

Total number of training batches in an epoch (for KL mini-batch scaling). Defaults to 1.

 1
kl_weight float

Additional scaling factor for KL divergence term. Defaults to 1.0.

 1.0
Source code in bensemble/losses.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def __init__(
    self,
    likelihood_model: nn.Module,
    alpha: float = 1.0,
    num_batches: int = 1,
    kl_weight: float = 1.0,
):
    """Initializes the VariationalLoss module.

    Args:
        likelihood_model: Module computing negative log-likelihood p(y|x).
        alpha: Rényi alpha parameter. When alpha=1.0, standard ELBO is used. Defaults to 1.0.
        num_batches: Total number of training batches in an epoch (for KL mini-batch scaling). Defaults to 1.
        kl_weight: Additional scaling factor for KL divergence term. Defaults to 1.0.
    """
    super().__init__()
    self.likelihood_model = likelihood_model
    self.alpha = alpha
    self.num_batches = num_batches
    self.kl_weight = kl_weight

forward 

forward(preds: Tensor, target: Tensor, kl_divergence: Tensor) -> torch.Tensor

Computes variational objective loss.

Parameters:

Name Type Description Default
preds Tensor

Predictions tensor of shape (K, batch_size, ...) or (batch_size, ...).

 required
target Tensor

Target tensor of shape (batch_size, ...).

 required
kl_divergence Tensor

Total KL divergence of model parameters.

 required

Returns:

Type Description
Tensor

torch.Tensor: Scalar loss value.
Source code in bensemble/losses.py 
 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def forward(
    self, preds: torch.Tensor, target: torch.Tensor, kl_divergence: torch.Tensor
) -> torch.Tensor:
    """Computes variational objective loss.

    Args:
        preds: Predictions tensor of shape (K, batch_size, ...) or (batch_size, ...).
        target: Target tensor of shape (batch_size, ...).
        kl_divergence: Total KL divergence of model parameters.

    Returns:
        torch.Tensor: Scalar loss value.
    """
    class_indices = not torch.is_floating_point(target)

    if class_indices:
        if preds.ndim == target.ndim + 1:
            preds = preds.unsqueeze(0)
        log_likelihoods = -torch.stack(
            [self.likelihood_model(sample, target) for sample in preds]
        )
    else:
        if preds.ndim == target.ndim:
            preds = preds.unsqueeze(0)
        elif target.shape == preds.shape[:-1]:
            target = target.unsqueeze(-1)
            preds = preds.unsqueeze(0)
        while target.ndim < preds.ndim - 1:
            target = target.unsqueeze(-1)
        log_likelihoods = -self.likelihood_model(preds, target)

    k_samples = preds.size(0)
    log_likelihoods = log_likelihoods.reshape(k_samples, -1).sum(dim=1)

    kl_scaled = (kl_divergence / self.num_batches) * self.kl_weight
    log_weights = log_likelihoods - kl_scaled

    if abs(self.alpha - 1.0) < 1e-6:
        return -(log_weights.mean())
    else:
        term = (1 - self.alpha) * log_weights
        log_sum_exp = torch.logsumexp(term, dim=0)
        loss = -(1 / (1 - self.alpha)) * (log_sum_exp - math.log(k_samples))
        return loss

Gaussian Likelihood

bensemble.losses.GaussianLikelihood 

GaussianLikelihood(init_log_sigma: float = -2.0)

Bases: Module

Gaussian Likelihood with learnable homoscedastic uncertainty.

Learns a global standard deviation (sigma) for data noise and computes negative log-likelihood via Gaussian NLL.

Initializes the GaussianLikelihood layer.

Parameters:

Name Type Description Default
init_log_sigma float

Initial value for log standard deviation. Defaults to -2.0.

 -2.0
Source code in bensemble/losses.py 
15
16
17
18
19
20
21
22
23
def __init__(self, init_log_sigma: float = -2.0):
    """Initializes the GaussianLikelihood layer.

    Args:
        init_log_sigma: Initial value for log standard deviation. Defaults to -2.0.
    """
    super().__init__()
    self.log_sigma = nn.Parameter(torch.tensor([init_log_sigma]))
    self.loss_fn = nn.GaussianNLLLoss(reduction="none")

sigma property 

sigma: float

float: Current standard deviation value.

forward 

forward(preds: Tensor, target: Tensor) -> torch.Tensor

Computes Gaussian negative log-likelihood loss.

Parameters:

Name Type Description Default
preds Tensor

Predicted tensor of shape (*).

 required
target Tensor

Target tensor of shape (*).

 required

Returns:

Type Description
Tensor

torch.Tensor: Element-wise negative log-likelihood tensor matching target shape.
Source code in bensemble/losses.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Computes Gaussian negative log-likelihood loss.

    Args:
        preds: Predicted tensor of shape (*).
        target: Target tensor of shape (*).

    Returns:
        torch.Tensor: Element-wise negative log-likelihood tensor matching target shape.
    """
    sigma = F.softplus(self.log_sigma) + 1e-6
    var = sigma.pow(2)
    var_expanded = var.expand_as(preds)

    return self.loss_fn(preds, target, var_expanded)