负二项分布

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

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

1. 定义

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

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

2. 概率质量函数

X=kX=k 时总共进行了 k+rk+r 次试验,最后一次为失败,故前 k+r1k+r-1 次试验总共成功了 kk 次,失败了 r1r-1 次。因此

f(k;r,p)Pr(X=k)=(k+r1k)pk(1p)rf(k; r, p)\equiv Pr(X=k)=\tbinom{k+r-1}{k}p^k(1-p)^r

3. 期望

根据定义

EX=k=0kf(k;r,p)=k=1kf(k;r,p)=k=1k(k+r1)!k!(r1)!pk(1p)r=rp1pk=1[(k1)+(r+1)1]!(k1)![(r+1)1]!pk1(1p)r+1=rp1pk=1f(k1;r+1,p)\begin{aligned} \mathbb{E}X &=\sum\limits_{k=0}^{\infty}kf(k;r,p)\\ &=\sum\limits_{k=1}^{\infty}kf(k;r,p)\\ &=\sum\limits_{k=1}^{\infty}k\frac{(k+r-1)!}{k!(r-1)!}p^k(1-p)^r\\ &=\frac{rp}{1-p} \sum\limits_{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\limits_{k=1}^{\infty}f(k-1;r+1,p) \end{aligned}

k=k1k'=k-1r=r+1r'=r+1,显然

k=1f(k1;r+1,p)=k=0f(k;r,p)=1\sum\limits_{k=1}^{\infty}f(k-1;r+1,p)=\sum\limits_{k'=0}^{\infty}f(k';r',p)=1

EX=rp1p\mathbb{E}X = \frac{rp}{1-p}

4. 方差

首先计算

EX2=k=0k2f(k;r,p)=rp1pk=1kf(k1;r+1,p)\begin{aligned} \mathbb{E}X^2 &=\sum\limits_{k=0}^{\infty}k^2f(k;r,p)\\ &=\frac{rp}{1-p}\sum\limits_{k=1}^{\infty}kf(k-1;r+1,p) \end{aligned}

k=k1k'=k-1r=r+1r'=r+1,考虑服从负二项分布的随机变量 YNB(r,p)Y\sim NB(r', p),其概率质量函数为 f(k;r,p)f(k';r',p),显然

k=1kf(k1;r+1,p)=k=0(k+1)f(k;r,p)=E(Y+1)=EY+1=rp1p+1=(r+1)p1p+1=rp+11p\begin{aligned} \sum\limits_{k=1}^{\infty}k f(k-1;r+1,p) &= \sum\limits_{k'=0}^{\infty}(k'+1)f(k';r',p)\\ &= \mathbb{E}(Y+1)\\ &= \mathbb{E}Y + 1\\ &= \frac{r'p}{1-p} + 1\\ &= \frac{(r+1)p}{1-p} + 1\\ &= \frac{rp+1}{1-p} \end{aligned}

EX2=rp1prp+11p=r2p2+rp(1p)2\mathbb{E}X^2 = \frac{rp}{1-p}\cdot \frac{rp+1}{1-p}= \frac{r^2p^2+rp}{(1-p)^2}

而根据定义

VarX=E(XEX)2=E[X22XEX+(EX)2]=EX2(EX)2=r2p2+rp(1p)2r2p2(1p)2=rp(1p)2\begin{aligned} \mathrm{Var}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^2p^2+rp}{(1-p)^2} - \frac{r^2p^2}{(1-p)^2}\\ &= \frac{rp}{(1-p)^2} \end{aligned}

我们在文章开头提到,负二项分布的 σ2μ\sigma^2\geq\mu。由于 0p10\leq p\leq1,这个结论是显而易见的。

5. 累积分布函数

负二项分布的累积分布函数可以表示为正则不完全 Beta 函数:

F(k;r,p)=I1p(r,k+1)F(k;r,p)=I_{1-p}(r, k+1)

证明如下:

F(k;r,p)P(Xk)=x=0kf(x;r,p)=x=0k(x+r1x)px(1p)r\begin{aligned} F(k;r,p) &\equiv P(X\leq k)\\ &=\sum_{x=0}^kf(x;r,p)\\ &=\sum_{x=0}^k \tbinom{x+r-1}{x}p^x(1-p)^r \end{aligned}

q=1pq=1-p,有

F(k;r,p)=F(k;r,1q)=x=0k(x+r1x)(1q)xqrF(k;r,p) = F(k;r,1-q) = \sum_{x=0}^k \tbinom{x+r-1}{x}(1-q)^xq^r

qq 求偏导,得

Fq=x=0k(x+r1x)[x(1q)x1qr+r(1q)xqr1]=x=0k(x+r1x)[x(1q)x1qr+r(1q)xqr1]=x=0k(x+r1x)[x[(1q)1](1q)x1qr1+r(1q)xqr1]=x=0k(x+r1x)[x(1q)x1qr1+(x+r)(1q)xqr1]=x=0kx(x+r1x)(1q)x1qr1+x=0k(x+r)(x+r1x)(1q)xqr1=x=1kx(x+r1x)(1q)x1qr1+x=0k(x+r)(x+r1x)(1q)xqr1=x=1k(x+r1)!(x1)!(r1)!(1q)x1qr1+x=0k(x+r)!x!(r1)!(1q)xqr1=rq2x=1k(x+r1)!(x1)!r!(1q)x1qr+1+rq2x=0k(x+r)!x!r!(1q)xqr+1=rq2x=0k1(x+r)!x!r!(1q)xqr+1+rq2x=0k(x+r)!x!r!(1q)xqr+1=rq2F(k1;r+1,1q)+rq2F(k;r+1,1q)=rq2f(k;r+1,1q)\begin{aligned} \frac{\partial F}{\partial q} & = \sum_{x=0}^k \tbinom{x+r-1}{x}\left[-x(1-q)^{x-1}q^r+r(1-q)^x q^{r-1}\right]\\ & = \sum_{x=0}^k \tbinom{x+r-1}{x}\left[-x(1-q)^{x-1}q^r+r(1-q)^x q^{r-1}\right]\\ & = \sum_{x=0}^k \tbinom{x+r-1}{x}\left[x[(1-q)-1](1-q)^{x-1}q^{r-1}+r(1-q)^x q^{r-1}\right]\\ & = \sum_{x=0}^k \tbinom{x+r-1}{x}\left[-x(1-q)^{x-1}q^{r-1}+(x+r)(1-q)^x q^{r-1}\right]\\ &= - \sum_{x=0}^kx \tbinom{x+r-1}{x}(1-q)^{x-1}q^{r-1}+\sum_{x=0}^k(x+r) \tbinom{x+r-1}{x}(1-q)^x q^{r-1}\\ &= - \sum_{x=1}^kx \tbinom{x+r-1}{x}(1-q)^{x-1}q^{r-1}+\sum_{x=0}^k(x+r) \tbinom{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'=0}^{k-1} \frac{(x'+r)!}{x' ! r!}(1-q)^{x' }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{aligned}

而根据正则不完全 Beta 函数的定义,有

I1p(r,k+1)=Iq(r,k+1)=B(q;r,k+1)B(r,k+1)=0qtr1(1t)kdtB(r,k+1)\begin{aligned} I_{1-p}(r, k+1) &= I_{q}(r, k+1)\\ &= \frac{B(q; r, k+1)}{B(r, k+1)}\\ &= \frac{ \int_0^qt^{r-1}(1-t)^k\mathrm dt}{B(r, k+1)} \end{aligned}

同样对 qq 求偏导,得

Iqq=qr1(1q)kB(r,k+1)=Γ(r+k+1)Γ(r)Γ(k+1)qr1(1q)k=(r+k)!(r1)!k!qr1(1q)k=rq2(r+k)!r!k!qr+1(1q)k=rq2f(k;r+1,1q)\begin{aligned} \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{aligned}

也就是说

qF(k;r;1q)=qIq(r,k+1)\frac{\partial}{\partial q} F(k;r;1-q)= \frac{\partial}{\partial q}I_q(r, k+1)

亦即

F(k;r;1q)=Iq(r,k+1)+CF(k;r;1-q) = I_q(r, k+1) + C

注意到 q=0q=0 时有

{F(k;r,1)=0I0(r,k+1)=0\begin{cases} F(k; r, 1) = 0\\ I_0(r, k+1) = 0 \end{cases}

解得常数 C=0C=0

证毕。

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

DeepAR 等模型中,网络的输出目标是概率分布的参数。例如正态分布的 μ\muσ\sigma。但对于负二项分布而言,让网络直接输出 μ\muσ\sigma 是不合适的,因为在训练过程中很难保证输出的值满足 σ2μ\sigma^2\geq\mu。那么让网络输出 rrpp 呢?似乎是可以的,只要保证 r>0r>00p10\leq p\leq 1 即可。前者可以使用 softplus 激活函数,后者可以使用 sigmoid 激活函数。有没有办法避免使用 sigmoid 呢?通常的做法是让网络输出 μ\muα=1/r\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.