たとえば動画的な連続フレームをRNNなどの時系列予測モデルに入力してその先を予測させたいということは環境モデルを考えているとよくある。
その場合、H×W×C的なデータを
- できるだけそのままH×W×Cで扱う(ConvLSTMなど)
- 256次元など1軸にflattenして扱う
の2通りがあり得るように思う。
このあたりをどうしているのか気になったので、最近出たDreamer V3
とAdA
についてそこだけを調べた。
Dreamer V3
Dreamer V3は根本的に上のようなアーキテクチャになっている。
- 入力画像からエンコードNNを通したZの方はまだH×W的な構造
- GRUの状態であるhの方は1軸にflattenされた構造
だと思われる。
付録Bにある記述を読むと
入力画像は4×4×CになるまでStride 2で畳み込みをどんどんかけられていき、その後flattenされGRUに渡されるようだ。
AdA
そもそもAdAのことをちゃんと理解できていないが、どうも表現をTransformer-XLに突っ込み、その後Muesliにさらに入力して予測を重ねていくらしい?
付録CにAgent Detailsが記載されている
- 入力 : 72×96×3 のRGB画像
- ネットワーク : ch数が[16, 32, 32]のResNet。各ブロックでは2×2のMaxPoolingでサイズが小さくなっていく? 最終出力は1軸の256サイズ
というわけで、この後Transformer-XLやLSTMなどが続いていくが、1軸にflattenされた表現であることには間違いなさそう。
所感
この2つだと1軸にfalttenする方が採用されているらしい。将棋だともとから盤面サイズが9×9なので、Encoderの最後で10chくらいにすればflatten化しても810次元くらいで、ギリギリ取り扱えるだろうか。そこまでch数を減らしてPolicy, Valueの性能が落ちないのかどうかはわからない。