ByteDance は推論コストを大幅に削減するまったく新しいアーキテクチャを導入しました。

なんと冷酷なのでしょう？

推論速度はMoEアーキテクチャより2〜6倍速く、推論コストは最大83％削減できます。

このまったく新しいスパース モデル アーキテクチャはUltraMemと呼ばれ、現在主流の MoE および PKM アーキテクチャの制限を効果的に解決します。

たとえば、MoE が推論を実行する場合、バッチ サイズが小さいとすべてのエキスパートがアクティブ化され、メモリ アクセスと推論の待ち時間が急激に増加します。PKM はメモリ アクセスのオーバーヘッドを削減しますが、その効果は低く、スケーラビリティも制限されます。

実験結果によると、 2,000 万の値でトレーニングされた UltraMem モデルは、同じコンピューティング リソースで業界をリードする推論速度とモデル パフォーマンスを実現でき、数十億の値やエキスパートを構築するための新しい道が開かれます。

この研究はICLR 2025に採択されました。

では、UltraMem アーキテクチャはどのようにしてこれを実現するのでしょうか? 続きを読みましょう。

メモリアクセスとモデルパフォーマンスのバランスをとる

大規模言語モデルの機能が向上するにつれて、必要なコンピューティング リソースが飛躍的に増加し、リアルタイム アプリケーションなどのリソースが制限された環境では大きな課題が生じます。

コンピューティング リソースのボトルネックを解決するために、研究者は MoE やプロダクト キー メモリ (PKM) などのさまざまなソリューションを提案してきました。

ただし、これらの方法にはそれぞれ制限があります。

MoE は、スパース アクティベーション エキスパートを使用して計算をパラメータから分離し、トレーニング中の計算を効果的に削減します。ただし、MoE のパフォーマンスは推論シナリオでは満足できるものではありません。

モデルは推論中に単語ごとに出力を生成する必要があるため、バッチ サイズとシーケンスの長さは通常小さく、一般に、少数のトークンでほぼすべてのエキスパートをアクティブ化できます。

この完全に専門家によって有効化されたモードでは、メモリ アクセス要件が急増し、深刻なメモリ アクセスのボトルネックが発生し、最終的には推論の待ち時間が大幅に増加します。

一方、PKM は、多数のスパース パラメーター値 (各値は実際にはベクトル) を含む大規模なメモリ レイヤーの概念を提案しました。

推論中、各トークンは「行ルーティング」と「列ルーティング」を通じて最高スコアの値に配置され、メモリ層の出力としてそれらに対して加重合計プーリングが実行されます。

PKMは推論中に各トークンがごく少数の値のみをアクティブ化するため、メモリアクセスのボトルネックを効果的に回避できます。しかし、優れたメモリアクセス効率にもかかわらず、PKMはモデル性能の面で劣っており、スケーリング能力も限られているため、大規模モデルのニーズを満たすことは困難です。

要約すると、MoE と PKM は計算リソースの問題をある程度軽減しましたが、推論効率、モデル パフォーマンス、スケーラビリティの点で依然として重大な欠点があり、大規模モデルの推論パフォーマンスをさらに最適化するための新しいソリューションが緊急に必要とされています。

一方、UltraMem は PKM の設計に基づいていますが、PKM の 3 つの欠点を解消し、より効率的なメモリ アクセス、より優れた値の取得を実現すると同時に、ビデオ メモリと展開コストを削減します。

1. モデル構造を最適化する

PKMの設計では、Transformer全体の中間層に挿入されたメモリ層が1つしかありません。これは大規模な学習には適しておらず、各残差接続には可能な限り多くのスパースパラメータが関与する必要があります。

そこで研究チームは、メモリ層を複数の小さなメモリ層に分割し、それらを一定の間隔でトランスフォーマー層に分散させ、現在のメモリ層の出力が後のトランスフォーマー層の出力に追加されるというスキップ層操作を追加しました。

これにより、モデルはメモリ層でのメモリアクセス操作とトランスフォーマー層での計算を並行して実行できるようになります。

