[논문 리뷰] Attention Scheme Inspired Softmax Regression

Yichuan Deng, Zhihang Li외 1명|arXiv (Cornell University)|2023. 04. 20.

Machine Learning and Algorithms인용 수 6

한 줄 요약

이 논문은 대규모 언어 모델의 어텐션 메커니즘을 영감으로 삼아 새로운 소프트맥스 회귀 프레임워크를 제안하며, 증명 가능하게 수렴 보장을 갖춘 탐욕적 최적화 알고리즘을 제안한다. 미묘한 조건 하에서 선형 수렴를 확립하며, 실용적으로 소프트맥스 기반 모델을 반복적 방법으로 훈련하는 데 이론적 근거를 제공한다.

ABSTRACT

Large language models (LLMs) have made transformed changes for human society. One of the key computation in LLMs is the softmax unit. This operation is important in LLMs because it allows the model to generate a distribution over possible next words or phrases, given a sequence of input words. This distribution is then used to select the most likely next word or phrase, based on the probabilities assigned by the model. The softmax unit plays a crucial role in training LLMs, as it allows the model to learn from the data by adjusting the weights and biases of the neural network. In the area of convex optimization such as using central path method to solve linear programming. The softmax function has been used a crucial tool for controlling the progress and stability of potential function [Cohen, Lee and Song STOC 2019, Brand SODA 2020]. In this work, inspired the softmax unit, we define a softmax regression problem. Formally speaking, given a matrix $A \in \mathbb{R}^{n imes d}$ and a vector $b \in \mathbb{R}^n$, the goal is to use greedy type algorithm to solve \begin{align*} \min_{x} \| \langle \exp(Ax), {\bf 1}_n angle^{-1} \exp(Ax) - b \|_2^2. \end{align*} In certain sense, our provable convergence result provides theoretical support for why we can use greedy algorithm to train softmax function in practice.

연구 동기 및 목표

• 대규모 언어 모델의 어텐션 메커니즘을 영감으로 삼아 지수 함수를 통한 정규화를 포함하는 새로운 소프트맥스 회귀 문제를 체계화한다.
• 실제로 소프트맥스 함수를 훈련하기 위해 탐욕적 알고리즘을 사용하는 데 이론적 근거를 제공한다. 특히 LLM 맥락에서의 적용을 고려한다.
• 행렬 노름과 정규화에 대한 현실적인 가정 하에서 소프트맥스 회귀 目표 함수에 대한 반복 알고리즘의 수렴 행동을 분석한다.
• 각 반복 단계에서 다항 시간 복잡도를 가지며, 대규모 응용에 적합한 증명 가능 수렴 알고리즘을 확립한다.
• 기존의 쌍곡선 회귀 연구를 확장하여, 소프트맥스 연산에 내재된 정규화 요소를 포함한다.

제안 방법

• 소프트맥스 회귀 문제를 $Ax$의 정규화된 지수 함수와 목표 벡터 $b$ 사이의 제곱 L2 거리 최소화 문제로 체계화하며, $\|\operatorname{diag}(w)Ax\|_2^2$를 포함한 정규화된 목표 함수를 제안한다.
• 각 단계에서 해를 갱신하기 위해 부분 샘플링된 대각 행렬을 통한 헤시안 근사치를 사용하는 탐욕적 반복 알고리즘(알고리즘 1)을 도입한다.
• 반복 계산 비용을 줄이면서도 수렴 보장을 유지하기 위해 부분 샘플링 기법(SubSample)을 활용한다.
• 근사 헤시안 $\widetilde{H} = A^T \widetilde{D} A$를 사용하는 선형 탐색 유사 갱신 규칙을 적용하며, 여기서 $\widetilde{D}$는 현재 기울기의 대각 행렬의 부분 샘플링 버전이다.
• 오차 $\|x_t - x^*\|_2$ 에 대한 귀납법을 통해 수렴성을 확립하며, 특정 조건 하에서 반복마다 수축 인자가 0.4임을 보인다.
• 헤시안의 리프시츠 연속성과 정규화 항 $\langle \exp(Ax), \mathbf{1}_n \rangle$ 의 하한을 이용해 오차 전파를 제어한다.

실험 결과

연구 질문

• RQ1 LLM의 어텐션 메커니즘에서 중심적인 역할을 하는 소프트맥스 회귀 목표 함수를 최적화하기 위해 탐욕적 알고리즘이 이론적으로 정당화될 수 있는가?
• RQ2 데이터 행렬 $A$, 목표 벡터 $b$, 정규화 가중치 $w$ 에 어떤 조건이 성립하면 반복 알고리즘이 선형 수렴를 보장하는가?
• RQ3 고차원 환경에서 수렴 보장을 유지하면서 헤시안 근사치를 효율적으로 계산할 수 있는 방법은 무엇인가?
• RQ4 수렴 속도는 행렬 $A$ 의 스펙트럼 노름, 최소 특이값 $\sigma_{\min}(A)$, 정규화 강도에 어떻게 의존하는가?
• RQ5 정규화 인자 $\langle \exp(Ax), \mathbf{1}_n \rangle^{-1}$ 는 최적화 과정의 안정성과 수렴성에 어떤 영향을 미치는가?

주요 결과

• 제안된 알고리즘은 제시된 가정 하에서 반복마다 수축 인자가 0.4인 선형 수렴를 달성한다.
• $\epsilon$-정확도에 도달하기 위한 반복 수는 $T = \log(\|x_0 - x^*\|_2 / \epsilon)$ 로 주어지며, 이는 로그 수렴 속도를 보장한다.
• 반복당 시간 복잡도는 $O((\operatorname{nnz}(A) + d^\omega) \cdot \operatorname{poly}(\log(n/\delta)))$ 이며, 여기서 $\omega \approx 2.373$ 은 행렬 곱셈 지수이다.
• 제시된 조건 하에서 손실 함수의 헤시안은 양의 정부호이자 리프시츠 연속임을 입증한다.
• 정규화 항 $\langle \exp(Ax), \mathbf{1}_n \rangle$ 에 대해 $\exp(-R^2)$ 의 하한이 확립되며, 이는 오차 제어에 핵심적이다.
• 알고리즘은 $T$ 반복 동안 높은 확률 $1 - \delta$ 로 성공하며, 유니온 바운드와 통제된 실패 확률을 갖는 부분 샘플링을 사용한다.