Loss Fisher

fisher

CLASS	DESCRIPTION
FisherSpectralLoss	Fisher spectral criterion for automatic spectral gap discovery.
VariationalSpectralLoss	Variational spectral loss with prior over number of states.

FUNCTION	DESCRIPTION
fisher_spectral_loss	Fisher spectral criterion for automatic spectral gap discovery.
variational_spectral_loss	Variational spectral loss with prior over number of states.

Classes

FisherSpectralLoss(min_states: int = 2, temperature: float = 1.0, var_floor: float = 0.005, sign: float = -1.0, context_prefix: str | None = None, **kwargs)

Bases: Loss

Fisher spectral criterion for automatic spectral gap discovery.

Treats the eigenvalue spectrum as a 1D clustering problem: for every possible split position k, the Fisher discriminant ratio measures how well the eigenvalues separate into a "signal" group (lambda_1, ..., lambda_k) and a "noise" group (lambda_{k+1}, ..., lambda_n).

The loss maximizes the best split quality via a smooth logsumexp approximation over all candidate positions.

Mathematically, for split position \(k\):

\[J(k) = \frac{(\mu_1 - \mu_2)^2} {\sigma_1^2 + \sigma_2^2 + \epsilon_{\text{var}}}\]

where \(\mu_1, \sigma_1^2\) are the mean and variance of the top \(k\) eigenvalues, and \(\mu_2, \sigma_2^2\) of the remaining \(n - k\).

The overall loss is the smooth maximum over all valid splits:

\[L = -T \cdot \text{logsumexp}(J(k) / T) \quad \text{for } k \in [k_{\min}, n{-}2]\]

Unlike gap-based losses that only act on the two eigenvalues adjacent to the gap, Fisher shapes the entire spectrum.

PARAMETER	DESCRIPTION
min_states	Minimum number of states to consider. Splits with fewer than min_states eigenvalues in the signal group are excluded. Must be >= 2 to avoid trivial solutions. TYPE: int, optional, by default 2 DEFAULT: 2
temperature	Temperature for logsumexp smooth maximum. Low temperature: sharper selection of best split (closer to hard argmax) High temperature: more uniform weighting across split positions TYPE: float, optional, by default 1.0 DEFAULT: 1.0
var_floor	Floor for within-class variance to prevent gradient explosion. When clusters become very tight (variance -> 0), the Fisher ratio diverges. The variance floor caps the ratio and stabilizes gradients. Values in [0.001, 0.01] work well for eigenvalues in [0, 1]. TYPE: float, optional, by default 0.005 DEFAULT: 0.005
sign	Loss sign multiplier. -1.0 for maximization (default). TYPE: float, optional, by default -1.0 DEFAULT: -1.0
**kwargs	Additional keyword arguments passed to Loss. DEFAULT: {}

ATTRIBUTE	DESCRIPTION
allowed_reduce	Allowed reduction operations (only ["none"] for this loss). TYPE: list[str]
loss_types	Loss function categories (SPECTRAL). TYPE: list[LossType]

Examples:

Automatic gap discovery (no predefined number of states):

>>> import torch
>>> from spectre.loss import FisherSpectralLoss
>>>
>>> loss_fn = FisherSpectralLoss(min_states=2, temperature=1.0)
>>> eigenvalues = torch.tensor([1.0, 0.95, 0.90, 0.30, 0.25, 0.20])
>>> context = {"eigenvalues": eigenvalues}
>>> loss = loss_fn(context)

METHOD	DESCRIPTION
forward	Compute Fisher spectral loss.

Source code in spectre/loss/fisher.py

