通过辅助损失实现解耦的分割学习
机器学习
2026-02-10 v2 人工智能
摘要
分割学习是一种分布式训练范式,其中神经网络在客户端和服务器之间进行划分,允许数据保留在客户端,仅共享中间激活值。传统的分割学习依赖于跨客户端-服务器分割点的端到端反向传播。这会导致巨大的通信开销(即每次迭代都需要交换前向激活值和反向梯度)和显著的内存使用(用于存储激活值和梯度)。本文开发了一种用于分割学习的超越反向传播的训练方法。在该方法中,客户端和服务器半独立地训练其模型分区,使用局部损失信号而非传播的梯度。具体来说,客户端的网络在分割点处增加了一个小型辅助分类器以提供局部误差信号,而服务器则使用真实损失函数在客户端传输的激活值上进行训练。这种解耦消除了发送反向梯度的需要,从而将通信成本大致减半,并减少了内存开销(因为每一方仅为其自身的反向传播存储局部激活值)。我们在CIFAR-10和CIFAR-100上评估了我们的方法。实验展示了两个关键结果。首先,所提出的方法实现了与使用反向传播的标准分割学习相当的性能。其次,它显著减少了通信量(传输激活值/梯度)达50%,并将峰值内存使用量降低了高达58%。
引用
@article{arxiv.2601.19261,
title = {Decoupled Split Learning via Auxiliary Loss},
author = {Anower Zihad and Felix Owino and Ming Tang and Chao Huang},
journal= {arXiv preprint arXiv:2601.19261},
year = {2026}
}