🤖 machine learning

LLMForge: Multi-Backend Hardware-Aware Neural Architecture Search with Infinite-Head Attention for Edge Language Models

LLMForge は、無限ヘッド・アテンションとマルチバックエンドのコストモデルを活用して、特定のエッジデバイスの制約に最適化されたサブ10 億パラメータ規模の言語モデルを自動的に生成するハードウェア意識型ニューラルアーキテクチャ探索フレームワークであり、既存のベースラインと比較してエネルギー効率、遅延、モデル精度において著しい改善を実現します。

原著者： Xinting Jiang, Junyi Luo, Ruichen Qi, Kauna Lei, Ben Laurie, Gregory Kielian, Mehdi Saligane

公開日 2026-05-19

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Xinting Jiang, Junyi Luo, Ruichen Qi, Kauna Lei, Ben Laurie, Gregory Kielian, Mehdi Saligane

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

スマートフォンのようなスマートフォンやスマートウォッチで動作する、完璧で小さなロボット脳を構築しようとしていると想像してください。言語を理解できるほど賢く、かつバッテリーを消費したり端末を過熱させたりすることなくポケットに入るほど小型であることが求められます。

問題は、これらの脳（トランスフォーマーと呼ばれる）の標準的な「設計図」が、巨大なスーパーコンピュータ向けに設計されていたことです。これらをエッジデバイス向けに縮小しようとすると、往々にして渋滞に巻き込まれたり、メモリ不足に陥ったり、エネルギーを大量に消費したりします。

LLMForgeは、この問題を解決するために研究者によって開発された新しいツールです。これは単に脳を縮小するのではなく、そのデバイスが動作する特定の「地形」に基づいて設計図を完全に再設計する、超賢くハードウェアを認識する建築家のようなものです。

LLMForge の仕組みを、3 つの簡単な部分に分けて解説します。

1. インフィニットヘッドアテンション（IHA）：規則の打破

従来のロボット脳には厳格な規則があります。思考の「ヘッド」（文の異なる部分を見るプロセッサ）の数を増やそうとすると、各ヘッドのサイズを小さくせざるを得ません。ピザを例に取れば、スライスを多く切れば切るほど、各スライスは小さくなります。これにより、脳の賢さには限界が生じます。

LLMForge の革新点： 厳格な規則書を捨て去ります。インフィニットヘッドアテンションにより、建築家は脳の各レイヤーごとに、スライスの数、スライスのサイズ、トッピングの種類を独立して変更できます。

比喩： 通常、部屋をすべて同じサイズで建設しなければならない建設チームを想像してください。LLMForge は彼らに魔法の道具を与え、同じフロアプラン内で、巨大なボールルームの隣に小さなクローゼット、そして中規模のオフィスを建設できるようにします。これにより、可能な設計図の数が400 倍に拡大し、特定のデバイスの制約に完璧に適合する形状を見つけることが可能になります。

2. フォージ・フューマー：水晶玉

新しい脳の設計図をゼロから構築し、テストするのは高価で時間がかかります。レシピが機能するか確認するためにケーキを焼き、その後それを捨てるようなものです。数千のレシピをテストする場合、破産してしまいます。

LLMForge の解決策： 彼らは水晶玉と呼ばれるフォージ・フューマーを構築しました。

仕組み： すべてのケーキを焼く代わりに、この水晶玉はレシピ（設計図）を見て、ケーキがどのくらい美味しくなるか（モデルがどのくらい賢くなるか）、そしてどれだけのエネルギーを消費するかを正確に予測します。
結果： この水晶玉は、従来の「推測」ツールよりもはるかに正確です。これにより、システムは以前は数種類しかテストできなかった時間を、数千種類の設計をテストするために使用できるようになり、莫大な時間とエネルギーを節約します。

3. フォージ-DSE：マルチツールナビゲーター

異なるデバイスには異なる課題があります。高性能なグラフィックカード（GPU）はデータ転送速度によって制限される可能性がありますが、スマートウォッチの小さなチップは発熱許容量によって制限される可能性があります。あるデバイスには完璧な設計が、別のデバイスにはひどい結果をもたらすこともあります。

LLMForge のアプローチ： フォージ-DSEエンジンが、マルチツールを備えたナビゲーターとして機能します。

