Transformer 自回归生成机制笔记-尧图网络科技

声明：这篇笔记仅为记录学习 Transformer 生成过程时的理解和实现，纯为技术分享，不含任何商业推广或培训导向。

1. 自回归解码的基本流程

无论是机器翻译中使用的编码器-解码器 Transformer，还是 GPT 这类纯解码器结构，文本生成的核心都是自回归：将已生成的序列当作新的输入，模型据此预测下一个词，拼接到尾部，如此循环。

步骤可以概括为：

提供初始输入（比如<sos>起始符，机器翻译还需要源语言句子）。
模型基于当前序列计算所有可能下一个词的概率分布。
从分布中选择一个词（贪心、采样等），附加到序列末尾。
重复第 2 步，直到出现<eos>或达到预设的最大长度。

每一步都能看到之前所有生成的内容，因此上下文是完整的。

2. 注意力机制及其实现

2.1 缩放点积注意力

注意力层的计算可以表述为：

Attention(Q,K,V)=softmax(QKTdk)V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) VAttention(Q,K,V)=softmax(dkQKT)V

除以dk\sqrt{d_k}dk是为了防止点积数值过大，让 softmax 的输出保持合理的梯度。

对应的 PyTorch 实现：

importtorchimporttorch.nn.functionalasFimportmathdefscaled_dot_product_attention(Q,K,V,mask=None):d_k=Q.size(-1)scores=torch.matmul(Q,K.transpose(-2,-1))/math.sqrt(d_k)ifmaskisnotNone:scores=scores.masked_fill(mask==0,float('-inf'))attn_weights=F.softmax(scores,dim=-1)returntorch.matmul(attn_weights,V),attn_weights

2.2 多头注意力

多头注意力将 Q、K、V 线性映射到多个子空间，各自做注意力运算，再将结果拼接、投影回原维度。

classMultiHeadAttention(nn.Module):def__init__(self,d_model,num_heads):super().__init__()self.num_heads=num_heads self.d_k=d_model//num_heads self.W_q=nn.Linear(d_model,d_model)self.W_k=nn.Linear(d_model,d_model)self.W_v=nn.Linear(d_model,d_model)self.W_o=nn.Linear(d_model,d_model)defsplit_heads(self,x):# x: (batch, seq_len, d_model) -> (batch, heads, seq_len, d_k)b,s,_=x.shapereturnx.view(b,s,self.num_heads,self.d_k).transpose(1,2)defcombine_heads(self,x):b,_,s,_=x.shapereturnx.transpose(1,2).contiguous().view(b,s,-1)defforward(self,Q,K,V,mask=None):Q=self.split_heads(self.W_q(Q))K=self.split_heads(self.W_k(K))V=self.split_heads(self.W_v(V))attn_out,_=scaled_dot_product_attention(Q,K,V,mask)returnself.W_o(self.combine_heads(attn_out))

3. 位置编码与前馈网络

Transformer 本身不包含循环或卷积，需要借助位置编码注入序列顺序信息。这里使用正弦-余弦固定编码：

classPositionalEncoding(nn.Module):def__init__(self,d_model,max_len=5000):super().__init__()pe=torch.zeros(max_len,d_model)position=torch.arange(0,max_len).unsqueeze(1).float()div_term=torch.exp(torch.arange(0,d_model,2).float()*-(math.log(10000.0)/d_model))pe[:,0::2]=torch.sin(position*div_term)pe[:,1::2]=torch.cos(position*div_term)pe=pe.unsqueeze(0)# (1, max_len, d_model)self.register_buffer('pe',pe)defforward(self,x):returnx+self.pe[:,:x.size(1),:]

每个注意力子层后面都有一个前馈网络，本质上是一个两层的全连接网络，中间用 ReLU 激活：

