LR ADMM

lr_admm.py

""" Fit logistic regression by ADMM. """
# Author: Vlad Niculae <v.niculae@uva.nl>
# License: MIT

import numpy as np
from sklearn.datasets import load_breast_cancer
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

from scipy.linalg import cho_factor, cho_solve
from scipy.special import expit, log_expit


def logloss(w, X, y_sign):
    return -log_expit(y_sign * (X @ w)).sum()


def lr_admm(X, y, rho=.01, n_iter=500):

    XtX = X.T @ X
    C_and_lower = cho_factor(XtX)

    n, d = X.shape

    w = np.zeros(d)
    s = np.zeros(n)
    u = np.zeros(n)

    for it in range(n_iter):
        print("iter", it, "loss", logloss(w, X, y))

        yxwpu = u- y * (X @ w)

        # update s by bisection:
        def f(ss):
            return -expit(-ss) + rho*(ss + yxwpu)

        s = np.ones(n) * -1000
        width = 2000
        # assert np.all(f(s) <= 0)
        # assert np.all(f(s+width) >= 0)

        for _ in range(30):
            width /= 2
            s[f(s+width) <= 0] += width

        # update w
        b = X.T @ (y * (s+u))
        w = cho_solve(C_and_lower, b)

        # update u
        u += s - y * (X @ w)

    return w


def main():
    X, y = load_breast_cancer(return_X_y=True)
    y_sign = 2*y-1

    scaler = StandardScaler()

    n = X.shape[0]
    X = scaler.fit_transform(X)
    X = np.column_stack((np.ones(n), X))
    lr = LogisticRegression(max_iter=10000, penalty=None, fit_intercept=False).fit(X, y)
    wtrue = lr.coef_[0]
    print("acc", lr.score(X, y))
    print("loss", logloss(wtrue, X, y_sign))

    w = lr_admm(X, y_sign)
    print("admm accuracy", np.mean((X @ w >= 0) == y))

    
if __name__ == '__main__':
    main()

lr_multi.py

""" Fit logistic regression by ADMM. """
# Author: Vlad Niculae <v.niculae@uva.nl>
# License: MIT


import numpy as np
from scipy.linalg import cho_factor, cho_solve
from scipy.special import softmax, log_softmax
from sklearn.datasets import load_wine
from sklearn.preprocessing import LabelBinarizer, StandardScaler
from sklearn.linear_model import LogisticRegression


def logloss(W, X, Y):
    S = X @ W
    return -np.sum(log_softmax(S, axis=-1) * Y)


def lr_admm(X, Y, rho=0.01, n_iter=100):
    XtX = X.T @ X
    C_and_lower = cho_factor(XtX)

    n, d = X.shape
    _, c = Y.shape

    W = np.zeros((d, c))
    S = np.zeros((n, c))
    U = np.zeros((n, c))

    for it in range(n_iter):
        print("iter", it, "loss", logloss(W, X, Y))

        # update S
        S = softmax_fixed_point(rho, Y + rho * (X @ W - U), S_init=S)

        # update W
        W = cho_solve(C_and_lower, X.T @ (S+U))

        # update U
        U += S - X @ W

    return W


def softmax_fixed_point(rho, B, S_init=None, n_iter=20):
    # find a solution of
    # softmax(S) + rho*S = B

    # F(S) = softmax(S) + rho*S - B
    # F'(S) = diag(P+rho) + PP'

    S = S_init if S_init is not None else np.zeros_like(B)

    for _ in range(n_iter):

        P = softmax(S, axis=-1)
        F = P + rho*S - B

        dot = np.sum(P * F / (P + rho), axis=-1)
        denom = 1 + np.sum(P ** 2 / (P + rho), axis=-1)
        coef = (dot/denom)[:, np.newaxis]
        JinvF = (F - coef*P) / (P + rho)
        S -= JinvF

    return S


def main():
    X, y = load_wine(return_X_y=True)

    scaler = StandardScaler()
    X = scaler.fit_transform(X)
    Y = LabelBinarizer().fit_transform(y)
    lr = LogisticRegression(max_iter=20,
                            penalty=None,
                            fit_intercept=False,
                            verbose=False).fit(X, y)
    print("acc", lr.score(X, y))
    wtrue = lr.coef_.T
    print("loss", logloss(wtrue, X, Y))

    W = lr_admm(X, Y)
    print("admm accuracy", np.mean(np.argmax(X @ W, axis=-1) == y))


if __name__ == '__main__':
    main()