# ニューラル数式ソルバーにおける途中結果の追跡と操作

松本悠太 ${ }^{1}$ Benjamin Heinzerling ${ }^{2,1}$ 吉川将司 ${ }^{1}$ 乾健太郎 1,2
1 東北大学 2 理化学研究所
yuta.matsumoto.q8@dc.tohoku.ac.jp benjamin.heinzerling@riken.jp
yoshikawa@tohoku.ac.jp kentaro.inui@tohoku.ac.jp

## 概要

言語モデルのより深い理解のためには、「モデル 内部でどのような処理が行われているか」という観点も重要である。我々は単純な数式とその途中結果 に着目することで、Transformer モデルが複数ステッ プに及ぶ処理を行っているかを検証する。途中結果 の情報が符号化されている箇所を特定する Tracing と、符号化されている箇所の状態を操作してモデル に対して因果的介入を行う Manipulation の二つの実験を行なった結果、内部表現の特定の方向が線形に 近い形で途中結果を符号化していること、そしてそ のような方向がモデルの推論結果に対して因果的に も関係していることを示す。

( www.github.com/cl-tohoku/trace-manipulate

## 1 はじめに

近年の言語モデルは数量推論のような複雑な処理 が必要だと思われる問題においてもある程度高い性能を出せることが知られている $[1]$ 。また、複数の 先行研究では線形代数や初等数学のような数値計算 のタスクを通じて言語モデルが持つ潜在的な能力を 測っており [2,3]、現在のモデルによって解ける問題や解けない問題の性質が分かりつつある。

一方で、モデルがこのような問題を解く際、その 内部で何が行われているかについての研究はほとん ど行われていない。しかし、モデルの能力をよりよ く理解するためには、入出力の結果だけでなくその 内部での処理を分析することも非常に重要だと考え られる。我々は過去にモデルが途中結果を介した複数ステップの推論を行なっているかを調べるため に、四則演算を学習した Transformer [4] の隠れ層を 分析し、四則演算の途中結果 ${ }^{1}$ がモデル内部の特定 の方向に保存 (符号化) されていることを示した [5]。
 は 154,290,116, 223 などである。本研究ではこの研究を拡張し、途中結果とモデル 内部状態の相関関係だけでなく、因果関係も調査 する。具体的には、四則演算を学習した Transformer の内部状態から主成分分析を用いて途中結果と相関 の高い、すなわち途中結果を符号化しているような 方向を抽出したのち、その部分のアクティベーショ ンを操作してモデルの予測結果の変化を観察する。 これによって、途中結果を符号化している部分がモ デルの推論に与える因果的な影響を調査する。その 結果、途中結果を符号化しているような方向の中に は実際にモデルの推論時にも使用されているものが あることを確認できた。また、相関が高くてもモデ ルの推論には使用されていない方向も確認されたた め、モデルの内部を分析する際に因果的な関係を調査することの重要性を示唆する結果となった。

本論文の貢献は、Transformer モデルの内部でどの ような処理が行われているのかを因果推論まで含め て詳細に分析できた点である。

## 2 関連研究

言語モデルと数値 複数の先行研究は、近年の言語モデルが数学的な問題をある程度解けることを 示している。Geva ら [1] は大規模言語モデルを使用 することで、数量推論タスクでほぼ最先端の性能を 達成できることを示した。また、数值計算に特化さ せた Transformer は初等数学や線形代数といった記号推論の問題をある程度解けることも示されてい る $[2,3]$ 。彼らの知見に基づき本研究では、モデル に解ける問題として四則演算を取り上げる。

モデル内部の分析 モデルの内部表現にどのよう な情報が符号化されているかについては、近年高い 関心が集まっている。このような研究の多くはプ ロービングの文脈で行われているが [6, 7]、Shibata ら [8] は我々と似た手法で実験を行っている。彼ら は形式言語で訓練した言語モデルを分析して、形式言語の深さと高い相関を持つアクティベーションの

図 1: 途中結果の追跡 (Tracing) と操作 (Manipulation) の概要。Tracing では PCAを用いて、数式中の途中結果 と相関が高いモデル内の方向を追跡する。Manipulation では、Tracing で追跡された方向に沿ってモデルのア クティベーションを操作した際のモデル予測の変化を観察する。

存在を示した。また、Elzar ら [9] はプロービングに よって抽出できる特徵は必ずしもモデルの推論には 使われていないことを明らかにし、内部表現への因果的介入を行う重要性を示した。本研究ではモデル の隠れ層を分析対象として、途中結果の情報を符号化している部分を抽出し、アクティベーションを連続的に変化させることで抽出された情報がモデルの 推論に使用されているかを調査する。

## 3 手法

本節では図 1 に示した本研究の提案手法について 記述する。まず四則演算のデータで Transformer を 訓練した後、下記に示す 2 つの方法でモデルの内部状態と数式の途中結果の関係を調査する。

## 3.1 途中結果の追跡 (Tracing)

数式 (大力)を変化させると、出力と同時に隠れ 層の表現も変化する。では、この表現は数式中の途中結果とどのような関係があるのだろうか。汎化し たモデルの表現空間はタスクの特徴を捉えたよう な構造になる [10] ことを考慮すると、数式中の途中結果は表現空間になんらかの形で符号化されて いるはずである。我々は一番単純な符号化の形とし て、表現空間中の特定の方向と途中結果の関係を調 べる。まず、モデルの各層に対してその表現空間中 で支配的な方向を見つける。モデルの $l$ 層目, $j$ 番目のトークンにおける表現 $h_{j}^{l}$ を全て連結した表現 $H^{l}=h_{1}^{l} \oplus h_{2}^{l} \oplus \cdots \oplus h_{n}^{l}$ に対して主成分分析 (PCA)を 行うことで、主成分 $p_{k}^{l}, k \in[1, \ldots, K]$ を得ることが

中のあるインスタンス $i$ に対してこの PCA モデル を適用すると、各主成分の重み $p_{i, k}^{l}$ を得ることが できる。次に、数式中の途中結果 $R_{i}^{j}$ と各主成分 $p_{k}^{l}$ の重みの相関 $\operatorname{corr}\left(R_{i}^{j}, p_{i, k}^{l}\right)$ を測る。これを各層に おける全ての途中結果と主成分の直積に対して行


(a) 1 層目

(b) 3 層目

(c) 4 層目

図 2: 層ごとの各主成分重みと数式の途中結果の相関関係のヒートマップ。各セルは第 $k$ 主成分の重み（列） とある途中結果（行）との相関係数の絶対値を表す。

うことで、ある途中結果と最も相関の高い主成分で ある最大相関方向 $\hat{p}_{k}^{l}\left(R^{j}\right):=\operatorname{argmax}_{k}\left(\operatorname{corr}\left(R_{i}^{j}, p_{i, k}^{l}\right)\right)$ を得る。仮に最大相関方向の相関が十分高ければ、対応する途中結果を符号化していると考えることが できる。例えば、図 1 のイメージは $a-(b-c)$ とい う数式における途中結果 $a$ や $b-c$ の値がある層の ある主成分に符号化されていることを表す。

## 3.2 途中結果の操作 (Manipulation)

本節では、Tracing によって得られた最大相関方向 がモデルの推論に使用されているのかを検証するた め、モデルに対して因果的介大を行う。具体的には 図 1 下部に示すように、主成分に沿ってモデルのア クティベーションを操作し、それによるモデルの予測結果の変化を観察する。定式化して表すと、 $l$ 層目の表現 $H^{l}$ に対して次の式に基づいて主成分 $p_{k}^{l}$ の 重みを $r$ 倍することで、操作後の表現 $H^{l^{\prime}}$ を得る。

$
H^{l^{\prime}} \leftarrow H^{l}+(r-1)\left(p_{k}^{l^{\top}} H^{l}\right) p_{k}^{l}
$

直観的には $r$ を変化させることで主成分 $p_{k}^{l}$ に沿っ てアクティベーション $H^{l}$ が移動し、それに伴って モデルの予測結果も変化する。

もし最大相関方向 $\hat{p}_{k}^{l}\left(R^{j}\right)$ が因果的にも途中結果 $R^{j}$ を符号化しているのであれば、 $p_{k}^{l}$ の重みを動か すことは数式の途中結果 $R^{j}$ を動かすことに対応し ていると考えられる。すなわち、 $R_{i}^{j}=f\left(p_{i, k}^{l}\right)$ とな るような関数 $f$ を定義できる。この時、逆関数を取 ることで途中結果から主成分の重みを予測する関数 $f_{p}=f^{-1}$ を定義することができる。「 $\hat{p}_{k}^{l}\left(R^{j}\right)$ が因果的にも途中結果 $R^{j}$ を符号化している」という仮説 を検証するためには、この $f_{p}$ が正しいかを確かめ れば良い。すなわち、実際に入力項を変更した際の 最大相関方向 $\hat{p}_{k}^{l}\left(R^{j}\right)$ の重みを $p_{i, k}^{l}=f_{a}\left(R_{i}^{j}\right)$ のよう に表せる実測関数 $f_{a}$ と $f_{p}$ を比較して確かめる。

## 4 実験設定

ニューラル数式ソルバーとして、Sajjad ら [11] を参考に 6 層の Transformer エンコーダを用意し、 [CLS] トークンから線形回帰を行う。足し算と引き 算、括弧を含む四則演算のデータをテンプレートか ら 20 万個自動生成し、3) そのうち 19 万個のデータで モデルを学習する。先行研究 [1] に従い、入力中の 数字は 1 桁ごとにトークン化する。例えば、“ 123 ” という数字は $1, \# \# 2, \# \# 3$ にトークン化される。訓練後、 1 万個の検証データにおいてこのモデルを評価 したところ決定係数 $R^{2}=0.9988$ で、異なるパター ンの数式を並列に学習させてもほぼ正しい答えを出力できることが確認された。

