最適化手法を最適化: GD by GD (LSTM Optimizer)

goose morning; (EleutherAI流の挨拶) 自称研究者の痛い幼気な中学生こと白上です。暇すぎるのでこれから論文紹介をNoteでしていこうと思います。今回紹介する論文は、Google BrainのM.Andrychowicz氏が率いる研究チームによるLearning to Learn by Gradient Descent by Gradient Descentです。2016年にNeurIPS採択なので若干古いですが、私の大好きな論文です。

さて、世の中にはOptimizerというものが沢山存在します。別名最適化手法、訓練アルゴリズム、更新則。周知の通り、これらは深層ニューラルネットワークの重みの更新、すなわちAIの学習に使われます。
沢山あると当然ヒエラルキーが出来るもので、Optimizer界隈ではAdam (ADAptive-Momentum) という手法が2014年に提案されて以降延々と玉座に座り続けています。Adamはその収束率の高さとあらゆる条件で汎化性能する汎用性から世界中の論文に引用され、最近流行りのChatGPTに至ってもほとんどの訓練段階で利用されているほどです。

しかし。Adamは人間が一からプログラムしたもので、AIの訓練に利用されるにも関わらずそれ自体は単なるルールベースのアルゴリズムでした。無論最適化できません。ルールベースですから。

ですがそんなことをする必要性もありません。Adamは既にAIを最適化する手法として最適ですからね。事実、Adamより優ると称して国際会議に採択された手法であるRAdamやAdaBelief、SAMやLionは結論として何の改善にもなっていなかったのですから。今のままがベスト。無理に変えようとするものじゃないです。

...… 本当に、そうでしょうか? (某自動車保険のCM)

という訳で本題に入ります。GD by GDとは、雑に説明すると「重み共有されたLSTMを別なモデルの訓練時にOptimizerとして活用しながらそれ自体も訓練する」というアイデアとなります。

より具体的には。LSTMの入力としては訓練対象モデルの今までの勾配を与え、出力結果をパラメータの更新量として扱います。ここで注意すべき点は、入出力系列の空間軸のサイズは常に1であり、訓練対象のパラメータの数だけ同じLSTMを共有するということです (従って訓練対象がどんなに大きなモデルだったとしてもOptimizer自体のモデルサイズは変わりません) 但し、個々のLSTMの隠れ状態は別なものを利用します。

では訓練時にどのような逆伝播をするか、という話になりますが。実は二階勾配を使います。というのもそれしか手段がなかったからですね。計算量の増大は仕方がないこととして割り切るのですかね。訓練対象の各イテレーションごとのloss値の重み付き平均をlossと見做して逆伝播するらしいです。

より詳細な実装ですが、実は私も知りません()。論文にも書いてなかったので仕方ありませんね()。なのでもし真面目に実装したい場合はDeepmind公式やyangsenius様の実装を参考にしてください()()。

さて、実際の戦績ですが割と良さげですね。単純な最適化問題においても、MNISTにおいてもAdamの1/2程のイテレーション数で収束していっていることが読み取れます。他にも、CIFARやスタイル転送のドメインにおいても割と悪くない成果を出したそうです。

個人的にはこの手法、強化学習やNeuroevolution(遺伝的アルゴリズムの応用)等と組み合わせたり最新のTransfomerでリファクタすることで、より良い成績の得られるであろう夢のある手法だと思っております。なので暇な方、もしよければ実験などやってみては如何でしょうか ()()。

いかがでしたでしょうか? 正直長い割に内容スカスカの記事なので、Twitterでお叱りを受ける気がしていてなりません。(なので今のうちに予防線を張っています) まあ今後も私の好きな論文について不定期的にNoteを出す予定なので、その際はよろしくお願いします。