1916年、Albert Einsteinは一般相対性理論の論文で、テンソル演算の添字を省略する記法を導入した。同じ添字が2回現れれば、その添字について総和を取る。この約束事だけで、何行もの総和記号が消える。今日、この「アインシュタインの縮約記法」はEinsum（Einstein summation）としてPyTorchやTensorFlow、NumPyに組み込まれ、深層学習の行列演算やテンソル演算を1行で記述する標準的な道具になっている。
Einsumの利点は簡潔さだ。MIT CSAILの博士課程学生Jaeyeon Wonは「コードの行数が少なくて済む」と述べている。ただし、この簡潔さには代償がある。Einsumは配列の中でどこがゼロなのかを把握しない。すべての要素に対して律儀に演算を実行する。ゼロに何を掛けてもゼロだという、小学生でも知っている事実を、GPUは何十億回も繰り返している。
この無駄が許容できない場面がある。現実世界のデータの多くは「スパース」、すなわち要素の大部分がゼロの配列として表される。SNSの人間関係グラフ、推薦システムの評価行列、分子シミュレーションの相互作用行列。Amazonのレビューデータをテンソルとして表現すると、成分数は$1.5 \times 10^{19}$個（8バイト換算で107エクサバイト）に達するが、非ゼロ成分は$1.7 \times 10^9$個（13ギガバイト）にすぎない。99.99999999%以上がゼロだ。
GPUは「密」な世界のために作られた
現代のGPUは密な行列、つまりゼロが少ない配列の演算に最適化されている。NVIDIAのTensor Coreは、密な行列積を1クロックあたり数千回の積和演算で処理する。スパース配列をそのまま渡すと、GPUはゼロを含む全要素に対して同じ演算を実行し、Tensor Coreの能力の大部分が空転する。
この問題を解決する試みは以前からあった。2017年、MITのFredrik KjolstadとSaman AmarasingheらはTACO（Tensor Algebra Compiler）を発表した。TACOは任意のスパースと密のテンソル演算に対して自動的にカーネル（GPUやCPU上で実行される計算の最小単位）を生成する初のコンパイラだった。スパース同士の演算では手書きの最適化カーネルに匹敵する性能を出す。
しかしTACOには構造的な限界があった。第一に、主にCPU向けに設計されており、GPU向けの高品質なコードを生成できない。第二に、スパースと密が混在する演算で、密な部分を効果的に最適化できない。Wonは「スパース同士の演算ではTACOが最良だが、スパースと密の掛け算にはまったく向かない」と指摘する。
その結果、スパースGPU計算の現場では、専門家が数百行から数千行の手書きCUDAコードを書くという力技が続いていた。グラフニューラルネットワークのライブラリや科学シミュレーションのコードベースには、こうした手作業カーネルが散在している。
一行の式からGPUカーネルを引き出す三段階の変換
WonとAhrens、Amarasinghe、Emerの4人が提案した解決策は、Einsumそのものを「書き換える」ことだった。彼らが考案したIndirect Einsum（間接Einsum）は、標準的なEinsumの簡潔さを保ちながら、スパース形式のメタデータ（非ゼロ要素の位置情報）を式の中に明示的に組み込む。
仕組みは三段階で動く。第一段階で、スパース配列をGroupCOOやBlockGroupCOOという新しい固定長形式に変換する。これらの形式は非ゼロ要素の値と座標を分離して格納し、GPUの規則的なメモリアクセスパターンに合うよう設計されている。第二段階で、形式情報を持たないEinsumを、間接インデックス（indirect indexing）を使って密なテンソル演算に書き換える。たとえば のような式は、「AIというインデックス配列を介してAとBの非ゼロ要素を取り出し、密な演算として実行し、結果をCの所定の位置に戻す」という意味になる。第三段階で、InsumコンパイラがこのIndirect EinsumをPyTorchコンパイラのFXグラフ（PyTorchプログラムの関数型中間表現）に変換し、TorchInductor経由で最適化されたGPUカーネルを生成する。
この設計の巧みさは、スパース計算専用の新しいハードウェア機構やコンパイラ基盤を必要としない点にある。既存の密行列演算コンパイラ（PyTorchコンパイラとTensor Core）を「そのまま」利用できる。スパース性の扱いは、データの変換と間接インデックスの挿入という前段階に集約されている。
さらに研究チームは、PyTorchコンパイラ自体にも拡張を加えた。Indirect Einsumの式パターンを認識してTensor Core命令を直接生成するパターンマッチングと、Tensor Core使用時の性能を向上させるLazy Broadcastingというコード生成技法だ。
手書きコード4,491行が1行に
研究チームは、グラフニューラルネットワークのメッセージ伝播、スパース行列と密行列の積、科学シミュレーションのカーネルなど多様なスパースGPUアプリケーションでInsumを評価した。結果は以下の通りだ。
| 指標 | 手書き最適化カーネル | Insum |
|---|---|---|
| 実行速度 | 基準（1.0×） | 1.14×〜3.81×高速 |
| コード行数 | 数百〜数千行 | 1行のIndirect Einsum式（202×〜4,491×の削減） |
| 対象ハードウェア | 特定GPU向けに個別最適化 | PyTorchコンパイラが対応する任意のGPU |
| スパース形式 | 手動で選択・実装 | GroupCOO / BlockGroupCOOを自動生成 |
最大3.81倍の高速化は、TechXploreの見出しにある「ほぼ4倍」の根拠だ。コード行数の削減幅は202倍から4,491倍に及ぶ。4,491行の手書きCUDAコードが、1行のIndirect Einsum式に置き換わるケースがあったということだ。
Metaのコンピュータ科学者Jason Anselは「グラフニューラルネットワークや科学シミュレーションのスパース計算は、現代のGPUで効率的に実行するのが極めて難しいことで知られる。Insumはこの問題を、より少ないコードでより速い実行を可能にする方法で緩和する」と評している。
比較: 既存手法とInsumの構造的な違い
| TACO (2017) | 手書きCUDAカーネル | Insum (2026) | |
|---|---|---|---|
| 主な対象 | CPU | GPU | GPU |
| スパース×スパース | 得意 | 可能だが労力大 | 拡張中（未対応） |
| スパース×密 | 不得意 | 可能だが労力大 | 得意 |
| Tensor Core活用 | なし | 手動で実装 | パターンマッチで自動生成 |
| 必要な専門性 | 中（形式記述） | 極めて高い | 低（式を1行書く） |
| コード量 | 数行 | 数百〜数千行 | 1行 |
自動化と「スパース×スパース」の壁
Emerは「今日のコンピュータは密な計算が非常に得意なので、この変換（標準EinsumからIndirect Einsumへの変換）が採算に合うには、ゼロの割合が十分高い必要がある」と述べている。ゼロが少数しかない配列では、変換のオーバーヘッドが計算節約を上回る可能性がある。
より根本的な制約は、この変換が現状では手作業だという点だ。研究者がスパース配列の形式と演算の構造を見て、適切なIndirect Einsum式を手で書き下す必要がある。Amarasingheは「このプロセスの自動化は進行中の研究課題だ」と述べている。
もう一つの未解決問題は、スパース配列同士の演算だ。Indirect Einsumへの変換は、スパースと密の組み合わせでは比較的素直に定義できるが、両方がスパースの場合は手順が格段に複雑になる。AmarasingheとEmerは現在、Ahrensおよびジョージア工科大学の学生と共同でこの問題に取り組んでいる。
TACOが2017年に「任意のテンソル演算をコンパイルできる」という普遍的な枠組みを提示してから9年。Insumはその枠組みが取りこぼしていた「GPU上のスパース×密演算」という実用上の空白を埋めた。次の空白は、変換の自動化とスパース同士の演算だ。この2つが埋まれば、スパース計算の専門知識を持たない開発者でも、GPUの性能をゼロの海から引き出せるようになる。
