core.distribution.utils

  1import math
  2from collections.abc import Callable, Iterator
  3
  4import torch
  5from torch import Tensor, nn
  6
  7from core.distribution import AbstractVariable
  8from core.layer.utils import LayerNameT, get_bayesian_torch_layers, get_torch_layers
  9
 10DistributionT = dict[LayerNameT, dict[str, AbstractVariable]]
 11
 12
 13def from_ivon(
 14    model: nn.Module,
 15    optimizer: "ivon.IVON",
 16    distribution: type[AbstractVariable],
 17    requires_grad: bool = True,
 18    get_layers_func: Callable[
 19        [nn.Module], Iterator[tuple[LayerNameT, nn.Module]]
 20    ] = get_torch_layers,
 21) -> DistributionT:
 22    """
 23    Construct a distribution from an IVON optimizer's parameters and Hessian approximations.
 24
 25    This function extracts weight and bias information (as well as Hessians) from the
 26    IVON optimizer, then creates instances of the specified `distribution` for each
 27    parameter. The newly created distributions can be used as a posterior or prior
 28    in a PAC-Bayes setting.
 29
 30    Args:
 31        model (nn.Module): The (deterministic) model whose parameters correspond to the IVON's weights.
 32        optimizer (ivon.IVON): An instance of IVON optimizer containing the Hessian approximations.
 33        distribution (Type[AbstractVariable]): The subclass of `AbstractVariable` to instantiate.
 34        requires_grad (bool, optional): If True, gradients will be computed on the newly created parameters.
 35        get_layers_func (Callable, optional): A function to retrieve model layers. Defaults to `get_torch_layers`.
 36
 37    Returns:
 38        DistributionT: A dictionary mapping layer names to dicts of {'weight': ..., 'bias': ...},
 39        each containing an instance of `AbstractVariable`.
 40    """
 41    distributions = {}
 42    i = 0
 43    shift = 0
 44    weights = optimizer.param_groups[0]["params"]
 45    hessians = optimizer.param_groups[0]["hess"]
 46    weight_decay = optimizer.param_groups[0]["weight_decay"]
 47    ess = optimizer.param_groups[0]["ess"]
 48    sigma = 1 / (ess * (hessians + weight_decay)).sqrt()
 49    rho = torch.log(torch.exp(sigma) - 1)
 50
 51    for name, layer in get_layers_func(model):
 52        if layer.weight is not None:
 53            weight_cutoff = shift + math.prod(layer.weight.shape)
 54            weight_distribution = distribution(
 55                mu=weights[i],
 56                rho=rho[shift:weight_cutoff].reshape(*layer.weight.shape),
 57                mu_requires_grad=requires_grad,
 58                rho_requires_grad=requires_grad,
 59            )
 60            shift = weight_cutoff
 61            i += 1
 62        else:
 63            weight_distribution = None
 64        if layer.bias is not None:
 65            bias_cutoff = shift + math.prod(layer.bias.shape)
 66            bias_distribution = distribution(
 67                mu=weights[i],
 68                rho=rho[shift:bias_cutoff].reshape(*layer.bias.shape),
 69                mu_requires_grad=requires_grad,
 70                rho_requires_grad=requires_grad,
 71            )
 72            shift = bias_cutoff
 73            i += 1
 74        else:
 75            bias_distribution = None
 76        distributions[name] = {"weight": weight_distribution, "bias": bias_distribution}
 77    return distributions
 78
 79
 80def from_flat_rho(
 81    model: nn.Module,
 82    rho: Tensor | list[float],
 83    distribution: type[AbstractVariable],
 84    requires_grad: bool = True,
 85    get_layers_func: Callable[
 86        [nn.Module], Iterator[tuple[LayerNameT, nn.Module]]
 87    ] = get_torch_layers,
 88) -> DistributionT:
 89    """
 90    Create distributions for each layer using a shared or flat `rho` array.
 91
 92    This function takes a model and a `rho` tensor/list containing values for all weight/bias
 93    elements in consecutive order. Each layer's `mu` is initialized from the current layer weights,
 94    and `rho` is reshaped accordingly for the layer's shape.
 95
 96    Args:
 97        model (nn.Module): The PyTorch model whose layers are converted into distributions.
 98        rho (Union[Tensor, List[float]]): A 1D tensor or list of floating values used to set `rho`.
 99        distribution (Type[AbstractVariable]): The subclass of `AbstractVariable` to instantiate.
100        requires_grad (bool, optional): If True, gradients will be computed for the distribution parameters.
101        get_layers_func (Callable, optional): Function for iterating over the model layers. Defaults to `get_torch_layers`.
102
103    Returns:
104        DistributionT: A dictionary of layer distributions keyed by layer names, each
105        containing 'weight' and 'bias' distributions if they exist.
106    """
107    distributions = {}
108    shift = 0
109    for name, layer in get_layers_func(model):
110        if layer.weight is not None:
111            # weight_cutoff = shift + layer.out_features * layer.in_features
112            weight_cutoff = shift + math.prod(layer.weight.shape)
113            weight_distribution = distribution(
114                mu=layer.weight,
115                rho=rho[shift:weight_cutoff].reshape(*layer.weight.shape),
116                mu_requires_grad=requires_grad,
117                rho_requires_grad=requires_grad,
118            )
119        else:
120            weight_distribution = None
121        if layer.bias is not None:
122            # bias_cutoff = weight_cutoff + layer.out_features
123            bias_cutoff = weight_cutoff + math.prod(layer.bias.shape)
124            bias_distribution = distribution(
125                mu=layer.bias,
126                rho=rho[weight_cutoff:bias_cutoff].reshape(*layer.bias.shape),
127                mu_requires_grad=requires_grad,
128                rho_requires_grad=requires_grad,
129            )
130            shift = bias_cutoff
131        else:
132            bias_distribution = None
133            shift = weight_cutoff
134        distributions[name] = {"weight": weight_distribution, "bias": bias_distribution}
135    return distributions
136
137
138def _from_any(
139    model: nn.Module,
140    distribution: type[AbstractVariable],
141    requires_grad: bool,
142    get_layers_func: Callable[[nn.Module], Iterator[tuple[LayerNameT, nn.Module]]],
143    weight_mu_fill_func: Callable[[nn.Module], Tensor],
144    weight_rho_fill_func: Callable[[nn.Module], Tensor],
145    bias_mu_fill_func: Callable[[nn.Module], Tensor],
146    bias_rho_fill_func: Callable[[nn.Module], Tensor],
147    weight_exists: Callable[[nn.Module], bool] = lambda layer: hasattr(layer, "weight"),
148    bias_exists: Callable[[nn.Module], bool] = lambda layer: hasattr(layer, "bias"),
149) -> DistributionT:
150    """
151    A helper function to create a distribution for each layer in a model using
152    user-provided fill functions for `mu` and `rho`.
153
154    Args:
155        model (nn.Module): The model to convert into distributions.
156        distribution (Type[AbstractVariable]): The type of variable distribution to instantiate.
157        requires_grad (bool): If True, gradients will be computed on the distribution parameters.
158        get_layers_func (Callable): A function to iterate over model layers and yield (name, layer) pairs.
159        weight_mu_fill_func (Callable): A function that returns a tensor for initializing the weight `mu`.
160        weight_rho_fill_func (Callable): A function that returns a tensor for initializing the weight `rho`.
161        bias_mu_fill_func (Callable): A function that returns a tensor for initializing the bias `mu`.
162        bias_rho_fill_func (Callable): A function that returns a tensor for initializing the bias `rho`.
163        weight_exists (Callable, optional): A predicate to check if a layer contains a weight attribute.
164        bias_exists (Callable, optional): A predicate to check if a layer contains a bias attribute.
165
166    Returns:
167        DistributionT: A dictionary mapping each layer to weight/bias distributions.
168    """
169    distributions = {}
170    for name, layer in get_layers_func(model):
171        if weight_exists(layer) and layer.weight is not None:
172            weight_distribution = distribution(
173                mu=weight_mu_fill_func(layer),
174                rho=weight_rho_fill_func(layer),
175                mu_requires_grad=requires_grad,
176                rho_requires_grad=requires_grad,
177            )
178        else:
179            weight_distribution = None
180        if bias_exists(layer) and layer.bias is not None:
181            bias_distribution = distribution(
182                mu=bias_mu_fill_func(layer),
183                rho=bias_rho_fill_func(layer),
184                mu_requires_grad=requires_grad,
185                rho_requires_grad=requires_grad,
186            )
187        else:
188            bias_distribution = None
189        distributions[name] = {"weight": weight_distribution, "bias": bias_distribution}
190    return distributions
191
192
193def from_random(
194    model: nn.Module,
195    rho: Tensor,
196    distribution: type[AbstractVariable],
197    requires_grad: bool = True,
198    get_layers_func: Callable[
199        [nn.Module], Iterator[tuple[LayerNameT, nn.Module]]
200    ] = get_torch_layers,
201) -> DistributionT:
202    """
203    Create a distribution for each layer with randomly initialized mean (using truncated normal)
204    and a constant `rho` value.
205
206    Args:
207        model (nn.Module): The target PyTorch model.
208        rho (Tensor): A single scalar tensor defining the initial `rho` for all weights/biases.
209        distribution (Type[AbstractVariable]): The class for creating each weight/bias distribution.
210        requires_grad (bool, optional): If True, allows gradient-based updates of `mu` and `rho`.
211        get_layers_func (Callable, optional): Function to iterate over model layers. Defaults to `get_torch_layers`.
212
213    Returns:
214        DistributionT: A dictionary containing layer-wise distributions for weights and biases.
215    """
216
217    def get_truncated_normal_fill_tensor(layer: nn.Module) -> Tensor:
218        t = torch.Tensor(*layer.weight.shape)
219        if hasattr(layer, "weight") and layer.weight is not None:
220            in_features = math.prod(layer.weight.shape[1:])
221        else:
222            raise ValueError(f"Unsupported layer of type: {type(layer)}")
223        w = 1 / math.sqrt(in_features)
224        return truncated_normal_fill_tensor(t, 0, w, -2 * w, 2 * w)
225
226    return _from_any(
227        model,
228        distribution,
229        requires_grad,
230        get_layers_func,
231        weight_mu_fill_func=get_truncated_normal_fill_tensor,
232        weight_rho_fill_func=lambda layer: torch.ones(*layer.weight.shape) * rho,
233        bias_mu_fill_func=lambda layer: torch.zeros(*layer.bias.shape),
234        bias_rho_fill_func=lambda layer: torch.ones(*layer.bias.shape) * rho,
235    )
236
237
238def from_zeros(
239    model: nn.Module,
240    rho: Tensor,
241    distribution: type[AbstractVariable],
242    requires_grad: bool = True,
243    get_layers_func: Callable[
244        [nn.Module], Iterator[tuple[LayerNameT, nn.Module]]
245    ] = get_torch_layers,
246) -> DistributionT:
247    """
248    Create distributions for each layer by setting `mu` to zero and `rho` to a constant value.
249
250    Args:
251        model (nn.Module): The PyTorch model.
252        rho (Tensor): A scalar defining the initial `rho` for all weights/biases.
253        distribution (Type[AbstractVariable]): Distribution class to instantiate.
254        requires_grad (bool, optional): Whether to track gradients for `mu` and `rho`.
255        get_layers_func (Callable, optional): Layer iteration function. Defaults to `get_torch_layers`.
256
257    Returns:
258        DistributionT: A dictionary mapping layer names to weight/bias distributions.
259    """
260    return _from_any(
261        model,
262        distribution,
263        requires_grad,
264        get_layers_func,
265        weight_mu_fill_func=lambda layer: torch.zeros(*layer.weight.shape),
266        weight_rho_fill_func=lambda layer: torch.ones(*layer.weight.shape) * rho,
267        bias_mu_fill_func=lambda layer: torch.zeros(*layer.bias.shape),
268        bias_rho_fill_func=lambda layer: torch.ones(*layer.bias.shape) * rho,
269    )
270
271
272def from_layered(
273    model: torch.nn.Module,
274    attribute_mapping: dict[str, str],
275    distribution: type[AbstractVariable],
276    requires_grad: bool = True,
277    get_layers_func: Callable[
278        [nn.Module], Iterator[tuple[LayerNameT, nn.Module]]
279    ] = get_torch_layers,
280) -> DistributionT:
281    """
282    Create distributions by extracting `mu` and `rho` from specified attributes in the model layers.
283
284    This function looks up layer attributes for weight and bias (e.g., "weight_mu", "weight_rho")
285    using `attribute_mapping`, then initializes each distribution accordingly.
286
287    Args:
288        model (nn.Module): The model whose layers contain the specified attributes.
289        attribute_mapping (dict[str, str]): A mapping of attribute names, for example:
290            {
291              "weight_mu": "mu_weight",
292              "weight_rho": "rho_weight",
293              "bias_mu": "mu_bias",
294              "bias_rho": "rho_bias"
295            }
296        distribution (Type[AbstractVariable]): The class used to create weight/bias distributions.
297        requires_grad (bool, optional): If True, gradients will be computed on `mu` and `rho`.
298        get_layers_func (Callable, optional): Layer iteration function.
299
300    Returns:
301        DistributionT: A dictionary of distributions keyed by layer names.
302    """
303    return _from_any(
304        model,
305        distribution,
306        requires_grad,
307        get_layers_func,
308        weight_exists=lambda layer: hasattr(layer, attribute_mapping["weight_mu"])
309        and hasattr(layer, attribute_mapping["weight_rho"]),
310        bias_exists=lambda layer: hasattr(layer, attribute_mapping["weight_mu"])
311        and hasattr(layer, attribute_mapping["weight_rho"]),
312        weight_mu_fill_func=lambda layer: layer.__getattr__(
313            attribute_mapping["weight_mu"]
314        )
315        .detach()
316        .clone(),
317        weight_rho_fill_func=lambda layer: layer.__getattr__(
318            attribute_mapping["weight_rho"]
319        )
320        .detach()
321        .clone(),
322        bias_mu_fill_func=lambda layer: layer.__getattr__(attribute_mapping["bias_mu"])
323        .detach()
324        .clone(),
325        bias_rho_fill_func=lambda layer: layer.__getattr__(
326            attribute_mapping["bias_rho"]
327        )
328        .detach()
329        .clone(),
330    )
331
332
333def from_bnn(
334    model: nn.Module,
335    distribution: type[AbstractVariable],
336    requires_grad: bool = True,
337    get_layers_func: Callable[
338        [nn.Module], Iterator[tuple[LayerNameT, nn.Module]]
339    ] = get_bayesian_torch_layers,
340) -> DistributionT:
341    """
342    Construct distributions by reading the attributes (e.g., mu_weight, rho_weight, mu_bias, rho_bias)
343    from layers typically found in BayesianTorch modules.
344
345    Args:
346        model (nn.Module): The Bayesian Torch model containing layer attributes such as mu_weight, rho_weight, etc.
347        distribution (Type[AbstractVariable]): The subclass of `AbstractVariable` for each parameter.
348        requires_grad (bool, optional): If True, allows gradient-based optimization of `mu` and `rho`.
349        get_layers_func (Callable, optional): A function that retrieves BayesianTorch layers. Defaults to `get_bayesian_torch_layers`.
350
351    Returns:
352        DistributionT: A dictionary mapping layer names to weight/bias distributions.
353    """
354    distributions = {}
355    for name, layer in get_layers_func(model):
356        if hasattr(layer, "mu_weight") and hasattr(layer, "rho_weight"):
357            weight_distribution = distribution(
358                mu=layer.__getattr__("mu_weight").detach().clone(),
359                rho=layer.__getattr__("rho_weight").detach().clone(),
360                mu_requires_grad=requires_grad,
361                rho_requires_grad=requires_grad,
362            )
363        elif hasattr(layer, "mu_kernel") and hasattr(layer, "rho_kernel"):
364            weight_distribution = distribution(
365                mu=layer.__getattr__("mu_kernel").detach().clone(),
366                rho=layer.__getattr__("rho_kernel").detach().clone(),
367                mu_requires_grad=requires_grad,
368                rho_requires_grad=requires_grad,
369            )
370        else:
371            weight_distribution = None
372        if hasattr(layer, "mu_bias") and hasattr(layer, "rho_bias"):
373            bias_distribution = distribution(
374                mu=layer.__getattr__("mu_bias").detach().clone(),
375                rho=layer.__getattr__("rho_bias").detach().clone(),
376                mu_requires_grad=requires_grad,
377                rho_requires_grad=requires_grad,
378            )
379        else:
380            bias_distribution = None
381        distributions[name] = {"weight": weight_distribution, "bias": bias_distribution}
382    return distributions
383
384
385def from_copy(
386    dist: DistributionT,
387    distribution: type[AbstractVariable],
388    requires_grad: bool = True,
389) -> DistributionT:
390    """
391    Create a new distribution by copying `mu` and `rho` from an existing distribution.
392
393    Args:
394        dist (DistributionT): A distribution dictionary to copy from.
395        distribution (Type[AbstractVariable]): The class to instantiate for each weight/bias.
396        requires_grad (bool, optional): If True, the new distribution parameters can be updated via gradients.
397
398    Returns:
399        DistributionT: A new distribution dictionary with the same layer structure,
400        but new `mu` and `rho` parameters cloned from `dist`.
401    """
402    distributions = {}
403    for name, layer in dist.items():
404        weight_distribution = distribution(
405            mu=layer["weight"].mu.detach().clone(),
406            rho=layer["weight"].rho.detach().clone(),
407            mu_requires_grad=requires_grad,
408            rho_requires_grad=requires_grad,
409        )
410        if layer["bias"] is not None:
411            bias_distribution = distribution(
412                mu=layer["bias"].mu.detach().clone(),
413                rho=layer["bias"].rho.detach().clone(),
414                mu_requires_grad=requires_grad,
415                rho_requires_grad=requires_grad,
416            )
417        else:
418            bias_distribution = None
419        distributions[name] = {"weight": weight_distribution, "bias": bias_distribution}
420    return distributions
421
422
423def compute_kl(dist1: DistributionT, dist2: DistributionT) -> Tensor:
424    """
425    Compute the total KL divergence between two distributions of the same structure.
426
427    Each corresponding layer's weight/bias KL is summed to produce a single scalar.
428
429    Args:
430        dist1 (DistributionT): The first distribution dictionary.
431        dist2 (DistributionT): The second distribution dictionary.
432
433    Returns:
434        Tensor: A scalar tensor representing the total KL divergence across all layers.
435    """
436    kl_list = []
437    for idx in dist1:
438        for key in dist1[idx]:
439            if dist1[idx][key] is not None and dist2[idx][key] is not None:
440                kl = dist1[idx][key].compute_kl(dist2[idx][key])
441                kl_list.append(kl)
442    return torch.stack(kl_list).sum()
443
444
445def compute_standard_normal_cdf(x: float) -> float:
446    """
447    Compute the cumulative distribution function (CDF) of a standard normal at point x.
448
449    Args:
450        x (float): The input value at which to evaluate the standard normal CDF.
451
452    Returns:
453        float: The CDF value of the standard normal distribution at x.
454    """
455    # TODO: replace with numpy
456    return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
457
458
459def truncated_normal_fill_tensor(
460    tensor: torch.Tensor,
461    mean: float = 0.0,
462    std: float = 1.0,
463    a: float = -2.0,
464    b: float = 2.0,
465) -> torch.Tensor:
466    """
467    Fill a tensor in-place with values drawn from a truncated normal distribution.
468
469    The resulting values lie in the interval [a, b], centered around `mean`
470    with approximate std `std`.
471
472    Args:
473        tensor (torch.Tensor): The tensor to fill.
474        mean (float, optional): Mean of the desired distribution. Defaults to 0.0.
475        std (float, optional): Standard deviation of the desired distribution. Defaults to 1.0.
476        a (float, optional): Lower bound of truncation. Defaults to -2.0.
477        b (float, optional): Upper bound of truncation. Defaults to 2.0.
478
479    Returns:
480        torch.Tensor: The same tensor, filled in-place with truncated normal values.
481    """
482    with torch.no_grad():
483        # Get upper and lower cdf values
484        l_ = compute_standard_normal_cdf((a - mean) / std)
485        u_ = compute_standard_normal_cdf((b - mean) / std)
486
487        # Fill tensor with uniform values from [l_, u_]
488        tensor.uniform_(l_, u_)
489
490        # Use inverse cdf transform from normal distribution
491        tensor.mul_(2)
492        tensor.sub_(1)
493
494        # Ensure that the values are strictly between -1 and 1 for erfinv
495        eps = torch.finfo(tensor.dtype).eps
496        tensor.clamp_(min=-(1.0 - eps), max=(1.0 - eps))
497        tensor.erfinv_()
498
499        # Transform to proper mean, std
500        tensor.mul_(std * math.sqrt(2.0))
501        tensor.add_(mean)
502
503        # Clamp one last time to ensure it's still in the proper range
504        tensor.clamp_(min=a, max=b)
505        return tensor