Parametric Regression

regression

CLASS	DESCRIPTION
RegressionProbabilistic	Probabilistic regression model with uncertainty estimation.

Classes

RegressionProbabilistic(*, model: ModelType | dict[str, ModelType], loss_kwargs: dict | None = None, optimizer_fn: type[torch.optim.Optimizer] | str | None = None, optimizer_kwargs: dict | None = None, lr_scheduler_fn: Any | str | None = None, lr_scheduler_kwargs: dict | None = None, device: Literal['cuda', 'cpu'] = 'cuda')

Bases: Parametric

Probabilistic regression model with uncertainty estimation.

This class implements a probabilistic regression approach that predicts both the mean and variance (scale) of target values. The model outputs parameters of a specified probability distribution, enabling uncertainty quantification in predictions.

The architecture consists of:

• A feature extraction network (the main model).
• A mean prediction layer (linear layer -> mean).
• A scale prediction layer (linear layer -> softplus -> positive scale).

The model optimizes the negative log-likelihood of the specified distribution, learning to predict both the expected value and uncertainty for each prediction.

The loss function used is the negative log-likelihood

\(L_{NLL} = -\frac{1}{N} \sum_{i=1}^{N} \log p(y_i | \mu_i, \sigma_i)\),

where \(p(y | \mu, \sigma)\) is the probability density function of the chosen distribution (e.g., Gaussian), \(\mu\) is the predicted mean, and \(\sigma\) is the predicted scale (standard deviation).

