Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration

기술 요약: 파워 반복을 통한 부분 직교화로 제로차 스펙트럼 최적화 가속화

문제 제기
제로차 (ZO) 최적화는 메모리 제약이 있는 엣지 장치에서 대규모 언어 모델 (LLM) 을 미세 조정하기 위한 핵심 기법으로 부상했습니다. 이는 1 차 (FO) 방법이 요구하는 메모리 집약적인 역전파를 회피하기 때문입니다. 그러나 ZO 방법은 몬테카를로 섭동으로 인해 기울기 추정치의 분산이 높아, FO 방법에 비해 수렴 속도가 현저히 느린 문제가 있습니다.

최근 ZO 최적화를 가속화하려는 시도는 저차원 부분공간 탐색 (예: LOZO, SubZero) 과 Muon 과 같은 스펙트럼 최적화기에 집중되어 왔으며, 이들은 직교화를 통해 행렬 기하학을 활용합니다. Muon 은 기울기를 직교화하여 약한 스펙트럼 방향을 활용함으로써 FO 환경에서 AdamW 를 능가하지만, 저자들은 Muon 을 ZO 환경에 직접 적용할 경우 성취도가 제한적임을 관찰했습니다. 핵심 문제는 ZO 기울기 추정치가 본질적으로 노이즈가 많고 신뢰할 수 없다는 점입니다. FO 기울기가 지배적인 특이 방향에 집중되는 것과 달리, ZO 추정치는 "더 무겁고 평평한 스펙트럼 꼬리"를 보입니다 (그림 1). 결과적으로 전체 활성 스펙트럼에 걸쳐 에너지를 균등하게 만드는 Muon 의 표준 뉴턴 - 슈르츠 (NS) 직교화는 이러한 노이즈가 많고 부정확한 방향을 부주의하게 증폭시켜, 추정 오해를 해로운 업데이트 구성 요소로 전환시킵니다.

방법론: ZO-MOPI
저자들은 ZO 최적화를 가속화하기 위해 지배적인 스펙트럼 방향을 선택적으로 추출하면서 노이즈를 억제하도록 설계된 ZO-MOPI(Zeroth-Order Momentum with Partial Orthogonalization) 를 제안합니다. 이 접근법은 세 가지 핵심 요소를 통합합니다:

1. 스트리밍 파워 반복 (SPI) 을 통한 부분 직교화:
   전체 스펙트럼 직교화 (NS 반복과 같은) 를 수행하는 대신, ZO-MOPI 는 "부분 직교화" 전략을 채택합니다. 이는 노이즈가 많은 스펙트럼 꼬리를 폐기하고 상위 $k$개 지배적인 특이 방향에만 집중합니다.

   • 메커니즘: 이 방법은 NS 절차를 더 빠르고 집중된 **스트리밍 파워 반복 (SPI)**으로 대체합니다. SPI 는 훈련 단계 전반에 걸쳐 선도적인 $k$차원 부분공간을 추적합니다.
   • 효율성: 오른쪽 특이 부분공간 ($V_{t-1}$) 을 캐시하고 이를 다음 반복의 웜 스타트로 사용하여, ZO-MOPI 는 전체 SVD 나 비싼 행렬 곱셈을 피합니다. 계산 비용은 전체 스펙트럼 방법의 비용에 비해 $O(rnk + nk^2)$로 감소하며, 여기서 $k \ll r \ll \min(m, n)$입니다.
   • 안정성: 이 방법은 지배적인 스펙트럼 방향이 모멘텀 평활화 하에서 매끄럽게 진화한다는 관찰에 의존하므로, 이전 부분공간이 현재 단계에 대한 신뢰할 수 있는 초기화가 됩니다.

2. 부분공간 ZO 추정:
   분산을 추가로 줄이기 위해 기울기 추정은 저랭크 부분공간으로 제한됩니다. 무작위 섭동 행렬 $Z$는 $AB$로 분해되는데, 여기서$A$는 고정된 부분공간 행렬이고$B$는 해당 부분공간 내의 무작위 섭동입니다. 이는 기울기 추정기의 유효 탐색 차원을$mn$에서$rn$으로 줄여 분산을 현저히 낮춥니다.

