Loss Composite

composite

CLASS	DESCRIPTION
CompositeLoss	Combine multiple loss functions with configurable weights.

Classes

CompositeLoss(losses: list[Loss] | dict[str, Loss], weights: list[float] | dict[str, float] | None = None, context_fn: Callable[[dict[str, Any]], dict[str, Any]] | None = None, context_maps: dict[str, Callable[[dict[str, Any]], dict[str, Any]]] | None = None)

Bases: Loss

Combine multiple loss functions with configurable weights.

Enables composition of multiple loss functions into a single weighted sum, with support for named losses and individual component tracking. Useful for multi-objective optimization in parametric methods.

PARAMETER	DESCRIPTION
losses	Loss functions to combine. Can be a list for unnamed losses with equal or sequential weights, or a dict for named losses with custom weights. TYPE: list[Loss] | dict[str, Loss]
weights	Weights for each loss function. If None, all weights default to 1.0. For list losses: provide list of floats (same length as losses) For dict losses: provide dict with matching keys If None: all losses weighted equally at 1.0 TYPE: list[float] | dict[str, float] | None, by default None DEFAULT: None
context_fn	Optional function to compute shared context for all losses. Called once per forward pass to avoid redundant computation. Receives the merged context dict (context + kwargs from forward()). Should return dict of values to merge with the provided context. Signature: context_fn(context: dict[str, Any]) -> dict[str, Any] TYPE: callable or None, by default None DEFAULT: None
context_maps	Optional per-loss context mapping functions for resolving key collisions. Maps loss names to functions that transform the merged context into loss-specific context. Useful when multiple losses need different tensors under the same key name (e.g., both need 'K' but for different purposes). Keys must match loss names. Functions receive merged context and return dict with keys expected by that specific loss. Signature: context_maps[name](context: dict[str, Any]) -> dict[str, Any] TYPE: dict[str, callable] or None, by default None DEFAULT: None

ATTRIBUTE	DESCRIPTION
losses	Internal storage of loss functions as name-to-loss mapping. TYPE: dict[str, Loss]
weights	Weights for each loss function. TYPE: dict[str, float]
context_fn	Context computation function. TYPE: callable or None
loss_components	Most recently computed individual loss components. Populated after each forward pass for logging purposes. Initialized as empty dict. TYPE: dict[str, Tensor]

Examples:

Combine losses with equal weights:

>>> from spectre.loss import CompositeLoss, EigenvalueLoss, OrthogonalLoss
>>> composite = CompositeLoss(
...     losses=[EigenvalueLoss(reduce="gap", n_eigen=2), OrthogonalLoss()],
... )

Combine losses with custom weights:

>>> composite = CompositeLoss(
...     losses=[EigenvalueLoss(reduce="gap", n_eigen=2), OrthogonalLoss()],
...     weights=[1.0, 0.1],
... )

Combine named losses with custom weights:

>>> composite = CompositeLoss(
...     losses={
...         "spectral": EigenvalueLoss(reduce="gap", n_eigen=2),
...         "orthogonal": OrthogonalLoss(),
...     },
...     weights={"spectral": 1.0, "orthogonal": 0.1},
... )
>>>
>>> # Access individual components after forward pass
>>> context = {"eigval": eigvals, "Z": latent_space}
>>> total_loss = composite(context)
>>> print(composite.loss_components)
{'spectral': tensor(0.5), 'orthogonal': tensor(0.02)}

Use with context function for shared computation:

>>> def compute_context(context, **kwargs):
...     X = context["X"]
...     z = encoder(X)
...     pdist = pairwise_fn(z)
...     kernel = kernel_fn(pdist)
...     eigvals = torch.linalg.eigvalsh(kernel)
...     return {"eigval": eigvals, "Z": z, "kernel": kernel}
>>>
>>> composite = CompositeLoss(
...     losses={"spectral": EigenvalueLoss(...), "ortho": OrthogonalLoss()},
...     weights={"spectral": 1.0, "ortho": 0.1},
...     context_fn=compute_context,
... )

Resolve key collisions with context_prefix:

>>> composite = CompositeLoss(
...     losses={
...         "kl": KLDivergenceLoss(context_prefix="kl"),
...         "distill": DistillationLoss(context_prefix="distill"),
...     },
...     weights={"kl": 1.0, "distill": 0.5},
... )
>>> context = {
...     "kl_P": P_distribution,
...     "kl_Q": Q_distribution,
...     "distill_T": teacher_output,
...     "distill_S": student_output,
... }
>>> loss = composite(context)

Resolve key collisions with context_maps:

>>> composite = CompositeLoss(
...     losses={
...         "kl": KLDivergenceLoss(),
...         "distill": DistillationLoss(),
...     },
...     context_maps={
...         "kl": lambda ctx: {"P": ctx["P_prob"], "Q": ctx["Q_prob"]},
...         "distill": lambda ctx: {"T": ctx["teacher"], "S": ctx["student"]},
...     },
... )

METHOD	DESCRIPTION
forward	Compute weighted sum of all loss functions.
exists	Check if composite loss contains a loss component matching the pattern.

Source code in spectre/loss/composite.py