2. 値取得方法を最適化する

検索時には、スコアが最も高いm値のみがアクティブ化されます。PKMのスコアは、「行スコア」と「列スコア」を組み合わせて得られます。

研究チームはさらに、より複雑な乗算方法である「Tucker Decomposed Query-Key Retrieval (TDQKR)」を調査しました。

この手法はタッカー分解に着想を得ています。具体的には、形状(n, n, h)を持つ値（hは隠れたサイズ）が与えられた場合、値のスコアS_gridは以下のように分解できます。

ここで、Srow、Scol∈Rr×n、C∈Rr×rは学習可能なタッカーコアです。この構造では、各値のスコアはr行r列のスコアの乗算と加算によって得られるため、複雑度が高くなります。

3. 暗黙的に拡張されたスパースパラメータ

通常、パラメータを多くするとパフォーマンスは向上しますが、パラメータが多すぎるとメモリやデプロイメントに問題が発生する可能性があります。

これに対処するため、研究チームはスパースパラメータを暗黙的に拡張する暗黙値拡張（IVE）法を提案し、仮想メモリと物理メモリの概念を導入しました。

4倍拡張を例に挙げると（下図参照）、仮想メモリの容量は物理メモリの4倍になります。複数のペア（スコア、インデックス）が与えられた場合、まず仮想メモリアドレステーブルを参照して検索を行います。4つの仮想ブロックは同じ物理メモリテーブルを参照します。次に、それぞれが重み付け合計プーリングを実行し、異なる線形層を通過します。最後に合計を計算し、結果を出力します。

最終的なLinearと値の取得の間には非線形演算がないため、各Linearを物理メモリテーブルとマージして、完全に新しいメモリテーブルを生成できます。この例では、値の数は暗黙的に4倍に拡張されます。

MoEより最大6倍高速

1. モデル性能評価

研究チームは、151M、680M、1.6Bの3つのサイズで活性化パラメータに関する広範な実験を実施しました。MoE、PKM、UltraMemの合計スパースパラメータは、活性化パラメータの12倍のままでした。

下の表に示すように、UltraMem は 680M および 1.6B で大幅なパフォーマンス上の利点があります。

スパースパラメータが増加すると、UltraMem のパフォーマンスと推論速度はどのように変化しますか?

• 下の図(b)は、UltraMemのパフォーマンスの変化を示しています。横軸はスパースパラメータと密パラメータの比率を表し、色付きの線はそれぞれスパース性のレベルを表しています。スパース性は、値の数/各トークンによってアクティブ化された値の数として定義されます。観察結果から、スパースパラメータを継続的に増加させることは、損失の減少と対数関係にあることがわかります。また、スパース性が小さいほど、モデルのパフォーマンスは向上します。ただし、スパース性を継続的に減少させることの利点は徐々に飽和します。
• 下の図(c)は、UltraMemの推論時間の変化を示しています。横軸はスパースパラメータと密パラメータの比率を表しています。スパースパラメータの数が増えてもUltraMemの推論時間はほぼ一定であるのに対し、MoEは顕著な増加傾向を示していることがわかります。

2. アブレーションテスト

研究チームは、1億5100万回のアクティベーションと合計15億のパラメータを持つスパースモデルで包括的なアブレーション実験を実施しました。

オリジナルの PKM から始めて、上で提案したいくつかのトリックと構造的な改善を徐々に追加することで、スパース パラメーターと計算コストをほとんど変えずに、最終的に C4 検証損失で-0.092という大幅な向上を達成できます。

まとめると、研究チームが提案した UltraMem はメモリアクセスが非常に少ないため、MoE と比較して最大 6 倍の速度向上を実現し、推論コストを最大 83% 削減します。

一方、パフォーマンスの面では、モデル容量の増加により、UltraMem は同じパラメータと計算負荷で MoE を上回り、より強力なスケーラビリティを備えていることがわかります。

この研究は、より効率的でスケーラブルな言語モデルを開発するための有望な方向性を示していると言えます。

論文の宛先:
https://arxiv.org/abs/2411.12364