Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

arXiv cs.LG Papers
compression pipeline-parallelism machine-learning llm-training orthogonal-projection stiefel-manifold

Summary

This paper introduces MAPL, a method for learned orthogonal compression of activations in pipeline parallelism, reducing communication overhead while maintaining performance via Stiefel manifold constraints and per-stage factorized anchor embeddings.

arXiv:2606.05484v1 Announce Type: new Abstract: 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.
Original Article
View Cached Full Text

Cached at: 06/05/26, 08:11 AM 

# Learned Subspace Compression for Communication-Efficient Pipeline Parallelism
Source: [https://arxiv.org/abs/2606.05484](https://arxiv.org/abs/2606.05484)
[View PDF](https://arxiv.org/pdf/2606.05484)

> Abstract: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\.

## Submission history

From: Paul Janson \[[view email](https://arxiv.org/show-email/b9761b3f/2606.05484)\] **\[v1\]**Wed, 3 Jun 2026 22:10:16 UTC \(203 KB\)

Similar Articles

Prompt Compression via Activation Aggregation

arXiv cs.CL

This paper proposes compressing instruction prompts into a single activation vector via learned weighted sums of intermediate layer activations, achieving under 2% accuracy drop and revealing insights into LLM activation space structure.

SHiPPO: Recurrent Memory with Transported Polynomial Projections

arXiv cs.LG

SHiPPO extends HiPPO by transporting polynomial projection coefficients into a moving channel frame, enabling selective state-space models to recover order-sensitive memory signals. The paper provides theoretical foundations and diagnostics supporting its transported-memory prior.