## 5 結果

## 5.1 途中結果の追跡

今回は $a-(b-c)$ という数式を分析の対象とする。各途中結果と上位 10 主成分の重みの関係を測定し た結果を図 2 に示す。図 2(a), 2(b) から、中間層にお いて数式中の途中結果と相関が非常に高い最大相関方向が複数存在することがわかる。例えば、 1 層目 の第 3 主成分の重みと $b$ の値の相関係数は $0.96 、 3$層目の第 2 主成分の重みと $b-c$ の値の相関係数は 0.94 である。

ここから、数式の途中結果はモデル内部の表現空間における特定の方向に線形的、局所的に符号化さ れていると考えることができる。

一方で図 2(c) からは $p_{1}^{4}$ が最終結果とほぼ 1 の相関を持っていることがわかり、 4 層目の時点で数値計算処理が終了していると推察できる。
程式の種類は全部で 87797 種類である。

(a) 重み操作によるモデル予測とそこから導いた途中結果の変化。

(b) 途中結果と主成分の重みに関する実測関数 $f_{a}(b)$ と予測関数 $f_{p}(b)$ 。

図 3: $\hat{p}_{3}^{1}(b)$ に対する Manipulation の結果。影の部分 はデータセット中に現れない重みの範囲を表す。

## 5.2 途中結果の操作

