gprob note

Probabilistic programs define statistical models as compositions of primitive distributions conditioned on observations. Inference of general conditional distributions is a hard task. The computational overhead of general inference methods, such as Markov chain or Hamiltonian Monte-Carlo, can make it difficult for probabilistic programs to compete with problem-specific approaches. For Gaussian random variables, however, this is not necessarily the case—as the composition of distributions and their conditioning can be performed analytically, the overhead for their computation can be minimal.

This note presents the background theory for gprob, a probabilistic programming package that specializes to Gaussian random variables.

1 Distributions defined via linear maps.

Any Gaussian probabilistic program [1] can be expressed as a linear map of a number of independent identically distributed (iid) normal variables $e_i$ with zero mean and unit variance,

$$e_i\sim 𝒩(0,1),i=1,\mathrm{\dots },n.$$ (1)

In gprob, all random arrays are represented in this way — as affine maps of latent variables. E.g., a random array named $x$ indexed by a multi-dimensional index $ijk\mathrm{\dots }$ is represented as

$$x_{ijk\mathrm{\dots }}=\sum _{\lambda }{e}_{\lambda }{a}_{\lambda ,ijk\mathrm{\dots }}+b_{ijk\mathrm{\dots }},$$ (2)

where $\{e_{\lambda }\sim 𝒩(0,1)\}$, $\lambda \in \mathbb{N}$ is a set of latent univariate normal variables, and $a$ and $b$ are numerical arrays with real or complex coefficients. The mean of $x_{ijk\mathrm{\dots }}$ equals $b_{ijk\mathrm{\dots }}$, and the covariance is

$$C_{ijk\mathrm{\dots },rst\mathrm{\dots }}=\sum _{\lambda }{a}_{\lambda ,ijk\mathrm{\dots }}{a}_{\lambda ,rst\mathrm{\dots }}^{\ast },$$ (3)

where $rst\mathrm{\dots }$ run over the same set of indices as $ijk\mathrm{\dots }$ and ^∗ denotes complex conjugation.

New latent variables are allocated during calls to normal(). For example, normal(0, 1) produces one latent variable, and normal(size=sz), where sz is integer, produces sz new latent variables. A global inventory of all latent variables is maintained, but their joint distributions are only materialized within individual arrays. This has an effect on the memory footprint. Roughly speaking, if one random variable takes $m$ bytes, then two arrays of $n$ variables each will take $m\times (2\times n^2)$ bytes, and not $m\times {(2\times n)}^2$ bytes.

Linear operations in gprob, such as addition or concatenation, produce new random arrays. The set of the latent variables of the result is a union of the sets of the latent variables of the operands. The $b$ of the operation result is simply the operation applied to the $b$s of the operands. The $a$ of the operation result is found by first extending the $a$s of the operands along their first dimension by zero entries, such that they all map the same union vector of latent variables, and then performing the operation iteratively over the first dimension of the extended $a$s. (In practice, no iteration takes place, but the operations on $a$s are vectorized.) To give a concrete example, the sum of two one-dimensional random variables $𝐱$ and $𝐲$ defined by the maps written in the matrix form as

(4)

is another variable, $𝐳=𝐱+𝐲$, whose map is

(5)

2 Conditioning.

In a probabilistic language that maintains a joint distribution over all latent variables, such as the one presented in Ref. [1], conditioning is a statement that updates the latent distribution. One can think of the conditioning in this case as a post-selection to the realizations of the latent variables that ensures that all the conditions are fulfilled. In contrast, in gprob, conditioning of $x$ on $y=0$ is an expression that produces a new random variable, whose latent space is the union of the latent spaces of $x$ and $y$. In this case, one can think of the conditioning as creating a version $x$ that is feed-forward corrected based on the measurement of $y$.

To show how the conditioning is implemented, we consider the one-dimensional real case. Any random variable can be transformed to this representation by flattening and, if the variable is complex, by separating its real and imaginary parts. Let the conditioned variable be $𝐱$, the condition be $𝐲=\mathrm{𝟎}$, and the result of conditioning be $𝐳\equiv 𝐱|(𝐲=\mathrm{𝟎})$. All more general cases can be reduced to this: if a non-zero value of $𝐲$ was observed, this value can be subtracted from $𝐲$, and if there is more than one observation, they can be stacked together in one $𝐲$ vector.