classFeedForward(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.linear1=nn.Linear(d_model,d_ff)self.linear2=nn.Linear(d_ff,d_model)defforward(self,x):returnself.linear2(F.relu(self.linear1(x)))

4. 编码器层与解码器层

4.1 编码器层

编码器层包含多头自注意力 + 前馈网络，每个子层后面有残差连接和层归一化。

classEncoderLayer(nn.Module):def__init__(self,d_model,num_heads,d_ff,dropout=0.1):super().__init__()self.self_attn=MultiHeadAttention(d_model,num_heads)self.ff=FeedForward(d_model,d_ff)self.norm1=nn.LayerNorm(d_model)self.norm2=nn.LayerNorm(d_model)self.dropout=nn.Dropout(dropout)defforward(self,x,mask=None):a=self.self_attn(x,x,x,mask)x=self.norm1(x+self.dropout(a))f=self.ff(x)x=self.norm2(x+self.dropout(f))returnx

4.2 解码器层

解码器层在自注意力之后还插入了一个交叉注意力，它用解码器的表示去查询编码器的输出。同时，自注意力部分必须使用因果掩码，防止当前词“偷看”后面内容。

classDecoderLayer(nn.Module):def__init__(self,d_model,num_heads,d_ff,dropout=0.1):super().__init__()self.self_attn=MultiHeadAttention(d_model,num_heads)self.cross_attn=MultiHeadAttention(d_model,num_heads)self.ff=FeedForward(d_model,d_ff)self.norm1=nn.LayerNorm(d_model)self.norm2=nn.LayerNorm(d_model)self.norm3=nn.LayerNorm(d_model)self.dropout=nn.Dropout(dropout)defforward(self,x,enc_output,src_mask=None,tgt_mask=None):# 带掩码的自注意力s=self.self_attn(x,x,x,tgt_mask)x=self.norm1(x+self.dropout(s))# 交叉注意力 (query来自解码器, key/value来自编码器)c=self.cross_attn(x,enc_output,enc_output,src_mask)x=self.norm2(x+self.dropout(c))# 前馈f=self.ff(x)x=self.norm3(x+self.dropout(f))returnx

5. 因果掩码与推理循环

因果掩码是一个下三角矩阵，保证位置 i 只能关注 0 到 i 的信息。实现方式很简单：

defcreate_causal_mask(sz):mask=torch.tril(torch.ones(sz,sz))# 下三角为1，上三角为0returnmask.unsqueeze(0).unsqueeze(0)# (1, 1, sz, sz)

在推理阶段，源序列经过编码器生成记忆memory（仅计算一次），然后解码器逐词生成：

defgreedy_decode(model,src,src_mask,max_len,start_sym,end_sym):model.eval()memory=model.encode(src,src_mask)ys=torch.ones(1,1).fill_(start_sym).long()# 初始序列，仅含<sos>for_inrange(max_len-1):tgt_mask=create_causal_mask(ys.size(1)).to(src.device)out=model.decode(memory,src_mask,ys,tgt_mask)# (1, seq_len, vocab_size)prob=out[:,-1,:]# 取最后一步的 logits_,next_word=torch.max(prob,dim=1)next_word=next_word.item()ys=torch.cat([ys,torch.ones(1,1).fill_(next_word).long()],dim=1)ifnext_word==end_sym:breakreturnys

实践中这里会加入KV 缓存以提升效率，但核心逻辑不变。

6. 完整模型与生成方法

将以上模块组合起来，得到一个标准的 Transformer 模型。下面展示encode、decode和generate三个核心方法（省略完整的__init__以节省篇幅，结构就是堆叠若干编码器层和解码器层，并加上词嵌入和位置编码）。

classTransformer(nn.Module):# __init__ 中定义:# encoder_embed, decoder_embed, pos_encoder, pos_decoder,# encoder_layers (ModuleList), decoder_layers (ModuleList),# output_proj (nn.Linear)# 细节见前文各个组件。defencode(self,src,src_mask=None):x=self.encoder_embed(src)*math.sqrt(self.d_model)x=self.pos_encoder(x)forlayerinself.encoder_layers:x=layer(x,src_mask)returnxdefdecode(self,memory,src_mask,tgt,tgt_mask=None):x=self.decoder_embed(tgt)*math.sqrt(self.d_model)x=self.pos_decoder(x)forlayerinself.decoder_layers:x=layer(x,memory,src_mask,tgt_mask)returnself.output_proj(x)@torch.no_grad()defgenerate(self,src,src_mask,max_len,start_sym,end_sym,temperature=1.0,top_k=0):self.eval()memory=self.encode(src,src_mask)ys=torch.ones(src.size(0),1).fill_(start_sym).long().to(src.device)for_inrange(max_len-1):tgt_mask=create_causal_mask(ys.size(1)).to(src.device)logits=self.decode(memory,src_mask,ys,tgt_mask)logits=logits[:,-1,:]/temperature# 温度缩放iftop_k>0:v,_=torch.topk(logits,top_k)logits[logits<v[:,-1:]]=float('-inf')probs=F.softmax(logits,dim=-1)next_word=torch.multinomial(probs,num_samples=1)# 随机采样# 贪心解码可替换为: next_word = torch.argmax(probs, dim=-1, keepdim=True)ys=torch.cat([ys,next_word],dim=1)if(next_word==end_sym).all():breakreturnys