负二项分布

本文链接：个人站 | 简书 | 优快云
版权声明：除特别声明外，本博客文章均采用 BY-NC-SA 许可协议。转载请注明出处。

之前在介绍 DeepAR 等时间序列预测模型时，为了简单起见，我们使用了大家比较熟悉的正态分布作为示例。在实际应用中，需要根据数据本身的特点选择合适的分布。泊松分布、二项分布、以及负二项分布都可以用来刻画计数类数据。其中，泊松分布的 $\mu ={\sigma }^2$，二项分布的 $\mu \ge {\sigma }^2$，负二项分布的 $\mu \le {\sigma }^2$。在我日常接触的业务场景中，$\mu \le {\sigma }^2$ 较为常见，为此免不了要跟负二项分布打交道。

虽然没什么必要，但是本着「有困难要上，没困难创造困难也要上」的精神，我们还是来推导一下负二项分布的相关公式。

1. 定义

一个成功概率为 $p$ 的伯努利试验，不断重复，直至失败 $r$ 次。此时成功的次数为一个随机变量，用 $X$ 表示。称 $X$ 服从负二项分布，记作 $X\sim NB(r,p)$。

需要注意的是，负二项分布的定义并不唯一。例如 tensorflow_probability 使用的定义与本文一致，而 scipy 则将 $X$ 定义为伯努利试验成功 $r$ 次时的失败次数。使用前一定要先看清楚，别问我怎么知道的。此外，Wikipedia 词条不同段落使用的定义竟然也不完全一致，或许是由不同的人编辑的。

2. 概率质量函数

$X=k$ 时总共进行了 $k+r$ 次试验，最后一次为失败，故前 $k+r-1$ 次试验总共成功了 $k$ 次，失败了 $r-1$ 次。因此
$$f(k;r,p)\equiv Pr(X=k)=\left(\genfrac{}{}{0ex}{}{k+r-1}{k}\right)p^k(1-p{)}^r$$

3. 期望

根据定义
$$\begin{array}{rl}\mathbb{E}X& =\sum _{k=0}^{\infty }kf(k;r,p)\\ & =\sum _{k=1}^{\infty }kf(k;r,p)\\ & =\sum _{k=1}^{\infty }k\frac{(k+r-1)!}{k!(r-1)!}{p}^k(1-p{)}^r\\ & =\frac{rp}{1-p}\sum _{k=1}^{\infty }\frac{[(k-1)+(r+1)-1]!}{(k-1)![(r+1)-1]!}{p}^{k-1}(1-p{)}^{r+1}\\ & =\frac{rp}{1-p}\sum _{k=1}^{\infty }f(k-1;r+1,p)\end{array}$$
令 $k^{\prime }=k-1$、$r^{\prime }=r+1$，显然
$$\sum _{k=1}^{\infty }f(k-1;r+1,p)=\sum _{k^{\prime }=0}^{\infty }f(k^{\prime };r^{\prime },p)=1$$
故
$$\mathbb{E}X=\frac{rp}{1-p}$$

4. 方差

首先计算
$$\begin{array}{rl}\mathbb{E}{X}^2& =\sum _{k=0}^{\infty }{k}^{2}f(k;r,p)\\ & =\frac{rp}{1-p}\sum _{k=1}^{\infty }kf(k-1;r+1,p)\end{array}$$
令 $k^{\prime }=k-1$、$r^{\prime }=r+1$，考虑服从负二项分布的随机变量 $Y\sim NB(r^{\prime },p)$，其概率质量函数为 $f(k^{\prime };r^{\prime },p)$，显然
$$\begin{array}{rl}\sum _{k=1}^{\infty }kf(k-1;r+1,p)& =\sum _{k^{\prime }=0}^{\infty }(k^{\prime }+1)f(k^{\prime };r^{\prime },p)\\ & =\mathbb{E}(Y+1)\\ & =\mathbb{E}Y+1\\ & =\frac{r^{\prime }p}{1-p}+1\\ & =\frac{(r+1)p}{1-p}+1\\ & =\frac{rp+1}{1-p}\end{array}$$
故
$$\mathbb{E}{X}^2=\frac{rp}{1-p}\cdot \frac{rp+1}{1-p}=\frac{r^2{p}^2+rp}{(1-p{)}^2}$$

而根据定义
$$\begin{array}{rl}\mathrm{V}\mathrm{a}\mathrm{r}X& =\mathbb{E}(X-\mathbb{E}X{)}^2\\ & =\mathbb{E}[X^2-2X\mathbb{E}X+(\mathbb{E}X{)}^2]\\ & =\mathbb{E}{X}^2-(\mathbb{E}X{)}^2\\ & =\frac{r^2{p}^2+rp}{(1-p{)}^2}-\frac{r^2{p}^2}{(1-p{)}^2}\\ & =\frac{rp}{(1-p{)}^2}\end{array}$$

我们在文章开头提到，负二项分布的 ${\sigma }^2\ge \mu$。由于 $0\le p\le 1$，这个结论是显而易见的。

5. 累积分布函数