数式 $617-(555-602)$ とその途中結果 $b=555$ を 例に取り、Manipulationを行う。5.1 節から、数式 $a-(b-c)$ において、最大相関方向 $\hat{p}_{3}^{1}(b)$ は 0.96 と いう高い相関を持つ。ここで、 $p_{3}^{1}$ に沿ってモデル のアクティベーションを操作すると、モデルの予測結果は図 3(a) に示されるように、元の予測結果 (664) から約 100-1200 まで変化することがわかる。 この時、予測結果の変化をモデル内の途中結果 $b$ の表現を変えたものによるものだと仮定すると、図 3(a) の右軸を得ることができ、軸を反転するこ とで $p_{i, 3}^{1}=f_{p}(b)$ となるような関数 $f_{p}$ が得られる。 この予測関数 $f_{p}$ を、実際に $b$ の値を変化させた時 の重みを $p_{i, 3}^{1}=f_{a}(b)$ と表せるような実測関数 $f_{a}$ と 比較した結果を図 3(b) に示す。図 3(b) から、予測関数 $f_{p}$ と実測関数 $f_{a}$ は概して一致しており、特に 元データの周りでは一致度が高いことがわかる。こ れは、 $\hat{p}_{3}^{1}(b)$ が $b$ の値を一定の範囲では因果的にも 符号化している証拠だと考えることができる。同じ 実験を異なる値を持つ $a-(b-c)$ の 1000 インスタ

