加权随机采样

发表于 | 分类于 | | 阅读数
字数统计: 2.7k | 阅读时长 ≈ 10

加权随机采样在推荐系统中随处可见，既可能用在模型训练数据处理过程中，也可能用于一些规则式的推荐策略里。典型的场景例如：

  1. 在新用户冷启动时，我们可以通过某些指标评估出内容的质量，并根据质量得分来将内容加权随机推荐给新用户，质量越高的内容，被曝光给新用户的概率也越大。
  2. 在样本采样时，有一种方法是对每条正样本，随机从所有的内容中选取 $k$ 个负样本，而每个内容被选为负样本的概率与其热度成正比（例如 word2vec 的 negative sampling 技术）。
  3. 基于用户历史行为可以构建内容的有向图，当用户行为较稀疏时，我们可以使用 deepwalk 之类的算法在图中随机游走，生成内容的序列，再基于 word2vec 等算法生成这些内容的 embedding。在随机游走时，比如当前到达节点 $v$，那么下一次游走到其他节点的概率，与有向边的权重正相关。

以上这些加权采样的场景往往都不可避免的面对大数据量挑战，因此对性能要求较高。而再细分一下，上面的场景 1, 2 都属于无放回采样，也就是需要采样的内容都不相同；而场景 3 则属于有放回采样，即允许采到相同的样本。本文就这两类加权随机采样问题分别探讨高效的解法。

常规方案——基数法

先再把问题抽象并更正式一点的描述一下：假设列表 $S$ 中有 $n$ 个元素，初始时，已知每个元素 $i$ 被抽取的概率为 ，则：

  1. 无放回加权随机采样是指，随机从 $S$ 中不放回的抽取 $m$ 个元素，每个元素 $i$ 被抽取的概率为 ，其中 $S’$ 表示该次抽样时所有剩余元素的集合；
  2. 有放回加权随机采样是指，随机从 $S$ 中有放回的抽取 $m$ 个元素，每个元素 $i$ 被抽取的概率为 $p_i$。

一般来讲，我们更希望获得一种时间复杂度低的抽样算法，也就是能够用尽可能少的步骤来得到 $m$ 个抽样的结果。在不考虑极致的精度和极端场景时，一种简单粗暴的方案是基于基数的抽样算法，这也是 word2vec 官方实现$^{[1]}$的策略：

  1. 生成一个超大的数组 $A$，数组长度为 $k$，例如 $k=10^8$，或者 $k=\text{ceil}(1/\min(p_i))$；
  2. 将数组中的元素填充为集合元素的 index，元素 $i$ 的 index 被填充的数量为 ；
  3. 每次生成随机整数 $r: 0\le r < k$，并将 $A[r]$ 对应的元素作为采样结果；
  4. 对于无放回的情况，还需要记录已经采样的结果集合，如果采样的结果出现在集合中，则重新采样，直到得到 $m$ 个不同的采样结果。

我们将这种方法称为基数法。这种方法在 $S$ 中元素的数量远远超过 $m$，并且每个元素被抽取的概率远远小于 $1$ 时有较高的效率，是一种空间换时间的策略。但是在极端情况下，例如某个元素 $i$ 的 $p_i$ 接近 $1$ 时，无放回的抽样效率可能极低，因为每次抽样绝大概率抽到 $i$，进而需要重新抽样。另外，由于第 $2$ 步实际是一种近似策略，只能保证大体上数据 $A$ 中的元素 index 数量与其抽样概率成正相关，而不能保留精确的比例。因此，在需要更高效率和更高精度的条件下，我们可以考虑下面的方案。

有放回加权随机采样——Alias Method

虽然基数法应用在有放回加权随机采样时，每次采样的时间复杂度为 $O(1)$，但是我们获得的实际上是一种非精确的采样，并且基数法的空间复杂度较高。相对而言，Alias Method$^{[2]}$ 是一种空间和时间都极为高效的算法，它的主要流程分为初始化和采样两步：首先用 $O(n)$ 的复杂度初始化两个长度为 $n$ 的数组，然后基于这两个数组进行采样，每次采样的复杂度为 $O(1)$。下面详细介绍下这种方法的流程。

首先将所有元素的抽样概率都乘以 $n$，得到一个平均概率为 $1$ 的新的采样概率 $Q=[q_1, q_2, \cdots,q_n]$，然后对这些概率进行两两组合：

  1. 选择一个概率不超过 $1$ 的元素 $i$：$q_i\le 1$，设置 $Prob[i]=q_i$；
  2. 选择一个概率不小于 $1$ 的元素 $j$：$q_j\ge 1$，设置 $Alias[i] = j$，并将 $q_j:=q_j-(1-q_i)$，即用 $q_j$ 来补足 $q_i$ 少于 $1$ 的部分；

