🤖 머신러닝

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

본 논문은 랭크-1 업데이트와 Sherman-Morrison 기반 역행렬 계산을 활용하여 계산 복잡도를 3차에서 2차로, 통신 복잡도를 3차에서 1차로 각각 줄임으로써, 최신 1차 및 2차 최적화 방법들과 비교하여 대규모 언어 모델의 학습 및 수렴을 크게 가속화하는 모멘텀 기반 2차 최적화 알고리즘인 MKOR을 소개한다.

원저자: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

게시일 2024-01-30✓ Author reviewed ⓘ

📖 5 분 읽기🧠 심층 분석

원저자: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

현대 인공지능의 광활한 풍경 속에서, 가장 강력한 시스템들은 패턴을 인식하고, 언어를 번역하며, 텍스트를 생성하는 법을 배우는 복잡한 수학적 구조인 심층 신경망을 기반으로 구축됩니다. 이러한 시스템을 가르치기 위해 연구자들은 거대한 최적화 문제를 해결해야 합니다. 즉, 네트워크가 정확한 예측을 할 수 있도록 하는 내부 숫자, 즉 가중치(weights)의 완벽한 집합을 찾는 것입니다. 수십 년 동안 표준적인 접근 방식은 오차를 가장 빠르게 줄이는 방향으로 이 가중치들을 미세하게 조정하는 것이었으며, 이는 안개 낀 산에서 발밑의 경사를 확인하며 길을 찾아 내려가는 등산객의 방식과 유사합니다. 이 1차 미분법(first-order method)은 신뢰할 수 있고 널리 사용되지만, 골짜기의 바닥에 도달하기 위해 종종 수천 번의 단계를 거쳐야 하므로 느릴 수 있습니다.

2차 최인화(second-order optimization)로 알려진 더 야심 찬 전략은 등산객에게 지면의 곡률을 포함한 전체 지형의 지도를 제공하여 훨씬 더 빠른 하강을 가능하게 하려는 시도입니다. 경사가 어떻게 변하는지 이해함으로써, 이러한 방법들은 이론적으로 훨씬 적은 단계만으로 최적의 해에 도달할 수 있습니다. 그러나 이 접근 방식은 역사적으로 오늘날 사용되는 거대 모델에 적용하기에는 계산 비용이 너무 많이 들었습니다. 대규모 네트워크에 필요한 곡률 정보를 계산하는 데는 엄청난 처리 능력과 메모리가 필요하며, 이는 종종 이 방법이 대체하고자 하는 더 단순한 접근 방식보다 더 느리게 만들기도 합니다. 과제는 지도를 지고 가는 무거운 부담 없이, 지도를 읽는 등산객의 속도를 포착하는 것이었습니다.

토론토 대학교, 러트거스 대학교, 그리고 텍사스 고급 컴퓨팅 센터의 연구팀은 이 딜레마를 해결하는 MKOR이라는 새로운 도구를 개발했습니다. 그들은 모든 세부 사항을 계산하지 않고도 지형의 형태를 근사화하기 위해 특정 수학적 지름길을 사용하는 모멘텀 기반 옵티마이저를 만들었습니다. 손실 지형(loss landscape)의 완전하고 복잡한 곡률을 계산하려고 시도하는 대신, MKOR은 랭크-1 업데이트(rank-1 updates)라고 불리는 기술을 사용합니다. 이 접근 방식은 변화의 가장 유의미한 방향에 집중함으로써 계산을 단순화하며, 계산 비용을 효과적으로 3차 관계에서 2차 관계로 줄여줍니다. 실질적인 관점에서, 이는 모델이 커짐에 따라 방법론이 훨씬 더 잘 확장됨을 의미하며, 이전에 2차 최적화 방법이 거대 언어 모델에 사용되는 것을 막았던 병목 현상을 피할 수 있게 해줍니다.

