NeurIPS 2024 | 如何防禦對抗性提示攻擊？AdvUnlearn讓圖片生成風險驟降

AIxiv專欄是機器之心發佈學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯繫報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文第一作者為密歇根州立大學計算機系博士生張益萌，為 OPTML 實驗室成員，指導教師為劉思佳助理教授。OPtimization and Trustworthy Machine Learning (OPTML) 實驗室的研究興趣涵蓋機器學習 / 深度學習、優化、計算機視覺、安全、信號處理和數據科學領域，重點是開發學習算法和理論，以及魯棒且可解釋的人工智能。

擴散模型（Diffusion Models, DMs）已經成為文本到圖像生成領域的核心技術之一。憑藉其卓越的性能，這些模型可以生成高質量的圖像，廣泛應用於各類創作場景，如藝術設計、廣告生成等。然而，隨著擴散模型的日益普及，其帶來的安全問題也逐漸顯現。模型在處理開放式互聯網數據時，可能會在生成過程中輸出有害的、不適當的內容，例如裸露、暴力、侵犯版權的圖像等，這為其實際應用帶來了道德和法律上的挑戰。

為應對這些問題，研究者提出了機器遺忘（Machine Unlearning）技術，也稱為概念擦除（Concept Erasing）[1-3]。通過這一技術，我們能夠有選擇性地 「抹去」 擴散模型中的某些不應生成的概念。然而，儘管這一技術能夠在非對抗性環境中取得一定效果，現有的概念擦除方法在面對對抗性提示攻擊（Adversarial Prompt Attacks）時仍然存在明顯的脆弱性。攻擊者可以通過對輸入提示的微小修改，誘導模型生成本應被擦除的內容，這使得現有的概念擦除技術不夠魯棒。

對抗性提示攻擊：機器遺忘中的隱患

對抗性提示攻擊的關鍵在於通過改變文本輸入提示的細微細節，誘使模型生成不當的內容。這類攻擊可能會通過調整拚寫、加入特定符號、或是改變句子結構來規避模型的防護。例如，即便模型已經接受了擦除 「裸露」 相關內容的訓練，但通過輸入一些巧妙設計的提示詞，攻擊者依然可能讓模型重新生成被擦除的圖像。這種攻擊的成功率顯示了現有擦除技術在對抗惡意輸入時的薄弱性。

為瞭解決這一問題，我們提出了 AdvUnlearn 框架。這是一個將對抗性訓練（Adversarial Training, AT）與概念擦除相結合的框架，旨在提升擴散模型在概念擦除任務中的魯棒性，防止模型在面對對抗性攻擊時再次生成被遺忘的內容。

• 論文題目：Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models

• 論文地址：https://arxiv.org/abs/2405.15234  

• 代碼地址：https://github.com/OPTML-Group/AdvUnlearn  

•  Unlearned Diffusion Model Benchmark:  https://huggingface.co/spaces/Intel/UnlearnDiffAtk-Benchmark 

AdvUnlearn 框架：對抗性訓練與概念擦除的融合

AdvUnlearn 框架的獨特之處在於，它系統性地結合了對抗性訓練與概念擦除方法，從而提升模型在對抗惡意輸入時的安全性和魯棒性。傳統的對抗性訓練主要應用於圖像分類任務，而 AdvUnlearn 首次將其引入到擴散模型的概念擦除中。通過雙層優化策略，AdvUnlearn 能夠在增強模型魯棒性的同時保持其圖像生成質量。

雙層優化（Bi-level Optimization）：AdvUnlearn 框架的優化分為兩個主要層次：

下層優化（Lower-level Optimization）：這一層的任務是生成對抗性提示，即通過對輸入文本提示的細微擾動，生成能夠誘使模型生成本應被擦除的內容的提示。這個過程的目標是找到這些對抗性提示，使得模型在面對這些提示時生成不合規內容。

上層優化（Upper-level Optimization）：這一層的任務是通過最小化模型對目標概念的響應來強化模型的魯棒性，同時確保模型生成質量不受影響。我們通過更新模型參數 θ，優化概念擦除的效果，同時保留模型的正常生成能力。

其中

為生成攻擊的目標函數。

為遺忘目標函數，

為目標遺忘內容，

保留效用的正則化（Utility-Retaining Regularization）

一個重大挑戰是，直接引入對抗性訓練可能會導致圖像生成質量下降。這是因為模型在面對對抗性提示時，往往會過於專注於抵禦不當輸入，而忽視了對其他正常內容的生成能力。為此，我們提出了保留效用的正則化方法。

該方法的核心思想是引入一個 「保留集」 （

），這是一組與目標概念無關的文本提示。這些提示用於確保模型在訓練過程中保留其生成正常內容的能力。例如，在處理 「裸露」 相關概念時，保留集中的提示可能涉及描述風景、動物、建築等非敏感內容。

