企业官网建设流程全解析

LLM 解码里的 Top-K，通常可以拆成两步：先保留分数最高的一组候选，再在候选集合上重新归一化概率。本文只实现这个概率变换，不实际随机抽取 token，也不运行语言模型。

这篇笔记来自我整理 AI 编程练习时对数值实现与测试的检查。所有代码和断言都在正文里，可以直接在本地 Python 环境运行。

一、先把接口约定写清楚

输入logits是非空有限实数列表；k是1 <= k <= len(logits)的整数。

这里额外约定：

1. 精确保留 k 个 token 索引。
2. logit 相同时，优先较小的 token 索引。
3. 只对保留的候选计算 Softmax。
4. 返回值与原词表等长，被排除位置的概率为 0。
5. 输出顺序与原词表一致，不修改输入。

同分规则属于这个练习的接口约定，不能假定所有框架都采用相同策略。下面的实现针对已满足这些前置条件的输入，不包含额外的输入校验。

二、为何不能只把尾部概率清零？

以logits = [2, 4, 1, 3]、k = 2为例，选中的索引是 1 和 3。

先对完整词表做 Softmax，可得到近似分布：

[0.087144, 0.643914, 0.032059, 0.236883]

如果只将索引 0、2 清零：

[0, 0.643914, 0, 0.236883]

保留下来的概率总和约为0.880797，还不是归一化的候选分布。除以这部分保留概率质量之后，才得到：

[0, 0.731059, 0, 0.268941]

对保留集合直接做 Softmax，与对完整分布截断后重新归一化，在精确算术下等价。直接实现前者更容易对齐本题的接口。

三、稳定 Softmax：先减去最大 logit

直接计算math.exp(1000)会溢出，但同一个常数平移所有 logits，不改变归一化结果。

用 S 表示保留索引集合，令m = max(logits[i] for i in S)。对每个保留位置计算exp(logits[i] - m)，再除以这些权重的总和。每个指数的输入都不大于 0，最大候选的权重为 1。

极小的相对权重仍可能在浮点计算中下溢为零，因此“被选中”和“显示出的概率大于零”不是在所有浮点输入上都能互相替代的判据。选择集合应由排序决定。

四、完整 Python 实现

importmathdeftop_k_probabilities(logits,k):selected=sorted(range(len(logits)),key=lambdai:(-logits[i],i),)[:k]peak=max(logits[i]foriinselected)weights=[math.exp(logits[i]-peak)foriinselected]total=sum(weights)result=[0.0]*len(logits)forindex,weightinzip(selected,weights):result[index]=weight/totalreturnresult

排序键(-logits[i], i)同时表达了两个要求：先按 logit 从大到小，再按索引从小到大。排序的是索引，原始logits列表保持不变。

词表大小为 V 时，这个版本的时间复杂度为O(V log V)，额外空间为O(V)。大型框架可能使用其他选择算法；这里优先让实现与约定容易检查。

五、测试不仅要覆盖普通输入

下面的测试包含 Top-1、全词表、同分边界、大 logits、原顺序和输入不变性。

defcheck_close(actual,expected,tolerance=1e-12):assertlen(actual)==len(expected)assertall(abs(a-b)<=tolerancefora,binzip(actual,expected))# Top-1：只有最大候选获得概率 1。asserttop_k_probabilities([1,3,2],1)==[0.0,1.0,0.0]# 边界同分：保留索引 1、2，而不是三个同分候选都保留。asserttop_k_probabilities([0,2,2,2],2)==[0.0,0.5,0.5,0.0]# k 等于词表大小：退化为完整 Softmax。asserttop_k_probabilities([0,0,0],3)==[1/3,1/3,1/3]# 大 logits：稳定求指数，概率顺序仍正确。p=top_k_probabilities([1000,1001,999],2)assertp[2]==0.0assertp[1]>p[0]>0.0assertabs(sum(p)-1.0)<1e-12# 统一平移：保留集合和概率不变。check_close(top_k_probabilities([2,4,1,3],2),top_k_probabilities([1002,1004,1001,1003],2),)# 返回原词表顺序，且不修改输入。x=[2,4,1,3]before=x.copy()p=top_k_probabilities(x,2)assertx==beforeassertp[0]==p[2]==0.0assertp[1]>p[3]>0.0# 极小权重可能下溢为零，但输出仍是有效的归一化结果。asserttop_k_probabilities([0,-1000],2)==[1.0,0.0]

这里的check_close使用逐元素绝对误差；概率限制在 0 到 1 范围内，这个小示例使用统一的1e-12容差即可。工程代码仍应按数据类型和任务要求选择容差。

六、用一个反例抓住“只清零尾部”的错误

测试普通值时，要主动选择被删掉的尾部仍有明显概率质量的输入。下面这个错误版本在 k 等于词表大小时看起来正常，但截断时会露出问题。

defwrong_zero_tail(logits,k):peak=max(logits)weights=[math.exp(value-peak)forvalueinlogits]total=sum(weights)full=[weight/totalforweightinweights]selected=set(sorted(range(len(logits)),key=lambdai:(-logits[i],i))[:k])return[pifiinselectedelse0.0fori,pinenumerate(full)]bad=wrong_zero_tail([2,4,1,3],2)good=top_k_probabilities([2,4,1,3],2)assert0.88<sum(bad)<0.89assertabs(sum(good)-1.0)<1e-12assertany(abs(a-b)>0.01fora,binzip(bad,good))

同理，只检查总概率约为 1 也不充分：一个“永远给最大项概率 1”的错误实现，照样满足归一化。需要同时检查正确的候选集合、被排除位置，以及多个保留候选之间的相对权重。

参考资料

• Hugging Face 官方图解：How to generate text — Top-K Sampling。它介绍候选筛选后的概率重分配，并继续执行文本采样；本文只计算概率。
• Python 官方文档：math.exp 与 排序稳定性。
• 参考题目来源：实现稳定的 Top-K Token 概率。