断想:系列入力ベースの強化学習

 最近は状態や報酬などを系列データとして扱う強化学習に興味が出ている。端的に言えばDecision Transformer1 のことになる。

 特に、エピソードをまたいだ(across-episodicな)長い系列を入れることに可能性を感じる。着目点は違うが、やっていることとしてはAlgorithm Distillation2に近い。個人的に期待するacross-episodicの良い点としては、体験した成功例・失敗例をそのまま有効活用してサンプル効率を高めることにある。明示的な体験の参照ができない状況では、ニューラルネットのパラメータに勾配法で知識が反映されるまで例を活かすことができないが、それはどうしても遅くなってしまうのではないかというところが気になっている。

 改善されると嬉しい別の例をコンピュータ将棋で挙げると、(ランダム性を入れずパラメータ更新もない状況の想定でも)全く同じ棋譜で全く同じ負け方をすることを回避できるようになることだと思う。完全に同じ失敗を繰り返すのはあまりに知的でない振る舞いで、見ていてとても気分が悪い。これをなんとかしたい。

 体験を外部情報の一種として捉えると、LLMにおける検索の利用(Retrieval Augmented Generation)と類似した側面もあるのかもしれない。知識をニューラルネットのパラメータへ反映させるには学習量が必要そうなので、プロンプトの参照情報として与えてしまうという発想になる。

課題

 この方針の課題は、今のニューラルネット(Transformer)では入れられる系列長が短いことだと思う。GPT-3.5を例にすると16Kが上限となっている。将棋でどの程度過去の棋譜を入れられるか概算すると、

1局だいたい200手と考えて、片方の手番だけを入れるとして1局あたり100局面で、各局面は9x9トークンになり、たとえば10局分入れようとすると81000トーク

となる。単純にやるとまともに入れられないことがわかる。それに、コンピュータ将棋では推論速度も重要になるので、なんとか動くけど遅いという状況ではメリットがほとんどなさそうだ。できるだけ速度は落とさないままに過去の情報を活用できる必要がある。

対策1: 系列長に対して線形な計算量のアーキテクチャを使う

 もちろんTransformerの、系列長に対して2乗の計算量が必要になってしまう性質は多くの人が改善したいと考えており、様々に研究がされている。最近だとMamba3であるとか、Based4であるといったものが有力なのではないかと囁かれているらしい。

 論文などを多少見てはみたが、状態空間モデルをいろいろこねくり回して上手いことやる手法は数式・概念とかが難しく、さらっと読んですぐに理解できるというものではなかった。ここに関しては、多少興味はあるが自分でめちゃくちゃ試行錯誤していくほどのモチベーションはないし、多くの人が取り組んでいくだろうから、その成果だけを後々享受させてもらうという姿勢で良い気もしている。

対策2: 入れる系列の選び方を工夫する

 先程は単純に1局から100局面を全て入力するという考えで系列長を試算したが、それはあまりにも愚直なやり方ではある。RAGのように、検索概念を導入して、現局面と類似する過去の局面を探して入力に追加するというのも有望だろう。

 また、直感的にヒント情報として重要なのは、良い方向であれ悪い方向であれ予測を大きく外した局面だと思われるので、過去の体験に対してValueの予測ハズレ具合で重み付け(優先順序付け)をするのも可能性があるのではないかと感じる。これは経験再生における優先順位付け(Prioritized Experience Replay5)の発想そのままではある。

 いずれにしても、少量の効果的な状態だけ上手く系列の一部に取り入れることで、計算効率をそこまで下げずに精度に寄与できると面白いと思う。とりあえずはこちらの方向性で考えてみたいか、という気分。