保留集的選擇尤為重要。我們用 COCO 和 ImageNet 等數據集中的非敏感的物體類別製作成文本提示作為保留集，並通過大模型進行篩選，確保保留集中的提示與待擦除概念無關。這些提示集通過以下正則化項加入優化目標：

其中為未經過機器遺忘的初始模型參數。通過該機制，AdvUnlearn 框架能夠有效解決對抗性訓練帶來的質量下降問題，確保生成圖像的多樣性和精度。

最佳優化模塊的選擇 （Modularity Exploriation）:  文本編碼器 vs. UNet

擴散模型通常由多個核心模塊組成，其中文本編碼器（Text Encoder）和 UNet 是兩個主要組件。UNet 通常用於處理圖像的逐步去噪過程，而文本編碼器則將文本提示轉換為適合生成圖像的表示。

在傳統的概念擦除方法中，UNet 被視為模型生成圖像的核心，因此大多數現有技術都集中在 UNet 上進行優化。然而，在我們的研究中，我們發現文本編碼器在應對對抗性提示攻擊時表現更佳。文本編碼器在處理輸入提示的過程中，負責將文本信息轉化為圖像生成的關鍵特徵。因此，對文本編碼器進行優化，可以更好地控制模型對輸入中細微變化的反應，從而提升模型在面對惡意輸入時的魯棒性。

具體來說，我們發現，優化文本編碼器能夠顯著降低對抗性攻擊的成功率，同時保留生成高質量圖像的能力。此外，經過對抗性訓練優化的文本編碼器可以作為一個 「即插即用」 的模塊，適用於不同類型的擴散模型，從而提高了框架的擴展性和通用性。

實證結果：AdvUnlearn 在多個任務中的表現

為了驗證 AdvUnlearn 框架的有效性，我們進行了廣泛的實驗證明。實驗涵蓋了多個概念擦除任務，包括裸露概念擦除、藝術風格擦除和對象擦除。UnlearnDiffAtk [4] 作為預設的對抗性提示攻擊來檢查遺忘模型的魯棒性。

裸露概念擦除

在處理敏感的 「裸露」 概念時，AdvUnlearn 展現了卓越的性能。相比傳統的概念擦除方法，AdvUnlearn 顯著降低了對抗性提示攻擊的成功率（ASR）。例如，在常見的對抗性提示攻擊場景下，ASR 從原始模型的 100% 下降至 21.13%，表明模型能夠有效抵禦惡意輸入。同時，通過引入保留集，模型在生成無關內容（如風景或物體圖像）時保持了高質量，FID 僅比標準模型略微上升。

藝術風格擦除

我們還對 「梵高風格」 的擦除進行了實驗。在這一任務中，AdvUnlearn 同樣表現出了優異的性能。經過對抗性訓練後的模型，能夠有效抵禦對抗性提示攻擊，成功率降至 2%。相比之下，未經優化的模型在攻擊下幾乎無法抵禦，生成的圖像仍保留了梵高的風格。此外，AdvUnlearn 在這一任務中成功保持了高質量的圖像生成能力。

對象概念擦除

在對象擦除任務中，我們對 「教堂」 這一具體對象進行了擦除實驗。AdvUnlearn 不僅能夠有效抹去生成 「教堂」 圖像的能力，還能在對抗性提示攻擊下維持圖像生成的質量。與其他擦除方法相比，AdvUnlearn 框架能夠在保證生成質量的前提下，顯著提升魯棒性，使得模型在面對惡意輸入時更加可靠。

總結與展望

AdvUnlearn 框架是首個結合對抗性訓練和概念擦除技術的系統性方法，旨在提升擴散模型在概念擦除任務中的魯棒性。通過對文本編碼器進行優化，以及引入保留效用的正則化機制，AdvUnlearn 成功平衡了魯棒性和模型效用。我們的實驗結果顯示，AdvUnlearn 能夠在多種場景下有效抵禦對抗性提示攻擊，並保持高質量的圖像生成。

隨著生成式 AI 的應用範圍不斷擴大，如何確保模型的安全性和可靠性成為一個關鍵問題。AdvUnlearn 框架為解決這一問題提供了新的路徑，並為未來的生成模型技術提供了重要的參考。如果你想深入瞭解 AdvUnlearn 框架的技術細節或實驗結果，歡迎訪問 GitHub 項目頁面（https://github.com/OPTML-Group/AdvUnlearn）。

[1] Gandikota R, Materzynska J, Fiotto-Kaufman J, et al. Erasing concepts from diffusion models

[2] Zhang Y, Fan C, Zhang Y, Yao Y, et al. Unlearncanvas: A stylized image dataset to benchmark machine unlearning for diffusion models, NeurIPS 2024.

[3] Fan C, Liu J, Zhang Y, et al. Salun: Empowering machine unlearning via gradient-based weight saliency in both image classification and generation, ICLR 2024.

[4] Zhang Y, Jia J, Chen X, et al. To generate or not? safety-driven unlearned diffusion models are still easy to generate unsafe images… for now, ECCV 2024.