脳科学がAIの解釈性を向上させる: BIMT(Brain-Inspired Modular Training)

## 要約
- 脳の神経ネットワークから着想を得た、ニューラルネットワーク（NN）をモジュール化させ、スパースになるように促す新しい訓練方法、Brain-Inspired Modular Training (BIMT)を提案
- BIMTを用いることでNNの解釈性向上


![BIMT1.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/367eb6b4-23eb-ba53-e5af-e9a6a4b09fde.gif)


## 導入
### 解決したい問題
NNの解釈性の問題
> ディープニューラルネットワーク（DNN）は大きな成功を収めたが、メカニズム的に解釈することは依然として非常に難しい（Olah et al., 2020; Olsson et al., 2022; Michaud et al., 2023; Elhage et al.） 

### 着想
脳に着想を得て、NNをモジュールに分解することで解釈性が上がるのではないか
> NNをより小さなモジュールに分解することができれば（Olah et al.、2020）、解釈可能性はより容易になるかもしれない。

### 背景知識
脳の神経ネットワークと人工のNNには違いがあります。

1. モジュール
脳はモジュール化されている。（Bear et al., 2020）
1. ニューロンの間の距離
人工のNNは、層内及び層間のニューロンの位置関係は考慮していません。
一方で脳の場合は2つの生体ニューロンを接続するコストは、その距離がどれだけ離れているかに依存します。
また軸索はこの距離を移動する必要があるため、エネルギーと脳容積を消費し、時間遅れを引き起こします。

筆者らは、モジュールとニューロンの間の距離について考慮してアルゴリズムを開発しています。

### 新規手法の概要
よりモジュール化された解釈可能なNNの発見を促進するため、BrainInspired Modular Training（BIMT）を提案しました。脳から着想を得て、距離が定義された幾何学的空間にニューロンを埋め込み、各ニューロンの接続長に比例するコストとして損失関数に追加しています。


## 手法(BIMT)
### 手法の概要と詳細
手法の概念図は以下です。

![ 2023-06-18 23.52.43.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/28747980-61c9-6258-93d3-aed2214a963e.png)

(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を用いることで、通常よりも少し精度は下がりますが、大きな精度低下はありませんでした。
精度が低下する理由は、正則化項によるものです。

![ 2023-06-19 0.23.14.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/c25cb645-1085-18f0-a63b-0016d4b33da2.png)

次に式の回帰と２値分類の学習過程をgifで示します、

![BIMT2.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/a034c9d8-2cfc-0db3-159a-45a46d5f1b94.gif)

![BIMT1.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/f81ab432-02de-874a-afb1-87419d0191c0.gif)


図の赤線が負の重み、青線が正の重みです。
学習が進行するにつれて(正則化の効果により)、重みの数が減少していることがわかります。(重みの大きさが限りなく0に近づいている）
また残った重みには局所性があることがわかります。

学習後のNNをミクロで見ると、どのニューロンが活性化しているか、入力を促進しているか抑制しているかがわかり、マクロで見ると、学習によってどのような構造が出現しているかがわかります。


MNISTの学習過程について以下に示します。

![ 2023-06-19 10.03.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/5def2407-2a4f-a410-8c60-b072c86f1af6.png)

MNISTの場合でも同様に、学習が進行するにつれて重みの数が減少し、局所性が生じていることがわかります。

加えて注目すべきは、BIMTが常に0になる周辺画素を刈り取ることを学習するため、入力層の受容野が縮小していることがわかります。BIMTは入力の局所性は保証していないので興味深い結果といえます。(CNNは入力の局所性をカーネルで保証しています。)

さらに層ごとに比較してみると、ある層では負の重みの数が多く、ある層では正の重みの数が多いことがわかります。これは学習器が、パターンマッチング/ミスマッチングの戦略をとっていると解釈できます。

最後に最終層を拡大して見てみます。

![ 2023-06-19 10.50.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/487820/92c85480-390e-10a0-ef18-e93bd16e537c.png)



少し見にくいですが、学習の進行によって(とりわけニューロン(ユニット)の入れ替えによって)「1と7ラベル」「0と8のラベル」が近づいていることがわかります。これは「1と7ラベル」「0と8のラベル」のラベルが類似しているためだと考えられます。



## 結論
脳の神経ネットワークから着想を得た、ニューラルネットワーク（NN）をモジュール化させ、スパースになるように促す新しい訓練方法、Brain-Inspired Modular Training (BIMT)を提案しました。


いくつかの比較的小規模なタスクでテストした結果、BIMTは解釈可能な洞察を与える能力があることが示されました。

## 今後の展望
筆者らが考えている今後の展望は以下です。


- BIMTがより大規模なタスク、例えば大規模言語モデル（LLM）に対しても有効であるかどうかを確認したい。
- 解釈可能なNNを構築することで、AIをより制御しやすく、より信頼性の高い、より安全なものにしたい。


## 備考
- 論文

https://arxiv.org/abs/2305.08746

- GitHub

https://github.com/KindXiaoming/BIMT

- 著者
筆頭著者の方は、MITのPhDである。

https://kindxiaoming.github.io/