Feature Selection Prune

prune

CLASS	DESCRIPTION
PrunedLinear	Linear layer with integrated Lasso-based feature selection and pruning.

FUNCTION	DESCRIPTION
prune_model	Apply structured pruning to an existing model.

Classes

PrunedLinear(in_features: int, out_features: int, mask: torch.Tensor | None = None, n_cv: int = 5, alphas: torch.Tensor | None = None, max_iter: int = 1000, tol: float = 0.0001, eps: float = 0.001, selection: Literal['cyclic', 'random'] = 'cyclic', n_jobs: int | None = None, random_state: int | None = None, threshold: float = 1e-05, bias: bool = True)

Bases: Module

Linear layer with integrated Lasso-based feature selection and pruning.

Automatically performs Lasso regression to select important features, applies binary mask to input features, and initializes weights with Lasso coefficients for optimal performance.

Call fit() method to perform automatic Lasso feature selection before training and prediction.

PARAMETER	DESCRIPTION
in_features	Number of input features. TYPE: int
out_features	Number of output features. TYPE: int
mask	Binary mask of shape (in_features,) where True = keep feature, False = prune. If None, all features are initially selected and Lasso fitting is required. TYPE: torch.Tensor or None, optional, by default None DEFAULT: None
n_cv	Number of folds for cross-validation in Lasso. TYPE: int, optional, by default 5 DEFAULT: 5
alphas	Grid of alpha values for cross-validation. If None, automatically generated. TYPE: torch.Tensor or None, optional, by default None DEFAULT: None
max_iter	Maximum number of iterations for coordinate descent. TYPE: int, optional, by default 1000 DEFAULT: 1000
tol	Tolerance for optimization convergence. TYPE: float, optional, by default 1e-4 DEFAULT: 0.0001
eps	Length of regularization path (ratio of min/max alpha). TYPE: float, optional, by default 1e-3 DEFAULT: 0.001
selection	Feature selection strategy for coordinate descent. TYPE: (cyclic, random) DEFAULT: "cyclic"
n_jobs	Number of parallel jobs for cross-validation. TYPE: int or None, optional, by default None DEFAULT: None
random_state	Random seed for reproducibility. TYPE: int or None, optional, by default None DEFAULT: None
threshold	Threshold for considering coefficients as non-zero. TYPE: float, optional, by default 1e-5 DEFAULT: 1e-05
bias	Whether to include bias term. TYPE: bool, optional, by default True DEFAULT: True

ATTRIBUTE	DESCRIPTION
mask	Binary mask of shape (in_features, ) where True = selected features. TYPE: Tensor
coefficients	Lasso coefficients from feature selection. TYPE: Tensor
features	Indices of selected features. TYPE: Tensor
lasso_model	Fitted Lasso model for feature selection. TYPE: Lasso
is_fitted	Whether the Lasso model has been fitted. TYPE: bool

Examples:

Create layer and fit with data:

>>> layer = PrunedLinear(in_features=100, out_features=1)
>>> layer.fit(X_train, y_train)
>>> output = layer(X_test)

METHOD	DESCRIPTION
fit	Fit Lasso model and update feature mask and weights.
forward	Apply pruned linear transformation.
check_fitted	Check if model is fitted, raise error if not.

