Top-K関数を用いたスパース化の実装
Q-Sparseが行う最も重要な操作は、入力テンソルにTop-K スパース化関数を適用することです。
具体的には、Transformerアーキテクチャは注意層とフィードフォワード層の両方でnn.Linear線形層(行列乗算)を用いて投影を行い、Y=X·W^Tと表現できます。(ここでXは入力テンソル、Wはその重み、Yは出力テンソルを表します)
Q-Sparseでは、入力活性化テンソルXに対して、まずその絶対値|X|を計算してソートし、絶対値が最大のK個の要素を見つけます。
ここでのKは事前に設定されたハイパーパラメータで、スパース化の程度を決定します。
その後、Q-SparseはXと同じ形状のバイナリマスクテンソルMを作成し、|X|の中で絶対値が最大のK個の要素に対応する位置のMの対応する位置を1に、残りの位置を0に設定します。
次に、入力テンソルXとマスクテンソルMのアダマール積(要素ごとの乗算)を行い、スパース化されたテンソルX_sparseを得ます。
順伝播過程では、スパース化されたテンソルX_sparseが元の入力テンソルXの代わりに後続の計算(行列乗算など)に参加します。
X_sparseの大部分の要素がゼロに設定されているため、計算量とメモリ帯域幅の要求を大幅に削減できます。
逆伝播過程では、Q-SparseはStraight-Through Estimator (STE)を使用してTop-K関数の勾配を計算します。
従来の学習方法では、通常、損失関数に対するネットワークパラメータの勾配を計算し、勾配降下法を用いてパラメータを更新して損失を最小化します。
しかし、ネットワークに量子化やTop-Kなどの微分不可能な操作が存在する場合、これらの操作の出力に対する入力の勾配がほとんどの点でゼロになるため、勾配の計算に問題が生じ、勾配が効果的に伝播できなくなります。
STEは勾配をスパース化前のテンソルに直接伝達することで、勾配消失の問題を回避します。
一般的な逆伝播では、損失関数Lのxに対する勾配∂L/∂x=∂L/∂y⋅∂y/∂xですが、微分不可能なため直接計算できません。
STEの解決策は、損失関数のスパース化テンソルyに対する勾配のみを計算し、それを元のテンソルxに直接コピーする、つまり∂L/∂yを∂L/∂xの推定値として直接使用することです。
フィードフォワード層では、Q-Sparseは通常のReLU活性化関数の代わりに二乗ReLU関数を使用し、二乗演算によって活性化のスパース性をさらに高めます(⊙はアダマール積を表します)。
また、量子化モデルに適応するため、Q-SparseはTop-Kスパース化を適用する前に入力テンソルを量子化し、スパース化操作が量子化表現と互換性があることを確認します。その関数表現は以下の通りです：
ここで、εは小さな定数で、分母がゼロになるのを避けるために使用されます。
特に、1ビット量子化の重みに対して、Q-Sparseは以下の量子化関数を使用します。ここでαは重みテンソルWの平均絶対値です。
60%の活性化パラメータで同等の効果を達成
比較実験によると、スパース率とモデルのパフォーマンスの両方において、Q-Sparseは以前のReLU手法を大幅に上回っています。
Q-Sparseの具体的な効果について、著者はゼロからの学習、継続学習、ファインチューニングの3つのタスクにおけるパフォーマンスを評価しました。
ゼロからの学習実験では、Llamaモデルを使用し、700Mと7Bモデルで70% top-K(つまり40%の全体的なスパース率)のQ-Sparseが密なベースラインと同等の学習損失を達成できることが示されました。
継続学習の目的は密なモデルをスパース化することで、ここでの実験対象はMistral-7Bでした。
結果として、活性化パラメータが2.9Bと3.8Bの場合、モデルのARC、MMLUなどのデータセットでのスコアに顕著な低下は見られませんでした。
ファインチューニング実験では、Qwen-7BとMistral-7Bの2つのモデルに対して、Q-Sparseは継続学習と同様の結果を示し、約60%の活性化パラメータで密なモデルと非常に近いパフォーマンスを実現しました。
これらの結果は、同じパフォーマンスの下で、密なモデルと比較して、スパース活性化モデルは推論プロセス中に活性化パラメータを大幅に削減し、それによってFLOPSの消費量を減らすことができることを意味します。
量子化モデルについては、チームは自社開発のBitNet b1.58モデルにQ-Sparseを適用し、複数のデータセットで学習と評価を行いました。
700Mと7Bの両方のスケールで、Q-Sparseを使用した量子化モデルの収束速度と最終的な損失関数値が、Q-Sparseを使用していない量子化モデル(BitNet b1.58)と同等であることが分かります。
これは、Q-Sparseが量子化モデルにシームレスに統合でき、モデルの学習と収束に大きな影響を与えないことを示しています。
これにより著者は、Q-Sparseと量子化技術を組み合わせることで、大規模言語モデルの推論段階の効率をさらに向上させることができると考えています。
推論最適化の新しい「スケーリング則」の発見
これらのモデルがスパース活性化を採用した際のパフォーマンスを評価するだけでなく、著者はモデルのパフォーマンス、規模、スパース率の3者間の関係を探究し、いくつかの新しい発見をしました。
スパース活性化モデルのパフォーマンススケーリング則: 著者は、密なモデルと同様に、スパース活性化モデルのパフォーマンスもべき乗則のスケーリング関係に従うことを発見しました。
具体的には、スパース率Sが与えられた場合、収束時のモデルの損失関数値L(N,S)は以下の式で近似できます：