The Discrete-Log Clock: How a Transformer Learns Modular Multiplication

技術要約：離散対数クロック（The Discrete-Log Clock）

問題提起
トランスフォーマーが加算（$a + b \mod p$）を学習する際の内部メカニズムは、疎なフーリエ表現（「クロック」アルゴリズム）として完全に逆解析されている一方で、乗算（$a \cdot b \mod p$）のメカニズムはこれまで解明されていなかった。先行する経験的研究（Doshi et al., 2024; Furuta et al., 2024）では、トランスフォーマーが乗算を学習する際、加算で見られた疎な構造とは対照的に、すべての周波数を必要とする「密な」フーリエスペクトルを学習することが報告されている。これは、乗算が根本的に異なる、解釈性の低いアルゴリズムに依存している可能性を示唆していた。

手法
著者らは、$a, b \in \{1, \dots, 112\}$ の入力に対して $a \cdot b \mod 113$ を予測する1層のデコーダーのみのトランスフォーマーを訓練する。モデルはデータの30%で訓練され、残りの70%に対して汎化しなければならない。ここでは、長い暗記期間の後に汎化が突如として発生する「グロッキング（grokking）」現象が見られる。

本研究の核心的な手法の革新は、分析基底の転換にある。トークン埋め込みに対して標準的な加法的離散フーリエ変換（DFT）を適用する代わりに、著者らは**乗法的指標変換（multiplicative character transform）**を適用する。これには以下の手順が含まれる：

1. 素数 $p$ を法とする非ゼロ整数が巡回群 $(\mathbb{Z}/p\mathbb{Z})^*$ を形成することを認識する。
2. 原始根 $g$（ここでは $g=3$）を用いて、入力をその離散対数 $\alpha = \log_g(a)$ および $\beta = \log_g(b)$ に写像する。
3. これらの離散対数に従って、埋め込み行列の行を並べ替える。
4. この並べ替えられたデータに対して標準的なDFTを適用する。

このアプローチは、元の空間における乗算が、対数空間における加算に対応するという群同型（$a \cdot b \equiv g^{\alpha + \beta} \mod p$）を利用している。

主な結果

• スペクトルの疎性: 加法的基底で分析した場合、埋め込みスペクトルは平坦で密に見える（ジニ係数 0.07）。しかし、乗法的指標基底では、スペクトルは極めて疎になり（ジニ係数 0.58）、エネルギーはわずか 4つの主要な周波数 ($\{2, 8, 47, 56\}$) に集中する。
• ニューロンのチューニング: 乗法的基底において、96.9%（496/512）のMLPニューロンが単一の周波数に「きれいにチューニング」されており、活性化エネルギーの85%以上がそこに集中している。加法的基達では、この閾値を満たすニューロンはゼロであった。
• 構造の可視化: ニューロンの活性化ヒートマップを離散対数によって並べ替えると、明確な2次元周期的な対角ストライプが現れる。これは、ニューロンが $\cos(k(\alpha + \beta))$ の形式の関数を計算しており、実質的に対数空間での加算を実行していることを示している。
• SVD分析: 離散対数によって再配置された埋め込み行列の主成分は、明確な正弦波構造を示し、乗法的指標の符号化を裏付けている。

「離散対数クロック」アルゴリズム
論文は、トランスフォーマーが加算のための「クロック」アルゴリズムに類似したアルゴリズム、すなわち**離散対数クロック（Discrete-Log Clock）**を実装していると結論付けている。メカニズムは以下の通りである：

1. 埋め込み: 整数は、その離散対数の正弦波関数として埋め込まれる。
2. ルーティング: アテンション機構がこれらの埋め込みを出力位置へとルーティングする。
3. 合成: MLPは三角関数の恒等式を適用して、離散対数の和（実質的に $\alpha + \beta$）を計算する。
4. デコーディング: モデルは、計算された対数空間における和と整合させることで、候補となる出力をスコアリングする。そこでは、正しい答えが主要な周波数全体にわたって建設的干渉を実現する。

意義と主張
本研究の主要な貢献は、「密なスペクトル」として観察されていたものは、データを誤った基底で分析したことによるアーティファクトであり、アルゴリズムの複雑さを反映したものではないことを示した点にある。分析基底を（標準的な加法的DFTではなく）対象となる群の代数的構造に一致させることで、従来のツールではノイズと見なされていた場所に、解釈可能な構造を明らかにできる。

この手法の汎用性については、以下の点が示されている：

• 疎な乗法的スペクトルは、10種類の異なる素数（$p=59$ から $113$）にわたって確認された。
• メカニズムはランダムシードによらず普遍的である（約80%のシードで成立）。
• 予備的な結果は、2層のトランスフォーマーを用いることで、このアプローチが剰余冪（$a^b \mod p$）にも拡張可能であることを示唆している。

著者らは、あらゆる代数的タスクにおいて、標準的な加法的DFTではなく、関連する群上のフーリエ変換を適用することが、標準的なツールでは構造が見えなかった場所で解釈可能な構造を明らかにするために不可欠であると強調している。