Source code in spectre/feature_selection/prune.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    mask: torch.Tensor | None = None,
    n_cv: int = 5,
    alphas: torch.Tensor | None = None,
    max_iter: int = 1000,
    tol: float = 1e-4,
    eps: float = 1e-3,
    selection: Literal["cyclic", "random"] = "cyclic",
    n_jobs: int | None = None,
    random_state: int | None = None,
    threshold: float = 1e-5,
    bias: bool = True,
):
    super().__init__()

    self.in_features = in_features
    self.out_features = out_features

    self.linear = torch.nn.Linear(in_features, out_features, bias=bias)

    self.lasso_model = Lasso(
        n_cv=n_cv,
        alphas=alphas,
        max_iter=max_iter,
        tol=tol,
        eps=eps,
        selection=selection,
        n_jobs=n_jobs,
        threshold=threshold,
        random_state=random_state,
    )

    if mask is not None:
        self.register_buffer("mask_buff", mask)
        self.register_buffer("coefficients_buff", torch.ones(in_features))
        self.register_buffer("features_buff", torch.where(mask)[0])

        # Apply pruning because the mask is provided
        weight_mask = mask.unsqueeze(0).expand(out_features, -1).float()
        prune.custom_from_mask(self.linear, name="weight", mask=weight_mask)

        self._is_fitted = True
    else:
        self.register_buffer("mask_buff", torch.ones(in_features, dtype=torch.bool))
        self.register_buffer("coefficients_buff", torch.zeros(in_features))
        self.register_buffer("features_buff", torch.zeros(0, dtype=torch.long))

        self._is_fitted = False

Attributes

features: torch.Tensor property

Get indices of selected features.

coefficients: torch.Tensor property

Get feature importance scores (absolute Lasso coefficients).

mask: torch.Tensor property

Get feature mask.

sparsity_ratio: float property

Get ratio of selected features to total features.

is_fitted: bool property

Whether the Lasso model has been fitted.

Functions

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

Fit Lasso model and update feature mask and weights.

PARAMETER	DESCRIPTION
X	Training features matrix. TYPE: torch.Tensor of shape (n_samples, n_features)
weights	Sample weights for training. TYPE: torch.Tensor of shape (n_samples,) or None DEFAULT: None
target	Target values. TYPE: torch.Tensor of shape (n_samples,) or (n_samples, n_targets) DEFAULT: None

RETURNS	DESCRIPTION
PrunedLinear	Self for method chaining.

Source code in spectre/feature_selection/prune.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> "PrunedLinear":
    """
    Fit Lasso model and update feature mask and weights.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, n_features)
        Training features matrix.

    weights : torch.Tensor of shape (n_samples,) or None, optional
        Sample weights for training.

    target : torch.Tensor of shape (n_samples,) or (n_samples, n_targets)
        Target values.

    Returns
    -------
    PrunedLinear
        Self for method chaining.
    """
    check_2d(X)

    if X.shape[1] != self.in_features:
        raise ValueError(
            f"Expected X to have {self.in_features} features (in_features), "
            f"got {X.shape[1]}"
        )

    if target is None:
        raise ValueError("Target values must be provided for Lasso regression.")

    check_same_len(X, target)

    # Check target dimensionality matches out_features
    if target.ndim == 1:
        if self.out_features != 1:
            raise ValueError(
                f"Expected target to have {self.out_features} outputs (out_features), "
                f"got 1-dimensional target. Use target.unsqueeze(1) or reshape target "
                f"to shape (n_samples, {self.out_features})"
            )
    elif target.ndim == 2:
        if target.shape[1] != self.out_features:
            raise ValueError(
                f"Expected target to have {self.out_features} outputs (out_features), "
                f"got {target.shape[1]}"
            )
    else:
        raise ValueError(
            f"Expected target to be 1D or 2D, got shape {target.shape}"
        )

    if weights is not None:
        check_1d(weights)
        check_same_len(X, weights)

    # Perform Lasso regression
    self.lasso_model.fit(X, weights=weights, target=target)

    # Extract results from fitted model
    self.coefficients_buff = self.lasso_model.coefficients.to(X.device)
    self.features_buff = self.lasso_model.features.to(X.device)

    # Create mask for pruning
    self.mask_buff = torch.zeros(X.shape[1], dtype=torch.bool, device=X.device)
    self.mask_buff[self.features_buff] = True

    # Initialize weights with Lasso coefficients
    out_features, _ = self.linear.weight.shape
    init_weights = torch.zeros_like(self.linear.weight)

    if self.coefficients_buff.ndim == 1:
        init_weights[0, self.features_buff] = self.coefficients_buff[
            self.features_buff
        ]
    else:
        for i in range(min(out_features, self.coefficients_buff.shape[0])):
            init_weights[i, self.features_buff] = self.coefficients_buff[
                i, self.features_buff
            ]

    self.linear.weight.data = init_weights

    # Apply PyTorch pruning to the linear layer
    weight_mask = self.mask_buff.unsqueeze(0).expand(out_features, -1).float()
    prune.custom_from_mask(self.linear, name="weight", mask=weight_mask)

    self._is_fitted = True

    return self

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

