要約
- 脳の神経ネットワークから着想を得た、ニューラルネットワーク（NN）をモジュール化させ、スパースになるように促す新しい訓練方法、Brain-Inspired Modular Training (BIMT)を提案
- BIMTを用いることでNNの解釈性向上
導入
解決したい問題
NNの解釈性の問題
ディープニューラルネットワーク（DNN）は大きな成功を収めたが、メカニズム的に解釈することは依然として非常に難しい（Olah et al., 2020; Olsson et al., 2022; Michaud et al., 2023; Elhage et al.）
着想
脳に着想を得て、NNをモジュールに分解することで解釈性が上がるのではないか
NNをより小さなモジュールに分解することができれば（Olah et al.、2020）、解釈可能性はより容易になるかもしれない。
背景知識
脳の神経ネットワークと人工のNNには違いがあります。
- モジュール
脳はモジュール化されている。（Bear et al., 2020） - ニューロンの間の距離
人工のNNは、層内及び層間のニューロンの位置関係は考慮していません。
一方で脳の場合は2つの生体ニューロンを接続するコストは、その距離がどれだけ離れているかに依存します。
また軸索はこの距離を移動する必要があるため、エネルギーと脳容積を消費し、時間遅れを引き起こします。
筆者らは、モジュールとニューロンの間の距離について考慮してアルゴリズムを開発しています。
新規手法の概要
よりモジュール化された解釈可能なNNの発見を促進するため、BrainInspired Modular Training（BIMT）を提案しました。脳から着想を得て、距離が定義された幾何学的空間にニューロンを埋め込み、各ニューロンの接続長に比例するコストとして損失関数に追加しています。
手法(BIMT)
手法の概要と詳細
手法の概念図は以下です。
(1): NNのニューロン(ユニット)一つ一つを幾何学的空間に埋め込みます。その理由は導入で述べたニューロンの間の距離を考えるためです。
(2): 幾何学的空間(ここではユークリッド空間)に基づいて距離(ここではマンハッタン距離)を計算し、距離が遠ければ遠いほどペナルティを付与します。
(3): (2)で計算した距離に基づくペナルティを小さくするために、NNのニューロン（ユニット）を入れ替えます。(2)のペナルティについて局所極小値にはまってしまった場合でも、入れ替え（離散探索）ならそれを回避できます。
また加えて筆者らは人間の脳はモジュール化されており、スパースであることが、効率的である理由であることは間違いないとして、L1正則化項を加えています。
つまりBIMTを一言で説明すると、L1正則化+幾何学的空間に基づいた距離による正則化項+ニューロン(ユニット)の入れ替えといえます。
L1正則化を導入することでNNがスパースとなるため、ニューロン(ユニット)の入れ替えは、(幾何学空間の中での)局所性を高める効果を持ちます。
BIMTの強み
NNをミクロに見るとそれぞれのニューロン(ユニット)がどのくらい有用かを判断することができ、マクロに見るとニューロンがどのような構造を持っているかを判断できる点です。つまりNNの解釈性が向上するという点です。
実験
実験手法
タスク
比較的小規模なタスク(式の回帰、２値分類、MNISTの10値分類)
アーキテクチャ
3-5層の全結合層から構成されるNN(Appendix含む)
パラメータの初期分布
一様分布 (GitHub参照、pytorchが提供しているnn.Linearでデフォルトで初期値している)
誤差関数
MSE(回帰)、クロスエントロピー(分類)
最適化手法
AdamW
結果
BIMTを用いることで、通常よりも少し精度は下がりますが、大きな精度低下はありませんでした。
精度が低下する理由は、正則化項によるものです。
次に式の回帰と２値分類の学習過程をgifで示します、
図の赤線が負の重み、青線が正の重みです。
学習が進行するにつれて(正則化の効果により)、重みの数が減少していることがわかります。(重みの大きさが限りなく0に近づいている）
また残った重みには局所性があることがわかります。
学習後のNNをミクロで見ると、どのニューロンが活性化しているか、入力を促進しているか抑制しているかがわかり、マクロで見ると、学習によってどのような構造が出現しているかがわかります。
MNISTの学習過程について以下に示します。
MNISTの場合でも同様に、学習が進行するにつれて重みの数が減少し、局所性が生じていることがわかります。
加えて注目すべきは、BIMTが常に0になる周辺画素を刈り取ることを学習するため、入力層の受容野が縮小していることがわかります。BIMTは入力の局所性は保証していないので興味深い結果といえます。(CNNは入力の局所性をカーネルで保証しています。)
さらに層ごとに比較してみると、ある層では負の重みの数が多く、ある層では正の重みの数が多いことがわかります。これは学習器が、パターンマッチング/ミスマッチングの戦略をとっていると解釈できます。
最後に最終層を拡大して見てみます。
少し見にくいですが、学習の進行によって(とりわけニューロン(ユニット)の入れ替えによって)「1と7ラベル」「0と8のラベル」が近づいていることがわかります。これは「1と7ラベル」「0と8のラベル」のラベルが類似しているためだと考えられます。
結論
脳の神経ネットワークから着想を得た、ニューラルネットワーク（NN）をモジュール化させ、スパースになるように促す新しい訓練方法、Brain-Inspired Modular Training (BIMT)を提案しました。
いくつかの比較的小規模なタスクでテストした結果、BIMTは解釈可能な洞察を与える能力があることが示されました。
今後の展望
筆者らが考えている今後の展望は以下です。
- BIMTがより大規模なタスク、例えば大規模言語モデル（LLM）に対しても有効であるかどうかを確認したい。
- 解釈可能なNNを構築することで、AIをより制御しやすく、より信頼性の高い、より安全なものにしたい。
備考
- 論文
- GitHub
- 著者
筆頭著者の方は、MITのPhDである。