1 Introduction↩︎

Training large-scale foundation models across geographically distributed, heterogeneous hardware introduces communication challenges that centralized distributed training systems [2] were not designed to handle [3]–[8]. Centralized systems assume tightly coupled accelerator clusters with high-bandwidth interconnects and large memory for each accelerator, but real-world low-bandwith deployments often rely on lower resource accelerators as their backbone and must operate over commodity wide-area networks with limited bandwidth and memory. Since, modern models have grown well beyond the memory capacity of these individual accelerators, model parallelism is a practical necessity. Pipeline parallelism [9] responds by partitioning parameters across devices, allowing GPUs to easily host a few layers each. However, this introduces a new bottleneck that every micro-batch triggers activation exchanges across stages in both the forward and backward passes inducing a significant communication cost.

This communication cost motivates low-bandwidth pipeline-parallel training, in which intermediate activations are aggressively compressed before transmission [5]. Beyond accessibility, compressing inter-stage activations is empirically well-motivated. There is growing evidence that foundation models converge to intrinsically low-rank solutions [5], [12]–[14], suggesting that the full-rank activation tensors exchanged between pipeline stages carry significant redundancy. Exploiting this structure allows us to eliminate communication overhead and transmit only the information that actually drives learning. Activation compression, however, is fundamentally harder than compression in data-parallel training [6], [15], where compressed gradients are computed independently across identical model replicas, compression errors can average out across workers, and distortions remain largely decoupled from the forward computation. In pipeline parallelism, stages hold complementary, non-overlapping subsets of the model, so the compressed activation produced by one stage becomes the direct input to the next; any transmission distortion therefore propagates through subsequent forward layers and backward gradients, accumulating across stages and potentially corrupting the learning signal even when the artifact is small.

Recent work on Subspace Networks (SSNs) [5] addresses this by constraining each activation of width \(d\) to a fixed, shared low-rank subspace of rank \(r\) defined by an orthogonal matrix \(U_r \in \mathbb{R}^{d \times r}\). Specifically, they transmit activations \(h \in \mathbb{R}^d\) as low-dimensional coordinates \(z = U_r^\top h\) and reconstruct them locally as \(\hat{h} = U_r z\). Notably [5] does not strictly constrain the activations to lie in the subspace thus the fixed projections cannot be removed at test time without performance degradation (See §7). In this sense SSNs represent an architectural modification. While SSNs confirm that architectural modification is a viable design space for communication-efficient training, the approach requires intrusive constrained optimization: model weights are restricted to a common low-rank subspace, a modified AdamW [11], [16] optimizer is needed to maintain a weight in the subspace, and a static embedding offset is applied at each stage (See §6). Consequently, we observe substantial accuracy degradation under a token-matched setting relative to uncompressed baselines.

Rather than prescribing a global basis for all layers, we let each pipeline stage learn its own communication subspace jointly with the task objective. Intermediate activations exhibit substantial low-rank structure even under modern optimizers [10], [17], suggesting that communication-efficient representations can naturally emerge end-to-end. However naively learning these projectors degrades the performance because standard gradient updates carry projectors off the Stiefel manifold (the set of all matrices with orthonormal columns), destroying the orthonormality on which isometric compression depends. Once a projector escapes the manifold, the model begins encoding features outside the intended subspace, leading to severe performance degradation. Crucially, we observe that these unconstrained models underperform even fixed orthogonal baselines, confirming that naive learning without manifold awareness is strictly worse than no learning at all (see §8). We identify this manifold escape as the principal failure mode and address it using manifold-constrained steepest-descent updates [18], which keep projectors strictly on the Stiefel manifold throughout training. Furthermore, removing the global subspace eliminates the need for SSN-style embedding decompositions. Instead, we introduce factorized anchor embeddings [19], a low-rank factorization \(E = E_p^{\mathrm{small}}P_p\) with fully trainable matrix \(E_p^{\mathrm{small}} \in \mathbb{R}^{V \times r}\) and a frozen matrix \(P_p \in \mathbb{R}^{r \times d}\), where \(V\) is the vocabulary size. This approach keeps the parameter count low while permitting the effective embedding to recover its full rank at each stage. Finally, we further reduce bandwidth via vector quantization (VQ) on the low-rank manifold, employing a lightweight dictionary-stream protocol that amortizes codebook synchronization across many activation exchanges.

