[논문 리뷰] Zero-th Order Algorithm for Softmax Attention Optimization

Yichuan Deng, Zhihang Li외 2명|arXiv (Cornell University)|2023. 07. 17.

Topic ModelingComputer Science인용 수 3

한 줄 요약

이 논문은 대규모 언어 모델(Large Language Models, LLMs)의 소프트맥스 어텐션 기반 최적화를 위해 설계된 제로스터드 최적화 알고리즘을 제안한다. 이는 오직 정방향 전파만을 통해 기울기 추정을 가능하게 하며, 전역 수렴성을 보장한다. 수렴 속도는 $ O(M \cdot (1 + d^{1.5} \kappa^2(A)/k) \cdot \mu^{-2} B^{-1} \log((L(x_0)-L^*)/\epsilon)) $ 로 표현되며, 여기서 $ M = \exp(O(R^2 + \log n)) $ 이다. 이는 고차원 설정에서 백프로파게이션 없이도 효율적인 훈련이 가능함을 보여준다.

ABSTRACT

Large language models (LLMs) have brought about significant transformations in human society. Among the crucial computations in LLMs, the softmax unit holds great importance. Its helps the model generating a probability distribution on potential subsequent words or phrases, considering a series of input words. By utilizing this distribution, the model selects the most probable next word or phrase, based on the assigned probabilities. The softmax unit assumes a vital function in LLM training as it facilitates learning from data through the adjustment of neural network weights and biases. With the development of the size of LLMs, computing the gradient becomes expensive. However, Zero-th Order method can approximately compute the gradient with only forward passes. In this paper, we present a Zero-th Order algorithm specifically tailored for Softmax optimization. We demonstrate the convergence of our algorithm, highlighting its effectiveness in efficiently computing gradients for large-scale LLMs. By leveraging the Zeroth-Order method, our work contributes to the advancement of optimization techniques in the context of complex language models.

연구 동기 및 목표

• 기하학적 파라미터 수가 막대한 대규모 LLM에서 백프로파게이션의 높은 계산 비용 문제를 해결하기 위해.
• 백프로파게이션을 통한 기울기 계산 없이도 소프트맥스 어텐션 기반 최적화를 효율적으로 수행하기 위해.
• 오직 정방향 전파 평가만을 사용해 기울기를 근사하는 제로스터드 방법을 개발하여 메모리 및 계산 오버헤드를 감소시키기 위해.
• 모델 파라미터와 손실 곡면에 대한 현실적인 가정 하에 제안된 알고리즘의 이론적 수렴 보장을 확립하기 위해.
• 기존의 백프로파게이션을 사용할 수 없는 대규모 LLM에 적합한 확장 가능한 최적화 프레임워크를 제공하기 위해.

제안 방법

• 동시 섭동 확률적 근사(Simultaneous Perturbation Stochastic Approximation, SPSA) 기반의 제로스터드 기울기 추정기법을 제안하며, 오직 정방향 손실 평가만을 사용한다.
• 편향 없는 기울기 추정을 위해, 변형된 파라미터와 원본 파라미터 간의 손실 차이를 활용한다: $ \widehat{g}(x_0)_i = \frac{1}{2\epsilon p_i}(L(x_0 + \epsilon p) - L(x_0 - \epsilon p)) $.
• 소프트맥스 어텐션 목적 함수 최적화를 위해 기울기 추정을 사영된 기울기 하강(GD) 프레임워크에 통합한다.
• 손실 함수 $ L(x) $ 는 정규화된 지수 출력과 목표 벡터 간의 제곱 오차로 정의되며, 어텐션을 소프트맥스 회귀 문제로 모델링한다.
• 기울기 추정의 분산을 제어하기 위해 농도 및 행렬 노름 경계를 사용하며, 부드러움과 PL 조건 하에서 수렴을 보장한다.
• 손실 감소의 반복 단위 기대치를 포함하는 재귀적 기대치 경계를 통해 수렴을 확립하며, $ \mu $-PL 조건과 헤시안 공분산의 트레이스를 활용한다.

실험 결과

연구 질문

• RQ1 제로스터드 최적화 방법이 대규모 언어 모델의 소프트맥스 어텐션에 대해 전역 수렴을 달성할 수 있는가?
• RQ2 모델 파라미터에 대한 현실적인 가정 하에 어텐션 기반 최적화의 샘플 복잡도는 어떻게 되는가?
• RQ3 전통적인 백프로파게이션과 비교해 본다면, 제안된 방법의 계산 효율성과 수렴 속도는 어떠한가?
• RQ4 고차원적이고 비볼록적인 어텐션 최적화 문제에서 기울기 추정에 대해 어떤 이론적 보장을 제공할 수 있는가?
• RQ5 조건 수 $ \kappa(A) $, 차원 $ d $, 배치 크기 $ B $ 와 같은 모델 파라미터가 수렴에 어떤 영향을 미치는가?

주요 결과

• 제안된 제로스터드 알고리즘은 $ O(M \cdot (1 + d^{1.5} \kappa^2(A)/k) \cdot \mu^{-2} B^{-1} \log((L(x_0)-L^*)/\epsilon)) $ 반복 내에 $ \epsilon $-최적 해로의 전역 수렴을 달성하며, 여기서 $ M = \exp(O(R^2 + \log n)) $ 이다.
• 이 방법은 $ \mu $-PL, $ l $-부드러움, 유한한 헤시안 효과적 랭크 조건 하에서 수렴을 보장하며, SPSA 편향을 통해 분산을 제어한다.
• 수렴 속도는 PL 상수 $ \mu $, 배치 크기 $ B $, 조건 수 $ \kappa(A) $ 의 역수에 의존하며, 더 큰 $ B $ 는 샘플 효율성을 향상시킨다.
• 알고리즘은 백프로파게이션을 완전히 회피하며 오직 정방향 전파 평가에 의존하므로, 블랙박스 또는 기울기 가림이 발생하는 LLM에 적합하다.
• 이론적 분석을 통해 기울기 추정기의 분산이 $ \operatorname{tr}[\Sigma(x)] \leq \alpha (L(x) - L^*) $ 로 유계임을 보여주며, 안정적인 하강을 가능하게 한다.
• 최종 수렴 경계는 $ \exp(O(R^2)) $ 의 비율로 척도가 조정되며, 여기서 $ R $ 은 파라미터 노름의 상한이다. 이는 최악의 경우 모델 크기에 대해 지수적 의존성을 나타낸다.