深入学习 torch.distributions

0. 引言

前几天分几篇博文精细地讲述了《von Mises-Fisher 分布》, 以及相应的 PyTorch 实现《von Mises-Fisher Distribution (代码解析)》, 其中以 Uniform 分布为例简要介绍了 torch.distributions 包的用法. 本以为已经可以了, 但这两天看到论文 The Power Spherical distribution 的代码, 又被其实现分布的方式所吸引.

Power Spherical 分布与 von Mises Fisher 分布类似, 只不过将后者概率密度函数中的指数函数换成了多项式函数: $$\begin{array}{rl}{f}_p(\mathbf{x};\mathbf{\mu },\kappa )& \propto exp(\kappa {\mathbf{\mu }}^{⊺}\mathbf{x})\\ & \Downarrow \\ f_p(\mathbf{x};\mathbf{\mu },\kappa )& \propto (1+{\mathbf{\mu }}^{⊺}\mathbf{x}{)}^{\kappa }\end{array}$$ 采样框架基本一致, 且这么做可以使边缘 $t$ 的线性变换 $\frac{t+1}{2}\sim Beta(\frac{p-1}{2}+\kappa ,\frac{p-1}{2})$, 从而避免了接受-拒绝采样过程.

当然, 按照之前的 VonMisesFisher 的写法, 这个 t 的采样大概是这样:

z = beta.sample(sample_shape)
t = 2 * z - 1

但现在我遇到了这种写法:

class MarginalTDistribution(tds.TransformedDistribution):arg_constraints = {'dim': constraints.positive_integer,'scale': constraints.positive,}has_rsample = Truedef __init__(self, dim, scale, validate_args=None):self.dim = dimself.scale = scalesuper().__init__(tds.Beta(  # 用 Beta 分布转换, z 服从 Beta(α+κ,β)(dim - 1) / 2 + scale, (dim - 1) / 2, validate_args=validate_args),transforms=tds.AffineTransform(loc=-1, scale=2),  # t=2z-1 是想要的边缘分布随机数)

然后就可以进行对 $t$ 的采样了.

上述代码的解构图; 浅蓝色代表抽象基类, 绿色代表实类; 虚线代表继承, 实线代表参数输入

我们可以看到其基本架构, 本文将详细解析其内部的具体细节.