[論文レビュー] Diffusion Models Beat GANs on Image Classification

Soumik Mukhopadhyay, Matthew Gwilliam+6 more|arXiv (Cornell University)|Jul 17, 2023

Generative Adversarial Networks and Image Synthesis被引用数 16

ひとこと要約

本論文は、事前学習済みの拡散モデルが統一的な自己監督表現として機能し、高い画像分類性能を達成できることを示し、256x256解像度で生成と分類の両方においてBigBiGANを上回り、FGVC転送も競争力を持つことを実証している。

ABSTRACT

While many unsupervised learning models focus on one family of tasks, either generative or discriminative, we explore the possibility of a unified representation learner: a model which uses a single pre-training stage to address both families of tasks simultaneously. We identify diffusion models as a prime candidate. Diffusion models have risen to prominence as a state-of-the-art method for image generation, denoising, inpainting, super-resolution, manipulation, etc. Such models involve training a U-Net to iteratively predict and remove noise, and the resulting model can synthesize high fidelity, diverse, novel images. The U-Net architecture, as a convolution-based architecture, generates a diverse set of feature representations in the form of intermediate feature maps. We present our findings that these embeddings are useful beyond the noise prediction task, as they contain discriminative information and can also be leveraged for classification. We explore optimal methods for extracting and using these embeddings for classification tasks, demonstrating promising results on the ImageNet classification task. We find that with careful feature selection and pooling, diffusion models outperform comparable generative-discriminative methods such as BigBiGAN for classification tasks. We investigate diffusion models in the transfer learning regime, examining their performance on several fine-grained visual classification datasets. We compare these embeddings to those generated by competing architectures and pre-trainings for classification tasks.

研究の動機と目的

• 識別タスクと生成タスクの両方を支える統一的な自己教師付き表現学習を動機づける。
• 拡散モデルの埋め込みが高精度な画像分類のために十分識別性を持つことを示す。
• 拡散特徴からの効果的な抽出とプーリング戦略を分類のために検討する。
• FGVCにおける拡散由来特徴の転移学習能力を評価する。
• 拡散層と時間ステップ全体の表現を特徴づけ、CKAを用いて他の事前学習手法と比較する。

提案手法

• 事前学習済みの無条件ガイド拡散モデル（ADM U-Net、256x256）を用いて、中間ブロックおよび拡散時間ステップから特徴を抽出する。
• 特徴抽出 f_theta(x0,t,b) を、ノイズ付き x_t を U-Net に通したときブロック b の直後の活性化として定義する。
• 線形プロービング、MLP/CNN/アテンションヘッド、さまざまなプーリング戦略を評価して、特徴マップを分類用のベクトルへ変換する。
• ImageNet-1k における精度と FID の観点で、拡散ベースの分類器を BigBiGAN および MAGE のベースラインと比較する。
• FGVC データセットへの転移性を評価し、CKA (Centered Kernel Alignment) を用いて表現を分析する。
• 時間ステップ t、ブロック指標 b、プーリングサイズのアブレーションを実施して、最適な特徴抽出設定を特定する。

実験結果

リサーチクエスチョン

• RQ1 拡散モデルの埋め込みを拡散モデルのファインチューニングを行わずに、識別的な画像分類に再利用できるか。
• RQ2 拡散特徴をどのようにプーリングし、分類するのが最適か（線形、MLP、CNN、アテンションヘッドなど）。
• RQ3 ImageNetおよびFGVCタスクにおいて、拡散由来の表現はGANベースおよび自己教師付きのベースラインとどのように比較されるか。
• RQ4 拡散特徴は細分類データセットへよく転移するか、タスク間でハイパーパラメータに対してどれくらい感度があるか。
• RQ5 CKA により測定されたように、拡散表現は層と拡散時間ステップに沿ってどのように変化するか。

主な発見

方法	精度	FID
BigBiGAN*	60.8%	28.54
MAGE	78.9%	9.10
U-Net Encoder	64.32%	n/a
GD (L, pool 1x1)	61.95%	26.21
GD (L, pool 2x2)	64.96%	26.21
GD (Attention)	71.89%	26.21

• 拡散モデルは、固定特徴を用いた注意機構ヘッド（b=24、t=90）で ImageNet-1k において 61.95% の精度を達成し、分類で BigBiGAN を上回る。
• GD variants は同じ FID 26.21 で 64.96%（L、プール 2x2）および 71.89%（Attention）の精度に達し、いくつかの統合ベースラインを上回る。
• 拡散特徴を用いた線形プロービングは、ImageNet-1k で注意ヘッドを用いて 61.86% の精度（b=24、t=150）を得る。安定した拡散特徴も分類をサポートする。
• FGVC データセットでは拡散特徴が競争力のある性能を示し、Aircraft は複数のヘッドで SimCLR ベースのベースラインを上回ることが多いが、他のデータセットでは依然差が残る。
• CKA 分析は、初期層がモデル間でより類似する一方、ボトルネック層が ResNet/ViT に類似した識別的表現を生み出すことを示す。