ED法+交差エントロピーをTF/Torchで実装してみた（おまけでBitNet×ED法を検証）

ED法の記事その３です。

その１：金子勇さんのED法を実装してMNISTを学習させてみた
その２：ED法を高速化してその性能をMNISTで検証してみた

もう記事を書く予定はなかったんですが、思ったより頭から離れなかったので記事に起こしました。

今回はED法をTorchとTFで実装し交差エントロピーとReLUを実装しました。
なので正規のMNISTを学習してみたいと思います。
最後にBitNet×ED法を試して考察しています。

本記事を含めたED法のコードをGithubに作成しました。

追加の実装

ED法+交差エントロピー

まずは今までの平均二乗誤差（MSE;Mean Squared Error）のイメージです。

MSEのlossは（正解データ-出力）で表され、それをED法で学習します。
次に交差エントロピーのイメージです。

各出力に対してED法ニューラルネットを作成し、その出力達をsoftmaxで確率に変換します。
変換した確率とOne-hotエンコーディングした正解ラベルの差がlossとなります。
※（正解ラベル-出力確率）とlossが綺麗な形になるのはsoftmax+交差エントロピーの相性がいいからで、必ずこの形になるわけではありません

ED法では各出力に対してとりあえず別々に重みを持つ実装を作りました。
どうにか重みを共有できないか考えていますが、いいアイデアは今のところ思い浮かんでいません…。

追記：ちょうど複数出力の方法を投稿した方がいたので追記しておきます。

ED法+ReLU

ReLUですが、少し右か左にずらすと学習できました。
どうやら0がいけないらしく、なぜかはわかりません…。

プログラムでは以下です。

def ed_relu(x, move=0.1):
    return np.where(x > move, x - move, 0)


def ed_relu_derivative(x, move=0.1):
    return np.where(x > move, 1, 0)

また、このED ReLUはどうやらUnit数とLayer数で最適なmoveの数字が変わるようで、学習が難しくなる印象を受けています。

MNISTの学習

今回TFとTorch両方で実装しました。
実装自体は重みの一括更新までは実装していないので高速化の余地は十分にあります。
ですがそれを置いてもフレームワークの恩賜を受けられるようになるのはやはり大きい…。

特にTFはKerasモデルで実行できるように実装したのでそのまま実行できます。
コード例は以下です。

(x_train, y_train), (x_test, y_test) = MNISTのデータセットを取得

# EDModelを生成(EDModelはGithubに上げてあります)
model = EDModel(
    input_num=28 * 28,
    output_num=10,
    layers=[  # 64x2層モデル
        (64, "sigmoid"),
        (64, "sigmoid"),
    ],
    out_type="linear",  # 出力はlinear
    training_mode="ce", # 交差エントロピーで学習
    lr=0.01,
)

# lossとoptimizerはBP用の話なので指定する必要はありません
model.compile(metrics=["accuracy"])

# 学習
model.fit(x_train, y_train, epochs=10, batch_size=512, validation_data=(x_test, y_test))

# 評価
model.evaluate(x_test, y_test)

現状、活性化関数は "sigmoid" "relu" "linear" しかありません。
また、training_mode も "mse" と "ce" のみとなります。

いろいろな実行結果

・実行環境
windows11
CPUx1: Core i9-12900 2.4GHz
GPUx1: NVIDIA GeForce RTX 3060 12GB
memory 32GB

ベースライン(TF)

レイヤー数：3
ユニット数：32
活性化関数：sigmoid
学習率：0.0005
バッチサイズ：512
エポック数：20

今回実装では出力が1つ固定のモデルと複数出力があるモデルを別に作っています。（速度が気になって…）
BPもED法も正解率は90%前後とちゃんと学習できていますね。
ただED法の方が少し低い結果になっています。

次は速度です。

当たり前ですがTFは早いですね…orz

ベースライン(Torch)

レイヤー数：3
ユニット数：32
活性化関数：sigmoid
学習率：0.001
バッチサイズ：512
エポック数：10

同じくTorchです。
TorchはTFより遅かったので少し学習率を上げてエポック数を減らしています。

BitNetの時にも思いましたが、同じ内容でもTFとTorchで結構違いました（デフォルト値が違うのかな？）
こちらはBPよりED法のほうが学習できているように見えます。

以下速度です。

フレームワークで結構変わりますね。

ちょっと多層モデル(TF)

レイヤー数：5
ユニット数：32
活性化関数：sigmoid
学習率：0.0005
バッチサイズ：512
エポック数：20

ベースラインからレイヤー数を5にした形です。

多層モデルの学習の速さは流石のED法ですね。
BP法は層が増えると勾配消失がなくても学習が遅くなります。

ReLUモデル(Torch)

レイヤー数：5
ユニット数：32
活性化関数：relu
学習率：0.001
バッチサイズ：512
エポック数：10

こっちはTorchで見てみました。
比較よりかはED版ReLUでちゃんと学習できるかの確認の意味合いが強いです。

