注釈付きJEPA

この記事は、Joint Embedding Predictive Architectures、略してJEPAの、ゼロからの段階的で注釈付きのウォークスルーです。JEPAにとって、TransformerがTransformerにとって行ったことと同じことを行うことを目指します。つまり、完全なオブジェクトを構築し、すべての可動部分を説明し、動作するトレーニングループで終わらせることです。 JEPAは、自己教師あり学習における根本的な問いに対するYann LeCunが提案した答えです。ラベルなしで、自明な解に崩壊せず、無関係な詳細に容量を浪費することなく、モデルに世界を理解させるにはどうすればよいでしょうか？その答えは、原理的にはエレガントで、実践では微妙ですが、表現または潜在空間での予測です。 議論を具体的に保つために、主な実行例はI-JEPA、つまり画像インスタンシエーションです。なぜテキストではなく画像とビデオなのか？LeCunは、言語はすでに高度に圧縮された離散的な知識表現であると主張しています。次のトークンを予測するには、物理的現実ではなく、人間のコミュニケーションパターンをモデル化する必要があります。対照的に、視覚的予測は、永続性、隠蔽、およびダイナミクスを理解する必要があります。JEPAは、ピクセルレベルの再構築が自明な解に崩壊する問題が発生しないドメイン向けに設計されています。これは、離散トークンでは同じようには発生しない問題です。私たちは、これについては終盤で再び取り上げます。これは、マスクされた領域の表現を可視コンテキストから予測することによって意味的な画像表現を学習する自己教師あり手法です。I-JEPAは、完全に手作りのデータ拡張を回避する非生成的アプローチです。私たちは、ゼロからI-JEPAを構築し、次にビデオへの拡張であるV-JEPAおよびV-JEPA 2、そしてエンジニアリングのヒューリスティックを分布正規化器に置き換える最新の試みであるLeJEPAを検討します。 以下は教育的な目的で書かれています。実装では、大規模トレーニングを可能にするFlashAttention、勾配チェックポインティング、混合精度、およびバッチ処理戦略は省略されています。これらは、プロダクションコードベースを支配するエンジニアリングの選択ですが、数学からは容易に分離可能です。 問題 自己教師あり表現学習は、ラベルなしで有用な特徴をどのように学習するかを問います。意味のある構造を捉える目的関数が必要ですが、ラベルなしで実際に機能するものを見つけることが、この分野の中心的な困難です。JEPAの答えは、予測によってトレーニングするが、表現空間で予測することです。 しかし、これはそもそもなぜうまくいくのでしょうか？画像の一部、コンテキスト（x）を見て、表現を学習したいとします。画像内のどこかに、見えないターゲット領域（y）があります。エンコーダーはyを表現syにマッピングします。予測器はコンテキストのエンコーディングを受け取り、syの推測であるsŷを出力します。トレーニングは、sŷとsyの距離D(sŷ, sy)を最小化します。さて、予測器はいつ成功できるでしょうか？コンテキストエンコーディングsxがsyを決定するのに十分な情報を含んでいる場合にのみ成功できます。車のボンネットを見た場合、ホイールの表現を予測するには、ボンネットのエンコーディングがこれが車であることを捉える必要があります。顔を見た場合、髪の表現を予測するには、エンコーディングがアイデンティティ、ポーズ、および照明を捉える必要があります。予測器は、コンテキストエンコーディングに欠けている構造を幻視することはできません。これが強制関数です。コンテキストエンコーダーは、xからyの表現を予測可能な特徴を抽出することを学習する必要があります。これらはまさに意味的、構造的な特徴です。オブジェクトのアイデンティティ、空間的関係、物理的制約です。xのピクセルレベルのノイズはsyの予測に役立たないため、エンコーダーはそれを無視することを学習します。残ったものが一般化するものです。 ターゲットエンコーダーには補完的な圧力があります。その出力syは、コンテキストから予測可能である必要があります。syがランダムな高周波テクスチャをエンコードした場合、コンテキストはそれを予測するのに役立ちません。したがって、ターゲットエンコーダーは、xとyの間の共有構造、予測を可能にする構造を捉える表現を出力することを学習します。単にyの特異な詳細を捉えるのではなく。 LeCunのポジションペーパーは、これをエネルギーベースの定式化としてフレーム化しています。xとyを表現にエンコードし、一方を他方から予測し、エネルギーをその抽象空間での予測誤差として定義します。アーキテクチャは、実装および分析可能なコンポーネントに分解されます。2つのエンコーダー、1つの予測器、1つの距離関数です。 JEPAテンプレート JEPAは、ペアになった意味的に関連する世界のビューから始まります。最も一般的な形式では、トリプル(x, y, z)を想像してください。xは観察するもの、yは予測したいもの、zは予測を多峰性にする未知の要因を捉えるオプションの潜在変数です。(x, y)のペアは、観察に関する同時分布から引き出されます。画像事前学習では、xは可視パッチ、yは同じ画像のマスクされたパッチである可能性があります。ビデオでは、xはクリップのプレフィックス、yは継続です。クロスモーダル設定では、xはオーディオ、yは対応するビデオである可能性があります。唯一の構造的要件は、xを知ることがyが何であるかを制約することです。 LeCunの定式化は、多峰性を予測関係で処理するために、オプションの潜在変数zを追加します。同じxに対して複数のyの値が一致する場合、予測器はzを条件としてそれらを選択します。これは時間的予測では重要ですが、将来は実際に不確実です。マスクされた画像モデリングでは、コンテキストが通常ターゲットをノイズまで決定するため、それほど重要ではありません。 テンプレートには3つの部分があります。両方の観察を共有表現空間にエンコードします。 sx = fθ(x), sy = f¯θ(y) エンコーダーfθとf¯θは、アーキテクチャまたはパラメータ共有で異なる場合があります。xとyが異なるモダリティに属する場合、それらは異なる必要があります。同じモダリティの場合、重み共有は、帰納バイアスと柔軟性のトレードオフです。 コンテキスト表現からターゲット表現を予測します。 ŷsy = gφ(sx, z) 表現空間での予測誤差を最小化します。 L(θ, ¯θ, φ) = D(ŷsy, sy) この記事の残りの部分は、これらの抽象化を具体化します。画像を具体的にxとyを定義し、エンコーダーと予測器のアーキテクチャを選択し、距離関数Dを指定し、すべての表現が定数に崩壊するという自明な解を防ぐメカニズムを実装します。 予測 vs. 再構築 なぜピクセル空間ではなく表現空間で予測するのでしょうか？意味的には同じで、同じオブジェクト、同じ意味を持つ2つの画像パッチを考えます。ただし、照明、テクスチャ、またはJPEGアーティファクトのために数百ピクセル異なります。再構築目的関数は、それらのすべての違いを説明する必要があります。モデルは、高エントロピーノイズをモデル化するために容量を浪費します。JEPAを使用すると、ターゲットエンコーダーは、不要な詳細を破棄する表現を出力することを学習できます。その表現に回帰すると、エンコーダーが保持するものを捉えるようにモデルにプッシュします。ピクセルが偶然含むものではなく、I-JEPAはこれを明確にします。アーキテクチャは生成モデルに似ていますが、損失は入力空間ではなく埋め込み空間にあります。これは知覚的損失（LPIPS）とピクセル損失（MSE）に類似しています。表現空間での予測により、モデルは意味的に意味のある特徴に焦点を当てることができます。 崩壊 マッチング損失を最小化する最も簡単な方法は何でしょうか？すべてに同じベクトルを出力することです。そうすれば、損失はゼロになります。しかし、表現は...（本文はここで途切れています）