ターゲットを絞った進化による自動アーキテクチャ合成

モデルアーキテクチャ設計はAIの基本的な柱であり、スケーリング機能や効率性から、事前トレーニング、調整、推論の基礎まで、あらゆるものを形作ります。アーキテクチャ設計における重要な課題は、AI システムをさまざまな環境に効果的に導入できるように、品質とハードウェアの制約 (特にレイテンシとメモリコスト) とのバランスを取ることです。

性能の高いアーキテクチャを設計することは、GPU上でのみ高速に動作するように最適化されたモデルに検索を制限したとしても、非常に重要な組み合わせ問題であることに変わりはありません。手動ヒューリスティックではパフォーマンスのトレードオフを予測するには限界があるため、この複雑さから、AI ラボや企業は特定の設計に早期から取り組むことが多くなっています。さらに、要件はアプリケーションドメインによって大きく異なります。たとえば、エッジユースケース向けに設計された言語モデルは、メモリフットプリントとエネルギー消費量が少なく、特定のターゲットハードウェアで優れたパフォーマンスを発揮する必要があります（CPUでの高速プリフィリングなど）。一方、クラウド使用向けに設計された言語モデルは通常、メモリフットプリントよりも品質とGPUレイテンシを優先します。実際には、要求はさらに微妙で複雑になります。

アーキテクチャのパフォーマンスの基盤は、その構築元となる計算ユニットと、これらがどのように相互に接続されているかによって決まります。現在のディープラーニングアーキテクチャの大半は、アテンション演算子とゲート付き線形単位の逐次インターリーブによって構築されています。これらの層は、私たちが線形と呼ぶ、はるかに大きなクラスの計算単位の単純な例です。 入力可変システム (LIV)。LIV は構造化演算子であり、その動作は入力自体によって点ごとに変調され、アテンションバリアント、線形アテンション、(ゲーテッド) 畳み込み、線形状態遷移を伴う (ゲート) 再帰、状態空間層、(ゲート) 線形ユニットなど、さまざまなクラスの計算単位を一般化するための抽象化を提供します。

「ビヨンド・トランスフォーマー」アーキテクチャの可能性を示す証拠はすでに存在しています。ゲーテッドコンボリューションやリカーレンスなどの単純なLIVと、縞模様のハイブリッドパターンにおける自己注意を組み合わせたモデルでは、適度な品質の向上、より長いシーケンスへのスケーリングの効率化、推論の高速化が実証されました（今年初めに報告しました）。 ハイブリッドの最初のスケーリング法則)。

新しい計算ユニットと相互接続戦略は、モデル性能の新たなフロンティアを実現する一方で、可能な設計の数が非常に多いため、アーキテクチャ設計にとって重要な課題でもあります。計算ユニットの特定のサブクラスに適用される手動の最適化やヒューリスティック（注意や畳み込みなど）や相互接続戦略（ウェイトシェアリング、KV共有、並列相互接続）に頼るのではなく、LIVに合わせた進化的アルゴリズムを活用しています。

STARのコアイノベーションの1つは、モデルアーキテクチャをSTARゲノムと呼ばれる階層的な数値シーケンスとして表現することです。これは、進化的最適化の原理を使用して進化したものです。このプロセスは反復的に行われます。ゲノムを具体的なアーキテクチャにコンパイルして評価し、次に最もパフォーマンスの高いアーキテクチャを選択して組み換え、次世代のアーキテクチャを作成します。

重要なのは、進化のプロセスは静的な目標と動的な目標の両方によって導かれるということです。静的目標は、パラメータ数やキャッシュサイズなど、アーキテクチャの特定の構成によって決まります。一方、動的な目標では、特定のデータセットでトレーニングした後の複雑さや、ターゲットハードウェアでのレイテンシーを測定するなどして、アーキテクチャを評価する必要があります。

