TorchJD：用Jacobian下降法解决多任务学习的梯度冲突问题

TorchJD是一个PyTorch扩展库，实现了Jacobian下降算法，专门用于解决多任务学习中多个损失函数之间的梯度冲突问题。本文将围绕其背景、核心方法、使用方式、应用场景等展开介绍，帮助读者理解这一工具的价值与意义。

在深度学习领域，多任务学习让单个神经网络同时处理多个相关任务，但不同任务的损失函数常产生相互冲突的梯度方向。当两个梯度内积为负时，简单平均会导致某一任务性能下降。例如视觉模型中，分类任务倾向全局特征，定位任务关注局部细节，梯度方向相反时传统方法顾此失彼，模型易陷入次优解。

TorchJD引入的Jacobian下降算法改变了多任务优化范式。与传统梯度下降处理标量损失不同，它直接操作损失向量对应的Jacobian矩阵——矩阵每行代表一个损失函数对模型参数的梯度。通过分析矩阵结构识别梯度冲突，采取更智能的聚合策略，数学上更严谨，实践中更有效。

TorchJD提供10+种梯度聚合器，代表性的UPGrad（无冲突梯度投影）核心思想是：聚合前将每个梯度投影到对偶锥（包含与原始梯度非负内积的方向）。这种投影保证足够小学习率下，每一次参数更新对所有任务都是有益的，不会损害任何任务性能。

TorchJD贴合PyTorch用户习惯，仅需将传统loss.backward()替换为torchjd.autojac.backward(losses)，再用jac_to_grad()转换Jacobian为梯度。针对多任务场景的mtl_backward()函数，会分别计算任务特定参数的损失梯度，以及共享参数的Jacobian矩阵，既保留独立优化又解决共享参数冲突。

TorchJD不仅支持传统多任务学习，还适用于实例级风险最小化（如个性化推荐、联邦学习）。库中torchjd.autojac.jac()函数允许不存储.jac字段直接计算Jacobian，为复杂自定义算法提供基础。

TorchJD为PyTorch生态带来重要工具，不仅解决多任务学习的梯度冲突问题，更拓展了神经网络优化理论的可能性。对于寻求多目标平衡的深度学习从业者，这是值得深入探索的库。