Skip to content

Variance Bayessian Loss

LogUniformVarKLLoss

Bases: VarDistLoss

KL loss between factorized variational distribution and LogUniform prior. Works only with modules with LogUniformVarDist posterior

Source code in src/methods/bayes/variational/optimization.py 
 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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
class LogUniformVarKLLoss(VarDistLoss):
    """KL loss between factorized variational distribution and LogUniform prior. 
    Works only with modules with LogUniformVarDist posterior"""
    def __init__(self):
        """_summary_"""
        super().__init__()

    def forward(
        self, posterior: dict[str, LogUniformVarDist], **kwargs
    ) -> torch.Tensor:
        """
        Computes KL loss between factorized variational distribution and LogUniform prior

        Args:
            posterior (dict[str, LogUniformVarDist]): factorized normal variational distribution
                    with hidden variable that is used with LogUniform prior
        """
        k1 = torch.tensor(0.63576)
        k2 = torch.tensor(1.87320)
        k3 = torch.tensor(1.48695)
        KL_z = torch.tensor(0)
        for dist in posterior.values():
            KL_z_element = (
                k1 * F.sigmoid(k2 + k3 * dist.scale_alphas_log)
                - 0.5 * F.softplus(-dist.scale_alphas_log)
                - k1
            )
            KL_z = KL_z + KL_z_element.sum()
        KL_w = torch.tensor(0)
        for dist in posterior.values():
            KL_w_element = 0.5 * (
                torch.log(1 / torch.exp(dist.param_std_log) ** 2)
                + torch.exp(dist.param_std_log) ** 2
                + dist.param_mus**2
                - 1
            )
            KL_w = KL_w + KL_w_element.sum()

        return -KL_z + KL_w

    def aggregate(
        self, fit_losses: list, dist_losses: list, beta: float
    ) -> VarDistLoss.AggregationResult:
        """
        This method aggregate dist_lossed and fit_losses for whole
        sampled parameters.

        Args:
            fit_losses (list):
                    list of data loss of each sample
            dist_losses (list):  {param_name: ParamDist}
                    list of distance loss of each sample
            beta:  {param_name: ParamDist}
                    sacle parameter of distance loss
        Returns:
            VarDistLoss.AggregationResult: Aggretion result for whole samples
        """
        fit_loss = torch.mean(torch.stack(fit_losses))
        dist_loss = torch.stack(dist_losses)[0]
        total_loss = fit_loss + beta * dist_loss
        return VarDistLoss.AggregationResult(total_loss, fit_loss, dist_loss)

__init__()

summary

Source code in src/methods/bayes/variational/optimization.py 
83
84
85
def __init__(self):
    """_summary_"""
    super().__init__()

aggregate(fit_losses, dist_losses, beta)

This method aggregate dist_lossed and fit_losses for whole sampled parameters.

Parameters:

Name Type Description Default
fit_losses list 
list of data loss of each sample
required
dist_losses list

{param_name: ParamDist} list of distance loss of each sample

 required
beta float

{param_name: ParamDist} sacle parameter of distance loss

 required

Returns: VarDistLoss.AggregationResult: Aggretion result for whole samples

Source code in src/methods/bayes/variational/optimization.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def aggregate(
    self, fit_losses: list, dist_losses: list, beta: float
) -> VarDistLoss.AggregationResult:
    """
    This method aggregate dist_lossed and fit_losses for whole
    sampled parameters.

    Args:
        fit_losses (list):
                list of data loss of each sample
        dist_losses (list):  {param_name: ParamDist}
                list of distance loss of each sample
        beta:  {param_name: ParamDist}
                sacle parameter of distance loss
    Returns:
        VarDistLoss.AggregationResult: Aggretion result for whole samples
    """
    fit_loss = torch.mean(torch.stack(fit_losses))
    dist_loss = torch.stack(dist_losses)[0]
    total_loss = fit_loss + beta * dist_loss
    return VarDistLoss.AggregationResult(total_loss, fit_loss, dist_loss)

