UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

ロボットに世界を理解させる方法を教えようとしている場面を想像してみてください。かつて、科学者たちはロボットに何百万枚もの画像を見せ、「次に何が来るか」を推測させたり、動画を見せてフレームごとに未来を予測させたりすることで、これを行おうとしました。これは、物語を理解する前に、すべての本のすべての文字を暗記しようとしたり、シーンを完璧に模写しようとしたりすることに似ています。これには膨大な時間とコンピュータの計算能力が必要です。

最近では、「Joint-Embedding Predictive Architecture（または短縮してJEPA）」と呼ばれる、より賢いアイデアが登場しました。JEPAは、画像全体を書き直そうとする代わりに、物事の「本質」を理解するようにロボットを教えます。これは、俳優のシャツの色をすべて覚える必要はなく、映画のプロット（筋書き）を学ぶようなものです。ロボットは、目にするものを効率的な小さな「メンタルノート（潜在空間）」へと圧縮し、そのノートが将来どのように変化するかを予測することを学びます。しかし、これまでは、明るさを調整したときに写真がどう変わるかを理解するためのロボットと、ビデオが時間とともにどう動くかを理解するためのロボットというように、異なる仕事のために異なるロボットを作る必要がありました。それらは、互いに会話できない二つの別々の図書室のようなものでした。

そこで、これらのメンタルノートの「ユニバーサル翻訳機」として機能する新しいアプローチ、UniJEPAが登場しました。この論文の研究者たちは、次のような問いを立てました。「写真の照明効果（明るさや色の変化など）による変化と、シーンが時間とともに動く様子を、同時に理解できるたった一つのロボットの脳を構築できるだろうか？」彼らは、これが可能であるだけでなく、これらを同時に行うことで、ロボットはより賢く、より速くなることを発見しました。単一の統一されたルールを用いることで、UniJEPAは、生のピクセルの煩雑な詳細を再構成することなく、シーンの「見た目」と「動き」の両方を予測することを学習します。単一の脳がこれら両方のタスクを処理でき、行動を即座に計画できる柔軟な世界の捉え方を学習できることが判明したのです。

大きなアイデア：一つの脳、二つの超能力

AIの従来のトレーニング方法を、教室にいる二人の生徒に例えて考えてみましょう。一人の生徒を「フォト・プレディクター（写真予測者）」と呼びましょう。彼女は静止画のみを学習することを許可されています。もし彼女に猫の写真を見せて、その猫の照明が変わったらどうなるかを想像させたら、彼女はゼロからそれを学ばなければなりません。もう一人の生徒「ビデオ・プレディクター（動画予測者）」は、別の部屋に座って動いている動画だけを勉強しています。彼はボールが転がる様子や人が歩く様子を学びますが、写真の色が変わったらどうなるかについては全く知りません。

問題は、これら二人の生徒は同じ世界を学んでいるのに、異なる言語を話していることです。彼らはノートを共有できません。もし、写真の理解と未来の動きの予測の両方ができるロボットが欲しい場合、二つの別々の高価なモデルを訓練し、それらがうまく連携することを祈るしかありません。

UniJEYAは、両方の生徒を同じ部屋に入れ、一つの教科書を与えることで、このゲームのルールを変えます。この論文は、単一のモデルで以下の二つのことを同時に行えることを示しています。

1. フォトメトリック予測（光度予測）： 写真の明るさを上げたり、色相を変えたりすることを想像してください。UniJEPAは、画面上のピクセルを実際に変更することなく、その写真の「メンタルノート」が変化後にどのような状態になるかを予測することを学びます。
2. テンポラル予測（時間的予測）： ボールが転がる動画を見ている場面を想像してください。UniJEPAは、次のフレームの「メンタルノート」を予測し、ボールが次にどこにあるかを判断します。

魔法のような点は、UniJEPAがこれら二つのスキルを「同じメンタルスペース」の中で学習することです。それはまるで、ロボットが「光を変えること」と「時間を前に進めること」が、同じ根底にある現実に対する二つの異なる関わり方であることを理解しているかのようです。

仕組み：アンチ・コラップス（崩壊防止）のトリック

「もしロボットに一度に二つの異なることを学ばせたら、混乱して、あらゆるものに対して同じ退屈な答えを出力してしまうのではないか？」と疑問に思うかもしれません。AIの用語では、これは「コラップス（崩壊）」と呼ばれ、モデルが学習を停止し、平坦な出力しか出さなくなる状態を指します。

この論文では、**ガウス正規化（Gaussian regularizer）**と呼ばれる巧妙なセーフティネットを導入しています。これは、生徒たちが全員同じ椅子に座らないように監視する厳しい教師のようなものです。教師は、生徒たちが独自の、そして区別されたメンタルノートを持つように強制します。著者たちは、この単純なルールがあれば、複雑なトリック（「勾配停止」などの一般的なが、扱いにくい回避策）や、学習済みの脳を使用することなく、ロボットの崩壊を防ぐのに十分であることを数学的に証明しました。これは、学習を健全に保つための、クリーンで単一のルールです。

結果：より速く、より賢く、よりシンプルに

研究者たちは、画像（有名なImageNetデータセットなど）、動画（料理をしている人や手を使ったジェスチャーなど）、および制御タスク（迷路をナビゲートするロボットなど）を用いてUniJEPAをテストしました。得られた結果は以下の通りです。

• 万能選手であること： UniJEPAは、画像専用、あるいは動画専用として訓練された特化型ロボットと同等、あるいはそれ以上の性能を発揮しました。画像テストにおいて、UniJEPAは**74.9%の精度を記録し、特化型の「フォト・プレディクター」の73.5%**を上回りました。動画テストでは、**78.1%に達し、特化型の「ビデオ・プレディクター」の77.3%**を上回りました。
• プランニング（計画）マシンであること： 真のテストは計画能力でした。ロボットは目標に到達する方法を理解できるでしょうか？過去のデータを用いた追加のトレーニングの後、UniJEPAは、人間が道を示さなくても、視覚的な目標（例：「赤い四角形に到達する」）に到達するための計画を立てることができました。その成功率は**75.8%**でした。
• 驚異的な速さ： これが最大の勝利です。UniJEPAは、生のピクセルを毎回描き直そうとするのではなく、コンパクトな「メンタルノート」の中で予測を行うため、非常に高速です。論文によると、UniJEPAは生成的なワールドモデル（すべてのピクセルを描こうとするモデル）よりも最大44倍速く計画を立てることができます。他のモデルが動きを計画するのに数秒かかる一方で、UniJEPAはそれを一瞬で行います。

なぜこれが重要なのか

この論文は、あらゆる仕事に対して異なる脳を用意する必要はないと主張しています。機械への「見方」と「予測の仕方」の教え方を統一することで、より効率的で、訓練が容易（調整すべき主要な設定が一つだけで済む）で、より柔軟なシステムが得られるのです。

また、著者たちは、ロボットが何に集中して学習するかをコントロールできることも示しました。もしロボットに「写真」の部分に注意を払うよう指示すれば、照明の変化を無視すること（不変性）に長けたモデルになります。もし「動画」の部分に集中するよう指示すれば、動きの追跡（等変性）に長けたモデルになります。用途に合わせて、最適なバランスを見つけるために、このノブをスライドさせることができるのです。

要約すると、Uni-JEPAは、単一の統一された脳が、写真の静的な美しさと動画のダイナミックな流れの両方を理解することを学習でき、同時に驚異的な速度で次の行動を計画できることを証明しました。これは、細部に囚われることなく、私たちの世界を真に理解し、ナビゲートできるAIエージェントの構築に向けた一歩となります。