To this end, we introduce Manifold Aware Projection Learning (MAPL), which treats inter-stage communication as a learnable geometric projection rather than a fixed architectural constraint. We make the following contributions. We show that enforcing a shared global subspace across pipeline stages degrades learning in token-matched settings and instead propose allowing each stage to learn its own compression subspace on the Stiefel manifold, which we find is better suited to low-bandwidth activation transmission (§ 3.2). Second, we introduce a low-rank embedding factorization that replaces the fixed token embedding decomposition used in SSN-style approaches with stage-specific learnable embeddings. Combined with vector quantization and a streaming codebook synchronization protocol, this design further reduces inter-stage bandwidth while requiring only integer token IDs to be transmitted between stages (§ 3.5). Third, we evaluate MAPL on models from 150M to 1B parameters, where it recovers downstream performance to within 1% of the uncompressed baseline in token-matched settings and outperforms SSNs by 5% across all evaluated scales (§ 4.2). As Fig. 1 shows, MAPL traces the Pareto frontier of cross-entropy versus compression, improving the achievable tradeoff over prior methods.

2 Related works↩︎

Low Bandwidth Pre-training. The growing utility of large AI systems has motivated efforts to democratize large-scale model training across decentralized bodies of participants connected over bandwidth-constrained networks such as the open internet. Early works on decentralized optimization established theoretical foundations: [20] showed that decentralized SGD can match centralized convergence rates, while subsequent analyses unified gossip-based optimization with compressed communication and local updates [21], [22]. To make distributed training communication-efficient, a rich line of work explored gradient compression through sparsification [23]–[25], scalar and sign-based quantization [26]–[29], low-rank approximation [30], and error-feedback mechanisms that correct biased compressors [31], [32]. Building on these foundations, volunteer-style collaborative training was demonstrated by DeDLOC [33] and CocktailSGD [34] for training over slow networks. Recent work has focused on data-parallel pre-training over the internet: DiLoCo-style methods [3], [4], [6], [35] adapt federated-averaging [36] to drastically reduce synchronization frequency. DeMo [37] exploits fast orthonormal transforms, top-k sparsification, and error feedback to enable multi-datacenter training. ACCO [38] overlaps the synchronization of delayed gradients with the computation of new gradients to increase GPU utilization. SparseLoCo [6] unifies error feedback, sparsification, and quantization to achieve internet-scale decentralized training [7]. However, all of these methods assume that each accelerator can hold a complete model replica. Pipeline-parallel training exhibits fundamentally different dynamics from DDP: SWARM [39] identified the square-cube law, showing that computation scales cubically while communication scales quadratically, and prior work explored decentralized model parallelism via mixture-of-experts routing [40] and reversible architectures that decouple forward and backward passes across stages [41]. Several works expanded on compressing activations [42], [43] and the change in activations [44]. Subspace Networks [5] exploited the empirical rank collapse of transformer residual streams to project boundary activations onto a fixed orthonormal basis. While effective, this imposes a single representational space across all transformer layers. In our work, we challenge this assumption to learn specialized, per-stage compression.

Efficient language models. A parallel line of work has investigated low-rank geometric structure for memory-efficient gradient compression and model compression, motivated by the empirical observation that deep networks converge to low-rank solutions [13], [14], [45]–[48]. GaLore [49] style methods [50], [51] projected stochastic gradients onto low-rank subspaces to reduce optimizer state memory, enabling full-parameter pre-training under tight memory budgets. Several works [12], [52], [53] further explored training low-rank factorized weights to reduce memory and communication requirements. LTE [54] used parallel LoRA-style adapters to pre-train models from scratch with reduced memory and bandwidth across parallel accelerators. Pufferfish [55] modified the architecture to obtain trained factorized layers that reduce communication costs, but focused on data parallelism. Our work similarly modifies the transformer architecture, but with the complementary goal of enabling efficient inter-stage activation communication for pipeline parallelism.

3 Manifold Aware Projection Learning↩︎

Figure 2 summarizes MAPL. At each of the inter-stage boundaries, We subtract the token embeddings and project the boundary activations to a low-dimensional subspace via a learnable orthogonal projector, transmitted, and reconstructed at the receiving stage. We derive this design bottom-up — beginning from an empirical observation about the structure of boundary activations (§3.1), which motivates the construction (§3.2) and its optimization procedure (§3.3). We validate each component ablatively in §3.4 and sketch a composable vector-quantization extension in §3.5.

3.1 Boundary activations are intrinsically low-rank↩︎

