Sum-and-sample estimator for learning a stochastic Bernoulli.

bernoulli.py

"""Sum-and-sample estimator for learning a stochastic Bernoulli.

Reproduces Experiment 1 from
    Liu et al, Rao-Blackwellized Stochastic Gradients for Discrete Distributions
    https://arxiv.org/abs/1810.04777
"""

# author: vlad niculae <vlad@vene.ro>
# license: mit

from itertools import product

import torch
from torch.distributions import Bernoulli, Categorical

import matplotlib.pyplot as plt


def get_cfgs(n_vars):
    return torch.stack([torch.Tensor(cfg)
                        for cfg in product(*([0, 1] for _ in range(n_vars)))])


class StochasticToy(torch.nn.Module):

    def __init__(self, target):
        super().__init__()
        self.cfgs = get_cfgs(len(target))
        self.target = target
        self.theta = torch.nn.Parameter(torch.tensor(-4.0))

    def loss_at(self, b):
        if b.ndim == 1:
            b = b.unsqueeze(dim=0)

        return torch.sum((b - self.target) ** 2, dim=1)

    def ll(self, b):
        # log likelihood at b
        logp = torch.nn.functional.logsigmoid(self.theta)
        log1mp = torch.nn.functional.logsigmoid(-self.theta)

        npos = b.sum(dim=1).detach()  # number of positive draws
        nneg = (1 - b).sum(dim=1).detach()  # negative

        return npos * logp + nneg * log1mp

    def loss(self):
        ll = self.ll(self.cfgs)
        return torch.dot(torch.exp(ll), self.loss_at(self.cfgs))


class SampledStochasticToy(StochasticToy):

    def __init__(self, target, sum_top=0, method='sf', control=False):
        super().__init__(target)
        self.method = method

        self.control = control
        self.run_avg = 0
        self.loss_evals = 0
        self.sum_top = sum_top

    def estimator(self, b):
        # compute a surrogate loss such that its grad is SFE at b

        if b.ndim == 1:
            b = b.unsqueeze(dim=0)

        ctrl = 0
        if self.control == 'plus':
            ctrl = self.loss_at(self.sample(b.shape[0])).detach()
        elif self.control == 'avg':
            ctrl = self.run_avg

        loss_at_b = self.loss_at(b)
        lv = (loss_at_b.detach() - ctrl) * self.ll(b) + loss_at_b

        # update moving average
        if self.control == 'avg':
            n_evals = b.shape[0]
            self.loss_evals += n_evals
            self.run_avg += ((loss_at_b.sum().item() - n_evals * self.run_avg)
                             / self.loss_evals)

        return lv

    def sample(self, n=1):
        bern = Bernoulli(logits=self.theta)
        b = bern.sample(sample_shape=(n, 3))
        return b

    def loss(self):

        if self.sum_top == 0:  # no summing, pure stochastic
            b = self.sample()
            return self.estimator(b)

        else:
            # sort configurations by ll
            lls = self.ll(self.cfgs)
            ix = lls.argsort(descending=True)
            ix_top = ix[:self.sum_top]
            ix_bot = ix[self.sum_top:]

            p_top = torch.exp(lls[ix_top].detach())
            loss_top = torch.dot(p_top, self.estimator(self.cfgs[ix_top]))

            p_bot = 1 - p_top.sum()

            if p_bot.item() < 1e-8:
                return loss_top  # basically fully deterministic

            # sample from the tail
            p_rest = torch.exp(lls[ix_bot].detach()) / p_bot
            tail_sample = Categorical(p_rest).sample()
            b = self.cfgs[ix_bot[tail_sample]]

            return loss_top + p_bot * self.estimator(b)



def train_plot(model, axes, max_iter=1000,
               n_trials=10, plot_args=None):

    # opt = torch.optim.Adam(model.parameters(), lr=0.1)
    opt = torch.optim.SGD(model.parameters(), lr=1)

    avg_loss = torch.zeros(max_iter)
    avg_theta = torch.zeros(max_iter)


    for _ in range(n_trials):
        losses = []
        thetas = []
        # restart
        model.theta.data = torch.tensor(-4.0)
        for it in range(max_iter):

            # log exact loss and current learned logit
            losses.append(StochasticToy.loss(model).item())
            thetas.append(model.theta.item())

            opt.zero_grad()
            loss = model.loss()
            loss.backward()
            opt.step()

        avg_loss += torch.tensor(losses)
        avg_theta += torch.tensor(thetas)

    avg_loss /= n_trials
    avg_theta /= n_trials

    ax_loss, ax_theta = axes
    ax_loss.plot(avg_loss, **plot_args)
    ax_theta.plot(avg_theta, **plot_args)


def main():

    torch.manual_seed(42)
    # torch.set_default_dtype(torch.double)

    target = torch.Tensor([.6, .51, .48])

    _, axes = plt.subplots(1, 2, constrained_layout=True)
    train_plot(StochasticToy(target), axes, plot_args={'label': 'exact'})

    control = 'avg'
    # control = 'plus'
    # control = None

    for k in range(4):
        train_plot(SampledStochasticToy(target, sum_top=k, method='sf',
                                        control=control),
                   axes, plot_args={'label': f'SF-{k}'})

    axes[1].legend()
    axes[0].set_ylabel("expected loss")
    axes[1].set_ylabel("$\\theta$")
    axes[0].set_xlabel("steps")
    axes[1].set_xlabel("steps")
    plt.show()


if __name__ == '__main__':
    main()