[論文レビュー] Ambient Diffusion: Learning Clean Distributions from Corrupted Data

Giannis Daras, Kulin Shah+4 more|arXiv (Cornell University)|May 30, 2023

Generative Adversarial Networks and Image Synthesis被引用数 9

ひとこと要約

拡散ベースのフレームワークは、拡散中に追加の計測歪みを注入することにより、クリーンなデータの分布を高く汚染されたサンプルのみから学習し、クリーンデータなしに真の分布からサンプリングを可能にする。

ABSTRACT

We present the first diffusion-based framework that can learn an unknown distribution using only highly-corrupted samples. This problem arises in scientific applications where access to uncorrupted samples is impossible or expensive to acquire. Another benefit of our approach is the ability to train generative models that are less likely to memorize individual training samples since they never observe clean training data. Our main idea is to introduce additional measurement distortion during the diffusion process and require the model to predict the original corrupted image from the further corrupted image. We prove that our method leads to models that learn the conditional expectation of the full uncorrupted image given this additional measurement corruption. This holds for any corruption process that satisfies some technical conditions (and in particular includes inpainting and compressed sensing). We train models on standard benchmarks (CelebA, CIFAR-10 and AFHQ) and show that we can learn the distribution even when all the training samples have $90\%$ of their pixels missing. We also show that we can finetune foundation models on small corrupted datasets (e.g. MRI scans with block corruptions) and learn the clean distribution without memorizing the training set.

研究の動機と目的

• 未調査・高コストな場合にデータ分布を学習する動機づけ。
• corrupted data を用いてクリーンデータの条件付き期待値を学習する拡散フレームワークを開発。
• 一般的な腐敗プロセスの下で汚染計測から学習するための理論的保証を提供。
• 標準ベンチマークと実世界の汚染データセットで実用的な有効性を示す。
• memorization を抑制しつつ、小規模な汚染データセット上で基盤拡散モデルの微調整を示す。

提案手法

• 拡散中に追加の計測歪みを導入し、より汚染された画像から元の汚染された画像を予測するように学習。
• スコアネットワーク h_theta を訓練し、与えられたさらに汚染された入力と A に対して E[x0 | Ãx_t, Ã] を推定。
• 訓練目的 J_corr(theta) を用い、||A(h_theta(Ã, Ãx_t, t) - x0)||^2 を (x0, x_t, A, Ã) に渡って最小化。
• 広範な腐敗設定で A^T A が全秩条件の下、 Ã に対して最適モデルが E[x0 | Ãx_t, Ã] を回復する理論を適用。
• 選択された A および Ã に対して E[x0 | x_t] を E[x0 | Ãx_t, Ã] で近似することにより、実用的な生成を可能にする。
• 拡散中にマスク間の一貫性を確保する Reconstruction Guidance によってサンプリングを強化。

Figure 1: Left panel: Baseline method of vanilla finetuning Deepfloyd IF using $3000$ images from CelebA-HQ. We show generated sample images and nearest neighbors from the finetuning set. As shown, the generated samples are often near-identical copies from training data. This verifies related work [

実験結果

リサーチクエスチョン

• RQ1 汚染された測定値 y0 = Ax0 のみが利用可能なとき、条件付き期待 E[x0 | y_t, A] を学習できるか。
• RQ2 もし可能なら、学習された復元を用いて E[x0 | x_t] を近似し、クリーンデータなしに真の分布からサンプルを得られるか。
• RQ3 訓練中に小さな追加汚染 Ã を導入することが、学習された分布やファインチューニング済みモデルの memorization にどのように影響するか。
• RQ4 標準ベンチマーク全体で、汚染レベルとジェネレータ品質の間にどんなトレードオフがあるか。
• RQ5 小規模な汚染データセット（例：医用画像）上で基盤拡散モデルをファインチューニングしつつ memorization を回避できるか。

主な発見

Dataset	Corruption Probability	Method	LPIPS	PSNR	NFE
CelebA-HQ	0.6	Ours	0.037	31.51	1
CelebA-HQ	0.6	DPS	0.053	28.21	100
CelebA-HQ	0.6	DDRM	0.139	25.76	35
CelebA-HQ	0.8	Ours	0.084	26.80	1
CelebA-HQ	0.8	DPS	0.107	24.16	100
CelebA-HQ	0.8	DDRM	0.316	22.96	99
CelebA-HQ	0.9	Ours	0.152	23.34	1
CelebA-HQ	0.9	DPS	0.168	20.89	100
CelebA-HQ	0.9	DDRM	0.332	18.74	99
AFHQ	0.4	Ours	0.030	33.27	1
AFHQ	0.4	DPS	0.020	34.06	100
AFHQ	0.4	DDRM	0.122	25.18	35
AFHQ	0.6	Ours	0.062	29.46	1
AFHQ	0.6	DPS	0.051	30.03	100
AFHQ	0.6	DDRM	0.246	20.76	35
AFHQ	0.8	Ours	0.124	25.37	1
AFHQ	0.8	DPS	0.107	25.30	100
AFHQ	0.8	DDRM	0.525	14.56	35
CIFAR-10	0.0	Ours	1.85	N/A	1
CIFAR-10	0.2	Ours	11.70	7.97	100
CIFAR-10	0.4	Ours	18.85	7.45	100
CIFAR-10	0.6	Ours	28.88	6.88	100
CIFAR-10	0.8	Ours	46.27	6.14	100
CIFAR-10	0.9	Ours	46.27	6.14	100

• 本手法は一般的な腐敗 A および Ã の下で条件付き期待 E[x0 | Ãx_t, Ã] を学習し、追加の歪み delta は小さくても正となり得る。
• 汚染データで訓練された拡散モデルは、クリーン分布に近いサンプルを生成でき、多くの設定で AmbientGAN やベースラインを上回る。
• CelebA-HQ、AFHQ、CIFAR-10 の実験で、最大 90% の欠損ピクセルでも合理的な生成品質を示す（FIDおよびInception Scores で測定）。
• learned restoration により、ベースラインが数百ステップを要するのに対して、少ないサンプリングステップで競合的または優れた性能を発揮。
• highly corrupted data でのファインチューニングにより memorization を低減し、生成サンプルの訓練画像との類似性が低下する一方で、脳 MRI などのドメイン移動を可能にする。
• コードとモデルはオープンソース化されている（論文内に GitHub リンクを提供）。

Random inpainting