Setup Pipeline parallelism partitions the \(L\) layers of a decoder-only Transformer [57] into \(P\) contiguous stages \(\mathcal{S}_1, \dots, \mathcal{S}_P\), each residing on a dedicated device. Let \(E\) denote the embedding table and \(t_{\mathrm{ids}}\) the input token indices. At each inter-stage boundary \(p \in \{1, \dots, P{-}1\}\), the forward pass communicates the boundary activation \(X^{b_p} \in \mathbb{R}^{B \times T \times d}\), while the backward pass transmits its gradient; \(B\), \(T\), and \(d\) denote batch size, sequence length, and hidden dimension, respectively.

Motivating example We profile a 150M-parameter LLaMA model [1] trained for 3,000 steps on DCLM [56] with the Muon optimizer [10] at \(P{=}8\) pipeline stages. To isolate representational structure of the residual stream, we subtract the token embeddings and analyze the residual \(X_{res} = X^{b_p} - E[t_{\mathrm{ids}}]\), reshaped to \((B{\cdot}T) \times d\), via singular value decomposition (SVD). We quantify spectral concentration through the cumulative energy ratio: \[\mathcal{E}(r) = \frac{\sum_{i=1}^{r} \sigma_i^2}{\sum_{i=1}^{d} \sigma_i^2}.\] As shown in Figure 3, boundary activations exhibit pronounced low-rank structure across all pipeline stages: a rank of \(r \approx 250\) suffices to retain \({\geq}99\%\) of the total activation energy, despite the ambient dimension \(d{=}1024\). The inter-stage signal therefore concentrates near a submanifold of effective dimension roughly one quarter of the full representation space.

This finding extends [5], which observed rank collapse in projection matrices and responded by explicitly constraining weight rows to a shared low-rank subspace. Crucially, our setting requires no such contraint on weights: the low-rank structure emerges organically during training, without regularization or weight constraints. We therefore define a rank-\(r\) projector (\(r < d\)) as information-preserving for this signal, and argue that a learned projection can discover—rather than impose—the latent subspace the activations already inhabit, enabling principled compression of inter-stage communication.

3.2 Compression via per-stage learned orthogonal projectors↩︎

We introduce, at each pipeline-stage boundary \(p\), a learnable orthogonal projector \(A_p\) that lies on the Stiefel manifold \[\mathrm{St}(d, r) = \{A \in \mathbb{R}^{d \times r} : A^\top A = I_r\}.\] We equip each boundary with a per-stage anchor \(E_p^{\mathrm{small}}[t_{\mathit{ids}}]\) and a companion projector \(P_p\) that absorbs the high-rank, token-driven offset of the residual stream. Forward compression and reconstruction then proceed as \[Z^{b_p} = \bigl(X^{b_p} - E_p^{\mathrm{small}}[t_{\mathit{ids}}] P_p \bigr)\, A_p \;\in\; \mathbb{R}^{B \times T \times r}, \label{eq:project}\tag{1}\] \[\hat{X}^{b_p} = Z^{b_p}\, A_p^\top + E_{p+1}^{\mathrm{small}}[t_{\mathit{ids}}] P_{p+1} \;\in\; \mathbb{R}^{B \times T \times d}. \label{eq:reconstruct}\tag{2}\] By enforcing the Stiefel constraint, we guarantee that \(A_p^\top\) is the exact inverse of \(A_p\). Projection and reconstruction are therefore isometric on the column space of \(A_p\) and the method achieves an exact compression ratio of \(r/d\). In the backward pass we symmetrically route the gradient \(\partial \mathcal{L}/\partial X^{b_p}\) through the same orthogonal projector \(A_p\).

The rank-\(r\) projector \(A_p\) captures the dominant low-rank component of the boundary activation \(X^{b_p}\) while deliberately leaving the token-frequency offset unaddressed; this offset is inherently high-rank and would otherwise exhaust projector capacity. In SSN [5], the embedding table is decoupled into static high rank and learnable low rank components. Each stage added a static high rank offset as the learnable token embeddings forced to be in the same subspace as the weights [5]. We instead add a learnable offset. To reduce the parameter pressure at each stage, we factorize the offset as \[E_p^{\mathrm{small}}[\mathit{ids}]\, P_p, \qquad E_p^{\mathrm{small}} \in \mathbb{R}^{V \times r}, \qquad P_p \in \mathbb{R}^{r \times d}, \label{eq:anchor}\tag{3}\] where \(E_p^{\mathrm{small}}\) is a trainable embedding table and \(P_p\) is a fixed random orthonormal matrix. Only the integer token IDs cross the communication channel for this reconstruction. Consequently, the anchor adds negligible bandwidth overhead.

3.3 Optimization via SPEL↩︎

