xbtorch.loss.landscapes

Tools for analyzing and visualizing neural network loss landscapes.

This module provides functions to reduce high-dimensional weight spaces to 2D directions, perturb model weights along those directions, and compute the corresponding loss grid for visualization or analysis.

Functions

• reduction_pca() : Reduce weight matrix to 2D using principal component analysis.

• reduction_random() : Reduce weight matrix to 2D using random orthogonal directions.

• compute_grid_loss() : Compute loss values on a 2D grid of weight perturbations.

Functions

compute_grid_loss(model, full_weights, ...)	Compute the loss landscape on a 2D grid along two given directions in weight space.
reduction_pca(matrix[, seed])	Perform PCA on a weight matrix to reduce it to two principal directions.
reduction_random(matrix[, seed])	Generate two random orthogonal directions for dimensionality reduction.

xbtorch.loss.landscapes.compute_grid_loss(model, full_weights, param_name, x_range, y_range, reduced_dirs, data_loader, criterion, device, epoch, resolution=1, fast=True, log=False)

Compute the loss landscape on a 2D grid along two given directions in weight space.

Parameters:

• model (torch.nn.Module) – Neural network model.

• full_weights (dict) – Dictionary of stored weight tensors by layer name and epoch. Format: full_weights[layer_name][epoch].

• param_name (str) – Name of the parameter to perturb in the model.

• x_range (tuple(float, float)) – Minimum and maximum scaling along the first direction.

• y_range (tuple(float, float)) – Minimum and maximum scaling along the second direction.

• reduced_dirs (np.ndarray) – 2 x num_features array of perturbation directions.

• data_loader (torch.utils.data.DataLoader) – Data loader for evaluating loss.

• criterion (torch.nn.Module) – Loss function.

• device (torch.device) – PyTorch device to perform computations on.

• epoch (int) – Epoch of the stored weights to initialize model.

• resolution (int, optional) – Number of steps per unit interval in the grid (default: 1).

• fast (bool, optional) – If True, perform a faster evaluation during training (default: True).

• log (bool, optional) – If True, print progress logs during computation (default: False).

Returns:

• xvnp.ndarray

  2D grid of x-axis values.

• yvnp.ndarray

  2D grid of y-axis values.

• loss_gridnp.ndarray

  Grid of computed loss values at each perturbation point.

Return type:

tuple

xbtorch.loss.landscapes.reduction_pca(matrix, seed=0)

Perform PCA on a weight matrix to reduce it to two principal directions.

Parameters:

• matrix (np.ndarray) – Input weight matrix of shape (num_samples, num_features).

• seed (int, optional) – Random seed for PCA initialization (default: 0).

Returns:

• path_2dnp.ndarray

  Transformed matrix projected onto the first two principal components.

• reduced_dirsnp.ndarray

  The principal component directions (2 x num_features).

• infodict

  Dictionary containing explained variance ratio: {‘explained_vars’: array([var1, var2])}.

Return type:

tuple

xbtorch.loss.landscapes.reduction_random(matrix, seed=None)

Generate two random orthogonal directions for dimensionality reduction.

Parameters:

• matrix (np.ndarray) – Input weight matrix of shape (num_samples, num_features).

• seed (int, optional) – Random seed for reproducibility.

Returns:

• None

  Placeholder for 2D path (not computed).

• reduced_dirsnp.ndarray

  Two random orthonormal directions (2 x num_features).

• infodict

  Empty dictionary (no explained variance available).

Return type:

tuple