概要

1　はじめに
近年の言語モデルは数量推論のような複雑な処理が必要だと思われる問題においてもある程度高い性能を出せることが知られている[1]。また、複数の先行研究では線形代数や初等数学のような数値計算のタスクを通じて言語モデルが持つ潜在的な能力を測っており[2, 3]、現在のモデルによって解ける問題や解けない問題の性質が分かりつつある。一方で、モデルがこのような問題を解く際、その内部で何が行われているかについての研究はほとんど行われていない。しかし、モデルの能力をよりよく理解するためには、入出力の結果だけでなくその内部での処理を分析することも非常に重要だと考えられる。我々は過去にモデルが途中結果を介した複数ステップの推論を行なっているかを調べるために、四則演算を学習したTransformer [4]の隠れ層を分析し、四則演算の途中結果りがモデル内部の特定の方向に保存(符号化)されていることを示した[5]。本研究ではこの研究を拡張し、途中結果とモデル内部状態の相関関係だけでなく、因果関係も調査する。具体的には、四則演算を学習したTransformerの内部状態から主成分分析を用いて途中結果と相関の高い、すなわち途中結果を符号化しているような方向を抽出したのち、その部分のアクティベーションを操作してモデルの予測結果の変化を観察する。これによって、途中結果を符号化している部分がモデルの推論に与える因果的な影響を調査する。その結果、途中結果を符号化しているような方向の中には実際にモデルの推論時にも使用されているものがあることを確認できた。また、相関が高くてもモデルの推論には使用されていない方向も確認されたため、モデルの内部を分析する際に因果的な関係を調査することの重要性を示唆する結果となった。本論文の貢献は、Transformerモデルの内部でどのような処理が行われているのかを因果推論まで含めて詳細に分析できた点である。
2   関連研究
言語モデルと数値複数の先行研究は、近年の言語モデルが数学的な問題をある程度解けることを示している。Gevaら[1]は大規模言語モデルを使用することで、数量推論タスクでほぼ最先端の性能を達成できることを示した。また、数値計算に特化させたTransformerは初等数学や線形代数といった記号推論の問題をある程度解けることも示されている[2, 3]。彼らの知見に基づき本研究では、モデルに解ける問題として四則演算を取り上げる。モデル内部の分析 モデルの内部表現にどのような情報が符号化されているかについては、近年高い関心が集まっている。このような研究の多くはプロービングの文脈で行われているが[ 6 , 7 ]、Shibataら[8]は我々と似た手法で実験を行っている。彼らは形式言語で訓練した言語モデルを分析して、形式言語の深さと高い相関を持つアクティベーションのクティベーションを操作した際のモデル予測の変化を観察する。存在を示した。また、Elzarら[9]はプロービングによって抽出できる特徴は必ずしもモデルの推論には使われていないことを明らかにし、内部表現への因果的介入を行う重要性を示した。本研究ではモデルの隠れ層を分析対象として、途中結果の情報を符号化している部分を抽出し、アクティベーションを連続的に変化させることで抽出された情報がモデルの推論に使用されているかを調査する。
3手法
本節では図1に示した本研究の提案手法について記述する。まず四則演算のデータでTransformerを訓練した後、下記に示す２つの方法でモデルの内部状態と数式の途中結果の関係を調査する。
途中結果の追跡(Tracing)
3.1
数式(入力)を変化させると、出力と同時に隠れ層の表現も変化する。では、この表現は数式中の途中結果とどのような関係があるのだろうか。汎化したモデルの表現空間はタスクの特徴を捉えたような構造になる[10]ことを考慮すると、数式中の途中結果は表現空間になんらかの形で符号化されているはずである。我々は一番単純な符号化の形として、表現空間中の特定の方向と途中結果の関係を調べる。まず、モデルの各層に対してその表現空間中で支配的な方向を見つける。モデルのl層目, j番目のトークンにおける表現h ′を全て連結した表現H l = h l ⊕ h l ⊕ ··· ⊕ h lに対して主成分分析(PCA)を行うことで、主成分p k , k ∈ [1,...,K]を得ることができる。nは入力列の長さを表す²)。データセット中のあるインスタンスiに対してこのPCAモデルを適用すると、各主成分の重みp i kを得ることができる。次に、数式中の途中結果R iと各主成分p iの重みの相関corr( R j , p j , k )を測る。これを各層における全ての途中結果と主成分の直積に対して行図2:層ごとの各主成分重みと数式の途中結果の相関関係のヒートマップ。各セルは第k主成分の重み（列）とある途中結果（行）との相関係数の絶対値を表す。うことで、ある途中結果と最も相関の高い主成分である最大相関方向ρ ( R j ) := argmax k (corr( R j , p j , ))を得る。仮に最大相関方向の相関が十分高ければ、対応する途中結果を符号化していると考えることができる。例えば、図1のイメージはa – (b – c)という数式における途中結果aやb − cの値がある層のある主成分に符号化されていることを表す。
3.2 途中結果の操作(Manipulation)
本節では、Tracingによって得られた最大相関方向がモデルの推論に使用されているのかを検証するため、モデルに対して因果的介入を行う。具体的には図1下部に示すように、主成分に沿ってモデルのアクティベーションを操作し、それによるモデルの予測結果の変化を観察する。定式化して表すと、I層目の表現H ′に対して次の式に基づいて主成分p ′の重みをr倍することで、操作後の表現H rを得る。
直観的にはrを変化させることで主成分p kに沿ってアクティベーションH 1が移動し、それに伴ってモデルの予測結果も変化する。もし最大相関方向ρ ( R ′ )が因果的にも途中結果R」を符号化しているのであれば、p 1の重みを動かすことは数式の途中結果R ′を動かすことに対応していると考えられる。すなわち、R i = f (p i k )となるような関数fを定義できる。この時、逆関数を取ることで途中結果から主成分の重みを予測する関数fp = f-¹を定義することができる。 ρ k (R ′ )が因果的にも途中結果R ′を符号化している」という仮説を検証するためには、このƒρが正しいかを確かめれば良い。すなわち、実際に入力項を変更した際の最大相関方向ρ (R i )の重みをρ i k = f α (R i )のように表せる実測関数f aとf pを比較して確かめる。
4
実験設定
ニューラル数式ソルバーとして、Sajjadら[11]を参考に6層のTransformerエンコーダを用意し、 [CLS]トークンから線形回帰を行う。足し算と引き算、括弧を含む四則演算のデータをテンプレートから20万個自動生成し、³）そのうち19万個のデータでモデルを学習する。先行研究[1]に従い、入力中の数字は1桁ごとにトークン化する。例えば、“123”という数字は1, ##2, ##3にトークン化される。訓練後、1万個の検証データにおいてこのモデルを評価したところ決定係数R² = 0.9988で、異なるパターンの数式を並列に学習させてもほぼ正しい答えを出力できることが確認された。
結果
5
途中結果の追跡
5.1
今回はa − (b − c)という数式を分析の対象とする。各途中結果と上位10主成分の重みの関係を測定した結果を図2に示す。図2(a), 2(b)から、中間層において数式中の途中結果と相関が非常に高い最大相関方向が複数存在することがわかる。例えば、1層目の第３主成分の重みとbの値の相関係数は0.96、3層目の第2主成分の重みとb − cの値の相関係数は0.94である。ここから、数式の途中結果はモデル内部の表現空間における特定の方向に線形的、局所的に符号化されていると考えることができる。一方で図2(c)からはp 4が最終結果とほぼ1の相関を持っていることがわかり、4層目の時点で数値計算処理が終了していると推察できる。図3:  ρ 1 ( b )に対するManipulationの結果。影の部分はデータセット中に現れない重みの範囲を表す。
5.2途中結果の操作
数式617 – (555 – 602)とその途中結果b = 555を例に取り、Manipulationを行う。5.1節から、数式a – (b – c)において、最大相関方向β (b)は0.96という高い相関を持つ。ここで、p 1に沿ってモデルのアクティベーションを操作すると、モデルの予測結果は図3(a)に示されるように、元の予測結果(664)から約100–1200まで変化することがわかる。この時、予測結果の変化をモデル内の途中結果bの表現を変えたものによるものだと仮定すると、図3(a)の右軸を得ることができ、軸を反転することでp i 3 = f ρ (b)となるような関数f ρが得られる。この予測関数ƒρを、実際にbの値を変化させた時の重みをp i 3 = f a (b)と表せるような実測関数f aと比較した結果を図3(b)に示す。図3(b)から、予測関数ƒρと実測関数ƒρは概して一致しており、特に元データの周りでは一致度が高いことがわかる。これは、ρ 1 ( b )がbの値を一定の範囲では因果的にも符号化している証拠だと考えることができる。同じ実験を異なる値を持つa − (b − c)の1000インスタンスで行った時の| f ρ ( b ) – f α ( b ) |の中央値を図4(a)に示す。図4(a)から、元データの表現から離れると必ずしもρ 1 (b)がbの値を因果的に符号化しているとは言えないことが定量的にも観察される。逆に、 Manipulationによって因果的には途中結果を符号化していない最大相関方向を特定することもできる。例えば、β 4 (b)はbと0.81の高い相関を持つが、この方向に対して重みの操作を行なった際の誤差中央値は図4(b)のようになり、誤差が非常に大きいことが観測される。これは図２から推測される「4層目までで途中結果を使用した数値計算処理はほぼ完了している」という仮説と対応するものである。
おわりに
6
本研究では言語モデルが数値計算を行う時に内部で行なっている処理を明らかにするため、数式中の途中結果に着目し追跡と操作を行なった。その結果、「数値計算ができるモデル」の内部では途中結果の情報が存在し、かつその情報はモデルの推論時にも使われていることを解明した。本研究の手法はアクティベーションを連続的に変化させることでモデルに対して詳細な因果推論を行うことができる。今後はこの手法をより一般的な数量推論モデルにも適用し、分析道具としての有用性を示したい。本研究はJST CREST JPMJCR20D2、JSPS科研費21K17814の助成を受けたものです。 [1]  Mor Geva, Ankit Gupta, and Jonathan Berant.  Inject- ing numerical reasoning skills into language models.  In Proceedings of the 58th Annual Meeting of the As- sociation for Computational Linguistics , pp. 946–958, Online, July 2020. Association for Computational Linguis- tics. David Saxton, Edward Grefenstette, Felix Hill, and Push- [2] meet Kohli. Analysing mathematical reasoning abilities of neural models.  In International Conference on Learn- ing Representations , 2019. François Charton. Linear algebra with transformers, 2021. [3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob [4] Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vish- wanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , Vol. 30. Curran As- sociates, Inc., 2017.松本悠太,吉川将司, Benjamin Heinzerling,乾健太郎. [5]四則演算を用いたtransformerの再帰的構造把握能力の調査言語処理学会年次大会, 2022. Ganesh Jawahar, Benoît Sagot, and Djamé Seddah.  What [6] does BERT learn about the structure of language?  In Pro- ceedings of the 57th Annual Meeting of the Associ- ation for Computational Linguistics , pp. 3651–3657, Florence, Italy, July 2019. Association for Computational Linguistics. [7] Alessandro Raganato and Jörg Tiedemann. An analysis of encoder representations in transformer-based machine translation.  In Proceedings of the 2018 EMNLP Work- shop BlackboxNLP: Analyzing and Interpreting Neu- ral Networks for NLP , pp. 287–297, Brussels, Belgium, November 2018. Association for Computational Linguis- tics. Chihiro Shibata, Kei Uchiumi, and Daichi Mochihashi. [8] How LSTM encodes syntax: Exploring context vectors and semi-quantization on natural text. In Proceedings of the 28th International Conference on Computational Linguistics , pp. 4033–4043, Barcelona, Spain (Online), December 2020. International Committee on Computa- tional Linguistics. Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav [9] Goldberg.  Amnesic probing:  Behavioral explanation with amnesic counterfactuals.  Transactions of the Associ- ation for Computational Linguistics , Vol. 9, pp. 160– 175, 2021. Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J Michaud, [10] Max Tegmark, and Mike Williams. Towards understanding grokking:  An effective theory of representation learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Informa- [11] Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov.  On the effect of dropping layers of pre-trained transformer models, 2021.四則演算データは、最大で5段の深さを持つ。各項xの大きさは、訓練、検証データ中では0 < x < 1000、 5.1, 5.2節での分析に使用したデータでは桁数を揃えるために100 <= x < 1000とした。訓練に使用した数式の例を表1に示す。
ニューラル数式ソルバーは表2に示す設定で訓練を行った。表2:ハイパーパラメータの一覧