We update each \(A_p\) jointly with the model weights via SPEL (Spectral Steepest Descent on the Stiefel Manifold) [18], using the task loss as the only signal. Given the Euclidean gradient \(g_t = \partial \mathcal{L} / \partial A_p\) at step \(t\) and learning rate \(\alpha\), update consists of: \[\begin{align} g_t^R &= g_t - A_p\, \mathrm{sym}(A_p^\top g_t), && \text{(tangent projection)} \tag{4} \\ m_t &= \beta\, m_{t-1} + (1 - \beta)\, g_t^R, && \text{(heavy-ball momentum)} \tag{5} \\ d_t &= \mathrm{PolarExpress}(m_t), && \text{(spectral LMO direction)} \tag{6} \\ A_p &\gets \mathrm{PolarExpress}\bigl(A_p - \alpha\, d_t\bigr), && \text{(retraction onto \mathrm{St}(d,r))} \tag{7} \end{align}\] where \(\mathrm{sym}(M) = \tfrac{1}{2}(M + M^\top)\). Polar express [58] calculates the spectral norm LMO direction which is a similar parameter update as Muon [10] and retracts the matrix back to the stiefel manifold. SPEL inherits the \(\mathcal{O}(1/\sqrt{T})\) convergence rate of first-order methods on smooth manifolds [18]. We send the updated \(A_p\) to the next stage after every optimizer step. This is a minimal communication cost when compared with activationa and gradient communication.

Unlike SSN’s Grassmann update that refresh every \({\sim}500\) steps [5], SPEL updates \(A_p\) at every step, letting the subspace track the evolving activation geometry. In practice we choose a learning rate an order smaller than parameter updates (\(\times0.1\)). Algorithm 4 summarizes the full per-boundary update.

3.4 Empirical Validation↩︎

We validate three properties of our method (§3.2) and show learning the projection is naturally more favorable compression than having a fixed projector. We show that per-boundary projectors discover geometrically distinct subspaces for each different stage. The learned \(A_p\) captures substantially more activation energy than a fixed orthonormal basis of equal rank, and the learned projection faithfully preserves pairwise token similarity after compression.
Each stage learns a distinct subspace. Figure 5 (a) reports pairwise mean principal angles between the seven learned projectors of all boundaries, computed as \(\bar{\theta} = \frac{1}{r}\sum_{k=1}^{r}\arccos\bigl(\sigma_k(A_i^\top A_j)\bigr)\). Off-diagonal angles range from \(53^\circ\) to \(72^\circ\), with adjacent stages exhibiting the greatest overlap and distant stages approaching near-orthogonality. This change in subspace alignment is consistent with the residual stream transitioning from lexical to task-specific representations across depth [59].
The learned projector captures substantially more energy. Figure 5 (b) compares the residual energy retained under projection, \(\|X_{\mathrm{res}} A_p\|_F^2 / \|X_{\mathrm{res}}\|_F^2\), between a fixed random orthonormal basis and the jointly trained \(A_p\) at matching rank \(r{=}128\). Note that \(\|A\|_F^2 = \sum_{i} \sigma_i^2\) . The fixed basis plateaus near \(36\%\) of captured energy while the learned projector reaches \({\sim}80\%\) within \(1{,}500\) steps and continues to improve—a factor-of-two gain in effective rank utilization. Inspection of weight stable rank confirms that the preceding attention out projection and MLP projection layers
adapt to this space defined by \((A_p)\) (Appendix 9). We see that the compression actively induces the low rank structure.
The learned projection preserves pairwise token geometry. Figure 5 (c) plots the cosine similarity between projected token pairs, \(\cos(x_i A_p,\, x_j A_p)\), against their full-dimensional counterpart \(\cos(x_i, x_j)\). The Pearson correlation of \(r = 0.992\) with the identity line demonstrates that \(A_p\) behaves as a near-isometry over the empirical token distribution and the relative angular structure in the activation space is preserved under compression. This property is not enforced explicitly and it emerges from Stiefel optimization alone and provides geometric evidence that the transmitted representation retains the relational information downstream stages require for computation.

3.5 Composing with vector quantization↩︎

