Loss Kl Divergence

kl_divergence

CLASS	DESCRIPTION
KLDivergenceLoss	Kullback-Leibler divergence loss with optional sample weights.

FUNCTION	DESCRIPTION
kl_divergence_loss	Compute Kullback-Leibler divergence loss with optional weighting.

Classes

KLDivergenceLoss(log_prob: bool = False, reduce: Literal['mean', 'sum'] = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps, context_prefix: str | None = None)

Bases: Loss

Kullback-Leibler divergence loss with optional sample weights.

Computes \(\mathrm{KL}(P \| Q) = \sum_k w_k P(x_k) \log(P(x_k) / Q(x_k))\), where \(P\) and \(Q\) are probability distributions and \(w_k\) are optional sample weights.

PARAMETER	DESCRIPTION
log_prob	Whether inputs are log probabilities. TYPE: bool, optional, by default False DEFAULT: False
reduce	Reduction method. TYPE: Literal["mean", "sum"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
eps	Small value for numerical stability. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

METHOD	DESCRIPTION
forward	Compute KL divergence loss.

Source code in spectre/loss/kl_divergence.py

def __init__(
    self,
    log_prob: bool = False,
    reduce: Literal["mean", "sum"] = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
    context_prefix: str | None = None,
) -> None:
    super().__init__(reduce=reduce, sign=sign, context_prefix=context_prefix)

    if not isinstance(log_prob, bool):
        raise TypeError("`log_prob` must be a boolean.")
    self.log_prob = log_prob

    if not isinstance(eps, float):
        raise TypeError(f"eps must be a float, got {type(eps).__name__}.")
    check_in_interval(eps, "(0, inf)")
    self.eps = eps

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute KL divergence loss.

PARAMETER	DESCRIPTION
context	Context dictionary containing the tensors P and Q. TYPE: dict
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	KL divergence loss.

Source code in spectre/loss/kl_divergence.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute KL divergence loss.

    Parameters
    ----------
    context : dict
        Context dictionary containing the tensors `P` and `Q`.

    **kwargs
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        KL divergence loss.
    """
    P = self._get_context(context, "P", None)
    Q = self._get_context(context, "Q", None)
    weights = self._get_context(context, "weights", None)

    if P is None or Q is None:
        raise ValueError("`KLDivergenceLoss` requires `P` and `Q` in context.")

    if weights is None:
        weights = torch.ones(P.shape[0], device=P.device, dtype=P.dtype)

    return _kl_divergence_loss_impl(
        P,
        Q,
        weights,
        reduce=self.reduce,
        log_prob=self.log_prob,
        sign=self.sign,
    )

Functions

kl_divergence_loss(P: torch.Tensor, Q: torch.Tensor, weights: Optional[torch.Tensor] = None, reduce: str = 'mean', log_prob: bool = False, sign: float = 1.0) -> torch.Tensor

Compute Kullback-Leibler divergence loss with optional weighting.

Computes \(\mathrm{KL}(P \| Q) = \sum_k w_k P(x_k) \log(P(x_k) / Q(x_k))\), where \(P\) and \(Q\) are probability distributions and \(w_k\) are optional sample weights.

PARAMETER	DESCRIPTION
P	Reference distribution (2D tensor). TYPE: Tensor
Q	Target distribution (2D tensor). TYPE: Tensor
weights	Sample weights (1D tensor). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
reduce	Reduction method. Options: "sum", "mean". TYPE: str, optional, by default "mean" DEFAULT: 'mean'
log_prob	Whether inputs are log probabilities. TYPE: bool, optional, by default False DEFAULT: False
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	KL divergence loss.

Source code in spectre/loss/kl_divergence.py

def kl_divergence_loss(
    P: torch.Tensor,
    Q: torch.Tensor,
    weights: Optional[torch.Tensor] = None,
    reduce: str = "mean",
    log_prob: bool = False,
    sign: float = 1.0,
) -> torch.Tensor:
    """
    Compute Kullback-Leibler divergence loss with optional weighting.

    Computes $\\mathrm{KL}(P \\| Q) = \\sum_k w_k P(x_k) \\log(P(x_k) / Q(x_k))$,
    where $P$ and $Q$ are probability distributions and $w_k$ are optional sample
    weights.

    Parameters
    ----------
    P : torch.Tensor
        Reference distribution (2D tensor).

    Q : torch.Tensor
        Target distribution (2D tensor).

    weights : torch.Tensor | None, optional, by default None
        Sample weights (1D tensor).

    reduce : str, optional, by default "mean"
        Reduction method. Options: "sum", "mean".

    log_prob : bool, optional, by default False
        Whether inputs are log probabilities.

    sign : float, optional, by default 1.0
        Loss sign multiplier.

    Returns
    -------
    torch.Tensor
        KL divergence loss.
    """
    check_same_shape(P, Q)
    check_2d(P)
    check_2d(Q)

    if weights is not None:
        check_1d(weights)
        check_same_len(P, weights)
        check_all_non_negative(weights)
    else:
        weights = torch.ones(P.shape[0], device=P.device, dtype=P.dtype)

    if reduce not in ["sum", "mean"]:
        raise ValueError("Unknown reduce operation. Options: 'sum', 'mean'.")

    return _kl_divergence_loss_impl(P, Q, weights, reduce, log_prob, sign)