新フォーマットTwELL、スパースLLMの学習を高速化
Sakana AI
2026年5月10日 (日)
- •Sakana AIとNVIDIAがスパース・トランスフォーマー・モデル向けの新しいフォーマット「TwELL」を発表した。
- •TwELLはカスタムCUDAカーネルを活用し、トークンを動的にルーティングすることでメモリ効率を向上させる。
- •ベンチマークでは、数十億パラメータ規模のモデルにおいて20%以上の高速化とメモリ消費量の削減を実証した。
本稿では、スパース・アーキテクチャを採用した大規模言語モデル(LLM)を最適化するための新しい技術的アプローチ「TwELL(Tile-wise ELLPACK)」を紹介する。スパース・モデルとは、入力データに対してモデル内の限られた一部のニューロンのみが活性化される設計のモデルを指す。現代のLLMは本質的にスパース性を備えているものの、標準的なハードウェアでは、この計算に伴う不規則なメモリ・アクセス・パターンを効率的に処理することが困難であった。
このハードウェア上の不整合を解決するため、研究チームは「ハイブリッド」形式を開発した。TwELLは、大半のスパースなトークンを高速な実行パスへ動的に振り分けつつ、複雑で計算負荷の高いトークンには密行列(dense matrix)のバックアップを用いる手法を採る。この戦略により、GPUがスパースな計算に伴う不規則な構造に悩まされる事態を防ぐことができる。
研究チームは、複数のSparse matrix乗算を融合し、ハードウェアのスループットを最大化して活性化サイズを最小限に抑えるカスタムCUDAカーネルを導入した。数十億パラメータ規模のモデルを用いた学習およびベンチマークの結果、本手法は20%を超える高速化に加え、メモリとエネルギー効率の大幅な向上を実現した。この成果は、ICML 2026で発表される予定である。