単に「最も賢い」モデルを探すのではなく、賢さ、速度、エネルギー効率の間の最適なトレードオフ（パレートフロンティア）を探します。
設計を4 種類の異なるハードウェア（高速 GPU、専用チップ、リング型チップなど）に対してテストします。
結果： システムは「一つのサイズがすべてに適合するわけではない」ことを認識します。各デバイスに対して、異なる固有の設計図を生成します。
- 速度重視のデバイスに対しては、幅広で浅い脳を構築します。
- エネルギー重視のデバイスに対しては、深く細い脳を構築します。

結果：実世界での勝利

研究者たちは、このシステムを約 1 億パラメータと 3 億パラメータの 2 つのモデルサイズでテストし、SmolLM2 や Qwen などの既存の人気モデルと比較しました。

「精度」モデル： 彼らの新しい設計の一つは、パラメータ数が少ない（つまり小型である）にもかかわらず、競合他社よりも賢く（誤り率が低く）動作しました。
「エネルギー」モデル： もう一つの設計は、標準モデルと比較して、生成される単語あたりのエネルギー消費を40% 削減しました。
「速度」モデル： 3 つ目の設計は、話し始めるまでの時間（Time-to-First-Token）と文の完了において、43% 高速でした。

まとめ

LLMForgeは、すべてのデバイスに「一つのサイズがすべてに適合する」脳を押し付けようとするのをやめ、代わりに柔軟な設計言語、高速な予測水晶玉、そして賢いナビゲーターを使用して、あらゆる特定のハードウェアに合わせたカスタム脳を設計するシステムです。その結果、ポケットに入るサイズでありながら、より小型で、高速で、エネルギー効率の高い AI が実現しました。

技術概要：LLMForge

問題定義

サブ10億パラメータのトランスフォーマー言語モデルは、プライバシー、レイテンシ、コストの優位性を活用するため、エッジデバイスへの展開が増加しています。しかし、オンデバイス推論は、厳格なメモリ帯域幅、エネルギー、熱予算によって制約されています。既存のアプローチは、量子化やプルーニングなどの学習後圧縮に依存することが多く、これにより基盤となるアーキテクチャの自由度が十分に探求されないままになっています。さらに、トランスフォーマー向けの現在のハードウェア対応ニューラルアーキテクチャサーチ（NAS）パイプラインには、以下の3つの限界があります：

制限された探索空間: アテンションのパラメータ化は、しばしばヘッドごとの次元を固定し、キー/バリューの容量を結合するため、実行可能な形状設定を単一のヘッド数選択に限定しています。
評価のボトルネック: 候補をゼロから訓練して評価することは、計算コストが高すぎます。
粗いコストモデリング: ハードウェアコストは、単一のアクセラレータクラスやパラメータ数などの粗い代理変数を通じてモデル化されることが多く、基盤に依存した候補の再ランキングが隠蔽されています。

手法

著者らは、3つの構成可能な貢献を単一のNSGA-II進化的探索ループに統合するハードウェア対応NASフレームワーク「LLMForge」を提示します。

1. 無限ヘッドアテンション（IHA）

IHAは、アテンション形状パラメータ（ $n_h$ , $n_{kv}$ , $d_{qk}$ , $d_v$ ）を、独立した層ごとの変数として解離します。標準的なマルチヘッドアテンション（MHA）やグループ化クエリアテンション（GQA）が、可除性の制約（ $d_{model} = n_h \cdot d_h$ ）とQ/K–Vの結合（ $d_h = d_{qk} = d_v$ ）を強制するのに対し、IHAはこれらの制約を緩和します。ただし、GQAのグループ化条件（ $n_{kv} | n_h$ ）は維持されます。

影響: これにより、定義された探索範囲内において、GQA下での27の構成からIHA下での11,250の構成へと、実行可能な層ごとのアテンション設定空間が約400倍拡大します。

2. Forge-Former（精度代理モデル）

評価のボトルネックに対処するため、著者らはトランスフォーマーエンコーダーベースの代理モデル「Forge-Former」を導入します。

アーキテクチャ: これは、アーキテクチャを層トークンの系列としてエンコードし、各トークンには層のIHAパラメータが連結されます。スカラーフィールドは線形層を通じてリフトされ、位置エンコーディングと加算されます。事前LNトランスフォーマーエンコーダーは、層内相互作用（フィードフォワードサブレイヤーを通じて）と層間相互作用（自己アテンションを通じて）を捕捉します。
性能: IHAテスト分割において、Forge-Formerはスピアマンの $\rho = 0.75$ およびケンドールの $\tau = 0.58$ を達成し、MLPおよびランダムフォレストのベースラインを1.4倍から2.5倍上回ります。これにより、ゼロから訓練するコストの断片でNSGA-IIの選択が可能になります。

