RikkaLayer -RNN型SSMベースでTransformerを超えた?（超えてませんでした。）パフォーマンスを発揮する新手法-

はじめに

失敗作です。ごめんなさい。
3/6追記 大規模化しても性能があまり向上しなかったので、超えれてませんでした。
2017年に”Attention is all you need.”という論文にて提唱されたTransformerというモデルがある。Transformerは、同一文章内の単語間の関係を上手く捉えることが可能となり、以前のRNNベースのAIとは一線を画す革新的な性能を実現した。
ただ、Transformerには問題がある。RNNの計算量は単語数に比例して計算量が増えていくのに対し、Transformerは、単語数2乗に比例して計算量が増えるという問題があり、昨今大規模化がますます加速していく中で学習に使われる電力消費量故の環境負担が大きいという課題がある。本文では、その環境負担が大きいという問題を解決すべく、Transformerと同等の性能を発揮しながら、計算量が少なく済む構造を模索した結果を述べていく。

１：Transformerの性能の理由

なぜ、TransformerはRNNベースのモデルと比べ、革新的な性能を発揮するのかについて考察をした。RNNは、入力内容に応じて潜在変数を変化させ、次の入力時に入力を引き継ぐという構造であり、視点を変えれば、入力回数=DNNの階層数とも取れる。
この構造では、入力回数が増えるごとに層が深くなってゆくので、入力回数が増えるに連れ、勾配消失が起きてゆくという問題が出てくるのは必然であり、明白である。
しかし、Transformerは、層を跨がずに全ての入力を1層で処理するので、勾配消失が起こりにくいのがわかる。層数が入力回数によって変化せず、勾配消失を起こさない層数で固定できるので、Transformerは高い性能が発揮できる。Transformerが高性能なのは、単に勾配を効率よく伝え、学習できるからであると考える。

2：RNNベースの利点と問題点の解消

Transformerは、Attention部分で入力と入力の対応表のようなものを作り、フィルターを生成するという構造をしてあり、表を作るということは、128個の1次元入力を入れた場合、縦128個、横128個の2次元表になるということで、計算量が入力数の2乗増えるのは明らかである。しかし、RNNでは、1次元のまま処理するので、計算量は入力数との単純な比例の関係になるので計算量が少なくなる。しかし今のRNNでは、先にあげたように勾配消失により学習できないという問題が付き纏ってしまう。そこで私は残差結合を、RNNレイヤー内の引き継ぎ変数に、残差結合を適用するというアイデアを提唱する。

３：提案手法

私が提案する手法は、ResNetで使用された残差結合を、従来の深さ方向に加え、時間軸方向にも使用するという手法である。ベースとなるモデルには、状態空間モデルの定数A,B,C,D,Kの代わりに全結合層を2層重ねたものを使用した手法を用いる。この構造を以後NNSSMと呼ぶ。
NNSSMの変数$y$と、$x$に残差結合を適用する。
計算式にすると、以下の通りとなる。これがRikkaLayerである。

$I_{t} = 入力$
$y_{t} = 出力$
$x_{t} =状態変数$

$LinearBlock = Linear(GELU(Linear(x)))$
$N() = Normalize()$
$A(),B(),C(),D(),K() = LinearBlock()$

$x_{t-0} = A(N(x_{t-1})) + B(N(I_{t-0})) + x_{t-1}$
$y_{t-0} = C(N(x_{t-1})) + D(N(I_{t-0})) + I_{t-1} + x_{t-1}$

構造図

2024 3/3　構造を最適化しました。

3:性能

検証が不十分なのは否めませんが、Transformerと比較した
おおよその性能は以下の通り。
Transformer = 12層の隠れ層128次元
RikkaLayer = 隠れ層128次元のRikkaLayer12層
メモリ消費量(32単語時) = $\frac{1}{2}$倍
計算速度 = 2倍
最終性能　＝　1倍

4:実際のデータ

学習データ量が非常に少ないので、完全な状態ではない。
事前学習2000文
ファインチューニング2000文
Human: **ちゃんが好きな食べ物は？
AI : 中華が一番だよ！
Human:　嫌いな人は？
AI : 自分？。えのきだよ。
Human : 大っ嫌いだ！！
AI : 好き？。お酒は？。

終わりに

Transformerの強みを、入力が長くなっても勾配消失が起きないことと仮定し、RNNに勾配消失対策として、残差結合を使用した所、すぐれた性能、特性を見せたところから、Transformerの性能の由来はMetaFormerなどで提唱された構造が由来という訳ではなく、は勾配消失が起きにくいことが性能の由来ではないかと見られる。

参考文献

ResNet : arXiv:1512.03385
Transformer : arXiv:1706.03762
SSM : https://ja.wikipedia.org/wiki/%E7%8A%B6%E6%85%8B%E7%A9%BA%E9%96%93_(%E5%88%B6%E5%BE%A1%E7%90%86%E8%AB%96)
MetaFormer : arXiv:2111.11418

お願い

本記事で私の提唱する新手法を、もし商用利用する場合は、私に連絡をいただきたいです。
実績として扱わせていただきたいのでご協力お願いします。あと単純に嬉しいので。

サンプルコード