PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training

기술 요약: PRISM

문제 정의

제곱근, 역제곱근, 직교화(극 분해)와 같은 행렬 함수는 신경망 학습(예: Shampoo, Muon)을 위한 전처리된 경사 하강법(preconditioned gradient methods)의 핵심입니다. GPU 가속을 위해 SVD(특이값 분해)나 고유값 추정에 의존하는 방식보다 GEMM(일반 행렬 곱셈)에 기반한 반복 알고리즘이 선호되지만, 기존의 최첨단 방식들은 다음과 같은 중대한 한계를 가집니다:

1. 적응성 부족: 최근의 방법들(예: PolarExpress, CANS)은 종종 사전에 정의된 스펙트럼 구간(예: $[\ell, u]$)에 최적화된 고정된 다항식 업데이트에 의존합니다. 만약 입력 행렬의 실제 특이값 분포가 이 구간에서 벗어나면 수렴 성능이 급격히 저하됩니다.
2. 사전 지식에 대한 의존성: 이러한 방법들은 대개 특이값 분포에 대한 사전 지식이나 최소 및 최대 특이값에 대한 타이트한 경계값이 필요하지만, 이는 행렬 함수 자체를 계산하는 것만큼이나 비용이 많이 들거나 구하기 어려운 경우가 많습니다.
3. 좁은 범위: 최근의 많은 가속 기술들은 다양한 행렬 함수를 위한 일반적인 프레임워크보다는 특정 사용 사례(예: Muon 옵티마이저에서의 극 분해 전용)에 맞춰져 있습니다.

방법론: PRISM 프레임워크

저자들은 PRISM(행렬 함수 계산을 위한 다항식 피팅 및 무작위 반복 스케칭, Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation)을 제안합니다. 이는 명시적인 스펙트럼 경계 없이도 반복적인 행렬 함수 계산을 가속화할 수 있도록 설계된 일반적인 메타 알고리즘입니다. PRISM은 두 가지 별개의 부분으로 작동합니다.

파트 I: 기본 설정 (반복적 다항식 근사)

PRISM은 행렬 함수 $T(A)$의 계산을 반복적인 다항식 근사 문제로 프레임화합니다.

• 스칼라 $x$가 $T(a)$를 근사하는 경우, 문제는 $\xi$가 해(solution)와의 거리를 측정하는 잔차 함수 $f(\xi)$를 근사하는 문제로 재구성됩니다.
• 표준적인 반복 방법(Newton-Schulz 등)은 $\xi=0$ 주변의 고정 차수 테일러 다항식 $f_d(\xi)$를 사용하여 유도됩니다.
• 행렬 반복은 $X_{k+1} = X_k f_d(R_k)$ 형태를 취하며, 여기서 $R_k$는 잔차 행렬(예, sign 함수의 경우 $I - X_k^2$)입니다.

파트 II: 가속 (적응형 피팅 및 스케칭)

잔차 행렬 $R_k$의 스펙트럼이 0에 가깝지 않을 때 발생하는 고정 테일러 다항식의 느린 초기 수렴 문제를 극ОВ하기 위해, PRISM은 두 가지 메커즘을 도입합니다.

1. 적응형 다항식 피팅 (Adaptive Polynomial Fitting):
   고정된 테일러 다항식을 사용하는 대신, PRISM은 최고 차수의 계수를 학습 가능한 파라미터 $\alpha$로 대체합니다. 다항식은 현재 잔차 행렬 $R_k$의 스펙트럼에 맞게 매 반복마다 업데이트됩니다.

• 업데이트 다항식은 $g_d(\xi; \alpha) = f_{d-1}(\xi) + \alpha \xi^d$로 정의됩니다.
• 최적의 $\alpha_k^*$는 $\alpha$에 대해 잔차 노름 $\|I - X_{k+1}^2\|_F^2$(또는 다른 함수의 경우 그에 상응하는 값)를 최소화함으로써 구해집니다. 이는 최소 자승(least-squares) 의미에서 $R_k$의 고유값에 다항식을 적합시키는 효과를 줍니다.
• 안정성을 보장하기 위해 $\alpha_k$는 고정된 구간 $[\ell, u]$(예: $d=1$일 때 $[1/2, 1]$) 내로 제한되며, 이는 입력 스펙트럼과 무관하게 수렴 특성을 보장합니다.

