論文要約：Structured State-Space Regularization for Compact and Generation-Friendly Image Tokenization

本論文は、現代のビジョンモデル（特に拡散モデル）の中核をなす「画像トークナイザー（Image Tokenizer）」の改良を目的とした研究です。著者らは、状態空間モデル（SSM: State-Space Models）の隠れ状態のダイナミクスを模倣させるための新しい正則化手法「構造化状態空間正則化（Structured State-Space Regularization）」を提案し、画像の潜在空間（Latent Space）を「コンパクト」かつ「生成に親和的（Generation-Friendly）」な特性を持つように導くことに成功しました。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

1. 背景と問題定義

近年の生成 AI や世界モデルは、高次元のピクセル空間ではなく、学習された低次元の潜在空間で処理を行うパラダイムへ移行しています。この際、画像を潜在表現に変換する「トークナイザー」の性能がシステム全体の効率と品質を決定づけます。

理想的なトークナイザーの潜在空間には、以下の 2 つの相反する要件が求められます。

1. コンパクト性（Compactness）: 画像の本質的な情報を効率的に圧縮し、限られた次元で最大限の情報を保持すること。
2. 生成への親和性（Generation-Friendliness）: 拡散モデルなどの生成モデルが、その潜在空間から高品質な画像を生成しやすいこと。

既存の研究では、これら 2 つの要件のトレードオフ（再構成精度を上げると生成性能が落ちる傾向）が課題となっていました。特に、生成モデルは画像の低周波成分から高周波成分へ順次合成するプロセスと解釈されるため、潜在空間が周波数構造を適切に捉えていることが重要ですが、従来のトークナイザーはこれを明示的に保証していませんでした。

2. 提案手法：構造化状態空間正則化

著者らは、画像トークナイザーの潜在空間に、状態空間モデル（SSM）の「周波数認識（Frequency Awareness）」という特性を転移させることを提案しました。

2.1 理論的基盤：SSM の一般化

SSM は通常、隠れ状態 $x(t)$ が入力 $u(t)$ を受け取り、行列 $A$ によって更新されるシステムとして定義されます。著者らは、SSM の核心を以下の 2 つのコンポーネントの組み合わせとして再定義しました。

1. 基底射影（Basis Projection, $c(\cdot)$）: 入力信号を基底関数の係数に圧縮する操作。
2. 入力変換（Input Transformation, $\theta(\cdot)$）: 入力信号に対する時間的な変化（例：画像のぼかし）。

SSM の隠れ状態は、この基底係数のダイナミクス（時間変化）を模倣するように学習されます。このとき、基底係数は信号の周波数成分に対応するため、SSM は自然に周波数情報を捉える能力（周波数認識）を獲得します。

2.2 具体的な正則化の導出

提案手法では、画像トークナイザーのエンコーダ $E$ に対して、以下のプロセスを正則化項として導入します。

1. 基底の選択: 2 次元の直交基底（フーリエ基底、チェビシェフ多項式など）を選択し、画像をその係数 $c(I)$ として表現します。
2. 入力変換の選択: 画像に対してガウスぼかし（Gaussian Blur）を適用するプロセス $\theta$ を選択します。ガウスぼかしは熱方程式の解と等価であり、高周波成分から順に減衰していく性質を持ちます。
3. 更新則の導出: ガウスぼかしを施した画像 $I_\tau$ の基底係数 $c(I_\tau)$ の時間微分を計算すると、以下の微分方程式が得られます。
   $$\frac{d}{d\tau} c(I_\tau) = \mathbf{A} c(I_\tau)$$
   ここで、$\mathbf{A}$ は基底関数のラプラシアン演算子による内積から構成される行列です。離散化すると、$c(I_t) = \bar{\mathbf{A}} c(I_{t-1})$ という更新則が得られます。
4. 正則化損失の定義: トークナイザーのエンコーダ $E$ に対して、ぼかしの異なる 2 枚の画像 $I_{\tau_1}, I_{\tau_2}$ を入力し、それぞれの潜在表現 $z_1, z_2$ を得ます。$z_1$ に行列 $\bar{\mathbf{A}}$ を作用させた値 $\hat{z}_2$ と、実際の $z_2$ の距離を最小化するように学習します。
   $$\mathcal{L} = d_z(z_2, \bar{\mathbf{A}} z_1) + \lambda_I d_I(I_{\tau_2}, D(\hat{z}_2))$$
   これにより、潜在空間の更新が「ガウスぼかしに伴う周波数成分の減衰」という物理的なダイナミクスに従うように強制されます。

2.3 実装上の工夫

• 空間平均中心化（Spatial Mean-Centering）: ぼかしが進むと潜在表現のノルムが 0 に収束してしまう問題を回避するため、潜在値から空間平均を引く操作 $\mathcal{C}$ を導入し、トレーニングの安定性を確保しています。
• 基底の選択: 実験ではフーリエ基底から導出された行列 $\mathbf{A}$ が最も優れた結果をもたらしました。

3. 主要な貢献

1. 新しい正則化手法の提案: 理論的に動機づけられた「構造化状態空間正則化」を提案し、画像トークナイザーの潜在空間を SSM の特性（周波数認識）に近づけました。
2. SSM と画像トークナイゼーションの統合: SSM の原理（基底射影と入力変換）を画像の 2 次元データに適用する初の理論的枠組みを提供しました。
3. 実証的な有効性: 広く利用されているトークナイザー（Flux, Cosmos）に適用し、再構成精度をほぼ維持しつつ、拡散モデルによる生成性能を大幅に向上させることを示しました。

4. 実験結果

ImageNet-1K データセットを用いた評価において、以下の結果が得られました。

• 生成性能の向上: 提案手法を適用した Flux トークナイザーを用いた拡散モデル（DiT-B/2）は、ベースライン（Flux）や既存の正則化手法（EQVAE, FT-SE）と比較して、生成 FID（gFID）や Inception Score（IS）で顕著な改善を示しました。
  • 例：Flux の gFID は 16.08 から 13.33 へ改善。
• 再構成精度の維持: 生成性能の向上に伴い、PSNR や LPIPS などの再構成指標はわずかに低下するものの、実用上は許容範囲内であり、トレードオフを最小限に抑えています。
• 潜在空間の可視化:
  • 周波数構造の獲得: 潜在チャネルを低周波から高周波の順に順次復元（Unmasking）すると、画像が粗い輪郭から細部へと段階的に再構成されることが確認されました。これは、各チャネルが特定の周波数帯域を効率的にエンコードしていることを示しています。
  • ダイナミクスの一致: 画像がぼやける過程における潜在表現の変化が、理論的に導出されたフーリエ係数の減衰軌跡とよく一致していました。

5. 意義と結論

本論文は、「状態空間モデル（SSM）」の数学的構造を「画像トークナイザー」の設計に応用するという新たな視点を開拓しました。

• 理論的意義: SSM が持つ「周波数分解能」や「長期的な依存関係のモデル化」という特性が、画像の潜在表現においても有効であることを示しました。
• 実用的意義: 生成 AI システムにおいて、より高品質な画像生成を可能にする効率的な潜在空間を提供します。これは、計算コストの削減と生成品質の向上を両立させるための重要なステップです。
• 将来展望: 本手法は、SSM を画像表現学習の強力なツールとして位置づけ、今後の生成モデルや表現学習フレームワークへの SSM 原理の統合への道を開いています。

要約すれば、この研究は「画像を周波数成分として捉え、その時間的変化（ぼかし）に従って潜在空間を整理する」というアプローチにより、生成 AI がより自然で高品質な画像を生成できる土台を整えた画期的なものです。