3. 투사된 모멘텀:
   메모리를 절약하기 위해 모멘텀은 전체 매개변수 공간이 아닌 축소된 $r \times n$ 좌표계 내에서 유지됩니다. 이 방법의 핵심 구성 요소는 부분공간 $A$가 주기적으로 갱신될 때 (매 $\nu$ 반복마다) 모멘텀을 투사하는 것입니다. 모멘텀 $M_{old}$는 $M_{new} \leftarrow A_{new}^\top A_{old} M_{old}$를 통해 새로운 부분공간 $A_{new}$로 투사됩니다.

   • 필요성: 저자들은 모멘텀이 단순히 분산 감소를 위한 것이 아니라 SPI 의 안정성에 필수적임을 입증했습니다. 모멘텀이 없으면 노이즈가 많은 ZO 추정치로 인해 지배적인 부분공간이 급격하게 요동쳐 파워 반복이 불안정해집니다.

주요 기여

• 스펙트럼 불일치 식별: 이 논문은 전체 스펙트럼 직교화 (뉴턴 - 슈르츠) 가 ZO 환경에서 해롭다고 강조합니다. 이는 ZO 기울기 추정치에 내재된 노이즈를 증폭시키기 때문이며, 반면 FO 기울기는 자연스럽게 집중되어 있습니다.
• 부분 직교화: 신뢰할 수 있는 스펙트럼 성분 (상위 $k$개) 만 보존하고 노이즈가 많은 꼬리는 무시하는 선택적 직교화 전략을 도입했습니다.
• 스트리밍 파워 반복 (SPI): 시간적 연속성과 웜 스타팅을 활용하여 비싼 계산을 피하는 지배적인 특이 방향 추적을 위한 경량 효율 알고리즘입니다.
• 투사된 모멘텀 통합: 부분공간 갱신 전반에 걸쳐 모멘텀 안정성을 유지하는 메커니즘으로, 이는 노이즈가 많은 환경에서 스펙트럼 추적의 효과성에 결정적인 것으로 입증되었습니다.

실험 결과
저자들은 SuperGLUE 벤치마크 (SST-2, RTE, BoolQ, SQuAD) 에서 OPT-13B, LLaMA3-8B, Gemma2-2B 모델을 사용하여 LLM 미세 조정 작업에서 ZO-MOPI 를 평가했습니다.

• 수렴 속도: ZO-MOPI 는 현재 최첨단 ZO-Muon 대비 1.5 배에서 4 배 빠른 수렴을 달성합니다. 예를 들어, RTE 및 SQuAD 작업에서 ZO-Muon 은 ZO-MOPI 와 동일한 정확도에 도달하는 데 약 3 배의 월클럭 시간이 필요합니다.
• 최종 정확도: 이 방법은 MeZO, LOZO, ZO-Muon 과 같은 강력한 베이스라인과 비교하여 경쟁력 있거나 우수한 최종 정확도를 달성합니다. 여러 경우 (예: SST-2 및 RTE 에서의 OPT-13B) ZO-MOPI 는 새로운 최고 ZO 정확도 기록을 세웠습니다.
• 효율성: 이 방법은 부분공간 및 파워 반복 상태를 저장하기 위한 무시할 수 있는 추가 메모리 오버헤드만 발생시키며, ZO 미세 조정의 메모리 효율성을 유지합니다.
• 애블레이션 연구: 실험은 모멘텀을 제거하거나 부분공간 갱신 간격이 너무 짧을 경우 성능이 현저히 저하됨을 확인시켜, 제안된 모멘텀 투사 및 지연 부분공간 샘플링 전략의 필요성을 검증했습니다.

의의 및 주장
이 논문은 ZO-MOPI 가 노이즈가 많은 기울기 추정치의 특정 스펙트럼 특성을 해결함으로써 LLM 미세 조정을 위한 ZO 와 FO 최적화 간의 성능 격차를 성공적으로 해소했다고 주장합니다. 저자들은 ZO 최적화를 가속화하는 열쇠가 부분공간 투사를 통한 분산 감소뿐만 아니라 노이즈를 증폭시키지 않는 선택적 스펙트럼 정렬을 채택하는 데 있다고 봅니다. 공격적인 전체 스펙트럼 직교화를 파워 반복을 통한 부분 직교화로 대체함으로써, 이 방법은 고분산 ZO 영역에서도 스펙트럼 최적화기가 효과적으로 작동하도록 합니다. 이 연구는 노이즈가 많은 추정치 하에서도 스펙트럼 구조를 더 잘 활용할 수 있음을 시사하며, 역전파에 의존하지 않고 대규모 모델을 효율적이고 메모리 제약이 있는 환경에서 미세 조정할 수 있는 길을 제시합니다.