From 630a9c4636875f94cf106d6991152e4de7435bd1 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Wed, 22 Jul 2026 08:01:33 -0400 Subject: [PATCH] =?UTF-8?q?=E5=B1=823=20E2:=20estimator.py=E2=80=94?= =?UTF-8?q?=E2=80=94=E5=BC=8F(4)=20=E5=87=A0=E4=BD=95=E5=9D=87=E5=80=BC?= =?UTF-8?q?=E5=85=88=E9=AA=8C=20chunk=5Fprior=20+=20=E5=BC=8F(5)=20?= =?UTF-8?q?=E8=B4=9D=E5=8F=B6=E6=96=AF=E7=9B=AE=E6=A0=87=20bayesian=5Ftarg?= =?UTF-8?q?et,17=20=E5=8D=95=E6=B5=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 对应 docs/04 §4 E2。detach 双防线(chunk_prior 内为本质防线、 bayesian_target 末尾为防御性第二道,分别对应参考 trainer:2196/2205); log 域抗下溢;clamp 下限守定理 4.1(b);k_sem∈[0,N] 口径校验; 方差收缩 toy 模拟对拍 validate_chunk_mc_estimator.py 精神 (MSE_bayes < MSE_freq,含劣先验 ±0.1 稳健性)。 Co-Authored-By: Claude Opus 4.8 (1M context) --- ars_opd/estimator.py | 89 ++++++++++++++++++++++ tests/test_estimator.py | 162 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 251 insertions(+) create mode 100644 ars_opd/estimator.py create mode 100644 tests/test_estimator.py diff --git a/ars_opd/estimator.py b/ars_opd/estimator.py new file mode 100644 index 0000000..26b5750 --- /dev/null +++ b/ars_opd/estimator.py @@ -0,0 +1,89 @@ +"""MC 估计与 Dirichlet 贝叶斯平滑——论文 §3.2.2 式(4)(5)。 + +把 similarity.py 产出的软计数 k_sem(外部 teacher 信号)与学生自身的 +chunk 置信度 π̄(内部先验)融合成有界目标 π̂ ∈ (0, 1],供层 5 的 chunk +损失当乘子:loss_c = −π̂ · mean(log p)。定理 4.1 三性质由此获得: +(a) π̂ 有界 → 无白盒式(2) 的梯度爆炸;(b) π̂ > 0 → k_sem=0 也不塌缩; +(c) 先验收缩 → 方差小于频率估计 k/N。 + +纯逻辑模块(CLAUDE.md §2):只依赖 torch,toy 张量本地 CPU 可测。 +""" + +import torch + + +def chunk_prior(log_probs: torch.Tensor) -> torch.Tensor: + """式(4):π̄ = exp((1/C)·Σ_t log p_t)——学生对整个 chunk 的几何均值置信度。 + + C 个 token 概率的几何均值,充当式(5) 的贝叶斯先验:teacher 采样(k_sem) + 是主信号,π̄ 只是"学生自己觉得这段有多稳"的地板,防 k_sem=0 时目标归零。 + + 参数: + log_probs: 学生对 chunk 内各 token 的对数概率,shape (C,),值 ≤ 0。 + (调用方从 log_softmax 后 gather 标签位置所得,层 5 负责。) + + 返回: + π̄,标量张量 shape ()、值 ∈ [1e-8, 1],**不带梯度**。 + + 实现细节: + - log 域先均值再 exp:直接连乘 C=50 个小概率会下溢 + (50 个 0.01 → 1e-100,超出 fp32 下限 ~1e-38),log 域安全。 + - detach 命门(参考实现 distillation_trainer.py:2196 同):π̄ 是学生 + 自身概率的函数,若保留梯度,优化器会发现"压低自己的 chunk 概率 + → π̄→0 → π̂ 变小 → 损失权重变小"这条逃逸路径——恰在 k_sem=0 + (teacher 否定)的 chunk 上最有利可图,这些 chunk 最先塌缩。 + π̄ 只能当常数先验,不能当优化变量。锁死断言见 + tests/test_estimator_detach.py。 + - clamp 下限 1e-8:极端负的均值 exp 后可能下溢为 0,而定理 4.1(b) + 的反塌缩要求 π̄ 严格为正。差异标注:参考实现 clamp(1e-8, 1.0), + 上限实为冗余——log p ≤ 0 ⇒ mean ≤ 0 ⇒ exp ≤ 1,此处省去。 + """ + if log_probs.numel() == 0: + raise ValueError("log_probs 为空:chunk 至少要含 1 个 token") + log_pi_bar = log_probs.detach().mean() # (C,) -> () + return log_pi_bar.exp().clamp(min=1e-8) + + +def bayesian_target( + k_sem: float, + pi_bar: torch.Tensor, + n_rollouts: int, + alpha: float, +) -> torch.Tensor: + """式(5):π̂ = (k_sem + α·π̄) / (N + α)——chunk 接受概率的贝叶斯估计。 + + 等价凸组合视角(论文式10): + π̂ = N/(N+α) · (k_sem/N) + α/(N+α) · π̄ + 即"teacher 频率估计"与"学生先验"的加权平均;默认 N=10、α=1 时权重 + 约 91% : 9%,teacher 主导,先验只兜底。 + + 参数: + k_sem: 式(3) 的软匹配计数,∈ [0, N](aggregate_similarity 产出)。 + pi_bar: 式(4) 的先验 π̄,标量张量(chunk_prior 产出)。 + n_rollouts: teacher rollout 数 N。**必须等于算 k_sem 时的 + len(teacher_rollouts)**——分子分母口径不一致会系统性偏移 π̂。 + alpha: 先验强度 α ≥ 0。α=0 退化为频率估计 k/N(层 6 的 + no_bayesian 消融,参考 config.py:315),失去定理 4.1(b) 保护。 + + 返回: + π̂,标量张量 shape ()、值 ∈ [1e-8, 1],**不带梯度**。 + + 实现细节: + - 差异标注:参考实现(distillation_trainer.py:2205)在此对 π̂ 整体 + detach;我们的 π̄ 在 chunk_prior 内已 detach,此处的 detach 是 + 第二道防线——防止将来有人把带梯度的张量传进 pi_bar。 + - clamp(1e-8, 1.0):下限防 α=0 且 k_sem=0 时 π̂=0(乘子归零则该 + chunk 完全失去监督);上限防 pi_bar 越界传入时 π̂ 溢出概率语义。 + """ + if n_rollouts < 1: + raise ValueError(f"n_rollouts 必须 ≥ 1,得到 {n_rollouts}") + if alpha < 0: + raise ValueError(f"alpha 必须 ≥ 0,得到 {alpha}") + if not 0.0 <= k_sem <= n_rollouts: + raise ValueError( + f"k_sem={k_sem} 越界 [0, {n_rollouts}]:检查是否与" + f" len(teacher_rollouts) 口径一致" + ) + # 式(5): π̂ = (k_sem + α·π̄) / (N + α) + pi_hat = (k_sem + alpha * pi_bar) / (n_rollouts + alpha) # () -> () + return pi_hat.clamp(1e-8, 1.0).detach() diff --git a/tests/test_estimator.py b/tests/test_estimator.py new file mode 100644 index 0000000..6950dac --- /dev/null +++ b/tests/test_estimator.py @@ -0,0 +1,162 @@ +"""estimator.py 单测——docs/04 §5.2:公式对拍 + 定理 4.1 性质 + 方差收缩。 + +对拍精神源自参考实现 validate_chunk_mc_estimator.py(比 MSE_freq vs +MSE_bayes),但全用 toy 数据本地 CPU 跑,不连真 teacher。 +detach 命门的两个世界断言在 tests/test_estimator_detach.py(E3)。 +""" + +import math + +import pytest +import torch + +from ars_opd.estimator import bayesian_target, chunk_prior + +# ------------------------------------------------------------- chunk_prior + + +def test_prior_is_geometric_mean(): + # 式(4) 手算:p = [0.9, 0.1] → π̄ = exp((log .9 + log .1)/2) = √0.09 = 0.3 + log_probs = torch.log(torch.tensor([0.9, 0.1])) + assert math.isclose(chunk_prior(log_probs).item(), 0.3, rel_tol=1e-6) + + +def test_prior_uniform_probs(): + # 全同概率的几何均值 = 该概率本身 + log_probs = torch.full((50,), math.log(0.5)) + assert math.isclose(chunk_prior(log_probs).item(), 0.5, rel_tol=1e-6) + + +def test_prior_shape_and_range(): + pi_bar = chunk_prior(torch.log(torch.rand(50).clamp(1e-6, 1.0))) + assert pi_bar.shape == () # (C,) -> 标量 + assert 0.0 < pi_bar.item() <= 1.0 + + +def test_prior_log_domain_survives_underflow(): + # 50 个 p=0.01 直接连乘 = 1e-100(fp32 下溢为 0);log 域算出 0.01 + log_probs = torch.full((50,), math.log(0.01)) + assert math.isclose(chunk_prior(log_probs).item(), 0.01, rel_tol=1e-4) + + +def test_prior_clamp_floor(): + # 极端负 log 均值 → exp 下溢,clamp 兜到 1e-8 保持严格为正(定理 4.1b 前提) + log_probs = torch.full((5,), -1e9) + assert chunk_prior(log_probs).item() == pytest.approx(1e-8) + + +def test_prior_is_detached(): + # detach 命门:π̄ 不带梯度(逃逸机制的完整断言在 test_estimator_detach.py) + log_probs = torch.log(torch.tensor([0.5, 0.5], requires_grad=True)) + pi_bar = chunk_prior(log_probs) + assert not pi_bar.requires_grad + + +def test_prior_empty_raises(): + with pytest.raises(ValueError, match="为空"): + chunk_prior(torch.tensor([])) + + +# --------------------------------------------------------- bayesian_target + + +def test_target_formula_hand_computed(): + # 式(5) 手算:k=8, π̄=0.5, N=10, α=1 → π̂ = (8 + 0.5)/11 = 0.77272… + pi_hat = bayesian_target(8.0, torch.tensor(0.5), n_rollouts=10, alpha=1.0) + assert math.isclose(pi_hat.item(), 8.5 / 11, rel_tol=1e-6) + + +def test_target_convex_combination_identity(): + # 式(10) 恒等:π̂ = N/(N+α)·(k/N) + α/(N+α)·π̄,任取参数逐点核对 + k, pi_bar, n, alpha = 3.7, torch.tensor(0.42), 10, 1.5 + direct = bayesian_target(k, pi_bar, n, alpha).item() + convex = (n / (n + alpha)) * (k / n) + (alpha / (n + alpha)) * pi_bar.item() + assert math.isclose(direct, convex, rel_tol=1e-6) + + +def test_target_anti_collapse_at_k_zero(): + # 定理 4.1(b):k=0(teacher 全否定)时 π̂ = α·π̄/(N+α) > 0,监督不归零 + pi_hat = bayesian_target(0.0, torch.tensor(0.3), n_rollouts=10, alpha=1.0) + assert math.isclose(pi_hat.item(), 0.3 / 11, rel_tol=1e-6) + assert pi_hat.item() > 0 + + +def test_target_full_score_shrinks_below_one(): + # 贝叶斯收缩:k=N 满分时 π̂ = (N+α·π̄)/(N+α) < 1(只要 π̄<1)—— + # 先验把估计从两端往中间拉,这正是方差收缩的来源 + pi_hat = bayesian_target(10.0, torch.tensor(0.5), n_rollouts=10, alpha=1.0) + assert math.isclose(pi_hat.item(), 10.5 / 11, rel_tol=1e-6) + assert pi_hat.item() < 1.0 + + +def test_target_bounded_in_unit_interval(): + # 定理 4.1(a):任意合法参数下 π̂ ∈ (0, 1] + for k in [0.0, 2.5, 10.0]: + for p in [1e-8, 0.5, 1.0]: + v = bayesian_target(k, torch.tensor(p), 10, 1.0).item() + assert 0.0 < v <= 1.0 + + +def test_target_alpha_zero_is_frequency_estimate(): + # α=0 退化为 k/N(no_bayesian 消融);k=0 时被 clamp 兜到 1e-8 而非 0 + assert math.isclose( + bayesian_target(7.0, torch.tensor(0.5), 10, 0.0).item(), 0.7, rel_tol=1e-6 + ) + assert bayesian_target(0.0, torch.tensor(0.5), 10, 0.0).item() == pytest.approx( + 1e-8 + ) + + +def test_target_is_detached_even_with_grad_input(): + # 第二道防线:pi_bar 带梯度传入,π̂ 仍必须 detach + pi_bar = torch.tensor(0.5, requires_grad=True) + pi_hat = bayesian_target(5.0, pi_bar, 10, 1.0) + assert not pi_hat.requires_grad + + +def test_target_validation_raises(): + pi_bar = torch.tensor(0.5) + with pytest.raises(ValueError, match="n_rollouts"): + bayesian_target(0.0, pi_bar, 0, 1.0) + with pytest.raises(ValueError, match="alpha"): + bayesian_target(0.0, pi_bar, 10, -0.1) + with pytest.raises(ValueError, match="越界"): + bayesian_target(11.0, pi_bar, 10, 1.0) # k_sem > N:口径不一致 + with pytest.raises(ValueError, match="越界"): + bayesian_target(-0.5, pi_bar, 10, 1.0) + + +# --------------------------------------------- 方差收缩(定理 4.1c,toy 模拟) + + +def test_variance_shrinkage_beats_frequency_estimate(): + """toy 模拟对拍 validate_chunk_mc_estimator.py 的 MSE_freq vs MSE_bayes。 + + 设真值 μ:每次试验采 N=10 个相似度 sim_i(均值 μ 的噪声), + 频率估计 = mean(sim) = k/N,贝叶斯估计 = (k + α·π̄)/(N+α)。 + 先验 π̄ = μ(理想先验)时,收缩纯降方差、零偏差代价,MSE 必更小。 + """ + torch.manual_seed(0) + mu, n, alpha = 0.7, 10, 1.0 + trials = 2000 + # (trials, N) 的相似度样本:均值 μ、截断到 [0,1] + sims = (mu + 0.25 * torch.randn(trials, n)).clamp(0.0, 1.0) + k = sims.sum(dim=1) # (trials,) 每次试验的 k_sem + freq = k / n + bayes = (k + alpha * mu) / (n + alpha) + mse_freq = ((freq - mu) ** 2).mean().item() + mse_bayes = ((bayes - mu) ** 2).mean().item() + assert mse_bayes < mse_freq + + +def test_variance_shrinkage_robust_to_imperfect_prior(): + # 先验偏离真值(π̄ = μ±0.1)仍应赢:α=1、N=10 时先验权重仅 1/11, + # 引入的偏差平方远小于省下的方差(定理 4.1c 在论文设定下的稳健性) + torch.manual_seed(1) + mu, n, alpha, trials = 0.6, 10, 1.0, 2000 + sims = (mu + 0.25 * torch.randn(trials, n)).clamp(0.0, 1.0) + k = sims.sum(dim=1) + mse_freq = ((k / n - mu) ** 2).mean().item() + for prior in [mu - 0.1, mu + 0.1]: + bayes = (k + alpha * prior) / (n + alpha) + assert ((bayes - mu) ** 2).mean().item() < mse_freq