【Python】乗数法を用いた等式制約付き LASSO 回帰

ある論文[^1]のアルゴリズムを実装する際に，線形等式制約付きの LASSO 回帰問題

```math
\begin{align}
\mathop{\text{minimize}}_\beta\quad &\frac{1}{2}\|y-X\beta\|_2^2 + \rho \|\beta\|_1 \quad\\
\text{subject to}\quad &A\beta = b
\end{align}
```

を解く場面があった．以前実装した[等式・不等式制約付き LASSO](https://qiita.com/Mrrmm252/items/79423cf423cb3ebb92cf) では計算が重いので，等式制約のみのバージョンを新たに実装した．
実装にあたって，$Q=X^TX,\ p=-X^Ty$ とおいて，等価な問題

```math
\begin{align}
\mathop{\text{minimize}}_\beta\quad &\frac{1}{2}\beta^TQ\beta + p^T\beta + \rho \|\beta\|_1 \quad\\
\text{subject to}\quad &A\beta = b
\end{align}
```

を解くことを考えた．

## 乗数法
線形等式制約付き最適化問題

```math
\begin{align}
\mathop{\text{minimize}}_{\beta,z}\quad &f(\beta) + g(z) \quad\\
\text{subject to}\quad &M\beta + Fz = c
\end{align}
```

を扱う手法として，乗数法が知られている．乗数法では，拡張ラグランジュ関数

```math
\mathcal{L}_\tau (\beta,z,\nu) = f(\beta) + g(z) + \nu^T(M\beta + Fz - c) + \frac{\tau}{2}\|M\beta + Fz - c\|_2^2
```

を $\beta,z$ について交互に最小化するように変数を更新する．

```math
\newcommand{\argmin}{\mathop{\rm arg\,min}\limits}
\begin{align}
\beta^{(t+1)} &\leftarrow \argmin_\beta \mathcal{L}_\tau (\beta, z^{(t)}, \nu^{(t)})\\
z^{(t+1)} &\leftarrow \argmin_z \mathcal{L}_\tau (\beta^{(t+1)}, z, \nu^{(t)})\\
\nu^{(t+1)} &\leftarrow \nu^{(t)} + \tau(M\beta^{(t+1)}+Fz^{(t+1)}-c) \tag{1}
\end{align}
```

## 乗数法の等式制約付き LASSO 回帰への適用

```math
\begin{align}
\mathop{\text{minimize}}_\beta\quad &\frac{1}{2}\beta^TQ\beta + p^T\beta + \rho \|z\|_1 \quad\\
\text{subject to}\quad &\left(\begin{array}{c}
A\\
I
\end{array}\right)\beta + \left(\begin{array}{c}
O\\
-I
\end{array}\right) z = \left(\begin{array}{c}
b\\
0
\end{array}\right)
\end{align}
```

と問題を書き換えることで，

```math
\begin{align}
f(\beta)=\frac12 \beta^TQ\beta + p^T\beta,\ 
g(z) = \rho\|z\|_1,\\
M=\left(\begin{array}{c}
A\\
I
\end{array}\right),\ 
F=\left(\begin{array}{c}
O\\
-I
\end{array}\right),\ 
c=\left(\begin{array}{c}
b\\
0
\end{array}\right)
\end{align}
```
として，乗数法を適用することができる．その際，$\nu=(u^T, v^T)^T$ として，拡張ラグランジュ関数を

```math
\mathcal{L}_\tau (\beta,z,\nu) = \frac12 \beta^TQ\beta + p^T\beta + \rho\|z\|_1 + u^T(A\beta - b) + v^T(\beta - z) + \frac{\tau}{2}\|A\beta - b\|_2^2 + \frac{\tau}{2}\|\beta - z\|_2^2
```

と書くことができる．

### β の更新式
最適性条件

```math
\frac{\partial \mathcal{L}_\tau}{\partial \beta}(\beta, z^{(t)}, \nu^{(t)}) = Q\beta + p + A^Tu^{(t)} + v^{(t)} + \tau A^T(A\beta - b) + \tau(\beta - z^{(t)}) = 0 \tag{2}
```
を解いて，

```math
\beta^{(t+1)} = \left(Q+\tau(A^TA+I)\right)^{-1}\left(-p-A^Tu^{(t)}-v^{(t)}+\tau A^Tb+\tau z^{(t)}\right)
```

で更新する．

### z の更新式
最適性条件

```math
\partial_z \mathcal{L}_\tau(\beta^{(t+1)}, z, \nu^{(t)}) = \rho\partial_z \|z\|_1 - v^{(t)} - \tau\left(\beta^{(t+1)} - z\right)\ni0
```

から，更新式

```math
z^{(t+1)} = s\left(\beta^{(t+1)} + \frac1\tau v^{(t)}; \frac\rho\tau\right)
```

を得る．ここで，$s(\cdot; \cdot)$ は軟閾値関数

```math
s(x; \lambda) = \mathrm{sign}\,(x)\max\left\{x - \lambda, 0\right\}
```

を表す．

### u, v の更新式
(1) 式より，$u$ は

```math
\begin{align}
u^{(t+1)} \leftarrow u^{(t)} + \tau\left(A\beta^{(t+1)}-b\right)
\end{align}
```

で更新する．これは (2) 式

```math
\begin{align}
&Q\beta^{(t+1)} + p + A^Tu^{(t)} + v^{(t)} + \tau A^T\left(A\beta^{(t+1)} - b\right) + \tau\left(\beta^{(t+1)} - z^{(t)}\right)\\
 =\ &Q\beta^{(t+1)} + p + A^T\left(u^{(t)} + \tau \left(A\beta^{(t+1)} - b\right)\right) + v^{(t)} + \tau\left(\beta^{(t+1)}-z^{(t)}\right) = 0

\end{align}
```

とラグランジュ関数 

```math
\mathcal{L}\left(\beta, z, u, v\right) = \frac12 \beta^TQ\beta + p^T\beta + \rho\|z\|_1 + u^T(A\beta - b) + v^T(\beta - z)
```

の $\beta$ 微分

```math
\frac{\partial \mathcal{L}}{\partial \beta}(\beta^{(t+1)}, z^{(t)}, \nu) = Q\beta^{(t+1)} + p + A^Tu + v = 0
```

を比較することにより得られる．$v$ の更新式についても同様．

### 変数の初期化
```math
\begin{align}
\mathop{\text{minimize}}_\beta\quad &\|\beta\|_1 \quad\\
\text{subject to}\quad &A\beta= b
\end{align}
```

を解くことで制約条件を満たす $\beta^{(0)}$ を得る．
この問題は，新たな変数 $\gamma$ を導入することで線形計画問題とみなすことができる．

```math
\begin{align}
\mathop{\text{minimize}}_{\beta,\gamma}\quad &\sum_{i=1}^p \gamma_i \quad\\
\text{subject to}\quad &A\beta= b,\\
&\ \gamma\geq0,\ \gamma\geq\beta,\ \gamma\geq-\beta
\end{align}
```

### アルゴリズムまとめ
1. $\beta^{(0)}$ を初期化し，$z^{(0)}\leftarrow\beta^{(0)},\ u^{(0)}\leftarrow0,\ v^{(0)}\leftarrow0,\ t\leftarrow 0$ とする．
1. 停止条件が満たされるなら，$z^{(t)}$ を出力
1. $\beta^{(t+1)} \leftarrow \left(Q+\tau(A^TA+I)\right)^{-1}\left(-p-A^Tu^{(t)}-v^{(t)}+\tau A^Tb+\tau z^{(t)}\right)$
1. $z^{(t+1)} \leftarrow s\left(\beta^{(t+1)} + \frac1\tau v^{(t)}; \frac\rho\tau\right)$
1. $u^{(t+1)} \leftarrow u^{(t)} + \tau\left(A\beta^{(t+1)}-b\right)$
1. $v^{(t+1)} \leftarrow v^{(t)} + \tau\left(\beta^{(t+1)}-z^{(t+1)}\right)$
1. $t\leftarrow t+1$ として，ステップ 2 へ．

## ソースコード
Python で実装した．
$\beta$ の初期化には `scipy.optimize.linprog` を用いた．

```python
import numpy as np
from scipy.optimize import linprog


class EqualityConstrainedL1QuadraticProgramming:
    """
    Problem
    ----------
    minimize 0.5 * βQβ + pβ + ρ||β||_1
    subject to Aβ = b

    Algorithm
    ----------
    Alternating Direction Method of Multipliers (ADMM)

    Parameters
    ----------
    A : np.ndarray, optional (default=None)
        The equality constraint matrix.

    b : np.ndarray, optional (default=None)
        The equality constraint vector.

    rho : float, optional (default=1.0)
        Constant that multiplies the L1 term.

    tau : float, optional (default=None)
        Constant that used in augmented Lagrangian function.

    tol : float, optional (default=1e-4)
        The tolerance for the optimization.

    max_iter : int, optional (default=300)
        The maximum number of iterations.

    extended_output : bool, optional (default=False)
        If set to True, objective function value will be saved in `self.f`.

    Attributes
    ----------
    f : a list of float
        objective function values.

    coef_ : array, shape (n_features,) | (n_targets, n_features)
        parameter vector.

    n_iter_ : int | array-like, shape (n_targets,)
        number of iterations run by admm to reach the specified tolerance.
    """
    def __init__(
            self,
            A: np.ndarray,
            b: np.ndarray,
            rho: float = 1.0,
            tau: float = 1.0,
            tol: float = 1e-4,
            max_iter: int = 300,
            extended_output: bool = False
    ):
        self.A = A
        self.b = b
        self.rho = rho
        self.tau = tau
        self.tol = tol
        self.max_iter = max_iter
        self.extended_output = extended_output

        self.f = list()
        self.coef_ = None
        self.n_iter_ = None

    def _linear_programming(self, n_features: int) -> np.ndarray:
        """
        Solve following problem.

        Problem:
            minimize ||β||_1 subject to Aβ=b

        Solver:
            scipy.optimize.linprog

        Parameters
        ----------
        n_features : int
            The dimension of decision variables

        Returns
        -------
        : np.ndarray, shape = (n_features, )
            The values of the decision variables that minimizes the objective function while satisfying the constraints.
        """
        # equality constraint matrix and vector
        c = np.hstack((np.zeros(n_features), np.ones(n_features)))
        A_eq = np.hstack((self.A, np.zeros_like(self.A)))
        b_eq = self.b

        # inequality constraint matrix and vector
        eye = np.eye(n_features)
        A_ub = np.vstack((
            np.hstack((eye, -eye)),
            np.hstack((-eye, -eye)),
            np.hstack((np.zeros((n_features, n_features)), -eye))
        ))
        b_ub = np.zeros(n_features * 3)

        return linprog(c=c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq)['x'][:n_features]

    def _prox(self, x: np.ndarray, rho: float) -> np.ndarray:
        """
        Proximal operator for L1 constraint

        Parameters
        ----------
        x : np.ndarray, shape = (n_features, )
            1D array

        rho : float
            a threshold
        """
        return np.sign(x) * np.maximum((np.abs(x) - rho), 0)

    def _objective_function(self, x: np.ndarray, Q: np.ndarray, p: np.ndarray, rho: float):
        """
        Return 1 / 2 * x^T Qx + p^T x + rho * ||x||_1
        """
        return 0.5 * x.dot(Q.dot(x)) + p.dot(x) + np.sum(np.abs(x)) * rho

    def fit(self, Q: np.ndarray, p: np.ndarray) -> None:
        """
        Parameters
        ----------
        Q : np.ndarray, shape = (n_features, n_features)
            Quadratic coefficient.

        p : np.ndarray, shape = (n_features, )
            Linear coefficient.
        """
        n_features = Q.shape[0]
        tau_inv = 1 / self.tau

        # initialize constants
        Q_inv = np.linalg.inv(Q * tau_inv + self.A.T.dot(self.A) + np.eye(n_features))
        p_tau = p * tau_inv
        Ab = self.A.T.dot(self.b)
        rho_tau = self.rho * tau_inv

        # initialize variables
        beta = self._linear_programming(n_features)
        z = np.copy(beta)
        u = np.zeros_like(self.b)
        v = np.zeros_like(beta)

        cost = self._objective_function(beta, Q, p, self.rho)
        # save objective function value if necessary
        if self.extended_output:
            self.f.append(cost)

        # main loop
        for k in range(self.max_iter):
            r = - p_tau - self.A.T.dot(u) - v + Ab + z
            beta = np.dot(Q_inv, r)
            z = self._prox(beta + v, rho_tau)
            u = u + self.A.dot(beta) - self.b
            v = v + beta - z

            pre_cost = cost
            cost = self._objective_function(beta, Q, p, self.rho)
            # save objective function value if necessary
            if self.extended_output:
                self.f.append(cost)
            if np.abs(cost - pre_cost) < self.tol:
                break

        # save result
        self.coef_ = z
        self.n_iter_ = k
```

### テスト
次の線形等式制約付き最適化問題を扱う．

```math
\begin{align}
\mathop{\text{minimize}}_\beta\quad &\frac{1}{2}\|y-X\beta\|_2^2 + \rho \|\beta\|_1 \quad\\
\text{subject to}\quad &\sum_i\beta_i=1
\end{align}
```

雑に $y = X\beta + \varepsilon$ でデータを生成し，Scikit-Learn の無制約の Lasso と比較した．

```python
from sklearn.linear_model import Lasso

n_samples, n_features = 1000, 100
param = 0.5

# データ生成
np.random.seed(0)
X = np.random.randn(n_samples, n_features)
beta = np.zeros(n_features)
beta[:n_features//10] = np.random.randn(n_features//10)
beta = beta / np.sum(beta)
y = X.dot(beta) + np.random.randn(n_samples)

# 実装したもの
clf = EqualityConstrainedL1QuadraticProgramming(A = np.ones((1, n_features)), b = np.array([1.]), rho=param, tau=n_features)
clf.fit(X.T.dot(X) / n_samples, -X.T.dot(y) / n_samples)

# Scikit-Learn
lasso = Lasso(alpha=param)
lasso.fit(X, y)
```

実験結果: $\beta$ の値をプロット，$\sum_i\beta_i,\ \frac12\|y-X\beta\|_2^2$ の値を表示．
Scikit-learn と比較して，制約条件を満たす $\beta$ が得られていることが確認できる．

```python
label_true = f'True: sum = {round(np.sum(beta), 5)}, func = {round(0.5 * np.linalg.norm(y - X.dot(beta)), 5)}'
label_equality_constrained = f'Equality Constrained: sum = {round(np.sum(clf.coef_), 5)}, func = {round(0.5 * np.linalg.norm(y - X.dot(clf.coef_)), 5)}'
label_sklearn = f'Scikit-learn: sum = {round(np.sum(lasso.coef_), 5)}, func = {round(0.5 * np.linalg.norm(y - X.dot(lasso.coef_)), 5)}'

plt.figure(figsize=(15, 10))
plt.plot(beta, 'go-', label=label_true)
plt.plot(clf.coef_, 'ro-', label=label_equality_constrained)
plt.plot(lasso.coef_, 'bo-', label=label_sklearn)
plt.grid(True)
plt.legend()
plt.show()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/394475/d00c914b-7141-c687-4aad-a88fe1835417.png)


[^1]: [Ho, J., Xie, Y., & Vemuri, B. (2013, February). On a nonlinear generalization of sparse coding and dictionary learning. In International conference on machine learning (pp. 1480-1488).](http://proceedings.mlr.press/v28/ho13a.pdf)