forward(posterior, **kwargs)

Computes KL loss between factorized variational distribution and LogUniform prior

Parameters:

Name Type Description Default
posterior dict[str, LogUniformVarDist]

factorized normal variational distribution with hidden variable that is used with LogUniform prior

 required
Source code in src/methods/bayes/variational/optimization.py 
 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
115
116
117
118
def forward(
    self, posterior: dict[str, LogUniformVarDist], **kwargs
) -> torch.Tensor:
    """
    Computes KL loss between factorized variational distribution and LogUniform prior

    Args:
        posterior (dict[str, LogUniformVarDist]): factorized normal variational distribution
                with hidden variable that is used with LogUniform prior
    """
    k1 = torch.tensor(0.63576)
    k2 = torch.tensor(1.87320)
    k3 = torch.tensor(1.48695)
    KL_z = torch.tensor(0)
    for dist in posterior.values():
        KL_z_element = (
            k1 * F.sigmoid(k2 + k3 * dist.scale_alphas_log)
            - 0.5 * F.softplus(-dist.scale_alphas_log)
            - k1
        )
        KL_z = KL_z + KL_z_element.sum()
    KL_w = torch.tensor(0)
    for dist in posterior.values():
        KL_w_element = 0.5 * (
            torch.log(1 / torch.exp(dist.param_std_log) ** 2)
            + torch.exp(dist.param_std_log) ** 2
            + dist.param_mus**2
            - 1
        )
        KL_w = KL_w + KL_w_element.sum()

    return -KL_z + KL_w

NormVarKLLoss

Bases: VarDistLoss

KL loss between factorized normals. Works only with modules with NormalReparametrizedDist posterior

Source code in src/methods/bayes/variational/optimization.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class NormVarKLLoss(VarDistLoss):
    """KL loss between factorized normals. 
    Works only with modules with NormalReparametrizedDist posterior"""
    def __init__(self):
        """_summary_"""
        super().__init__()
    def forward(
        self, posterior: dict[str, NormalReparametrizedDist], **kwargs
    ) -> torch.Tensor:
        r"""
        Computes KL loss between factorized normals

        Args:
            posterior_params (nn.ParameterList): factorized normal variational distribution
            prior_parmeter (Optional[nn.ParameterList]): assumed fixed N($\mu$, $\sigma$) for all paramteres.
                As it is possible to analitically find optimal ($\mu$, $\sigma$), this parameter is ignored here.
        """
        # optimal \alpha parameters
        # sigma_opt here is the squared sigma
        mu_opt = 0
        sigma_opt = 0
        # number of posterior random variables = num of parameters
        n_params = 0

        # compute mu
        for param in posterior.loc:
            mu_opt += param.sum()
            n_params += param.numel()
        mu_opt /= n_params

        # compute sigma in two steps
        for param in posterior.loc:
            sigma_opt += torch.sum((param - mu_opt) ** 2) / n_params
        for param in posterior.log_scale:
            sigma_opt += torch.exp(param).sum() / n_params

        # compute kl-loss between posterior and prior in two steps
        kl_loss = 0
        for param in posterior.loc:
            kl_loss += torch.sum(0.5 * (1 / sigma_opt) * (param - mu_opt) ** 2)
        for param in posterior.log_scale:
            kl_loss += -0.5 * param.sum() + param.numel() * 0.5 * torch.log(sigma_opt)
            kl_loss += torch.sum(-0.5 + 0.5 * torch.exp(param) / sigma_opt)

        return kl_loss

    def aggregate(
        self, fit_losses: list, dist_losses: list, beta: float
    ) -> VarDistLoss.AggregationResult:
        """
        This method aggregate dist_lossed and fit_losses for whole
        sampled parameters.

        Args:
            fit_losses (list):
                    list of data loss of each sample
            dist_losses (list):  {param_name: ParamDist}
                    list of distance loss of each sample
            beta:  {param_name: ParamDist}
                    sacle parameter of distance loss
        Returns:
            VarDistLoss.AggregationResult: Aggretion result for whole samples
        """
        fit_loss = torch.mean(torch.stack(fit_losses))
        dist_loss = torch.stack(dist_losses)[0]
        total_loss = fit_loss + beta * dist_loss
        return VarDistLoss.AggregationResult(total_loss, fit_loss, dist_loss)

