Research map / Architectures and efficiency / Attention and transformers
Mixtures of Subspaces for Bandwidth Efficient Context Parallel Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a compression scheme for context-parallel attention that transmits low-rank key/value activations under a learned per-chunk rotation of a warm-started subspace, enabling decentralized long-context training over 300 Mbps links.
Problem
Context-parallel training of LLMs requires all-gathering keys and values across devices every layer and step, which is prohibitively communication-heavy over low-bandwidth decentralized internet links.
Approach
Factorize attention weights so K/V outputs lie in dynamic mixtures of low-dimensional subspaces: a fixed warm-started orthonormal basis is rotated by a per-chunk rotation predicted from the chunk's mean activation (parameterized by a single skew-symmetric generator with a small angle), so only compressed activations plus a scalar are broadcast; the auxiliary heads are removed after training.
Key results
- At 300Mbps with ~96.5% K/V compression, the 800M model reaches validation perplexity 17.06 on FineWeb vs 17.18 for centralized 100Gbps CP, with TPS 55K vs 2.7K for uncompressed decentralized CP.
- Ablations at 132K context show fixed subspace (26.57) and no warm start (26.63) degrade perplexity relative to the method (22.64) on FineWeb.
- Removing the second-order approximation leaves perplexity unchanged, and removing reparameterization drops TPS from 55K to 37K.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| compress activations with the proposed method instead of centralized CP | 800M Llama 3 · FineWeb validation | perplexity: 17.18 -> 17.06 | Helped |
| compress activations with the proposed method instead of centralized CP | 800M Llama 3 · C4 validation | perplexity: 17.51 -> 17.47 | Helped |
| compress activations with the proposed method instead of centralized CP | 800M Llama 3 · BookCorpus validation | perplexity: 17.88 -> 17.81 | Helped |
| use a fixed projection subspace instead of learned, adaptive subspaces (ablation) | Ours · FineWeb validation, 10K steps | perplexity: 22.64 -> 26.57 | Hurt |
| use random rotations instead of learned rotations (ablation) | Ours · FineWeb validation, 10K steps | perplexity: 22.64 -> 24.93 | Hurt |
| remove the second-order approximation (ablation) | Ours · FineWeb validation, 10K steps | perplexity: 22.64 -> 22.64 | No clear effect |
| add warm-start initialization of principal directions (ablation) | Ours · FineWeb validation, 10K steps | perplexity: 26.63 -> 22.64 | Helped |
| compress both K and V by 98% (ablation) | Ours · FineWeb validation, 10K steps | perplexity: 22.64 -> 24.74 | Hurt |
| compress K by 99% and V by 95% (ablation) | Ours · FineWeb validation, 10K steps | perplexity: 22.64 -> 24.68 | Hurt |
| use 300 warmup steps instead of 500 (ablation) | Ours · FineWeb validation | perplexity: 22.64 -> 22.66 | No clear effect |
| reparameterize U (ablation) | Ours · decentralized CP | TPS: 37K -> 55K | Helped |
| use the second-order approximation (ablation) | Ours · decentralized CP | TPS: 30K -> 55K | Helped |
Surprising findings
- Even a single rotation angle (k=1) per chunk suffices to match a fixed global rotation in bandwidth efficiency and stability.
- Uncompressed decentralized CP at 300Mbps would take over 150 days to train, so only throughput is reported for it.
Limitations
- Theoretical explanation for good minima under very low-dimensional rotation search is lacking.
- Alternative reparameterizations beyond subspace rotations are not explored.
- Experiments are limited to ~800M-parameter models (with a 3B scaling run) and up to 132K context.
Experimental setup
Datasets: FineWeb C4 BookCorpus
Benchmarks: FineWeb C4 BookCorpus
Metrics: validation perplexity tokens per second (TPS) wall-clock convergence memory per GPU
Baselines: Top-10% sparsification 4-bit quantization BigBird CosFormer Uncompressed centralized CP (100Gbps) Uncompressed decentralized CP (300Mbps)
Builds on: Ring Attention Power Gossip Llama 3 architecture Chinchilla scaling law
Closest papers by meaning
- Subspace Networks: Scaling Decentralized Training with Communication-Efficient Model Parallelism (NeurIPS 2025 · similarity 0.91)
- Towards a Better Theoretical Understanding of Independent Subnetwork Training (ICML 2024 · similarity 0.87)
- Compressed Context Memory for Online Language Model Interaction (ICLR 2024 · similarity 0.87)
- Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models (NeurIPS 2024 · similarity 0.87)
- CO2: Efficient Distributed Training with Full Communication-Computation Overlap (ICLR 2024 · similarity 0.86)
- Tensor-Parallelism with Partially Synchronized Activations (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.