[논문 리뷰] A Fast Optimization View: Reformulating Single Layer Attention in LLM Based on Tensor and SVM Trick, and Solving It in Matrix Multiplication Time

Yeqi Gao, Zhao Song외 2명|arXiv (Cornell University)|2023. 09. 14.

Topic ModelingComputer Science인용 수 3

한 줄 요약

이 논문은 텐서 대수와 SVM 유사 희소화 기법을 활용해 단일 레이어 자기주의(attention)의 최적화 프레임워크를 재구성함으로써, 행렬 곱셈 시간 내에 수렴하는 새로운 최적화 프레임워크를 제안한다. 이는 $\widetilde{O}((\mathcal{T}_{\mathrm{mat}}(n,n,d) + \mathcal{T}_{\mathrm{mat}}(n,d,d) + d^{2\omega})\log(1/\epsilon))$ 시간 내에 $\epsilon$-정확도를 확보하는 증명 가능 보장을 제공한다.

ABSTRACT

Large language models (LLMs) have played a pivotal role in revolutionizing various facets of our daily existence. Solving attention regression is a fundamental task in optimizing LLMs. In this work, we focus on giving a provable guarantee for the one-layer attention network objective function $L(X,Y) = \sum_{j_0 = 1}^n \sum_{i_0 = 1}^d ( \langle \langle \exp( \mathsf{A}_{j_0} x ) , {\bf 1}_n angle^{-1} \exp( \mathsf{A}_{j_0} x ), A_{3} Y_{*,i_0} angle - b_{j_0,i_0} )^2$. Here $\mathsf{A} \in \mathbb{R}^{n^2 imes d^2}$ is Kronecker product between $A_1 \in \mathbb{R}^{n imes d}$ and $A_2 \in \mathbb{R}^{n imes d}$. $A_3$ is a matrix in $\mathbb{R}^{n imes d}$, $\mathsf{A}_{j_0} \in \mathbb{R}^{n imes d^2}$ is the $j_0$-th block of $\mathsf{A}$. The $X, Y \in \mathbb{R}^{d imes d}$ are variables we want to learn. $B \in \mathbb{R}^{n imes d}$ and $b_{j_0,i_0} \in \mathbb{R}$ is one entry at $j_0$-th row and $i_0$-th column of $B$, $Y_{*,i_0} \in \mathbb{R}^d$ is the $i_0$-column vector of $Y$, and $x \in \mathbb{R}^{d^2}$ is the vectorization of $X$. In a multi-layer LLM network, the matrix $B \in \mathbb{R}^{n imes d}$ can be viewed as the output of a layer, and $A_1= A_2 = A_3 \in \mathbb{R}^{n imes d}$ can be viewed as the input of a layer. The matrix version of $x$ can be viewed as $QK^ op$ and $Y$ can be viewed as $V$. We provide an iterative greedy algorithm to train loss function $L(X,Y)$ up $ε$ that runs in $\widetilde{O}( ({\cal T}_{\mathrm{mat}}(n,n,d) + {\cal T}_{\mathrm{mat}}(n,d,d) + d^{2ω}) \log(1/ε) )$ time. Here ${\cal T}_{\mathrm{mat}}(a,b,c)$ denotes the time of multiplying $a imes b$ matrix another $b imes c$ matrix, and $ω\approx 2.37$ denotes the exponent of matrix multiplication.

연구 동기 및 목표

• LLM의 단일 레이어 어텐션 최적화를 위한 증명 가능 효율적인 방법을 제공함으로써, 훈련 및 추론에서의 계산 병목 현상을 해결한다.
• 텐서 대수와 크로네cker乘법을 활용해 어텐션 목적함수를 재구성함으로써 더 빠른 계산을 가능하게 한다.
• 증명 가능 수렴 보장을 갖는 탐욕적 반복 알고리즘을 활용해 행렬 곱셈 시간 내에 수렴을 달성한다.
• 고급 희소화 기법을 통해 어텐션 회귀의 계산 복잡도를 이차원에서 근사 최적의 행렬 곱셈 시간으로 감소시킨다.
• 안정적인 수렴을 보장하는 $(l,M)$-좋은 손실 함수와 잘 초기화된 점을 갖춘 이론적 프레임워크를 구축한다.

