【DL輪読会】Masked Diffusion as Self-supervised Representation Learner　

DEEP LEARNING JP [DL Papers] Masked Diffusion as Self-supervised Representation Learner Yuta Oshima, Matsuo Lab http://deeplearning.jp/

• http://deeplearning.jp/

Masked Diffusion as Self-supervised Representation Learner 書誌情報 著者 • Zixuan Pan, Jianxu Chen, Yiyu Shi 概要 • 従来の拡散モデルに変更を加え，表 現学習に特化したモデルを作成． • 表現を用いた下流タスク（主にセグ メンテーション）の性能が良い． 2

背景 Denoising Diffusion Probalistic Model (DDPM) [1] • データx0 に徐々にガウシアンノイズを加えていき，完全なガウシアンノイズ xT ~𝑁(xT ; 0, 𝐼)にすることを考える（拡散過程）． • 生成は，拡散過程を逆向きに辿る（逆過程）．逆過程はNNでモデル化する． • 拡散モデルは高品質な生成のアプリケーションがよく研究されているが，本研 究では，拡散モデルを用いた表現学習が目的． 3

背景 Masked Autoencoders (MAE) [2] • 自己教師あり表現学習手法の一つ． • 入力画像をパッチに分割し，固定のマ スク率でパッチを[mask]する． • 可視部分を用いて，欠損部分を予測す る（マスク予測）ことで事前学習． • 本研究でもマスク予測により表現学習 を行なっている． 4

背景 Diffusion Models as Masked Autoencoder (DiffMAE) [3] • MAEのマスク予測をノイズ除去過程と みなして表現学習を行う． • MAEよりも生成品質は向上したが，表 現を用いた分類やセグメンテーション で，MAEにわずかに劣る． • 一方，本研究では，セマンティックセ グメンテーションにおいて，MAEを凌 駕． 5

手法 Masked Diffusion Model (MDM) を提案 • 従来のガウスノイズによる拡散モデルをやめ，マスク付き拡散モデルを導入． • 入力画像をパッチに分け，タイムステップ𝑡 において，マスク率 𝑅𝑚 = 𝑡 𝑇+1 でマスキング．𝑡 が小さくなると，ガウシアンノイズではなく，マスクが減る． • 表現学習に特化したモデルであり，DDPM等で行われる多段階の生成はできない． 6

手法 Masked Diffusion Model (MDM) を提案 • 従来のガウスノイズによる拡散モデルをやめ，マスク付き拡散モデルを導入． • MAEとは異なり，マスクを含めて画像全体をUNetに入力する． • 再構成誤差にはMSE Loss ではなく，SSIM Loss を用いる． • SSIMを用いることで，より意味のある表現を獲得できたと主張している． 7

手法 Masked Diffusion Model (MDM) を提案 • 表現抽出器が学習されたら，重みを凍結し，下流タスクを学習する． • 本論文では，下流タスクとしてセグメンテーションを主に扱う． • セグメンテーションの学習に用いるデータは，必ずしもMDMの学習に用いたデータの部分集 合でなくても良い．ドメイン（医療画像，顔画像，etc…）は同一でないといけない． 8

実験 先行研究との比較 • 医療画像のセグメンテーション • 全体の10%のラベルしか利用できない場合でも，高精度を達成． 9

実験 先行研究との比較 • 顔画像のセグメンテーション • DatasetGAN，DatasetDDPMはGANやDDPMを用いた教師ありのセグメンテーション手法． 10

実験 先行研究との比較 • まとめ • MDMは4つ全てのデータセットで，DDPMより優れる． • 拡散モデルによる表現学習においては，必ずしもノイズ除去は必要ない可能性． • MDMは，MAEよりも優れる． • tを変えることでさまざまなマスク率で学習できることが，頑健で意味のある表現学習に有 利な可能性． • MDMは利用可能なラベルが少ない場合でもうまく動作する． • 医療画像のセグメンテーションなど，アノテーションコストが高い場面で有益． 11

実験 アブレーション • MDMとDDPMにおいてtを固定して検証． • MDMでは𝑡 = 50 ，DDPMではt = 250で固定する． • MDMの方が，DDPMより， tを動かすことで大きな恩恵を受けるうえ，性能も高い． • 拡散モデルで表現学習をしたい場合，必ずしもノイズ除去である必要はない可能性を強調． 12

実験 アブレーション • 損失関数について • SSIM Loss と MSE Lossを比較し，SSIM Lossの優位性を確認． • また，DDPMに関しては，ノイズ予測型の損失関数の方が，𝑥0 を直接予測する型の損失関数(†) よりも，表現学習において優れることを示した． 13

まとめ • 表現学習に特化した拡散モデルである MDMを提案． • マスク予測で学習される拡散モデル． • タイムステップtが変化すると，ノイズの大き さではなくマスク率が変化． • 主にセグメンテーションで表現の性能を 評価． • 医療画像，顔画像のセグメンテーションで先 行研究を上回る結果． • 少数データでも高性能． 14

感想 興味深かった点 • 少数の医療画像で高い性能を出していたので応用面で期待が高そう． 考える余地のある点 • “Masked Diffusion as Self-supervised Representation Learner” というタイトルだが，評価に用いたタス クはセグメンテーションだけ． • 表現学習というからには，複数のタスク（分類，重要点推定，深度推定など）で優位性を見たかった． • 筆者はこれまで医療系の深層学習の研究を多く出していたので，多分筆者的にはセグメンテーションがうまくいった から満足だったの可能性がある． • ノイズ除去は表現学習に必要ないと主張していた． • デノイジングにより低周波から高周波の情報を扱えることが拡散モデルの強みだと考えていた． • 表現学習においては必ずしもそうではないのか． 15

参考文献 [1] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and HsuanTien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html. [2] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollar, and Ross B. Girshick. Masked ´ autoencoders are scalable vision learners. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 15979–15988. IEEE, 2022. doi: 10.1109/CVPR52688.2022.01553. URL https://doi.org/10.1109/ CVPR52688.2022.01553. [3] Chen Wei, Karttikeya Mangalam, Po-Yao Huang, Yanghao Li, Haoqi Fan, Hu Xu, Huiyu Wang, Cihang Xie, Alan Yuille, and Christoph Feichtenhofer. Diffusion models as masked autoencoder. arXiv preprint arXiv:2304.03283, 2023. 16

背景 Denoising Diffusion Probalistic Model (DDPM) • データx0に徐々にガウシアンノイズを加えていき，完全なガウシアンノイズ xT~𝑁(𝑥𝑇; 0, 𝐼)にすることを考える（拡散過程）． • 生成は，拡散過程を逆向きに辿る（逆過程）．逆過程はNNでモデル化する． 17