AttentionでもMambaでもない。「位相」で記憶する言語モデルが動いた（Phase-Associative Memory）

Phase-Associative Memory: Sequence Modeling in Complex Hilbert Space 解説

概要

複素数の位相を使って過去を記憶する言語モデルが、実際に学習できています。

提案手法のPhase-Associative Memory、PAMは、token間の関連を複素行列へ書き込み、queryとkeyの位相が合う情報だけを読み出す系列モデルです。softmax attentionは使いません。Mambaのような実数値のState Space Modelとも構造が違います。

この論文でまず驚くのは、PAMが既存モデルを上回ったことではありません。実際、5Mから100Mまでの全規模で、実数値の比較対象SAMに負けています。

それでも、別の記憶原理だけで言語モデルを最後まで動かした。そこが、この論文の一番の見どころです。

規模を増やすとPAMの方が速く改善し、両者の差は単調に詰まりました。著者らは、この傾向が続けば、さらに大きな規模で逆転する可能性があると報告しています。

はじめに

現在の言語モデルで系列を扱う方法は、大きく二つあります。

一つ目はTransformerのAttentionです。過去のtokenを直接参照できるため、必要な情報を精密に引き出せます。一方、学習時の計算量は系列長を T とすると O(T^2)。生成時には過去のkeyとvalueを保存するKV cacheも必要です。

二つ目はMambaなどのState Space Modelです。過去の情報を固定サイズの状態へ順番に畳み込みます。生成時の計算量をtokenあたり一定にできますが、多数の関連を限られた状態へ押し込むので、必要な情報を正確に取り出すことが難しくなります。

では、過去のtokenを直接見返さず、一本の状態ベクトルへ無理に詰め込むこともなく、情報を保存する方法はないのか。

PAMの答えは、状態を複素行列にして、連想記憶として関連を書き込むというものです。

複素数には、値の大きさを表す振幅と、向きを表す位相があります。PAMでは、位相がそろった記憶は強め合い、位相がずれた記憶は打ち消し合います。softmaxで候補を選ぶ代わりに、複素数の干渉そのものが検索になります。

PAMが既存研究と完全に切り離された仕組みというわけではありません。状態へvalueとkeyの外積を加える更新式は、linear attention、fast-weight programmer、mLSTMなどと地続きです。Mamba-2も構造化State Space Modelとlinear attentionの接点を論じています。

PAMが違うのは、その行列状態を最初から最後まで複素数で扱い、共役内積で記憶を読み出す点です。

図1: 問題の例

論文のFigure 1です。WikiText-103を使い、約10MパラメータのPAMと実数値版SAMを10 epoch学習した結果です。上段がvalidation loss、下段がperplexity。

図の左側と右側で何が違うか
横軸は学習の進み方で、左が学習初期、右が学習後半です。星印のSAMは早い段階で損失が下がります。三角印のPAMも安定して学習できていますが、収束は遅く、最後までSAMとの差が残ったまま。
どこを見るべきか
曲線が途中で交差するかどうかではなく、10 epoch後にも残る縦方向の差です。perplexityはSAMが40.20、PAMが58.71。
この図から何が問題だと分かるか
言語モデルを複素数に置き換えただけでは、性能は上がりません。新しい計算原理で学習自体は成立していますが、小規模では実数値モデルの方が明確に強い。

提案手法

PAMの中心は、複素行列へ関連を書き込み、位相の一致で読み出す連想記憶です。

各headは、複素数からなる行列状態を持ちます。新しいtokenが入るたびに、valueとkeyの複素共役から外積を作り、行列状態へ加えます。

考え方だけを書けば、次の二段階です。

蓄積：前の状態を少し忘れ、valueとkeyの関連を書き込む
読み出し：行列状態へqueryを掛け、位相が合う記憶を取り出す

keyとqueryの位相がそろっていれば、対応するvalueが強く出ます。位相がずれていれば、複素数の成分が互いに打ち消し合う。

Attentionは、各tokenとの類似度を計算し、softmaxで重みを鋭くして候補を絞ります。PAMにはこのsoftmaxがありません。位相が合わない候補は干渉で勝手に消える。確率として並べて選ぶのとは仕組みが違います。

Mambaとの違いは記憶の形です。PAMはheadごとに複素行列を持ちます。行列には d^2 個の複素成分があり、keyが互いに直交していれば、最大 d 個の関連を劣化させずに保持できます。一本のベクトルに全部を畳み込むMambaとは、容量の確保の仕方が違います。

古い情報をどれだけ残すかは、入力から計算する減衰率と保護gateが決めます。重要な状態は保護し、不要になった関連は減衰させる。

モデル全体では、tokenの埋め込み、channel方向の変換、PAM、出力層まで複素数値のまま処理します。活性化には位相を壊さないmodReLU、正規化は振幅にだけ作用させます。位置情報はComplex RoPEで位相に埋め込む。

