\documentclass[12pt]{article}

\input{preamble}

\begin{document}

\begin{flushleft}
\textbf{Propagator BBVI: a flexible framework for rapid design of BBVI
  optimization algorithms} \\
Adam Massmann (akm2203) \\
\today
\end{flushleft}

\textit{Abstract:}

Here I present \texttt{propagator-bbvi}: a BBVI inference approach
built on top of a propagator computing framework. Propagators consist
of stateful cells that accumulate information about a value, and
autonomous machines that add information about a value. In the case of
BBVI, cells are each variational factor, and the autonomous machines
are the computations that calculate gradients from the joint and
update the variational factors. This framework allows us to rapidly
develop and test many novel stochastic gradient update schemes: each
variational factor can use different step size functions and
convergence specifications, the propagator network can optimize
subsets of variational factors based on subsets of the log joint, and
we can mix both re-parameterization and score-based gradient estimates
on portions of the model that are differentiable or not. Additionally,
the propagator based framework has the potential to scale to massively
parallel computation with no changes to model specification.

\section{Introduction}

Box's loop allows us to iteratively refine assumptions about our
data. However, efficient movement around the loop requires that
computationally tractable inference algorithms are easy to derive from
a model encoding our assumptions. Black box variational inference
(BBVI) performs inference on a very broad set of probabilistic models
without difficult model-specific derivations, and is a path towards
quick iterations over Box's loop \citep[e.g.,][]{ranganath-2014,
  kuc-2017}. Here I build a new approach to BBVI using a propagator
computing framework. The \texttt{propagator-bbvi} Haskell library aims
for:

\begin{itemize}
\item A balance between the ease of specification and the flexibility
  of the inference optimization: the probabilistic model should be easy
  to specify on the computer and BBVI inference should follow
  automatically from the specification. However, the library is also
  designed for easy tuning and/or novel development of the optimization
  algorithms underlying the inference methods. So, there is a trade-off
  between automatic inference from the model specification and the
  flexibility to tune the inference algorithm for tractability. Here
  we favor a balance between adding complexity in what the user must
  specify in order to run inference, and granting the user the ability
  to try many different optimization strategies underlying
  inference. Sensible defaults can mitigate some of the added
  complexity of user specification.
\item Extensibility: the library should be easy for users to
  extend. For example, adding new variational distributions or
  optimization algorithms should be simple and additive to code
  (e.g. no need to rewrite or modify existing code).
\item Ability to scale to massive data and models. Through stochastic
  gradient ascent \citep[e.g.,][]{hoffman-2013, ranganath-2014,
    kuc-2017} and parallel execution we hope to build a framework that
  scales. In keeping with the extensibility goal, development of new
  parallel scheduling and data sub-sampling techniques should be
  possible without significant modification of the larger codebase,
  most importantly the machinery for specifying models.
\end{itemize}

As I describe the library design and its novelty for developing new
approaches to BBVI, I will refer back to these three goals:
\textbf{\emph{optimization flexibility}}, \textbf{\emph{extensibility}}, and \textbf{\emph{scalability}}.

\section{Introduction to propagators}

The propagator computing framework \citep{radul-2009} consists of
stateful ``cells'' that accumulate information about a value and
stateless autonomous ``propagators'' that do computations and add
information to cells. In other words, cells do no computation and
propagators store no data. Propagators are attached to
cells to build algorithms, and there are two primary attachment types for propagators:
\texttt{watch} and \texttt{write} attachments. When cells acquire new
information, they trigger any propagators watching them to execute
their computations with the new data, and those propagators will write
the result of their computations to any cells they have a
\texttt{write} attachment to. If any of these writes are new
information for a given cell, then they will trigger their own
watching propagators to fire. In this manner, propagators and cells
form an asynchronous computation network that runs until no cells in
the network can acquire any new information (``quiescence'').

