Implement Lasso Regression using ISTA

#50 · Machine Learning · Medium

Problem

Implement Lasso Regression using the Iterative Shrinkage-Thresholding Algorithm (ISTA). Lasso adds an L1 penalty to the loss, which promotes sparsity in the learned weights.

Solution

import numpy as np

def soft_threshold(x, threshold):
    return np.sign(x) * np.maximum(np.abs(x) - threshold, 0)

def lasso_ista(X, y, alpha=0.1, lr=0.01, n_iterations=100):
    X = np.array(X, dtype=np.float64)
    y = np.array(y, dtype=np.float64).reshape(-1)
    n_samples, n_features = X.shape
    weights = np.zeros(n_features)

    for _ in range(n_iterations):
        predictions = X @ weights
        residual = predictions - y
        gradient = (1 / n_samples) * (X.T @ residual)

        # Gradient step
        weights_temp = weights - lr * gradient

        # Proximal (soft-thresholding) step
        weights = soft_threshold(weights_temp, alpha * lr)

    return weights.tolist()

Explanation

Initialize weights to zero.
At each iteration, compute the gradient of the MSE loss.
Take a gradient step to get a temporary weight vector.
Apply the soft-thresholding (proximal) operator, which shrinks values toward zero and sets small values exactly to zero.
The threshold is alpha * lr, where alpha controls the L1 regularization strength.

Complexity

Time: O(n_iterations n d) where n is samples and d is features
Space: O(d) for the weight vector

← #49 #51 →