問題なく学習できていますね。

多層モデル(TF)

レイヤー数：100 # あまり増やすとメモリにのらなかったり…
ユニット数：8
活性化関数：sigmoid
学習率：0.0005
バッチサイズ：512
エポック数：10

BPでは学習できていませんがED法ではちゃんと学習できていますね。（精度も90%超えました）
以下速度です。

速度ですが、TFは多層モデルにそこまで対応していなそうでそこまで離れていません（そもそも高速化できないのかな？）

実行に使ったコードはGithubにあるので興味のある方は数字を色々変えて試してみてください。

考察1：シナプスの強化と山登り法

前回Q&Aに書いたシナプスの強化に関する疑問について1つ仮説を書いておきます。

疑問ですが、シナプスの強化はプログラム的にはプラス方向の足し算しかしておらず、減少することはありません。
興奮性ニューロンと抑制性ニューロンで相殺が起こると無限に増えそうですが、なぜか増えないという疑問です。

これは山登り法で両側から学習しているとイメージすると合ってそうな気がしました。

書いてみて思いましたが雑な考察ですね…、一応残しておきます。

考察2: BitNet+ED法

BitNetは過去に記事を書いているのでそちらも参考にどうぞ

内容を簡単に言うと、ニューラルネットの重みを小数ではなく[-1,0,1]の3値で表現したら、（十分なユニット数がある場合）性能が下がらずに逆に上がったという内容です。
また、これは[-1,1]の2値では性能が上がらなかったという前の論文の話と繋がっています。
この結果をED法に置き換えると面白い内容となります。

ED法を2値で表現しようとした場合、正のシナプスを[0,1]、負のシナプスを[0,-1]と考えるのが自然です。
これは[-1,0,1]の3値に繋がりますね。
要するにBP法の3値化=ED法の2値化と同じ意味合いを持つと解釈できます。
[-1,1]の2値で学習できなかったのはそもそもニューロンが保持する情報の最小が[0,1],[0,-1]だったと考えると納得できます。

また、各シナプスが2値しか持たないという事は発火する・しないといった本来の神経系の動作に近いモデルになるのでは？とか思ったり。

実験

コード上は forward の重みを足すタイミングで weights を2値化すれば実現できます。
（入力の8bit化は省略します）

ここですが、BitNetに倣うなら重みの平均を引くところですが、間違えて割る形で実装したらなぜか学習できました。
（引くと逆にできなくなった）
（全重みではなくとりあえずニューロンに接続されている重みの平均を割ったのでそこら辺が影響？）

コード例は以下です。（Githubのコードにも実装してあります）

# --- listの例
def forward(self, x, training=False):
    (略)

    # 重みの平均
    a = sum(self.weights) / len(self.weights)
    if a < 0.0001:
        a = 0.0001

    y = 0
    for i in range(len(self.weights)):
        # y += x[i] * self.weights[i] * self.forward_ope_list[i]
        # ↓
        w = 1 if self.weights[i] / a > 0.5 else 0
        y += x[i] * w * self.forward_ope_list[i]

    (略)
    return y

XORは無事に学習できました。

--- result ---
[0, 0] ->  0.00005, target [0]
[1, 0] ->  0.99033, target [1]
[0, 1] ->  0.98907, target [1]
[1, 1] ->  0.00372, target [0]

MNISTの結果は以下です。

レイヤー数：5
ユニット数：64
活性化関数：sigmoid
学習率：0.001
バッチサイズ：256
エポック数：20(Torchは10)

・TF版

・Torch版

ゆっくりですが学習は出来ていそうです。
これ重みが[0,1]or[0,-1]の2値しかとっていないんだぜ…

ED法+2値化が実現すると

もしED法ニューラルネットが2値([0,1] or [0,-1])で表現できるとすると学習で面白いことができるかもしれません。

考察1で書いたとおりED法では強化する方向が決まっています。
なので各重みは「0→1」または「0→-1」という変化しかしなくなります。
要するに学習時に更新の方向と更新の量が一位に決まるので、計算自体が不要という事に…。
・・・すごいことになりそうだ

１つ試した内容は、学習でランダムに１つの重みを0から1に変えるという方法です。
学習則の計算も無視して学習できるので大幅な高速化が見込めます。
山登り法より多分収束するだろうという仮定でかなり強引な方法ですが、残念ながら今のところ学習できていません。
遺伝的アルゴリズムとか利用するのがいいのかな（本末転倒かも

最後に

フレームワークに実装したので比較的実行しやすい環境を整えることができたのかなと思います。
また、今回の結果はMNISTだけの結果なのでちゃんと見るにはもっと他のデータセットを見てみる必要があります。
（が、個人の趣味の範疇を超えるのでやりませんが…）

また、プログラムを走らせているとやはりBP法とは違った性質を感じました。
これ今でも十分に論文を出せるレベルじゃなかろうか（私は書かないですけど…）