ENH add Anderson acceleration to group_bcd_solver (#29)

Author: Badr-MOUFAD

skglm/estimators.py

@@ -940,7 +940,7 @@ def path(self, X, y, alphas, coef_init=None, return_n_iter=True, **params):
             Target vector relative to X.
 
         alphas : array
-            Values of regularization strenghts for which solutions are
+            Values of regularization strengths for which solutions are
             computed.
 
         coef_init : array, shape (n_features,), optional

skglm/solvers/group_bcd_solver.py

@@ -1,7 +1,7 @@
 import numpy as np
 from numba import njit
 
-from skglm.utils import check_group_compatible
+from skglm.utils import AndersonAcceleration, check_group_compatible
 
 
 def bcd_solver(X, y, datafit, penalty, w_init=None, p0=10,
@@ -65,53 +65,54 @@ def bcd_solver(X, y, datafit, penalty, w_init=None, p0=10,
     all_groups = np.arange(n_groups)
     p_objs_out = np.zeros(max_iter)
     stop_crit = 0.  # prevent ref before assign when max_iter == 0
+    accelerator = AndersonAcceleration(K=5)
 
     for t in range(max_iter):
-        if t == 0:  # avoid computing grad and opt twice
-            grad = _construct_grad(X, y, w, Xw, datafit, all_groups)
-            opt = penalty.subdiff_distance(w, grad, all_groups)
-            stop_crit = np.max(opt)
+        grad = _construct_grad(X, y, w, Xw, datafit, all_groups)
+        opt = penalty.subdiff_distance(w, grad, all_groups)
+        stop_crit = np.max(opt)
+
+        if verbose:
+            p_obj = datafit.value(y, w, Xw) + penalty.value(w)
+            print(
+                f"Iteration {t+1}: {p_obj:.10f}, "
+                f"stopping crit: {stop_crit:.2e}"
+            )
 
-            if stop_crit <= tol:
-                break
+        if stop_crit <= tol:
+            break
 
         gsupp_size = penalty.generalized_support(w).sum()
         ws_size = max(min(p0, n_groups),
                       min(n_groups, 2 * gsupp_size))
         ws = np.argpartition(opt, -ws_size)[-ws_size:]  # k-largest items (no sort)
 
         for epoch in range(max_epochs):
+            # inplace update of w and Xw
             _bcd_epoch(X, y, w, Xw, datafit, penalty, ws)
 
+            w_acc, Xw_acc = accelerator.extrapolate(w, Xw)
+            p_obj = datafit.value(y, w, Xw) + penalty.value(w)
+            p_obj_acc = datafit.value(y, w_acc, Xw_acc) + penalty.value(w_acc)
+
+            if p_obj_acc < p_obj:
+                w, Xw = w_acc, Xw_acc
+                p_obj = p_obj_acc
+
+            # check sub-optimality every 10 epochs
             if epoch % 10 == 0:
                 grad_ws = _construct_grad(X, y, w, Xw, datafit, ws)
                 opt_in = penalty.subdiff_distance(w, grad_ws, ws)
                 stop_crit_in = np.max(opt_in)
 
                 if max(verbose - 1, 0):
-                    p_obj = datafit.value(y, w, Xw) + penalty.value(w)
                     print(
                         f"Epoch {epoch+1}: {p_obj:.10f} "
                         f"obj. variation: {stop_crit_in:.2e}"
                     )
 
                 if stop_crit_in <= 0.3 * stop_crit:
                     break
-
-        p_obj = datafit.value(y, w, Xw) + penalty.value(w)
-        grad = _construct_grad(X, y, w, Xw, datafit, all_groups)
-        opt = penalty.subdiff_distance(w, grad, all_groups)
-        stop_crit = np.max(opt)
-
-        if verbose:
-            print(
-                f"Iteration {t+1}: {p_obj:.10f}, "
-                f"stopping crit: {stop_crit:.2e}"
-            )
-
-        if stop_crit <= tol:
-            break
-
         p_objs_out[t] = p_obj
 
     return w, p_objs_out, stop_crit
@@ -137,7 +138,6 @@ def _bcd_epoch(X, y, w, Xw, datafit, penalty, ws):
         for idx, j in enumerate(grp_g_indices):
             if old_w_g[idx] != w[j]:
                 Xw += (w[j] - old_w_g[idx]) * X[:, j]
-    return
 
 
 @njit

skglm/tests/test_group.py

@@ -8,7 +8,7 @@
 from skglm.datafits.group import QuadraticGroup
 from skglm.solvers.group_bcd_solver import bcd_solver
 
-from skglm.utils import grp_converter, make_correlated_data
+from skglm.utils import grp_converter, make_correlated_data, AndersonAcceleration
 from celer import GroupLasso, Lasso
 
 
@@ -60,8 +60,7 @@ def test_alpha_max(n_groups, n_features, shuffle):
         alpha=alpha_max, grp_ptr=grp_ptr,
         grp_indices=grp_indices, weights=weights)
 
-    w = bcd_solver(
-        X, y, quad_group, group_penalty, max_iter=10000, tol=0)[0]
+    w = bcd_solver(X, y, quad_group, group_penalty, tol=1e-12)[0]
 
     np.testing.assert_allclose(norm(w), 0, atol=1e-14)
 
@@ -82,7 +81,7 @@ def test_equivalence_lasso():
         alpha=alpha, grp_ptr=grp_ptr,
         grp_indices=grp_indices, weights=weights)
 
-    w = bcd_solver(X, y, quad_group, group_penalty, max_iter=10000, tol=1e-12)[0]
+    w = bcd_solver(X, y, quad_group, group_penalty, tol=1e-12)[0]
 
     celer_lasso = Lasso(
         alpha=alpha, fit_intercept=False, tol=1e-12, weights=weights).fit(X, y)
@@ -123,5 +122,45 @@ def test_vs_celer_grouplasso(n_groups, n_features, shuffle):
     np.testing.assert_allclose(model.coef_, w, atol=1e-5)
 
 
+def test_anderson_acceleration():
+    # VAR: w = rho * w + 1 with |rho| < 1
+    # converges to w_star = 1 / (1 - rho)
+    max_iter, tol = 1000, 1e-9
+    n_features = 2
+    rho = np.array([0.5, 0.8])
+    w_star = 1 / (1 - rho)
+    X = np.diag([2, 5])
+
+    # with acceleration
+    acc = AndersonAcceleration(K=5)
+    n_iter_acc = 0
+    w = np.ones(n_features)
+    Xw = X @ w
+    for i in range(max_iter):
+        w, Xw = acc.extrapolate(w, Xw)
+        w = rho * w + 1
+        Xw = X @ w
+
+        if norm(w - w_star, ord=np.inf) < tol:
+            n_iter_acc = i
+            break
+
+    # without acceleration
+    n_iter = 0
+    w = np.ones(n_features)
+    for i in range(max_iter):
+        w = rho * w + 1
+
+        if norm(w - w_star, ord=np.inf) < tol:
+            n_iter = i
+            break
+
+    np.testing.assert_allclose(w, w_star)
+    np.testing.assert_allclose(Xw, X @ w_star)
+
+    np.testing.assert_array_equal(n_iter_acc, 13)
+    np.testing.assert_array_equal(n_iter, 99)
+
+
 if __name__ == '__main__':
     pass

skglm/utils.py

@@ -254,3 +254,47 @@ def check_group_compatible(obj):
                 f"'{obj_name}' is not block-separable. "
                 f"Missing '{attr}' attribute."
             )