To minimize communication overhead during pipeline-parallel training, we employ Multi-Codebook Vector Quantization (MCVQ) for compressing both forward activations and backward gradients. The method operates on a low-rank projected representation, decomposed into \(G\) groups, and applies \(R\) rounds of residual quantization to achieve high compression ratios while preserving representational fidelity.
Quantization Scheme Formally, the projected representation is compressed via residual vector quantization using a per-stage codebook. \(Z^{b_p} \in \mathbb{R}^{B \times T \times r}\) is compressed using a per-stage codebook \(\mathcal{C}_p \in \mathbb{R}^{r \times K}\). Each stage quantizes the residual error from the previous round, enabling progressive refinement of the compressed representation across \(R\) rounds.
Codebook Synchronization Sender and receiver nodes synchronize \(\mathcal{C}_p\) through a streaming dictionary update protocol: the codebook is partitioned into random subsets, with a \(\frac{1}{K}\) fraction of codes transmitted per micro-batch. This design is motivated by the empirical observation of [60] that VQ codebooks evolve slowly over the course of training, ensuring that the staleness introduced by incremental updates remains negligible with respect to overall bandwidth and convergence.

4 Experiments↩︎

4.1 Experimental Setup↩︎

Architecture and scales. We evaluate our method on decoder-only transformers based on the LLaMA [1] architecture across three parameter scales: 150M, 500M, and 1B. All models employ a context length of 2048 tokens and are trained with the LLaMA tokenizer (vocabulary size 32,000) in bf16 precision. All models are pre-trained on the DCLM corpus [56], with 5M tokens held out uniformly at random as a validation set. Following the Chinchilla compute-optimal scaling regime [61], each model is trained on a token budget of 20 tokens per parameter. Implementation details are provided in Appendix 10. All the methods are compared at equal token budgets.

Optimization We adopt a hybrid optimizer configuration in which the 2D hidden weight matrices are updated using Muon [10], while embeddings, biases, and the output projection are updated with AdamW [16]. The Muon learning rate is set to \(\eta_\mu = 0.02\) for the 150M, 500M sizes and \(\eta_\mu = 0.01\) for the 1B case, and the AdamW learning rate is coupled to it via a multiplier of \(0.5\), yielding \(\eta_{\text{adam}} = 0.5 \cdot \eta_\mu\). Training uses a global batch size of 512. We choose the multiplier of \(0.1\) for learning rate we use to update the projector using § 3.3. Our primary experiments are conducted under pipeline-parallel configurations with \(P \in \{4, 8\}\) stages.

Baselines. We compare against methods while holding all shared hyperparameters fixed to ensure a fair comparison. Uncompressed transmits activations and gradients across pipeline boundaries at the full model width and therefore serves as an upper bound on attainable quality. SSN [5] projects inter-stage tensors onto a learned low-rank subspace(§ 6). SSN (AdamW version) use SSN with AdamW optimizer for all the parameters. MAPL applies our proposed compression scheme to inter-stage activations and gradients. Finally, MAPL+ VQ augments our approach with vector quantization of the projected representation, effectively doubling the compression ratio.

Evaluation We report cross-entropy validation loss on the held-out DCLM split as our primary metric, together with the relative degradation \(\Delta\%\) with respect to the Uncompressed baseline at the same scale. We additionally evaluate zero-shot downstream accuracy on HellaSwag [62], PIQA [63], ARC-Easy, and ARC-Challenge [64] using lm-evaluation-harness [65]. We also report average accuracy across all the tasks.

4.2 Main Results↩︎

MAPL closes the gap to the uncompressed baseline across all scales, Table 1 shows that, at \(4\times\) compression on the 150M model, MAPL achieves validation losses of \(3.156\) (\(P{=}4\)) and \(3.165\) (\(P{=}8\)). This is just \(0.84\%\) and \(1.11\%\) above the uncompressed reference of \(3.13\). At this scale, we can see that SSN degrades much higher in performance in both pipeline cases of pipeline split. Scaling to 500M parameters under \(6\times\) compression, MAPL tracks the baseline within \(1.90\%\) (\(P{=}4\)) and \(0.08\%\) (\(P{=}8\)); at 1B under \(8\times\) compression, it holds losses of \(2.72\) and \(2.73\), degrading by only \(1.38\%\) and \(2.02\%\) respectively. At all scales the SSN performs much worse and looses a maximum of nearly 14% performance compared to uncompressed baseline.

We attribute the stability of MAPL across model scales to the per-stage Stiefel-constrained projectors, which faithfully capture the intrinsic low-rank geometry of the activations identified in §3.1. Composing MAPL with residual vector quantization (§3.5) doubles the compression ratio while incurring only marginal additional loss. At 1B parameters, MAPL+VQ achieves \(16\times\) inter-stage compression at \(3.01\%\) (\(P{=}4\)) and \(2.30\%\) (\(P{=}8\)) degradation. We trace this behavior to the empirical near-isometry of the learned projector (Figure 5c).By preserving pairwise similarity on the low-rank manifold, the codebook in \(\mathbb{R}^r\) inherits a well-conditioned distribution that aligns itself naturally for vector quantization.

