Transformer 里的残差连接很多人已经当作“结构标配”来看待了子层输出后面接一个x sublayer(x)要么就在MultiHeadAttention和FFN外面各自包一层。做模型轻量化或者手写训练脚本时经常会有人问这个问题这个残差连接到底能不能去掉去掉之后模型还能不能正常训出来如果你也在考虑剪掉 Transformer 的残差连接这篇博客建议直接收藏。我会从三篇论文出发把残差连接的真正作用、剪掉后会遇到的坑、替代方案以及可复现的实验代码完整铺开。注意我推荐的三篇论文不是最新的三篇而是理解“为什么可以剪”的底层文献ResNet、ReZero、LayerScale。它们分别解释了残差连接为什么有效、如何用零初始化缩放替代显式相加、以及如何用可学习缩放保住深层表达能力。本文内容包括先看残差连接在 Transformer 中承担的任务再分析剪掉残差连接的可能收益和风险接着梳理三篇论文的核心思想与启示随后给出一套修改 PyTorch 代码去掉残差连接的通用模板最后是实验验证思路、性能观察方法和常见问题排查。适合关注 Transformer 结构优化、模型训练稳定性、显存与内存优化的读者。1. 三篇论文技术要点速览维度说明核心问题Transformer 的残差连接能否剪掉剪掉后如何保持训练稳定第一篇论文ResNet提出残差连接解决深层网络的退化问题和梯度传播困难第二篇论文ReZero残差分支乘以 0 初始化的缩放参数让网络初始为恒等映射加速收敛第三篇论文LayerScale残差分支乘以一个可学习小标量用可学习缩放替代固定连接剪掉残差连接的替代思路初始化重置、缩放参数、重参数化、调整子层输出分布主要收益减少恒等分支的存储与计算简化结构可能降低显存和内存占用主要风险训练不稳定、梯度消失或爆炸、表征退化、收敛变慢验证方式小规模同参数量模型对比训练观察 loss、梯度范数、验证精度适合读者研究 Transformer 内部结构、做模型轻量化、复现论文实验的工程师这里先把结论放在前面完全剪掉残差连接在标准 Transformer 上直接训练基本都会出问题但借助 ReZero、LayerScale 这类“缩放代替连接”的思路可以在形式上把显式x f(x)改成x * scale f(x)甚至让scale0起步达到类似“剪掉”的效果。实际做结构改动时一定要有对应的稳定化手段不是简单把加法删掉就完事。2. 残差连接在 Transformer 中承担的任务Transformer 的原始结构里每个子层都包含一个残差连接output LayerNorm(x Sublayer(x))不少实现会把 LayerNorm 放在前面形式变成output x Sublayer(LayerNorm(x))这两种写法在训练稳定性上有差异但残差连接本身的作用是一致的。第一残差连接提供恒等映射路径。网络在初始化时子层输出接近零整体模块近似一个恒等函数输入信息可以直接流向后层。这个过程保证了深层网络的初始状态不会因为堆叠而产生剧烈扰动。第二残差连接改善梯度传播。反向传播时梯度可以直接通过短接线传到前面层避免因为一连串矩阵乘法导致梯度幅度指数级衰减。尤其是在深层 Transformer 中梯度主要依赖残差路径回传子层自身的梯度贡献反而有限。第三残差连接配合 LayerNorm 维持特征分布。Transformer 的特征流不是普通卷积网络里的图像特征而是序列化的 token 向量。残差连接让每一层的输入保持在可预测的尺度附近后续 LayerNorm 才能稳定计算均值与方差。第四残差连接让“层间变化”局部化。子层网络只需要学习相对于输入的增量而不是从头重构整个输出。这种做法降低了优化难度也让模型可以堆到几十上百层。如果直接去掉残差连接相当于强迫每个子层从头学习完整映射。这个映射不仅包含特征变换还包含上一层的已有信息优化难度会成倍增加。小型模型可能还能勉强收敛深层模型大概率出现训练不稳定的问题。3. 为什么有人想剪掉残差连接残差连接看起来只是一次加法但它在实现和运行中并不是“零成本”的。第一额外的显存和内存占用。残差连接在反向传播时要求保存短接线上的中间张量。虽然加法本身计算量小但为了梯度回传前向过程中x的副本在部分框架中会被保留到反向阶段。对于长序列、大 hidden size 的训练场景这一部分显存占用并不算小。第二多分支结构限制硬件算子融合。残差连接是典型的“合并两个分支”结构主分支x和子分支f(x)需要元素级相加。编译器或推理引擎在做算子融合时要处理两个输入张量的对齐和拷贝不像单一分支那样容易展开成连续内存操作。第三结构冗余。很多工作发现残差连接在训练后期更多承担“恒等映射保存”的作用子层真正贡献的特征变化集中在残差之外的部分。既然如此能不能用更轻量的缩放机制替代显式加法同时保留恒等映射的稳定性这正是后面三篇论文的切入点。第四为了结构搜索和重参数化。剪掉残差连接会让网络变成严格的链式结构大幅简化部署和推理。如果能在训练阶段保留某种等价的缩放路径在推理阶段再把缩放合入相邻线性层就能得到一个没有残差分支的网络。需要注意的是以上收益大多来自结构简化和条件释放而不是“残差连接本身算力很高”。如果你想通过剪掉残差来明显提升推理速度必须先测量当前实现中残差分支带来的算子执行开销再判断这个优化值不值得做。4. 三篇论文的核心思路与启发4.1 ResNet为什么要引入残差连接ResNet 一篇实际上解决的是层数加深后出现的“退化问题”网络层数增加训练误差不仅没有下降反而升高。作者给出的方案就是残差连接把基础映射H(x)改写成F(x)x让网络去学习残差F(x)。对 Transformer 来说ResNet 的启发很直接没有残差连接深层网络很难训练这不是 Transformer 独有的问题。任何堆叠子层结构都需要一条跨距离传播梯度的路径。所以如果你想剪掉残差连接就必须用其他方式提供至少一条“梯度高速公路”否则深层模型几乎必然退化。4.2 ReZero让残差分支初始化为零ReZero 的核心做法是在每个残差分支的输出上乘一个可学习的标量alpha并且把alpha初始化为 0。这样网络初始时完全等价于恒等映射每个子层什么都不做。随着训练推进alpha从 0 逐渐变大子层开始慢慢学习。这种设计对剪掉残差连接的启发最大。因为 ReZero 证明了一个关键事实在 Transformer 中真正稳定训练的不一定是“残差加和”这个结构而是“初始时子层输出为 0”带来的恒等映射特性。也就是说你可以把x f(x)改成x alpha * f(x)并且让alpha初始为 0从效果上等效于“第一个训练阶段剪掉了残差分支”但随着训练推进又保留了学习的可能性。ReZero 的另一个价值是加速收敛。因为网络初始为恒等映射梯度可以像在浅层网络中一样顺畅传播甚至可以使用更大的学习率。如果你想做“剪掉残差”实验ReZero 是一个必须读的对照组设计。4.3 LayerScale用可学习缩放替代固定连接LayerScale 来自视觉 Transformer 的改进工作核心做法是对每个残差分支的输出乘一个可学习标量一般初始化为很小的值例如1e-6到1e-4。它和 ReZero 的区别在于初始化不为 0且缩放参数随层独立学习。LayerScale 的启发在于残差连接可以拆成“恒等映射 子层输出 放缩系数”三个部分。如果放缩系数足够小子层输出在初始阶段几乎不参与前向传播如果放缩系数最后训练到接近 0说明该层的残差分支被模型自动“剪掉”了。也就是说LayerScale 提供了一种自动化剪除残差分支的动态机制而无需我们手动修改结构。从这三篇论文可以总结出一条清晰的主线残差连接的实质作用不是“加法”而是“提供跨层恒等路径”这个路径可以用一个可学习的缩放参数替代缩放参数初始化为 0 或很小值时网络初始状态和没有残差分支几乎等价如果训练结束后某些缩放参数趋近于 0那对应层在推理时完全可以剪掉残差分支。5. 剪掉残差连接的具体实验方法下面给出一个通用的 PyTorch 修改思路。注意这是实验模板不是某篇论文的精确复现路径和超参数需要按实际项目调整。5.1 标准 Transformer 子层标准的 Pre-LN Transformer 子层通常长这样import torch import torch.nn as nn class TransformerLayer(nn.Module): def __init__(self, d_model, d_ff, nhead, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(d_model, nhead, dropoutdropout, batch_firstTrue) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x): # Pre-LN 写法 h x self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]) h h self.dropout(self.ffn(self.norm2(h))) return h这里每个子层都做了x sublayer(x)也就是残差连接。5.2 去掉残差连接引入零初始化的缩放参数借鉴 ReZero 和 LayerScale我们可以把加法改写成“恒等项 缩放后的子层输出”class NoResidualTransformerLayer(nn.Module): def __init__(self, d_model, d_ff, nhead, dropout0.1, init_scale0.0): super().__init__() self.attn nn.MultiheadAttention(d_model, nhead, dropoutdropout, batch_firstTrue) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) # 两个可学习缩放参数 self.alpha1 nn.Parameter(torch.full((), init_scale)) self.alpha2 nn.Parameter(torch.full((), init_scale)) def forward(self, x): # 形式上仍然是 x alpha * f(x) h x self.alpha1 * self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]) h h self.alpha2 * self.dropout(self.ffn(self.norm2(h))) return h这个模板看起来还保留着加法但如果你把alpha1和alpha2初始化为 0第一个训练 step 中前向输出就是x相当于子层完全被跳过。随着训练推进alpha可能逐渐离开 0子层才真正参与学习。这正好符合“剪掉残差连接”的初始状态。如果你想做“完全无残差”的硬实验直接把x 删掉class HardNoResidualTransformerLayer(nn.Module): def __init__(self, d_model, d_ff, nhead, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(d_model, nhead, dropoutdropout, batch_firstTrue) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x): h self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]) h self.dropout(self.ffn(self.norm2(h))) return h这种结构只适合极小规模的 toy model。直接用在标准 BERT/Large 规模上大概率出现训练不收敛或严重退化。它的存在价值是作为“完全去掉连接”的底线对照组用来证明“没有恒等路径”会带来多大的训练困难。6. 实验设计与训练验证剪掉残差连接不是一个“改改模型就能跑”的操作必须设计一套可对比的验证流程。6.1 最小的可对比实验建议先做一组相同参数规模的小模型对比Transformer 层数建议 4 到 8 层。对比组至少包含Baseline标准残差连接ReZero 组alpha初始化为 0LayerScale 组alpha初始化为1e-4硬剪除组完全删除加法。使用同一个数据集、同一种学习率调度、同样的 batch size 和训练步数。重点观察三个指标训练 loss 是否正常下降最后一层梯度范数是否稳定验证集准确率是否达到 Baseline 的 95% 以上。6.2 训练脚本模板下面给出一段通用的训练脚本骨架主要用来跑 toy task例如文本分类或语言模型小规模测试import torch from torch import nn, optim def train_epoch(model, dataloader, optimizer, loss_fn): model.train() total_loss 0 for batch in dataloader: x, y batch optimizer.zero_grad() logits model(x) loss loss_fn(logits, y) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader) model NoResidualTransformerLayer( d_model128, d_ff512, nhead4, dropout0.1, init_scale0.0 ) # 注意这里只定义了一层实验时应该用多层堆叠 WordEmbedding 输出层组成完整模型训练时还可以加一个梯度范数打印def log_grad_norm(model): total_norm 0 for p in model.parameters(): if p.grad is not None: total_norm p.grad.norm().item() ** 2 return total_norm ** 0.5观察梯度范数很有用。硬剪除残差连接后梯度范数很容易在深层位置掉到1e-6以下这说明梯度传播路径断了。ReZero 组由于初始化为 0梯度可以直接流经恒等路径梯度范数通常比较稳定。6.3 判断实验是否成功判断标准不要只看最终 accuracy。对于“剪掉残差连接”这项结构实验至少满足以下三条才值得继续训练 loss 在相同步数内下降到与 Baseline 同一数量级验证集指标下降不超过 2% 到 3%训练过程中的梯度范数没有出现持续衰减或剧烈震荡。如果硬剪除组能在一个 4 层小 Transformer 上跑通说明该模型规模较小线性层之间的耦合还不够深换到 12 层以上问题会迅速暴露。7. 资源占用与训练性能观察很多人剪残差连接的主要目的是省显存。要验证这个收益需要对比同一个模型在“有残差”和“无残差”两种情况下的显存占用和训练速度。观察方式用nvidia-smi查看训练进程的显存占用用训练脚本统计每个 epoch 的平均耗时输入同样 batch size、同样序列长度确保对比条件一致如果显存下降明显要继续验证是“省掉了残差副本”还是“batch size 变小了”。需要注意x sublayer(x)中和x等大的中间张量在反向传播时是否被保留取决于实现框架和算子融合策略。有些框架会自动复用内存有些则会额外分配显存。因此剪掉残差连接带来的显存收益必须实测不能凭理论推算。如果你用的是长序列场景例如 4096 或 8192 长度残差连接的显存影响会更明显。这时可以采用 ReZero 方式的零初始化缩放在初始阶段让子层输出为 0反向时该分支的梯度贡献极小从而让框架有机会释放部分中间张量。不过最终能省多少仍然以本机nvidia-smi实测为准。8. 常见问题与排查方法问题现象可能原因排查方式解决方案训练 loss 不下降完全删除残差连接后梯度无法传播到浅层打印各层梯度范数检查是否趋近于 0改用 ReZero/LayerScale保留恒等路径并加上零初始化或小初始化缩放模型收敛速度极慢子层输出在初期扰动太大特征分布不稳定观察首个 step 的 loss 与 Baseline 差异将缩放参数初始化为 0 或1e-6让网络从恒等映射起步验证集指标大幅下降模型表达能力不足删掉连接后信息流受限比较每个子层输出幅值检查是否某些层完全失效尝试只删部分层的残差连接保留关键层的残差路径训练后期 loss 震荡alpha学习率设置不合理打印alpha的梯度与数值变化为alpha设置独立的小学习率或加权重衰减显存没有下降框架没有释放恒等分支中间张量对比有/无残差两种实现的峰值显存记录使用torch.utils.checkpoint或手动优化前向传播中的张量生命周期推理速度提升不明显残差连接加法本身计算量很低吞吐瓶颈在矩阵乘法使用 profiler 分析各算子耗时先量化矩阵乘法或优化 Attention再考虑结构剪除如果你在实验中发现 ReZero 组的最终准确率低于 Baseline不必奇怪。因为 Residual connection 在标准 Transformer 中不仅是稳定性工具还可能提供了一层“特征加法”的表示能力。ReZero 和 LayerScale 在部分任务上需要额外调参才能追平 Baseline。9. 最佳实践与使用边界剪掉残差连接在工程上真正有价值的不是“一刀切删除”而是“让残差分支可以被学习地关闭”。建议从以下路径开始先用 ReZero 或 LayerScale 替换显式残差加和保持训练稳定训练结束后检查alpha的数值分布找到趋向于 0 的层把这些层的残差分支从推理脚本中删除实现结构简化如果要继续压缩模型可以将缩放参数合并进相邻线性层的权重中得到真正无残差分支的部署模型。不推荐直接在完整 BERT/Large 或 GPT 规模上硬删残差连接。除非你同时调整初始化方法、学习率策略和调度方式否则训练失败概率非常高。另外做这类实验时请使用自己生成或已获授权的数据。如果涉及第三方文本、图像或语音数据务必确认版权和许可边界。结构优化实验和具体部署之间的差距很大最终效果要以目标硬件上的实测为准。10. 总结与下一步想剪掉 Transformer 的残差连接先不要去删代码里的。先读三篇论文ResNet 让你理解残差连接为什么出现ReZero 告诉你零初始化缩放可以替代残差分支LayerScale 教你用可学习缩放决定每层残差是否保留。在这三条思路的基础上再设计“删除残差”的实验方案才可能得到稳定可用的结果。下一步可以这么做先搭建一个 4 层小 Transformer跑通 Baseline然后改成 ReZero 形式观察收敛速度和最终效果再做硬剪除实验记录梯度范数和 loss 曲线。当你看到硬剪除导致训练崩溃就会明白残差连接设计的底层价值当你把 ReZero 跑通就知道剪掉残差连接的正确姿势是“用缩放代替连接”而不是“直接删掉”。
阅读完成 · 觉得有帮助?