3. Forge-DSE（設計空間探索エンジン）

Forge-DSEは、Forge-Formerをマルチバックエンドハードウェアコストモデルと組み合わせる多目的探索エンジンです。

探索ループ: 非支配ソートと混雑距離生存を用いたNSGA-IIにより、IHAパラメータ化されたアーキテクチャの集団を進化させます。カスタム変異演算子（削除、複製、回転、数値的摂動）と実行可能性修復ステップを採用します。
代理共進化: フレームワークは、現在の集団から有益なサンプルを実際に再訓練し、経験再生を通じて代理モデルを微調整することで、分布のシフトを修正しながら、探索中にForge-Formerを任意に洗練させます。
ハードウェアバックエンド: システムは3つの異なるバックエンドをサポートします：
1. バックエンドA: ZEUSを使用したNVIDIA A100 GPU上での測定コスト。
2. バックエンドB: Timeloopによる解析コスト。シストリック（Gemmini）、行固定（Eyeriss）、結合アテンション（FLAT）のデータフローを対象とします。
3. バックエンドC: エッジアクセラレータ向けのマルチチップリングデータフローシミュレータ（rDXE）。アーキテクチャとチップの共同探索を実行します。

主要な結果

代理モデルの検証

IHAデータセットにおいて、Forge-Formerは、ランキング忠実度とトップX%の回復においてベースラインを大幅に上回ります。MLPおよびランダムフォレストのベースラインと比較して、上位5%のアーキテクチャの平均絶対誤差（MAE）を約5倍削減します。t-SNE分析は、Forge-Formerの埋め込みが性能の類似性をよりよく反映し、高損失の尾部を密にクラスタリングすることを確認します。

マルチバックエンド探索性能

4つのハードウェア基盤（A100、Gemmini、Eyeriss、FLAT）において、探索は、各基盤固有のコストボトルネックを追跡する、異なるパレートフロントアーキテクチャに収束します：

A100（帯域幅豊富）: 平坦なMLP幅と最も広いアテンションヘッドをもたらします。
Eyeriss/Gemmini（KVキャッシュ制約）: $n_{kv}$ を積極的なグループ化クエリアテンションに縮小し、より広いMLP容量で補償します。
FLAT（柔軟）: 浅い深さと狭いヘッドを持つ最小のモデルを生成します。
改善点: SmolLM2-135Mベースラインと比較して、探索により検証損失が5–7%、トークンあたりのエネルギーが16–63%、レイテンシ指標（TTFT/TPOT）が16–59%削減されました。

rDXE上でのスケーリング検証

著者らは、rDXEマルチチップリング基盤上でアーキテクチャとチップの共同探索を行い、選択された候補をFineWeb-Edu-10BT（131億トークン）で再訓練し、SmolLM2およびQwenベースラインに対して評価しました：

精度変種（3.47億パラメータ）: SmolLM2-3.6億パラメータおよびQwen-0.5Bよりも少ないパラメータで、最低の検証損失（2.798）と競争力のあるベンチマーク性能を達成しました。
エネルギー変種（2.94億パラメータ）: トークンあたりのエネルギーを**40%**削減しました。
レイテンシ変種（3.65億パラメータ）: 最初のトークンまでの時間（TTFT）および出力トークンあたりの時間（TPOT）を**43%**削減しました。

意義と主張

本論文は、LLMForgeが、単一の「最良」アーキテクチャではなく、基盤条件付きパレートフロントを生成できることを実証していると主張しています。IHAを通じてアテンションパラメータを解離し、高忠実度なエンコーダーベースの代理モデルを利用することで、このフレームワークは、特定のハードウェアボトルネックに一致するように、層ごとの幅と深さのプロファイルを効果的に再形成します。

著者らは、LLMForgeを、展開を意識したアーキテクチャとチップの構成ペアの発見を可能にする方法論的貢献として位置づけています。エネルギーとレイテンシの削減は modest（ modest）であるものの、それらは主に反復的なエッジ推論の運用コストと炭素プロファイルに利益をもたらすと指摘しています。この研究は、代理モデルの訓練コーパス生成の計算コストにより、サブ10億パラメータ規模（ $\le$ 500M）に限定されており、10億パラメータモデルへのスケーリングは将来のステップとして特定されています。