LoRDO: Distributed Low-Rank Optimization with Infrequent Communication

기술 요약: LoRDO – 빈번하지 않은 통신을 활용한 분산 저차원 최적화

1. 문제 정의

분산 데이터 병렬(DDP)을 통한 파운데이션 모델의 분산 학습은 상호 연결 대역폭과 옵티마이저 상태(예: Adam 모멘타)를 저장하는 데 필요한 메모리 오버헤드에 의해 점점 더 제약을 받고 있습니다. 빈번하지 않은 통신(로컬 업데이트)을 사용하는 전략은 동기화 빈도를 줄여주지만, 여전히 풀-랭크(full-rank) 옵티마이저 상태의 메모리 및 통신 비용에 의해 병목 현상이 발생합니다.

저차원 적응형 옵티마이저(예: GaLore, LDAdam)는 그래디언트를 저차원 부분 공간으로 투영함으로써 이러한 제약을 완화할 수 있는 경로를 제시합니다. 그러나 이러한 방법들을 빈번하지 않은 동기화 환경으로 일반화하는 것은 다음과 같은 결정적인 최적화 과제를 안겨줍니다:

1. 로컬 투영(Local Projections): 개별 워커에서 확률적 그래디언트를 사용하여 로컬 저차원 투영을 계산하는 것은 작은 유효 배치 크기로 인해 상당한 노이즈를 유발합니다. 이는 부적절한 투영 부분 공간을 초래하고 성능 저하를 일으킵니다.
2. 글로벌 투영(Global Projections): DDP와 같이 집계된 의사 그래디언트(pseudo-gradient)로부터 도출된 글로벌 투영 행렬을 사용하는 것은 더 안정적인 기저(basis)를 제공합니다. 그러나 로컬 업데이트 환경에서 이 접근 방식은 최적화 궤적을 고정된 저차원 부분 공간 내로 영구적으로 제한합니다. 풀-랭크 신호가 없으면, 투영 기저가 초기 부분 공간 외부의 새로운 방향을 "갱신"하거나 탐색하지 못하기 때문에 최적화가 정체됩니다.

핵서 질문은 다음과 같습니다: 통신 효율적인 학습을 위해, 어떻게 하면 낮은 메모리와 통신 오버헤드를 유지하면서도 부분 공간 정체(subspace stagnation)를 피할 수 있는 고성능 저차원 옵티마이저를 설계할 수 있는가?

2. 방법론: LoRDO

저자들은 저차원 최적화와 빈번하지 않은 동기화를 통합하는 프레임워크인 LoRDO(Low-Rank Distributed Optimization)를 제안합니다. 이 방법론은 세 가지 핵심 메커니즘을 통해 투영 안정성과 부분 공간 탐색 사이의 긴장을 해결합니다.

A. 글로벌 투영 행렬 (Global Projection Matrices)

각 워커가 자신의 로컬 확률적 그래디언트에 기반하여 투영 행렬 $Q_m$을 결정하는 대신(이는 높은 분산을 초래함), LoRDO는 글로벌 투영 행렬 $Q_t$를 사용합니다. 이 행렬은 $K$번의 로컬 스텝 동안의 모델 파라미터 총 변화량을 나타내는 집계된 의사 그래디언트($\Delta_t$)를 사용하여 동기화 경계에서 계산됩니다.

• 이점: 의사 그래디언트는 유효 배치 크기가 $M \times B$($M$은 워커 수)이므로, 로컬 그래디언트에 비해 분산이 적고 더 안정적인 신호를 생성합니다. 또한 의사 그래디언트에 내재된 곡률 정보를 포함하므로 더 고품질의 기저 추정을 가능하게 합니다.

B. 풀-랭크 준-쌍곡선 모멘텀 (Full-Rank Quasi-Hyperbolic Momentum, QHM)

글로벌 투영에서 발생하는 "정체된 부분 공간" 문제를 해결하기 위해, LoRDO는 로컬 업데이트에 풀-랭크 준-쌍곡선 모멘텀 항을 주입합니다.

• 메커니즘: 업데이트 규칙은 저차원 모멘텀과 두 번째 모멘텀으로 스케일링된 풀-랭크 그래디언트 신호를 결합합니다.
  $$\text{Update} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• 효과: 이 주입은 집계된 의사 그래디언트가 고정된 rank-$r$ 부분 공간 내에 갇혀 있는 것을 방지합니다. 이를 통해 서로 다른 워커들로부터의 rank-$r$ 부분 공간을 시간적으로 집계하여, 결국 완전한 부분 공간 탐색을 가능하게 합니다. 이는 통신 오버헤드를 늘리지 않고도 전체 솔루션 공간을 탐색하는 능력을 복원합니다.

C. 에러 피드백 및 모멘타 회전 (Error Feedback and Momentum Rotation)