연구진은 BERT-Large-Uncased로 알려진 거대 언어 모델을 대상으로 64개의 고성능 그래픽 프로세서에서 실험을 수행하여 이 새로운 옵티마이저를 테스트했습니다. 그들은 MKOR을 기존의 가장 우수한 1차 방법 및 사용 가능한 가장 진보된 2차 구현 방식들과 비교했습니다. 결과는 놀라웠습니다. 질문 답변 작업에서 표준 1차 방법은 특정 정확도 수준에 도달하기 위해 1,536단계가 필요했습니다. 새로운 MKOR 방식은 단 600단계 만에 동일하거나 더 높은 정확도에 도달했습니다. 훈련 완료에 필요한 총 시간을 측정했을 때, MKOR은 3시간 조금 넘게 걸린 반면, 표준 방식은 거의 8시간이 걸렸습니다. 이는 2.57배의 속도 향상을 나타내며, 실질적으로 훈련 시간을 절반 이상 단라한 것입니다.

순수 속도 외에도, 이 연구는 2차 최적화 방법에서 흔히 발생하는 치명적인 안정성 문제를 다루었습니다. 이러한 방법들은 복잡한 수학적 요인에 기반하여 학습 과정을 조정하기 때문에, 조정값이 너무 커져서 훈련 과정이 중단되거나 발산하는 현상인 "기울기 폭주(exploding gradients)"가 발생하기 쉽습니다. 이전의 해결책들은 안전을 유지하기 위해 학습률을 크게 늦춰야 하는 경우가 많았고, 이는 속도의 이점을 상쇄했습니다. MKOR은 이러한 조정값의 크기를 모니터링하는 내장 안정화 장치를 포함하고 있습니다. 값이 너무 커지면 시스템이 자동으로 이를 재조정(rescale)하여, 2차 접근 방식의 공격적인 속도를 유지하면서도 훈련이 폭주하는 것을 방지합니다. 이를 통해 옵티마이저는 불안정성의 위험 없이 빠르게 이동할 수 있습니다.

연구진은 또한 모델이 데이터의 광범위한 윤곽을 배우는 훈련 초기 단계에서 2차 방법의 이점이 가장 두드렷하다는 것을 관찰했습니다. 모델이 최적의 해에 가까워질수록 복잡한 곡률 계산의 이점은 줄어듭니다. 이를 해결하기 위해 그들은 MKOR-H라고 불리는 하이브리드 버전의 옵티마이저를 도입했습니다. 이 버전은 강력한 2차 방법을 사용하여 초기 학습 단계를 가속화한 다음, 훈련 손실이 급격히 감소하는 것을 멈추면 자동으로 더 단순한 1차 방법으로 전환됩니다. 이 하이브리드 접근 방식은 시스템이 훨씬 더 적은 단계 내에 수렴하도록 하여, 표준 방법보다 절반도 안 되는 시간 만에 동일한 최종 정확도를 달 수 있게 했습니다.

연구는 또한 왜 이 특정 수학적 지름길이 이토록 잘 작동하는지도 조사했습니다. 연구진은 신경망 내의 데이터 구조를 분석하여, 네트워크의 서로 다른 부분들이 어떻게 상호작용하는지를 설명하는 공분산 행렬(covariance matrices)이 종종 낮은 랭크(low-rank)를 가진다는 것을 발견했습니다. 이는 데이터의 차원이 높음에도 불구하고 정보가 몇 개의 핵심적인 방향에 집중되어 있음을 의미합니다. 이러한 행렬들을 랭크-1 업데이트를 사용하여 근사함으로써, MKOR은 전체 행렬 계산의 오버헤드 없이 최적화에 필요한 필수 정보를 포착합니다. 연구팀은 광범위한 테스트를 통해 이 근사 방식이 매우 적은 오차만을 유발한다는 것을 검증하였으며, 이러한 단순화가 정확도의 희생을 수반하지 않음을 확인했습니다.

