[論文レビュー] Attention Scheme Inspired Softmax Regression

Yichuan Deng, Zhihang Li+1 more|arXiv (Cornell University)|Apr 20, 2023

Machine Learning and Algorithms被引用数 6

ひとこと要約

本稿では、大規模言語モデルにおけるアテンションメカニズムにインspiredされた、新たなソフトマックス回帰フレームワークを提案する。グリーディ最適化アルゴリズムを提案し、収束保証を理論的に示す。弱い条件下でも線形収束を確立し、実用的なソフトマックスベースのモデル学習における反復的手法の理論的裏付けを提供する。

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.

研究の動機と目的

• 大規模言語モデルにおけるアテンションメカニズムにインspiredされた、指数関数による正規化を組み込んだ新たなソフトマックス回帰問題を形式化すること。
• 特に大規模言語モデルの文脈において、ソフトマックス関数の学習にグリーディアルゴリズムを用いることの理論的妥当性を裏付けること。
• 行列ノルムと正則化に関する現実的な仮定の下で、ソフトマックス回帰目的関数に対する反復アルゴリズムの収束挙動を分析すること。
• 各反復で多項式時間計算量を確保する、理論的に収束することが保証されたアルゴリズムを確立すること。これは大規模応用に適している。
• ソフトマックス操作に内在する正規化因子を含めた、ハイパボリック回帰に関する先行研究の拡張すること。

提案手法

• 正規化された指数関数 $\exp(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 大規模言語モデルのアテンションメカニズムにおける中心的役割を果たすソフトマックス回帰目的関数を最適化するグリーディアルゴリズムが、理論的に正当化可能か？
• 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)$ であり、対数的収束速度を保証する。
• 1反復あたりの時間コストは $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$ でアルゴリズムが成功する。