__init__()

summary

Source code in src/methods/bayes/variational/optimization.py 
146
147
148
def __init__(self):
    """_summary_"""
    super().__init__()

aggregate(fit_losses, dist_losses, beta)

This method aggregate dist_lossed and fit_losses for whole sampled parameters.

Parameters:

Name Type Description Default
fit_losses list 
list of data loss of each sample
required
dist_losses list

{param_name: ParamDist} list of distance loss of each sample

 required
beta float

{param_name: ParamDist} sacle parameter of distance loss

 required

Returns: VarDistLoss.AggregationResult: Aggretion result for whole samples

Source code in src/methods/bayes/variational/optimization.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def aggregate(
    self, fit_losses: list, dist_losses: list, beta: float
) -> VarDistLoss.AggregationResult:
    """
    This method aggregate dist_lossed and fit_losses for whole
    sampled parameters.

    Args:
        fit_losses (list):
                list of data loss of each sample
        dist_losses (list):  {param_name: ParamDist}
                list of distance loss of each sample
        beta:  {param_name: ParamDist}
                sacle parameter of distance loss
    Returns:
        VarDistLoss.AggregationResult: Aggretion result for whole samples
    """
    fit_loss = torch.mean(torch.stack(fit_losses))
    dist_loss = torch.stack(dist_losses)[0]
    total_loss = fit_loss + beta * dist_loss
    return VarDistLoss.AggregationResult(total_loss, fit_loss, dist_loss)

forward(posterior, **kwargs)

Computes KL loss between factorized normals

Parameters:

Name Type Description Default
posterior_params ParameterList

factorized normal variational distribution

 required
prior_parmeter Optional[ParameterList]

assumed fixed N(\(\mu\), \(\sigma\)) for all paramteres. As it is possible to analitically find optimal (\(\mu\), \(\sigma\)), this parameter is ignored here.

 required
Source code in src/methods/bayes/variational/optimization.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def forward(
    self, posterior: dict[str, NormalReparametrizedDist], **kwargs
) -> torch.Tensor:
    r"""
    Computes KL loss between factorized normals

    Args:
        posterior_params (nn.ParameterList): factorized normal variational distribution
        prior_parmeter (Optional[nn.ParameterList]): assumed fixed N($\mu$, $\sigma$) for all paramteres.
            As it is possible to analitically find optimal ($\mu$, $\sigma$), this parameter is ignored here.
    """
    # optimal \alpha parameters
    # sigma_opt here is the squared sigma
    mu_opt = 0
    sigma_opt = 0
    # number of posterior random variables = num of parameters
    n_params = 0

    # compute mu
    for param in posterior.loc:
        mu_opt += param.sum()
        n_params += param.numel()
    mu_opt /= n_params

    # compute sigma in two steps
    for param in posterior.loc:
        sigma_opt += torch.sum((param - mu_opt) ** 2) / n_params
    for param in posterior.log_scale:
        sigma_opt += torch.exp(param).sum() / n_params

    # compute kl-loss between posterior and prior in two steps
    kl_loss = 0
    for param in posterior.loc:
        kl_loss += torch.sum(0.5 * (1 / sigma_opt) * (param - mu_opt) ** 2)
    for param in posterior.log_scale:
        kl_loss += -0.5 * param.sum() + param.numel() * 0.5 * torch.log(sigma_opt)
        kl_loss += torch.sum(-0.5 + 0.5 * torch.exp(param) / sigma_opt)

    return kl_loss

