【DL輪読会】Grokking Group Multiplication with Cosets

DEEP LEARNING JP [DL Papers] Grokking Group Multiplication with Cosets Ryutaro Yamauchi, Matsuo Iwasawa Lab http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 Grokking Group Multiplication with Cosets 著者 出典 ：Dashiell Stander, Qinan Yu, Honglu Fan, Stella Biderman (EleutherAI) ：https://arxiv.org/abs/2312.06581 内容： – 置換群の演算で grokking を起こした NN をリバースエンジニアリングした – 剰余類回路 (coset circuit) を発見した 選定理由： – Grokking 時に何が生じているのか気になっていたため – 個人的に馴染みのない研究だったため 2

• https://arxiv.org/abs/2312.06581

Grokking [1] 算術的データセットで訓練されたモデルが突然ルールを理解（grok）し， テストデータに完全に汎化する現象． データセットの例． • a=0, b=2, c=1, d=4, e=3 • X★Y = X^2+XY+Y^2+X (mod5) 1. Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets https://arxiv.org/abs/2201.02177 3

貢献 1. 置換群 𝑆5 , 𝑆6 において訓練された隠れ層一つの全結合NNをリバース エンジニアリングする． 2. 群フーリエ変換を回路レベルの機械的解釈に応用． 3. 因果的介入を用いて、提案する回路の全ての特性を徹底的にテスト する． 4. 機械的解釈可能性に関する現在の研究を調査し，我々の研究の困難 さと、この分野におけるより広範な課題との関連性を示す． 4

置換と対称群 要素を並べ替える操作を置換 (permutation) と呼ぶ． 例えば4つの要素の順序を逆にする置換 𝜎 = (4 3 2 1) によって と置換される． 置換 𝜏 = (3 2 1 4) で並べ替えたのち 𝜎 で並べ替えるとすれば，これを置換の積 𝜎𝜏 と考えることができ， より 𝜎𝜏 = (4 1 2 3) である． 置換と置換の積はまた置換であり，単位元や逆元も考えられるから，𝑛 要素の置換 の全体は群をなす．このような群を対称群 𝑆𝑛 という． 5

剰余類 (coset) とは 𝐺 が群で，𝐻 がその部分群，𝑔 は 𝐺 の元とする．このとき， 𝑔𝐻 = {𝑔ℎ: ℎ ∈ 𝐻} を 𝐺 における 𝐻 の左剰余類 (left coset) といい， 𝐻𝑔 = {ℎ𝑔: ℎ ∈ 𝐻} を 𝐺 における 𝐻 の右剰余類 (right coset) という． [参考] 剰余類：https://ja.wikipedia.org/wiki/剰余類 6

• https://ja.wikipedia.org/wiki/

剰余類 (coset) とは | 具体例で考える 𝑆3 について考える． という具合で，𝐻1 の左剰余類は の３つである（𝜏(𝑖𝑗) は 要素 𝑖, 𝑗 のみを入れ替え ここで を考えると， る置換（互換））． これ自体が群になっている（𝑆3 の部分群）． 要素 1 を固定したものと思えば，これは 𝑆2 と準 これらは共通部分を持たないから，𝑆3 を分解す 同型である． るものになっている． 𝐻1 に 𝑆3 の元をそれぞれ左から作用させると 同様に他の部分群 についても同様に左剰余類を考えることができ る． 7

NNが置換群の積を計算する仕組み 上記の表からわかるように，どの剰余類に入るかを特定できれば，その元を特定できる．したがって，剰 余類 𝐻 に含まれる元 𝜎 と剰余類 𝐾 に含まれる元 𝜏 の積 𝜎𝜏 が，どの剰余類に含まれるかの情報を持ってい れば，次のような手順で計算結果を求めることができる． 1. 𝜎, 𝜏 がどの剰余類に含まれるかを調べる 2. 剰余類同士の積についての情報から，積 𝜎𝜏 がどの剰余類に含まれるかを調べる 3. 2 で得られた条件を満たす元を求める ※ 剰余類同士の積が常に剰余類になるわけではないが，決定できる場合がある．そしてすべてが決定できる必要はな く，たとえば 𝐻1 と 𝐻2 にともに含まれる元だということが分かれば (1 2 3) だと判断できる． 8

