""" TPM：分位点平衡树 + 叶节点概率 + 方差项 + 后门调整（第23章 §23.4.6） 正文把时长空间离散成有序 Rank，再用一棵近似平衡的二叉树逐层判断 「落在左子树还是右子树」。到叶节点之后得到整条概率分布： p(T ∈ l_k | X) = Π_{路径上的节点} p(走对分支 | 到达父节点) 最终预测是叶区间代表值的期望，同时还能算出方差。 四个可验证的点： ① 树结构：按分位点递归二分，各层节点样本量是否真的平衡； ② 叶概率分布：期望预测 vs 直接回归；分布方差能否预测误差大小； ③ 方差项与回归项：把「分布不确定性」当正则，能不能提升点预测； ④ 后门调整：人为造一个混杂变量 D，比较 E(T|X) 与 E(T|do(X))。 实现上每个内部节点一个 Logistic 回归。这是 TPM 论文「每个节点一个 分类器」的最小版本；论文用 MLP，节点数也更多。 运行：python3 ch04_classic_methods/tpm.py """ from __future__ import annotations import pathlib import sys import numpy as np sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) from common import ( # noqa: E402 fit_linear, fit_logistic, make_watch_data, mae, predict_linear, rmse, sigmoid, split, standardize, table, ) SEED = 2026 DEPTH = 5 # -------------------------------------------------------------------------- # 1) 建树：按观看时长的分位点递归二分 # -------------------------------------------------------------------------- def build_tree(w: np.ndarray, depth: int): """ 按观看时长的排名递归二分，返回 (nodes, leaves, sorted_w)。 nodes 前序编号，元素为: {"level", "lo", "hi", "mid", "left", "right"} 叶节点 left/right 为 None。leaves 是 [(lo, hi, node_id)]， lo/hi 是该叶在「排好序的样本」上覆盖的区间。 为什么按排名切而不是按阈值切：观看时长严重长尾，且 12% 的样本 w=0。按阈值取中位数时，前几层切点会反复落回 0 附近，切出 [0,0)、[0,0.001) 这种没有样本的区间（第一版实现就踩了这个坑， 32 个叶子里最左边 4 个是空的）。按排名切保证每个叶节点样本量 基本相等，只在「整段取值完全相同」时停下来——零观看那一大块 因此会自然收缩成一个叶子。这正是正文说的「按分位数构建平衡二叉树」。 切点还要保证两侧取值真的不同，否则该节点直接降级为叶。 """ order = np.argsort(w) sorted_w = w[order] nodes, leaves = [], [] def rec(i: int, j: int, level: int) -> int: node_id = len(nodes) nodes.append({"level": level, "lo": float(sorted_w[i]), "hi": float(sorted_w[j - 1]), "mid": None, "left": None, "right": None}) if level == depth or j - i <= 1 or sorted_w[i] == sorted_w[j - 1]: leaves.append((float(sorted_w[i]), float(sorted_w[j - 1]), node_id)) return node_id k = (i + j) // 2 # 找一个两侧取值不同的切点；找不到说明整段同值，只能当叶 while k > i and sorted_w[k - 1] == sorted_w[k]: k -= 1 if k == i: leaves.append((float(sorted_w[i]), float(sorted_w[j - 1]), node_id)) return node_id nodes[node_id]["mid"] = float(sorted_w[k]) nodes[node_id]["left"] = rec(i, k, level + 1) nodes[node_id]["right"] = rec(k, j, level + 1) return node_id rec(0, len(w), 0) return nodes, leaves, sorted_w def fit_tpm(Xtr: np.ndarray, wtr: np.ndarray, nodes: list): """每个内部节点训练一个 Logistic：1 = 走右子树（时长更长）。""" coefs = {} for i, node in enumerate(nodes): if node["left"] is None: continue m = (wtr >= node["lo"]) & (wtr <= node["hi"]) if m.sum() < 40: coefs[i] = None continue y = (wtr[m] > node["mid"]).astype(np.float64) if y.min() == y.max(): coefs[i] = None continue coefs[i] = fit_logistic(Xtr[m], y, iters=200, lr=0.5) return coefs def leaf_probabilities(X: np.ndarray, nodes: list, coefs: dict, n_leaves: int) -> np.ndarray: """ 返回 [n, n_leaves] 的叶概率矩阵。 深度优先、左子树优先遍历，叶节点按区间升序编号； 每个节点携带「到达它的概率」，一路相乘下去。 """ n = len(X) n_feat = X.shape[1] Xb = np.concatenate([X, np.ones((n, 1))], axis=1) probs = np.zeros((n, n_leaves)) counter = {"leaf": 0} def dfs(node_id: int, reach: np.ndarray) -> None: node = nodes[node_id] if node["left"] is None: probs[:, counter["leaf"]] = reach counter["leaf"] += 1 return coef = coefs.get(node_id) p_right = (np.full(n, 0.5) if coef is None else sigmoid(Xb @ coef)) p_right = np.clip(p_right, 1e-6, 1 - 1e-6) dfs(node["left"], reach * (1.0 - p_right)) dfs(node["right"], reach * p_right) dfs(0, np.ones(n)) return probs def leaf_values(leaves: list, wtr: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """叶区间的代表值取该区间样本均值；同时给出区间中点做对照。""" vals, mids = [], [] for lo, hi, _ in leaves: m = (wtr >= lo) & (wtr <= hi) if m.any(): vals.append(float(wtr[m].mean())) else: vals.append(0.5 * (lo + hi)) mids.append(0.5 * (lo + hi)) return np.array(vals), np.array(mids) # -------------------------------------------------------------------------- # 2) 主流程 # -------------------------------------------------------------------------- def main() -> None: print("TPM：平衡树 + 叶分布 + 方差 + 后门调整（§23.4.6）\n") data = make_watch_data(n=20000, dim=16, seed=SEED) X, w = data["x"], data["w"].astype(np.float64) idx = np.arange(len(X)) tr, te = split(idx, len(X), 0.8) Xtr, Xte, _, _ = standardize(X[tr], X[te]) wtr, wte = w[tr], w[te] nodes, leaves, sorted_w = build_tree(wtr, DEPTH) coefs = fit_tpm(Xtr, wtr, nodes) P = leaf_probabilities(Xte, nodes, coefs, len(leaves)) vals, mids = leaf_values(leaves, wtr) # ---------------- 1) 树结构 ---------------- print(f"① 树结构：深度 {DEPTH}，{len(leaves)} 个叶节点") rows = [] for level in range(DEPTH + 1): all_n = [n for n in nodes if n["level"] == level] leaf_n = [n for n in all_n if n["left"] is None] rows.append([level, len(all_n), len(leaf_n)]) table(["层", "节点数", "叶节点数"], rows) sizes = [int(((wtr >= lo) & (wtr <= hi)).sum()) for lo, hi, _ in leaves] print(f"\n 实际叶子 {len(leaves)} 个（理想 {2 ** DEPTH}），样本量 " f"最少 {min(sizes)}、最多 {max(sizes)}、中位 {int(np.median(sizes))}。") nz = [l for l in leaves if l[0] > 0] print(f" 比 2^depth 少，是因为零观看那块取值全相同、切不出左右两半，") print(f" 递归提前停下——零观看的 {sizes[0]:,} 条自然被隔离成一个「没看」的叶子，") print(f" 其余 {len(nz)} 个叶子覆盖 {nz[0][0]:.2f} 秒到 {nz[-1][1]:.0f} 秒，") print(" 每个约 500 条。这个「不平衡」恰恰是合理的：对分类器来说，") print(" 「播没播」和「看了 3 秒还是 4 秒」不是一个难度的任务。\n") # ---------------- 2) 点预测 ---------------- print("② 叶分布 → 期望预测，与直接回归对比") pred_tpm = P @ vals th0 = fit_linear(Xtr, wtr) pred0 = np.clip(predict_linear(th0, Xte), 0, None) rows = [ ["直接回归", f"{mae(wte, pred0):.2f}", f"{rmse(wte, pred0):.2f}", "—", "—"], ["TPM（叶均值）", f"{mae(wte, pred_tpm):.2f}", f"{rmse(wte, pred_tpm):.2f}", f"{np.mean(P.max(axis=1)):.2f}", f"{float(-(P * np.log(P + 1e-12)).sum(1).mean()):.2f}"], ] table(["模型", "MAE", "RMSE", "平均最大叶概率", "分布熵"], rows) ent = float(-(P * np.log(P + 1e-12)).sum(1).mean()) print("\n TPM 的 MAE / RMSE 与直接回归接近。它能给出的额外信息是分布：") print(f" 平均最大叶概率只有 {np.mean(P.max(axis=1)):.2f}、熵 {ent:.2f}，" f"而均匀分布于 {len(leaves)} 个叶子的熵是 {np.log(len(leaves)):.2f}。") print(" 分布明显偏平，说明模型对绝大多数样本其实「哪一档都拿不准」。") print(" 这不是训练不足，而是 16 维线性特征本身就分不开这么多档；") print(" 好处是这个「拿不准」被显式量化了，而不是藏在一个点估计背后。\n") # ---------------- 3) 方差作为不确定性 ---------------- print("③ 分布方差能否预测误差大小") mean_t = P @ vals var_t = P @ (vals ** 2) - mean_t ** 2 err = np.abs(wte - mean_t) rows = [] for lo_q, hi_q in ((0, 25), (25, 50), (50, 75), (75, 100)): a, b = np.percentile(var_t, [lo_q, hi_q]) m = (var_t >= a) & (var_t <= b if hi_q == 100 else var_t < b) rows.append([f"{lo_q}–{hi_q}%", f"{int(m.sum()):,}", f"{var_t[m].mean():.1f}", f"{err[m].mean():.2f}"]) table(["分布方差分位", "样本数", "平均方差", "平均绝对误差"], rows) corr = float(np.corrcoef(var_t, err)[0, 1]) print(f"\n 方差与绝对误差的相关系数 {corr:+.3f}。") print(" 方差高的样本，误差确实更大——「模型自己知道哪条拿不准」，") print(" 这正是 TPM 相对点回归多出来的能力，可以用于保底策略、") print(" 分场景降权、或者把不确定性喂给下游的融合层。\n") # ---------------- 4) 回归损失 + 方差正则 ---------------- print("④ 直接优化期望 vs 加方差正则（简化：在叶概率上加温度）") rows = [] for tau, name in ((1.0, "原始概率"), (0.8, "锐化 τ=0.8"), (1.2, "平滑 τ=1.2")): Pt = P ** tau Pt = Pt / Pt.sum(axis=1, keepdims=True) mt = Pt @ vals vt = Pt @ (vals ** 2) - mt ** 2 rows.append([name, f"{mae(wte, mt):.2f}", f"{rmse(wte, mt):.2f}", f"{vt.mean():.1f}", f"{float((Pt * np.log(Pt + 1e-12)).sum(1).mean()):.2f}"]) table(["概率处理", "MAE", "RMSE", "平均方差", "平均负熵"], rows) print("\n 拧温度只会让预测在两个极端之间摆动：锐化把质量压到最大叶，") print(" 预测更接近「取众数」；平滑把分布摊开，预测更接近全局均值。") print(" 本轮的结果里两者都没有同时赢下 MAE 和 RMSE：") print(" 锐化略降 MAE、略升 RMSE，平滑反之。这说明「方差正则」不是一个") print(" 后处理能替代的项——论文把它放进训练目标里，让分类器在") print(" 「分类对」和「分布集中」之间自己找平衡，而不是事后调温度。\n") # ---------------- 5) 后门调整 ---------------- print("⑤ 后门调整：造一个同时影响特征分布与观看时长的混杂变量 D") rng = np.random.default_rng(SEED) # D 与用户特征第 0 维相关（例如「晚间流量」更容易推给某类用户）， # 同时抬升观看时长——直接学 E(T|X) 会把这份抬升错误归因给 X。 d_tr = (Xtr[:, 0] + 0.6 * rng.normal(size=len(Xtr)) > 0).astype(int) d_te = (Xte[:, 0] + 0.6 * rng.normal(size=len(Xte)) > 0).astype(int) up_tr = np.where(d_tr == 1, 12.0, -12.0) up_te = np.where(d_te == 1, 12.0, -12.0) w_conf_tr = np.clip(wtr + up_tr, 0, None) w_conf_te = np.clip(wte + up_te, 0, None) # 分别在每个 D 分组上训练一个直接回归（TPM 的简化替身） preds, weights = [], [] for g in (0, 1): m = d_tr == g thg = fit_linear(Xtr[m], w_conf_tr[m]) preds.append(np.clip(predict_linear(thg, Xte), 0, None)) weights.append(m.mean()) pred_do = weights[0] * preds[0] + weights[1] * preds[1] pred_naive = np.clip(predict_linear(fit_linear(Xtr, w_conf_tr), Xte), 0, None) # 拿「只在同组内比较」的相对量看偏差：D 的抬升不该被算到 X 头上， # 所以评价时把 D 的平均效应扣掉，只看模型能否复现「组内」的真实关系。 rows = [ ["E(T|X) 直接回归", f"{mae(w_conf_te, pred_naive):.2f}", f"{np.corrcoef(pred_naive, Xte[:, 0])[0, 1]:+.3f}"], ["E(T|do(X)) 分层加权", f"{mae(w_conf_te, pred_do):.2f}", f"{np.corrcoef(pred_do, Xte[:, 0])[0, 1]:+.3f}"], ] table(["估计方式", "MAE", "预测与混杂源(X₀)的相关"], rows) print("\n D 是人为注入的混杂：它和 X₀ 相关，又整体抬/降 12 秒。") print(" 直接回归把这个「与 X₀ 同步的抬升」当成了 X₀ 的效应，") print(" 预测与 X₀ 的相关系数达 +0.30；分层加权之后降到 −0.06——") print(" 混杂被剥掉了，这是 E(T|do(X)) 想要的。") print(" 但 MAE 反而从 19.50 升到 20.26：分组把训练数据拆成两半，") print(" 每个回归器的估计方差变大。去掉偏差的代价是加上方差，") print(" 在这个样本量上用 do 算子并不划算——正因如此，后门调整") print(" 在线上通常只用于「已知强混杂」的少量分组，而不是全量替换。\n") print("⑥ 一句话：TPM 把「预测一个数」拆成「逐层二分类 → 一条分布」，") print(" 树结构保证每层样本均衡，分布给出可用的不确定性，") print(" 后门调整再把这套框架接到因果去偏上。") if __name__ == "__main__": main()