LoRDO는 저차원 압축 과정에서 손실된 정보를 처리하기 위해 로컬에서 에러 피드백을 포함합니다. 또한, 새로운 투영 행렬을 계산한 후에는 옵티마이저 상태(모멘타)를 **회전(rotate)**시켜, 모멘타가 동일한 부분 공간에 계속 축적되도록 보장합니다. 이는 LDAdam에서 채택된 기술입니다.

3. 주요 기여

1. 투영 실패 모드 분석: 저자들은 로컬 투영이 작은 워커 배치 크기로 인한 높은 분산 때문에 성능을 해치며, 풀-랭크 신호가 없는 글로벌 투영은 부분 공간 정체를 유발한다는 것을 입증했습니다.
2. 전체 부분 공간 탐색 복원: LoRDO는 글로벌 투영을 사용하여 투영 안정성을 확보하는 동시에, 로컬 업데이트에 풀-랭크 그래디언트 신호(QHM을 통해)를 주입하는 저차원 옵티마이저를 도입했습니다. 이는 통신이나 메모리 오버헤드를 증가시키지 않고도 정체를 방지합니다.
3. DDP와의 동등성 및 효율성: 이 프레임워크는 동기식 저차원 DDP와 거의 대등한 성능(perplexity 차이 < 1%)을 달면서도, 옵티마이저 메모리와 통신을 최대 $8\times$–$12\times$까지 줄입니다. 저차원이 필수적인 극심한 메모리 제약 상황에서 LoRDO는 심지어 DDP보다 우수한 perplexity를 보여줍니다.

4. 실험 결과

실험은 SmolLM2 데이터 믹스처를 사용하여 16M, 125M, 720M 파라미터 규모로 확장된 디코더 전용 트랜스포머를 대상으로 수행되었습니다.

• 정체 vs. 탐색: 16M 모델에서 QHM이 없는 글로벌 투영은 학습에 실패(perplexity 46.8)하여 DDP(36.6)에 비해 성능이 낮았으며, 이는 부분 공간 정체를 확인시켜 주었습니다. 풀-랭크 QHM을 추가하면 성능이 복원되어 LoRDO-Global이 로컬 변형 모델들보다 우수한 성능을 보였습니다.
• 투영 품질: LoRDO-Global은 특히 낮은 랭크($r=8, 16$)에서 집계된 의사 그래디언도로부터 도출된 더 높은 품질의 투영 기저 덕분에 우수한 안정성과 성능을 보여주었습니다. 또한 로컬 배치 크기 변화에 대해 로컬 투영 방식보다 낮은 민감도를 보였습니다.
• 720M으로의 확장: 720M 규모에서 LoRDO-Global ($r=256$)은 저차원 DDP ($r=256$)와 거의 대등한 성능(perplexity 차이 10.41 vs. 10.34 미만)을 달성했습니다. 또한 DDP와 일치하는 다운스트림 태스크 정확도(ARC, HellaSwag, MMLU, PIQA)를 보여주었습니다.
• 효율성 이득:
  • 통신: LoRDO는 저차원 DDP 대비 약 10배, 풀-랭크 DDP 대비 25배 적은 통신 오버헤드를 기록했습니다.
  • 메모리: 풀-랭크 대응 모델 대비 옵티마이저 상태 메모리 오버헤드를 8배 줄였습니다.
  • 실제 실행 시간(Wall-Clock Time): 저대역폭 환경(10 Gbps로 모델링)에서 LoRDO는 감소된 동기화 빈도 덕분에 DDP보다 월등히 나은 실제 실행 시간을 보여주었습니다.

5. 의의 및 주장

본 논문은 Lo-RDO가 투영 안정성과 부분 공간 탐색 사이의 근본적인 긴장을 해결한다고 주장합니다. 집계된 의사 그래디언드를 사용하여 투영 안정성을 확보하고, 풀-랭크 준-쌍곡선 모멘텀을 통해 탐색을 수행함으로써 LoRDO는 다음을 가능하게 합니다:

• 확장 가능한 학습: 통신 및 메모리 비용을 크게 줄이면서도 대규모(최대 720M 파라미터)에서 경쟁력 있는 성능을 제공합니다.
• 자원 효율성: 풀-랭크 옵티마이저를 사용할 수 없는 제한된 자원(낮은 메모리/대역폭)의 하드웨어에서도 사전 학습을 가능하게 합니다.
• 분산 학습: 모델 학습을 단일 데이터 센터 너머로 확장함으로써 효율적인 분산 환경에서의 학습을 위한 길을 엽니다.

저자들은 공식적인 수렴 증명이 부족하다는 점과 현재 연구가 매트릭스 기반 옵티마이저(예: Muon)보다는 적응형 옵티마이저(Adam)에 집중되어 있다는 한계점을 언급하며, 이를 향후 과제로 남겨두었습니다. 그러나 이들은 작은 모델에서 관찰된 경향성이 더 큰 규모에서도 유지될 것이라고 긍정합니다.