def __init__(
    self,
    min_states: int = 2,
    temperature: float = 1.0,
    var_floor: float = 0.005,
    sign: float = -1.0,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce="none", sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(min_states, int):
        raise TypeError(f"min_states must be int, got {type(min_states)}")
    check_in_interval(min_states, "[2, inf)")
    self.min_states = min_states

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

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

Functions

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

Compute Fisher spectral loss.

PARAMETER	DESCRIPTION
context	Context dictionary, must contain "eigenvalues" key. TYPE: dict
**kwargs	Additional keyword arguments. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Scalar Fisher spectral loss.

Source code in spectre/loss/fisher.py

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

    Parameters
    ----------
    context : dict
        Context dictionary, must contain ``"eigenvalues"`` key.

    **kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Scalar Fisher spectral loss.
    """
    eigenvalues = self._get_context(context, "eigenvalues", None)

    if eigenvalues is None:
        raise ValueError("FisherSpectralLoss requires eigenvalues in context.")

    return fisher_spectral_loss(
        eigenvalues,
        min_states=self.min_states,
        temperature=self.temperature,
        var_floor=self.var_floor,
        sign=self.sign,
    )

VariationalSpectralLoss(min_states: int = 2, temperature: float = 1.0, var_floor: float = 0.005, prior: Literal['uniform', 'geometric'] | torch.Tensor = 'uniform', prior_param: float | None = None, prior_weight: float = 0.1, sign: float = -1.0, context_prefix: str | None = None, **kwargs)

Bases: Loss

Variational spectral loss with prior over number of states.

Extends FisherSpectralLoss by treating the number of states k as a discrete latent variable with a learned posterior and a prior. The loss maximizes a variational lower bound (ELBO):

\[L = -\tau \cdot \text{logsumexp}(J(k) / \tau) + \beta \cdot \text{KL}(q(k) \| p(k))\]

where:

• \(J(k)\) is the Fisher discriminant ratio at split position \(k\)
• \(q(k | \lambda) = \text{softmax}(J(k) / \tau)\) is the posterior over number of states, derived from the eigenvalue spectrum
• \(p(k)\) is a prior encoding domain knowledge
• \(\tau\) is the temperature and \(\beta\) controls prior strength

The first term is the Fisher spectral loss (data fit). The KL term regularizes the posterior toward the prior, resolving ambiguity when multiple split positions have similar Fisher scores.

The posterior \(q(k)\) can be monitored during training to track how the model's belief about the number of states evolves.

PARAMETER	DESCRIPTION
min_states	Minimum number of states to consider. TYPE: int, optional, by default 2 DEFAULT: 2
temperature	Temperature for logsumexp and posterior softmax. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
var_floor	Floor for within-class variance (see FisherSpectralLoss). TYPE: float, optional, by default 0.005 DEFAULT: 0.005
prior	Prior distribution over number of states. "uniform": no preference (KL reduces to entropy regularization) "geometric": \(p(k) \propto \alpha^{k - k_{\min}}\), preferring fewer states (Occam's razor). Controlled by prior_param. torch.Tensor: explicit log-probabilities (unnormalized) with shape matching the number of valid splits. TYPE: ``"uniform"`` | ``"geometric"`` | torch.Tensor DEFAULT: 'uniform'
prior_param	Parameter for parametric priors. For "geometric": the decay rate \(\alpha \in (0, 1)\). Smaller values penalize larger state counts more strongly. Defaults to 0.5 if not specified. TYPE: float | None, optional, by default None DEFAULT: None
prior_weight	Weight \(\beta\) for the KL divergence term. \(\beta = 0\): pure Fisher loss (no prior influence) \(\beta > \tau\): prior dominates over data TYPE: float, optional, by default 0.1 DEFAULT: 0.1
sign	Loss sign multiplier. -1.0 for maximization. TYPE: float, optional, by default -1.0 DEFAULT: -1.0
**kwargs	Additional keyword arguments passed to Loss. DEFAULT: {}

ATTRIBUTE	DESCRIPTION
allowed_reduce	Allowed reduction operations (only ["none"]). TYPE: list[str]
loss_types	Loss function categories (SPECTRAL). TYPE: list[LossType]

Examples:

With geometric prior (prefer fewer states):

>>> import torch
>>> from spectre.loss import VariationalSpectralLoss
>>>
>>> loss_fn = VariationalSpectralLoss(
...     prior="geometric",
...     prior_param=0.5,
...     prior_weight=0.5,
... )
>>> eigenvalues = torch.tensor([1.0, 0.9, 0.8, 0.3, 0.2, 0.1])
>>> context = {"eigenvalues": eigenvalues}
>>> loss = loss_fn(context)

With uniform prior (equivalent to Fisher + entropy bonus):

>>> loss_fn = VariationalSpectralLoss(
...     prior="uniform",
...     prior_weight=0.1,
... )

METHOD	DESCRIPTION
forward	Compute variational spectral loss.

Source code in spectre/loss/fisher.py

def __init__(
    self,
    min_states: int = 2,
    temperature: float = 1.0,
    var_floor: float = 0.005,
    prior: Literal["uniform", "geometric"] | torch.Tensor = "uniform",
    prior_param: float | None = None,
    prior_weight: float = 0.1,
    sign: float = -1.0,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce="none", sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(min_states, int):
        raise TypeError(f"min_states must be int, got {type(min_states)}")
    check_in_interval(min_states, "[2, inf)")
    self.min_states = min_states

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

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

    if isinstance(prior, str) and prior not in ("uniform", "geometric"):
        raise ValueError(
            f"prior must be 'uniform', 'geometric', or a Tensor, got '{prior}'"
        )
    self.prior = prior

    if prior_param is not None:
        if not isinstance(prior_param, (int, float)):
            raise TypeError(
                f"prior_param must be numeric or None, got {type(prior_param)}"
            )
        if isinstance(prior, str) and prior == "geometric":
            check_in_interval(float(prior_param), "(0, 1)")
    self.prior_param = float(prior_param) if prior_param is not None else None

    if not isinstance(prior_weight, (int, float)):
        raise TypeError(f"prior_weight must be numeric, got {type(prior_weight)}")
    check_in_interval(prior_weight, "[0, inf)")
    self.prior_weight = float(prior_weight)

Functions

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

Compute variational spectral loss.

PARAMETER	DESCRIPTION
context	Context dictionary, must contain "eigenvalues" key. TYPE: dict
**kwargs	Additional keyword arguments. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Scalar variational spectral loss.

Source code in spectre/loss/fisher.py

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

    Parameters
    ----------
    context : dict
        Context dictionary, must contain ``"eigenvalues"`` key.

    **kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Scalar variational spectral loss.
    """
    eigenvalues = self._get_context(context, "eigenvalues", None)

    if eigenvalues is None:
        raise ValueError("VariationalSpectralLoss requires eigenvalues in context.")

    return variational_spectral_loss(
        eigenvalues,
        min_states=self.min_states,
        temperature=self.temperature,
        var_floor=self.var_floor,
        prior=self.prior,
        prior_param=self.prior_param,
        prior_weight=self.prior_weight,
        sign=self.sign,
    )

Functions

fisher_spectral_loss(eigenvalues: torch.Tensor, min_states: int = 2, temperature: float = 1.0, var_floor: float = 0.005, sign: float = -1.0) -> torch.Tensor

Fisher spectral criterion for automatic spectral gap discovery.

Evaluates the Fisher discriminant ratio at every possible split position in the eigenvalue spectrum, then takes a smooth maximum via logsumexp.

For each split at position k, the eigenvalues are partitioned into a signal group (top k) and a noise group (remaining n - k). The Fisher ratio measures the quality of this partition:

\[J(k) = \frac{(\mu_1 - \mu_2)^2} {\sigma_1^2 + \sigma_2^2 + \epsilon_{\text{var}}}\]

The loss is:

\[L = \text{sign} \cdot T \cdot \text{logsumexp}(J(k) / T)\]

The gradient with respect to each eigenvalue lambda_i acts on all eigenvalues simultaneously:

• Signal group eigenvalues are pushed toward their group mean (tightening)
• Noise group eigenvalues are pushed toward their group mean (tightening)
• The gap between group means is widened

This provides a much richer learning signal than gap-based losses, which only act on the two eigenvalues adjacent to the spectral gap.

PARAMETER	DESCRIPTION
eigenvalues	Eigenvalues in descending order, shape (n_eigenvalues,). TYPE: Tensor
min_states	Minimum number of states. Splits below this are excluded. TYPE: int, optional, by default 2 DEFAULT: 2
temperature	logsumexp temperature. Lower values focus on the best split. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
var_floor	Floor for within-class variance to prevent gradient explosion when clusters become very tight. TYPE: float, optional, by default 0.005 DEFAULT: 0.005
sign	Loss sign (-1 for maximization). TYPE: float, optional, by default -1.0 DEFAULT: -1.0

RETURNS	DESCRIPTION
Tensor	Scalar loss value.

RAISES	DESCRIPTION
ValueError	If eigenvalues are not sorted descending or the valid split range is empty.

Source code in spectre/loss/fisher.py

def fisher_spectral_loss(
    eigenvalues: torch.Tensor,
    min_states: int = 2,
    temperature: float = 1.0,
    var_floor: float = 0.005,
    sign: float = -1.0,
) -> torch.Tensor:
    """
    Fisher spectral criterion for automatic spectral gap discovery.

    Evaluates the Fisher discriminant ratio at every possible split position in
    the eigenvalue spectrum, then takes a smooth maximum via `logsumexp`.

    For each split at position `k`, the eigenvalues are partitioned into a
    signal group (top `k`) and a noise group (remaining `n - k`). The Fisher
    ratio measures the quality of this partition:

    $$J(k) = \\frac{(\\mu_1 - \\mu_2)^2}
    {\\sigma_1^2 + \\sigma_2^2 + \\epsilon_{\\text{var}}}$$

    The loss is:

    $$L = \\text{sign} \\cdot T \\cdot \\text{logsumexp}(J(k) / T)$$

    The gradient with respect to each eigenvalue `lambda_i` acts on **all**
    eigenvalues simultaneously:

    - Signal group eigenvalues are pushed toward their group mean (tightening)
    - Noise group eigenvalues are pushed toward their group mean (tightening)
    - The gap between group means is widened

    This provides a much richer learning signal than gap-based losses, which
    only act on the two eigenvalues adjacent to the spectral gap.

    Parameters
    ----------
    eigenvalues : torch.Tensor
        Eigenvalues in descending order, shape ``(n_eigenvalues,)``.

    min_states : int, optional, by default 2
        Minimum number of states. Splits below this are excluded.

    temperature : float, optional, by default 1.0
        `logsumexp` temperature. Lower values focus on the best split.

    var_floor : float, optional, by default 0.005
        Floor for within-class variance to prevent gradient explosion
        when clusters become very tight.

    sign : float, optional, by default -1.0
        Loss sign (-1 for maximization).

    Returns
    -------
    torch.Tensor
        Scalar loss value.

    Raises
    ------
    ValueError
        If eigenvalues are not sorted descending or the valid split range is
        empty.
    """
    check_1d(eigenvalues)
    check_eigenvalues_sorted(eigenvalues, descending=True, raise_error=True)

    n = len(eigenvalues)

    # Need at least min_states in signal group and 2 in noise group
    # to compute meaningful variance for both groups
    if n < min_states + 2:
        raise ValueError(
            f"Need at least {min_states + 2} eigenvalues for min_states={min_states}, "
            f"got {n}."
        )

    device, dtype = eigenvalues.device, eigenvalues.dtype

    # Vectorized computation via cumulative sums: O(n) instead of O(n^2)
    cumsum = torch.cumsum(eigenvalues, dim=0)
    cumsum_sq = torch.cumsum(eigenvalues**2, dim=0)
    total = cumsum[-1]
    total_sq = cumsum_sq[-1]

    # Split positions: k = min_states, ..., n-2
    # k eigenvalues in signal group, n-k in noise group
    ks = torch.arange(min_states, n - 1, device=device, dtype=dtype)
    ki = ks.long() - 1  # 0-indexed for cumsum lookup
    n_k = n - ks  # noise group sizes

    # Group 1 (signal): first k eigenvalues
    sum1 = cumsum[ki]
    sum_sq1 = cumsum_sq[ki]
    mu1 = sum1 / ks
    var1 = torch.clamp(sum_sq1 / ks - mu1**2, min=0.0)

    # Group 2 (noise): remaining n-k eigenvalues
    sum2 = total - sum1
    sum_sq2 = total_sq - sum_sq1
    mu2 = sum2 / n_k
    var2 = torch.clamp(sum_sq2 / n_k - mu2**2, min=0.0)

    # Fisher discriminant ratio with variance floor
    between = (mu1 - mu2) ** 2
    within = var1 + var2 + var_floor
    fisher_scores = between / within

    # Smooth maximum via logsumexp
    loss = temperature * torch.logsumexp(fisher_scores / temperature, dim=0)

    return sign * loss

variational_spectral_loss(eigenvalues: torch.Tensor, min_states: int = 2, temperature: float = 1.0, var_floor: float = 0.005, prior: Literal['uniform', 'geometric'] | torch.Tensor = 'uniform', prior_param: float | None = None, prior_weight: float = 0.1, sign: float = -1.0) -> torch.Tensor

Variational spectral loss with prior over number of states.

Combines the Fisher spectral criterion with a KL divergence term that regularizes the posterior \(q(k)\) toward a prior \(p(k)\):

\[L = \text{sign} \cdot \left[ \tau \cdot \text{logsumexp}(J(k)/\tau) - \beta \cdot \text{KL}(q \| p) \right]\]

The posterior is \(q(k) = \text{softmax}(J(k) / \tau)\) where \(J(k)\) is the Fisher discriminant ratio at split position \(k\).

When prior_weight=0, this reduces to fisher_spectral_loss.

PARAMETER	DESCRIPTION
eigenvalues	Eigenvalues in descending order, shape (n_eigenvalues,). TYPE: Tensor
min_states	Minimum number of states. TYPE: int, optional, by default 2 DEFAULT: 2
temperature	Temperature for logsumexp and posterior. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
var_floor	Floor for within-class variance. TYPE: float, optional, by default 0.005 DEFAULT: 0.005
prior	Prior distribution. See VariationalSpectralLoss for details. TYPE: ``"uniform"`` | ``"geometric"`` | torch.Tensor DEFAULT: 'uniform'
prior_param	Parameter for parametric priors (e.g., decay rate for "geometric"). TYPE: float | None, optional, by default None DEFAULT: None
prior_weight	Weight \(\beta\) for the KL term. TYPE: float, optional, by default 0.1 DEFAULT: 0.1
sign	Loss sign (-1 for maximization). TYPE: float, optional, by default -1.0 DEFAULT: -1.0

RETURNS	DESCRIPTION
Tensor	Scalar loss value.

RAISES	DESCRIPTION
ValueError	If eigenvalues are not sorted descending or the valid split range is empty.

Examples:

>>> import torch
>>> eigenvalues = torch.tensor([1.0, 0.9, 0.6, 0.5, 0.2, 0.1])
>>> loss = variational_spectral_loss(
...     eigenvalues,
...     prior="geometric",
...     prior_param=0.5,
...     prior_weight=0.5,
... )

Source code in spectre/loss/fisher.py

def variational_spectral_loss(
    eigenvalues: torch.Tensor,
    min_states: int = 2,
    temperature: float = 1.0,
    var_floor: float = 0.005,
    prior: Literal["uniform", "geometric"] | torch.Tensor = "uniform",
    prior_param: float | None = None,
    prior_weight: float = 0.1,
    sign: float = -1.0,
) -> torch.Tensor:
    """
    Variational spectral loss with prior over number of states.

    Combines the Fisher spectral criterion with a KL divergence
    term that regularizes the posterior $q(k)$ toward a prior $p(k)$:

    $$L = \\text{sign} \\cdot \\left[
    \\tau \\cdot \\text{logsumexp}(J(k)/\\tau)
    - \\beta \\cdot \\text{KL}(q \\| p)
    \\right]$$

    The posterior is $q(k) = \\text{softmax}(J(k) / \\tau)$ where
    $J(k)$ is the Fisher discriminant ratio at split position $k$.

    When ``prior_weight=0``, this reduces to `fisher_spectral_loss`.

    Parameters
    ----------
    eigenvalues : torch.Tensor
        Eigenvalues in descending order, shape ``(n_eigenvalues,)``.

    min_states : int, optional, by default 2
        Minimum number of states.

    temperature : float, optional, by default 1.0
        Temperature for `logsumexp` and posterior.

    var_floor : float, optional, by default 0.005
        Floor for within-class variance.

    prior : ``"uniform"`` | ``"geometric"`` | torch.Tensor
        Prior distribution. See `VariationalSpectralLoss` for details.

    prior_param : float | None, optional, by default None
        Parameter for parametric priors (e.g., decay rate for
        ``"geometric"``).

    prior_weight : float, optional, by default 0.1
        Weight $\\beta$ for the KL term.

    sign : float, optional, by default -1.0
        Loss sign (-1 for maximization).

    Returns
    -------
    torch.Tensor
        Scalar loss value.

    Raises
    ------
    ValueError
        If eigenvalues are not sorted descending or the valid split
        range is empty.

    Examples
    --------
    >>> import torch
    >>> eigenvalues = torch.tensor([1.0, 0.9, 0.6, 0.5, 0.2, 0.1])
    >>> loss = variational_spectral_loss(
    ...     eigenvalues,
    ...     prior="geometric",
    ...     prior_param=0.5,
    ...     prior_weight=0.5,
    ... )
    """
    check_1d(eigenvalues)
    check_eigenvalues_sorted(eigenvalues, descending=True, raise_error=True)

    n = len(eigenvalues)

    if n < min_states + 2:
        raise ValueError(
            f"Need at least {min_states + 2} eigenvalues for "
            f"min_states={min_states}, got {n}."
        )

    # Fisher scores and split positions
    scores, ks = _compute_fisher_scores(
        eigenvalues, min_states=min_states, var_floor=var_floor
    )

    # Fisher LSE term (data fit)
    fisher_lse = temperature * torch.logsumexp(scores / temperature, dim=0)

    # KL divergence term (prior regularization)
    kl = torch.tensor(0.0, device=eigenvalues.device, dtype=eigenvalues.dtype)
    if prior_weight > 0:
        log_p = _build_log_prior(ks, prior, prior_param)
        q = torch.softmax(scores / temperature, dim=0)
        log_q = torch.log(q + 1e-10)
        kl = (q * (log_q - log_p)).sum()

    # ELBO: maximize data fit, minimize KL to prior
    # L = sign * (fisher_lse - prior_weight * kl)
    total = fisher_lse - prior_weight * kl

    return sign * total