建築候補が斬新で高性能であることを保証するために、 ゲノムエンコーディングは、ニューラルネットワークの新しい一般的な計算単位である線形入力変化システム（LIV）を導入した設計理論に基づいています。。ディープラーニングの最新の計算ユニットが入力コンテキストに基づいて計算をどのように変調するかを制御する基本的なメカニズムを特定し、抽象化しました。テンソルネットワークとシステム理論に基づいたこのフレームワークは、LIVを2つの重要な側面から特徴づけます。 構造 (オペレータのトークンとチャネルのミキシング構造) と フィーチャライゼーション (オペレータにおける入力依存の関数形式)。高度なアーキテクチャ設計を表現できるようにするため、このフレームワークでは演算子の構成を第一級の問題として扱い、レイヤーの連続的な積み重ねを超えたアーキテクチャ構築の新しい道筋を開いています。設計理論の詳細については後ほど説明します。

STARゲノムにより、LIVデザイン空間を進化的最適化に適した階層的な数値エンコーディングにマッピングできます。これにより、エンコードされたアーキテクチャで採用されている各計算ユニットの特性と、これらのユニットの相互接続方法が定義されます。

私たちはまず、改良された言語モデリングアーキテクチャの設計においてSTARを評価し、3つの目的が混在する場合の最適化を行いました。i.) 品質 （トレーニング後の困惑）、ii。） 品質とパラメータ効率、およびIII。) 品質とキャッシュ効率。

わずか2、3ラウンドの進化の後、 ほとんどのアーキテクチャは、Transformerや強力なハイブリッドベースラインなどの定番アーキテクチャよりも優れていますより多くのラウンドが実行されるにつれて、一貫した改善が行われています。特に、品質のみを最適化する場合、評価対象のSTARが進化したアーキテクチャはすべて、下流の評価ベンチマークで注意反復ハイブリッドよりも優れており、トランスフォーマーよりもハイブリッドのベンチマークの2倍のベンチマークで改善が見られました。この結果は、我々の設計空間における進化的探求の有効性を裏付ける強力な証拠であることがわかった。ハイブリッドは多大なリソースを投資して手作業で設計および改良されてきましたが、 STARは1日足らずでアーキテクチャを生成でき、ヒット率は 90% を超えています。

検索はサポートできます 多目的最適化。品質とモデルサイズを共同で最適化すると、進化したアーキテクチャはトランスフォーマーとストライプハイブリッドの両方よりも常に優れていると同時に、パラメーター数を減らすことができます。これにより、エッジ環境やリソースに制約のある環境向けにモデルをコンパクト化できます。進化と評価の結果をスケールを超えて伝えるためにさまざまな方法を試しましたが、一般的には、アーキテクチャを最適化するよりも、薄くて深いアーキテクチャの候補で最適化する方が良い伝達が得られることがわかりました。 モチーフ 目標幅で。

同様に、推論スタックで直接プロファイリングすることで得られるターゲットハードウェアの品質、モデルサイズ、レイテンシのバランスを取るために STAR を使用しました。これが可能なのは以下の場合です。 STAR はメトリックのグラデーションを必要としません。以下と互換性があります 静的指標と動的指標の混合 ゲノムからコンパイルされたアーキテクチャで計算され、 詳細なプロファイリングパスを含む レイテンシーと通信オーバーヘッドを最小限に抑えます。

STARは、特定の目的に合わせてアーキテクチャを最適化するだけでなく、進化の過程で繰り返し現れるアーキテクチャモチーフを特定し、観察されたパフォーマンスの向上を促進する分析ツールも提供します。興味深いことに、KVシェアリングや何らかの形のウェイトシェアリングなど、以前に提案されていた手動による相互接続パターンは、まったく新しいパターンと並行して自然に出現します。

STARで実証した機能は、その可能性を最大限に引き出すためのものです。あらゆる指標の組み合わせを最適化できることと、LIV の多様性が組み合わさったおかげで、合成設計の多様性と品質の両方が継続的に向上しているのを目の当たりにしています。どのパターンと目的が共存するかについての理解が深まったので、今度はSTARの進化的アルゴリズムと初期集団をさらに改良して、すべての世代が前世代より良くなるようにすることを検討しています。これには、ゲノム階層の最下位レベルで最適化を実行することも含まれます。また、モジュール設計空間を構築できる他の分野にも同様の方法論を適用することに関心があります。

すべての詳細については、論文を参照してください。」STAR: テーラードアーキテクチャの統合」。