技術要約：ニューラル損失がいかにVAEの潜在変数（Latents）を形作るか

問題提起

現代の変分オートエンコーダ（VAE）は、通常、$\beta$-VAEの目的関数が示唆する点ごとの尤度（一般的には平均二乗誤差、すなわち画素単位のSSE）を用いて訓練されることは稀である。実際には、画素損失を、知覚的損失（例：LPIPS、VGG特徴量）や敵対的損失（例：PatchGAN）といった複雑なニューラル目的関数で拡張または置き換える手法へとシフトしている。これらの修正は再構成の品質と知覚的な忠実度を向上させるが、これらの特定の歪度（distortion）指標が、潜在空間の根本的なダイナミクスをどのように変化させるかについての理論的な理解が不足している。具体的には、点ごとの損失からニューラル損失へと移行することで、レート歪み（rate-distortion, RD）のトレードオフや、学習された事後分布の内部幾何学がどのように再形成されるのかが不明である。

手法

著者らは、歪度関数を単なる静的な実装の詳細としてではなく、レート歪み最適化問題における主要な変数として再定義することで、この問題にアプローチしている。彼らの手法は、シャノンのレート歪み理論の理論的分析と、訓練されたVAEを用いた経験的検証を組み合わせたものである。

1. 理論的枠組み:

   • 著者らは、**点ごとのアフィン優位性（pointwise affine domination）**を用いて、歪度関数間の「弱さ」という概念を定式化している。歪度 $d_2$ が $d_1$ よりも弱いとされるのは、$d_2(x, \hat{x}) \leq c \cdot d_1(x, \hat{x}) + b$ が成立する場合である。
   • この定義を用いて、著者らは、もしある歪度関数がより弱い場合、最適なシャノン・レート歪み曲線 $R(\Delta)$ は、より強い歪度の曲線よりも下側、あるいは同等になることを証明している。これは、弱い歪度の方が、与えられた歪度予算に対してより少ない情報量（より低いレート）を必要とすることを意味する。
   • これを特定のニューラル損失に適用する：
     • 知覚的損失: 特徴抽出器がリプシッツ連続であるという仮定の下で、特徴空間のMSE（例：VGG）が画素空間のMSEによって支配されることを証明する。
     • 敵対的損失: ディスクリミネータのスコアに基づくPatchGAN生成器の損失が、画素SSEによってアフィン優位となる有効な歪度指標として書き換えられることを証明する。

2. 経験的検証:

   • レート分析: 著者らは、pythae や diffusers のアーキテクチャを用いたVAEを、統一された目的関数 $L = B(\lambda D + (1-\lambda)\text{SSE}) + \beta \text{KL}$ を用いてデータセット（CelebA, Tiny-ImageNet）上で訓練している。ここで、$D$ はニューラル歪度（LPIPS, DINOv2, またはPatchGAN）を表し、$B$ は大きさを正規化するための損失バランサーである。彼らは $\lambda$（ニューラル損失の重み）と $\beta$ をスイープし、収束するKLダイバージェンス（レート）を観察する。
   • 幾何学分析: 情報レート全体とは独立して、歪度が幾何学に与える影響を分離するために、訓練中に $\beta$ を動的に調整して、モデルを特定の目標レート（$\hat{R}$）に強制的に収束させる GECO アルゴリズムを採用している。そして、サンプルごとの事後分布の異方性（per-sample posterior anisotropy）、すなわち潜在次元間の対数分散の分散（$A_{post}(x) = \text{Var}[\log \sigma^2_{\phi,i}(x)]$）を測定する。
   • トイモデル: 線形ガウスモデルを用い、歪度係数が分散の「ウォーターフィリング（水詰め）」配分にどのように影響するかを解析的に示し、歪度の幾何学がスカラーのレートが固定されている場合でも事後分布の異方性を決定することを確認する。

主な貢献

1. ニューラル損失は潜在変数に含まれる情報量を減少させる

本論文は、一般的な知覚的および敵対的目的関数が、画素単位の二乗誤差よりも「弱い」歪度尺度であることを証明し、検証している。