속도와 안정성 외에도, 이 새로운 방법은 훈련 중에 서로 다른 프로세서 간에 전송되어야 하는 데이터의 양을 크게 줄여줍니다. 수천 개의 프로세서가 함께 작동하는 분산 컴퓨팅 환경에서 통신은 주요 병목 현상이 될 수 있습니다. MKOR은 이러한 통신 복잡도를 2차 관계에서 선형 관계로 줄여주며, 이는 더 많은 프로세서가 추가되더라도 데이터 전송 지연으로 인해 속도가 저하되지 않고 시스템이 효율적으로 확장됨을 의미합니다. 이러한 효율성 덕분에 이 방법은 현대의 거대 언어 모델이 요구하는 거대한 배치 크기를 처리할 수 있으며, 이는 이전의 2차 접근 방식들이 가졌던 한계였습니다.

이 작업의 함의는 단 하나의 특정 모델을 훈련하는 것을 넘어섭니다. 2차 최적화를 대규모 시스템에 실용적으로 만듦으로써, MKOR은 인공지능의 더 효율적인 훈련을 향한 경로를 제시합니다. 연구진은 그들의 방법이 머신러닝 종사자들에게 상당한 양의 에너지와 시간을 절약하는 데 도움이 될 수 있다고 언급했습니다. 기술 자체는 중립적인 도구이지만, 의료 진단에서 기후 모델링에 이르는 분야에서의 적용은 과학적 발견을 가속화할 수 있습니다. 그러나 저자들은 또한 효율적인 최적화 알고리즘이 부정적인 사회적 영향을 미치는 모델을 포함한 모든 모델에 적용될 수 있음을 인정하며, 이러한 기술의 책임 있는 사용은 그것을 사용하는 사람들의 선택에 달려 있음을 강조했습니다.

궁극적으로, 이 논문은 딥러닝의 오랜 문제에 대한 구체적인 해결책을 제시합니다. 모멘텀, 랭크-1 근사, 그리고 하이브리드 전환 메커니즘을 결합함으로써, 연구진은 이전 모델들보다 더 빠르고, 더 안정적이며, 더 확장 가능한 옵티마이저를 만들어냈습니다. 실험은 2차 방법의 강력한 힘을 활용하면서도, 그동안 2차 방법의 사용을 막았던 과도한 계산 비용을 피하는 것이 가능하다는 것을 입증합니다. 인공지능 분야가 더 크고 복잡한 모델을 향해 계속 나아감에 따라, MKOR과 같은 도구는 이러한 시스템을 더 효율적으로 훈련할 수 있는 방법을 제공하며, 한때 이론적인 이점에 머물렀던 것을 실질적인 현실로 바꾸어 놓습니다.

기술 요약: MKOR – Rank-1 업데이트를 사용하는 모멘텀 기반 크로네커 인자 기반 옵티마이저

1. 문제 정의

Natural Gradient Descent (NGD) 및 Kronecker-Factored Approximate Curvature (KFAC)와 같은 2차 최적화 방법은 곡률 정보(Hessian 또는 Fisher Information Matrix 근사치)를 활용함으로써 1차 방법(예: SGD, Adam, LAMB)보다 우수한 수렴 속도를 제공한다. 그러나 대규모 딥 뉴럴 네트워크(DNN), 특히 트랜스포머(Transformer) 및 대규모 언어 모델(LLM) 학습에 있어 이들의 채택은 세 가지 주요 병목 현상으로 인해 저해된다:

계산 복잡도: 전통적인 KFAC 방법은 $d \times d$ 크기(여기서 $d$ 는 레이어 차원)의 공분산 행렬 역행렬을 필요로 하며, 이는 $O(d^3)$ 의 복잡도를 초래한다. 마찬가지로, Sherman-Morrison-Woodbury (SMW) 기반 방법(SNGD)은 종종 배치 크기 $b$ 에 따라 $O(b^3)$ 로 스케일링된다. 배치 크기가 시퀀스 길이에 따라 확장되는 트랜스포머 모델에서 이러한 복잡도는 매우 부담스럽다.
통신 오버헤드: 분산 학습에서는 워커(worker) 간에 2차 정보(공분산 행렬 및 그 역행렬)를 동기화해야 한다. KFAC는 $O(d^2)$ 의 데이터를 동기화해야 하는 반면, SNGD는 $O(bd + b^2)$ 를 요구하여 확장성을 제한하는 심각한 통신 병목을 생성한다.
업데이트 빈도 및 안정성: 높은 역행렬 계산 비용으로 인해, 최신 구현체(KAISA 등)는 "오래된(stale)" 인자를 사용하여 2차 정보를 드물게 업데이트한다(예: 100~1000회 반복마다). 이는 곡률 정보의 효과를 감소시킨다. 또한, 2차 방법은 수치적 불안정성(그래디언트 폭주, 특이 행렬)에 취ñ약하며, 댐핑 인자(damping factors)나 고정밀 연산을 필요로 하는 경우가 많아 계산 비용을 더욱 증가시킨다.