VarDistLoss

Bases: BaseLoss

Abstract class for Distribution losses. Your distribution loss should be computed using prior and posterior classes and parameters, sampled from posterior.

In forward method loss should realize logic of loss for one sampled weights.

In aggregate method loss aggregate the data losses and distribution losses for samples.

Aggregation returns VarDistLoss.AggregationResult

Source code in src/methods/bayes/variational/optimization.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class VarDistLoss(BaseLoss):
    """
    Abstract class for Distribution losses. Your distribution loss should
    be computed using prior and posterior classes and parameters, sampled from posterior.

    In forward method loss should realize logic of loss for one sampled weights.

    In aggregate method loss aggregate the data losses and distribution losses for samples.

    Aggregation returns `VarDistLoss.AggregationResult`

    """

    @dataclass
    class AggregationResult:
        total_loss: torch.Tensor
        fit_loss: torch.Tensor
        dist_loss: torch.Tensor

    def __init__(self):
        """_summary_"""
        super().__init__()

    @abstractmethod
    def forward(
        self,
        *,
        param_sample_dict: Dict[str, nn.Parameter],
        posterior: Dict[str, ParamDist],
        prior: Dict[str, ParamDist],
    ) -> torch.Tensor:
        """
        This method computes loss for one sampled parameters.

        Args:
            param_sample_dict: {param_name: nn.Parameter}
                    sampled parameters on network.
            posterior:  {param_name: ParamDist}
                    posterior distribution of net parameters.
            prior:  {param_name: ParamDist}
                    prior distribution of net parameters.

        Returns:
            torch.Tensor: distanse loss for one sampled parameters
        """
        ...

    @abstractmethod
    def aggregate(
        self, fit_losses: list, dist_losses: list, beta: float
    ) -> AggregationResult:
        """
        This method aggregate dist_lossed and fit_losses for whole
        sampled parameters.

        Args:
            fit_losses (list):
                    list of data loss of each sample
            dist_losses (list):  {param_name: ParamDist}
                    list of distance loss of each sample
            beta:  {param_name: ParamDist}
                    sacle parameter of distance loss
        Returns:
            AggregationResult: Aggretion result for whole samples
        """
        ...

__init__()

summary

Source code in src/methods/bayes/variational/optimization.py 
31
32
33
def __init__(self):
    """_summary_"""
    super().__init__()

aggregate(fit_losses, dist_losses, beta) abstractmethod

This method aggregate dist_lossed and fit_losses for whole sampled parameters.

Parameters:

Name Type Description Default
fit_losses list 
list of data loss of each sample
required
dist_losses list

{param_name: ParamDist} list of distance loss of each sample

 required
beta float

{param_name: ParamDist} sacle parameter of distance loss

 required

Returns: AggregationResult: Aggretion result for whole samples

Source code in src/methods/bayes/variational/optimization.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
@abstractmethod
def aggregate(
    self, fit_losses: list, dist_losses: list, beta: float
) -> AggregationResult:
    """
    This method aggregate dist_lossed and fit_losses for whole
    sampled parameters.

    Args:
        fit_losses (list):
                list of data loss of each sample
        dist_losses (list):  {param_name: ParamDist}
                list of distance loss of each sample
        beta:  {param_name: ParamDist}
                sacle parameter of distance loss
    Returns:
        AggregationResult: Aggretion result for whole samples
    """
    ...

forward(*, param_sample_dict, posterior, prior) abstractmethod

This method computes loss for one sampled parameters.

Parameters:

Name Type Description Default
param_sample_dict Dict[str, Parameter]

{param_name: nn.Parameter} sampled parameters on network.

 required
posterior Dict[str, ParamDist]

{param_name: ParamDist} posterior distribution of net parameters.

 required
prior Dict[str, ParamDist]

{param_name: ParamDist} prior distribution of net parameters.

 required

Returns:

Type Description
Tensor