The propagator model is useful in many computational applications,
including functional reactive programming, dependency directed search,
and constraint satisfaction. For a full discussion I recommend
\citet{radul-2009}. However, why propagators for BBVI? In short, they
are a path towards massively parallel computation and flexible
optimization strategies. Because all computation is asynchronous and
there is no global state, the ordering and scheduling of computation
is completely independent of the propagator network. So, once a
problem is coded in terms of propagators, one can alter the scheduling
of the propagator computations to suit their goals without touching
the propagator network. In this way propagators directly lend
themselves to my design goals: (1) \textit{scalability} through
parallel computation, and (2) \textit{extensibility}: computation
scheduling is independent of the algorithm code, so new parallel
schedulers can be developed without modification of any other pieces
of the library.

\section{BBVI with Propagators}

Variational factors are represented as cells: they accumulate
information moving them closer (on average) to the true
posterior. Computations calculating gradients and updating the
variational factors are the propagators. Variational factor cells are
initialized from priors on the latent variables, variational
factor-specific step-size functions, and definitions of convergence
(e.g. maximum number of steps, small change in variational
factor). Gradient propagators are derived from the log joint of the
model. Using the new BBVI algorithm built on propagators consists of
the following specification:

\begin{enumerate}
\item Specify the log joint: $\log p(x, z)$
\item Initialize variational factor cells for each $z$ (defaults
  automate all but prior specification)
\item Transform the log joint to a gradient propagator, using
  existing \texttt{propagator-bbvi} transformation functions or
  developing a new approach with Haskell abstractions.
\item Specify the attachment type of the propagators to the
  variational factors (defaults automate this)
\end{enumerate}

Steps 2-4 may seem opaque now, and I will elaborate on the
consequences and benefits of this algorithmic approach to BBVI.