负二项分布的累积分布函数可以表示为正则不完全 Beta 函数：
$$F(k;r,p)=I_{1-p}(r,k+1)$$
证明如下：
$$\begin{array}{rl}F(k;r,p)& \equiv P(X\le k)\\ & =\sum _{x=0}^{k}f(x;r,p)\\ & =\sum _{x=0}^k\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)p^x(1-p{)}^r\end{array}$$
令 $q=1-p$，有
$$F(k;r,p)=F(k;r,1-q)=\sum _{x=0}^k\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)(1-q{)}^x{q}^r$$
对 $q$ 求偏导，得
$$\begin{array}{rl}\frac{\partial F}{\partial q}& =\sum _{x=0}^k\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)\left[-x(1-q{)}^{x-1}{q}^r+r(1-q{)}^x{q}^{r-1}\right]\\ & =\sum _{x=0}^k\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)\left[-x(1-q{)}^{x-1}{q}^r+r(1-q{)}^x{q}^{r-1}\right]\\ & =\sum _{x=0}^k\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)\left[x[(1-q)-1](1-q{)}^{x-1}{q}^{r-1}+r(1-q{)}^x{q}^{r-1}\right]\\ & =\sum _{x=0}^k\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)\left[-x(1-q{)}^{x-1}{q}^{r-1}+(x+r)(1-q{)}^x{q}^{r-1}\right]\\ & =-\sum _{x=0}^{k}x\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)(1-q{)}^{x-1}{q}^{r-1}+\sum _{x=0}^k(x+r)(\genfrac{}{}{0ex}{}{x+r-1}{x})(1-q{)}^x{q}^{r-1}\\ & =-\sum _{x=1}^{k}x\left(\genfrac{}{}{0ex}{}{x+r-1}{x}\right)(1-q{)}^{x-1}{q}^{r-1}+\sum _{x=0}^k(x+r)(\genfrac{}{}{0ex}{}{x+r-1}{x})(1-q{)}^x{q}^{r-1}\\ & =-\sum _{x=1}^k\frac{(x+r-1)!}{(x-1)!(r-1)!}(1-q{)}^{x-1}{q}^{r-1}+\sum _{x=0}^k\frac{(x+r)!}{x!(r-1)!}(1-q{)}^x{q}^{r-1}\\ & =-\frac{r}{q^2}\sum _{x=1}^k\frac{(x+r-1)!}{(x-1)!r!}(1-q{)}^{x-1}{q}^{r+1}+\frac{r}{q^2}\sum _{x=0}^k\frac{(x+r)!}{x!r!}(1-q{)}^x{q}^{r+1}\\ & =-\frac{r}{q^2}\sum _{x^{\prime }=0}^{k-1}\frac{(x^{\prime }+r)!}{x^{\prime }!r!}(1-q{)}^{x^{\prime }}{q}^{r+1}+\frac{r}{q^2}\sum _{x=0}^k\frac{(x+r)!}{x!r!}(1-q{)}^x{q}^{r+1}\\ & =-\frac{r}{q^2}F(k-1;r+1,1-q)+\frac{r}{q^2}F(k;r+1,1-q)\\ & =\frac{r}{q^2}f(k;r+1,1-q)\end{array}$$
而根据正则不完全 Beta 函数的定义，有
$$\begin{array}{rl}{I}_{1-p}(r,k+1)& =I_q(r,k+1)\\ & =\frac{B(q;r,k+1)}{B(r,k+1)}\\ & =\frac{{\int }_0^q{t}^{r-1}(1-t{)}^k\mathrm{d}t}{B(r,k+1)}\end{array}$$
同样对 $q$ 求偏导，得
$$\begin{array}{rl}\frac{\partial I_q}{\partial q}& =\frac{q^{r-1}(1-q{)}^k}{B(r,k+1)}\\ & =\frac{\Gamma (r+k+1)}{\Gamma (r)\Gamma (k+1)}{q}^{r-1}(1-q{)}^k\\ & =\frac{(r+k)!}{(r-1)!k!}{q}^{r-1}(1-q{)}^k\\ & =\frac{r}{q^2}\frac{(r+k)!}{r!k!}{q}^{r+1}(1-q{)}^k\\ & =\frac{r}{q^2}f(k;r+1,1-q)\end{array}$$
也就是说
$$\frac{\partial }{\partial q}F(k;r;1-q)=\frac{\partial }{\partial q}{I}_q(r,k+1)$$
亦即
$$F(k;r;1-q)=I_q(r,k+1)+C$$
注意到 $q=0$ 时有
$$\{\begin{array}{l}F(k;r,1)=0\\ I_0(r,k+1)=0\end{array}$$
解得常数 $C=0$。

证毕。

6. 在时间序列预测模型中的使用

DeepAR 等模型中，网络的输出目标是概率分布的参数。例如正态分布的 $\mu$ 和 $\sigma$。但对于负二项分布而言，让网络直接输出 $\mu$ 和 $\sigma$ 是不合适的，因为在训练过程中很难保证输出的值满足 ${\sigma }^2\ge \mu$。那么让网络输出 $r$ 和 $p$ 呢？似乎是可以的，只要保证 $r>0$，$0\le p\le 1$ 即可。前者可以使用 softplus 激活函数，后者可以使用 sigmoid 激活函数。有没有办法避免使用 sigmoid 呢？通常的做法是让网络输出 $\mu$ 和 $\alpha =1/r$，只要使用 softplus 激活函数确保二者均为正数即可。

参考文献

1. Negative binomial distribution - Wikipedia
2. Beta function - Wikipedia
3. Patil G P. On the evaluation of the negative binomial distribution with examples[J]. Technometrics, 1960, 2(4): 501-505.