Apply pruned linear transformation.

Uses PyTorch's pruning mechanism to zero out non-selected features. Must call :meth:fit() before using.

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

RETURNS	DESCRIPTION
Tensor	Output after applying pruned linear layer.

Source code in spectre/feature_selection/prune.py

def forward(self, X: torch.Tensor) -> torch.Tensor:
    """
    Apply pruned linear transformation.

    Uses PyTorch's pruning mechanism to zero out non-selected features.
    Must call :meth:`fit()` before using.

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

    Returns
    -------
    torch.Tensor
        Output after applying pruned linear layer.
    """
    self.check_fitted()
    return self.linear(X)

check_fitted()

Check if model is fitted, raise error if not.

Source code in spectre/feature_selection/prune.py

def check_fitted(self):
    """Check if model is fitted, raise error if not."""
    if not self._is_fitted:
        raise RuntimeError("Model must be fitted before calling this method.")

Functions

prune_model(model: torch.nn.Module, mask: torch.Tensor, layer_name: str = 'weight', make_permanent: bool = False) -> torch.nn.Module

Apply structured pruning to an existing model.

Uses PyTorch's pruning utilities to mask input features based on input selection.

PARAMETER	DESCRIPTION
model	Neural network to prune. TYPE: Module
mask	Binary mask where True = keep feature, False = prune feature. TYPE: Tensor
layer_name	Parameter name to prune. TYPE: str, optional, by default "weight" DEFAULT: 'weight'
make_permanent	Whether to permanently remove pruned weights. TYPE: bool, optional, by default False DEFAULT: False

RETURNS	DESCRIPTION
Module	Pruned model with masked features.

Examples:

Lasso feature selection:

>>> from spectre.feature_selection import lasso, prune_model
>>> results = lasso(X, target=target)
>>> mask = torch.zeros(X.shape[1], dtype=torch.bool)
>>> mask[results["features"]] = True

Apply to existing model:

>>> model = torch.nn.Linear(X.shape[1], 1)
>>> pruned_model = prune_model(model, mask)

Source code in spectre/feature_selection/prune.py

def prune_model(
    model: torch.nn.Module,
    mask: torch.Tensor,
    layer_name: str = "weight",
    make_permanent: bool = False,
) -> torch.nn.Module:
    """
    Apply structured pruning to an existing model.

    Uses PyTorch's pruning utilities to mask input features based on
    input selection.

    Parameters
    ----------
    model : torch.nn.Module
        Neural network to prune.

    mask : torch.Tensor
        Binary mask where True = keep feature, False = prune feature.

    layer_name : str, optional, by default "weight"
        Parameter name to prune.

    make_permanent : bool, optional, by default False
        Whether to permanently remove pruned weights.

    Returns
    -------
    torch.nn.Module
        Pruned model with masked features.

    Examples
    --------
    Lasso feature selection:

    >>> from spectre.feature_selection import lasso, prune_model
    >>> results = lasso(X, target=target)
    >>> mask = torch.zeros(X.shape[1], dtype=torch.bool)
    >>> mask[results["features"]] = True

    Apply to existing model:

    >>> model = torch.nn.Linear(X.shape[1], 1)
    >>> pruned_model = prune_model(model, mask)
    """
    for module in model.modules():
        if isinstance(module, torch.nn.Linear):
            if module.in_features == len(mask):
                # Create weight mask [out_features, in_features].
                weight_mask = mask.unsqueeze(0).expand(module.out_features, -1)

                prune.custom_from_mask(module, name=layer_name, mask=weight_mask)

                if make_permanent:
                    prune.remove(module, layer_name)

                break

    return model