采样过程源码解析：从 logits 到 token 的采样策略
采样过程源码解析：从 logits 到 token 的采样策略
在自然语言处理（NLP）领域，尤其是生成式模型中，采样过程是将模型输出的 logits 转换为实际可读的 token 的关键步骤。这一过程不仅决定了生成文本的多样性，还影响着模型输出的质量和实用性。本文将深入解析采样过程的源码实现，探讨从 logits 到 token 的多种采样策略，帮助读者更好地理解这一核心环节。
1. Logits 的生成与理解
在深度学习模型中，尤其是基于 Transformer 的架构，模型的最后一层通常会输出一个形状为 (batch_size, sequence_length, vocab_size) 的张量，其中 vocab_size 是词汇表的大小。这个张量中的每个值，我们称之为 logits，代表了模型对每个位置上每个可能 token 的预测分数。
Logits 本身并不直接代表概率，它们需要通过 softmax 函数进行归一化，转换为概率分布。然而，在采样过程中，我们并不总是直接使用 softmax 后的概率，而是基于 logits 应用各种采样策略，以平衡生成文本的准确性和多样性。
2. 采样策略概览
采样策略的选择直接影响生成文本的质量。常见的采样策略包括贪婪采样（Greedy Sampling）、随机采样（Random Sampling）、温度采样（Temperature Sampling）、Top-k 采样（Top-k Sampling）和 Top-p（Nucleus）采样（Top-p Sampling）。下面，我们将逐一解析这些策略的源码实现。
2.1 贪婪采样
贪婪采样是最简单的采样策略，它选择每个位置上概率最高的 token 作为输出。虽然这种方法能保证生成文本的确定性，但往往缺乏多样性，容易陷入重复或模式化的输出。
import torch import torch.nn.functional as F defgreedy_sample(logits):# 应用 softmax 获取概率分布 probs = F.softmax(logits, dim=-1)# 选择每个位置上概率最高的 token _, sampled_tokens = torch.max(probs, dim=-1)return sampled_tokens 2.2 随机采样
随机采样，也称为多项式采样，根据 softmax 后的概率分布随机选择 token。这种方法能增加输出的多样性，但也可能导致生成不连贯或无意义的文本。
defrandom_sample(logits): probs = F.softmax(logits, dim=-1)# 根据概率分布随机采样 sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)return sampled_tokens 2.3 温度采样
温度采样通过调整 softmax 函数的“温度”参数来控制输出的多样性。温度参数 T 越大，概率分布越平滑，采样结果越多样；T 越小，概率分布越尖锐，采样结果越接近贪婪采样。
deftemperature_sample(logits, temperature=1.0):# 调整 logits 的“温度” adjusted_logits = logits / temperature probs = F.softmax(adjusted_logits, dim=-1) sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)return sampled_tokens 2.4 Top-k 采样
Top-k 采样限制了采样的范围，只从概率最高的 k 个 token 中进行选择。这种方法能在保证一定多样性的同时，避免选择到概率极低、可能不合理的 token。
deftop_k_sample(logits, k=10): probs = F.softmax(logits, dim=-1)# 获取概率最高的 k 个 token 的索引 top_k_probs, top_k_indices = torch.topk(probs, k=k, dim=-1)# 重新归一化 top-k 概率 top_k_probs /= top_k_probs.sum(dim=-1, keepdim=True)# 从 top-k token 中随机采样 sampled_tokens = torch.multinomial(top_k_probs, num_samples=1).squeeze(-1)# 将采样到的索引映射回原始 token# 注意：这里需要额外的步骤来获取实际 token，简化示例中省略return sampled_tokens # 实际实现中需处理索引映射2.5 Top-p（Nucleus）采样
Top-p 采样，也称为 Nucleus 采样，是一种更灵活的采样策略。它选择概率累积和超过预设阈值 p 的最小 token 集合，然后在这个集合中进行随机采样。这种方法能自适应地调整采样范围，既保证了多样性，又避免了低概率 token 的干扰。
deftop_p_sample(logits, p=0.9): probs = F.softmax(logits, dim=-1) sorted_probs, sorted_indices = torch.sort(probs, descending=True) cumulative_probs = torch.cumsum(sorted_probs, dim=-1)# 找到满足累积概率 >= p 的最小索引 mask = cumulative_probs < p # 在满足条件的 token 中随机采样# 注意：这里需要处理边界情况，简化示例中省略# 实际应用中，可能需要更复杂的逻辑来确保至少选择一个 token selected_indices = sorted_indices[mask] selected_probs = sorted_probs[mask]# 重新归一化选中的概率 selected_probs /= selected_probs.sum(dim=-1, keepdim=True) sampled_tokens = torch.multinomial(selected_probs, num_samples=1).squeeze(-1)# 将采样到的索引映射回原始 token# 注意：这里需要额外的步骤来获取实际 token，简化示例中省略return sampled_tokens # 实际实现中需处理索引映射3. 结论
采样过程是自然语言生成模型中的关键环节，它直接决定了生成文本的质量和多样性。从简单的贪婪采样到复杂的 Top-p 采样，每种策略都有其独特的优势和适用场景。在实际应用中，我们往往需要根据具体任务的需求，灵活选择或组合这些采样策略，以达到最佳的生成效果。
通过深入解析这些采样策略的源码实现，我们不仅能更好地理解它们的工作原理，还能为模型优化和定制提供有力的支持。希望本文能为读者在自然语言生成领域的探索提供有益的参考。
