Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

기술 요약: 통신 효율적인 파이프라인 병렬성을 위한 학습된 부분 공간 압축 (Learned Subspace Compression for Communication-Efficient Pipeline Parallelism)

문제 정의
지리적으로 분산되고 이기종 하드웨어로 구성된 환경에서 거대 언어 모델(LLM)을 학습시키는 것은 특히 파이프라인 병렬성(pipeline parallelism)에서 상당한 통신 병목 현상에 직면하게 합니다. 파이프라인 병렬성은 메모리 제한을 극0복하기 위해 모델 파라미터를 여러 장치에 분할하지만, 이는 순전파(forward) 및 역전파(backward) 과정 모두에서 단계(stage) 간의 중간 활성화 값(intermediate activations) 교환을 필연적으로 수반합니다. 대역폭이 낮은 네트워크에서는 이러한 단계 간 통신이 주요 병목 지점이 됩니다.

Subspace Networks(SSNs)와 같은 기존 방식은 직교 행렬(orthogonal matrix)로 정의된 고정된 공유 저차원 부분 공간(low-rank subspace)으로 활성화 값을 투영하여 이를 완화하려고 시도합니다. 그러나 이 방법은 가중치를 공유 부분 공간 내에 유지하기 위한 제약 최적화와 수정된 옵티마이저의 사용을 포함하여 침습적인 아키텍처 수정이 필요합니다. 경험적으로, 이러한 고정된 글로벌 부분 공간은 압축률이 높아질수록 성능 저하가 심각해지는데, 이는 훈련 과정 중 변화하는 활성화 값의 기하학적 구조에 적응하지 못하기 때문입니다. 또한, 제약 없이 투영 행렬을 학습시키면 투영 행렬이 스티펠 다양체(Stiefel manifold, 직교 행렬의 집합)를 벗어나게 되어, 효과적인 압축에 필요한 등거리 성질(isometric properties)을 파괴하고 심각한 성능 저하를 초사합니다.

방법론: 다양체 인식 투영 학습 (Manifold Aware Projection Learning, MAPL)
저자들은 **MAPL(Manifold Aware Projection Learning)**을 제안하며, 이는 단계 간 압축을 고정된 아키텍처 제약이 아닌 학습 가능한 기하학적 투영으로 취급합니다. MAPL의 핵심 구성 요소는 다음과 같습니다:

1. 단계별 학습 가능한 직교 투영기 (Per-Stage Learnable Orthogonal Projectors): 단일 글로벌 부분 공간을 강제하는 SSN과 달리, MAPL은 각 파이프라인 단계 경계에 고유한 학습 가능한 직교 투영기 $A_p \in \mathbb{R}^{d \times r}$를 배치합니다. 이 투영기들은 스티펠 다양체($St(d, r)$) 내로 제한되어$A_p^\top A_p = I_r$을 보장합니다. 이를 통해 각 단계는 자신만의 작업 최적화된 압축 부분 공간을 발견하고 지속적으로 적응할 수 있습니다.
2. 다양체 제약 최적화 (SPEL): 투영기가 표준 경사 하강 업데이트 중에 다양체를 벗어나는 것을 방-지하기 위해, 저자들은 **스티펠 다양체 상의 스펙트럴 급경사 하강법(Spectral Steepest Descent on the Stiefel Manifold, SPEL)**을 채택했습니다. 이 알고리즘은 매 최적화 단계마다 접사 투영(tangent projections)과 리트랙션 연산(Polar Express 방식을 통한)을 사용하여 투영기를 업데이트함으로써 직교 제약 조건이 엄격하게 유지되도록 합니다.
3. 요인화된 앵커 임베딩 (Factorized Anchor Embeddings): 투영 용량을 소진하지 않으면서 잔차 스트림(residual streams)의 고차원 토큰 기반 오프셋을 처리하기 위해, MAPL은 요인화된 앵커 임베딩을 도입합니다. 오프셋은 학습 가능한 저차원 행렬 $E_{small}^p$와 고정된 직교 행렬 $P_p$로 분해됩니다. 수신 단계에서 이 오프셋을 재구성하기 위해 정수 토큰 ID만을 전송하며, 이는 무시할 만한 수준의 통신 오버헤드만을 추가합니다.
4. 벡터 양자화 (Vector Quantization, VQ): 추가적인 압축을 위해, 이 방법은 저차원 투영된 표현에 대한 잔차 벡터 양자화(RVQ)를 지원합니다. 스트리밍 코드북 동기화 프로토콜을 사용하여 딕셔너리 통신 비용을 분산시킴으로써, 최소한의 대역폭 비용으로 높은 압축률(최대 16배)을 달성합니다.