NNアーキテクチャ 実験で用いたのは右図のようなNN． – 入力 𝑥𝑔 は長さ 𝐺 の one-hot ベクトル． – Embedding layer: 𝐄𝑙 , 𝐄𝑟 (𝑑, 𝐺 ) – Liner layer: 𝐖 = [𝐋 𝐑] (𝑤, 2𝑑) – Unembedding layer: 𝐔 ( 𝐺 , 𝑤) 9

符号回路 sign circuit 𝑆𝑛 上の偶置換全体の集合を 𝐴𝑛 とすると，これも部分群であり，剰余類を持つ．𝐴𝑛 の剰余類は 𝐴𝑛 それ自身と奇置換 𝜏𝐴𝑛 である．ここで置換 𝜎 の符号を と定義する． この時 sgn(𝜎𝜏) = sgn(𝜎)sgn(𝜏)である．この関係を NN は以下のよう に表現する． 1. 入力された 𝜎, 𝜏 に対して 𝑥sgn(𝜎), −𝑥sgn(𝜏) という値を作る． • 右図の赤いところ 2. 二つの値を足し合わせる．結果が 0 なら 𝜎𝜏 は 𝐴𝑛 に属す （偶置換）と判定する．そうでない場合 (2𝑥, −2𝑥 の場合) は 奇置換となる． 3. ただし，ReLU がある場合 2𝑥, −2𝑥 のどちらかは 0 になり 偶置換と混同してしまう．この場合 𝑥 の正負が異なる同じ 回路が２つあれば問題は回避できる． 10

剰余類回路 coset circuit 𝐴𝑛 以外の部分群を法とする剰余類の場合も，基本的な方法は同じ． ※しかし coset 間のすべての積が別の coset になるわけではない点に 注意が必要（右下図）． 剰余類回路の構成： • Embedding : 置換 𝜎 を受け取って，それがどの剰余類に含まれ ているかを示すベクトルを作る． • Linear：符号回路で言うところの 𝑥sgn(𝜎), −𝑥sgn(𝜏) を様々な 剰余類に対して作る． 例： – 左■：𝜎 ∈ 𝐻1 のとき 𝑥, 𝜎 ∉ 𝐻1 のとき −𝑥 – 右■：𝜏 ∈ 𝐻5 𝜏 15 のとき − 𝑥, それ以外で 𝑥 この時，■を足し合わせた値が 0 ならば 𝜎𝜏 ∈ 𝜏 15 𝐻1 • Unembedding： 𝜎𝜏 が含まれる剰余類一覧が得られて いるので，それを集約して logit を出力する． 11

剰余類回路の発見と同定 • NNに対して群フーリエ変換を適用していたら見つかった • NN内のあるニューロンを関数 𝑓 ∶ 𝐺 → ℝ とみて，以下の指標 𝐶𝐻 𝑓 を計算 – 特定の剰余類に属する置換を 𝑓 に入力した場合のニューロンの活性値の分散 / 全置換を入力したときのニューロンの活性値の分散． – これが十分小さいとき，そのニューロンは剰余類回路になっている． 12

剰余類回路の根拠 1. 剰余類回路に関係するニューロンを削除すると精度が悪化する が，それ以外のニューロンを削除しても精度は下がらない． 2. NNに対して様々な介入を行い，実験結果と剰余類回路の性質か ら予測される結果と比較する – Embedding Exchange : 悪化するはず – Switch Permutation Sign : 両方変えれば変化しないはず – Absolute Value Non-linearity：良くなるはず • – ReLU を絶対値関数に置き換える Distribution Change：0 周りで非対称なら悪化する • ノイズを加える 13

機械的解釈の方法論 1. タスクと相関のある参照データセットまたは分布を選ぶ． 2. モデルの計算グラフの中で，神経回路として機能するノードのサブ セットを特定する． 3. この回路がタスクのパフォーマンスの大部分を占めていることを確 認する． 4. 回路の機能を忠実に抽象化した回路モデルを考える． 5. モデルの動作がデータ分布全体にわたって忠実であることを因果的 にテストする． 14

まとめと感想 • 置換の積に対して grokking した NN を解析すると「剰余類回路」が見出される • 剰余類回路モデルに期待される振舞いと実際のNNの振る舞いを比較することで， 回路が実際に機能していることを確かめた – 介入の仕方がこれでいいのか疑問 • 剰余類回路のようなものが勝手に生じるのは面白い • 対象と対象の関係を考えるときに，対象を性質の束に還元した上で，性質と性質 の関係を考える，というのはとても自然なようにも思われる 15