torch.Tensor: distanse loss for one sampled parameters
Source code in src/methods/bayes/variational/optimization.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
@abstractmethod
def forward(
    self,
    *,
    param_sample_dict: Dict[str, nn.Parameter],
    posterior: Dict[str, ParamDist],
    prior: Dict[str, ParamDist],
) -> torch.Tensor:
    """
    This method computes loss for one sampled parameters.

    Args:
        param_sample_dict: {param_name: nn.Parameter}
                sampled parameters on network.
        posterior:  {param_name: ParamDist}
                posterior distribution of net parameters.
        prior:  {param_name: ParamDist}
                prior distribution of net parameters.

    Returns:
        torch.Tensor: distanse loss for one sampled parameters
    """
    ...

VarRenuiLoss

Bases: VarDistLoss

Loss with Renui divergence. https://arxiv.org/pdf/1602.02311v1

Source code in src/methods/bayes/variational/optimization.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
class VarRenuiLoss(VarDistLoss):
    """
    Loss with Renui divergence.
    https://arxiv.org/pdf/1602.02311v1
    """

    def __init__(self, alpha: float = -1, aggregation: str = "weighting"):
        """_summary_

        Args:
            alpha (float): alpha
            aggregation (str): aggregation
        """
        assert aggregation in ["weighting", "sample"]
        self.alpha = alpha
        self.aggregation = aggregation
        super().__init__()

    def forward(
        self,
        param_sample_dict: dict[str, nn.Parameter],
        posterior: dict[str, ParamDist],
        prior: dict[str, ParamDist],
    ) -> torch.Tensor:
        """
        This method computes loss for one sampled parameters.

        Args:
            param_sample_dict: {param_name: nn.Parameter}
                    sampled parameters on network.
            posterior:  {param_name: ParamDist}
                    posterior distribution of net parameters.
            prior:  {param_name: ParamDist}
                    prior distribution of net parameters.

        Returns:
            torch.Tensor: distanse loss for one sampled parameters
        """
        prior_likelihood: torch.tensor = 0.0
        posterior_likelihood: torch.tensor = 0.0

        for parameter, posterior_distr, prior_distr in zip(
            param_sample_dict.values(), posterior.values(), prior.values()
        ):
            prior_likelihood += prior_distr.log_prob(parameter).sum()
            posterior_likelihood += posterior_distr.log_prob(parameter).sum()
            # if key == "fc2.bias":
            # print(posterior_distr.loc)
            # print(prior_likelihood)

        return prior_likelihood - posterior_likelihood

    def aggregate(self, fit_losses, dist_losses, beta_param) -> torch.Tensor:
        """
        This method aggregate dist_lossed and fit_losses for whole
        sampled parameters.

        Args:
            fit_losses (list):
                    list of data loss of each sample
            dist_losses (list):  {param_name: ParamDist}
                    list of distance loss of each sample
            beta:  {param_name: ParamDist}
                    sacle parameter of distance loss
        Returns:
            VarDistLoss.AggregationResult: Aggretion result for whole samples
        """
        fit_losses = torch.stack(fit_losses)
        dist_losses = torch.stack(dist_losses)

        stat_tensor = -fit_losses + beta_param * dist_losses

        if self.aggregation == "sample":
            positions = F.gumbel_softmax(
                ((1 - self.alpha) * stat_tensor.detach()), hard=True, dim=0
            )
        elif self.aggregation == "weighting":
            positions = F.softmax(
                ((1 - self.alpha) * stat_tensor.detach()), dim=0
            ).detach()
        else:
            raise ValueError()

        loss = -(positions * (stat_tensor)).sum()

        # print(loss)
        return VarDistLoss.AggregationResult(
            total_loss=loss, fit_loss=fit_losses.mean(), dist_loss=-dist_losses.mean()
        )

__init__(alpha=-1, aggregation='weighting')

summary

Parameters:

Name Type Description Default
alpha float

alpha

 -1
aggregation str

aggregation

 'weighting'
