A new method for parameter estimation in probabilistic models: Minimum probability flow

Jascha Sohl-Dickstein, Peter Battaglino, Michael R. DeWeese

Acknowledgments

We would like to thank Javier Movellan, Tamara Broderick, Miroslav Dudík, Gašper Tkačik, Robert E. Schapire, William Bialek for sharing work in progress and data; Ashvin Vishwanath, Jonathon Shlens, Tony Bell, Charles Cadieu, Nicole Carlson, Christopher Hillar, Kilian Koepsell, Bruno Olshausen and the rest of the Redwood Center for many useful discussions; and the James S. McDonnell Foundation (JSD, PB, JSD) and the Canadian Institute for Advanced Research - Neural Computation and Perception Program (JSD) for financial support.

References

Appendix A Taylor Expansion of KL Divergence

The minimum probability flow learning objective function K(θ)K\left(\theta\right) is found by taking up to the first order terms in the Taylor expansion of the KL divergence between the data distribution and the distribution resulting from running the dynamics for a time ϵ\epsilon:

where we used the fact that ∑i∈Dpi(t)+∑i∉Dpi(t)=1\sum_{i\in\mathcal{D}}p_{i}^{(t)}+\sum_{i\notin\mathcal{D}}p_{i}^{(t)}=1. This implies that the rate of growth of the KL divergence at time t=0t=0 equals the total initial flow of probability from states with data into states without.

Appendix B Convexity

As observed by Macke and Gerwinn Macke and Gerwinn (2009), the MPF objective function is convex for models in the exponential family.

The first term in the Hessian is a weighted sum of outer products, with non-negative weights 14Γijpi(0)\frac{1}{4}\Gamma_{ij}p_{i}^{(0)}, and is thus positive semidefinite. The second term is for models in the exponential family (those with energy functions linear in their parameters).

Parameter estimation for models in the exponential family is therefore convex using minimum probability flow learning.

Appendix C Relationship of MPF to other techniques

Score matching, developed by Aapo Hyvärinen Hyvärinen (2005), is a method that learns parameters in a probabilistic model using only derivatives of the energy function evaluated over the data distribution (see Equation (C-5)). This sidesteps the need to explicitly sample or integrate over the model distribution. In score matching one minimizes the expected square distance of the score function with respect to spatial coordinates given by the data distribution from the similar score function given by the model distribution. A number of connections have been made between score matching and other learning techniques Hyvärinen (2007b); Sohl-Dickstein and Olshausen (2009); Movellan (2008); Lyu (2009). Here we show that in the correct limit, MPF also reduces to score matching.

For a dd-dimensional, continuous state space, we can write the MPF objective function as

where the sum ∑x∈D\sum_{x\in\mathcal{D}} is over all data samples, and NN is the number of samples in the data set D\mathcal{D}. Now we assume that transitions are only allowed from states xx to states yy that are within a hypercube of side length ϵ\epsilon centered around xx in state space. (The master equation will reduce to Gaussian diffusion as ϵ→0\epsilon\rightarrow 0.) Thus, the function g(y,x)g(y,x) will equal 1 when yy is within the xx-centered cube (or xx within the yy-centered cube) and 0 otherwise. Calling this cube CϵC_{\epsilon}, and writing y=x+αy=x+\alpha with α∈Cϵ\alpha\in C_{\epsilon}, we have

If we Taylor expand in α\alpha to second order and ignore cubic and higher terms, we get

which, removing a constant offset and scaling factor, is exactly equal to the score matching objective function,

Score matching is thus equivalent to MPF when the connectivity function g(y,x)g(y,x) is non-zero only for states infinitesimally close to each other. It should be noted that the score matching estimator has a closed-form solution when the model distribution belongs to the exponential family Hyvärinen (2007a), so the same can be said for MPF in this limit.

C.2 Contrastive divergence

Contrastive divergence Welling and Hinton (2002); Carreira-Perpiñán and Hinton (2004) is a variation on steepest gradient descent of the maximum (log) likelihood (ML) objective function. Rather than integrating over the full model distribution, CD approximates the partition function term in the gradient by averaging over the distribution obtained after taking a few, or only one, Markov chain Monte Carlo (MCMC) step away from the data distribution (Equation C-7). Qualitatively, one can imagine that the data distribution is contrasted against a distribution which has evolved only a small distance towards the model distribution, whereas it would be contrasted against the true model distribution in traditional MCMC approaches. Although CD is not guaranteed to converge to the right answer, or even to a fixed point, it has proven to be an effective and fast heuristic for parameter estimation MacKay (2001); Yuille (2005).

