Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

Technische Samenvatting: Learned Subspace Compression voor Communicatie-efficiënte Pipeline Parallelisme

Probleemstelling
Het trainen van grote taalmodellen (LLM's) over geografisch verspreide, heterogene hardware vormt aanzienlijke communicatie-bottlenecks, met name in pipeline parallelisme. Hoewel pipeline parallelisme modelparameters over apparaten verdeelt om geheugenlimieten te overwinnen, vereist het uitwisseling van intermediaire activaties tussen stadia tijdens zowel de forward als de backward pass. Op netwerken met een lage bandbreedte wordt deze inter-stage communicatie de dominante bottleneck.

Bestaande benaderingen, zoals Subspace Networks (SSNs), proberen dit te verzachten door activaties te projecteren op een vaste, gedeelde laag-rang subspace gedefinieerd door een orthogonale matrix. Echter, deze methode vereist ingrijpende architecturale aanpassingen, inclusief geconstreerde optimalisatie om gewichten binnen een gedeelde subspace te houden en het gebruik van aangepaste optimizers. Empirisch gezien leiden deze vaste globale subspaces tot significante prestatievermindering, vooral naarmate de compressieratio's toenemen, omdat ze er niet in slagen zich aan te passen aan de evoluerende geometrie van activaties tijdens de training. Bovendien zorgt het naïef leren van projectiematrices zonder restricties ervoor dat ze de Stiefel-variëteit (de verzameling orthogonale matrices) verlaten, wat de isometrische eigenschappen die nodig zijn voor effectieve compressie vernietigt en leidt tot ernstige prestatiedalingen.

Methodologie: Manifold Aware Projection Learning (MAPL)
De auteurs stellen Manifold Aware Projection Learning (MAPL) voor, een methode die inter-stage compressie behandelt als een leerbare geometrische projectie in plaats van een vaste architecturale beperking. De kerncomponenten van MAPL zijn:

1. Per-Stage Leerbare Orthogonale Projectoren: In tegen tegenstelling tot SSNs die een enkele globale subspace afdwingen, rust MAPL elk pipeline-stadium grensvlak uit met zijn eigen leerbare orthogonale projector $A_p \in \mathbb{R}^{d \times r}$. Deze projectoren zijn geconstreerd aan de Stiefel-variëteit ($St(d, r)$), wat garandeert dat $A_p^\top A_p = I_r$. Dit stelt elk stadium in staat om zijn eigen taak-optimale compressie-subspace te ontdekken en continu aan te passen.
2. Manifold-Constrained Optimization (SPEL): Om te voorkomen dat projectoren de variëteit verlaten tijdens standaard gradiënt-updates, gebruiken de auteurs Spectral Steepest Descent on the Stiefel Manifold (SPEL). Dit algoritme werkt projectoren bij via tangentiële projecties en retractie-operaties (via de Polar Express methode) bij elke optimalisatiestap, waardoor de orthogonaliteitsrestrictie strikt wordt gehandhaafd.
3. Gefactoriseerde Anchor Embeddings: Om hoog-rang, token-gestuurde offsets in residual streams te verwerken zonder de capaciteit van de projector uit te putten, introduceert MAPL gefactoriseerde anchor embeddings. De offset wordt gedecomposeerd in een leerbare laag-rang matrix $E_{small}^p$ en een bevroren orthogonale matrix $P_p$. Alleen integer token ID's worden verzonden om deze offsets aan de ontvangende zijde te reconstrueren, wat een verwaarloosbare communicatie-overhead toevoegt.
4. Vector Quantization (VQ): Voor verdere compressie ondersteunt de methode residual vector quantization (RVQ) op de laag-rang gepjecteerde representaties. Er wordt een streaming codebook synchronisatieprotocol gebruikt om de communicatie van de dictionary te amorteren, waardoor hoge compressieratio's (tot 16×) mogelijk zijn met minimale bandbreedtekosten.

Belangrijkste Bijdragen

• Per-Stage Subspace Learning: Het artikel toont aan dat het afdwingen van een gedeelde globale subspace het leren in token-gematchte settings verslechtert. In plaats daarvan staat MAPL elk stadium toe om een afzonderlijke, door de Stiefel-variëteit geconstreerde compressie-subspace te leren, die beter aansluit bij de intrinsieke laag-rang structuur van activaties op verschillende dieptes.
• Manifold-Constrained Optimization: De auteurs identificeren "manifold escape" als de primaire faalmodus van naïef projector-leren. Ze introduceren SPEL om orthogonaliteit te handhaven, waarbij ze aantonen dat onbeperkt leren strikt slechter is dan het gebruik van vaste baselines.
• Efficiënte Reconstructie: Door leerbare projectoren te combineren met gefactoriseerde anchor embeddings, bereikt MAPL volledige reconstructie van activaties met een verwaarloosbare communicatie-overhead, waardoor de noodzaak voor de ingrijpende architecturale decomposities van SSNs vervalt.
• Verbetering van de Pareto-front: De methode brengt de Pareto-front in kaart van validatie cross-entropy versus communicatiekosten, waarbij het een hoge compressie bereikt met een verwaarloosbare prestatievermindering vergeleken met eerdere methoden.

Experimentele Resultaten
De auteurs evalueerden MAPL op decoder-only LLaMA-modellen variërend van 150M tot 1B parameters, getraind op de DCLM corpus met pipeline configuraties van $P=4$ en $P=8$ stadia.

• Prestaties vs. Compressie: MAPL presteert consequent beter dan SSN-baselines. Voor een 150M model bij 4× compressie bereikte MAPL een validatieverlies binnen 0,84% van de ongecomprimeerde baseline, terwijl SSN-baselines met 8–12% degradeerden.
• Schaalbaarheid: Bij de schaal van 1B parameters met 8× compressie behield MAPL een verliesdegradatie van slechts 1,38–2,02%, terwijl SSN-baselines degradaties tot wel 26% leden.
• Downstream Taken: In zero-shot evaluaties (HellaSwag, PIQA, ARC) volgde MAPL de ongecomprimeerde baselines nauwgezet (binnen ~1 punt gemiddelde accuratesse bij 500M/1B schalen), terwijl SSN-baselines degradaties van tot wel 9 punten leden.
• Extreme Compressie: Wanneer gecombineerd met vector quantization (MAPL+VQ), bereikte de methode 16× compressie met een bescheiden extra degradatie (bijv. ~2,3% toename in verlies bij 1B parameters), hoewel de accuratesse op downstream taken een meer uitgesproken daling vertoonde vergeleken met niet-gekwantiseerde MAPL.

Betekenis en Claims
Het artikel claimt dat inter-stage transformer activaties van nature adaptieve, laag-dimensionale structuren toelaten die gezamenlijk met de taak-objective geleerd kunnen worden, in plaats van opgelegd te worden via vaste globale subspaces. Door compressie te behanden als een leerbare geometrische projectie op de Stiefel-variëteit, maakt MAPL agressieve communicatiereductie mogelijk (4× tot 16×) terwijl de kwaliteit van de training en de downstream prestaties behouden blijven. Het werk benadrukt manifold-constrained representatieleer als een praktisch mechanisme om communicatiekosten te verlagen, en biedt een superieur evenwicht tussen prestaties en compressie vergeleken met bestaande subspace-gebaseerde benaderingen.

Beperkingen
De auteurs merken op dat de experimenten momenteel beperkt zijn tot modellen tot 1B parameters, en dat verdere evaluatie op grotere schaal en onder real-world heterogene netwerkomstandigheden vereist is. Daarnaast, hoewel MAPL de degradatie aanzienlijk vermindert, neemt de prestatie nog steeds af onder extreme compressieratio's, met name wanneer vector quantization wordt toegepast.