PARAMETER	DESCRIPTION
model	Neural network model for feature extraction. The model output features are used as input to the mean and scale prediction layers. TYPE: ModelType | dict[str, ModelType]
loss_kwargs	Additional keyword arguments to pass to the LogLikelihoodLoss constructor. "distribution" : Literal["normal", "laplace"], by default "normal" Choice of probability distribution for modeling targets. "reduce" : Literal["sum", "mean"], by default "mean" Reduction method for loss computation. TYPE: dict, optional, by default None DEFAULT: None
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) TYPE: type[torch.optim.Optimizer] | str | None, optional, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation (e.g., {"lr": 1e-3, "weight_decay": 1e-4}). TYPE: dict | None, optional, by default None DEFAULT: None
lr_scheduler_fn	Learning rate scheduler specification. None: No scheduler (default) str: Scheduler name from torch.optim.lr_scheduler (e.g., "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, optional, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. TYPE: dict | None, optional, by default None DEFAULT: None
device	Computation device. Automatically falls back to "cpu" if CUDA unavailable. TYPE: Literal['cuda', 'cpu'] DEFAULT: "cuda"

Notes

• The model automatically adds a torch.nn.Sigmoid output activation if none exists
• Mean and scale layers are added on top of the feature extraction model
• Scale predictions are constrained to be positive using torch.nn.Softplus activation
• Sign for the loss is set to -1.0 to convert log-likelihood maximization. Even if provided otherwise in loss_kwargs, it will be overridden to -1.0

METHOD	DESCRIPTION
forward_step	Forward pass through the feature extraction model.
loss	Compute probabilistic regression loss.
score_step	Compute negative log-likelihood as the scoring metric.
training_step	Training step for PyTorch Lightning with probabilistic loss.
predict_step	Prediction step returning distribution parameters.

Source code in spectre/parametric/regression.py

def __init__(
    self,
    *,
    model: ModelType | dict[str, ModelType],
    loss_kwargs: dict | None = None,
    optimizer_fn: type[torch.optim.Optimizer] | str | None = None,
    optimizer_kwargs: dict | None = None,
    lr_scheduler_fn: Any | str | None = None,
    lr_scheduler_kwargs: dict | None = None,
    device: Literal["cuda", "cpu"] = "cuda",
) -> None:
    super().__init__(
        model=model,
        device=device,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

    if self.model.has_output_fn is False:
        self.model.sequential.add_module("act_output", torch.nn.Sigmoid())

    if self.model.out_features < 2:
        warnings.warn(
            "`RegressionProbabilistic` has two additional layers: mean and scale "
            "of shape [model.out_features, 1]. Consider increasing the output size "
            "of the model."
        )

    self.mean_layer = torch.nn.Linear(self.model.out_features, 1)
    self.scale_layer = torch.nn.Sequential(
        torch.nn.Linear(self.model.out_features, 1),
        torch.nn.Softplus(),
    )

    self.mean_layer = self.mean_layer.to(device=device)
    self.scale_layer = self.scale_layer.to(device=device)

    if loss_kwargs is None:
        loss_kwargs = {}
    if not isinstance(loss_kwargs, dict):
        raise TypeError(
            "Additional loss parameters `loss_kwargs` must be a dictionary."
        )
    # Override provided sign, taking 1.0 does not make sense for likelihood.
    loss_kwargs["sign"] = -1.0

    self.loss_fn = LogLikelihoodLoss(**loss_kwargs)

Functions

forward_step(X: torch.Tensor) -> torch.Tensor

Forward pass through the feature extraction model.

PARAMETER	DESCRIPTION
X	Input features. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, out_features)	Extracted features for mean and scale prediction.

Source code in spectre/parametric/regression.py

def forward_step(self, X: torch.Tensor) -> torch.Tensor:
    """
    Forward pass through the feature extraction model.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input features.

    Returns
    -------
    torch.Tensor of shape (n_samples, out_features)
        Extracted features for mean and scale prediction.
    """
    X = self.model(X)
    return X

loss(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute probabilistic regression loss.

Computes the negative log-likelihood of the target values under the predicted distribution parameters (mean and scale).

PARAMETER	DESCRIPTION
X	Input features. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. TYPE: torch.Tensor of shape (n_samples,) or None DEFAULT: None
target	Target values. TYPE: torch.Tensor of shape (n_samples, 1) | None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar loss value (negative log-likelihood).

Source code in spectre/parametric/regression.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute probabilistic regression loss.

    Computes the negative log-likelihood of the target values under the
    predicted distribution parameters (mean and scale).

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input features.

    weights : torch.Tensor of shape (n_samples,) or None, default None
        Sample weights for weighted loss computation.

    target : torch.Tensor of shape (n_samples, 1) | None, default None
        Target values.

    Returns
    -------
    torch.Tensor
        Scalar loss value (negative log-likelihood).
    """
    X = self.forward_step(X)

    mean = self.mean_layer(X)
    scale = self.scale_layer(X)

    context = {"X": X, "Y": target, "mean": mean, "scale": scale}
    if weights is not None:
        context["weights"] = weights

    loss = self.loss_fn(context)

    return loss

score_step(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute negative log-likelihood as the scoring metric.

For probabilistic regression, the natural scoring metric is the negative log-likelihood, which measures how well the predicted distribution fits the observed data.

PARAMETER	DESCRIPTION
X	Input features. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted scoring. TYPE: torch.Tensor of shape (n_samples,) or None DEFAULT: None
target	Target values for scoring. TYPE: torch.Tensor of shape (n_samples, 1) or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Negative log-likelihood score (lower is better).

Source code in spectre/parametric/regression.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute negative log-likelihood as the scoring metric.

    For probabilistic regression, the natural scoring metric is the
    negative log-likelihood, which measures how well the predicted
    distribution fits the observed data.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input features.

    weights : torch.Tensor of shape (n_samples,) or None, default None
        Sample weights for weighted scoring.

    target : torch.Tensor of shape (n_samples, 1) or None, default None
        Target values for scoring.

    Returns
    -------
    torch.Tensor
        Negative log-likelihood score (lower is better).
    """
    if target is None:
        raise ValueError("Target values are required for scoring.")

    score = self.loss(X, weights=weights, target=target)
    return score

training_step(batch: DataBatch, batch_idx: int) -> torch.Tensor

Training step for PyTorch Lightning with probabilistic loss.

Computes the negative log-likelihood loss and logs it for monitoring during training and validation.

PARAMETER	DESCRIPTION
batch	Input batch containing (X, weights, target) where: X : torch.Tensor of shape (n_samples, in_features) weights : torch.Tensor of shape (n_samples,) target : torch.Tensor of shape (n_samples, 1) TYPE: Batch
batch_idx	Batch index. TYPE: int

RETURNS	DESCRIPTION
Tensor	Scalar training loss (negative log-likelihood).

Source code in spectre/parametric/regression.py

def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Training step for PyTorch Lightning with probabilistic loss.

    Computes the negative log-likelihood loss and logs it for monitoring
    during training and validation.

    Parameters
    ----------
    batch : Batch
        Input batch containing (X, weights, target) where:

        - X : torch.Tensor of shape (n_samples, in_features)
        - weights : torch.Tensor of shape (n_samples,)
        - target : torch.Tensor of shape (n_samples, 1)

    batch_idx : int
        Batch index.

    Returns
    -------
    torch.Tensor
        Scalar training loss (negative log-likelihood).
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss)

    return loss

predict_step(batch: DataBatch, batch_idx: int, dataloader_idx: int | None = None) -> torch.Tensor

Prediction step returning distribution parameters.

Predicts both the mean and scale (standard deviation) of the target distribution for uncertainty quantification.

PARAMETER	DESCRIPTION
batch	Input batch containing input features. TYPE: Batch
batch_idx	Batch index (required by PyTorch Lightning but unused). TYPE: int
dataloader_idx	Dataloader index (required by PyTorch Lightning but unused). TYPE: int or None DEFAULT: None

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, 2)	Concatenated predictions where: Column 0: Predicted mean values Column 1: Predicted scale (std dev) values

Source code in spectre/parametric/regression.py

def predict_step(
    self,
    batch: DataBatch,
    batch_idx: int,
    dataloader_idx: int | None = None,
) -> torch.Tensor:
    """
    Prediction step returning distribution parameters.

    Predicts both the mean and scale (standard deviation) of the target
    distribution for uncertainty quantification.

    Parameters
    ----------
    batch : Batch
        Input batch containing input features.

    batch_idx : int
        Batch index (required by PyTorch Lightning but unused).

    dataloader_idx : int or None, default None
        Dataloader index (required by PyTorch Lightning but unused).

    Returns
    -------
    torch.Tensor of shape (n_samples, 2)

        Concatenated predictions where:

        - Column 0: Predicted mean values
        - Column 1: Predicted scale (std dev) values
    """
    X = self.model(batch.data)
    mean = self.mean_layer(X)
    scale = self.scale_layer(X)

    return torch.concat([mean, scale], dim=1)