这样组合以后，我们得到了两个数组：原始概率 $Prob$ 和组合元素索引 $Alias$。从构造的过程中，我们可以发现这两个数组的长度都是 $n$，并且 $Prob$ 对应的类别顺序与 $P$ 一致，而 $Alias$ 则保存着跟原始类别进行概率组合的类别编号。进一步的，对于每个数组下标 $i$，它一定会对应一个原始的元素 $i$，以及至多一个组合元素 $Alias[i]$，这个组合元素是用来跟 $i$ 一起将概率凑成 $1$ 的。

基于这两个数组，在每次采样时，只需要生成两个随机数：

  1. 第一个随机数范围是 $r_1\in[1,n]$，用于确定原始元素，假设 $r_1=i$；
  2. 第二个随机数范围是 $r_2\in [0,1]$，如果 $r_2< Prob[i]$，本次采样的类别就是原始元素 $i$，否则本次采样的类别是其组合元素 $j=Alias[i]$；

到这里，大家肯定会好奇，在两两组合概率时，是不是一定能找到一种方案，使得对于所有的元素 $i$，一定能为其找到最多一个组合类别，使得它们的概率之和为 $1$。答案当然是肯定的，我们可以用归纳法进行证明：

  1. 当 $n=1$ 时，$Prob[1]=1$, $Alias[1]=null$，命题显然成立；
  2. 对于任意正整数 $k$，假设当 $n=k$ 时命题成立，则当 $n=k+1$ 时，我们一定能找到两个类别 $c_i, c_j$，满足 $q_i\le 1, q_j\ge 1$，则我们设置 $Prob[i]=q_i$，$Alias[i] = j$，$q_j:=q_j-(1-q_i)$ 后，得到除 $c_i$ 外的 $k$ 个类别，它们的平均概率仍为 $1$，因此根据假设，$n=k+1$ 时，命题仍然成立。

因此，我们按上面的组合策略，一定能成功构造出一个 $Alias$ 和 $Prob$ 的数组。

下面提供了一个 python 版本的实现，仅供参考。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import numpy as np
import numpy.random as npr

class AliasMethodSampling:
    def __init__(self, p):
        self.prob, self.alias = self._setup_alias(p)
    
    def _setup_alias(self, p):
        small, large = [], []
        n = len(p)
        prob, alias = np.zeros(n), np.zeros(n, dtype=np.int)
        # init small and large array
        for i, pi in enumerate(p):
            prob[i] = pi * n
            small.append(i) if prob[i] < 1.0 else large.append(i)
        # fill a small element each iteration
        while small and large:
            il, ig = small.pop(), large.pop()
            alias[il] = ig
            prob[ig] = (prob[ig] + prob[il]) - 1
            small.append(ig) if prob[ig] < 1 else large.append(ig)
        # handle numerical instability
        while large:
            print('only large exists', prob, alias, large, small)
            prob[large.pop()] = 1
        while small:
            print('only small exists', prob, alias, large, small)
            prob[small.pop()] = 1
        return prob, alias
        
    def take_sample(self):
        i_rand = npr.randint(len(self.prob))
        return i_rand if npr.rand() < self.prob[i_rand] else self.alias[i_rand]
    
    
probs, cnt = [0.1, 0.5, 0.2, 0.05, 0.15], np.zeros(5)
alias = AliasMethodSampling(probs)

for i in range(1000000):
    cnt[alias.take_sample()] += 1
    
cnt / 1000000

无放回加权随机采样——A-ExpJ 算法

在上面讨论的有放回加权随机采样情形中，每次抽样时，每个元素的抽样概率实际上是不发生变化的，我们也不需要知道之前抽样的结果是什么。但是在无放回情形中，每次抽到一个元素后，后面就不能再抽到相同的元素，也就意味着每次抽样后，元素的抽样概率会发生变化；并且我们需要记录之前抽样的结果有哪些，来防止抽到重复的元素。

前面也提到了，应用基数法来进行无放回加权随机采样的主要问题是可能重复抽到同一个元素，而生成随机数的成本其实是很高的，极端场景下，我们需要生成随机数的次数远大于 $n$。

