Transformer输出解码
在之前的学习中，我们已经在从零搭建Transformer-CS336作业以及cs336Assignment1 训练组件中分别实现了一个完整的Decoder-only的Transfomer模型，同时搭建了用于训练模型的各个组件，包括数据加载，优化器，损失函数等等，并构建了一个完整的训练流程。
后面将实现如何实现将输入给模型后，对模型的输出进行解码，得到我们期望的单个文字输出，进而得到输出的完整句子。
当然，我们已经知道，对于我们实现的Transformer，在进行完Embedding后，输入维度为为\((B, S, D)\)，输出维度为为\((B, S, V)\)，其中：
- \(B\)：批次大小，Batch size
- \(S\)：序列长度，Sequence length
- \(D\)：Embedding维度，Embedding dimension，即一个token进行Embedding后得到的向量维度
- \(V\)：词表大小，Vocabulary size，即我们在训练模型时，使用的词表中所有token的数量。
在输出的张量中，我们固定到一个Batch中，则每个Batch索引着一个大小为\((S, V)\)的张量矩阵，对于第\(i\)个位置上的输出向量，代表着模型吃进去前面\(i\)个token后，对第\(i+1\)个token的预测。这个预测对词表中的每个词进行打分，所有词表中所有词的分数构成了这个长度为\(V\)的向量。而我们此时的目标就是对这个长度为\(V\)的向量进行解码，得到我们期望的输出token。
当然，我们给模型一段文本，并不是仅仅让他预测一个token，而是希望模型能够预测出一段完整的文本，就如同现在的GPT，DeepSeek等模型一样，对于我们的输入，他会给我们一段完整的输出文本。这就要求我们在解码时，将解码后的token再次输入到模型中，继续预测下一个token，直到我们得到一个完整的句子。这个过程可以见下面的图示：
当然，我们也不是期望模型无限制的输出下去，而是希望模型在输出到一个特定的token时，停止输出，这个token就是我们在训练模型时，使用的<eos>，即End of Sentence的token。当检测到生成了这个token时，我们就停止解码。
解码过程分为如下几个步骤：
- 在输出中取出末端的向量，之后的操作将围绕其进行
- 对上述的向量进行温度缩放以及softmax
- Top-p采样，得到目标的token序列
- 归一化与采样以及回拼
在代码中即为
prompt
↓
模型预测下一个 token 的分布
↓
温度调整
↓
Top-p 筛选
↓
随机采样
↓
追加到已有序列
↓
继续生成，直到结束
- 推理输入并对输出取出末端的向量
对于我们的输出的\((B, S, V)\)张量，在每个批次上，我们实际上只需要取出最后一个位置的向量，因为为了预测下一个位置的token，我们实际上只需要对这个向量操作即可。
import torch
def decode(model, prompt, eot_token_id=None, max_new_tokens=256, temperature=0.5, top_p=0.9, context_length=256)-> list[int]:
ids = list(prompt)
# prompt: 输入的token序列，提示词
x = torch.tensor(ids).unsqueeze(0)
# unsqueeze(0)的作用是将(S,)的一维张量变为(1, S)的二维张量，准备作为模型的输入(B, S)
# B = 1
model.eval() # 切换到评估模式，关闭dropout等训练时特有的操作
torch.set_grad_enabled(False) # 关闭梯度计算，节省内存和计算资源
for i in range(1, max_new_tokens + 1):
x_input = x[:, -context_length:]
# x[:, -context_length:],第一个":", 取出所有的batch，第二个":", 倒着取出最多context_length个token，保证输入的长度不超过模型的最大上下文长度
x_input = x[:, -context_length:]
这一句会在所有的输入x中，取出所有的batch（这里只有一个batch），对于每个batch上的输入，由于用的是-context_length:，所以会从最后一个token开始，向前取出最多context_length个token，保证输入的长度不超过模型的最大上下文长度。如
context_length = 256
完整序列 x： 300 个 token
模型输入 x_input：最后 256 个 token
接着
import torch
def decode(model, prompt, eot_token_id=None, max_new_tokens=256, temperature=0.5, top_p=0.9, context_length=256)-> list[int]:
ids = list(prompt)
x = torch.tensor(ids).unsqueeze(0)
model.eval()
torch.set_grad_enabled(False)
for i in range(1, max_new_tokens + 1):
x_input = x[:, -context_length:] # shape: (B, S)
# 接上
pred_y = model(x_input)[:, -1, :]
# y = model(x_input) -> (B, S, V)，这是推理的结果
# 而 pred_y = y[:, -1, :]，取出最后一个位置的向量，shape: (B, V)，即每个batch上最后一个位置的预测向量，这里是(1, V)
- 温度调整
对每个位置上的token进行分数归一化时，我们采用一种叫做温度调整的策略：
其中：
- \(z_j\)：第\(j\)个token的logit分数
- \(\tau\)：温度参数，temperature，控制分布的平滑程度
- \(p_j\)：第\(j\)个token的归一化概率
不难看出，这个操作实际上就是在普通的\(softmax(z_i) = \frac{\exp(z_i)}{\sum_k \exp(z_k)}\)的基础上，让每个参与运算的分数\(z_i\)都除以温度参数\(\tau\)。我们可以想象：
-
当\(\tau=1\)时，\(\frac{z_i}{\tau} = z_i\)，此时就是普通的softmax归一化
-
当\(\tau>1\)时，\(\frac{z_i}{\tau} < z_i\)，此时每个分数都被缩小了，不妨类比一下正则化的作用，此时所有的分数都以同样比例缩小，即较高的分数会缩小的多，而较低的分数会缩小的少，这样在归一化后，概率分布会更加平滑。
-
当\(\tau<1\)时，\(\frac{z_i}{\tau} > z_i\)，此时每个分数都被放大了，此时所有的分数都以同样比例放大，即较高的分数会放大的多，而较低的分数会放大的少，这样在归一化后，概率分布会更加尖锐。
下面这张图演示了不同温度\(\tau\)下的softmax归一化的概率分布情况：
为什么要调整这个呢？
一个前提是，我们最后选取目标token并不像一些图像分类任务一样，类别输出选取的max概率的类别，而是希望这个token在选择的时候根据自身的概率去做随机抽样，保留一定的多样化生成。
而在\(\tau\)比较小的时候，概率分布会比较尖锐，意味着模型会更倾向于选择概率最高的token，此时生成的文本会比较确定。而当\(\tau\)比较大的时候，概率分布会比较平滑，此时选择各个token的概率会更加接近，模型在选择token时会有更多的随机性，生成的文本会更加多样化。
接着
import torch
def decode(model, prompt, eot_token_id=None, max_new_tokens=256, temperature=0.5, top_p=0.9, context_length=256)-> list[int]:
ids = list(prompt)
x = torch.tensor(ids).unsqueeze(0)
model.eval()
torch.set_grad_enabled(False)
for i in range(1, max_new_tokens + 1):
x_input = x[:, -context_length:] # shape: (B, S)
# 接上
pred_y = model(x_input)[:, -1, :]
# 接上
probs = torch.softmax(
pred_y / temperature,
dim=-1,
).squeeze(0)
# 对pred_y的每个最后一个维度进行温度调整和softmax归一化->(1, V)
# squeeze(0)的作用是将(1, V)的二维张量变为(V,)的一维张量，方便后续操作
- Top-p 筛选
所谓的Top-p筛选，也叫做核采样（Nucleus Sampling），是指在生成文本时，我们并不是直接从整个词表中进行采样，而是先根据概率分布筛选出一个子集，然后再从这个子集中进行采样。当然，类比到图像分类领域，Top-k筛选是比较常见的。
这个子集的筛选方式为： 保留从最高概率的token开始累计概率达到阈值的候选token
为了实现这样的筛选需要去对候选的向量：
- 从高到低将概率排序
- 保留排序前后token的对应索引，确保在后续的采样中，能够将排序后的token映射回原始的token
- 对排序后的概率求前缀和
- 有top-p阈值，找到首次达到或者超过阈值的索引位置，保留该索引位置之前的所有token作为候选token
给定一个例子：
probs = [0.15, 0.50, 0.10, 0.25] # 每个元素为每个token的概率分布
sorted_probs = [0.50, 0.25, 0.15, 0.10] # 对上述分布从高到低排序，保证top-p从最高的开始累计
sorted_indices = [1, 3, 0, 2] # 排序后的元素在排序前的索引位置，即probs[sorted_indices] = sorted_probs
cumsum = [0.50, 0.75, 0.90, 1.00] # 对排序后的分布进行累计求和，得到累计概率分布
# 0.50 = 0.50
# 0.75 = 0.50 + 0.25
# 0.90 = 0.50 + 0.25 + 0.15
# 1.00 = 0.50 + 0.25 + 0.15 + 0.10
# 若阈值为 top_p = 0.8，则找到首次到达或者超过阈值的索引位置
cutoff_idx = 2 # 因为cumsum_probs[2] = 0.90 >= 0.8
# 此时保留前两个token，也就是将后面的所有token，在原来的概率分布中置为0，得到一个新的概率分布
# 因此我们使用一个掩码取出后面token的索引
sorted_mask = (
torch.arange(cumsum.shape[-1], device=probs.device)
# 生成一个(0, 1, ..., cumsum.shape[-1]-1)的张量，这里是(0, 1, 2, 3)
# cumsum.shape[-1] 即为所有token的数量
<= cutoff_idx # 注意这里是小于等于号，不是箭头
)
# torch.arange(cumsum.shape[-1], device=probs.device)<= cutoff_idx的结果是一个bool量的Tensor，表示每个位置的索引是否小于等于cutoff_idx，因此0, 1, 2, ..., cutoff_idx位置的值为True，其他位置为False
# 在我们的例子中[0, 1, 2, 3] <= [2, 2, 2, 2]，得到的结果为[True, True, True, False]
# 别忘了，我们的目标是让False位置上的值为0，因此我们需要将这个掩码取反，得到一个新的掩码
# 然后用sorted_indices[~sorted_mask]取出原始概率分布中对应的索引位置，将这些位置上的值置为0
probs[sorted_indices[~sorted_mask]] = 0
即：
import torch
def decode(model, prompt, eot_token_id=None, max_new_tokens=256, temperature=0.5, top_p=0.9, context_length=256)-> list[int]:
ids = list(prompt)
x = torch.tensor(ids).unsqueeze(0)
model.eval()
torch.set_grad_enabled(False)
for i in range(1, max_new_tokens + 1):
x_input = x[:, -context_length:] # shape: (B, S)
# 接上
pred_y = model(x_input)[:, -1, :]
# 接上
probs = torch.softmax(
pred_y / temperature,
dim=-1,
).squeeze(0)
# 接上
sorted_probs, sorted_indices = torch.sort(probs, descending=True)
# 对probs进行从高到低排序，得到排序后的概率分布sorted_probs，以及排序后的索引sorted_indices，保证top-p从最高的开始累计
cumsum = sorted_probs.cumsum(dim=-1)
# 对排序后的分布进行累计求和，得到累计概率分布cumsum（前缀和）
cutoff_idx = torch.searchsorted(cumsum, top_p) # 找到首次到达或者超过阈值top_p的索引位置
sorted_mask = torch.arange(cumsum.shape[-1], device=probs.device) <= cutoff_idx
probs[sorted_indices[~sorted_mask]] = 0
probs /= probs.sum() # 由于丢掉了一些token的概率，因此需要重新归一化，使得所有token的概率和为1
- 随机采样与回拼
结束上述工作后，我们就得到了最终的概率分布probs，随后的工作就是对其进行随机采样，得到目前的预测token，并将其追加到已有的token序列中，继续进行下一轮的预测。
import torch
def decode(model, prompt, eot_token_id=None, max_new_tokens=256, temperature=0.5, top_p=0.9, context_length=256)-> list[int]:
ids = list(prompt)
x = torch.tensor(ids).unsqueeze(0)
model.eval()
torch.set_grad_enabled(False)
for i in range(1, max_new_tokens + 1):
x_input = x[:, -context_length:] # shape: (B, S)
# 接上
pred_y = model(x_input)[:, -1, :]
# 接上
probs = torch.softmax(
pred_y / temperature,
dim=-1,
).squeeze(0)
# 接上
sorted_probs, sorted_indices = torch.sort(probs, descending=True)
cumsum = sorted_probs.cumsum(dim=-1)
cutoff_idx = torch.searchsorted(cumsum, top_p)
sorted_mask = torch.arange(cumsum.shape[-1], device=probs.device) <= cutoff_idx
probs[sorted_indices[~sorted_mask]] = 0
probs /= probs.sum()
# 接上
next_token_id = torch.multinomial(probs, num_samples=1).item()
# 通过torch.multinomial对归一化后的概率分布probs进行随机采样，得到一个token的索引next_token_id
# num_samples=1表示采样一个token，.item()将得到的张量(1,)转换为Python的int类型
ids.append(next_token_id)
# 将采样得到的token索引next_token_id追加到已有的token序列ids中，准备进行下一轮的预测
x = torch.cat([x, torch.tensor([[next_token_id]], device=x.device)], dim=1)
# x -> (1, S)
# torch.tensor([[next_token_id]], device=x.device) -> (1, 1)
# torch.cat([x, torch.tensor([[next_token_id]], device=x.device)], dim=1) -> (1, S+1)
# dim = 1 在 S 维度上进行拼接，即在序列长度的维度上进行拼接，得到新的输入x，准备进行下一轮的预测
if eot_token_id is not None and next_token_id == eot_token_id:
break
# 当结束符号存在，同时下一个token的索引等于结束符号的索引时，停止解码，跳出循环
return ids
浙公网安备 33010602011771号