+
+
+class AndersonAcceleration:
+    """Abstraction of Anderson Acceleration.
+
+    Extrapolate the asymptotic VAR ``w`` and ``Xw``
+    based on ``K`` previous iterations.
+
+    Parameters
+    ----------
+    K : int
+        Number of previous iterates to consider for extrapolation.
+    """
+
+    def __init__(self, K):
+        self.K, self.current_iter = K, 0
+        self.arr_w_, self.arr_Xw_ = None, None
+
+    def extrapolate(self, w, Xw):
+        """Return ``w`` and ``Xw`` extrapolated."""
+        if self.arr_w_ is None or self.arr_Xw_ is None:
+            self.arr_w_ = np.zeros((w.shape[0], self.K+1))
+            self.arr_Xw_ = np.zeros((Xw.shape[0], self.K+1))
+
+        if self.current_iter <= self.K:
+            self.arr_w_[:, self.current_iter] = w
+            self.arr_Xw_[:, self.current_iter] = Xw
+            self.current_iter += 1
+            return w, Xw
+
+        U = np.diff(self.arr_w_, axis=1)  # compute residuals
+
+        # compute extrapolation coefs
+        try:
+            inv_UTU_ones = np.linalg.solve(U.T @ U, np.ones(self.K))
+        except np.linalg.LinAlgError:
+            return w, Xw
+        finally:
+            self.current_iter = 0
+
+        # extrapolate
+        C = inv_UTU_ones / np.sum(inv_UTU_ones)
+        # floating point errors may cause w and Xw to disagree
+        return self.arr_w_[:, 1:] @ C, self.arr_Xw_[:, 1:] @ C