Black Box Variational Inference: Bayesian Linear Regression

bblinreg.py

import os

import numpy as np


def generate_data(n, p, seed):
    np.random.seed(seed)
    beta = np.random.normal(size=p, scale=2.0)
    X = np.random.normal(size=(n,p))
    y = np.dot(X, beta)
    return X, y, beta


def save_data(X, y, beta, dirname):
    if not os.path.exists(dirname):
        os.makedirs(dirname)
    with open(os.path.join(dirname, "betas.txt"), "w") as f:
        for b in beta:
            f.write("%f\n" % b)
    with open(os.path.join(dirname, "y.txt"), "w") as f:
        for yy in y:
            f.write("%f\n" % yy)
    with open(os.path.join(dirname, "X.txt"), "w") as f:
        for x in X:
            f.write("%s\n" % (" ".join([str(xx) for xx in x])))


class BBGaussian(object):
    def __init__(self, eta=1.0):
        self.params = np.array([0.0, 1.0])
        self.gt = np.array([0.0, 0.0])
        self.eta = eta

    def nparam(self):
        return len(self.params)

    def setparams(self, params):
        self.params = params
        if self.params[1] <= 0.0:
            print "UHOH! sigmasq < 0: %f" % self.params[1]
            self.params[1] = 1e-6

    def applygrad(self, paramgrad, rho=None):
        if rho is None:
            # AdaGrad
            self.gt += np.square(paramgrad)
            self.params += paramgrad * self.eta / np.sqrt(self.gt)
        else:
            # standard Robbins-Monro
            self.params += rho * paramgrad
        if self.params[1] <= 0.0:
            print "UHOH! sigmasq < 0: %f" % self.params[1]
            self.params[1] = 1e-6

    def getparams(self):
        return self.params

    def sample(self, size=1):
        mu, sigmasq = self.params
        return np.random.normal(loc=mu, scale=np.sqrt(sigmasq), size=size)

    def loggrad(self, x):
        mu, sigmasq = self.params
        mu_grad = (x - mu) / sigmasq
        sigmasq_grad = 0.5 * (np.square(x - mu)/sigmasq - 1.0) / sigmasq
        return np.array([mu_grad, sigmasq_grad])

    def logprob(self, x):
        mu, sigmasq = self.params
        return -0.5 * np.log(2 * np.pi * sigmasq) - 0.5 * np.square(x - mu) / sigmasq


class BBLinearRegression(object):
    def __init__(self, X, y, S=100, tau=0.7, kappa=1.0, sigmasq=1.0, epssq=1.0, eta=1.0, ada=True):
        p = X.shape[1]
        self.X = X
        self.y = y
        self.qbetas = [BBGaussian(eta=eta) for i in xrange(p)]
        self.sigmasq = sigmasq
        self.epssq = epssq
        self.S = S
        self.kappa = kappa
        self.tau = tau
        self.ada = ada

    def logjoint(self, betas):
        n, p = self.X.shape
        quad_beta = 0.5 * np.sum(np.square(betas)) / self.sigmasq
        norm_beta = 0.5 * p * np.log(2.0 * np.pi * self.sigmasq)
        quad_lik = 0.5 * np.sum(np.square(self.y - np.dot(self.X, betas))) / self.epssq
        norm_lik = 0.5 * n * np.log(2.0 * np.pi * self.epssq)
        return -(quad_beta + norm_beta + quad_lik + norm_lik)

    def qsample(self):
        samples = [qb.sample(size=self.S) for qb in self.qbetas]
        p = len(self.qbetas)
        samps = np.empty((self.S, p))
        for j in xrange(p):
            for s in xrange(self.S):
                samps[s,j] = samples[j][s]
        return samps

    def run_iter(self, rho=None):
        p = X.shape[1]
        samples = self.qsample()
        S = len(samples)
        grads = []
        for i, qb in enumerate(self.qbetas):
            f = np.empty((S, qb.nparam()))
            h = np.empty((S, qb.nparam()))
            for s, z in enumerate(samples):
                h[s] = qb.loggrad(z[i])
                f[s] = h[s] * (self.logjoint(z) - qb.logprob(z[i]))
            grad = np.empty(qb.nparam())
            for d in xrange(qb.nparam()):
                fd = f[:,d]
                hd = h[:,d]
                cov = np.cov(f[:,d], h[:,d])
                a = cov[0,1] / cov[1,1]
                grad[d] = np.mean(f[:,d] - a * h[:,d])
            grads.append(grad)
        for grad, qb in zip(grads, self.qbetas):
            if rho is None:
                qb.applygrad(grad)
            else:
                qb.applygrad(grad, rho)

    def run(self, niter):
        for i in xrange(niter):
            print "Iteration %d" % (i + 1)
            if self.ada:
                self.run_iter()
            else:
                rho = np.power(i + self.tau, -self.kappa)
                self.run_iter(rho)

    def print_params(self, outdir):
        if not os.path.exists(outdir):
            os.makedirs(outdir)
        mufile = os.path.join(outdir, "params.mu")
        sigmasqfile = os.path.join(outdir, "params.sigmasq")
        with open(mufile, "w") as fmu, open(sigmasqfile, "w") as fsigmasq:
            for qb in self.qbetas:
                params = qb.getparams()
                fmu.write("%f\n" % params[0])
                fsigmasq.write("%f\n" % params[1])


if __name__ == '__main__':
    X, y, beta = generate_data(100, 10, 20)
    datadir = os.path.abspath("data")
    save_data(X, y, beta, datadir)

    # standard Robbins-Monro (does terribly...):
    #lr = BBLinearRegression(X, y, S=1000, tau=20.0, kappa=0.55, ada=False)
    # AdaGrad:
    lr = BBLinearRegression(X, y, S=1000, eta=1.0)
    lr.run(100)
    outdir = os.path.abspath("output")
    lr.print_params(outdir)