ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

双头模型架构实战:共享底座与多任务分支的深度拆解

双头模型架构实战:共享底座与多任务分支的深度拆解 双头模型这个概念我在很多项目里被反复问到尤其是刚接触多任务学习的同学总以为“双头”就是把两个模型塞进同一个网络里甚至连输入都要拆成两路。其实真正落地的时候绝大多数情况都是共享底座、顶上接两个分支。这个系列写到现在第25章我把这个架构彻底拆开讲一遍从为什么这么设计、到底解决什么问题到代码怎么落、踩过哪些坑一次说清楚。这篇内容适合下面这几类人想给同一个模型同时做分类和回归的比如既预测点击率又预测转化金额在做对比学习或者双塔召回需要两个分支输出向量做相似度计算的以及手里数据量不算大又不想为每个任务单独训练一套完整模型的人。看完这篇你至少能自己动手搭出一个不算水的双头模型并且知道训练过程中常见的问题该怎么查。1 整体设计与思路拆解1.1 双头模型到底长什么样一句话讲清楚双头模型 一个共享底座 两个任务分支。底座负责把原始输入变成一份通用的特征表达两个分支在这个特征之上各自干各自的活。听着简单但很多人理解偏了以为双头就是两个全连接层叠在最后其实完全不是一回事。先看一个我实际做过的场景电商场景里同时预测用户会不会点击二分类和如果点击了会花多少钱回归。如果照老思路那得训练两个完全独立的模型一个做分类一个做回归。每个模型都要自己去学商品特征、用户特征、上下文特征的组合规律。可现实是用户行为底层的那套特征规律是共通的——高转化商品的图像、描述、价格区间无论对点击还是对金额都有相似的表达逻辑。双头模型的思路就是抓住这份共通性底座把输入映射到一个表达能力很强的特征向量两个轻量级分支分别在这个向量上做分类头和回归头。底座的参数是两个任务共享的只在最后几层分开。对应到代码上结构大致是这样class DualHead(nn.Module): def __init__(self, backbone, head_a, head_b): super().__init__() self.backbone backbone # 共享底座 self.head_a head_a # 分支A self.head_b head_b # 分支B def forward(self, x): feat self.backbone(x) out_a self.head_a(feat) out_b self.head_b(feat) return out_a, out_b这里关键点是backbone的输出是所有后续分支的唯一输入。分支之间没有交互底座是唯一的信息交换节点。这跟某些把任务输入拼在一起再送进模型的“伪多任务”做法有本质区别。1.2 它解决的核心问题双头模型存在的理由说到底就是三件事省资源、抗过拟合、利用任务关联性。省资源很好理解。底座是网络里参数最重的那部分两个任务共用一份底座参数量从“两套模型”变成了“一套底座 两个小头”训练和推理的显存占用、耗时都少了将近一半。我在任务里测过同样一个文本分类加意图判断的双头模型跟两套完全独立的模型比单次推理时间从 4.2ms 降到了 2.3ms显存占用更是只有原来的六成左右。抗过拟合这一点常被忽略。底座学到的特征是经过两个头共同“审查”的这本身就是一个天然的正则化机制。一个头只对某一类特征敏感另一个头会对特征的不同维度施加约束底座为了同时满足两个头就不得不学到更通用的表达而不是过拟合到单任务上那些细枝末节的噪声特征。比方说表情符号分类头可能觉得某个表情跟“投诉”强相关回归头却认为它跟“退款金额”没什么关系底座最终学到的就会是一个更稳健的语义向量。利用任务关联性也很直观。两个任务往往共享底层信息最典型的例子就是搜索排序里的“相关性判断”和“点击率预估”。相关性高不一定点击率高但两者都依赖于对“用户查询意图”和“文档内容”之间语义关系的理解。如果分开建模型这份理解就要学两遍即使数据量够计算量也是白白浪费。共享底座让两个任务的梯度共同塑造底层语义建模最终两个指标都比独立训练时有一定提升。1.3 为什么不是“两个独立模型结果融合”你可能会问那我保留两套模型最后把两个结果加权融合效果是不是一样不一样而且是本质上的不一样。独立模型是两份完全独立的特征空间。A模型内部的向量和B模型内部的向量之间没有对齐关系融合只能发生在最后的概率得分层比如把 pctr 和 cvr 线性加权。这个融合是“事后诸葛亮”丢失了特征层面的交互信息。双头模型不同它强制两个任务在共享底座输出的特征空间里对齐。底座给出来的每个特征维度必须对两个分支都有信息量这会让特征表达同时兼容两种任务的“口味”。举个例子用户历史点击序列里“最近是否反复查看同类商品”这个特征对点击率和转化金额都有意义但在独立模型里两个模型对它的编码方式可能完全不同。双头模型则让这个特征在两种任务中被同时激活和修正表达一致性更高。尤其在做向量检索或者需要特征 embedding 复用的场景下这种对齐价值更大。我做过一个推荐召回模型底座输出的是一个 128 维向量一个分支做分类一个分支做 embedding 匹配。这样同一个向量既能参与分类打分又能做近邻检索训练时两份梯度共同优化它最后召回效果比单独训一个 embedding 模型再硬塞分类头要稳定得多。2 核心细节解析与实操要点2.1 共享底座怎么选底座没有统一的模板得看你的输入是什么类型以及任务对特征抽象层级的要求。处理结构化表格数据时底座一般就是几层全连接加激活函数输出特征向量在 16 到 64 维之间。这时候底座不要设计得太深否则两个头会同时遭受梯度消失的困扰我实测超过 5~6 层全连接底座某些任务上训练速度反而显著下降。处理图像时底座可以是 ResNet 系列或者 Vision Transformer 的前若干层。注意这里有个具体选择是取倒数第二层特征还是取 Transformer 最后一层加个全局池化。我的经验是如果两个分支一个是分类、一个是分割那么 VIT 的输出需要保留空间维度不应该直接全局池化成一条向量。如果两个分支都是全局层面的任务比如粗分类 细分类那直接把[CLS]token 拿出来就好。处理文本时底座通常是一个预训练语言模型分歧点在于分支从哪一层特征接入。这里给一个很实用的经验分类头适合用最后 1~2 层输出而做实体级别或者细粒度情感头的建议从倒数第二、第三层接入文本里有些细粒度信息太接近顶层时会被抽象掉这点很反直觉但确实频繁出现。底座参数规模的决策准则是底座决定了模型对原始信息的保留量分支决定了高层面向任务的表达能力。数据量大、任务复杂底座就做大数据量有限底座做小、分支多堆几层抗过拟合能力会好很多。2.2 两个分支的结构设计要点分支的设计核心是把底座输出特征“翻译”成任务需要的输出形式。分类分支最常见的写法是几层 MLP 后接 Softmax输出类别概率分布。这部分没什么花活但有几个细节要提醒分类分支的最后一层不要加 Dropout否则推理时概率分布会显得过于平滑这一点在线上 A/B 测试时看得很明显。回归分支要注意激活函数的选择。如果预测目标是金额、时长这种天然非负的值最后输出层不要用裸的 Linear最好接一个Softplus或者ReLU防止模型在训练初期输出负数导致 loss 震荡。我踩过一个大坑预测订单金额时输出层用了纯线性前几百个 batch 模型疯狂输出负金额反向传播的梯度方向全被带歪了后来接上Softplus才稳定下来。嵌入分支或对比分支一般输出一个固定维度的向量然后通过 L2 归一化也就是F.normalize把向量约束到单位超球面上。这一步极其关键如果不归一化模型会钻空子把向量模长拉大来减小 loss聚类和召回效果直接崩掉。做这个分支的时候我还建议在归一化前加一个带偏置的线性层别小看这个偏置它对冷启动特征很少的样本help 很大。另外两个分支内部的第一层宽度一般不相等也没关系但它们各自跟底座输出特征的匹配非常关键。底座输出的特征如果本身已经比较高维比如 512分支第一层直接压到个位数维度信息损失会非常严重。我一般建议第一层压缩比不超过 4:1。2.3 损失函数的设计与权重分配双头模型的核心难点之一就是两个分支的损失怎么合并。最常见的做法是加权相加loss w1 * loss_head_a w2 * loss_head_b这里 w1、w2 不是随便设的。分类分支的交叉熵损失量级通常很小0.1 ~ 1回归分支的 MSE 损失量级可能很大几百到几千。如果直接把两个 loss 加起来哪怕回归损失只占权重 0.001它也会在反向传播时主导梯度分类分支完全没法学起来。我常用的两个方案第一个方案是手动归一化把回归任务的标签做标准化减均值除标准差让回归损失变得相对集中在 0.1 ~ 2 这个区间。这样两个 loss 的量级就比较接近了权重可以简单设成一样或者根据业务重要性微调。第二个方案是学习一个噪声参数。这个做法在学术上叫不确定性加权核心思路是给每个任务分配一个可学习的参数用来在训练中自动调整任务权重避免手动去调。代码实现也不复杂log_var_a torch.zeros((1,), requires_gradTrue) log_var_b torch.zeros((1,), requires_gradTrue) loss_a nn.functional.cross_entropy(out_a, label_a) loss_b nn.functional.mse_loss(out_b, label_b) loss_total 0.5 * torch.exp(-log_var_a) * loss_a log_var_a loss_total 0.5 * torch.exp(-log_var_b) * loss_b log_var_b注意 log 方差初始值最好是 0避免一开始就压低某个任务的权重。这个方案在任务间量级差异极大的场景里非常管用我拿它处理过一个“分类 回归”双头模型比手动调权重节省了至少十几次训练迭代。另一种组合玩得也比较多叫对比损失加分类损失。对比损失让同类的样本向量靠近、异类样本相互推开分类损失则直接用交叉熵监督类别。共享底座同时收到来自两个分支的不同梯度信号这种组合能显著改善特征空间的判别力。2.4 梯度的隔离和反向传播双头模型的梯度流跟单头模型完全一样都是从头端传到分支再从分支汇聚到底座。但这里有个容易被忽视的问题两个分支回传到共享底座时梯度是累加的。如果两个分支数值范围差得离谱比如 A 分支梯度的范数在 1e-3B 分支梯度范数在 1e2累加之后 B 分支完全盖住 A 分支共享底座几乎只按 B 分支的需求更新。解决这个问题的思路就两个方向要么控制损失权重要么直接对分支梯度做裁剪或者归一化。我在实践中更推荐控制损失权重 监控梯度范数双管齐下。训练时每隔固定步数打印一次各分支最后一个线性层梯度的 L2 范数把数值记录到日志里一旦发现某个分支的梯度范数长期高出另一个分支 10 倍以上就着手调整该分支的损失权重。不夸张地说双头模型百分之七八十的训练不稳定问题最后都能追溯到梯度范数失衡上。另外还有一个常用的操作叫detach或者stop_gradient。什么时候用当你希望某个分支的特征表达不要反向影响到底座时。比如说你做一个图像分类 风格迁移的双头模型你想让分类分支帮助底座提取语义信息却不想让风格分支过于强力的梯度干扰底座的语义表达。那就在该分支输入到底座输出特征之前调detach()让这个分支的梯度只更新它自己的参数不回传到共享底座。x torch.cat([a, b.detach()], dim-1)注意detach是个双刃剑用多了会丢失共享底座对分支的调节能力务必只在你确实需要隔离任务梯度时使用。3 实操过程与核心环节实现3.1 从一个真实场景出发为了讲清楚全过程我拿一个非常常见的场景来走一遍假设你有一批用户的浏览历史和行为数据特征是 32 维的向量目标是同时预测两个东西——这个用户接下来会不会购买某类商品二分类以及如果购买的话预计消费金额回归。这种任务非常适合用双头模型。第一个头分类预测购买概率第二个头回归预测金额。共享底座需要做到的是把用户行为特征抽象成能够同时支撑两个目标的表达。前三步数据标准化、划分训练测试集、建立基线模型。基线模型我用的是一个单任务模型只预测购买概率的 AUC 和只预测金额的 RMSE后面才能对比双头模型有没有实际收益。3.2 完整代码实现我用 PyTorch 写一个尽量简洁但足够完整的版本。底座用三层全连接分支分别做分类和回归。import torch import torch.nn as nn import torch.nn.functional as F class DualHeadModel(nn.Module): def __init__(self, in_dim32): super().__init__() # 共享底座 self.backbone nn.Sequential( nn.Linear(in_dim, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.3), nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.3), nn.Linear(64, 16), nn.ReLU(), ) # 分类分支 self.cls_head nn.Sequential( nn.Linear(16, 8), nn.ReLU(), nn.Dropout(0.2), nn.Linear(8, 2), # 购买/不购买 ) # 回归分支 self.reg_head nn.Sequential( nn.Linear(16, 8), nn.ReLU(), nn.Linear(8, 1), nn.Softplus(), # 保证输出非负 ) def forward(self, x): feat self.backbone(x) out_cls self.cls_head(feat) out_reg self.reg_head(feat).squeeze(-1) return out_cls, out_reg model DualHeadModel(in_dim32)关于这个结构有几个细节我想强调一下。BatchNorm1d放在全连接层之间有时候会引发小 batch size 下的训练不稳定。如果你训练时 batch size 小于 32建议把 BatchNorm 去掉换成 LayerNorm 或者干脆不用归一化层。我以前踩过坑batch size 设成 16一批样本的均值和方差统计量波动特别大双头模型训练曲线抖得像锯齿换成 LayerNorm 后明显平滑。回归分支的Softplus是刻意加上的。很多教程里回归头就是一层 Linear但对金额这种非负目标裸线性输出会在训练初期产生大量负预测值损失函数会被少数离群点主导加 Softplus 是一种非常实用的稳健化处理。3.3 训练循环和损失合并下面是最关键的训练部分。我在这个例子里用不确定性加权来合并损失省去手动调权重的麻烦。from torch.utils.data import DataLoader # 假设 train_loader 中每个 batch 是 (x, cls_label, reg_label) # cls_label 是 0/1reg_label 是标准化后的消费金额 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) optimizer torch.optim.Adam(model.parameters(), lr1e-3) # 不确定性参数 log_var_cls torch.zeros((1,), requires_gradTrue, devicedevice) log_var_reg torch.zeros((1,), requires_gradTrue, devicedevice) for epoch in range(30): model.train() total_loss 0.0 for x, y_cls, y_reg in train_loader: x, y_cls, y_reg x.to(device), y_cls.to(device), y_reg.to(device) out_cls, out_reg model(x) loss_cls F.cross_entropy(out_cls, y_cls) loss_reg F.mse_loss(out_reg, y_reg) # 不确定性加权合并 loss 0.5 * torch.exp(-log_var_cls) * loss_cls log_var_cls loss 0.5 * torch.exp(-log_var_reg) * loss_reg log_var_reg optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch:2d} | Loss {total_loss / len(train_loader):.4f} | fbias_cls {log_var_cls.item():.2f} | bias_reg {log_var_reg.item():.2f})这个训练循环里有个需要特别注意的地方log_var_cls和log_var_reg是放在model.parameters()之外的所以 optimizer 构造时不会更新它们。如果你想把它们也交给自己优化需要把两个参数加入 optimizer或者干脆用torch.optim.Adam(list(model.parameters()) [log_var_cls, log_var_reg], lr1e-3)。我推荐前者也就是让这两个参数保持独立用一个稍微小一点的学习率防止它们在训练前期波动过大打乱两个任务的平衡。损失合并时torch.exp(-log_var)的物理含义就是自动给损失大的任务降权、给损失小的任务升权这个机制在实战里非常稳。如果你的任务量级差异不大也可以直接写loss loss_cls 0.5 * loss_reg权重按业务重要程度调。3.4 推理阶段的分支操作训练完成后模型做推理时有几种用法你自己按业务取舍。第一种两个分支结果一起用。比如搜索引擎里分类相关性和回归点击率同时参与最后的排序打分。第二种只用其中一个分支。比如你当初引入第二个分支只是为了帮助底座学到更好的特征线上部署时只保留分类头。这种情况可以直接删掉另一个分支的权重或者推理代码里只调对应分支的 forward避免多余的计算。第三种把底座输出特征向量直接导出作为样本的 embedding 使用。这在召回、向量检索场景里最常见。底座输出 16 或者 64 维向量存入向量数据库线上用向量检索的方式做召回分类和回归分支在离线得到用来调试线上根本不用跑。推理代码需要注意的是如果模型里用了 Dropout 或者 BatchNorm记得调用model.eval()否则每个批次的特征分布都不同会导致同一个样本在不同 batch 下得到完全不同的预测结果。3.5 一套完整的调参顺序双头模型比单头模型多出来的调参空间主要就在任务权重和分支深浅。我个人的调参顺序是这样的帮你省点时间第一轮先调底座的深度和宽度两个分支保持简单。这个阶段看总 loss 有没有收敛不细分单个任务效果。第二轮开始调分支的宽度和深度同时观察两个分支各自在验证集上的表现。如果分类分支效果差优先增加分类分支的层数而不是动底座。第三轮调损失权重或者不确定性参数这时候如果某个分支的 loss 一直压不下去再调权重。最后一轮才动学习率、batch size、dropout 这些通用超参数。顺序千万不要反过来我见过很多人一上来就调损失权重结果共享底座本身就没设计好怎么调都白费。4 常见问题与排查技巧实录4.1 两个分支相互打架总 loss 震荡不降这是双头模型最经典的病。总 loss 要么在某个数值附近剧烈震荡要么反复回升到很高点再降下来。本质原因是两个分支的梯度方向在共享底座处发生冲突一个分支的更新方向把底座推往对另一个分支不利的区域。我的排查顺序是先看两个分支单独损失曲线。如果分类损失降得很快回归损失还在高位乱撞问题大概率出在回归分支梯度的量级过大。用我前面说的梯度范数监控打印出两个分支的梯度范数确认谁在主导更新。再看两个任务的相关性。有时候两个任务本质上是矛盾的比如既要推荐高点击率内容又要推荐高收藏率内容它们对内容的偏好确实存在冲突。这种情况下共享底座天然会学到一个“折中表达”两个任务指标都半吊子。应对策略是调整两个分支接入底座的位置。比如分类分支从底座最后一层接回归分支从倒数第二层接让两个分支享有不完全相同的特征层级缓解冲突。如果真的发现任务矛盾比较大我个人建议放弃双头模型回到两个独立模型后期再做结果融合。这没什么可丢人的双头模型不是万能的它更适合相关任务而不是矛盾任务。4.2 一个分支收敛快另一个收敛慢这种情况特别常见尤其是一个分支是简单的二分类另一个分支是连续回归任务。根本原因是两个任务的复杂度天然不同。怎么办第一招给收敛慢的分支分配一个独立的子网络把该分支自己变成更深一点的 MLP弥补底座特征无法覆盖的信息。第二招将收敛快的分支冻结一部分训练步数先专注训练慢的分支。具体实现就是在训练循环里判断当前的 epoch 数小于某个阈值时把收敛快的分支的参数requires_grad设为 False只让慢分支和共享底座更新。更平滑的做法是用渐进式权重调度训练初期把容易任务的权重调低后面逐步恢复正常权重。这样模型前期优先搞定难任务的特征表达后面再由易任务帮它精调。4.3 推理时发现底座特征质量很差底座输出的特征向量看起来数值分布比较怪比如大量维度的值都挤在 0 附近或者某些维度数值巨大但这些现象本身不一定说明底座坏了。需要先分情况判断。如果底座特征在分类分支上表现不错但直接拿去检索效果差那问题在于特征没有经过 L2 归一化距离度量不稳定。解决办法是在底座输出后面加一个归一化层而不是让分支各自去处理。如果底座的中间层激活值大量死亡比如 ReLU 输出全为 0是梯度消失导致底座学不出来。排查时看看是不是底座初始学习率过大或者底座结构太深。我遇到过因为底座权重初始化不合理导致双头模型前 10 个 epoch 特征向量几乎全是 0后来换用 Kaiming 初始化才缓解。另一个比较隐蔽的原因是 BatchNorm 导致的分布偏移训练集和验证集底层分布不一致时底座输出的统计量也不一致特征在训练集上表现正常在新数据上完全变了。这种情况可以对底座输出做分布对齐或者把 BatchNorm 换成 LayerNormLayerNorm 对样本量变化更不敏感。4.4 多分类头存在严重类别不平衡双头模型其中一个分支是多分类且类别分布极其不均比如 99.5% 的样本属于 A 类、0.5% 属于 B 类。这种时候交叉熵几乎是 crippled 的模型学到最后只会输出全 A 的概率分布。我的处理方式是给分类分支的交叉熵加类别权重直接用F.cross_entropy(..., weightclass_weight)class_weight 可以是类别频率的倒数也可以根据业务收益动态定义。这个操作只影响分类分支不会直接污染回归分支。更强硬的做法是把多分类分支换成 Focal Loss让模型聚焦于难分样本。关键是 Focal Loss 要求模型在训练前期能够对简单样本给出高置信预测否则所有样本都变得“难分”loss 反而更大。所以我的习惯是先跑几个 epoch 普通交叉熵再切换 Focal Loss双头模型整体收敛效果比一开始就用 Focal Loss 好不少。4.5 训练结束之后的稳定性和可复现性双头模型训练过程中我建议固定随机种子、固定数据加载顺序、固定 CUDA 的确定性模式。原因很简单双头模型有两个任务同时在影响底座更新训练过程的随机性会被放大稍有不慎同一套数据训出来的模型效果差异就很大。我在实际项目里测过不固定随机种子时双头模型训练三次分类分支的 AUC 在 0.78 到 0.83 之间浮动回归分支的 RMSE 也有十个点左右的波动。这个波动幅度对线上部署是完全不可接受的。另外保存模型时一定要把两个分支的 state_dict 和 log_var 参数一起保存。很多人只存了 model.state_dict()结果推理时发现不确定性权重丢了损失函数没法复现好几天之后想回看训练过程中的任务权重变化找不到记录非常痛苦。我的习惯是把超参数、两个 log_var 的值、训练停机轮数、随机种子全都写进训练日志或者直接打包成 json 存到模型同目录下。5 从双头到多头一次说清扩展思路5.1 三个及以上分支怎么处理双头模型代码写顺了以后扩到多头非常容易。只需要把两个头的结构改成字典映射self.heads nn.ModuleDict({ cls: nn.Sequential(...), reg: nn.Sequential(...), embed: nn.Sequential(...), })难点不在代码而在损失合并和任务平衡。三个分支意味着三份梯度的冲突概率更高不确定性加权升级成多任务版后每个任务都得维护一份 log_var可学习的参数变多了训练稳定性需要格外留心。我的建议是优先做两两配对训练。如果三个任务分别是 A、B、C可以先训练共享底座 A 分支再训练共享底座 B 分支最后再接上 C。每个阶段的底座参数用上一阶段初始化让底座的成熟度逐步提升最后再合并成三头模型微调。这个“顺序热身法”比从头直接训三头稳定得多效果也更可控。5.2 两个分支同时回传还是交替回传训练双头模型时每个 step 两个分支的梯度是同时回传的这是标准操作。但某些场景下交替回传可能更好当两个任务的样本是分离的时比如不可能一条样本同时带分类标签和回归标签你可以把一条 batch 的数据按任务拆开分两次 forward、累加梯度之后再 step。交替回传的收益是每个任务都能看到足够多的样本。损失合并时可以用简单的累加只要两个任务的 batch 大小接近梯度量级就不会差太多。缺点也很明显训练速度慢了将近一半而且共享底座参数更新的步频减半整体收敛时间反而变长了。我个人的经验是如果两个任务的标签在同一批样本里都存在同时回传如果标签是互斥的交替回传。5.3 共享底座不同层的特征是否要拼接有些做双头模型的同学会倾向于把底座中间某层特征和顶层特征拼接起来喂给分支觉得这样“信息量更大”。我的建议是谨慎操作。拼接操作会让分支的输入维度翻倍但中间层特征的语义和顶层特征的语义并不一致拼接之后分支需要自己学会怎么组合它们。如果任务本身对两种粒度的特征都有需求比如同时预测物体类别和物体位置那拼接确实有帮助。但如果是标准的分类 回归任务拼接往往只会增加过拟合风险并不会带来有效提升。从工程角度讲拼接还会增加梯度传播路径的复杂度。底座中间层和顶层都会收到分支的梯度两条路径的梯度尺度本来就不同合并时更容易失衡。所以在你没有充分理由之前默认只把底座最后一层输出给分支是最保守也最稳的做法。5.4 有没有必要预训练共享底座这个问题几乎每次都会被问到。如果你的双头模型是在解决一个全新的领域底座没有可用的预训练权重那我建议先用单任务跑一个简单的基线模型把它作为底座的初始化权重再开始搭双头。直接随机初始化共享底座并同时训练两个任务底座很容易陷入到一个不温不火的局部最优。如果你在NLP场景中做双头直接用预训练语言模型做底座然后让两个分支接着训练这已经是常规操作。有一点提醒预训练语言模型的顶层参数在微调的时候两个分支的梯度会同时作用改动幅度过大有时候会把语言模型原本很稳的语义空间搞乱。我的做法是底座最上面两层用较小的学习率分支完全放开用正常学习率。PyTorch 里可以对不同参数组设置不同的学习率操作起来并不麻烦。5.5 双头模型能否做增量任务这个扩展思路我觉得挺实用。你已经有一个训练好的单任务模型突然多了第二个任务需求就可以把原来的模型前半部分直接作为共享底座新增一个随机初始化的分支再去训练第二个任务。这里有个关键的坑新增分支的梯度在早期非常剧烈因为它的参数是随机的它会把底座往错误方向猛推。优化办法是给新增分支一个较小的初始学习率或者让新增分支的梯度冻结若干个 epoch先只训练它的参数等分支本身不会乱跳了再开放底座参数。这个方法我在一个文本分类项目上验证过新增任务分支后原有任务的准确率几乎没有下降新任务也从随机初始化一路爬到了一个可用的水平。做增量训练时还有一个细节旧分支的参数最好在训练新任务时保持较小学习率或者干脆冻结。因为你并不希望底座在适配新任务时把已经学好的旧任务特征彻底冲掉这样保留一份“记忆”的隐性正则化机制在序列化多任务场景里特别有用。6 写在后面的一点个人体会双头模型这个东西资深工程师操作起来很顺但新手经常认死理觉得“多一个头就是多一份能力”或者认为加了第二个分支就一定能免费提升第一个任务的指标。我做过一个系统回归纯从算法角度讲双头模型就是一个带共享参数的组合模型。你说它玄学吧底层的原理其实清清楚楚梯度的相互作用、特征表达的复用、任务相关性的利用这些都是可测量、可验证的东西。你说它不玄吧真正把一个双头模型调好确实要在任务平衡、梯度冲突、数据分布这些细节上花不少精力。所以我个人的建议是如果两个任务高度相关、计算资源吃紧双头模型是一个值得优先尝试的架构。如果两个任务本质相悖别硬凑在一起老老实实分开建模。学习这个过程最好的方式不是看多少理论而是把一个简单双头模型从零搭起来打印每一层的梯度观察两个分支如何彼此拖拽然后你就自然理解它所有的优势跟局限了。这个系列虽然叫《模型不玄学》但说到底能把架构拆清楚、能把问题调明白的才是真正的高手。
返回列表