arXiv:2606.05484cs.LG2026-06中稿 · the 2nd Workshop o…

让流水线各阶段自适应学习压缩子空间,大幅降低通信开销且几乎不损失性能。

Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

论文配图:Learned Subspace Compression for Communication-Efficient Pipeline Parallelism
图 1 · 摘自论文原文
  • 在斯特夫尔流形约束下学习各阶段的可训练正交投影
  • 150M~1B参数模型上实现高压缩比,性能损失极小
  • 适合大规模语言模型训练中的高效分布式部署

流水线并行使大语言模型训练突破单设备内存限制,但在低带宽网络下跨阶段激活通信成为主要瓶颈。现有方法使用固定正交投影压缩激活值,仍导致显著性能下降,并需额外非标准优化调整。本文提出曼达帕意识投影学习(MAPL),将跨阶段压缩建模为显式施蒂费尔流形(正交矩阵)约束下的可学习正交投影。与固定全局子空间不同,MAPL让每个阶段通过流形约束最速下降法自主发现并动态适配任务最优压缩子空间。为恢复阶段边界处的令牌级信号,引入分块锚点嵌入,实现全秩激活重建且通信开销可忽略。此外,可在投影后结合残差向量量化,配合流式字典同步协议,摊销字典通信成本。在150M至1B参数的LLaMA模型上验证,MAPL可轻松集成至现有流水线架构,在显著提升性能-压缩权衡的同时,实现高压缩比且性能损失极小。

原文摘要 · Abstract (English)

Pipeline parallelism enables training of large language models that exceed single-device memory, yet inter-stage activation communication becomes the dominant bottleneck when trained on low-bandwidth networks. Recent work in this area has proposed using fixed orthogonal projections to compress activations. However, this still results in a significant performance degradation and requires a number of non-standard adaptations to constrain the optimization. A natural alternative is to learn a low rank projection for each pipeline stage, however maintaining the necessary orthogonality of these projectors during training remains a challenge. We present Manifold Aware Projection Learning (MAPL), a method that treats inter-stage compression as a learnable orthogonal projection under explicit Stiefel manifold (orthogonal matrices) constraints. Rather than prescribing a fixed global subspace, MAPL lets each pipeline stage discover and continuously adapt its own task-optimal compression subspace via manifold-constrained steepest descent. To recover token-specific signals at stage boundaries, we introduce per-stage factorized anchor embeddings that allow for full-rank activation reconstruction with negligible communication overhead. We further show that we can incorporate residual vector quantization after projection with a streaming codebook synchronization protocol that amortizes dictionary communication. Across LLaMA models from 150M to 1B parameters we show that MAPL can be easily applied to the existing pipeline and can achieve high compression with neglibile performance degradation with a drastically improved tradeoffs in performance vs. compression compared to Subspace Networks.

流水线并行通信压缩子空间学习大模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。