Source code in src/methods/bayes/variational/optimization.py 
218
219
220
221
222
223
224
225
226
227
228
def __init__(self, alpha: float = -1, aggregation: str = "weighting"):
    """_summary_

    Args:
        alpha (float): alpha
        aggregation (str): aggregation
    """
    assert aggregation in ["weighting", "sample"]
    self.alpha = alpha
    self.aggregation = aggregation
    super().__init__()

aggregate(fit_losses, dist_losses, beta_param)

This method aggregate dist_lossed and fit_losses for whole sampled parameters.

Parameters:

Name Type Description Default
fit_losses list 
list of data loss of each sample
required
dist_losses list

{param_name: ParamDist} list of distance loss of each sample

 required
beta

{param_name: ParamDist} sacle parameter of distance loss

 required

Returns: VarDistLoss.AggregationResult: Aggretion result for whole samples

Source code in src/methods/bayes/variational/optimization.py 
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
def aggregate(self, fit_losses, dist_losses, beta_param) -> torch.Tensor:
    """
    This method aggregate dist_lossed and fit_losses for whole
    sampled parameters.

    Args:
        fit_losses (list):
                list of data loss of each sample
        dist_losses (list):  {param_name: ParamDist}
                list of distance loss of each sample
        beta:  {param_name: ParamDist}
                sacle parameter of distance loss
    Returns:
        VarDistLoss.AggregationResult: Aggretion result for whole samples
    """
    fit_losses = torch.stack(fit_losses)
    dist_losses = torch.stack(dist_losses)

    stat_tensor = -fit_losses + beta_param * dist_losses

    if self.aggregation == "sample":
        positions = F.gumbel_softmax(
            ((1 - self.alpha) * stat_tensor.detach()), hard=True, dim=0
        )
    elif self.aggregation == "weighting":
        positions = F.softmax(
            ((1 - self.alpha) * stat_tensor.detach()), dim=0
        ).detach()
    else:
        raise ValueError()

    loss = -(positions * (stat_tensor)).sum()

    # print(loss)
    return VarDistLoss.AggregationResult(
        total_loss=loss, fit_loss=fit_losses.mean(), dist_loss=-dist_losses.mean()
    )

forward(param_sample_dict, posterior, prior)

This method computes loss for one sampled parameters.

Parameters:

Name Type Description Default
param_sample_dict dict[str, Parameter]

{param_name: nn.Parameter} sampled parameters on network.

 required
posterior dict[str, ParamDist]

{param_name: ParamDist} posterior distribution of net parameters.

 required
prior dict[str, ParamDist]

{param_name: ParamDist} prior distribution of net parameters.

 required

Returns:

Type Description
Tensor

torch.Tensor: distanse loss for one sampled parameters
Source code in src/methods/bayes/variational/optimization.py 
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def forward(
    self,
    param_sample_dict: dict[str, nn.Parameter],
    posterior: dict[str, ParamDist],
    prior: dict[str, ParamDist],
) -> torch.Tensor:
    """
    This method computes loss for one sampled parameters.

    Args:
        param_sample_dict: {param_name: nn.Parameter}
                sampled parameters on network.
        posterior:  {param_name: ParamDist}
                posterior distribution of net parameters.
        prior:  {param_name: ParamDist}
                prior distribution of net parameters.

    Returns:
        torch.Tensor: distanse loss for one sampled parameters
    """
    prior_likelihood: torch.tensor = 0.0
    posterior_likelihood: torch.tensor = 0.0

    for parameter, posterior_distr, prior_distr in zip(
        param_sample_dict.values(), posterior.values(), prior.values()
    ):
        prior_likelihood += prior_distr.log_prob(parameter).sum()
        posterior_likelihood += posterior_distr.log_prob(parameter).sum()
        # if key == "fc2.bias":
        # print(posterior_distr.loc)
        # print(prior_likelihood)

    return prior_likelihood - posterior_likelihood