The contrastive divergence update rule can be written in the form

where TijT_{ij} is the probability of transitioning from state jj to state ii in a single Markov chain Monte Carlo step (or a small number of steps). Equation C-7 has obvious similarities to the MPF learning gradient

Thus, steepest gradient descent under MPF resembles CD updates, but with the MCMC sampling/rejection step TijT_{ij} replaced by a weighting factor gijexp⁡[12(Ej(θ)−Ei(θ))]g_{ij}\exp\left[\frac{1}{2}\left(E_{j}\left(\theta\right)-E_{i}\left(\theta\right)\right)\right]. Note that this difference in form provides MPF with a well-defined objective function, and it guarantees consistency (i.e., there is a global minimum when model and data distributions agree).

Appendix D Sampling the connectivity matrix 𝚪𝚪\mathbf{\Gamma}

The MPF learning scheme is blind to regions in state space which data states don’t have any connectivity to - the flow at time 0 is only a function of the states that are directly connected to data states. To get the most informative learning signal, it seems sensible to encourage probability flow directly between data states and states that are probable under the model. That way the objective function is sensitive to the regions which are probable under the model. We believe nearest neighbor connectivity schemes are effective largely because as the parameters converge the regions around data states become the high probability regions for the model. We wish to try connectivity schemes other than nearest neighbors to allow probability to most efficiently flow between data states and high probability model states. In order to do so we need to slightly extend the MPF algorithm. We do this by allowing the connectivity pattern in Γ\Gamma to be sampled independently in every infinitesimal time step.

Since Γ\Gamma is now sampled, we will modify detailed balance to demand that, averaging over the choices for Γ\Gamma, the net flow between pairs of states is 0.

where the ensemble average is over the connectivity scheme for Γ\Gamma. We describe the connectivity scheme via a function gijg_{ij}, such that the probability of there being a connection from state jj to state ii at any given moment is gijg_{ij}. We also introduce a function FijF_{ij}, which provides the value Γij\Gamma_{ij} takes on when a connection occurs from jj to ii. That is, it is the probability flow rate when flow occurs -

F\mathbf{F} is underconstrained by the above equation. Motivated by symmetry and aesthetics, we choose as the form for the (non-zero, non-diagonal) entries in F\mathbf{F}

Similarly, its average value can be written as

So, we can use any connectivity scheme g\mathbf{g} in learning. We just need to scale the non-zero, non-diagonal entries in Γ\mathbf{\Gamma} by (gjigij)12\left(\frac{g_{ji}}{g_{ij}}\right)^{\frac{1}{2}} so as to compensate for the biases introduced by the connectivity scheme.

The full MPF objective function in this case is

where the inner sum is found by averaging over samples from gijg_{ij}.

Appendix E Additional information on Ising spin glass example from main text

The four competing techniques against which MPF was compared are: mean field theory (MFT) with Thouless-Anderson-Palmer (TAP) corrections Thouless DJ (1977), one-step and ten-step contrastive divergence Carreira-Perpiñán and Hinton (2004) (CD-1 and CD-10), and pseudolikelihood Besag (1975). Table 1 shows the relative performance at convergence for each technique in terms of convergence time and mean square error in coupling strengths and pairwise correlations.

MFT with TAP involves approximating the Gibbs free energy of the model with a second-order Plefka expansion T (1982). MFT+TAP is fast because it involves only an inversion of the magnetic susceptibility matrix, but it can perform poorly, for instance near criticality Fischer KH (1991).

Contrastive divergence approximates the term involving the partition function in \partial_{\theta}D_{KL}\left(\mathbf{p^{(0)}}\bigl{|}\!\bigr{|}\mathbf{p^{(\infty)}}\left(\theta\right)\right), via a Markov chain which is initialized at the data distribution p(0)\mathbf{p}^{(0)}, and then truncated after only a small number of sampling steps. It is commonly used in machine learning, and provides an effective and fast stochastic parameter update rule for learning in many probabilistic models. However, it is not guaranteed to converge to a fixed point, and it does not correspond exactly to an objective function. The relationship between our technique and contrastive divergence is discussed in the Supplemental Material.

Pseudolikelihood approximates the joint probability distribution of a collection of random variables with a product of conditional distributions, where each factor is the distribution of a single random variable conditioned on the others:

This approach often leads to surprisingly good estimates, despite the extreme nature of the approximation.

E.2 Optimization steps taken for parameter estimation algorithms

