摘要：Tensorflow Distributions提供了两类抽象：distributions和bijectors。distributions提供了一系列具备快速、数值稳定的采样、对数概率计算以及其他统计特征计算方法的概率分布。bijectors提供了一系列针对distribution的可组合的确定性变换。
1、Distributions
1.1 methods
一个distribution至少实现以下方法：sample、log_prob、batch_shape_tensor、event_shape_tensor；同时也实现了一些其他方法，例如：cdf、survival_function、quantile、mean、variance、entropy等；Distribution基类实现了给定log_prob计算prob、给定log_cdf计算log_survival_fn的方法。
1.2 shape semantics
将一个tensor的形状分为三个部分：sample shape、batch shape、event shape。
sample shape：描述从给定概率分布上独立同分布的采样形状；
batch shape：描述从概率分布上独立、非同分布的采样形状，也即，我们可以指定一组参数不同的相同分布，batch shape通常用来为机器学习中一个batch的样本每个样本指定一个分布；
event shape：描述从概率分布上单次采样的形状；
1.3 sampling
reparameterization：distributions拥有一个reparameterization属性，这个属性表明了自动化微分和采样之间的关系。目前包括两种：“fully reparameterized” 和 “not reparameterized”。
fully reparameterized：例如，对于分布dist = Normal(loc, scale)，采样y = dist.sample()的内部过程为x = tf.random_normal([]); y = scale * x + loc. 样本y是reparameterized的，因为它是参数loc、scale及无参数样本x的光滑函数。
not reparameterized：例如，gamma分布使用接收-拒绝的方式进行采样，是参数的非光滑函数。
end to end automatic differentiation：通过与tensorflow结合，一个fully reparameterized的分布可以进行端到端的自动微分。例如，要最小化分布Y的期望损失E [φ(Y)]，可以使用蒙特卡洛近似的方法最小化
这使得我们可以使用SN作为期望损失的估计，还可以使用ΔλSN作为梯度ΔλE [φ(Y)]的估计，其中λ是分布Y的参数。
1.4 high order distributions
TransformedDistribution：对一个基分布执行一个可逆可微分转换即可得到一个TransformedDistribution。例如，可以从一个Exponential分布得到一个标准Gumbel分布：
standard_gumbel = tfd.TransformedDistribution( distribution=tfd.Exponential(rate=1.), bijector=tfb.Chain([ tfb.Affine( scale_identity_multiplier=-1., event_ndims=0), tfb.Invert(tfb.Exp()), ])) standard_gumbel.batch_shape # ==> [] standard_gumbel.event_shape # ==> []
基于gumbel分布，可以构建一个Gumbel-Softmax(Concrete)分布：
alpha = tf.stack([ tf.fill([28 * 28], 2.), tf.ones(28 * 28)]) concrete_pixel = tfd.TransformedDistribution( distribution=standard_gumbel, bijector=tfb.Chain([ tfb.Sigmoid(), tfb.Affine(shift=tf.log(alpha)), ]), batch_shape=[2, 28 * 28]) concrete_pixel.batch_shape # ==> [2, 784] concrete_pixel.event_shape # ==> []
Independent：对batch shape和event shape进行转换。例如：
image_dist = tfd.TransformedDistribution( distribution=tfd.Independent(concrete_pixel), bijector=tfb.Reshape( event_shape_out=[28, 28, 1], event_shape_in=[28 * 28])) image_dist.batch_shape # ==> [2] image_dist.event_shape # ==> [28, 28, 1]
Mixture：定义了由若干分布组合成的新的分布，例如：
image_mixture = tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical( probs=[0.2, 0.8]), components_distribution=image_dist) image_mixture.batch_shape # ==> [] image_mixture.event_shape # ==> [28, 28, 1]
1.5 distribution functionals
functional以一个分布作为输入，输出一个标量，例如：entropy、cross entropy、mutual information、kl距离等。
p = tfd.Normal(loc=0., scale=1.) q = tfd.Normal(loc=-1., scale=2.) xent = p.cross_entropy(q) kl = p.kl_divergence(q) # ==> xent - p.entropy()
2、Bijectors
2.1 definition
Bijector API提供了针对distribution的可微分双向映射（differentialble, bijective map, diffeomorphism）转换接口。给定随机变量X和一个diffeomorphism F，可以定义一个新的随机变量Y，Y的密度可由下式计算：
其中DF-1是F的Jacobian的逆。（参考：https://zhuanlan.zhihu.com/p/100287713）
每个bijector子类都对应一个F，TransformedDistribution自动计算Y=F(X)的密度。bijector使得我们可以利用已有的分布构建许多其他分布。
bijector主要包含以下三个函数：
forward：实现x → F (x)，TransformedDistribution.sample函数使用该函数将一个tensor转换为另一个tensor；
inverse：forward的逆变换，实现y → F-1(y)，TransformedDistribution.log_prob使用该函数计算对数概率（上式）；
inverse_log_det_jacobian：计算log |DF−1(y)|，TransformedDistribution.log_prob使用该函数计算对数概率（上式）；
通过使用bijectors，TransformedDistribution可以自动高效地实现sample、log_prob、prob，对于具有恒定Jacobian的bijector，TransformedDistribution自动实现一些基础统计量，如mean、variance、entropy等。
以下实现了对Laplace的放射变换：
vector_laplace = tfd.TransformedDistribution( distribution=tfd.Laplace(loc=0., scale=1.), bijector=tfb.Affine( shift=tf.Variable(tf.zeros(d)), scale_tril=tfd.fill_triangular( tf.Variable(tf.ones(d * (d + 1) / 2)))), event_shape=[d])
由于tf.Variables，该分布是可学习的。
2.2 composability
bijectors可以构成高阶bijectors，例如Chain、Invert。
chain bijector可以构建一系列丰富的分布，例如创建一个多变量logit-Normal分布：
matrix_logit_mvn = tfd.TransformedDistribution( distribution=tfd.Normal(0., 1.), bijector=tfb.Chain([ tfb.Reshape([d, d]), tfb.SoftmaxCentered(), tfb.Affine(scale_diag=diag), ]), event_shape=[d * d])
Invert可以通过交换inverse和forward函数，高效地将bijectors数量翻倍，例如：
softminus_gamma = tfd.TransformedDistribution( distribution=tfd.Gamma( concentration=alpha, rate=beta), bijector=tfb.Invert(tfb.Softplus()))
2.3 caching
bijector自动缓存操作的输入输出对，包括log det jacobian。caching的意义时，当inverse计算很慢或数值不稳定或难以实现时，可以高效的执行inverse操作。当计算采样结果的概率是，缓存被触发。如果q(x)是x=f(ε)的密度，且ε~r，那么caching可以降低计算q(xi)的计算成本：
caching机制也可用来进行高效地重要性采样（importance sampling）：
3、 应用
3.1 核密度估计（KDE）
例如，可以通过以下代码构建一个由n个mvn_diag分布作为kernel的混合高斯模型，其中每个kernel的权重为1/n。注意，此时Independent会对分布的shape进行重定义（reinterpret），tfd.Normal(loc=x, scale=1.)创建了一个batch_shape = n*d, event_shape = []的分布，对其Independent之后，变为batch_shape = n, event_shape = d的分布。
Independent文档：https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/Independent?hl=zh-cn
f = lambda x: tfd.Independent(tfd.Normal( loc=x, scale=1.)) n = x.shape[0].value kde = tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical( probs=[1 / n] * n), components_distribution=f(x))
3.2 变分自编码器（VAE）
论文：https://arxiv.org/pdf/1312.6114.pdf
博客：https://spaces.ac.cn/archives/5253
def make_encoder(x, z_size=8): net = make_nn(x, z_size * 2) return tfd.MultivariateNormalDiag( loc=net[..., :z_size], scale=tf.nn.softplus(net[..., z_size:]))) def make_decoder(z, x_shape=(28, 28, 1)): net = make_nn(z, tf.reduce_prod(x_shape)) logits = tf.reshape( net, tf.concat([[-1], x_shape], axis=0)) return tfd.Independent(tfd.Bernoulli(logits)) def make_prior(z_size=8, dtype=tf.float32): return tfd.MultivariateNormalDiag( loc=tf.zeros(z_size, dtype))) def make_nn(x, out_size, hidden_size=(128, 64)): net = tf.flatten(x) for h in hidden_size: net = tf.layers.dense( net, h, activation=tf.nn.relu) return tf.layers.dense(net, out_size)
3.3 Edward概率编程
tfd是Edward的后端。以下代码实现一个随机循环神经网络（stochastic rnn），其隐藏状态是随机的。
stochastic rnn论文：https://arxiv.org/pdf/1411.7610.pdf
from edward.models import Normal z = x = [] z[0] = Normal(loc=tf.zeros(K), scale=tf.ones(K)) h = tf.layers.dense( z[0], 512, activation=tf.nn.relu) loc = tf.layers.dense(h, D, activation=None) x[0] = Normal(loc=loc, scale=0.5) for t in range(1, T): inputs = tf.concat([z[t - 1], x[t - 1]], 0) loc = tf.layers.dense( inputs, K, activation=tf.tanh) z[t] = Normal(loc=loc, scale=0.1) h = tf.layers.dense( z[t], 512, activation=tf.nn.relu) loc = tf.layers.dense(h, D, activation=None) x[t] = Normal(loc=loc, scale=0.5)