1. はじめに
Turing生成AIチームの荒居です。生成AIチームでは、完全自動運転の実現に向けて、現実世界の複雑な状況を理解し予測する動画生成AI、「世界モデル」の開発に取り組んできました。前回の私の記事では、その取り組みの中で調査したGAIA-1の紹介を行いました。今回は、Turingが開発した世界モデル"Terra"で利用している技術や課題・展望などについて紹介をしたいと思います。
この記事で解説している内容
- 世界モデルTerraは何ができるのか
- 世界モデルTerraはどのような仕組みで動いているのか
- どのような課題が残されているか
Terraで生成した動画。Terraは運転環境に特化しており、車載カメラの一人称視点動画の生成を行うことができる
2. 世界モデルTerraは何ができるのか？
Turingの世界モデルは大きく分けて二つのことができます。一つ目は短い動画を与えるとその続きを生成するvideo rolloutです。以下の動画では最初の10フレーム(赤枠がついているフレーム)を入れて続く45フレームを生成しています。この機能を使うと、生成されたフレームをコンテキストに入れてさらにその続きを生成する、ということができるため原理上は無限に長い動画の生成が可能です。
video rolloutの一例。赤い枠で囲まれている区間が実際の動画で、そのあとの部分が生成された動画になっている
二つ目は、短い動画とそのあとになぞって欲しい軌跡(trajectoryと呼びます)を与えて、その軌跡に沿って進んだ場合の動画を生成するaction-conditioned generationです。以下の例では、下の画像の緑の軌跡に沿って進んだ場合の動画が左側に、赤の軌跡に沿って進んだ場合の動画が右側に表示されています。
条件付けに与えている軌跡の例
(左) 上の画像の緑の軌跡に沿った場合の生成動画 (右) 上の動画の赤の軌跡に沿った場合の生成動画
この機能を用いることで、運転操作の良し悪しを評価する仕組みが作れると私たちは考えています。具体的には運転操作を行なった未来の動画を生成し、その動画に対して「先行車に近づきすぎているので評価は0.3 / 1.0」「安全な運転が継続できるため評価は0.8 / 1.0」のように評価値をつけるようなことを考えています。
3. 世界モデルTerraの仕組み
世界モデルTerraの基本的な構造はGAIA-1を参考にしており、画像フレームを離散トークンに変える Image Tokenizer、画像フレーム列を離散化したトークン列と条件付けのトークン列を受け取り、未来の画像トークン列を予測するWorld Model、そして予測された未来の画像トークン列を入力してフレーム列を出力するVideo Decoderからなっています。
Terraの大まかな構造
3.1 Image Tokenizer
Image Tokenizerは画像をパッチに分割し、各パッチに対して整数値を割り当てて離散トークンの列に変換する機能を担います。画像の離散化に用いられる代表的な手法としてVQVAEが挙げられます。GAIA-1でもおそらくVQVAEを用いている、という話は前回のブログ記事でも触れました。一方で、近年は画像生成の分野の発展に伴い様々な画像の離散化方式が提案されています。
VQVAEの特徴ベクトルが離散トークンに変換される模式図
例えば、MAGVITv2の論文中で提案されたLookup Free Quantization(LFQ)は、VQVAEのように事前に用意したコードブックに対して最近傍探索をすることで離散化をするのではなく、画像のエンコーダの出力のEmbeddingの各次元の値を2値化した上でEmbeddingをバイナリとして見立てることで離散化を行います。
LFQで特徴ベクトルが離散トークンに変換される模式図
また、TiTokでは、画像をパッチ列に直したあとVision Transformerを用いて特殊なトークンに情報を移すようにしながらAuto Encodingを行うことで、256×256の画像をわずか32トークンで表現できることを示し大幅な圧縮効率向上を実現しました。
TiTokの論文より引用。Vision Transformerに画像のパッチと一緒にlatent tokensという特殊トークンを入力し、Transformer Encoderの出力のlatent tokensに対応した部分のみを使って再構成を行う
今回は、学習済みモデルが公開されていて高い再構成性能を達成していたLFQを画像のトークナイズ方式として選びました。なお、この学習済みモデルはトークンの語彙数が262,144と桁違いに多く非常に高い表現力を持っています。
3.2 World Model
世界モデルの本体の部分のコアは一般的なLLMそのものなため、様々なアーキテクチャを用いることができます。今回は学習済み重みを用いた学習ではなく、乱数初期値からのスクラッチ学習になるため、モデルアーキテクチャの選択は性能面では重要ではないと考え、様々なパラメータ数の参照設定が存在するOPTアーキテクチャを選択しました。
Terraは条件付けのためにtrajectoryを受け取るため、trajectoryをトランスフォーマーに入力できる形にする必要があります。Terraでは一つのtrajectoryを3次元のベクトルが5つ連なっているデータ構造と定義し、それぞれのベクトルをprojectorでトランスフォーマーの入力の次元に変換した上で一つずつトークンとして入力します。3次元のベクトルは(車両中心座標系におけるx座標, 車両中心座標系におけるy座標, その画像フレームの時点を原点とした時刻t)になっています。
TerraのWorld Model部分の構造の模式図
上の図では画像の離散トークンもTrajectoryも1フレーム分しか描かれていませんが、実際には1フレーム分の画像トークン列と、その時刻からみた将来の軌跡トークン列が交互に入力されるようになっています。
学習はNext Token Predictionのタスクで行い、損失はCross Entropy Errorを用います。Trajectoryトークンの部分は損失計算を行わず、画像トークンの予測に対してのみ損失計算を行うようにしました。
また、推論時は自己回帰的に画像トークンを予測するようにして生成を行います。1フレーム分の画像トークンは576にしているため、576回分の推論で1フレーム分が生成されます。1フレーム分生成が行われた後は、Trajectoryトークンを挿入した上で次のフレームの画像トークンの生成を自己回帰的に行います。
3.3 Video Decoder
World Modelが予測した画像列の離散トークンを画像列に直し、人間が走行映像だと認識できるような動画に直す部分がVideo Decoderの役割です。これを実現する簡単な方法は、画像のトークナイザー(今回であればLFQ)のデコーダをそのまま用いることです。画像のトークナイザーはAutoEncoderのエンコーダの部分にあたるため、デコーダの部分が利用できます。
しかし、画像のトークナイザーは動画対応のものを選ばない限り、フレーム単位でしかデコードを行わないため、個々の再構成画像の品質が高く見えたとしても動画として繋いで見てみるとチラつきが目立つことがあります。以下の例では、画像のトークナイザーのデコーダをそのまま用いた場合と、動画に対応したデコーダを用いた場合になりますが、画像のトークナイザーのデコーダの方がチラつきが多く動画の品質が低く見えるのではないでしょうか？
LFQのデコーダを用いてデコードを行った場合
動画対応のデコーダを用いた場合
動画に対応したdecoderは例えば、Stable Video Diffusionなどで使われているTemporal Decoderが挙げられますが、これは離散値の潜在表現には対応していないものになっています。
そこで、TerraではStable Video Diffusion(SVD)をファインチューンし、画像のトークナイザのデコーダでデコードした画像を時系列的に滑らかに接続するようなモデルを作成しました。
SVDは画像を入力としてその画像を動画化するImage-to-Videoのモデルになっています。入力画像は条件付けのために使われますが、その際、デノイジング対象のノイズと、入力画像をVAEでエンコードしたものをチャンネル方向に繋いで入力することで画像中の物体位置などの構成を保った動画を生成することができるようになります。
SVDでは、生成するフレーム数分のノイズフレームに対して全て同じ画像の潜在表現を結合することで条件付けを行いますが、Terraでは世界モデルの予測したフレームごとの結果をVAEでエンコードしてノイズと結合します。
TerraのVideo Decoderの模式図
4. 課題と今後の展望
Terraは試行錯誤の過程の途中のものということもあり、改善の余地が多く残されています。これらの課題は大きく分けると以下の三つに分類されます。
- 生成速度の課題
- 生成された動画の品質の課題
- 世界モデルの活用可能性における課題
4.1 生成速度の課題
Terraは画像1フレーム分の生成を行うのに576回の推論が必要なため、1フレーム分(=0.1秒先の予測)に数秒がかかります。生成速度は世界モデルの活用可能性にも大きく関わるところであり、優先度高く対応したいところです。
生成速度の改善のためには、
- モデルのパラメータ数を減らす
- 1枚の画像の表現に用いるトークン数を減らす
- トークンの予測方式を変更し、1回の推論で複数のトークンを予測する
などの対策が考えられます。
まずモデルのパラメータ数を減らす方法についてですが、こちらは性能に大きく関わる部分ですので慎重に判断したいところです。現在のTerraは10億パラメータ程度のモデルですが、GAIA-1を開発したWayveによれば世界モデルにおいてもスケーリング則が成り立つためパラメータ数を減らすことは性能悪化に直結してしまいます。一方でcomma.aiの公開している世界モデルは3億パラメータ程度のGPT-2モデルを50,000時間の走行動画で学習させた結果安定した推論ができるようになっています。これを踏まえると、パラメータ数は少なめで、データの量・多様性を増すことで高い予測性能を出すことができる可能性はあり得ると考えています。
1枚の画像の表現に用いるトークン数を減らす方向は、かなり有望であると考えています。Image Tokenizerの節でも紹介したTiTokは非常に少ないトークン数で高い再構成性能を実現しており、1Dのtokenizationは今後の発展にも期待したいところです。また、iVideoGPTやΔ-IRISなど、フレーム間に強い相関があることに注目した動画特化のtokenizationも今後の発展の方向性としては注視しています。
トークンの予測方式の変更もかなり有望な方向性です。MaskGITやPhenakiなどはNon-Autoregressiveな生成方式で1度に複数のトークンを予測しながら生成を行うようにしており、生成速度面において大きなアドバンテージがあります。また、最近発表されたMasked Autoregressive model(MAR)も有望な方向性です。
4.2 生成された動画の品質の課題
Terraの生成動画の品質はまだ高いとは言えません。シーンによってはかなりリアリスティックな動画を生成できるようになってきていますが、生成動画であると容易に判別できるような動画を生成してしまうことは少なくありません。したがって、今後もしばらくは生成動画の品質向上につとめる予定ですが、そもそも生成動画の品質の評価の方法がかっちりと定まっていないという問題があります。
既存の世界モデル関連の研究では、FIDやFVDなどの評価指標が用いられることが多いですが、これらの指標は計算に時間がかかることや、画像/動画の多様なアスペクトを評価できないことなどが問題点として挙げられています。世界モデルと近い研究分野である、テキスト動画生成においては近年になってようやくT2VScoreやVBenchなどの評価手法が出てきましたが、これらの評価手法はすぐに世界モデルに適用できるようなものではなく今後の研究の余地が多く残されています。
生成動画の評価では、動画・画像の歪みの量などについての評価と、動画・画像の内容についての評価を切り分ける場合があります。生成された動画の品質を高める、と言った場合にもそれぞれについて別個の改善方法が考えられます。
まず、動画・画像の歪みの量はエンコーダ/デコーダに由来する部分もあります。Terraではエンコーダに学習済みのモデルを使っていますが、これにより画像の細部、特に運転において非常に重要な信号機や標識が潰れてしまうというような現象を確認しています。Terraの用いている画像のエンコーダはImageNetで学習されたもののため、交通環境の画像・動画をエンコードするのに適していない可能性があり、独自にエンコーダを学習する必要があるのではないかと考えられます。
(左)LFQの入力に用いた実画像 / (右)LFQの再構成画像。よく見ると信号機や標識などが潰れてしまって判読不能になっている。
また、デコーダについては現在はSVDが用いているデコーダに依存していることになりますが、こちらも独自にTemporal Decoderを学習するようにすれば、離散トークンを用いたモデルのデコーダをそのままVideo Decoderとして使うことができるようになると考えています。
一方で生成動画の内容の品質向上のためには、データの量・多様性・品質を増やした上で学習をすることが最も確実な打ち手であると考え、現在もこの方向性で学習のスケールアップを実施しています。
4.3 世界モデルの活用可能性における課題
Terraはまだ実用にたえるものではありませんが、実用化にあたっての目下最大の障害が、生成結果を評価する仕組みがないことです。ここでいう評価とは、生成された動画の品質の評価ではなく、生成された動画の品質が十分に高いという前提のもと、生成された動画が実際の運転の質として高いか低いかを評価するものを指しています。LLMで言えば、RLHFの報酬モデルのようなものに当たります。
Dreamer v3やIRISなど、シミュレータとしての利用がすでにうまくいっている世界モデルはゲーム環境をモデル化しており、報酬データを得ることができるため報酬モデルを作りやすいのですが、運転環境においては明示的な報酬データを得ることが難しいため、独自に報酬データをアノテーション等で収集する必要があります。
現時点での動画の生成品質は十分に高いとはまだまだ言えませんが、リアリスティックな動画は生成できるようにはなってきているため、報酬データの収集にも手を出していこうとしています。
5. 終わりに
本記事では、Turingが開発を進める世界モデルTerraについてその能力と技術的な概要について紹介を行いました。また、Terraが抱える課題や、その解決のための指針についても紹介しました。
Turingは「We Overtake Tesla」をミッションに完全自動運転の実現に向けて尽力しています。経済産業省/NEDOの競争力のある生成AI基盤モデルの開発を支援するGENIACプロジェクトにも採択され、生成AI・大規模基盤モデルの開発を進めています。完全自動運転に向けたマルチモーダルな基盤モデルの開発にご興味ある方、お気軽にコメントやX(Twitter)のDMにてお問い合わせください。
参考文献
- Ha, David, and Jürgen Schmidhuber. "World models." arXiv preprint arXiv:1803.10122 (2018).
- Hu, Anthony, et al. "Gaia-1: A generative world model for autonomous driving." arXiv preprint arXiv:2309.17080 (2023).
- Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." Advances in neural information processing systems 30 (2017).
- Yu, Lijun, et al. "Language Model Beats Diffusion--Tokenizer is Key to Visual Generation." arXiv preprint arXiv:2310.05737 (2023).
- Yu, Qihang, et al. "An Image is Worth 32 Tokens for Reconstruction and Generation." arXiv preprint arXiv:2406.07550 (2024).
- Zhang, Susan, et al. "Opt: Open pre-trained transformer language models." arXiv preprint arXiv:2205.01068 (2022).
- Blattmann, Andreas, et al. "Stable video diffusion: Scaling latent video diffusion models to large datasets." arXiv preprint arXiv:2311.15127 (2023).
- Wu, Jialong, et al. "iVideoGPT: Interactive VideoGPTs are Scalable World Models." arXiv preprint arXiv:2405.15223 (2024).
- Micheli, Vincent, Eloi Alonso, and François Fleuret. "Efficient World Models with Context-Aware Tokenization." arXiv preprint arXiv:2406.19320 (2024).
- Chang, Huiwen, et al. "Maskgit: Masked generative image transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.
- Villegas, Ruben, et al. "Phenaki: Variable length video generation from open domain textual descriptions." International Conference on Learning Representations. 2022.
- Li, Tianhong, et al. "Autoregressive Image Generation without Vector Quantization." arXiv preprint arXiv:2406.11838 (2024).
- Wu, Jay Zhangjie, et al. "Towards a better metric for text-to-video generation." arXiv preprint arXiv:2401.07781 (2024).
- Huang, Ziqi, et al. "Vbench: Comprehensive benchmark suite for video generative models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.
- Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).
- Micheli, Vincent, Eloi Alonso, and François Fleuret. "Transformers are sample-efficient world models." arXiv preprint arXiv:2209.00588 (2022).
Discussion