Both Minimum Probability Flow and Pseudolikelihood have well defined objective functions and gradients. Parameter estimation was thus performed by applying an off the shelf L-BFGS (quasi-Newton gradient descent) implementation Schmidt (2005) to their objective functions evaluated over the full training dataset D\mathcal{D}.

E.3 Contrastive Divergence

The CD update rule was computed using the full training dataset. The learning rate was annealed in a linear fashion from 3.0 to 0.1 to accelerate convergence.

Mean field theory requires the computation of the inverse of the magnetic susceptibility matrix, which, for strong correlations, was often singular. A regularized pseudoinverse was used in the following manner:

where II is the identity matrix, M+M^{+} denotes the Moore-Penrose pseudoinverse of a matrix MM, χ\chi is the magnetic susceptibility χij=<xixj>−<xi><xj>\chi_{ij}=\left<x_{i}x_{j}\right>-\left<x_{i}\right>\left<x_{j}\right>, and λ\lambda is a regularizing parameter. This technique is known as stochastic robust approximation Boyd and Vandenberghe (2004).

E.4 Dependence of computation time on sample size

For the Ising spin glass example described above and in the text, we measured both the time to evaluate the objective function and the time for the L-BFGS implementation in Section E.2.1 to converge as a function of batch size. As can be seen in Figure E-1, for large batch size the objective function evaluation time is linear, and the convergence time is approximately linear.

Appendix F Additional Ising spin glass comparison

The Ising model has a long and storied history in physics Brush (1967) and machine learning Ackley et al. (1985) and it has recently been found to be a surprisingly useful model for networks of neurons in the retina Schneidman et al. (2006); Shlens et al. (2006). The ability to fit Ising models to the activity of large groups of simultaneously recorded neurons is of current interest given the increasing number of these types of data sets from the retina, cortex and other brain structures.

We fit an Ising model (fully visible Boltzmann machine) of the form

to a set of NN dd-element iid data samples {x(i)∣i=1...N}\left\{x^{(i)}|i=1...N\right\} generated via Gibbs sampling from an Ising model as described below, where each of the dd elements of x\mathbf{x} is either 0 or 1. Because each xi∈{0,1}x_{i}\in\{0,1\}, xi2=xix_{i}^{2}=x_{i}, we can write the energy function as

The probability flow matrix Γ\mathbf{\Gamma} has 2N×2N2^{N}\times 2^{N} elements, but for learning we populate it extremely sparsely, setting

Figure F-1 shows the average error in predicted correlations as a function of learning time for 20,000 samples from a 40 unit, fully connected Ising model. The final absolute correlation error is 0.0058. The JijJ_{ij} used were graciously provided by Broderick and coauthors, and were identical to those used for synthetic data generation in the 2008 paper “Faster solutions of the inverse pairwise Ising problem” Broderick et al. (2007). Training was performed on 20,000 samples so as to match the number of samples used in section III.A. of Broderick et al. Note that given sufficient samples, the minimum probability flow algorithm would converge exactly to the right answer, as learning in the Ising model is convex (see Appendix B), and has its global minimum at the true solution. On an 8 core 2.33 GHz Intel Xeon, the learning converges in about 1515 seconds. Broderick et al. perform a similar learning task on a 100-CPU grid computing cluster, with a convergence time of approximately 200200 seconds.

Appendix G Continuous state space independent component analysis (ICA) Bell AJ (1995) model

Training was performed on 100,000 10×1010\times 10 pixel whitened natural image patches from the van Hateren database Hateren and Schaaf (1998). Minimization was performed by alternating between minimization of the objective function in Equation D-14 and updates to the continuous state space connectivity function g(xj,xi)g\left(\mathbf{x}_{j},\mathbf{x}_{i}\right), as described in more detail in Section H. Both training techniques were initialized with identical isotropic Gaussian noise (with variance 0.010.01, such that each receptive field was initialized to nearly unit length), and trained on the same image patches, which accounts for the similarity of individual filters found by the algorithms.

Appendix H Continuous state space learning with the connectivity function set via Hamiltonian Monte Carlo

Qualitatively, we desire to connect every data state qj∈D\mathbf{q}_{j}\in\mathcal{D} to the non data states qi\mathbf{q}_{i} which will be most informative for learning. The most informative states are those which have high probability under the model distribution p(∞)(q)p^{(\infty)}\left(\mathbf{q}\right). We therefore propose to populate g(qi,qj)g\left(\mathbf{q}_{i},\mathbf{q}_{j}\right) using a Markov transition function for the model distribution. Borrowing techniques from Hamiltonian Monte Carlo Neal (2010) we use Hamiltonian dynamics in our transition function, so as to effectively explore the state space.

