[論文レビュー] Training Multi-Layer Over-Parametrized Neural Network in Subquadratic Time

Zhao Song, Lichen Zhang+1 more|arXiv (Cornell University)|Dec 14, 2021

Advanced Neural Network Applications被引用数 6

ひとこと要約

この論文は、構造化された重み行列と効率的なデータ構造を活用することで、1イテレーションあたりの計算量を準平方根時間オーダーで学習可能な、深く過パラメータ化されたニューラルネットワークのための新規フレームワークを提示する。初期化段階の前処理と適応的勾配計算を用いることで、1イテレーションあたりのコストを O(m²⁻Ω(1)) に抑え、標準的な O(m²) より顕著に低減し、大規模言語モデル（LLM）の高速微調整を可能にする。

ABSTRACT

We consider the problem of training a multi-layer over-parametrized neural network to minimize the empirical risk induced by a loss function. In the typical setting of over-parametrization, the network width $m$ is much larger than the data dimension $d$ and the number of training samples $n$ ($m=\mathrm{poly}(n,d)$), which induces a prohibitive large weight matrix $W\in \mathbb{R}^{m imes m}$ per layer. Naively, one has to pay $O(m^2)$ time to read the weight matrix and evaluate the neural network function in both forward and backward computation. In this work, we show how to reduce the training cost per iteration. Specifically, we propose a framework that uses $m^2$ cost only in the initialization phase and achieves \emph{a truly subquadratic cost per iteration} in terms of $m$, i.e., $m^{2-Ω(1)}$ per iteration. Our result has implications beyond standard over-parametrization theory, as it can be viewed as designing an efficient data structure on top of a pre-trained large model to further speed up the fine-tuning process, a core procedure to deploy large language models (LLM).

研究の動機と目的

• 大規模な m×m 重み行列に起因する、学習における禁止的とされる O(m²) の1イテレーションあたりのコストを解消すること。
• 訓練コストを m に関して真正に準平方根時間、すなわち O(m²⁻Ω(1)) に低減する手法の開発。
• 過パラメータ化と構造化計算を活用することで、大規模言語モデル（LLMs）の効率的微調整を実現すること。
• 従来の手法が O(nm²) のコストを負担するか、入力次元 d に対して指数関数的に依存するという制限を克服すること。
• 事前学習済みモデルの上にデータ構造を設計し、微調整の高速化を図るとともに、収束保証を維持すること。

提案手法

• ニューロン活性化におけるスパarsity を誘発するため、シフトされた ReLU 活性化関数を用いることで、1層あたりの有効計算量を削減する。
• 活性化されたニューロン出力の分布とそのノルムを、切断されたガウス確率変数を用いてモデル化する。
• 切断されたカイ二乗分布およびサブガウス分布に対する集中不等式を適用し、層間におけるノルムのフラクチュエーションを制限する。
• 初期化段階で一度だけ O(m²) のコストを負担する前処理フェーズを設計し、以降のすべての訓練イテレーションで準平方根時間のコストを達成する。
• すべてのデータポイントと層に対してユニオンバウンドを適用することで、ネットワーク全体における高確率でのノルム安定性を保証する。
• 事前に計算された統計的性質に基づいて活性化されたニューロンをインデックス化するデータ構造を採用し、高速なフォワードおよびバックワードパスを実現する。

実験結果

リサーチクエスチョン

• RQ1 m×m 重み行列を有する深く過パラメータ化されたニューラルネットワークの学習において、1イテレーションあたりの計算量を準平方根時間オーダーに達成できるか？
• RQ2 過パラメータ化と活性化スパarsity を活用することで、標準的な O(m²) の壁を越えて計算量を低減できるか？
• RQ3 入力次元 d に指数関数的に依存しない形で、O(m²⁻Ω(1)) の1イテレーションあたりコストを実現する前処理スキームを設計できるか？
• RQ4 理論的な準平方根時間コストが、大規模言語モデルの微調整において実際の場面でどの程度実現可能か？
• RQ5 ランダム初期化下でも、層間におけるノルム安定性を維持しつつ、準平方根時間の学習コストを達成することは可能か？

主な発見

• 提案フレームワークは、深く過パラメータ化されたネットワークの学習において、1イテレーションあたり O(m²⁻Ω(1)) のコストを達成し、標準的な O(m²) より顕著に低減している。
• この手法は、初期化段階でのみ O(m²) のコストを負担し、以降のすべての訓練イテレーションは準平方根時間で実行される。
• 高確率で、すべての層およびデータポイントにおいて隠れ表現の ℓ₂ ノルムが [1−ε, 1+ε] の範囲内に保たれ、安定した学習が保証される。
• このフレームワークは、深さ L ≥ 2 かつ幅 m = poly(n, d) を満たすネットワークに適用可能であり、典型的な過パラメータ化条件を満たす。
• 過パラメータ化された構造を効果的に活用することで、大規模言語モデル（LLMs）の高速微調整が可能になる。
• 分析は、切断ガウス分布およびカイ二乗分布に対する集中不等式に依拠しており、ノルム安定性に対する理論的保証を確立している。