Additionally, the library is publicly available at:
\url{https://github.com/massma/propagator-bbvi}, and some sparse (but
growing) documentation is available at:
\url{http://www.columbia.edu/~akm2203/propagator-bbvi-0.1.0.0/}.

\subsection{Initializing variational factor cells}

Variational factor cells contain three pieces of information: the
variational factor and its parameterization (e.g. for a Gaussian, the
mean and variance), a \textit{time} integer corresponding to the
number of times the variational factor has been updated, and a
\textit{memory} vector equal in length to the number of parameters of
the variational factor (e.g. in the case of the Gaussian, a vector of
length 2). The \textit{time} and \textit{memory} data are useful for
step size calculations. For example, in \citet{kuc-2017} the step size
parameterization uses both the number of gradient update steps, as
well as a quantity determined by the history of gradient magnitudes.

In addition to the information stored in variational factor cells,
invariant data are also defined for each variational factor. These
include:

\begin{enumerate}
  \item The prior on the variational factor's latent variable,
  \item A step size function. Inputs: the cell's \textit{time} and
    \textit{memory} data; and outputs: a step size and change in the
    \textit{memory} parameter
  \item A count of the total number of factors in the model's log
    joint that contain the variational factor's corresponding latent
    variable.
\end{enumerate}

The use of this count will be elaborated on in Section \ref{gradient},
but it facilitates distributed updates by many propagators
representing subsets of the log joint to a single variational
factor. The count and step size function aid optimization flexibility:
the user can easily specify many different optimization approaches
and parameters to a variational factor independent of those used
for other variational factors.

\subsubsection{Variational factor representation}

A variational factor's parameters can be stored in any form. For
example, in \texttt{propagator-bbvi}'s current version the Dirichlet
distribution's parameters are stored as a vector, while the normal
distribution's parameters are stored in a record datatype. However,
all variational factors must implement the following functions to be a
part of the necessary typeclasses for use in
\texttt{propagator-bbvi}:

\begin{verbatim}
class DistUtil a where
  fromParamVector :: ParamVector -> a
  toParamVector :: a -> ParamVector
  nParams :: a -> Int

class DistUtil a => Dist a c where
  resample :: a -> GenST s -> ST s c
  logProb :: a -> c -> Double
  paramGradOfLogQ ::  a -> c -> ParamVector
\end{verbatim}

The typeclass \texttt{DistUtil}\footnote{a quick rundown on Haskell
  syntax in case you are not familiar: ``::'' means ``has type.'' so
  ``fromParamVector :: ParamVector -> a'' means the function
  ``fromParamVector'' has a type which takes in a parameter vector and
  outputs a distribution of arbitrary type ``a.'' But don't worry too
  much about it; I will explain all code snippits in normal language.}
provides functions that help with the necessary vector arithmetic for
calculating gradients: \texttt{fromParamVector} converts a vector of
parameters to a distribution, \texttt{toParamVector} converts a
distribution to a vector of its parameters, and \texttt{nParams} takes a
distribution and provides the number of parameters it contains. For a
Gaussian example, \texttt{nParams} would return 2, and
\texttt{fromParamVector} would read a vector containing the mean and
variance and return a Gaussian, while \texttt{toParamVector} would
take a Gaussian and return a vector containing the mean and
variance. These functions are generally trivial to implement.

The typeclass \texttt{Dist} provides the necessary functions to
compute score gradients \citep{ranganath-2014}. \texttt{resample}
takes a distribution and a random number generator, and returns a
sample from the distribution. \texttt{lobProb} takes a distribution
and an observation and returns the log of the probability/density of the
observation. \texttt{paramGradOfLogQ} calculates the gradient of the
log probability of an observation under a distribution, with respect to
the parameters of the distribution. The gradient must be in vector
form, and the locations of the gradients within the vector should
match the location of the parameters in \texttt{fromParamVector} and
\texttt{toParamVector}. This typeclass approach to representing
variational factors aids the design goal of extensibility: to add a
new variational factor for use in \texttt{propagator-bbvi}, one need
only define six functions, three of which are trivial.

However, one may also wish to use reparameterization gradient
updates \citep{kuc-2017} for a variational factor. Any distribution
that implements the following four functions can be used with
reparameterization propagators:

\begin{verbatim}
class (Dist a c) => Differentiable a c where
  epsilon ::  a -> GenST s -> ST s c
  transform :: a -> c -> c
  gradTransform :: a -> c -> ParamVector
  sampleGradOfLogQ ::  a -> c -> c
\end{verbatim}

\texttt{epsilon} generates an un-transformed sample: note that while
the function takes a distribution as input, this is only to maintain
correctness with Haskell's type system. \texttt{transform} transforms
a sample generated from \texttt{epsilon} into a sample from the variational
distribution. \texttt{gradTransform} takes the gradient of the
transformation with respect to the variational distribution's parameters, while
\texttt{sampleGradOfLogQ} takes the gradient of the log probability
function with respect to a transformed sample. In the
Gaussian example if \texttt{epsilon} samples from a standard normal,
\texttt{transform} would be $\mu + \sigma \epsilon$,
\texttt{gradTransform} would be $[1.0, \epsilon]$, and
\texttt{sampleGradOfLogQ} would be $-(z - \mu) / \sigma ^ 2$. To
reiterate, defining these functions for a variational factor is not
necessary for using a factor in \texttt{propagator-bbvi}, but they do
allow one to calculate gradients with the reparameterization
gradient. One of the strengths of
\texttt{propagators-bbvi} is that it allows us to easily mix and match
reparameterization updates with score updates on parts of the model
that are or are not differentiable, and on factors for which we can or
cannot implement the functions in the \texttt{Differentiable}
typeclass. This approach lends itself to the goal of optimization
flexibility: we can use reparameterization gradients (with their lower
variance) on factors where possible, and score gradients elsewhere.

\subsection{Transform the log joint to a gradient propagator }
\label{gradient}

A gradient propagator takes as input the information stored in
variational factor cells and observations and outputs gradient updates
for the variational factor cells. \texttt{propagator-bbvi} contains
functions to help build gradient propagators from the log joint, and
in some cases this can currently be automated. In other cases, users
can build their own transformations from a log joint to a gradient
propagator using \texttt{propagator-bbvi} functions stitched together
with Haskell abstractions. Currently, \texttt{propagator-bbvi} only
focuses on score-gradient updates with Rao-Blackwellization
\citep{ranganath-2014}, and reparameterization gradient updates
\citep{kuc-2017}. Rao-Blackwellization requires specification from the
user, but the resulting decrease in the variance of the gradients is
worth the cost. More importantly, Rao-Blackwellization identifies
which factors in the log joint are reliant on which variational
factors, which has benefits in the propagator framework. For example,
with propagators we can split up the gradient calculations from the
log joint into many separate propagators that make gradient updates
based on subsets of the log joint. If we track the number of factors
from the log joint that a propagator uses in its calculations, then we
can rescale the prior and entropy terms to calculate a gradient update
from a subset of the log joint. A specific example: if a propagator
only calculates the gradient based on one factor in the log joint, and
it is writing its update to a variational cell with 10 total factors
in the log joint, then the prior and entropy terms are scaled down by
$1/10$. Haskell's lazy evaluation also helps the user implement
flexible, distributed propagators. The user can specify the necessary
log joint calculations on the full model, and use this same
specification in very local propagators (for example, a gradient
propagator update based on a single factor in the log joint). When the
propagator writes its gradient to the variational cell, lazy
evaluation insures that only the pieces of the full log joint relevant
to the gradient propagation are actually calculated. So, in many cases
one can use the full joint specification in local propagators without
fear that unneeded likelihood computations will be executed. All of
these features allow us to break up the BBVI update computation into
many different propagators, which can all push updates to variational
cells asynchronously. These design decisions aid scalability.

\subsection{Specify the attachment type of gradient propagators}
% discuss rao-blackwellization here?

Current defaults automate specifying the attachment types. However,
small changes to the attachment types can significantly alter the
ordering and characteristics of the optimization, and are a powerful
abstraction tool for testing new optimization techniques. In addition
to the aforementioned \texttt{watch}, we can also attach propagators
using \texttt{with}, which reads content from a cell but is not
triggered to re-compute if the cell is changed. We can also read a
cell with \texttt{content} which waits until an individual cell has
reached quiescence before proceeding with further computation. By
combining these read attachments with selective decisions to
\texttt{write}, we can quickly develop very novel optimization
strategies and test their impact on rate of convergence (e.g. fully
optimizing some factors with others held fixed, or massive
asynchronous optimization where all distributed propagators
continuously update whenever there is a change to any attached cell).

\section{Summary}

\texttt{propagator-bbvi} is a new algorithmic approach to BBVI that is
extensible, scalable, and allows for easy specification of novel
gradient optimizations used in variational inference. To make best use
of space I focused on design decisions in this report, reflecting the
significant time and thought I dedicated to that process. However,
testing along the way suggested that non traditional gradient
optimizations (e.g. fully or partially optimizing some variational
factors with others held fixed) drastically affects convergence rates.

\appendix
\section{Appendix: Informal discussion of future work}

Most of the work so far on the library focused on designing the new
propagator based BBVI-algorithm, creating an extensible framework, and
testing it. In this process, I created only the necessary variational
distributions (normal, Dirichlet) for testing. However, now
\texttt{propagator-bbvi}'s extensible design and framework is
relatively finalized. Much of the motivation of the project was to
create a tool I can use in my future research, and I will continue to
extend the library as I use it. In particular, this will entail adding
more variational distributions and generalizable transformations from
log joints to gradient propagators. Also, currently
\texttt{propagator-bbvi} is built on the \texttt{propagators} library
which has a very simple serial scheduler. While still fast for serial
computations (and useful as a prototype), scaling up
\texttt{propagator-bbvi} as I intend requires a more sophisticated
parallel scheduler. The good news is that developing new underlying
schedulers is completely independent of the \texttt{propagator-bbvi}
code: any execution scheduler can be used without changing the
existing algorithms and framework.

\bibliography{references}

\end{document}
%%% Local Variables:
%%% mode: latex
%%% TeX-master: t
%%% End:
