AIはもっとスマートになれる──“重すぎる自己注意”を数学で軽くした話

「AIは賢くなっているけど、電気も食い過ぎてない？」

そんな素朴な疑問を抱いたことはないでしょうか？

今やAIは、チャットボットから画像生成、コード作成まで、私たちの暮らしの中に自然と入り込んでいます。
その多くのAIの頭脳には「Transformer（トランスフォーマー）」と呼ばれるモデルが使われており、その中核には「自己注意（self-attention）」という技術が組み込まれています。

でもこの自己注意、便利な一方で、使えば使うほどメモリも電力もガンガン消費するという大きな弱点があるのです。

今回紹介するのは、そんな自己注意の課題を解決するための画期的な新提案です。

「自己注意」ってそもそもなに？
AIの「記憶力」を支える代償：電力とメモリ
革命的な提案：トークンごとのコストが「一定」になる自己注意
「100万語」を読むAIが軽やかになる未来
「賢くなる」だけじゃなく、「やさしくなる」AIへ
まとめ：「無限の文脈」に向けて

「自己注意」ってそもそもなに？

まずは簡単におさらいしましょう。

自己注意とは、文章の中で「どの単語が、他のどの単語に注目すべきか」を決める仕組みです。

たとえば、

「花が咲いた。でも彼女はそれに気づかなかった」

という文があるとき「それ」は「花」を指していますよね。
この”つながり”を見抜くのが自己注意の役割です。

ところが、この注目（attention）を計算するためには、すべての単語の組み合わせを考えなければならないため、文が長くなればなるほど、計算が爆発的に増えてしまうのです。

AIの「記憶力」を支える代償：電力とメモリ

従来の自己注意では、文脈が長くなればなるほど、必要な「記憶」も「計算」も膨れ上がります。
研究では、このコストは文の長さに対して二乗の計算量（O(n²)）になるとされています。
全体としてn個のトークンを処理するのに、トークンあたりO(n)の空間と時間が必要になるのです。

つまり、1000語の文章を読むより、1万語の文章を読むほうが100倍ではなく10000倍大変になる、という世界です。

これはまるで、本を読むたびに脳のサイズも頭のエネルギー消費も大きくなるようなもの。
やがてはデータセンターの冷却ファンも音を上げてしまいます。

革命的な提案：トークンごとのコストが「一定」になる自己注意

そこで登場するのが、今回の論文が提案する新手法です。
「Symmetry-Aware Taylor Approximation（SATA）」による自己注意の近似計算です。

この手法では、自己注意のコアである指数関数（exp）の計算をテイラー展開という手法で近似します。

ここで「数学か…」と身構えなくても大丈夫。
例えるなら、

「曲がった道を進む代わりに、少しずつ曲げながら直進を繰り返して近づく」

といった感覚です。
完全な曲線（exp）ではなく、いくつかの直線でうまく真似をする、というイメージですね。

さらに、この近似を対称性をうまく活用して、重複する計算を省略することで、圧倒的な効率化を実現しています。
具体的には、テイラー展開の各項を対称的なテンソル積の連鎖として分解し、その対称的な構造から最小限の多項式基底を自然に導き出しているのです。

「100万語」を読むAIが軽やかになる未来

この手法の凄さは、次の一言に凝縮されます。

「1トークンあたりの計算量が一定になる」

つまり、どれだけ長い文章を読んでも、AIの負担は変わらないという夢のような状態が実現されるのです。

論文の図では、トークン数が増えるにつれて、従来の方法ではFLOPs（計算量）が直線的に増えていくのに対し、新手法ではほぼ水平のまま維持されているのがわかります。

この手法を使えば、AIが100万語の小説や数年分の対話ログを、今よりずっと少ないエネルギーで扱えるようになります。
隠れ状態のサイズは(dV + 1)×C(dK + P – 1, P – 1)という固定サイズで、トークンあたりのFLOPsも一定です。
ここで、dKとdVはキーと値のサイズ、Pはテイラー級数の項数、Cは組み合わせ関数です。

実際、論文では4つのテイラー項（P = 4）で、従来の注意機構をFloat16の精度レベルで再現できることが確認されています。