The variables $𝐱$ and $𝐲$ take values in ${ℝ}^{d_x}$ and ${ℝ}^{d_y}$, and are represented by the maps of the same vector of latent variables $𝐞=(e_1,\mathrm{\dots },e_n)$,

$$𝐱={𝐀}_x^T𝐞+{𝐛}_x,𝐲={𝐀}_y^T𝐞+{𝐛}_y,$$ (6)

where ${𝐀}_x$ and ${𝐀}_y$ are, respectively, $n\times d_x$ and $n\times d_y$ matrices, and ${𝐛}_x$ and ${𝐛}_y$ are $d_x$ and $d_y$-size mean vectors. The rank of the condition matrix, ${𝐀}_y$, is assumed to be full and equal to $d_y\le n$, which means that the conditions are neither incompatible nor tautological. If the rank of ${𝐀}_y$ is less than $d_y$, yet the conditions are compatible, some of them can be eliminated from ${𝐀}_y$.

Conditioning splits the latent space into two orthogonal subspaces—one in which the allowed values of the latent variables are fully determined to yield the observed value of $𝐲$, and one in which the values of the latent variables are unconstrained. The conditional mean is found by evaluating $𝐱$ at the value of the latent vector $\widetilde{𝐞}$ from the first subspace, which satisfies

$${𝐀}_y^T\widetilde{𝐞}+{𝐛}_y=0.$$ (7)

This value can be formally found as

$$\widetilde{𝐞}=-{\left({𝐀}_y^{+}\right)}^T{𝐛}_y,$$ (8)

where ${𝐀}_y^{+}$ is the Moore–Penrose pseudoinverse of ${𝐀}_y$. (As ${𝐀}_y$ is normally not a square matrix, it does not have an inverse). To obtain $\widetilde{𝐞}$ in practice, it is more convenient to directly solve the system of equations in Eq. (7) in the least square sense. The mean of the conditional variable $𝐳$, ${𝐛}_z$, is then given by

$${𝐛}_z={𝐛}_x-{𝐀}_x^T{\left({𝐀}_y^{+}\right)}^T{𝐛}_y.$$ (9)

The conditional map, ${𝐀}_z$, is found by left-multiplying ${𝐀}_x$ by the projector on the latent subspace unaffected by the conditioning, i.e. the subspace orthogonal to the columns of the constraint matrix ${𝐀}_y$. The operator that performs the desired projection is $𝐈-{𝐀}_y{𝐀}_y^{+}$, where $𝐈$ is the identity matrix, and the conditional map is

$${𝐀}_z={𝐀}_x-{𝐀}_y{𝐀}_y^{+}{𝐀}_x.$$ (10)

One can verify that the conditional mean and covariance that follow from Eqs. (9-10) agree with those obtained based on the momentum representation of $𝐱$ and $𝐲$ (given in Appendix A). To make this comparison, we can use the facts that the cross-covariance matrix between $𝐱$ and $𝐲$ is

$${𝐂}_{xy}={𝐀}_x^T{𝐀}_y,$$ (11)

that $𝐏={𝐀}_y{𝐀}_y^{+}$ is an orthogonal projector satisfying ${𝐏}^2=𝐏$ and ${𝐏}^T=𝐏$, and also that for a full-rank matrix ${𝐀}_y$ the block ${𝐂}_{yy}={𝐀}_y^T{𝐀}_y$ is invertible, with its inverse being

$${𝐂}_{yy}^{-1}={𝐀}_y^{+}{\left({𝐀}_y^{+}\right)}^T.$$ (12)

3 Causal conditioning of time series. Wiener filtering.

Sec. 2 described the conditioning of all elements of one random vector $𝐱$ on all elements of another vector $𝐲$. There is also another, more specialized, problem that arises sometimes—the conditioning of individual elements $x_i$, for $i=1,..,d_x$, on subsets $\{y_j:j\le s_i\}$ for some non-decreasing set of integer indices $s_i$. This problem appears, in particular, in the causal analysis of time series via Kalman and Wiener filtering, where the elements of $𝐱$ and $𝐲$ correspond to a range of sequential moments in time. (For an introduction to filtering see, for example, Ref. [2]). The time stamps on $𝐱$ and $𝐲$ are non-decreasing, but not necessarily equally separated, and can also repeat. The task can be for each $x_i$ to find its distribution conditioned only on those observations $\{y_j\}$ that came before $x_i$. We, therefore, call the corresponding problem “causal conditioning”.

