Tune Ridge With a Hypergradient

Tune the ridge penalty by gradient descent, through the linear solve.

For a penalty lam, ridge regression has a closed form. The weights are the solution of one linear system:

(X_train.T @ X_train + lam * I) @ w = X_train.T @ y_train

The validation loss is mean((X_val @ w - y_val) ** 2). Normally you would try a grid of lam values. But torch.linalg.solve is differentiable, so autograd can give the gradient of the validation loss with respect to lam itself. This is a hypergradient: a gradient for a hyperparameter.

Make log_lam a scalar tensor that requires a gradient, start it at 0 (so lam = 1), and use lam = exp(log_lam) so the penalty stays positive. Run steps steps of torch.optim.Adam with learning rate lr on the validation loss. Return the final lam and the ridge weights for that lam.

The result is graded: the validation loss at your lam must be within 1% of the best loss on a fine grid of penalties.

Function signature

def tune_ridge(X_train, y_train, X_val, y_val, steps, lr):

Example

Input

X = torch.eye(2, dtype=torch.float64)
tune_ridge(X, torch.tensor([2.0, 2.0], dtype=torch.float64),
           X, torch.tensor([0.5, 0.5], dtype=torch.float64), steps=500, lr=0.05)

Output

(3.0, tensor([0.5000, 0.5000], dtype=torch.float64))   # lam is about 3.0

Reasoning

With X = I the weights are w = 2 / (1 + lam). The validation targets are 0.5, and 2 / (1 + lam) = 0.5 at lam = 3. Gradient descent on log_lam moves from 0 to log(3).

Constraints and notes

• All tensors are float64. X_train is (n, d), y_train is (n,), and the same for the validation pair. There is no intercept.
• Return lam as a Python float, and w as a detached tensor of shape (d,).
• Use torch.linalg.solve, not an explicit inverse.