学習時には系列を並列に処理するため、計算量はAttentionと同じ O(T^2) です。一方、生成時には固定サイズの行列状態だけを更新するので、tokenあたり O(1) で処理でき、KV cacheは要りません。

図2: 提案手法の概念図

論文には、PAM全体の処理の流れを示す図は掲載されていません。そこで、位相を使う意味が最も分かりやすいFigure 3を示します。

この図は、学習後の複素token埋め込みについて、同義語、反義語、無作為な単語対の位相関係を調べたものです。横軸は二つの単語間の位相差、縦軸は位相の整合度。

提案手法では何を追加・変更しているか
token、key、query、value、行列状態を複素数に変えています。検索には複素共役を使った内積を用います。
それによって何が解決されるか
位相が一致する関連は強く残り、合わない関連は打ち消されます。無作為な単語対より、同義語や反義語の方が高い位相整合度を示しており、学習された位相が単なる雑音ではないことが読み取れます。
なぜこの方法がシンプルで効果的なのか
tokenごとの重みをsoftmaxで計算する必要がありません。書き込みは外積、読み出しは行列とqueryの積。生成時には過去の全tokenを保存せず、行列状態だけを更新すればよい。

位相だけで全情報を表しているわけではない点も報告されています。著者らはkeyとqueryの振幅を固定し、位相だけで検索する実験も行いました。validation lossは下がったものの、生成文が同じ単語を繰り返す状態へ崩れたそうです。位相がどの関連かを表し、振幅がその強さを補っている、という構図のようです。

評価実験

PAMと実数値版のSAMを、WikiText-103で比較しています。

SAMは、PAMと同じ行列状態と処理順序を持ちます。ただし、すべての複素数計算を実数へ置き換えたもの。複素線形層は実部と虚部に対応する二つの重みを持つため、SAM側は次元数とmemory bank数を増やし、総パラメータ数をそろえています。

比較した規模は5M、10M、25M、50M、100Mの5段階。学習率、batch size、系列長、epoch数は統一。主要実験はApple M4 Max上で行われました。

結果だけを見ると、PAMはまだ勝っていません。測定したすべての規模で、validation lossはSAMの方が低いまま。100Mでも、PAMが3.56、SAMが3.26です。

ただし、規模を増やしたときにPAMの方が速く改善しました。5Mで0.78あったloss差が、100Mでは0.30まで詰まっています。

図3: 実験結果

論文のFigure 2です。上段がvalidation loss、下段がvalidation perplexity。横軸はパラメータ数で、両軸とも対数表示。三角がPAM、星がSAM。

どのデータセット・タスクで評価したか
WikiText-103を使った次token予測。GPT-2のBPE tokenizerを使い、系列長512で学習しています。
提案手法あり・なしで何が変わったか
実測範囲では、実数値のSAMがすべて勝っています。ただし近似直線の傾きはPAMの方が急で、lossの傾きはPAMが-0.15、SAMが-0.12。
結果から何が分かるか
PAMは小規模では不利ですが、規模を増やすほどSAMとの差が詰まっています。近似線を延長すると、lossでは約4.5Bパラメータ、perplexityでは約550Mパラメータで交差する計算。

もちろん、4.5Bでの逆転は実測ではありません。最大100Mまでの5点から引いた外挿です。

それでも、5つの規模すべてで差が同じ方向へ縮んでいます。attentionでもState Space Modelでもない方法が、規模を増やすほど差を詰めてくるというのは、データとしては出ている。

学習後のPAMが使う行列状態の有効rankは、64次元中およそ10で頭打ちでした。大きな行列を用意しても、現在の文脈に関係する少数の関連だけが残る。使っているのは容量の1割程度です。

結論

著者らは、複素行列へtoken間の関連を書き込み、keyとqueryの位相一致で記憶を読み出すPhase-Associative Memoryを提案しました。Attentionのように過去の全tokenを参照するわけでもなく、Mambaのように情報を実数値のベクトル状態へ畳み込むわけでもない。複素数の干渉を、検索そのものとして使います。

実験では、5Mから100Mまでの全規模で実数値版SAMに負けました。一方、パラメータを増やすとPAMの方が速く改善し、差は一貫して詰まっています。

この論文は、PAMがTransformerを倒した話ではありません。attentionでもMambaでもない原理だけで、言語モデルを本当に学習させられた。まずは、そこが大きい。

実装、学習ログ、model weightは、著者のGitHub repositoryで公開されています。

参考文献:
「Phase-Associative Memory: Sequence Modeling in Complex Hilbert Space」, Gowrav Vishwakarma, Christopher J. Agostino, arXiv, 2026. DOI: 10.48550/arXiv.2604.05030.