清華大学の Zhu Jun 氏のチームによる新しい研究: 4 ビット整数を使用して Transformer をトレーニングすると、FP16 よりも 2.2 倍、35.1% 高速になり、AGI の到来が加速します。

アクティベーション、重み、勾配を 4 ビットに量子化すると、ニューラルネットワークのトレーニングが高速化されると期待されます。

ただし、既存の 4 ビットトレーニング方法では、最新のハードウェアではサポートされていないカスタム数値形式が必要です。

最近、清華大学の Zhu Jun 氏のチームは、INT4 アルゴリズムを使用してすべての行列乗算を実装する Transformer トレーニング方法を提案しました。

極めて低い INT4 精度でのトレーニングは非常に困難です。この目標を達成するために、研究者はTransformerの活性化と勾配の特定の構造を注意深く分析し、専用の量子化器を提案しました。

研究者らは、順方向伝播において外れ値の問題を特定し、外れ値を抑制するためにアダマール量子化器を提案した。

逆方向パスでは、ビット分割を提案することで勾配の構造的スパース性を活用し、分数サンプリング技術を使用して勾配を正確に量子化します。

この新しいアルゴリズムは、自然言語理解、機械翻訳、画像分類など、幅広いタスクで競争力のある精度を実現します。

プロトタイプの線形演算子は、同様の FP16 演算子よりも 2.2 倍高速に動作し、トレーニング速度は 35.1% 向上します。

写真

論文アドレス: https://arxiv.org/abs/2306.11987

コードアドレス: https://github.com/xijiu9/Train_Transformers_with_INT4

新しいINT 4トレーニングアルゴリズム

ニューラルネットワークのトレーニングには計算負荷がかかります。トレーニングに低精度演算を使用すると（完全量子化トレーニング/FQT）、計算効率とメモリ効率が向上することが期待されます。

FQT メソッドは、元の完全精度計算グラフにいくつかの量子化器と逆量子化器を追加し、より高価な浮動小数点演算をより安価な低精度浮動小数点演算に置き換えます。

FQT の研究は、収束速度や精度をあまり犠牲にすることなく、トレーニングの数値精度を下げることを目的としています。

必要な数値精度は FP16 から FP8、INT32+INT8、INT8+INT5 に削減されました。

FP8 トレーニングは、Transformer Engine を搭載した Nvidia H100 GPU に実装されており、大規模な Transformer のトレーニングを加速します。最近のトレーニングの数値精度は 4 桁に低下しました。

ただし、これらの 4 ビットトレーニングメソッドは、最新のハードウェアではサポートされていないカスタム数値形式を必要とするため、高速化に直接使用することはできません。

まず、順方向伝播における微分不可能な量子化器によって損失ランドスケープが凸凹になり、勾配ベースのオプティマイザーが簡単に局所最適状態に陥る可能性があります。

第二に、勾配は低い精度で近似的にしか計算されません。このような不正確な勾配は、トレーニングプロセスを遅くしたり、トレーニングが不安定になったり発散したりする原因となる可能性があります。

この研究では、研究者らはTransformer用の新しいINT4トレーニングアルゴリズムを提案しました。

写真

Transformer をトレーニングするためのすべての高価な線形演算は、行列乗算 (MM) の形式で記述できます。

この MM 形式により、Transformer のアクティベーション、重み、勾配の特定の構造を活用して、FP32 行列乗算をより適切に近似する、より柔軟な量子化器を設計できるようになります。

この量子化器は、乱数線形代数 (RandNLA) の分野における進歩を活用します。

研究者らは、フォワードパスの場合、アクティベーションの外れ値が精度低下の主な原因であることを発見した。

外れ値を抑制するために、彼らは活性化行列の変換バージョンを量子化するアダマール量子化器を提案しました。この変換は、外れ値に含まれる情報を行列の隣接するエントリに伝播し、外れ値の数値範囲を縮小するブロック対角アダマール行列です。

後方伝播では、活性化勾配の構造的スパース性を利用します。研究者たちは、いくつかのトークンの勾配が非常に大きいことを発見しました。同時に、他のほとんどのトークンの勾配は非常に均一であり、より大きな勾配の量子化された残差よりもさらに均一です。

写真

したがって、すべての勾配を計算するのではなく、より大きな勾配残差を計算して計算リソースを節約する方が適切です。

このスパース性を活用するために、研究者らは各トークンの勾配を上位 4 ビットと下位 4 ビットに分割するビット分割を提案しました。

次に、RandNLA の重要なサンプリング手法であるレバレッジスコアサンプリングを通じて、最も有益な勾配が選択されます。

写真

研究者らは、順方向および逆方向の伝播量子化技術を組み合わせて、Transformer 内のすべての線形演算を INT4MM を使用して実行するアルゴリズムを提案し、自然言語理解、質問応答、機械翻訳、画像分類などのさまざまなタスクで Transformer をトレーニングするためのアルゴリズムを評価しました。

彼らのアルゴリズムは、既存の 4 ビットトレーニングアルゴリズムと比較して、競争力のある、あるいはそれ以上の精度を実現します。

さらに、このアルゴリズムは、FP4 や対数形式などのカスタム数値形式を必要としないため、GPU などの最新のハードウェアと互換性があります。

このプロトタイプの量子化 + INT4 MM 演算子の実装は、FP16MM ベースラインよりも 2.2 倍高速で、トレーニング速度が 35.1% 向上します。