• 理論的結果: 画素レベルの近接性が特徴レベルの近接性を意味するため（リプシッツ連続性による）、画素SSEの予算を満たすことは、自動的に知覚的予算を満たすことになる（スケーリング後）。したがって、ニューラル損失の理想的なRD曲線は、画素SSEの曲線よりも下方に位置する。
• 経験的結果: 訓練されたVAEにおいて、ニューラル損失の重み（$\lambda$）を増加させると、固定された $\beta$ に対して収束時に到達するKLダイバージェンスが一貫して減少する。これは、ニューラル損失を用いることで、再構成目的を達成しながらも、インスタンス固有の情報を潜在表現により少なく蓄えることが可能であることを示している。

2. 歪度は潜在的な不確実性を形作る（幾何学）

著者らは、歪度関数の選択が、総情報レートとは独立して、事後分布の不確実性のプロファイルを支配することを実証している。

• 異方性 vs 等方性: 点ごとの二乗誤差（画素SSE）は、少数の高度に異方的な次元（アクティブなユニット）に精度を集中させ、他の次元を事前分布に近く残す。対照的に、知覚的および敵対的損失は、潜在次元における分散のより等方的な分布を促す。
• メカニズム: これは「ウォーターフィリング」の類推によって説明される。画素SSEはすべての画素誤差を等しく罰するため、エンコーダはデータ多様体の高分散方向を保持しようとし、結果として分散が集中する。一方、ニューラル損失は、画素レベルの方向が知覚的な等価クラスへと集約された特徴空間で動作するため、特定のパターンを記憶する圧力が取り除かれ、レートがより均一に分配されるようになる。
• 逆説的な発見: 画素SSEは画素空間において最も等方的（すべての画素を等しく重み付け）であるが、最も異方的な事後分布を生み出す。逆に、画素空間では異方的であるニューラル損失は、最も等方的な事後分布を生み出す。

結果

• レート歪み曲線: CelebAおよびTiny-ImageNetを用いた実験では、知覚的または敵対的損失の重みを増すと、レート歪みトレードオフのパレート最適フロンティアが下方へシフトすることが示された。固定された $\beta$ に対して、ニューラル損失が存在する場合、モデルはより低いKLレートで収束する。
• 事後分布の幾何学: GECOを用いて固定レートで訓練を行った際、ニューラル損失の重みを増すと、事後分布の異方性が単調に減少した。分散は潜在次元間でより均一に分布するようになる。この効果は、異なるアーキテクチャ（ベクトル潜在変数VAEおよび空間潜在変数AutoencoderKL）およびデータセットにわたって認められた。
• トイモデルによる検証: 線形ガウス型のトイモデルは、歪度指標を変更すること（等方的から異方的へ）が、同じスカラーの分散-KL予算を維持しながら、事後分散のプロファイルを再形成することを裏付けている。

意義と主張

本論文は、伝統的に $\beta$ パラメータの観点から捉えられてきたレート歪みトレードオフは、現代のVAEを理解するための不完全なレンズであると主張している。歪度関数の選択は、単なる再構成品質のためのハイパーパラメータではなく、潜在的な挙動を決定する主要な駆動要因である。

• 潜在設計: 著者らは、知覚的および敵対的損失を、容量（達成されるレートを下げる）および形状（情報を等方的に再分配する）の直接的なコントローラーとして位置づけている。これは、損失の選択が理論的な変数ではなく、単なるエンジニアリング上のヒューリスティックとして扱われてきた文献におけるギャップを埋めるものである。
• 生成モデリングへの示唆: これらの知見は、潜在空間の「形状」（等方性 vs 異方性）が、VAEの潜在変数上で拡散モデルを訓練するといったダウンストリームのタスクにとって重要な要因であることを示唆している。論文では、再構成品質（FID）と生成品質がしばしば相反することに触れており、ニューラル損失によって誘導される潜在幾地学が、潜在変数に蓄えられる情報量を減少させたとしても、生成には有益である可能性があることを示唆している。
• 限界: 著者らは、理論的証明が理想的なシャノンRD曲線と簡略化された仮定（リプシッツ連続性、線形デコーダなど）に基づいていることを認めている。SGDで訓練された有限かつ非凸なVAEにおける経験的結果は、これらの理論的傾向が実用においても成立することを示す証拠であるが、効果の正確な大きさは最適化のダイナミクスに依存する。

要約すると、本論文はVAEに対するメカニズム的なアプローチを提案している。すなわち、歪度指標は根本的に最適化問題を再形成し、どれだけの情報が蓄えられ、その情報が潜在空間内でどのように幾何学的に配置されるかを決定するというものである。