2. 방법론: MKOR 알고리즘

저자들은 위 문제들을 네 가지 핵심 메커니즘을 통해 해결하는 MKOR(Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates)를 제안한다:

Rank-1 근사 및 SM 기반 역행렬 계산:
전체 공분산 행렬을 계산하고 역행렬을 구하는 대신, MKOR는 샘플의 평균( $E[aa^T] \approx \bar{a}\bar{a}^T$ )으로부터 유도된 rank-1 행렬을 사용하여 활성화(activation)와 그래디언트의 공분산을 근사한다. 이러한 rank-1 업데이트에 Sherman-Morrison 항등식을 적용함으로써, MKOR는 크로네커 인자의 역행렬을 직접 업데이트한다. 이를 통해 인자 역행렬 계산의 복잡도를 $O(d^3)$ 에서 $O(d^2)$ 로 줄인다.
- 업데이트 규칙: 역행렬 인자는 모멘텀( $\gamma$ )과 Sherman-Morrison 공식을 사용하여 업데이트되며, 이를 통해 전체 행렬 역행렬 계산 비용 없이 빈번한 업데이트가 가능하다.
선형 통신 복잡도:
분산 환경에서 MKOR는 전체 공분산 행렬( $4d^2$ ) 대신 두 개의 rank-1 근사 벡터(크기 $2d$ )만을 워커 간에 동기화한다. 이는 통신 복잡도를 이차식 $O(d^2)$ 에서 선형 $O(d)$ 로 줄인다. 또한, MKOR는 이 벡터들에 대해 half-precision (FP16)을 사용하여 통신 비용을 다시 절반으로 줄이는데, 이는 수치적 민감도 때문에 다른 2차 방법들이 달성하기 어려운 성과이다.
안정화 메커니즘:
2차 방법에서 흔히 발생하는 수치적 불안정성과 그래디언트 폭주를 방지하기 위해:
- Norm 기반 안정화 장치: 역행렬 인자의 무한 노름(infinity norm)을 모니터링한다. 만약 임계값을 초과하면, 발산을 방지하기 위해 인자를 항등 행렬과 블렌딩한다 ( $\hat{L} = \zeta L + (1-\zeta)I$ ).
- 그래디언트 재스케일링: 프리컨디셔닝된(preconditioned) 그래디언트를 원래 그래디언트의 노름에 맞게 재스케일링하여, 러닝 레이트 스케줄러가 효과적으로 유지되도록 한다.
- 이론적 안정성: 저자들은 (Lemma 3.1을 통해) 인자들이 인위적인 댐핑 인자( $\mu I$ ) 없이도 양의 정치성(positive-definite)을 유지하며, 업데이트 규칙의 단일 스칼라 나눗셈이 0이 아님이 보장됨을 증명한다.
하이브리드 MKOR (MKOR-H):
학습이 수렴함에 따라(손실 지형이 identity에 가까워짐) 2차 방법의 이점이 종종 감소한다는 점을 인식하여, MKOR-H는 스위칭 메커즘을 도입한다. 이는 손실 감소율을 모니터링하고, 2차 업데이트가 더 이상 가속을 제공하지 못할 경우 자동으로 1차 옵티마이저로 전환하여 비용과 수렴 속도 사이의 트레이드오프를 최적화한다.

3. 주요 기여

