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

本日は、ニューラルネットワークのアーキテクチャ設計とカスタマイズの自動化における進展をご報告します。私たちは、新しい設計理論から導かれたモデルアーキテクチャの数値表現に進化的アルゴリズムを適用する形で、目的に合わせたアーキテクチャの合成（STAR）のためのアルゴリズムを開発しました。STAR はアーキテクチャの探索と最適化のプロセスを自動化し、エンドツーエンドのプロセスへと変えます。 これらの手法により、独自のタスク・評価指標・ハードウェアに合わせてアーキテクチャを最適化できるようになりました。STAR を用いて数百種類の設計を合成した結果、より小さなキャッシュとパラメータ数で、強力な Transformer やハイブリッド型アーキテクチャを品質で上回るものが得られました。

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

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

アーキテクチャの性能の基盤は、それを構成する計算ユニットと、それらがどのように相互接続されているかによって決まります。現在の深層学習アーキテクチャの大半は、アテンション演算子とゲート付き線形ユニットを逐次的に交互配置することで構築されています。これらの層は、私たちが線形入力可変システム（LIV）と呼ぶ、より広範な計算ユニット群の単純な例です。LIV は、その作用が入力自体によって点ごとに変調される構造化演算子であり、アテンションの各種バリアント、線形アテンション、（ゲート付き）畳み込み、線形状態遷移を伴う（ゲート付き）再帰、状態空間層、（ゲート付き）線形ユニットなど、多様なクラスの計算ユニットを一般化するための抽象化を提供します。

「Transformer を超える」アーキテクチャの可能性を示す証拠はすでに存在しています。ゲート付き畳み込みや再帰といった単純な LIV と自己アテンションを、層を交互に配置するストライプ型のハイブリッドパターンで組み合わせたモデルは、ある程度の品質向上、より長い系列への効率的なスケーリング、推論の高速化を実証しています（今年初めに、私たちはハイブリッドモデルに関する初のスケーリング則を報告しました）。

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

STAR の中核となるイノベーションの一つは、モデルアーキテクチャを STAR ゲノムと呼ばれる階層的な数値列として表現し、それを進化的最適化の原理に基づいて進化させることです。このプロセスは反復的に行われます。ゲノムを具体的なアーキテクチャにコンパイルして評価し、最も性能の高いアーキテクチャを選択・組み換えて次世代を生成します。

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

アーキテクチャ候補が新規かつ高性能であることを保証するため、ゲノムエンコーディングは、ニューラルネットワークのための新しい汎用的な計算ユニットのクラスである線形入力可変システム（LIV）を導入した私たちの設計理論に基づいています。私たちは、深層学習における現代の計算ユニットが入力コンテキストに基づいて計算をどのように変調するかを支配する基本的なメカニズムを特定し、それを中心に抽象化を構築しました。テンソルネットワークとシステム理論に基づくこのフレームワークは、LIV を 2 つの重要な側面から特徴づけます。構造（演算子におけるトークンとチャネルの混合構造）と特徴量化（演算子における入力依存性の関数形）です。高度なアーキテクチャ設計を表現できるよう、このフレームワークでは演算子の合成を第一級の関心事として扱い、層を逐次的に積み重ねるだけにとどまらないアーキテクチャ構築の新たな道を開いています。設計理論の詳細については、追ってご紹介します。

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

私たちはまず、改良された言語モデリングアーキテクチャの設計において STAR を評価し、3 つの目的の組み合わせについて最適化を行いました。i.) 品質（学習後のパープレキシティ）、ii.) 品質とパラメータ効率、iii.) 品質とキャッシュ効率です。

わずか 2〜3 ラウンドの進化の後には、ほとんどのアーキテクチャが Transformer や強力なハイブリッドのベースラインといった定番アーキテクチャを上回り、ラウンドを重ねるごとに一貫して改善しました。特に、品質のみを最適化した場合、評価した STAR 進化アーキテクチャはすべて、下流の評価ベンチマークにおいてアテンションと再帰を組み合わせたハイブリッドを上回り、その改善幅はハイブリッドが Transformer に対して示した改善幅の 2 倍に達しました。この結果は、私たちの設計空間における進化的探索の有効性を裏付ける強力な証拠であると考えています。ハイブリッドは多大なリソースを投じて手作業で設計・改良されてきたのに対し、STAR は 1 日足らずで、>90% のヒット率でアーキテクチャを生成できます。

この探索は多目的最適化にも対応しています。品質とモデルサイズを同時に最適化すると、進化したアーキテクチャはパラメータ数を削減しつつ、Transformer とストライプ型ハイブリッドの両方を一貫して上回り、エッジ環境やリソースに制約のある環境向けにモデルをコンパクト化できます。進化と評価の結果をスケールをまたいで転移させるためのさまざまな手法を試したところ、一般に、目標とする幅でアーキテクチャのモチーフを最適化するよりも、細く深いアーキテクチャ候補で最適化する方が良好な転移が得られることがわかりました。

同様に、推論スタック上で直接プロファイリングして得たターゲットハードウェアでのレイテンシ、品質、モデルサイズのバランスを取るためにも STAR を使用しました。これが可能なのは、STAR が評価指標の勾配を必要とせず、ゲノムからコンパイルしたアーキテクチャで計算される静的指標と動的指標の組み合わせに対応しているためです。これには、レイテンシと通信オーバーヘッドを最小化するための詳細なプロファイリングも含まれます。

STAR は、特定の目的に合わせてアーキテクチャを最適化するだけでなく、進化の過程で繰り返し現れ、観察された性能向上をもたらすアーキテクチャのモチーフを特定する分析ツールも提供します。興味深いことに、KV 共有や何らかの形の重み共有など、これまでに提案されてきた手動の相互接続パターンが、まったく新しいパターンとともに自然に現れます。

今回 STAR で実証した機能は、その可能性のほんの一端にすぎません。任意の評価指標の組み合わせを最適化できることと LIV の汎用性が相まって、合成される設計の多様性と品質の両方が継続的に向上しています。どのパターンと目的が同時に現れるかについての理解が深まったことを受け、今後は STAR の進化的アルゴリズムと初期集団をさらに改良し、ゲノム階層の最下位レベルでの最適化の実行も含め、世代を重ねるごとに前の世代より良くなるようにすることを目指しています。また、モジュール型の設計空間を構築できる他の領域にも、同様の方法論を適用することに関心を持っています。

詳細は論文「STAR: Synthesis of Tailored Architectures」をご覧ください。

引用

本研究を利用する場合は、以下のテクニカルレポートを引用してください。

@inproceedings{thomas2025star,
  title={STAR: Synthesis of tailored architectures},
  author={Thomas, Armin and Parnichkun, Rom and Amini, Alexander and Massaroli, Stefano and Poli, Michael},
  booktitle={International Conference on Learning Representations},
  volume={2025},
  pages={88489--88541},
  year={2025}
}