目录21. 解码策略应用温度调节复盘Top-K 截断复盘Top-P 核采样截断复盘正确的完整代码测试实现结果学习收获AI帮我总结的21. 解码策略应用温度调节我的第一次尝试defapply_temperature(logits:torch.Tensor,temperature:float)-torch.Tensor: 应用温度调节。注意通常 T1.0 意味着不改变T 越接近 0 越确定Greedy。 # # TODO 1: 温度下限与缩放# # temp ???temp1e-4returnlogits/tempAI 纠错之后应该这么写defapply_temperature(logits:torch.Tensor,temperature:float)-torch.Tensor:# 用传入的 temperature,只在它为 0 时用 1e-6 兜底tempmax(temperature,1e-6)returnlogits/temp复盘「防除零」这个工程坑抓到了。T0 会让 logits / 0 直接 inf整个 softmax 全 NaN。但是我压根没用传进来的 temperature 参数。。。审题不仔细啊纯粹低级错误。打比方用户调函数说我要 0.7 度的温水我永远给 1e-4 度的冰水。这等于直接把 Temperature 这个策略废掉了 —— 不管传什么 T行为都一样。而且 1e-4 太极端了logits [2.3, 2.0] 除以 1e-4 [23000, 20000]再softmax 后 ≈ [1.0, 0.0]永远贪心。参考答案用的是 1e-6只在 T0 时兜底正常调用时用真实值。Top-K 截断我的第一次尝试defapply_top_k(logits:torch.Tensor,top_k:int)-torch.Tensor: Top-K 截断。只保留值最大的 top_k 个其余置为 -inf。 iftop_k0ortop_klogits.size(-1):returnlogits# # TODO 2: Top-K 截断# # filter_value ???# kth_values ???# logits ???filter_valuefloat(-inf)kth_values,_torch.topk(logits,top_k,dim-1)kth_valueskth_values[-1]logitstorch.where(foriinlogits:ikth_values,i,-inf)returnlogitslogits这行有提示也写不对没写过不太清楚语法。AI 纠错之后应该这么写filter_valuefloat(-inf)# 1. 取 top_k 个最大值,kth_values 形状 [B, K]kth_values,_torch.topk(logits,top_k,dim-1)# 2. 取每行第 K 大的值(最后一个),保持维度 [B, 1] 以便广播kth_valueskth_values[...,-1:]# 3. 小于阈值的置 -inf,大于等于的保留原值logitstorch.where(logitskth_values,torch.tensor(filter_value,devicelogits.device),logits)复盘kth_values[-1] 取错了维度。假设 logits 是 [B2, V10]topk(…, dim-1, k3) 返回的 kth_values 是 [2, 3] —— 每行 3 个最大值。但是 kth_values[-1]直接取 batch 维度的最后一行形状变成 [3]肯定不对。实际想要的是每行的第 3 大值即 topk 维度的最后一个应该用 […, -1:] 保持维度。因为后面要和 [2, 10] 的 logits 做 torch.where 广播必须留一个能广播的形状 [2, 1]。for i in logits 直接语法都不对了。torch.where 接受布尔张量作为条件不是 Python for 循环。PyTorch 是整张量并行操作禁止循环。‘-inf’ 是字符串。torch.where 要的是同 dtype 的张量不是字符串 ‘-inf’。要写成 torch.tensor(float(‘-inf’)) 或直接 float(‘-inf’)。Top-P 核采样截断我的第一次尝试defapply_top_p(logits:torch.Tensor,top_p:float)-torch.Tensor: Top-p (Nucleus) 核采样截断。 iftop_p0.0ortop_p1.0:returnlogits# 1. 首先需要将 logits 从大到小排序# 注意我们需要记住原始的索引 (indices)因为截断完了还要把它复原回原来的位置sorted_logits,sorted_indicestorch.sort(logits,descendingTrue)# 2. 对排序后的概率 (需要先算一遍 Softmax) 计算累加和 (Cumulative Probability)cumulative_probstorch.cumsum(F.softmax(sorted_logits,dim-1),dim-1)# # TODO 3: Top-p 核心逻辑# # sorted_indices_to_remove ???# sorted_logits[sorted_indices_to_remove] ???# restored_logits ???sorted_indices_to_removetorch.where(cumulative_probstop_p,True,False)sorted_indices_to_remove[...,1:]sorted_indices_to_remove[...,0:-1].clone()sorted_indices_to_remove[...,0]Falsesorted_logits[sorted_indices_to_remove]torch.where(sorted_logitstop_p,torch.tensor(float(-inf),devicesorted_logits.device),sorted_logits)restored_logitstorch.zeros_like(logits).scatter_(dim-1,indexsorted_indices,srcsorted_logits)returnrestored_logitsAI 纠错之后应该这么写defapply_top_p(logits:torch.Tensor,top_p:float)-torch.Tensor:iftop_p0.0ortop_p1.0:returnlogits sorted_logits,sorted_indicestorch.sort(logits,descendingTrue)cumulative_probstorch.cumsum(F.softmax(sorted_logits,dim-1),dim-1)# 1. 找出累积概率超过 top_p 的位置sorted_indices_to_removecumulative_probstop_p# 2. 向右平移一位,保留首次超过阈值的那个 tokensorted_indices_to_remove[...,1:]sorted_indices_to_remove[...,:-1].clone()sorted_indices_to_remove[...,0]False# 3. 直接用 mask 赋 -inf(不需要 torch.where)sorted_logits[sorted_indices_to_remove]float(-inf)# 4. 按原索引恢复顺序restored_logitstorch.zeros_like(logits).scatter_(dim-1,indexsorted_indices,srcsorted_logits)returnrestored_logits复盘首先是代码写的有点冗余让人一看是外行。然后是冗余的代码把自己脑子搞晕拿 logits 和 top_p 比大小维度都不一样也没有归一化。sorted_logits 是未归一化的分数比如 [4.0, 3.1, 2.3, …]可以是任意实数。top_p 是概率阈值范围 (0, 1)比如 0.85。sorted_logits top_p 等于在问4.0 0.85 吗这样永远 False所以 torch.where 永远走第二个分支等于啥都没删。然后boolean mask 赋值不该嵌套 torch.where。sorted_logits[mask] torch.where(...)左边 sorted_logits[mask] 是一维张量只取 True 的位置右边 torch.where(…) 返回完整形状。形状对不上。直接sorted_logits[sorted_indices_to_remove] float(-inf)这样布尔 mask 索引赋值会自动只对 True 的位置赋值。正确的完整代码 Task 3 - 21. Decoding Strategies | 解码策略 完整实现 测试用例 学习路径: 1. apply_temperature - 温度缩放(防除零) 2. apply_top_k - Top-K 截断(torch.topk torch.where) 3. apply_top_p - Top-p 核采样(sort cumsum 平移掩码 scatter_) 4. decode_next_token - 组合三策略 multinomial 采样 importtorchimporttorch.nn.functionalasFdefapply_temperature(logits:torch.Tensor,temperature:float)-torch.Tensor:应用温度调节。T1 更确定,T1 更随机。# 防除零:用 1e-6 兜底,正常调用用真实值tempmax(temperature,1e-6)returnlogits/tempdefapply_top_k(logits:torch.Tensor,top_k:int)-torch.Tensor:Top-K 截断。只保留值最大的 top_k 个,其余置 -inf。iftop_k0ortop_klogits.size(-1):returnlogits filter_valuefloat(-inf)# 1. 取 top_k 个最大值,形状 [B, K]kth_values,_torch.topk(logits,top_k,dim-1)# 2. 取每行第 K 大的值(最后一个),保持维度 [B, 1] 以便广播kth_valueskth_values[...,-1:]# 3. 小于阈值的置 -inf,大于等于的保留logitstorch.where(logitskth_values,torch.tensor(filter_value,devicelogits.device,dtypelogits.dtype),logits,)returnlogitsdefapply_top_p(logits:torch.Tensor,top_p:float)-torch.Tensor:Top-p (Nucleus) 核采样截断。iftop_p0.0ortop_p1.0:returnlogits# 1. 降序排序 记住原始索引(后面要恢复)sorted_logits,sorted_indicestorch.sort(logits,descendingTrue)# 2. 计算累积概率cumulative_probstorch.cumsum(F.softmax(sorted_logits,dim-1),dim-1)# 3. 找出累积概率超过 top_p 的位置sorted_indices_to_removecumulative_probstop_p# 4. 向右平移一位,保留首次超过阈值的那个 tokensorted_indices_to_remove[...,1:]sorted_indices_to_remove[...,:-1].clone()sorted_indices_to_remove[...,0]False# 5. 直接用 mask 赋 -infsorted_logits[sorted_indices_to_remove]float(-inf)# 6. 按原索引恢复顺序restored_logitstorch.zeros_like(logits).scatter_(dim-1,indexsorted_indices,srcsorted_logits)returnrestored_logitsdefdecode_next_token(logits:torch.Tensor,temperature0.7,top_k50,top_p0.9):组合三策略 multinomial 采样。logitsapply_temperature(logits,temperature)logitsapply_top_k(logits,top_k)logitsapply_top_p(logits,top_p)probsF.softmax(logits,dim-1)next_tokentorch.multinomial(probs,num_samples1)returnnext_token# # 测试# deftest_decoding():torch.manual_seed(42)vocab_size10logitstorch.tensor([[0.1,2.3,0.4,1.2,-0.5,4.0,3.1,0.0,1.1,-1.0]])print(原始 Logits (前10个单词):,logits.squeeze().tolist())print()# 1. Temperaturet_logitsapply_temperature(logits.clone(),0.5)asserttorch.allclose(t_logits[0,5]-t_logits[0,6],(logits[0,5]-logits[0,6])*2),温度调节错误!print(✅ Temperature 热量调节通过!)print(f T0.5 时 logits[5] - logits[6] {t_logits[0,5].item():.2f}-{t_logits[0,6].item():.2f}{(t_logits[0,5]-t_logits[0,6]).item():.2f})print(f 原始差距 {(logits[0,5]-logits[0,6]).item():.2f}, 缩放后 2x {((logits[0,5]-logits[0,6])*2).item():.2f})print()# 2. Top-Kk_logitsapply_top_k(logits.clone(),3)valid_count(k_logits!float(-inf)).sum().item()assertvalid_count3,fTop-K 截断错误,保留了{valid_count}个valid_indices(k_logits!float(-inf)).nonzero(as_tupleTrue)[1].tolist()print(✅ Top-K 暴力截断通过!)print(f 保留的 token 索引:{valid_indices}(应为 [1, 5, 6] - 三个最大值))print()# 3. Top-pp_logitsapply_top_p(logits.clone(),0.8)valid_count(p_logits!float(-inf)).sum().item()assertvalid_count3,fTop-p 截断错误,保留了{valid_count}个valid_indices(p_logits!float(-inf)).nonzero(as_tupleTrue)[1].tolist()print(✅ Top-p (Nucleus) 核采样动态截断通过!)print(f 保留的 token 索引:{valid_indices}(应为 [1, 5, 6] - 累积概率 0.540.220.100.86 刚好超过 0.8))print()# 4. 完整管线next_tokendecode_next_token(logits.clone(),temperature0.7,top_k50,top_p0.9)assertnext_token.shape(1,1),解码的词张量维度不对print(f✅ All Tests Passed! 完整管线通过。)print(f 本次采样的下一个 token ID:{next_token.item()})print()# 5. 直观演示:不同温度下的采样分布print(*60)print( 不同温度下的采样分布(跑 1000 次统计):)print(*60)fortempin[0.3,0.7,1.0,1.5]:torch.manual_seed(0)countstorch.zeros(vocab_size)for_inrange(1000):tokdecode_next_token(logits.clone(),temperaturetemp,top_k50,top_p0.9)counts[tok.item()]1# 显示前 5 个最常被采样的 tokentop5counts.topk(5)dist_str, .join(fidx{int(i)}:{int(c)}forc,iinzip(top5.values,top5.indices))print(f T{temp:4}|{dist_str})if__name____main__:test_decoding()测试实现结果原始 Logits(前10个单词):[0.10000000149011612,2.299999952316284,0.4000000059604645,1.2000000476837158, -0.5,4.0,3.0999999046325684,0.0,1.100000023841858, -1.0]✅ Temperature 热量调节通过!T0.5时 logits[5]- logits[6]8.00-6.201.80原始差距0.90, 缩放后2x1.80✅ Top-K 暴力截断通过!保留的 token 索引:[1,5,6](应为[1,5,6]- 三个最大值)✅ Top-p(Nucleus)核采样动态截断通过!保留的 token 索引:[1,5,6](应为[1,5,6]- 累积概率0.540.220.100.86刚好超过0.8)✅ All Tests Passed!完整管线通过。 本次采样的下一个 token ID:5 不同温度下的采样分布(跑1000次统计):T0.3|idx5:1000, idx7:0, idx6:0, idx9:0, idx8:0T0.7|idx5:785, idx6:215, idx7:0, idx9:0, idx8:0T1.0|idx5:604, idx6:255, idx1:104, idx3:37, idx8:0T1.5|idx5:427, idx6:253, idx1:132, idx3:77, idx8:72学习收获AI帮我总结的三种策略改的都是 logits不是概率所有截断最后都靠 -inf softmax 归一化实现管线顺序Temperature → Top-K → Top-p → Softmax → Multinomial不能乱。[…, -1:] 保持维度的切片技巧比 squeeze/unsqueeze 更优雅面试高频Top-p 的三大坑平移掩码 / 保第 0 位 / scatter_ 恢复顺序缺一不可-inf 是 Transformer 的通用屏蔽语言mask、padding、Top-K、Top-p 都靠它