Causal conditioning can be performed trivially using the result of Sec. 2 by iterating over the elements of $𝐱$ and each time selecting an appropriate range from $𝐲$. This would be, roughly speaking, $d_x$ time more computationally expensive than the simple conditioning of all $x_i$ on all $y_j$. Fortunately, there is an algorithm that allows to causally condition all elements of $𝐱$ in one go.

To introduce this algorithm, we first note that in the case of regular conditioning both the conditional mean and map (given by Eq. (9) and Eq. (10), respectively), could be expressed using a $d_y\times d_x$ matrix $𝐁$,

$$𝐁={𝐀}_y^{+}{𝐀}_x$$ (13)

as

$${𝐛}_z^T={𝐛}_x^T-{𝐛}_y^T𝐁,{𝐀}_z={𝐀}_x-{𝐀}_y𝐁.$$ (14)

The $i$-th column of $𝐁$ consists of the coefficients that give the closest in Euclidean distance approximation of the $i$-th column of ${𝐀}_x$ via a linear combination of the columns of ${𝐀}_y$. From this point of view, it should be clear that for vectors whose components are indexed left to right, $𝐱=(x_1,\mathrm{\dots },x_{d_x})$ and $𝐲=(y_1,\mathrm{\dots },y_{d_y})$, a causal structure can be enforced in conditioning by constraining the closest-approximation matrix $𝐁$ to be upper-triangular,

(15)

Here, the list of indices $s_i$ is a non-decreasing sequence of integers from the range $[1,d_y]$ that is tailored to the causal structure of the problem. For example, in causal Wiener filtering of a scalar random process $x(t)$ with a scalar noisy measurement $y(t)$, discretized as $x_i$ and $y_i$, we could put $s_i=i-1$ to use all observations before the time $t_i$ for the prediction of $x(t_i)$.

What remains is to find a way of obtaining causally-constrained matrices $𝐁$. Towards this goal, we introduce the QR decomposition of ${𝐀}_y$,

$${𝐀}_y=\mathrm{𝐐𝐑},$$ (16)

where $𝐐$ is a $n\times d_y$ matrix whose columns is an orthonormal set of vectors, and $𝐑$ is a $d_y\times d_y$ non-degenerate upper triangular matrix. The left-multiplication of $𝐁$ by $𝐑$ preserves the triangular structure (meaning that the elements of $𝐁$ that are zero in Eq. (15) are also zero in the multiplication product), and, therefore, the least-square solution of the system of equations

$${𝐀}_y𝐁={𝐀}_x,$$ (17)

with the desired causal constraint on $𝐁$ can be expressed as

$$𝐁={𝐑}^{-1}ℳ\left[{𝐐}^T{𝐀}_x\right],$$ (18)

where $ℳ[\mathrm{\dots }]$ a masking operator that zeros for its argument $\mathrm{\dots }$ the matrix elements that are zero in Eq. (15). Finally, the causally conditioned mean and map are given by

${𝐛}_z^T={𝐛}_x^T-{𝐛}_y^T{𝐑}^{-1}𝐐ℳ\left[{𝐐}^{†}{𝐀}_x\right],$ ${𝐀}_z={𝐀}_x-𝐐ℳ\left[{𝐐}^{†}{𝐀}_x\right].$ (19)

A similar analysis can be performed when $𝐁$ is generalized lower-triangular, which in the Wiener filtering theory corresponds to anti-causal filtering. The main difference in this case is that ${𝐀}_y$ needs to be decomposed as ${𝐐}^{\prime }𝐋$, where $𝐋$ is lower-triangular.

Appendix A: Gaussian distributions defined via moments.

The probability density function of a $d$-dimensional Gaussian random variable with the mean vector $𝐦=(m_1,\mathrm{\dots },m_d)$ and covariance matrix $𝐂$ to take the value $𝐱\in {ℝ}^d$ is given by

$$p(𝐱)={(2\pi )}^{-d/2}{\left(\det (𝐂)\right)}^{-1/2}\exp \left(-\frac{1}{2}{(𝐱-𝐦)}^T{𝐂}^{-1}(𝐱-𝐦)\right),$$ (A1)

where ${𝐂}^{-1}$ is the inverse of the covariance. If $𝐂$ is degenerate and thus non-invertible, the possible realizations of the random variable do not cover the whole space ${ℝ}^d$. Our convention when defining the probability density function for such variables is to assign $p(𝐱)=0$ to impossible samples, and, when dealing with possible samples, to use the Moore–Penrose pseudoinverse ${𝐂}^{+}$ instead of ${𝐂}^{-1}$ and the determinant of the non-degenerate projection of $𝐂$ instead of $\det (𝐂)$ in Eq. (A1). This is consistent with the way degenerate distributions are handled in scipy.stats [3].