복잡도 감소: MKOR는 $O(d^2)$ 의 계산 복잡도와 $O(d)$ 의 통신 복잡도를 달성하여, KFAC의 $O(d^3)$ 및 $O(d^2)$ , SNGD의 $O(b^3)$ 병목을 완화한다.
고빈도 업데이트: 역행렬 계산 비용을 줄임으로써, MKOR는 매 반복(또는 매우 빈번하게)마다 2차 업데이트를 가능하게 하여 "오래된" 인자의 필요성을 제거하고 곡률 정보를 완전히 활용한다.
강건성: 노름 기반 안정화 및 그래디언트 재스케일링의 도입을 통해 MKOR는 댐핑 인자나 배정밀도(double-precision) 연산 없이도 안정적으로 작동하며, half-precision 하드웨어의 효율적인 사용을 가능하게 한다.
하이브리드 전략: MKOR-H 변형 모델은 전체 학습 생애 주기 동안 효율성을 유지하기 위해 2차 방법과 1차 방법 사이를 동적으로 전환한다.

4. 실험 결과

저자들은 64개의 A100 GPU를 사용하여 BERT-Large-Uncased에 대해 MKOR를 평가하였으며, 최신 1차 옵티마이저인 LAMB와 2차 구현체인 KAISA (KFAC) 및 HyLo (SNGD), 그리고 Eva와 비교하였다.

가속도: MKOR는 사전 학습(pre-training) 작업에서 LAMB 대비 2.57배, KAISA 대비 1.85배의 속도 향상을 달了. 구체적으로, BERT-Large-Uncased의 학습 시간을 8시간에서 3시간으로 단축시켰다.
수렴: SQuAD v1.1 질의응답 작업에서 MKOR-H는 600 스텝 만에 수렴하였으며(LAMB는 1,536 스텝), 대등하거나 더 나은 F1 점수(90.64 vs 90.44)를 달성했다.
GLUE 벤치마크: MKOR는 1,500 스텝 만에 GLUE 데이터셋에서 새로운 SOTA(state-of-the-art) 지표를 달성했다. 특히, KAISA는 1,563 스텝 이후에도 베이스라인 평균 정확도에 수렴하지 못한 반면, MKOR는 성공적으로 수렴했다.
확장성: 높은 메모리 오버헤드( $O(bd + b^2)$ )로 인해 발산하거나 GPU 메모리에 적재하지 못한 HyLo와 달리, MKOR는 메모리 문제 없이 큰 배치 크기(accumulation을 통한 최대 64k 유효 배치 크기)로 성공적으로 학습되었다.
근사 타당성: ResNet-50 및 BERT에 대한 실험은 공분산 행렬의 rank-1 근사가 무시할 만한 오차만을 발생시킨다는 것을 보여주었으며, 이는 핵심 근사 가정을 검증한다.

5. 의의 및 주장

본 논문은 MKOR가 2차 최적화의 이론적 이점과 현대적 대규모 DNN 학습의 실제적 제약 사이의 간극을 성공적으로 메웠다고 주장한다. 2차 업데이트의 복잡도를 (모델 크기에 따른 스케일링 측면에서) 1차 방법과 대등한 수준으로 줄이면서 수렴 이점을 유지함으로써, MKOR는 다음을 가능하게 한다:

더 빠른 학습: LLM의 실제 학습 시간(wall-clock time)을 대폭 단축한다.
더 나은 수렴: 다른 2차 방법(KAISA 등)이 실패하거나 정체되는 지점에서 더 높은 정확도 지표로 수렴할 수 있는 능력을 제공한다.
확장성: 이전에 2차 방법을 제한했던 통신 및 메모리 병목 없이 대규모 배치 크기와 모델 차원으로 확장할 수 있다.

저자들은 본 연구가 "대규모 모델" 시대에 2차 최적화를 실행 가능하게 만들며, 특히 트랜스포머 아키텍처에서 시퀀스 길이에 의존하는 배치 크기로 인해 발생하는 확장성 문제를 해결한다고 강조한다. 결론적으로 MKOR는 심층 신경망의 효율적이고 효과적인 학습을 위한 새로운 SOTA를 나타낸다.