주요 기여

• 단계별 부분 공간 학습: 본 논문은 공유된 글로벌 부분 공간을 강제하는 것이 토큰 매칭 설정에서 학습을 저해한다는 것을 입증합니다. 대신, MAPL은 각 단계가 서로 다른 스티펠 제약 기반의 압축 부분 공간을 학습할 수 있도록 하여, 다양한 깊이에서의 활성화 값의 고유한 저차원 구조에 더 잘 부합하도록 합니다.
• 다양체 제약 최적화: 저자들은 단순한 투영 학습의 주요 실패 모드로 "다양체 탈출(manifold escape)"을 식별했습니다. 그들은 직교성을 유지하기 위해 SPEL을 도입하였으며, 제약 없는 학습이 고정된 베이스라인보다 엄격하게 열등함을 보여주었습니다.
• 효율적인 재구성: 학습 가능한 투영기와 요인화된 앵커 임베딩을 결다면, MAPL은 무시할 만한 통신 오버헤드로 전체 차원의 활성화 값을 재구성할 수 있으며, SSN에서 요구되는 침습적인 아키텍처 분해를 제거합니다.
• 파레토 프런티어 개선: 이 방법은 검증 교차 엔트로피(validation cross-entropy)와 통신 비용 사이의 파레토 프런티어를 추적하며, 이전 방식들에 비해 성능 저하를 최소화하면서 높은 압축을 달ей합니다.

실험 결과
저자들은 $P=4$ 및 $P=8$ 단계의 파이프라인 구성으로 DCLM 코퍼스에서 학습된 150M에서 1B 파라미터 범위의 디코더 전용 LLaMA 모델에 대해 MAPL을 평가했습니다.

• 성능 대비 압축률: MAPL은 SSN 베이스라인을 일관되게 능가합니다. 4배 압축 시 150M 모델에 대해 MAPL은 검증 손실(validation loss)을 압축되지 않은 베이스라인의 0.84% 이내로 유지한 반면, SSN 베스라인은 **8~12%**의 성능 저하를 보였습니다.
• 확장성: 1B 파라미터 규모에서 8배 압축 시, MAPL은 손실 저하를 **1.38~2.02%**로 유지했으나, SSN 베이스라인은 최대 **26%**의 저하를 겪었습니다.
• 다운스트림 태스크: 제로샷 평가(HellaSwag, PIQA, ARC)에서 MAPL은 압축되지 않은 베이스라인을 밀접하게 추종했습니다(500M/1B 규모에서 평균 정확도 약 1포인트 차이). 반면 SSN 베이스라인은 최대 9포인트의 정확도 하락을 겪었습니다.
• 극한의 압축: 벡터 양자화(MAPL+VQ)와 결합했을 때, 이 방법은 16배 압축을 달성하면서도 완만한 추가 저하(예: 1B 파라미터에서 약 2.3% 손실 증가)를 보였습니다. 다만, 양자화가 적용된 경우 비양자화 MAPL에 비해 다운스트림 태스크 정확도가 더 뚜렷하게 하락했습니다.

의의 및 주장
본 논문은 단계 간 트랜스포머 활성화 값이 고정된 글로벌 부분 공간을 통해 강제되는 것이 아니라, 작업 목적 함수와 함께 공동으로 학습될 수 있는 적응형 저차원 구조를 자연스럽게 가진다고 주장합니다. 압축을 스티펠 다양체 상의 학습 가능한 기하학적 투영으로 다룸으로써, MAPL은 학습 품질과 다운스트림 성능을 보존하면서 공격적인 통신 감소(4배에서 16배)를 가능하게 합니다. 이 연구는 다양체 제약 표현 학습(manifold-constrained representation learning)이 분산된 파운데이션 모델 학습에서 통신 비용을 줄이기 위한 실질적인 메커니즘임을 강조하며, 기존의 부분 공간 기반 접근 방식보다 우수한 성능-압축 트레이드오프를 제공합니다.

한계점
저자들은 현재 실험이 최대 1B 파라미터 규모의 모델로 제한되어 있으며, 더 큰 규모와 실제 이기종 네트워크 환경에서의 추가적인 평가가 필요하다고 언급했습니다. 또한, MAPL이 성능 저하를 크게 줄이기는 하지만, 벡터 양자화가 적용되는 극한의 압축 비율에서는 여전히 성능이 하락하는 경향이 있습니다.