MAPL preserves downstream performance on various tasks. Validation-loss improvements translate directly to zero-shot performance as shown in Table 2. Across the tested downstream tasks, MAPL closely tracks the uncompressed baseline at every scale. At 500M parameters, the average accuracy gap is \(0.2\) points at \(P{=}4\) (\(41.8\) vs.\(42.0\)) and \(0.4\) points at \(P{=}8\) (\(41.6\) vs.\(42.0\)); at 1B, the corresponding gaps are \(0.8\) and \(1.5\) average points. SSN baselines, by contrast, suffer accuracy drops as large as \(8.8\) points at 1B. This shows that the learning capacity is restricted due to the global-subspace weight constraint and shows degraded performance at token matched comparisons. The MAPL+VQ variant incurs a more pronounced performance loss in downstream tasks.

5 Conclusion↩︎

We introduced MAPL, a communication-efficient pipeline-parallel training method that learns stage-specific orthogonal compression subspaces directly on the Stiefel manifold using manifold-constrained optimization. By combining adaptive low-rank projectors with lightweight factorized anchor embeddings and optional residual vector quantization, MAPL enables aggressive inter-stage activation compression while preserving training quality. Across LLaMA-style models ranging from 150M to 1B parameters and pipeline configurations with 4 and 8 stages, MAPL achieves \(4\times\)–\(8\times\) communication compression within approximately \(1\)–\(2\%\) of the uncompressed validation loss. It extends to \(16\times\) compression with only modest additional degradation when combined with vector quantization. The method consistently outperforms recent work subspace networks [5] while maintaining strong downstream zero-shot accuracy. Our results suggest that inter-stage transformer activations naturally admit adaptive low-dimensional structure that can be learned jointly with the task objective, rather than imposed through a fixed global subspace. More broadly, this work highlights manifold-constrained representation learning as a practical mechanism for reducing communication costs in distributed foundation-model training.

Limitations. Our experiments are limited to models up to 1B parameters, and further evaluation at larger scales and under real heterogeneous network conditions remains necessary. In addition, although MAPL substantially reduces the degradation associated with activation compression, performance still declines under extreme compression ratios, particularly when vector quantization is applied.

6 Background: Subspace networks↩︎

SSN [5] exploits a structural property of trained Transformers: the output projection weights \(W_2^\ell \in \mathbb{R}^{d_{\mathrm{ff}} \times d}\) and attention projection weights \(W_1^\ell \in \mathbb{R}^{d\times d}\) undergo rank collapse during training, with their effective row space converging to a low-dimensional subspace \(\mathcal{S} \subset \mathbb{R}^d\) of dimension \(k \ll d\). This collapse arises naturally from AdamW’s adaptive scaling, which attenuates gradient components aligned with negligible singular directions.

SSN operationalizes this observation via five steps. (1) A fixed orthonormal basis \(U_k \in \mathbb{R}^{d \times k}\) is constructed by QR decomposition of a random matrix. (2) The rows of each \(W_2^\ell\) and \(W_1^\ell\) are constrained to \(\mathrm{span}(U_k)\) by replacing the per-parameter adaptive scaling in AdamW with a row-wise constant variant. (3) Deterministic components—positional embeddings \(\mathrm{PE}\) and fixed token embeddings \(T_{\mathrm{fixed}}\)—are subtracted from \(X^{b_p}\) before compression. (4) The stage transmits the compressed activation \[Z^{b_p} = \bigl(X^{b_p} - \mathrm{PE} - T_{\mathrm{fixed}}\bigr)\, U_k \;\in\; \mathbb{R}^{B \times T \times k},\] achieving a compression ratio of \(k/d\). (5) The receiving stage reconstructs \(\hat{X}^{b_p} = Z^{b_p} U_k^\top + \mathrm{PE} + T_{\mathrm{fixed}}\).

Because weight confinement ensures \(X^{b_p} - \mathrm{PE} - T_{\mathrm{fixed}} \in \mathrm{span}(U_k)\) exactly, SSN’s compression is architecturally lossless on its modified model, though it incurs loss relative to an unconstrained baseline. The basis \(U_k\) is refined via Grassmann manifold updates every \({\sim}500\) steps.

