Transformer核心模块PyTorch实战从LayerNorm到FFN的工程实现细节1. 环境准备与基础概念在开始构建Transformer核心模块之前我们需要明确几个关键概念。Transformer架构之所以能够在各类序列任务中表现出色很大程度上归功于其精心设计的子模块协同工作。不同于简单地堆叠神经网络层Transformer通过LayerNorm、多头注意力机制和前馈网络(FFN)的有机组合实现了对序列数据的高效建模。首先配置基础环境import torch import torch.nn as nn import torch.nn.functional as F import math # 确保可复现性 torch.manual_seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu)Transformer中的LayerNorm与传统的BatchNorm有本质区别特性LayerNormBatchNorm统计量计算维度特征维度批量维度适用场景变长序列任务固定长度输入参数数量2×特征维度2×通道数推理行为与训练一致依赖运行统计量提示在序列任务中LayerNorm对每个时间步独立计算统计量这使得它天然适合处理语音、文本等变长数据。2. LayerNorm的PyTorch实现与调试技巧让我们从最基础的LayerNorm实现开始。虽然PyTorch提供了现成的nn.LayerNorm但理解其底层机制对调试复杂模型至关重要。class CustomLayerNorm(nn.Module): def __init__(self, normalized_shape, eps1e-5): super().__init__() if isinstance(normalized_shape, int): normalized_shape (normalized_shape,) self.normalized_shape normalized_shape self.eps eps self.weight nn.Parameter(torch.ones(normalized_shape)) self.bias nn.Parameter(torch.zeros(normalized_shape)) def forward(self, x): # 计算均值和方差 mean x.mean(dim[-d for d in range(1, len(self.normalized_shape)1)], keepdimTrue) var x.var(dim[-d for d in range(1, len(self.normalized_shape)1)], keepdimTrue, unbiasedFalse) # 归一化 x_normalized (x - mean) / torch.sqrt(var self.eps) return self.weight * x_normalized self.bias实际使用时会遇到几个典型问题维度不匹配当输入张量的最后几维与normalized_shape不一致时会引发运行时错误。例如# 错误示例 norm CustomLayerNorm((64,)) x torch.randn(32, 128) # 第二维128 ! 64 out norm(x) # 报错梯度消失当eps值设置过小如1e-10时在方差接近零的情况下可能导致数值不稳定。初始化问题如果权重初始化为全零会使得输出始终为零破坏信息流动。调试建议在实现自定义LayerNorm时建议添加以下验证代码def test_layernorm(): # 测试2D输入如序列任务 x torch.randn(10, 64) # (batch_size, feature_dim) custom_norm CustomLayerNorm(64) official_norm nn.LayerNorm(64) # 前向一致性检查 assert torch.allclose(custom_norm(x), official_norm(x), atol1e-6) # 反向传播检查 x.requires_grad_(True) out_custom custom_norm(x) out_official official_norm(x) out_custom.sum().backward() grad_custom x.grad.clone() x.grad.zero_() out_official.sum().backward() grad_official x.grad.clone() assert torch.allclose(grad_custom, grad_official, atol1e-6) print(所有测试通过)3. 多头注意力机制(MHSA)的工程实现多头注意力是Transformer的核心创新点其实现需要考虑效率、数值稳定性和可扩展性。我们先看基础的自注意力实现class MultiHeadAttention(nn.Module): def __init__(self, d_model512, n_heads8, dropout0.1): super().__init__() assert d_model % n_heads 0 self.d_k d_model // n_heads self.n_heads n_heads self.d_model d_model # 线性变换层 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) self.dropout nn.Dropout(dropout) self.scale 1.0 / math.sqrt(self.d_k) def forward(self, q, k, v, maskNone): batch_size q.size(0) # 线性变换并分头 q self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) k self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) v self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 计算注意力分数 scores torch.matmul(q, k.transpose(-2, -1)) * self.scale if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # 注意力权重 attn_weights F.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 上下文向量 context torch.matmul(attn_weights, v) context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.w_o(context)实际工程中需要注意的几个关键点内存优化当序列长度较大时如超过1024注意力矩阵会消耗大量内存。解决方案包括使用内存高效的注意力实现采用分块计算策略混合精度训练数值稳定性缩放因子1/sqrt(d_k)对防止梯度消失至关重要对softmax输入做减最大值处理x - x.max()可增强数值稳定性并行计算利用torch.baddbmm替代循环计算合理设置batch_size和n_heads以充分利用GPU并行能力性能优化技巧在推理阶段可以通过以下方式提升效率# 缓存K和V的计算结果适用于自回归解码 class CachedAttention(MultiHeadAttention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.cache_k None self.cache_v None def forward(self, q, k, v, maskNone, use_cacheFalse): if use_cache: if self.cache_k is None: # 首次调用 self.cache_k k self.cache_v v else: # 后续调用拼接新token self.cache_k torch.cat([self.cache_k, k], dim1) self.cache_v torch.cat([self.cache_v, v], dim1) k, v self.cache_k, self.cache_v return super().forward(q, k, v, mask)4. 前馈网络(FFN)的实现与变体Transformer中的前馈网络虽然结构简单但有几个关键设计值得关注class PositionwiseFFN(nn.Module): def __init__(self, d_model, d_ff2048, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.activation nn.GELU() # 比原始ReLU表现更好 def forward(self, x): return self.linear2(self.dropout(self.activation(self.linear1(x))))FFN的几种常见变体及其特点GELU激活高斯误差线性单元(Gaussian Error Linear Unit)在实践中表现优于ReLUdef gelu(x): return 0.5 * x * (1 torch.tanh(math.sqrt(2 / math.pi) * (x 0.044715 * x**3)))门控机制如GLUGated Linear Unit变体class GLUFFN(nn.Module): def __init__(self, d_model, d_ff): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear_gate nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) def forward(self, x): return self.linear2(F.silu(self.linear1(x)) * self.linear_gate(x))参数共享在多任务学习中可以共享部分FFN层参数FFN实现中的常见问题及解决方案问题现象可能原因解决方案输出值过大/过小初始化不当使用Kaiming初始化训练后期性能停滞梯度消失添加残差连接推理速度慢中间维度(d_ff)过大使用Bottleneck结构多GPU训练不稳定各卡统计量不一致使用同步LayerNorm5. 完整Transformer层的集成与调试将上述模块组合成完整的Transformer层时需要注意几个关键细节class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff2048, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn PositionwiseFFN(d_model, d_ff, dropout) self.norm1 CustomLayerNorm(d_model) self.norm2 CustomLayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, maskNone): # 自注意力子层 attn_output self.self_attn(x, x, x, mask) x x self.dropout1(attn_output) x self.norm1(x) # 前馈子层 ffn_output self.ffn(x) x x self.dropout2(ffn_output) x self.norm2(x) return x集成测试时需要关注的指标梯度检查def grad_check(): layer TransformerEncoderLayer(d_model512, n_heads8) x torch.randn(10, 32, 512, requires_gradTrue) # (batch, seq, dim) # 前向传播 out layer(x) # 反向传播检查 out.sum().backward() for name, param in layer.named_parameters(): if param.grad is None: print(f警告{name} 无梯度) elif torch.all(param.grad 0): print(f警告{name} 梯度全零)数值范围监控各层输出值应在合理范围内如[-10, 10]注意力权重应介于[0,1]之间计算图验证# 使用torchviz可视化计算图 from torchviz import make_dot make_dot(out, paramsdict(layer.named_parameters()))实际项目中遇到的典型问题案例问题描述在训练初期损失值波动剧烈然后突然变为NaN诊断过程检查各层输出范围发现LayerNorm输出出现极大值追踪发现注意力分数在softmax前达到1e30量级确认d_k计算错误缩放因子失效解决方案修正缩放因子计算添加梯度裁剪初始化权重缩小10倍6. 高级技巧与性能优化当Transformer模型规模扩大时需要考虑以下优化策略混合精度训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, targets in dataloader: optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()激活检查点减少内存消耗from torch.utils.checkpoint import checkpoint class MemoryEfficientEncoderLayer(TransformerEncoderLayer): def forward(self, x, maskNone): return checkpoint(super().forward, x, mask)自定义内核融合使用torch.jit.script或Triton编写融合操作torch.jit.script def fused_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float): scores torch.matmul(q, k.transpose(-2, -1)) * scale attn torch.softmax(scores, dim-1) return torch.matmul(attn, v)不同规模模型的实现差异参数量级实现重点典型优化手段100M开发效率原生PyTorch实现100M-1B计算效率混合精度、梯度检查点1B分布式训练模型并行、流水线并行7. 测试与验证策略为确保实现的正确性建议建立全面的测试套件一致性测试与参考实现对比def test_vs_huggingface(): # 对比HuggingFace实现 from transformers import BertModel our_model OurTransformer() hf_model BertModel.from_pretrained(bert-base-uncased) # 权重转换后比较输出 assert torch.allclose(our_model(x), hf_model(x)[0], atol1e-4)梯度检验from torch.autograd import gradcheck input torch.randn(10, 64, dtypetorch.double, requires_gradTrue) test gradcheck(CustomLayerNorm(64), input, eps1e-6, atol1e-4) print(梯度检验通过:, test)数值稳定性测试def test_numerical_stability(): model TransformerEncoderLayer(d_model512, n_heads8) with torch.no_grad(): for name, param in model.named_parameters(): param.data.normal_(mean0, std100) # 极端初始化 x torch.randn(2, 256, 512) * 100 # 大输入值 out model(x) assert not torch.isnan(out).any() assert not torch.isinf(out).any()性能基准测试from torch.utils.benchmark import Timer def benchmark(): model TransformerEncoderLayer(d_model512, n_heads8).cuda() x torch.randn(32, 128, 512).cuda() # 预热 for _ in range(10): _ model(x) timer Timer( stmtmodel(x), globals{model: model, x: x} ) print(timer.timeit(100))8. 实际应用中的经验分享在工业级应用中我们发现以下几个实践对提升Transformer模块效果至关重要预热学习率策略配合Adam优化器使用线性预热def get_lr(step, d_model, warmup_steps4000): return d_model**-0.5 * min(step**-0.5, step * warmup_steps**-1.5)残差连接缩放在某些深层模型中效果更好class ScaledResidual(nn.Module): def __init__(self, scale0.1): super().__init__() self.scale scale def forward(self, x, residual): return x self.scale * residual注意力模式混合结合不同注意力模式提升效果class MixedAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.local_attn LocalAttentionWindow(d_model, n_heads, window_size64) self.global_attn MultiHeadAttention(d_model, n_heads) self.gate nn.Linear(d_model, 2) def forward(self, q, k, v, maskNone): local_out self.local_attn(q, k, v, mask) global_out self.global_attn(q, k, v, mask) gate F.softmax(self.gate(q), dim-1) return gate[..., 0:1] * local_out gate[..., 1:2] * global_out动态FFN宽度根据输入动态调整中间维度class DynamicFFN(nn.Module): def __init__(self, d_model, max_ff2048): super().__init__() self.controller nn.Linear(d_model, 1) self.linear1 nn.Linear(d_model, max_ff) self.linear2 nn.Linear(max_ff, d_model) def forward(self, x): ratio torch.sigmoid(self.controller(x.mean(dim1))) # [0,1] current_ff int(ratio * self.linear1.out_features) h self.linear1(x) h h[..., :current_ff] # 动态切片 return self.linear2(h)
Transformer实战：从LayerNorm到FFN，手把手教你用PyTorch实现关键模块
Transformer核心模块PyTorch实战从LayerNorm到FFN的工程实现细节1. 环境准备与基础概念在开始构建Transformer核心模块之前我们需要明确几个关键概念。Transformer架构之所以能够在各类序列任务中表现出色很大程度上归功于其精心设计的子模块协同工作。不同于简单地堆叠神经网络层Transformer通过LayerNorm、多头注意力机制和前馈网络(FFN)的有机组合实现了对序列数据的高效建模。首先配置基础环境import torch import torch.nn as nn import torch.nn.functional as F import math # 确保可复现性 torch.manual_seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu)Transformer中的LayerNorm与传统的BatchNorm有本质区别特性LayerNormBatchNorm统计量计算维度特征维度批量维度适用场景变长序列任务固定长度输入参数数量2×特征维度2×通道数推理行为与训练一致依赖运行统计量提示在序列任务中LayerNorm对每个时间步独立计算统计量这使得它天然适合处理语音、文本等变长数据。2. LayerNorm的PyTorch实现与调试技巧让我们从最基础的LayerNorm实现开始。虽然PyTorch提供了现成的nn.LayerNorm但理解其底层机制对调试复杂模型至关重要。class CustomLayerNorm(nn.Module): def __init__(self, normalized_shape, eps1e-5): super().__init__() if isinstance(normalized_shape, int): normalized_shape (normalized_shape,) self.normalized_shape normalized_shape self.eps eps self.weight nn.Parameter(torch.ones(normalized_shape)) self.bias nn.Parameter(torch.zeros(normalized_shape)) def forward(self, x): # 计算均值和方差 mean x.mean(dim[-d for d in range(1, len(self.normalized_shape)1)], keepdimTrue) var x.var(dim[-d for d in range(1, len(self.normalized_shape)1)], keepdimTrue, unbiasedFalse) # 归一化 x_normalized (x - mean) / torch.sqrt(var self.eps) return self.weight * x_normalized self.bias实际使用时会遇到几个典型问题维度不匹配当输入张量的最后几维与normalized_shape不一致时会引发运行时错误。例如# 错误示例 norm CustomLayerNorm((64,)) x torch.randn(32, 128) # 第二维128 ! 64 out norm(x) # 报错梯度消失当eps值设置过小如1e-10时在方差接近零的情况下可能导致数值不稳定。初始化问题如果权重初始化为全零会使得输出始终为零破坏信息流动。调试建议在实现自定义LayerNorm时建议添加以下验证代码def test_layernorm(): # 测试2D输入如序列任务 x torch.randn(10, 64) # (batch_size, feature_dim) custom_norm CustomLayerNorm(64) official_norm nn.LayerNorm(64) # 前向一致性检查 assert torch.allclose(custom_norm(x), official_norm(x), atol1e-6) # 反向传播检查 x.requires_grad_(True) out_custom custom_norm(x) out_official official_norm(x) out_custom.sum().backward() grad_custom x.grad.clone() x.grad.zero_() out_official.sum().backward() grad_official x.grad.clone() assert torch.allclose(grad_custom, grad_official, atol1e-6) print(所有测试通过)3. 多头注意力机制(MHSA)的工程实现多头注意力是Transformer的核心创新点其实现需要考虑效率、数值稳定性和可扩展性。我们先看基础的自注意力实现class MultiHeadAttention(nn.Module): def __init__(self, d_model512, n_heads8, dropout0.1): super().__init__() assert d_model % n_heads 0 self.d_k d_model // n_heads self.n_heads n_heads self.d_model d_model # 线性变换层 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) self.dropout nn.Dropout(dropout) self.scale 1.0 / math.sqrt(self.d_k) def forward(self, q, k, v, maskNone): batch_size q.size(0) # 线性变换并分头 q self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) k self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) v self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 计算注意力分数 scores torch.matmul(q, k.transpose(-2, -1)) * self.scale if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # 注意力权重 attn_weights F.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 上下文向量 context torch.matmul(attn_weights, v) context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.w_o(context)实际工程中需要注意的几个关键点内存优化当序列长度较大时如超过1024注意力矩阵会消耗大量内存。解决方案包括使用内存高效的注意力实现采用分块计算策略混合精度训练数值稳定性缩放因子1/sqrt(d_k)对防止梯度消失至关重要对softmax输入做减最大值处理x - x.max()可增强数值稳定性并行计算利用torch.baddbmm替代循环计算合理设置batch_size和n_heads以充分利用GPU并行能力性能优化技巧在推理阶段可以通过以下方式提升效率# 缓存K和V的计算结果适用于自回归解码 class CachedAttention(MultiHeadAttention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.cache_k None self.cache_v None def forward(self, q, k, v, maskNone, use_cacheFalse): if use_cache: if self.cache_k is None: # 首次调用 self.cache_k k self.cache_v v else: # 后续调用拼接新token self.cache_k torch.cat([self.cache_k, k], dim1) self.cache_v torch.cat([self.cache_v, v], dim1) k, v self.cache_k, self.cache_v return super().forward(q, k, v, mask)4. 前馈网络(FFN)的实现与变体Transformer中的前馈网络虽然结构简单但有几个关键设计值得关注class PositionwiseFFN(nn.Module): def __init__(self, d_model, d_ff2048, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.activation nn.GELU() # 比原始ReLU表现更好 def forward(self, x): return self.linear2(self.dropout(self.activation(self.linear1(x))))FFN的几种常见变体及其特点GELU激活高斯误差线性单元(Gaussian Error Linear Unit)在实践中表现优于ReLUdef gelu(x): return 0.5 * x * (1 torch.tanh(math.sqrt(2 / math.pi) * (x 0.044715 * x**3)))门控机制如GLUGated Linear Unit变体class GLUFFN(nn.Module): def __init__(self, d_model, d_ff): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear_gate nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) def forward(self, x): return self.linear2(F.silu(self.linear1(x)) * self.linear_gate(x))参数共享在多任务学习中可以共享部分FFN层参数FFN实现中的常见问题及解决方案问题现象可能原因解决方案输出值过大/过小初始化不当使用Kaiming初始化训练后期性能停滞梯度消失添加残差连接推理速度慢中间维度(d_ff)过大使用Bottleneck结构多GPU训练不稳定各卡统计量不一致使用同步LayerNorm5. 完整Transformer层的集成与调试将上述模块组合成完整的Transformer层时需要注意几个关键细节class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff2048, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn PositionwiseFFN(d_model, d_ff, dropout) self.norm1 CustomLayerNorm(d_model) self.norm2 CustomLayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, maskNone): # 自注意力子层 attn_output self.self_attn(x, x, x, mask) x x self.dropout1(attn_output) x self.norm1(x) # 前馈子层 ffn_output self.ffn(x) x x self.dropout2(ffn_output) x self.norm2(x) return x集成测试时需要关注的指标梯度检查def grad_check(): layer TransformerEncoderLayer(d_model512, n_heads8) x torch.randn(10, 32, 512, requires_gradTrue) # (batch, seq, dim) # 前向传播 out layer(x) # 反向传播检查 out.sum().backward() for name, param in layer.named_parameters(): if param.grad is None: print(f警告{name} 无梯度) elif torch.all(param.grad 0): print(f警告{name} 梯度全零)数值范围监控各层输出值应在合理范围内如[-10, 10]注意力权重应介于[0,1]之间计算图验证# 使用torchviz可视化计算图 from torchviz import make_dot make_dot(out, paramsdict(layer.named_parameters()))实际项目中遇到的典型问题案例问题描述在训练初期损失值波动剧烈然后突然变为NaN诊断过程检查各层输出范围发现LayerNorm输出出现极大值追踪发现注意力分数在softmax前达到1e30量级确认d_k计算错误缩放因子失效解决方案修正缩放因子计算添加梯度裁剪初始化权重缩小10倍6. 高级技巧与性能优化当Transformer模型规模扩大时需要考虑以下优化策略混合精度训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, targets in dataloader: optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()激活检查点减少内存消耗from torch.utils.checkpoint import checkpoint class MemoryEfficientEncoderLayer(TransformerEncoderLayer): def forward(self, x, maskNone): return checkpoint(super().forward, x, mask)自定义内核融合使用torch.jit.script或Triton编写融合操作torch.jit.script def fused_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float): scores torch.matmul(q, k.transpose(-2, -1)) * scale attn torch.softmax(scores, dim-1) return torch.matmul(attn, v)不同规模模型的实现差异参数量级实现重点典型优化手段100M开发效率原生PyTorch实现100M-1B计算效率混合精度、梯度检查点1B分布式训练模型并行、流水线并行7. 测试与验证策略为确保实现的正确性建议建立全面的测试套件一致性测试与参考实现对比def test_vs_huggingface(): # 对比HuggingFace实现 from transformers import BertModel our_model OurTransformer() hf_model BertModel.from_pretrained(bert-base-uncased) # 权重转换后比较输出 assert torch.allclose(our_model(x), hf_model(x)[0], atol1e-4)梯度检验from torch.autograd import gradcheck input torch.randn(10, 64, dtypetorch.double, requires_gradTrue) test gradcheck(CustomLayerNorm(64), input, eps1e-6, atol1e-4) print(梯度检验通过:, test)数值稳定性测试def test_numerical_stability(): model TransformerEncoderLayer(d_model512, n_heads8) with torch.no_grad(): for name, param in model.named_parameters(): param.data.normal_(mean0, std100) # 极端初始化 x torch.randn(2, 256, 512) * 100 # 大输入值 out model(x) assert not torch.isnan(out).any() assert not torch.isinf(out).any()性能基准测试from torch.utils.benchmark import Timer def benchmark(): model TransformerEncoderLayer(d_model512, n_heads8).cuda() x torch.randn(32, 128, 512).cuda() # 预热 for _ in range(10): _ model(x) timer Timer( stmtmodel(x), globals{model: model, x: x} ) print(timer.timeit(100))8. 实际应用中的经验分享在工业级应用中我们发现以下几个实践对提升Transformer模块效果至关重要预热学习率策略配合Adam优化器使用线性预热def get_lr(step, d_model, warmup_steps4000): return d_model**-0.5 * min(step**-0.5, step * warmup_steps**-1.5)残差连接缩放在某些深层模型中效果更好class ScaledResidual(nn.Module): def __init__(self, scale0.1): super().__init__() self.scale scale def forward(self, x, residual): return x self.scale * residual注意力模式混合结合不同注意力模式提升效果class MixedAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.local_attn LocalAttentionWindow(d_model, n_heads, window_size64) self.global_attn MultiHeadAttention(d_model, n_heads) self.gate nn.Linear(d_model, 2) def forward(self, q, k, v, maskNone): local_out self.local_attn(q, k, v, mask) global_out self.global_attn(q, k, v, mask) gate F.softmax(self.gate(q), dim-1) return gate[..., 0:1] * local_out gate[..., 1:2] * global_out动态FFN宽度根据输入动态调整中间维度class DynamicFFN(nn.Module): def __init__(self, d_model, max_ff2048): super().__init__() self.controller nn.Linear(d_model, 1) self.linear1 nn.Linear(d_model, max_ff) self.linear2 nn.Linear(max_ff, d_model) def forward(self, x): ratio torch.sigmoid(self.controller(x.mean(dim1))) # [0,1] current_ff int(ratio * self.linear1.out_features) h self.linear1(x) h h[..., :current_ff] # 动态切片 return self.linear2(h)
