无放回两步采样场景下如何调整样本权重以匹配原始概率分布

问题核心解法 #

这个问题本质是通过两步采样的边际分布匹配原始概率向量p，我们可以通过全概率公式推导第二步的采样权重：

核心逻辑 #

设：

• 第一步采样得到子集I的概率为P(I)
• 第二步在子集I中采样到索引i的条件概率为w(i|I)
  我们需要满足对任意索引i，全概率等式成立：

sum_{I: i ∈ I} P(I) * w(i|I) = p_i
同时每个子集内的权重满足归一化约束：
sum_{i ∈ I} w(i|I) = 1，且所有w(i|I) ≥ 0

你给出的示例求解 #

你给出的场景中已知：

• p = [1/6, 1/3, 1/2]
• 第一步子集概率：P([0,1])=3/20，P([0,2])=4/15，P([1,2])=7/12
  我们对三个索引分别列方程：

1. 索引0：3/20 * w(0|[0,1]) + 4/15 * w(0|[0,2]) = 1/6
2. 索引1：3/20 * w(1|[0,1]) + 7/12 * w(1|[1,2]) = 1/3
3. 索引2：4/15 * w(2|[0,2]) + 7/12 * w(2|[1,2]) = 1/2
   加上三个子集的归一化约束：

• w(0|[0,1]) + w(1|[0,1]) = 1
• w(0|[0,2]) + w(2|[0,2]) = 1
• w(1|[1,2]) + w(2|[1,2]) = 1
  该线性方程组存在多组可行解，我们可以任选一组满足非负约束的解即可，比如取一组简单的可行解：
• I=[0,1] 权重：[1/2, 1/2]
• I=[0,2] 权重：[11/32, 21/32]
• I=[1,2] 权重：[31/70, 39/70]
  验证：
• 索引0总采样概率：3/20 * 1/2 + 4/15 * 11/32 = 1/6，符合要求
• 索引1总采样概率：3/20 * 1/2 + 7/12 * 31/70 = 1/3，符合要求
• 索引2总采样概率：4/15 * 21/32 + 7/12 * 39/70 = 1/2，符合要求

通用方案 #

• 如果第一步采样策略是固定的：直接按上述方式列线性方程组求解即可，需要唯一解时可以添加额外优化目标，比如最小化第二步权重和原始p在子集内归一化结果的偏差，降低采样方差。
• 如果第一步采样策略可以自定义：推荐使用Gumbel-top-K trick，无需手动计算权重：
  1. 对每个索引i，独立采样均匀随机数u_i ~ Uniform(0,1)，计算Gumbel噪声g_i = -log(-log(u_i))
  2. 第一步选g_i最大的K个索引作为子集I
  3. 第二步直接选I中g_i最大的索引，最终采样结果完全符合原始p的分布

内容的提问来源于stack exchange，提问作者borchero