Limitations. SSN’s losslessness is contingent on strict weight confinement. The subspace is fixed at initialization and only coarsely adapted; the model is forced to conform to a predetermined geometry rather than discovering a task-optimal one. In §3.1, we showed that pipeline boundary activations exhibit low-rank structure intrinsically—even when trained with a modern optimizer [10] without any architectural constraint—motivating a jointly learned subspace that adapts to the activation geometry rather than imposing it.

7 Evaluating SSN without Subspace Projection↩︎

We ran probe evaluations on 16 sequences of 2048 tokens drawn from the DCLM validation set [56]. In this ablation, we replaced the subspace projection matrix \(U_k\) with an identity matrix at inference time and replaced the decoupled low rank embedding table with its projections \(TE \mathbin{@} U_k^\top\).

As a reference, the uniform-distribution baseline (i.e., random guessing over the vocabulary) yields a loss of \(\approx 10.38\). Any value above this threshold indicates performance worse than chance.

The learned subspace projection matrices \(U_k\) do more than compress information—they form an integral part of the model’s learned representations. When we replace them with the identity matrix at inference time (while keeping all weights trained with active projection), the model suffers severe degradation across every tested configuration. Losses increase by between \(+3.44\) and \(+8.25\), and several \(P=8\) runs exceed the uniform baseline entirely. See Table 4.

This effect holds consistently across model scales (150M–1B), optimizers (Muon [10] and AdamW [11]), and degrees of parallelism. These results show that the network internalizes representations tightly coupled to the projected subspaces. Removing the projection does not simply degrade the model but breaks it.

6pt

8 Keeping Projectors in the Stiefel Manifold↩︎

In this section, we demonstrate the necessity of constraining subspace projectors to the Stiefel manifold during training. If the subspace projectors are trained naively, standard optimizer updates cause them to deviate from the manifold, thereby violating the orthonormality constraint. Consequently, the projection loses its isometric properties, allowing the model to learn representations outside the intended subspace.

To empirically validate this, we train a 150M parameter model under three distinct projection configurations. In the baseline setup, we initialize an orthogonal projector and freeze it for the duration of training. In the second configuration, the projector is updated using standard Euclidean gradients without any manifold constraints. As anticipated, these unconstrained updates cause the weight matrix to escape the manifold and lose its orthogonality, which actively degrades performance—yielding worse results than simply utilizing a fixed projection. Finally, enforcing the manifold constraint via SPEL [18] updates resolves this issue and achieves the lowest validation loss. These results are summarized in Table 5.

9 Rank Collapse in Weight Matrices↩︎

Figure 6 illustrates the induced rank collapse observed when co-training the projection matrix on the Stiefel manifold. Learning the projection encourages the model to develop a low-rank structure that aligns closely with the learned basis. Consequently, the model exhibits a significantly more pronounced rank collapse compared to a fixed orthogonal projection. This structural adaptation accounts for the higher percentage of residual energy captured within the subspace.

10 Implementation Details↩︎

In this section, we provide a comprehensive overview of the models, datasets, and hyperparameters used across all experiments. Our goal is to ensure full reproducibility of the pipeline-parallel compression results presented in the main text. We ran our experiments of nodes consisting of 4\(\times\) Nvidia H100 GPUs with precision bf16, we used Pytorch [66] framework with huggingface transformers [67] implementation of Llama 2 [57], with RoPE[68] positional embeddings. We used Python 3.10.

10.1 Model Architectures↩︎

All experiments evaluate decoder-only transformers based on the LLaMA 2 architecture [57]. We define three scales of models: 150M, 500M, and 1B parameters. To ensure stable comparisons across scales, all models use a context sequence length of 2048 tokens and process inputs using a vocabulary size of 32,000. The structural hyper-parameters for each model scale are as follows:

• 150M: \(d_{\mathrm{model}} = 1024\), \(n_{\mathrm{layers}} = 9\), \(n_{\mathrm{heads}} = 8\), \(d_{\mathrm{ffn}} = 2688\)

• 500M: \(d_{\mathrm{model}} = 1536\), \(n_{\mathrm{layers}} = 18\), \(n_{\mathrm{heads}} = 12\), \(d_{\mathrm{ffn}} = 4096\)

• 1B: \(d_{\mathrm{model}} = 2048\), \(n_{\mathrm{layers}} = 20\), \(n_{\mathrm{heads}} = 16\), \(d_{\mathrm{ffn}} = 5440\)

10.2 Dataset and token budget↩︎