(a) $\hat{p}_{3}^{1}(b)$ に対する Manipulation の結果。

(b) $\hat{p}_{2}^{4}(b)$ に対する Manipulation の結果。

図 4: 1000 個のデータに対して重みの操作を行なっ た際の誤差 $\left|f_{p}(b)-f_{a}(b)\right|$ の中央値。

ンスで行った時の $\left|f_{p}(b)-f_{a}(b)\right|$ の中央值を図 4(a) に示す。図 4(a) から、元データの表現から離れると 必ずしも $\hat{p}_{3}^{1}(b)$ が $b$ の值を因果的に符号化している とは言えないことが定量的にも観察される。逆に、

Manipulation によって因果的には途中結果を符号化 していない最大相関方向を特定することもできる。例えば、 $\hat{p}_{2}^{4}(b)$ は $b$ と 0.81 の高い相関を持つが、こ の方向に対して重みの操作を行なった際の誤差中央値は図 4(b) のようになり、誤差が非常に大きいこと が観測される。これは図 2 から推測される「4 層目 までで途中結果を使用した数値計算処理はほぼ完了 している」という仮説と対応するものである。

## 6 おわりに

本研究では言語モデルが数値計算を行う時に内部で行なっている処理を明らかにするため、数式中 の途中結果に着目し追跡と操作を行なった。その結果、「数値計算ができるモデル」の内部では途中結果の情報が存在し、かつその情報はモデルの推論時 にも使われていることを解明した。本研究の手法は アクティベーションを連続的に変化させることでモ デルに対して詳細な因果推論を行うことができる。今後はこの手法をより一般的な数量推論モデルにも 適用し、分析道具としての有用性を示したい。

## 謝辞

本研究は JST CREST JPMJCR20D2、JSPS 科研費 21K17814 の助成を受けたものです。

## 参考文献

[1] Mor Geva, Ankit Gupta, and Jonathan Berant. Injecting numerical reasoning skills into language models. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 946-958, Online, July 2020. Association for Computational Linguistics.

[2] David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. In International Conference on Learning Representations, 2019.

[3] François Charton. Linear algebra with transformers, 2021.

[4] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob 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. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, Vol. 30. Curran Associates, Inc., 2017.

[5] 松本悠太, 吉川将司, Benjamin Heinzerling, 乾健太郎.四則演算を用いた transformer の再帰的構造把握能力 の調査. 言語処理学会年次大会, 2022.

[6] Ganesh Jawahar, Benoît Sagot, and Djamé Seddah. What does BERT learn about the structure of language? In Proceedings of the 57th Annual Meeting of the Association 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 Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 287-297, Brussels, Belgium, November 2018. Association for Computational Linguistics.

[8] Chihiro Shibata, Kei Uchiumi, and Daichi Mochihashi. 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 Computational Linguistics.

[9] Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav Goldberg. Amnesic probing: Behavioral explanation with amnesic counterfactuals. Transactions of the Association for Computational Linguistics, Vol. 9, pp. 160175, 2021.

[10] Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J Michaud, 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- tion Processing Systems, 2022.

[11] Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of pre-trained transformer models, 2021.

## A 四則演算データセットの詳細

四則演算データは、最大で 5 段の深さを持つ。各項 $x$ の大きさは、訓練、検証データ中では $0<x<1000$ 、 5.1, 5.2 節での分析に使用したデータでは桁数を揃えるために $100<=x<1000$ とした。訓練に使用した数式 の例を表 1 に示す。

表 1: 四則演算データセットの例

## B ニューラル数式ソルバーの詳細

ニューラル数式ソルバーは表 2 に示す設定で訓練を行った。

表 2: ハイパーパラメータの一覧

4) https://huggingface.co/bert-base-uncased

5) https://pypi.org/project/pytorch-pretrained-bert
 松本悠太 ${ }^{1}$ Benjamin Heinzerling ${ }^{2,1}$ 吉川将司 ${ }^{1}$ 乾健太郎 1,2 1 東北大学 2 理化学研究所 yuta.matsumoto.q8@dc.tohoku.ac.jp benjamin.heinzerling@riken.jp yoshikawa@tohoku.ac.jp kentaro.inui@tohoku.ac.jp  