目录
- 1. Top-p 和 Top-k 采样
- 2. LayerNorm 和 RMSNorm
1. Top-p 和 Top-k 采样
概念讲解：
在自回归文本生成中，模型每一步会输出一个概率分布（logits 经过 softmax），我们需要从中采样下一个 token。直接使用整个词汇表采样（即 temperature 缩放后的随机采样）可能导致生成低概率 token，使结果不连贯。Top-k 采样 和 Top-p 采样 是两种常用的截断采样方法，用于限制候选 token 集合，提高生成质量。
- Top-k 采样：
- 做法：只保留概率最高的 k 个词，把剩下的词概率强制设为 0，然后重新归一化（让剩下的概率和为 1），再从中采样。
- 作用：直接砍掉长尾的低概率词，防止生成生僻字或乱码。
- 缺点：k 是固定的。如果模型很自信（某个词概率 90%），k 太大也会采样到噪音；如果模型很犹豫（概率很平），k 太小会限制多样性。
- Top-p (Nucleus) 采样：
- 做法：将词按概率从大到小排序，依次累加概率，直到累加和超过 p (比如 0.9)。保留这些词，剩下的截断，重新归一化，再采样。
- 作用：动态调整候选词数量。模型自信时候选词少，模型犹豫时候选词多。
- 现状：目前 LLM 推理中，Top-p 比 Top-k 更常用，或者两者结合。
两种方法可以结合使用（如先取 top-k 再取 top-p），但通常分别实现。
代码实现：
def top_k_filtering(logits, top_k=50, temperature=1.0, filter_value=-float('Inf')):"""logits: [vocab_size] 或 [batch, vocab_size]，模型输出的原始分数top_k: 保留概率最大的 k 个词temperature: 温度，大于 1 增加多样性，小于 1 增加确定性"""logits = logits / temperature# 找出所有小于第 k 大值的索引，将其设为 filter_value (即 -inf，softmax 后为 0)indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]logits[indices_to_remove] = filter_valuereturn logitsdef top_p_filtering(logits, top_p=0.9, temperature=1.0, filter_value=-float('Inf'), min_tokens_to_keep=1):"""logits: [vocab_size] 或 [batch, vocab_size]，模型输出的原始分数top_p: 保留累积概率超过 p 的最小词集temperature: 温度，大于 1 增加多样性，小于 1 增加确定性"""logits = logits / temperature# 按概率从大到小排序sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) # [batch, vocab]cumulated_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # [batch, vocab]# 创建 mask：累积概率 > p 的位置为 True（需要被移除）sorted_indices_to_remove = cumulated_probs > top_p # [batch, vocab]sorted_indices_to_remove[:, :min_tokens_to_keep] = False # 保留第一个超过 p 的 token（确保至少有一些 token）# 注意：sorted_logits 是排序后的，需要映射回原始位置indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)logits[indices_to_remove] = filter_value # 将被移除的 logits 设为 -inf（softmax 后概率为 0）return logitsif __name__ == "__main__":# 模拟一个 vocab_size=5 的 logitslogits = torch.tensor([[2.0, 1.0, 0.1, 0.1, 0.1]]) print("原始 Logits:", logits)filtered_logits = top_p_filtering(logits, top_p=0.8, temperature=1.0)print("过滤后 Logits:", filtered_logits)# 采样probs = F.softmax(filtered_logits, dim=-1)next_token = torch.multinomial(probs, num_samples=1)print("采样结果 Token ID:", next_token)
2. LayerNorm 和 RMSNorm
概念讲解：
- Layer Normalization (LayerNorm)：对每个样本的每个特征层进行归一化。给定输入
x形状(batch, seq_len, hidden_size)，对最后一个维度（hidden_size）计算均值和方差，然后标准化：(x - mean) / sqrt(var + eps)，再乘以可学习的缩放参数gamma并加上偏移beta。LayerNorm 广泛用于 Transformer，稳定训练。 - RMSNorm：Root Mean Square Layer Normalization 是 LayerNorm 的一个简化变体。它假设减去均值不是必需的，只使用均方根 (RMS) 进行归一化：
x / RMS(x)，其中RMS(x) = sqrt(mean(x^2) + eps)。同样乘以可学习的缩放参数gamma，但没有beta。RMSNorm 计算量更小，且在实验中性能与 LayerNorm 相当。
代码实现：
import torch
import torch.nn as nnclass LayerNorm(nn.Module):"""层归一化 (Layer Normalization)公式: y = gamma * (x - mean) / sqrt(var + eps) + beta"""def __init__(self, hidden_size, eps=1e-5):super().__init__()self.gamma = nn.Parameter(torch.ones(hidden_size))self.beta = nn.Parameter(torch.ones(hidden_size))self.eps = epsdef forward(self, x):# x: (batch, seq_len, hidden_size) 或 (batch, hidden_size)mean = x.mean(dim=-1, keepdim=True)var = x.var(dim=-1, keepdim=True, unbiased=False)x_norm = (x - mean) / torch.sqrt(var + self.eps)return self.gamma * x_norm + self.betaclass RMSNorm(nn.Module):"""RMS 归一化 (Root Mean Square Layer Normalization)公式: y = gamma * x / RMS(x), 其中 RMS(x) = sqrt(mean(x^2) + eps)"""def __init__(self, hidden_size, eps=1e-5):super().__init__()self.gamma = nn.Parameter(torch.ones(hidden_size))self.eps = epsdef forward(self, x):# x: (batch, seq_len, hidden_size) 或 (batch, hidden_size)rms = torch.sqrt(torch.mean(x, dim=-1, keepdim=True) + self.eps)x_norm = x / rmsreturn self.gamma * x_norm# 示例用法
if __name__ == "__main__":batch, seq, hidden = 2, 3, 4x = torch.randn(batch, seq, hidden)ln = LayerNorm(hidden)rms = RMSNorm(hidden)out_ln = ln(x)out_rms = rms(x)print("LayerNorm 输出形状:", out_ln.shape)print("RMSNorm 输出形状:", out_rms.shape)