2. 무작위 스케칭 (Randomized Sketching):
   정확한 최적의 $\alpha_k^*$를 계산하려면 잔차 행렬의 고차 거듭제곱의 트레이스(trace)를 계산해야 하는데, 이는 계산 비용이 매우 높습니다($O(n^3)$ 이상).

• PRISM은 (Gaussian 행렬과 같은 Oblivious Subspace Embeddings를 사용하는) 무작위 스케칭을 사용하여 최소 자승 문제를 근사합니다.
• 잔차의 스케치된 버전에 대해 최소화 문제를 해결함으로써, $\alpha_k$를 찾는 비용을 $O(n^2 \log n)$ (또는 $O(n^2 p)$ where $p \ll n$)으로 줄일 수 있으며, 이는 메인 반복문의 $O(n^3)$ 비용에 비해 무시할 수 있는 수준입니다.
• 이론적 보증(Theorem 2)에 따르면, 스케치 차원 $p = O(\log n)$일 때 수렴 속도는 정확한 최소화와 본질적으로 동일하게 유지됩니다.

주요 기여

1. 일반 가속 프레임워크: PRISM은 Newton-Schulz(제곱근, 역제곱근, sign, 극 분해용), Chebyshev 방법, Inverse Newton 방법 등 광범위한 반복 알고리즘에 적용 가능한 통합된 접근 방식을 제공합니다.
2. 분포 불변 적응성 (Distribution-Free Adaptivity): 이전 연구들과 달리, PRISM은 특이값의 경계나 분포에 대한 사전 지식 없이도 현재 반복값의 진화하는 스펙트럼에 동적으로 적응합니다.
3. 효율적인 구현: 무작위 스케칭의 통합은 적응형 다항식 피팅의 오버헤드를 최소화하여, GEMM 중심의 반복 과정이 가진 GPU 친화적인 특성을 보존합니다.
4. 이론적 보증: 본 논문은 스케치 차원이 충분히 클 경우, PRISM 가속 반복이 이차(quadratic) 또는 그 이상의 고차 수렴 속도를 유지한다는 수렴 증명을 제공합니다.

실험 결과

저자들은 합성 벤치마크와 신경망 옵티마이저 통합을 통해 PRISM을 검증했습니다.

• 수렴 속도: 다양한 종횡비와 스펙트럼 분포(사전 학습된 모델을 모사하는 Gaussian 랜덤 행렬 및 Heavy-Tailed Marchenko-Pastur 분포 포함)를 가진 행렬들에 대해, PRISM은 고전적인 Newton-Schulz 및 PolarExpress와 같은 고정 구간 방식보다 일관되게 우수한 성능을 보였습니다.
  • PolarExpress는 실제 특이값 범위가 최적화된 구간에서 벗어날 경우 수렴 속도가 느려지거나 실패하는 경향이 있습니다.
  • PRISM은 $\alpha_k$를 자동으로 조정함으로써 전체 $\sigma_{min}$ 범위에서 안정적인 가속을 유지합니다.
• 신경망 학습:
  • Shampoo 옵티마이저: CIFAR 데이터셋에 대한 ResNet-20/32의 역제곱근 전처리기를 계산하는 데 PRISM을 통합한 결과, 고유값 분해(eigen-decomposition) 및 PolarExpress 베이스라인보다 빠른 학습 시간을 기록했습니다.
  • Muon 옵티마이저: GPT-2 학습(FineWeb 데이터셋) 시 그래디언트 행렬의 극 분해에 PRISM을 적용한 결과, AdamW, PolarExpress 및 고정 차수 Newton-Schulz 변형 모델보다 개선된 검증 손실(validation loss)과 학습 효율성을 입증했습니다.

의의 및 주장

본 논문은 PRISM이 사례 적응적(instance-adaptive)이고 GPU 친화적인 행렬 함수 계산을 위한 실용적이고 일반적인 경로를 제공한다고 주장합니다. PRism의 주요 의의는 "스펙트럼 적응성"을 수동적인 스펙트럼 경계 설정이라는 튜닝 부담에서 벗어나, 신뢰할 수 있는 자동화된 알고리즘 프리미티브(primitive)로 전환했다는 점에 있습니다. 가속화를 사전 스펙트럼 지식으로부터 분리함으로써, PRISM은 다양한 스펙트럼 영역(예: 사전 학습된 모델의 헤비 테일 분포)에서 견고한 성능을 발휘할 수 있게 하여, 대규모 신경망 학습을 가속화하는 다재다능한 도구가 됩니다.