PeLAP-A: Adaptive Latent Pruning for Lightweight Latent Diffusion Models

技術要約: PeLAP-A (Adaptive Latent Pruning for Diffusion)

問題提起
潜在拡散モデル（LDM）は、変分オートエンコーダ（VAE）によって生成された圧縮された潜在空間で動作することにより、高品質な画像合成を実現している。しかし、すべての潜在チャネルが拡散プロセスに等しく寄与しているのか、あるいは大幅な冗長性が存在するのかについては、依然として未解決の問いである。既存の研究は、拡散モデルのアーキテクチャ自体（例：重みのプルーニング）やサンプリングプロセスの圧縮に焦点を当てているが、潜在表現自体の情報量についてはほとんど探索されていない。著者らは、生成品質を低下させることなく、冗長な潜在チャネルを特定し、選択的に抑制できるかどうかを調査している。

手法: PeLAP-A
本論文では、標準的な潜在拡散パイプラインに、学習可能で入力適応型のチャネル別重要度予測器を付加する軽量フレームワークである PeLAP-A (Adaptive Latent Pruning for Diffusion) を提案する。

• アーキテクチャ: わずか292個のパラメータを持つ2層の多層パーセプトロン（MLP）が、VAEエンコーダとデノイジングUNetの間に挿入される。
• メカニズム:
  1. VAEエンコーダからの潜在テンソル $z$ は、グローバル平均プーリングを経てチャネルごとの要約が生成される。
  2. この要約はMLPによって処理され、シグモイド活性化関数を通過することで、ソフトマスク $m \in (0, 1)^C$ を生成する。
  3. マスクは空間方向にブロードキャストされ、潜在テンソルに対して要素ごとに適用される（$z' = z \odot m$）。これはUNetに入る前に行われる。
  4. 予測器の最終線形層は、バイアス値 +3.0 で初期化されており、これによりすべてのチャネルがほぼ完全に開いた状態（$\sigma(3) \approx 0.95$）から始まり、時期尚早な抑制を防いでいる。
• 学習目的関数: システムは、以下の結合損失関数を用いてCIFAR-10上でエンドツーエンドで学習される：
  $$L = L_{vae} + L_{diff} + \lambda \cdot \frac{1}{C}\sum_{c=1}^{C} m_c$$
  ここで、$L_{vae}$ は再構成およびKLダイバージェンス損失、$L_{diff}$ は拡散ノイズ予測のMSEであり、第3項はマスクの平均値をゼロへと向かわせるスパース性正則化項である。

主な結果
CIFAR-10を用いた実験により、**「スパース性崩壊（sparsity collapse）」**と呼ばれる驚くべき現象が明らかになった：

• 崩壊のダイナミクス: 強力なスパース性正則化（$\lambda = 0.01$）の下では、重要度予測器が最初の2エポック以内にすべての潜在チャネルをほぼゼロの値へと駆動してしまう。
• デノイジング性能: 潜在情報の完全な抑制にもかかわらず、PeLAP-Aモデル内のデノイジングUNetは、プルーニングされていないベースラインよりも低い検証指標を達成した：
  • 拡散損失 (Diffusion Loss): 0.0236 (PeLAP-A) vs 0.0240 (Baseline)。
  • VAE再構成MSE: 22.59 (PeLAP-A) vs 24.67 (Baseline)。
• 生成の失敗: デノイジング経路は改善されたものの、生成出力の品質は崩壊した。FIDスコアは、278.1 (Baseline) から 362.6 (PeLAP-A) へと大幅に悪化した。定性的分析によれば、潜在値がほぼゼロになると、デコーダは構造化された画像ではなく、一様なグレーの出力を生成した。
• 残留構造: 崩壊後であっても、マスクの大きさはクラスに依存した順序（例：「自動車」クラスは「飛行機」クラスよりもわずかに高い平均活性化を保持していた）を維持しており、これは予測器が閾値以下の大きさであってもクラスに関連する情報を符号化していたことを示している。

意義と主張
本論文は、PeLAP-Aを潜在拡散学習におけるスパース性ダイナミクスの探索的研究として位置づけている。その主な主張と貢献は以下の通りである：

1. UNetの堅牢性: これらの知見は、デノイジングUNetが潜在チャネルの抑制に対して極めて堅牢であることを示している。UNetは、ほぼゼロに近い潜在空間に対して効果的な無条件事前分布を学習することができ、これはデノイジングプロセスにおける情報要件が、潜在空間の次元数が示唆するよりもはるかに低いことを示唆している。
2. 拡散損失の限界: 改善された拡散損失と悪化したFIDの間の乖離は、潜在空間が崩壊した場合、拡散損失単体では生成品質の不十分なプロキシ（代理指標）であることを浮き彫りにしている。潜在的なプルーニングを評価するには、FIDのような知覚的指標が不可欠である。
3. スパース性崩壊現象: 著者らは、スパース性の圧力が学習の初期段階において、有益な勾配の学習を圧倒してしまう急激な相転移を特定した。彼らは、$\lambda = 0.01$ がこの特定のアーキテクチャにとって臨界閾値を超えており、システムが構造化されたデータではなく、ゼロに近いテンソルをデノイズするように適応する局所解に陥ったと仮定している。
4. 限定的な範囲: 著者らは、実装された手法が、情報を保持するチャネルを選択的に残しながら他を抑制するという目標を達成できなかったことを明示している。本研究は、生成モデルにおけるスパース性の臨界閾値を理解するための基礎となるものであり、部分的なプルーニングを実現するためのウォームアップ・スケジュールの導入、離散緩和、あるいはより大きな潜在空間での学習といった将来の研究を動機付けるものである。