Logits、Softmax 与基础采样

从 hidden state 到一个 token id

模型最后一层 hidden state $h_t\in\mathbb{R}^{B\times H}$ 经 vocabulary projection 得到：

z_t=h_tW_{vocab}+b,\qquad z_t\in\mathbb{R}^{B\times V}

$z_t$ 是 logits：可为任意实数、未归一化，也不是概率。稳定 softmax 计算：

p_i=\frac{\exp(z_i-m)}{\sum_j\exp(z_j-m)},\qquad m=\max_j z_j

减去 $m$ 不改变分布，因为分子分母同时乘以同一常数，却能避免 exp(1000) overflow。Prefill/Decode 的每次模型前向产生下一位置 logits；sampling 决定实际追加哪个 token id。

Temperature、Top-k、Top-p 的准确语义

Greedy： $\arg\max_i z_i$ ，没有随机采样；“temperature=0”通常是框架的 greedy 分支，不应真的计算 $z/0$ 。
Temperature $\tau>0$ ：先用 $z/\tau$ 。 $\tau<1$ 通常让分布更尖， $\tau>1$ 更平；它不改变 logit 排名。
Top-k：只保留当前 logits 最大的 $k$ 个 token，再重新归一化。
Top-p / nucleus：按概率从大到小选择累计概率达到阈值 $p$ 的最小前缀，再重新归一化。保留数量随分布变化，不等同固定 $k$ 。

过滤顺序是算法定义的一部分。常见实现先 temperature，再 top-k/top-p；presence/frequency penalty、repetition penalty、bad words、grammar mask 等若加入，会在某个明确顺序上修改 logits。不同服务的同名参数不保证完全相同的处理顺序。

可运行实现

import platform
import torch

torch.manual_seed(41)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float32

def stable_softmax(logits):
    shifted=logits-logits.max(-1,keepdim=True).values
    exp=shifted.exp(); return exp/exp.sum(-1,keepdim=True)

def filter_logits(logits, temperature=1.0, top_k=None, top_p=None):
    assert temperature>0
    filtered=logits/temperature
    if top_k is not None:
        cutoff=filtered.topk(top_k,-1).values[...,-1,None]
        filtered=filtered.masked_fill(filtered<cutoff,-torch.inf)
    if top_p is not None:
        assert 0<top_p<=1
        sorted_logits,indices=filtered.sort(-1,descending=True)
        sorted_probs=stable_softmax(sorted_logits)
        # 保留跨过阈值的那个 token，所以先减当前 token 概率。
        remove=sorted_probs.cumsum(-1)-sorted_probs>top_p
        sorted_logits=sorted_logits.masked_fill(remove,-torch.inf)
        filtered=torch.full_like(filtered,-torch.inf).scatter(-1,indices,sorted_logits)
    return filtered

def main():
    logits=torch.tensor([[1000.,999.,997.,990.,980.]],device=DEVICE)
    probs=stable_softmax(logits)
    torch.testing.assert_close(probs,torch.softmax(logits,-1))
    torch.testing.assert_close(probs.sum(-1),torch.ones(1,device=DEVICE))
    top_k_logits=filter_logits(logits,temperature=.7,top_k=3)
    top_p_logits=filter_logits(logits,top_p=.90)
    assert torch.isfinite(top_k_logits).sum().item()==3
    generator=torch.Generator(device=DEVICE).manual_seed(42)
    sample=torch.multinomial(stable_softmax(top_p_logits),1,generator=generator)
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} logits={tuple(logits.shape)}")
    print(f"prob_sum={probs.sum().item():.8f} greedy={logits.argmax(-1).item()} sampled={sample.item()}")
    print(f"top_k_kept={torch.isfinite(top_k_logits).sum().item()}")
    print(f"top_p_kept={torch.isfinite(top_p_logits).sum().item()}")
    e1=torch.distributions.Categorical(logits=logits/.7).entropy()
    e2=torch.distributions.Categorical(logits=logits/1.5).entropy()
    print(f"entropy_tau_0.7={e1.item():.6f} entropy_tau_1.5={e2.item():.6f}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/sampling_reference.py。实际运行中稳定 softmax 概率和为 1.00000000；top-k 保留 3 个，top-p=0.90 保留 2 个；temperature 从 0.7 增至 1.5 时该例熵从 0.546166 增至 0.877431。

成本、精度与可复现性

Vocabulary projection 常是 $[B,H]\times[H,V]$ ，大词表下权重读取和 GEMM/GEMV 可能不可忽略。softmax、排序或 top-k selection 的成本随 $V$ 增长；生产 kernel 会避免完整通用排序，或融合 penalty/filter/sampling。

FP16/BF16 logits、并行归约顺序、backend、相等 logits 的 tie-breaking 和随机数流都会影响 token 级复现。即使使用同一 seed，不同 batch 组成或提前结束的序列也可能消费不同数量随机数。确定分布语义与 bitwise identical 是两个标准。

常见错误包括：

直接 exp(logits) 导致 overflow；
top-p 在累计概率超过阈值前就移除当前 token，导致集合少一个边界元素；
在过滤后不重新归一化就把值当概率；
把 temperature 当作改变词表、模型权重或事实正确性的保证；
用 argmax 实验声称验证了随机 sampling 分布。

与高级解码的边界

Speculative decoding 的 exact sampling 不仅是“draft top-k 后 target 选一个”：它要用 $\min(1,p(y)/q(y))$ 接受候选，拒绝时从校正分布采样。MTP 训练多个未来 token 目标也不自动改变最终 sampling 协议。这些高级方法必须从本页的目标分布 $p$ 出发证明语义。

Beam Search、约束解码与 Tree/Self Speculation进一步区分序列级搜索、按状态改变合法 token support，以及共享前缀候选的 tentative/committed 状态；它们不能都叫“sampling 优化”。

上下游关系

Prefill/Decode 产生每一步 logits，并把选出的 token 送回下一 decode step。
Tokenizer/Chat Template 定义 vocabulary ids 与 special-token 语义；采到 EOS id 可能结束请求，但 stop string 仍需 detokenization 层处理。
Speculative Decoding以本页概率为 target $p$ ，区分经验分布一致、接受长度和 target 串行步数。

参考资料

Holtzman et al., The Curious Case of Neural Text Degeneration — nucleus sampling。
PyTorch — softmax
PyTorch — multinomial