文献 $[3]$ 提出了一种基于代理特征来采样的方法，也就是对于每个元素 $i$，我们选取服从均匀分布的随机数 $u_i=\text{rand}(0,1)$，用 $k_i=u_i^{1/p_i}$ 来作为采样的关键值，并选择关键值最大的 $m$ 个样本作为采样结果。这样，我们至多需要生成 $n$ 个随机数，就能完成采样过程。至于 $k_i$ 选择的正确性可以参考文献 $[3]$ 中的证明。基于这种代理特征的好处是，不需要知道每个元素的采样概率，只需要知道其权重即可 (也就是说不需要知道所有元素的总权重，无需做概率的归一化)，这就特别适合流式采样的场景。

当然，这里选择关键值最大的 $m$ 个样本可以采用最大堆来实现，这样就只需要进行一次全量 $O(n)$ 扫描，同时只要保留 $m$ 个当前最大的结果即可，这就是 A-Res 算法。在绝大多数场景中，A-Res 算法已经足够高效了。

为了进一步减少随机数生成的数量，作者提出了 A-ExpJ 算法，能将随机数的生成量从 $O(n)$ 减少到 $O(m\log⁡(\frac{n}{m})))$。实际上就是用计算复杂度比较低的反向计算来代替复杂度高的随机数生成，并直接跳过一些关键值明显较小的元素。具体步骤如下：

  1. 将列表 $S$ 的前 $m$ 个元素放入结果集合 $R$；
  2. 对于结果集里的每个元素，计算关键值 $k_i=u_i^{(1/p_i)}$，其中 $u_i=\text{rand}(0,1)$；
  3. 将 $R$ 中小最的关键值记为阈值 $k_{min}$；
  4. 对剩下的元素重复以下步骤：
    1. 令 $r=\text{rand}(0,1)$ 且 $x_p=\log(r)/\log(t)$；
    2. 从当前元素 $c$ 开始跳过元素，直到遇到元素 $i$，满足 ；
    3. 使用 $i$ 替换 $R$ 中关键值最小的元素；
    4. 令 , $r_2=\text{rand}(t,1)$, $i$ 的关键值 $k_i=r_2^{(1/p_i)}$；
    5. 令新的阈值 $k_{min}$ 为此时 $R$ 中的最小关键值。

下面提供了一种 python 实现，仅供参考。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import random
import math
import heapq
import numpy as np


def a_expj_sample(prob, m):
    """ 根据 prob 数组无放回随机抽取 m 个元素 """
    topn = []
    for i, pi in enumerate(prob[:m]):
        heapq.heappush(topn, (random.random() ** (1/pi), i))
        
    thres, w_sum = topn[0][0], 0
    xw = math.log(random.random()) / math.log(thres)
    i = m
    for pi in prob[m:]:
        if w_sum + pi >= xw:
            tw = thres ** pi
            r2 = random.uniform(tw, 1)
            ki = r2 ** (1/pi)
            heapq.heappop(topn)
            heapq.heappush(topn, (ki, i))
            thres = topn[0][0]
            xw = math.log(random.random()) / math.log(thres)
            w_sum = 0
        else:
            w_sum += pi
        i += 1
    return [item[1] for item in topn]


probs = [0.1, 0.5, 0.2, 0.05, 0.15]

for k in range(1, 4):
    cnt = np.zeros(5)
    for i in range(1000000):
        for j in a_expj_sample(probs, k):
            cnt[j] += 1
    print(k, ':', cnt / 1000000 / k)

小结

本文主要介绍了三种加权采样的算法，其中，有放回加权随机采样推荐使用 Alias Method，无放回加权随机采样推荐使用 A-ExpJ 或者 A-Res，如果对精度要求没那么高，并且样本权重呈现非极端的分布，也可以使用简单的基数法。本篇讨论的算法主要应用场景在于我们已知每个元素的采样权重，进而可以设置其采样概率。但是有些情况下，我们是不知道这些元素的分布是什么样，而我们希望从这些未知分布中抽取样本，再利用这些样本对目标做一个估计。这时就需要考虑重要性采样、马尔可夫链蒙特卡罗方法、Gibbs 采样等采样算法，后面有机会再讨论吧。

参考文献

[1] Github: word2vec. dav: https://github.com/dav/word2vec/blob/master/src/word2vec.c.

[2] Darts,Dice, and coids. Keith. 2011: http://www.keithschwarz.com/darts-dice-coins.

[3] Efraimidis, Pavlos S. , and P. G. Spirakis . “Weighted random sampling with a reservoir.” Information Processing Letters97.5(2006):181-185.