"""gelcd.py: coordinate descent for group elastic net.
This implementation uses block-wise coordinate descent to solve the optimization
problem:
min_b (1/2m)||y - b_0 - sum_j A_j@b_j||^2 +
sum_j {sqrt{n_j}(l_1*||b_j|| + l_2*||b_j||^2)}.
The algorithm repeatedly minimizes with respect to b_0 and the b_js. For b_0,
the minimization has a closed form solution. For each b_j, the minimization
objective is convex, twice differentiable; and can be solved using gradient
descent, Newton's method etc. This implementation allows the internal minimizer
to be chosen. These are the block_solve_* functions. Each solves the above
optimization problem with respect to one of the b_js, while keeping the others
fixed. The gradient with respect to b_j is
(-1/m)A_j'@(y - b_0 - sum_j A_j@b_j) + sqrt{n_j}(l_1*b_j/||b_j|| +
2*l_2*b_j).
and the Hessian is
(1/m)A_j'@A_j + sqrt{n_j}(l_1*(I/||b_j|| - b_j@b_j'/||b_j||^3) + 2*l_2*I).
The coefficients are represented using a scalar bias b_0 and a matrix B where
each row corresponds to a single group (with appropriate 0 padding). The root of
the group sizes are stored in a vector sns. The features are stored in a list of
tensors each of size m x n_j where m is the number of samples, and n_j is the
number of features in group j. For any j, r_j = y - b_0 - sum_{k =/= j} A_k@b_k,
q_j = r_j - A_j@b_j, C_j = A_j'@A_j / m, and I_j is an identity matrix of size
n_j. Finally, a_1 = l_1*sns and a_2 = 2*l_2*sns.
"""
import torch
import tqdm
def _f_j(q_j, b_j_norm, a_1_j, a_2_j, m):
"""Compute the objective with respect to one of the coefficients i.e.
(1/2m)||q_j||^2 + a_1||b_j|| + (a_2/2)||b_j||^2.
"""
return (
((q_j @ q_j) / (2.0 * m))
+ (a_1_j * b_j_norm)
+ ((a_2_j / 2.0) * (b_j_norm ** 2))
)
def _grad_j(q_j, A_j, b_j, b_j_norm, a_1_j, a_2_j, m):
"""Compute the gradient with respect to one of the coefficients."""
return (A_j.t() @ q_j / (-m)) + (b_j * (a_1_j / b_j_norm + a_2_j))
def _hess_j(C_j, I_j, b_j, b_j_norm, a_1_j, a_2_j):
"""Compute the Hessian with respect to one of the coefficients."""
D_j = torch.ger(b_j, b_j)
return C_j + (a_1_j / b_j_norm) * (I_j - D_j / (b_j_norm ** 2)) + a_2_j * I_j
def block_solve_agd(
r_j,
A_j,
a_1_j,
a_2_j,
m,
b_j_init,
t_init=None,
ls_beta=None,
max_iters=None,
rel_tol=1e-6,
verbose=False,
zero_thresh=1e-6,
zero_fill=1e-3,
):
"""Solve the optimization problem for a single block with accelerated
gradient descent."""
b_j = b_j_init
b_j_prev = b_j
k = 1 # iteration number
t = 1 # initial step length (used if t_init is None)
pbar_stats = {} # stats for the progress bar
pbar = tqdm.tqdm(desc="Solving block with AGD", disable=not verbose, leave=False)
while True:
# Compute the v terms.
mom = (k - 2) / (k + 1.0)
v_j = b_j + mom * (b_j - b_j_prev)
q_v_j = r_j - A_j @ v_j
v_j_norm = v_j.norm(p=2)
f_v_j = _f_j(q_v_j, v_j_norm, a_1_j, a_2_j, m)
grad_v_j = _grad_j(q_v_j, A_j, v_j, v_j_norm, a_1_j, a_2_j, m)
b_j_prev = b_j
# Adjust the step size with backtracking line search.
if t_init is not None:
t = t_init
while True:
b_j = v_j - t * grad_v_j # gradient descent update
if ls_beta is None:
# Don't perform line search.
break
# Line search: exit when f_j(b_j) <= f_j(v_j) + grad_v_j'@(b_j -
# v_j) + (1/2t)||b_j - v_j||^2.
q_b_j = r_j - A_j @ b_j
b_j_norm = b_j.norm(p=2)
f_b_j = _f_j(q_b_j, b_j_norm, a_1_j, a_2_j, m)
b_v_diff = b_j - v_j
c2 = grad_v_j @ b_v_diff
c3 = b_v_diff @ b_v_diff / 2.0
if t * f_b_j <= t * (f_v_j + c2) + c3:
break
else:
t *= ls_beta
# Make b_j non-zero if it is 0.
if len((b_j.abs() < zero_thresh).nonzero()) == len(b_j):
b_j.fill_(zero_fill)
b_j_norm = b_j.norm(p=2)
b_diff_norm = (b_j - b_j_prev).norm(p=2)
pbar_stats["t"] = "{:.2g}".format(t)
pbar_stats["rel change"] = "{:.2g}".format(b_diff_norm / b_j_norm)
pbar.set_postfix(pbar_stats)
pbar.update()
# Check max iterations exit criterion.
if max_iters is not None and k == max_iters:
break
k += 1
# Check tolerance exit criterion. Exit when the relative change is less
# than the tolerance.
# k > 2 ensures that at least 2 iterations are completed.
if b_diff_norm <= rel_tol * b_j_norm and k > 2:
break
pbar.close()
return b_j
def block_solve_newton(
r_j,
A_j,
a_1_j,
a_2_j,
m,
b_j_init,
C_j,
I_j,
ls_alpha=0.5,
ls_beta=0.9,
max_iters=20,
tol=1e-8,
verbose=False,
):
"""Solve the optimization problem for a single block with Newton's
method."""
b_j = b_j_init
k = 1
pbar_stats = {} # stats for the progress bar
pbar = tqdm.tqdm(
desc="Solving block with Newton's method", disable=not verbose, leave=False
)
while True:
# First, compute the Newton step and decrement.
q_b_j = r_j - A_j @ b_j
b_j_norm = b_j.norm(p=2)
grad_b_j = _grad_j(q_b_j, A_j, b_j, b_j_norm, a_1_j, a_2_j, m)
hess_b_j = _hess_j(C_j, I_j, b_j, b_j_norm, a_1_j, a_2_j)
hessinv_b_j = torch.inverse(hess_b_j)
v_j = hessinv_b_j @ grad_b_j
dec_j = grad_b_j @ (hessinv_b_j @ grad_b_j)
# Check tolerance stopping criterion. Exit if dec_j / 2 is less than the
# tolerance.
if dec_j / 2 <= tol:
break
# Perform backtracking line search.
t = 1
f_b_j = _f_j(q_b_j, b_j_norm, a_1_j, a_2_j, m)
k_j = grad_b_j @ v_j
while True:
# Compute the update and evaluate function at that point.
bp_j = b_j - t * v_j
q_bp_j = r_j - A_j @ bp_j
bp_j_norm = bp_j.norm(p=2)
f_bp_j = _f_j(q_bp_j, bp_j_norm, a_1_j, a_2_j, m)
if f_bp_j <= f_b_j - ls_alpha * t * k_j:
b_j = bp_j
break
else:
t *= ls_beta
# Make b_j non-zero if it is 0.
if all(b_j.abs() < tol):
b_j.fill_(1e-3)
pbar_stats["t"] = "{:.2g}".format(t)
pbar_stats["1/2 newton decrement"] = "{:.2g}".format(dec_j / 2)
pbar.set_postfix(pbar_stats)
pbar.update()
# Check max iterations stopping criterion.
if max_iters is not None and k == max_iters and k > 2:
break
k += 1
pbar.close()
return b_j
def make_A(
As, ns, device=torch.device("cpu"), dtype=torch.float32
): # pylint: disable=unused-argument
"""Move the As to the provided device, convert to the provided dtype, and
return as A."""
As = [A_j.to(device, dtype) for A_j in As]
return As
def gel_solve(
A,
y,
l_1,
l_2,
ns,
b_init=None,
block_solve_fun=block_solve_agd,
block_solve_kwargs=None,
max_cd_iters=None,
rel_tol=1e-6,
Cs=None,
Is=None,
verbose=False,
):
"""Solve a group elastic net problem.
Arguments:
A: list of feature tensors, one per group (size m x n_j).
y: tensor vector of predictions.
l_1: the 2-norm coefficient.
l_2: the squared 2-norm coefficient.
ns: iterable of group sizes.
b_init: tuple (b_0, B) to initialize b.
block_solve_fun: the function to be used for minimizing individual
blocks (should be one of the block_solve_* functions).
block_solve_kwargs: dictionary of arguments to be passed to
block_solve_fun.
max_cd_iters: maximum number of outer CD iterations.
rel_tol: tolerance for exit criterion; after every CD outer iteration,
b is compared to the previous value, and if the relative difference
is less than this value, the loop is terminated.
Cs, Is: lists of C_j and I_j matrices respectively; used if the internal
solver is Newton's method.
verbose: boolean to enable/disable verbosity.
"""
p = len(A)
m = len(y)
device = A[0].device
dtype = A[0].dtype
y = y.to(device, dtype)
if block_solve_kwargs is None:
block_solve_kwargs = dict()
# Create initial values if not specified.
if b_init is None:
b_init = 0.0, torch.zeros(p, max(ns), device=device, dtype=dtype)
if not isinstance(ns, torch.Tensor):
ns = torch.tensor(ns)
sns = ns.to(device, dtype).sqrt()
a_1 = l_1 * sns
ma_1 = m * a_1
a_2 = 2 * l_2 * sns
b_0, B = b_init
b_0_prev, B_prev = b_0, B
k = 1 # iteration number
pbar_stats = {} # stats for the outer progress bar
pbar = tqdm.tqdm(
desc="Solving gel with CD (l_1 {:.2g}, l_2 {:.2g})".format(l_1, l_2),
disable=not verbose,
)
while True:
# First minimize with respect to b_0. This has a closed form solution
# given by b_0 = 1'@(y - sum_j A_j@b_j) / m.
b_0 = (y - sum(A[j] @ B[j, : ns[j]] for j in range(p))).sum() / m
# Now, minimize with respect to each b_j.
for j in tqdm.trange(
p, desc="Solving individual blocks", disable=not verbose, leave=False
):
r_j = y - b_0 - sum(A[k] @ B[k, : ns[k]] for k in range(p) if k != j)
# Check if b_j must be set to 0. The condition is ||A_j'@r_j|| <=
# m*a_1.
if (A[j].t() @ r_j).norm(p=2) <= ma_1[j]:
B[j] = 0
else:
# Otherwise, minimize. First make sure initial value is not 0.
if len((B[j, : ns[j]].abs() < 1e-6).nonzero()) == ns[j]:
B[j, : ns[j]] = 1e-3
# Add C_j and I_j to the arguments if using Newton's method.
if block_solve_fun is block_solve_newton:
block_solve_kwargs["C_j"] = Cs[j]
block_solve_kwargs["I_j"] = Is[j]
B[j, : ns[j]] = block_solve_fun(
r_j,
A[j],
a_1[j].item(),
a_2[j].item(),
m,
B[j, : ns[j]],
verbose=verbose,
**block_solve_kwargs,
)
# Compute relative change in b.
b_0_diff = b_0 - b_0_prev
B_diff = B - B_prev
delta_norm = (b_0_diff ** 2 + (B_diff ** 2).sum()).sqrt()
b_norm = (b_0 ** 2 + (B ** 2).sum()).sqrt()
pbar_stats["rel change"] = "{:.2g}".format(delta_norm.item() / b_norm.item())
pbar.set_postfix(pbar_stats)
pbar.update()
# Check max iterations exit criterion.
if max_cd_iters is not None and k == max_cd_iters:
break
k += 1
# Check tolerance exit criterion.
if delta_norm.item() <= rel_tol * b_norm.item() and k > 2:
break
b_0_prev, B_prev = b_0, B
pbar.close()
return b_0.item(), B