If the random variable consists of two stacked vectors, $𝐱$ with the dimension $d_x$, and $𝐲$ with the dimension $d_y$, and $𝐂$ is their joint covariance matrix expressed in the block form as,

(A2)

the conditional probability density of $𝐱$ given $𝐲=\widetilde{𝐲}$ is Gaussian with the mean and covariance given by

${𝐦}_{\mathrm{c}}={𝐦}_x+{𝐂}_{xy}{𝐂}_{yy}^{-1}(\widetilde{𝐲}-{𝐦}_y),$ (A3)

${𝐂}_{\mathrm{c}}={𝐂}_{xx}-{𝐂}_{xy}{𝐂}_{yy}^{-1}{𝐂}_{yx}.$ (A4)

Appendix B: Parametric derivatives and Fisher information.

A number of information-theoretic properties of distributions whose mean and covariance depend on some parameters $\theta$ are expressed via the logarithm of the probability density and its derivatives. For non-degenerate Gaussian distributions, they are given, respectively, by

$$\ln p(𝐱)=-\frac{1}{2}\mathrm{\Delta }{𝐱}^T{𝐂}^{-1}\mathrm{\Delta }𝐱-\frac{1}{2}\ln \det (𝐂),$$ (B1)

and

$$\frac{\partial \ln p(𝐱)}{\partial {\theta }_i}=\frac{\partial {𝐦}^T}{\partial {\theta }_i}{𝐂}^{-1}\mathrm{\Delta }𝐱+\frac{1}{2}\mathrm{\Delta }{𝐱}^T{𝐂}^{-1}\frac{\partial 𝐂}{\partial {\theta }_i}{𝐂}^{-1}\mathrm{\Delta }𝐱-\frac{1}{2}\mathrm{Tr}\left[{𝐂}^{-1}\frac{\partial 𝐂}{\partial {\theta }_i}\right],$$ (B2)

where

$$\mathrm{\Delta }𝐱=𝐱-𝐦.$$ (B3)

The expression for the derivatives was obtained using $\partial {𝐂}^{-1}/\partial \theta =-{𝐂}^{-1}(\partial 𝐂/\partial \theta ){𝐂}^{-1}$, and the fact that the matrix determinant equals the product of the eigenvalues. The components of the Fisher information matrix $F_{ij}$ are given by (e.g. [4])

$$F_{ij}=\frac{\partial {𝐦}^T}{\partial {\theta }_i}{𝐂}^{-1}\frac{\partial 𝐦}{\partial {\theta }_j}+\frac{1}{2}\mathrm{Tr}\left[{𝐂}^{-1}\frac{\partial 𝐂}{\partial {\theta }_i}{𝐂}^{-1}\frac{\partial 𝐂}{\partial {\theta }_j}\right],$$ (B4)

and related to the log probability density as

$$F_{ij}=⟨\frac{\partial \ln p(𝐱)}{\partial {\theta }_i}\frac{\partial \ln p(𝐱)}{\partial {\theta }_j}⟩=-⟨\frac{{\partial }^2\ln p(𝐱)}{\partial {\theta }_i\partial {\theta }_j}⟩.$$ (B5)

For completeness, here is also the expression for the second derivatives of the log probability density,

${\displaystyle \frac{{\partial }^2\ln p(𝐱)}{\partial {\theta }_i\partial {\theta }_j}}=$ $-F_{ij}+{\displaystyle \frac{{\partial }^2{𝐦}^T}{\partial {\theta }_i\partial {\theta }_j}}{𝐂}^{-1}\mathrm{\Delta }𝐱$

$+\mathrm{Tr}\left[\left({\displaystyle \frac{1}{2}}{𝐂}^{-1}{\displaystyle \frac{{\partial }^2𝐂}{\partial {\theta }_i\partial {\theta }_j}}{𝐂}^{-1}+{𝐂}^{-1}{\displaystyle \frac{\partial 𝐂}{\partial {\theta }_i}}{\displaystyle \frac{\partial {𝐂}^{-1}}{\partial {\theta }_j}}\right)\left(\mathrm{\Delta }𝐱\mathrm{\Delta }{𝐱}^T-𝐂\right)\right].$ (B6)