In order to implement Hamiltonian dynamics, we first extend the state space to include auxiliary momentum variables.

The initial data and model distributions are p(0)(q)p^{(0)}\left(\mathbf{q}\right) and

and the extended data and model distributions become

The initial (data) distribution over the joint space x\mathbf{x} can be realized by drawing a momentum v\mathbf{v} from a uniform Gaussian distribution for every observation q\mathbf{q} in the dataset D\mathcal{D}.

We connect every state xj\mathbf{x}_{j} to all states which satisfy one of the following 2 criteria,

All states which share the same position qj\mathbf{q}_{j}, with a quadratic falloff in g(xi,xj)g\left(\mathbf{x}_{i},\mathbf{x}_{j}\right) with the momentum difference vi−vj\mathbf{v}_{i}-\mathbf{v}_{j}.

The state which is reached by simulating Hamiltonian dynamics for a fixed time tt on the system described by H(x;θH)H\left(\mathbf{x};\theta_{H}\right), and then negating the momentum. Note that the parameter vector θH\theta_{H} is used only for the Hamiltonian dynamics.

where if x′=HAM⁡(x;θH)\mathbf{x}^{\prime}=\operatornamewithlimits{HAM}\left(\mathbf{x};\theta_{H}\right), then x′\mathbf{x}^{\prime} is the state that results from integrating Hamiltonian dynamics for a time tt and then negating the momentum. Because of the momentum negation, x=HAM⁡(x′;θH)\mathbf{x}=\operatornamewithlimits{HAM}\left(\mathbf{x}^{\prime};\theta_{H}\right), and g(xi,xj)=g(xj,xi)g\left(\mathbf{x}_{i},\mathbf{x}_{j}\right)=g\left(\mathbf{x}_{j},\mathbf{x}_{i}\right).

H.3 Discretizing Hamiltonian dynamics

It is generally impossible to exactly simulate the Hamiltonian dynamics for the system described by H(x;θH)H\left(\mathbf{x};\theta_{H}\right). However, if HAM⁡(x;θH)\operatornamewithlimits{HAM}\left(\mathbf{x};\theta_{H}\right) is set to simulate Hamiltonian dynamics via a series of leapfrog steps, it retains the important properties of reversibility and phase space volume conservation, and can be used in the connectivity function g(xi,xj)g\left(\mathbf{x}_{i},\mathbf{x}_{j}\right) in Equation H-9. In practice, therefore, HAM⁡(x;θH)\operatornamewithlimits{HAM}\left(\mathbf{x};\theta_{H}\right) involves the simulation of Hamiltonian dynamics by a series of leapfrog steps.

H.4 MPF objective function

The MPF objective function for continuous state spaces and a list of observations D\mathcal{D} is

For the connectivity function g(xi,xj)g\left(\mathbf{x}_{i},\mathbf{x}_{j}\right) given in Section H.2, this reduces to

Note that the first term does not depend on the parameters θ\theta, and is thus just a constant offset which can be ignored during optimization. Therefore, we can say

Parameter estimation is performed by finding the parameter vector θ^\hat{\theta} which minimizes the objective function K(θ;D,θH)K\left(\theta;\mathcal{D},\theta_{H}\right),

H.5 Iteratively improving the objective function

The more similar θH\theta_{H} is to θ\theta, the more informative g(xi,xj)g\left(\mathbf{x}_{i},\mathbf{x}_{j}\right) is for learning. If θH\theta_{H} and θ\theta are dissimilar, then many more data samples will be required in D\mathcal{D} to effectively learn. Therefore, we iterate the following procedure, which alternates between finding the θ^\hat{\theta} which minimizes K(θ;D,θH)K\left(\theta;\mathcal{D},\theta_{H}\right), and improving θH\theta_{H} by setting it to θ^\hat{\theta},

Set θ^t+1=argmin⁡θK(θ;D,θHt)\hat{\theta}^{t+1}=\operatornamewithlimits{argmin}_{\theta}K\left(\theta;\mathcal{D},\theta^{t}_{H}\right)

Set θHt+1=θ^t+1\theta^{t+1}_{H}=\hat{\theta}^{t+1}

θ^t\hat{\theta}^{t} then represents a steadily improving estimate for the parameter values which best fit the model distribution p(∞)(q;θ)p^{(\infty)}\left(\mathbf{q};\theta\right) to the data distribution p(0)(q)p^{(0)}\left(\mathbf{q}\right), described by observations D\mathcal{D}. Practically, step 1 above will frequently be truncated early, perhaps after 10 or 100 L-BFGS gradient descent steps.