제안 방법

• 크로네cker 곱을 활용해 단일 레이어 어텐션 목적함수를 행렬 최적화 문제로 재구성함: $\mathsf{A} = A_1 \otimes A_2 \in \mathbb{R}^{n^2 \times d^2}$, 이로 인해 텐서 기반 계산이 가능해진다.
• TensorSketch를 통한 SVM 기법의 텐서 버전을 적용해 어텐션 계산을 희소화함으로써, 정확도를 유지하면서도 효과적 차원을 감소시킨다.
• 근사 헤시안 역행렬을 활용한 탐욕적 반복 알고리즘을 사용해 $X$와 $Y$의 파라미터를 갱신함으로써, $(l,M)$-좋은 손실 조건 하에서 수렴을 보장한다.
• 행렬 곱셈 기반 프레임워크를 활용하여 시간 복잡도가 $\mathcal{T}_{\mathrm{mat}}(a,b,c)$로 제한되며, 이는 $a\times b$와 $b\times c$ 행렬을 곱하는 데 걸리는 시간이다.
• 초기값 조건으로서 $r_0 M \leq 0.1l$을 도입함으로써, 양의 정부호 헤시안을 갖는 최소값으로 국소 수렴을 보장한다.
• 레마 14.3를 통한 반복적 축소를 활용해 각 단계에서 최적해와의 거리가 기하급수적으로 감소함을 보장하며, $r_{t+1} \leq 0.4 r_t$를 확보한다.

실험 결과

연구 질문

• RQ1 LLM의 단일 레이어 어텐션 최적화 문제를 텐서 대수를 활용해 재구성하여 더 빠른 계산을 가능하게 할 수 있는가?
• RQ2 SVM 유사 희소화 기법을 조정하여 어텐션 계산의 복잡도를 낮추면서도 정확도를 유지할 수 있는가?
• RQ3 어떤 보장이 있더라도 어텐션 회귀 문제에 대해 행렬 곱셈 시간 내에 수렴을 달성할 수 있는가?
• RQ4 손실 함수에 어떤 조건이 요구되어 어텐션 기반 최적화에서 안정적이고 빠른 수렴이 보장되는가?
• RQ5 알고리즘 재구성과 희소화를 통해 어텐션 레이어의 계산 병목 현상을 어떻게 완화할 수 있는가?

주요 결과

• 제안된 알고리즘은 $\widetilde{O}((\mathcal{T}_{\mathrm{mat}}(n,n,d) + \mathcal{T}_{\mathrm{mat}}(n,d,d) + d^{2\omega})\log(1/\epsilon))$ 시간 내에 단일 레이어 어텐션 최적화 문제를 해결하며, $\epsilon$-정확도를 달성한다.
• 표준 가정 하에 손실 함수가 $(l,M)$-좋음으로 증명되어, 최소점에서 양의 정부호 헤시안을 갖는 국소 수렴이 보장된다.
• 알고리즘이 기하급수적 수렴을 보장함: $r_{t+1} \leq 0.4 r_t$이며, 반복 과정 전반에 걸쳐 $M \cdot r_t \leq 0.1l$이 유지된다.
• TensorSketch 기반 희소화를 활용함으로써, 전체 $n^2 \times n^2$ 행렬의 명시적 계산을 피할 수 있어 메모리 및 시간 비용을 감소시킨다.
• 표준 어텐션 계산보다 증명 가능하게 빠른 성능 향상을 이끌어내며, 비용이 많이 드는 소프트맥스 및 행렬 역행렬 연산의 의존도를 감소시킨다.
• 이 방법은 표준 LLM 구성 요소에 그대로 적용 가능하다: $X$는 $QK^T$에 대응하고, $Y$는 $V$에 대응하며, $B$는 타겟 출력에 대응하므로 트랜스포머 훈련에 직접적으로 관련된다.