Models were pre-trained on the DCLM dataset [56]. For all runs, we adhered to the Chinchilla compute-optimal scaling regime [61], defining our total training steps such that the models process 20 tokens per parameter. For example, the 150M model was trained on roughly 3 billion tokens, the 500M model on 10 billion tokens, and the 1B model on 20 billion tokens. A uniform 5M token split from DCLM was held out entirely for validation.

10.3 Optimizer Configuration↩︎

We employed a hybrid optimizer approach. The 2D hidden weight matrices (attention projections and MLP matrices) were updated using the Muon optimizer [10], using PolarExpress [58] algorithm for orthogonalizing the updates. The 1D parameters, including biases, layer normalizations, token embeddings, and the final language modeling head, were updated using AdamW [16].

Hyperparameter Sharing Strategy: To ensure a fair comparison between the baseline (uncompressed) models and our proposed methods and baselines (MAPL, SSN, etc.), we used the exact same optimizer learning rates and weight decay as the uncompressed baselines across all compressed experiments, and we did not tune them specifically for the compression techniques. For the main model parameters, we set the Muon learning rate to \(\eta_\mu = 0.02\) and the AdamW learning rate to \(\eta_{\mathrm{adam}} = 0.01\) (a \(0.5\times\) multiplier on the Muon rate). For the 1B model experiments , we used \(\eta_\mu = 0.01\) and \(\eta_{\mathrm{adam}} = 0.005\). We use a global batch size of 512 for all runs.

10.4 MAPL Compression Hyperparameters↩︎

Our method relies on learning an orthogonal projector \(A_p \in \mathrm{St}(d, r)\) per pipeline stage boundary. We optimized these Stiefel manifold projectors using the manifold constrained steepest descent algorithm [18] via a Polar Express retraction [58]. We found empirically that the projector optimization benefits from a slightly reduced learning rate compared to the main model weights; therefore, we set the projector learning rate to be \(0.1\times\) the main Muon learning rate.

For the Vector Quantized (MAPL + VQ) experiments, we used a codebook size of \(K=256\) and applied \(R=2\) residual quantization rounds with a group dimension of \(G=2\), utilizing an asynchronous streaming dictionary of steaming interval \(5\), to amortize codebook synchronization costs across micro-batches.

11 Amortizing VQ: Streaming Codebook Synchronization↩︎

For our vector-quantized variant (MAPL + VQ), both the sender and receiver stages must maintain identical codebooks. Transmitting the full codebook \(\mathcal{C}_p \in \mathbb{R}^{r \times K}\) across the network at every micro-batch would severely erode the bandwidth savings achieved by quantization.

Motivated by the observation that VQ codebooks evolve slowly over training [60], we introduced an asynchronous streaming dictionary update protocol (§3.5). Instead of full synchronization, the sender transmits only a small, randomly sampled subset of the codebook (e.g., \(1/K\) of the codes) alongside the quantized activations in each micro-batch.

To validate that this partial synchronization does not degrade representational fidelity, we conducted an ablation comparing full codebook synchronization against our streaming protocol on a 150M model (\(P=4\), \(8\times\) compression). The model trained with full synchronization achieved a validation loss of \(3.1647\), while the model trained with the streaming dictionary protocol achieved an essentially identical validation loss of \(3.1651\). This result confirms that continuous, fractional codebook updates are sufficient to maintain alignment between pipeline stages without incurring the communication penalty of full synchronization.

12 Ablation of Factorized Anchors↩︎

In Section 3.2, we hypothesized that the token-driven offset in the residual stream is inherently high-rank and consumes projector capacity if not explicitly modeled. To validate this, we ablate the anchor formulation on a 150M model across an 8-stage pipeline (\(P=8\), \(4\times\) compression), a depth where compression errors severely compound.

As shown in Table 6, removing the anchor entirely degrades the validation loss to \(3.209\) (a \(+2.39\%\) degradation relative to the uncompressed baseline). Adding a static, frozen embedding projection (analogous to the static offset in SSN) fails to improve performance (\(3.212\)). Conversely, utilizing a fully trainable per-stage anchor recovers performance to \(3.149\), almost matching the uncompressed baseline (\(3.134\)), but incurs an unacceptable parameter overhead (an additional full embedding table per stage).

Our proposed factorized anchor strikes the optimal balance: by learning a low-rank trainable component (\(E_p^{\mathrm{small}}\)) against a frozen random projection (\(P_p\)), it recovers the majority of the performance (\(3.165\), only \(+1.11\%\) degradation) while adding negligible parameter overhead to the pipeline stages.

13 Summary of Hyperparameters↩︎

4pt

References↩︎