TorchJD 迎来大更新：多损失训练不再靠玄学调权重

多损失训练这件事，做过多任务学习的都知道有多恶心：几个 loss 加权求和，权重怎么给？拍脑袋定的 0.3、0.5、0.2，训到一半发现某个任务完全没学起来，回去接着调。TorchJD 这次的更新，算是把这件事从玄学往工程化推了一大步。

这个专门在 PyTorch 里做 Jacobian descent 的开源库，刚刚合入了一批新的聚合器实现，把学术界这几年提出的多目标优化方法几乎收齐了。对于每天在多任务、约束优化、正则项之间纠结的开发者来说，值得认真看一眼。

先说清楚：Scalarization 和 Jacobian Descent 到底差在哪

训练一个模型，如果只有一个 loss，故事很简单——算梯度、反传、更新。但真实场景里，你手上往往有一堆 loss：主任务的、辅助任务的、约束项的、正则项的。

处理它们，业界两条路：

• Scalarization（标量化）：把 N 个 loss 揉成一个标量。最粗暴的是直接平均，稍微讲究一点的用 GradNorm、DWA、Uncertainty Weighting 这类动态加权。优点是便宜——一次反传搞定，显存友好。
• Jacobian Descent（雅可比下降）：对每个 loss 分别算梯度，得到一个 Jacobian 矩阵（每行一个 loss 的梯度），然后用某种聚合策略把这堆梯度合成一个更新向量。目标不是让"平均 loss"下降，而是让"每个 loss 都下降"。

差别在哪？举个直观的例子：假设你在做推荐系统的多目标优化，点击率和转化率两个任务。如果这俩任务的梯度方向差不多，标量化没啥问题，加权平均就行。但如果它们的梯度方向严重冲突——一个往东、一个往西，加权平均出来的更新方向可能哪个任务都没优化好，甚至互相拖后腿。

Jacobian descent 的思路是：既然梯度冲突了，那我在聚合时就得考虑冲突，找一个能让所有 loss 都下降（或者至少不让某个 loss 变差）的方向。经典方法比如 MGDA、PCGrad、CAGrad 都属于这个流派。

代价是显存和计算。Jacobian 意味着你得为每个 loss 单独做一次反传（或者用一些技巧近似），显存开销是 O(N × 参数量) 级别的。

这次更新收了哪些东西

从项目组的说法看，v0.2 之后他们持续在补齐文献里的主流聚合器。现在库里能直接用的方法覆盖面已经相当广：

• UPGrad：TorchJD 团队自己主推的聚合器，思路是把每个梯度投影到"不损害其他任务"的方向再聚合，理论上有 Pareto 最优性质。
• MGDA / MGDA-UB：多目标优化的经典方法，找一个所有梯度线性组合中范数最小的方向。
• PCGrad：Google 那篇经典论文，梯度冲突时把一个投影到另一个的正交补上。
• CAGrad：在平均梯度附近找一个能保证每个任务不变差的更新。
• GradDrop、Nash-MTL、Aligned-MTL、DualProj 等一系列后续工作。

加上传统的 Scalarization 一族（Mean、Sum、加权、Random Weighting、GradNorm 等），基本上你在 paper 里读到的、想复现的，都能一行 API 换过来对比。

这在实际研发里价值不小。以前想对比 PCGrad 和 CAGrad，你得分别去两个作者的仓库扒代码，接口不一样、依赖不一样、PyTorch 版本还可能对不上。现在统一到一个库、一个接口、一致的语义，A/B 测试的门槛低了一大截。

用起来是什么感觉

TorchJD 的 API 设计沿用了 PyTorch 的习惯，把 loss.backward() 换成库里的 backward，传一个 loss 向量和一个聚合器进去。伪代码大致是这样：

from torchjd import backward
from torchjd.aggregation import UPGrad

aggregator = UPGrad()

for batch in loader:
    losses = model.compute_losses(batch)  # 返回一个 loss 张量，形状 [num_losses]
    optimizer.zero_grad()
    backward(losses, aggregator)
    optimizer.step()

比起自己手撸 PCGrad——注册 hook、手动算每个任务的梯度、投影、写回参数——干净得多。而且如果你想换算法，只要改一行 aggregator = CAGrad() 就行。

对于共享 backbone + 多任务 head 的经典 MTL 结构，库里还提供了 mtl_backward，只对共享参数做 Jacobian descent，任务专属参数走普通反传，节省一大截显存和算力。这个细节能看出团队是真做过 MTL 工程的——学术版本的 PCGrad 直接对全模型算 Jacobian，实际跑起来显存爆炸。

适合谁用，不适合谁用

说点判断。

值得试的场景：

• 多任务学习：尤其是任务间目标有一定冲突的，比如推荐系统里的 CTR + CVR + 停留时长，或者自动驾驶感知里的检测 + 分割 + 深度。
• 约束优化：主 loss + 一堆硬性约束（比如公平性、安全性约束），传统拉格朗日乘子调起来很痛苦，Jacobian descent 的"每个都不能变差"语义更契合。
• 物理信息神经网络（PINN）：数据 loss + 物理方程残差 loss 出了名的难平衡，社区里已经有一批工作在用 Jacobian descent 类方法。
• 对比学习 + 辅助任务：主任务 loss + 各种辅助 loss，比起手调权重，让算法自己决定聚合方向靠谱得多。

别硬上的场景：

• 单一主 loss + 弱正则：正则项权重比主 loss 小两个数量级的那种，标量化足够了，上 Jacobian descent 是杀鸡用牛刀，还得多花显存。
• 超大模型：LLM 微调那种参数量级，做全参数 Jacobian 基本不现实，只能挑关键层做，收益有多少得实测。
• loss 数量特别多：几十上百个 loss，Jacobian 矩阵本身就成瓶颈了。

一个更大的背景

多损失训练这块，长期以来处在一个尴尬位置：paper 一堆、方法一堆，但工业界大多数团队还是在用最朴素的加权求和。原因不是方法不行，是工程门槛太高——每篇 paper 一套代码、复现困难、超参敏感、缺乏统一 benchmark。

TorchJD 这种"把一个子领域的方法收齐、统一接口、和 PyTorch 深度集成"的库，对推动方法落地是有实际意义的。有点像早年 timm 之于视觉模型、transformers 之于 NLP——不是发明了什么，但降低了使用成本，让更多人愿意去试。

从项目活跃度看，v0.2 到现在贡献者明显变多，issue 响应也快。文档写得不算特别好但够用，示例覆盖了 MTL、PINN、对抗训练等几个典型场景。开源协议是 MIT，商用无压力。

小结

如果你手上正在跑一个多损失的项目，标量化的权重调得心累，那 TorchJD 值得花半天时间接进去试一下。最坏结果是效果差不多，你确认了标量化够用；好一点的话，可能省下你后面一个月的调参时间。

方法归方法，工程归工程。这个库最大的价值不在于哪个聚合器多强——学术上哪个方法"最优"到现在都没有定论——而在于它把这条技术线从 paper 拉到了随手可用的距离。这对整个多任务、约束优化方向的研究和落地，都是件好事。

参考来源

• TorchJD GitHub 仓库 — 源码、文档与示例代码
• Reddit r/MachineLearning 讨论帖 — 作者本人发布的更新说明与社区讨论
• Reddit 早期版本讨论 — v0.2 发布时的社区反馈