def __init__(
    self,
    losses: list[Loss] | dict[str, Loss],
    weights: list[float] | dict[str, float] | None = None,
    context_fn: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    context_maps: dict[str, Callable[[dict[str, Any]], dict[str, Any]]]
    | None = None,
) -> None:
    # Skip Loss.__init__ validation since CompositeLoss does not use
    # Loss's parameters: reduce and sign
    torch.nn.Module.__init__(self)

    # Convert to internal dict representation
    if isinstance(losses, dict):
        self.losses = losses
        if weights is None:
            self.weights = {name: 1.0 for name in losses.keys()}
        elif isinstance(weights, dict):
            # Validate weight keys match loss keys
            if set(weights.keys()) != set(losses.keys()):
                raise ValueError(
                    f"Weight keys {set(weights.keys())} must match loss keys "
                    f"{set(losses.keys())}."
                )
            self.weights = weights
        else:
            raise TypeError(
                f"For dict losses, weights must be dict or None, got {type(weights)}."
            )

    elif isinstance(losses, list):
        # Convert list to dict with auto-generated names
        self.losses = {}
        self.weights = {}

        if weights is None:
            weights = [1.0] * len(losses)
        elif isinstance(weights, list):
            if len(weights) != len(losses):
                raise ValueError(
                    f"Number of weights ({len(weights)}) must match number of "
                    f"losses ({len(losses)})."
                )
        else:
            raise TypeError(
                f"For list losses, weights must be list or None, got {type(weights)}."
            )

        # Generate unique names for duplicate loss types
        for i, loss_fn in enumerate(losses):
            name = _make_unique_name(loss_fn, self.losses)
            self.losses[name] = loss_fn
            self.weights[name] = weights[i]
    else:
        raise TypeError(f"losses must be list or dict, got {type(losses)}.")

    if context_fn is not None and not callable(context_fn):
        raise TypeError(f"`context_fn` must be callable, got {type(context_fn)}.")
    self.context_fn = context_fn

    if context_maps is not None:
        if not isinstance(context_maps, dict):
            raise TypeError(
                f"`context_maps` must be dict or None, got {type(context_maps)}."
            )
        # Validate all values are callable
        for name, map_fn in context_maps.items():
            if not callable(map_fn):
                raise TypeError(
                    f"`context_maps['{name}']` must be callable, got {type(map_fn)}."
                )
        # Validate all keys exist in losses
        invalid_keys = set(context_maps.keys()) - set(self.losses.keys())
        if invalid_keys:
            raise ValueError(
                f"`context_maps` contains keys not in losses: {invalid_keys}. "
                f"Available loss names: {set(self.losses.keys())}."
            )
    self.context_maps = context_maps or {}

    self.loss_components: dict[str, torch.Tensor] = {}

Attributes

loss_types: list[LossType] property

Aggregate loss types from all constituent losses.

RETURNS	DESCRIPTION
list[LossType]	List of unique loss types across all child losses.

Functions

forward(context: dict[str, torch.Tensor] | None = None, **kwargs) -> torch.Tensor

Compute weighted sum of all loss functions.

Accepts either a context dictionary or kwargs (or both). All arguments are merged and passed to child losses.

PARAMETER	DESCRIPTION
context	Dictionary containing tensors required for child losses. If None, uses empty dict. TYPE: dict[str, torch.Tensor] or None, by default None DEFAULT: None
**kwargs	Additional keyword arguments merged into context for child losses. If context_fn is provided, the merged context is passed to it. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Scalar total loss as weighted sum of individual components.

Notes

Individual loss components are stored in self.loss_components.

Source code in spectre/loss/composite.py

def forward(
    self, context: dict[str, torch.Tensor] | None = None, **kwargs
) -> torch.Tensor:
    """
    Compute weighted sum of all loss functions.

    Accepts either a context dictionary or kwargs (or both). All arguments
    are merged and passed to child losses.

    Parameters
    ----------
    context : dict[str, torch.Tensor] or None, by default None
        Dictionary containing tensors required for child losses.
        If None, uses empty dict.

    **kwargs
        Additional keyword arguments merged into context for child losses.
        If `context_fn` is provided, the merged context is passed to it.

    Returns
    -------
    torch.Tensor
        Scalar total loss as weighted sum of individual components.

    Notes
    -----
    Individual loss components are stored in `self.loss_components`.
    """
    if context is None:
        context = {}

    # Merge kwargs into context for child losses
    # Child losses expect all arguments in the context dict
    merged_context = {**context, **kwargs}

    # Compute shared context if provided
    if self.context_fn is not None:
        computed_context = self.context_fn(merged_context)
        # Merge computed context with existing context
        merged_context = {**merged_context, **computed_context}

    total_loss = None
    self.loss_components = {}

    for name, loss_fn in self.losses.items():
        # Apply context mapping if provided for this loss
        if name in self.context_maps:
            loss_context = self.context_maps[name](merged_context)
        else:
            loss_context = merged_context

        component = self._call_loss_fn(loss_fn, loss_context)
        self.loss_components[name] = component
        weighted_component = self.weights[name] * component
        if total_loss is None:
            total_loss = weighted_component.clone()
        else:
            total_loss = total_loss + weighted_component

    # Handle edge case of empty losses list
    if total_loss is None:
        raise ValueError("CompositeLoss requires at least one loss function.")

    return total_loss

exists(pattern: str) -> bool

Check if composite loss contains a loss component matching the pattern.

PARAMETER	DESCRIPTION
pattern	Pattern to match against loss component names. Supports prefix matching (e.g., "mincut" matches "loss_mincut", "loss_mincut_1", etc.). TYPE: str

RETURNS	DESCRIPTION
bool	True if any loss component name contains the pattern, False otherwise.

Source code in spectre/loss/composite.py

def exists(self, pattern: str) -> bool:
    """
    Check if composite loss contains a loss component matching the pattern.

    Parameters
    ----------
    pattern : str
        Pattern to match against loss component names. Supports prefix matching
        (e.g., "mincut" matches "loss_mincut", "loss_mincut_1", etc.).

    Returns
    -------
    bool
        True if any loss component name contains the pattern, False otherwise.
    """
    # Check composite keys
    if any(pattern in name for name in self.losses.keys()):
        return True
    # Delegate to constituent losses
    return any(loss_fn.exists(pattern) for loss_fn in self.losses.values())