Transformer推理的基本流程

说实话，很多人刚开始接触Transformer推理时，总觉得跟训练差不多——喂数据，跑前向，出结果。但实际工程里，推理和训练完全是两码事。嗯，我来给你捋一遍我平时搭管道时的标准流程。

我个人习惯把Transformer推理拆成五个阶段。每个阶段都有坑，也有技巧。我们先看一张总图。

阶段一：输入预处理

为什么要先预处理？因为模型不认识原始文本，只认识数字。

我记得第一次做文本生成时，直接把英文句子喂给BertTokenizer，结果漏掉了special tokens，跑了半天生成的句子全是乱码。后来养成习惯：

• tokenize：用对应模型的分词器将文本转成token ID序列。
• 添加特殊标记：比如 [CLS]、[SEP]、<s>、</s>，根据模型要求来。
• padding / truncation：固定输入长度，或者动态batch时按最大长度padding。
• attention mask：告诉模型哪些位置是真实token，哪些是padding。

代码示例：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer(
    ["Hello, world!", "Transformers are awesome."],
    padding="max_length",
    truncation=True,
    max_length=128,
    return_tensors="pt"
)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

阶段二：词嵌入 + 位置编码

这一步模型把token ID映射成稠密向量，再叠加位置信息。说白了，就是让每个词从“孤立的符号”变成“有语义的向量”，同时还知道它在一段话里的位置。

我曾在项目中遇到过一个问题：模型用了相对位置编码（比如T5），但我在推理时忘了传入position_ids，导致生成的内容上下文位置完全错乱。从那以后我养成了习惯——无论使用哪种位置编码，都必须显式地检查输入参数。

常见位置编码方式对比：

类型	代表模型	特点
绝对位置编码	BERT, GPT-2	每个位置有固定embedding，长度有限制
相对位置编码	T5, GPT-Neo	只编码token间的相对距离，可处理更长序列
旋转位置编码（RoPE）	LLaMA, Mistral	通过旋转矩阵实现，支持长度外推
ALiBi	BLOOM, MPT	线性偏置，不参与学习，简单高效

在推理时，嵌入层和位置编码层通常合并成一个nn.Embedding或一个layer，我们可以直接调用模型的forward或者transformer.wte等子模块。

阶段三：自回归解码

这是推理的核心。Transformer生成文本是逐token的——你想想看，它每生成一个词，就把这个词再拼回输入，继续预测下一个词。

为什么不一次生成所有词？因为自回归模型本质上是条件概率链：P(y1,y2,...,yn|x) = P(y1|x) * P(y2|x,y1) * ...。每一步都依赖之前生成的结果。

具体循环：

1. 把当前已生成的token序列喂给模型。
2. 模型输出logits（形状[batch, seq_len, vocab_size]）。
3. 只取最后一个位置的logits（logits[:,-1,:]）。
4. 通过softmax得到概率分布。
5. 选择下一个token（greedy / sampling / beam search）。
6. 重复直到遇到结束符或达到最大长度。

代码框架：

import torch

def generate_step(model, input_ids, attention_mask):
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)
    logits = outputs.logits[:, -1, :]  # 只取最后一个token的logits
    probs = torch.softmax(logits, dim=-1)
    next_token = torch.argmax(probs, dim=-1, keepdim=True)  # greedy
    return next_token

input_ids = torch.tensor([[tokenizer.bos_token_id]])
attention_mask = torch.ones_like(input_ids)

for step in range(max_new_tokens):
    next_token = generate_step(model, input_ids, attention_mask)
    input_ids = torch.cat([input_ids, next_token], dim=-1)
    attention_mask = torch.cat([attention_mask, torch.ones_like(next_token)], dim=-1)
    if next_token.item() == tokenizer.eos_token_id:
        break

阶段四：输出投影

模型的最后一层通常是一个线性投影（lm_head），把隐藏状态映射到词表大小的logits上。这一步计算量不大，但很重要——它决定了最终词表的选择。

为什么叫“投影”？因为隐藏状态的维度（比如768维）可能远小于词表大小（比如32k）。通过一个权重矩阵，我们把每个token的表示“投影”到词表空间。

代码实现通常很简单：

hidden_states = model.transformer(input_ids)[0]  # [batch, seq, hidden_dim]
lm_logits = model.lm_head(hidden_states)           # [batch, seq, vocab_size]

有一点我一直提醒自己：有些模型（比如GPT-2）的lm_head和token embedding是共享权重的，所以当你修改embedding时，别忘了同步修改lm_head。

阶段五：后处理

后处理就是把logits转成可读的文本。这里面有好多坑。

• Top-k / Top-p采样：只从概率最高的k个里选，或者累积概率到p。用来控制生成多样性。
• 温度参数：logits除以温度（temperature）再softmax。温度越高越随机，越低越确定。
• 重复惩罚：对已出现过的token做惩罚，防止无限循环。
• EOS处理：生成到结束符提前停，但有时模型不输出EOS，需要设最大长度。

我曾经踩过一个坑：生成序列里出现了<unk>标记，一查发现tokenizer里unk token的ID没有被mask掉。从那以后，我总会显式设置bad_words_ids来过滤掉特殊标记。

后处理示例：

from transformers import AutoTokenizer, LogitsProcessorList, TopKLogitsWarper, TemperatureLogitsWarper

logits_processor = LogitsProcessorList([
    TemperatureLogitsWarper(temperature=0.7),
    TopKLogitsWarper(top_k=50),
])
logits = outputs.logits[:, -1, :]
logits = logits_processor(input_ids, logits)

probs = torch.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)

小结一张表：五大阶段回顾

阶段	输入	输出	关键操作
输入预处理	原始文本	input_ids, attention_mask	tokenize, padding, truncation
词嵌入+位置编码	token ID	嵌入向量序列	Embedding + 位置编码
自回归解码	嵌入向量序列	隐藏状态序列	多层Transformer block逐层计算（含KV Cache）
输出投影	隐藏状态	logits	线性投影+softmax
后处理	logits	最终token/文本	采样、惩罚、过滤、解码

好了，以上就是Transformer推理的基本流程。你想想看，其实每一步都不复杂，但合起来就是一个精密的链条。每个环节都可能成为性能瓶颈，我们在后面的章节会逐一拆解并优化。