Hotdry.
归档
systems-engineering

别名法：O(n) 预处理实现 O(1) 非均匀离散采样

详解别名法算法的预处理与采样过程，利用简单数组操作和均匀随机数生成器，实现高吞吐量非均匀离散分布采样，适用于模拟和机器学习场景。

在高性能计算场景中，如蒙特卡洛模拟、强化学习中的动作采样或生成模型中的分类分布抽样，非均匀离散分布采样是一个常见瓶颈。传统方法如累积分布函数（CDF）逆变换，需要 O (n) 时间排序概率并二分查找，导致每样本开销过高。当分布固定且采样次数巨大时，别名法（Alias Method）提供完美解：O (n) 预处理构建两个数组（Prob 和 Alias），后续每样本严格 O (1) 时间，仅需两次均匀随机数调用和简单数组访问。

别名法核心思想

别名法将概率质量重新分配到 n 个等宽 “柱子”（每个宽度 1/n，总宽 1），每个柱子高度 <=1，由左侧原概率质量填充（高度 Prob [i]），剩余部分（1 - Prob [i]）由某个 “别名” 事件填充（Alias [i]）。这样，采样时均匀选柱子 k，再掷均匀骰子 u ∈ [0,1)，若 u < Prob [k] 则返回 k，否则返回 Alias [k]。这保证了精确采样，且无分支预测失效问题。

Keith Schwarz 在其经典文章《Darts, Dice, and Coins》中详细阐述了该方法1。算法的关键是预处理阶段的 “平衡堆栈” 策略，避免复杂图匹配。

预处理算法（O (n) 时间与空间）

给定概率数组 p [0..n-1]（和为 1），构建 Prob [0..n-1] 和 Alias [0..n-1]：

  1. 初始化：

    Prob[i] = n * p[i]  // 缩放至 [0, n]
小堆 (L)：Prob[i] < 1 的 i
大堆 (G)：Prob[i] >= 1 的 i

  2. 平衡循环：

    while L 非空 且 G 非空:
    l = L.pop()
    g = G.pop()
    Alias[l] = g
    Prob[g] -= (1 - Prob[l])  // 转移质量给 l 的右侧
    if Prob[g] < 1:
        L.push(g)
    else:
        G.push(g)

  3. 剩余处理（由于浮点精度，总质量守恒）：

    while G 非空:
    g = G.pop()
    Prob[g] = 1.0  // 设为满柱

伪代码完整实现（Python 风格）：

def build_alias(probs):
    n = len(probs)
    prob = [0.0] * n
    alias = [0] * n
    small = []
    large = []
    for i in range(n):
        prob[i] = n * probs[i]
        if prob[i] < 1.0:
            small.append(i)
        else:
            large.append(i)
    
    while small and large:
        l = small.pop()
        g = large.pop()
        alias[l] = g
        prob[g] -= 1.0 - prob[l]
        if prob[g] < 1.0:
            small.append(g)
        else:
            large.append(g)
    
    while large:
        g = large.pop()
        prob[g] = 1.0
    
    return prob, alias

注意：输入 probs 需预归一化（sum=1），但实际工程中常从计数构建：probs [i] = counts [i] /total。

采样过程（严格 O (1)）

def sample(prob, alias, rng):
    n = len(prob)
    col = int(rng.uniform(0, n))  # 均匀选柱 [0, n)
    test = rng.uniform(0, 1)
    return col if test < prob[col] else alias[col]

仅需快速均匀 RNG（如 PCG 或 Xorshift），无循环、无排序。现代 CPU 上，单样本 <10 周期。

工程化参数与优化

  • 规模：n ≤ 10^7 内存友好（~80MB 双数组 double+int），预处理 O (n) 秒级。ML 中 softmax 后 n=vocab_size50k 常见。
  • RNG 选择：SIMD 友好如 wyhash 或 lemire RNG，支持批量采样（vectorize col 和 test）。
  • 精度：浮点误差 <1e-10，剩余 G 堆调整确保总质量 = 1。若严格整除，用整数实现（Vose 变体）：概率 ×n! 但 n 大不可行，故 float 为主。
  • 批量采样：预取 1024 样本，用 AVX2 并行 uniform+cmp+select。
  • 存储：序列化 Prob/Alias 为二进制，加载即用。动态更新：全重建（若 n 小）。
  • 监控点：
    指标 阈值 异常处理
    预处理时间 < n/1e8 秒 增大堆栈缓冲
    采样吞吐 >1e8/s/ 核 换 RNG / 向量化
    分布 KS 测试 p>0.01 重建表格
    内存驻留 <100MB 分页表

回滚策略：fallback 到轮询（round-robin）或二分 CDF（O (log n)）。

应用场景

  • 模拟：粒子物理中事件类型采样（darts-dice-coins 隐喻）。
  • ML：策略梯度中动作采样、Gumbel-softmax 替代、NLL 损失下的分类。
  • 游戏： loot box 稀有度抽取，高 QPS。 对比：vs multinomial（O (n) per sample），alias 加速 10-100x 于重复分布。

实际基准：在 M1 Max，n=1e6，采样 10^9 次～5s（单线程）。

总之，别名法以极简实现换取极致性能，是高吞吐采样标配。工程中，关注 RNG 质量与向量化即可落地。

Footnotes

  1. Keith Schwarz, "Darts, Dice, and Coins: Sampling from a Discrete Distribution", https://keithschwarz.com/darts-dice-coins/ ↩

查看归档