NeurIPS 2024 | 如何防御對抗性提示攻擊？AdvUnlearn讓圖片生成風險驟降
本文第一作者為密歇根州立大學計算機系博士生張益萌，為 OPTML 實驗室成員，指導教師為劉思佳助理教授。OPtimization and Trustworthy Machine Learning (OPTML) 實驗室的研究興趣涵蓋機器學習 / 深度學習、優(yōu)化、計算機視覺、安全、信號處理和數(shù)據(jù)科學領域，重點是開發(fā)學習算法和理論，以及魯棒且可解釋的人工智能。
擴散模型（Diffusion Models, DMs）已經(jīng)成為文本到圖像生成領域的核心技術之一。憑借其卓越的性能，這些模型可以生成高質(zhì)量的圖像，廣泛應用于各類創(chuàng)作場景，如藝術設計、廣告生成等。然而，隨著擴散模型的日益普及，其帶來的安全問題也逐漸顯現(xiàn)。模型在處理開放式互聯(lián)網(wǎng)數(shù)據(jù)時，可能會在生成過程中輸出有害的、不適當?shù)膬?nèi)容，例如裸露、暴力、侵犯版權的圖像等，這為其實際應用帶來了道德和法律上的挑戰(zhàn)。
為應對這些問題，研究者提出了機器遺忘（Machine Unlearning）技術，也稱為概念擦除（Concept Erasing）[1-3]。通過這一技術，我們能夠有選擇性地 “抹去” 擴散模型中的某些不應生成的概念。然而，盡管這一技術能夠在非對抗性環(huán)境中取得一定效果，現(xiàn)有的概念擦除方法在面對對抗性提示攻擊（Adversarial Prompt Attacks）時仍然存在明顯的脆弱性。攻擊者可以通過對輸入提示的微小修改，誘導模型生成本應被擦除的內(nèi)容，這使得現(xiàn)有的概念擦除技術不夠魯棒。
對抗性提示攻擊：機器遺忘中的隱患
對抗性提示攻擊的關鍵在于通過改變文本輸入提示的細微細節(jié)，誘使模型生成不當?shù)膬?nèi)容。這類攻擊可能會通過調(diào)整拼寫、加入特定符號、或是改變句子結構來規(guī)避模型的防護。例如，即便模型已經(jīng)接受了擦除 “裸露” 相關內(nèi)容的訓練，但通過輸入一些巧妙設計的提示詞，攻擊者依然可能讓模型重新生成被擦除的圖像。這種攻擊的成功率顯示了現(xiàn)有擦除技術在對抗惡意輸入時的薄弱性。
為了解決這一問題，我們提出了 AdvUnlearn 框架。這是一個將對抗性訓練（Adversarial Training, AT）與概念擦除相結合的框架，旨在提升擴散模型在概念擦除任務中的魯棒性，防止模型在面對對抗性攻擊時再次生成被遺忘的內(nèi)容。
- 論文題目：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 框架的獨特之處在于，它系統(tǒng)性地結合了對抗性訓練與概念擦除方法，從而提升模型在對抗惡意輸入時的安全性和魯棒性。傳統(tǒng)的對抗性訓練主要應用于圖像分類任務，而 AdvUnlearn 首次將其引入到擴散模型的概念擦除中。通過雙層優(yōu)化策略，AdvUnlearn 能夠在增強模型魯棒性的同時保持其圖像生成質(zhì)量。
雙層優(yōu)化（Bi-level Optimization）：AdvUnlearn 框架的優(yōu)化分為兩個主要層次：
- 下層優(yōu)化（Lower-level Optimization）：這一層的任務是生成對抗性提示，即通過對輸入文本提示的細微擾動，生成能夠誘使模型生成本應被擦除的內(nèi)容的提示。這個過程的目標是找到這些對抗性提示，使得模型在面對這些提示時生成不合規(guī)內(nèi)容。
- 上層優(yōu)化（Upper-level Optimization）：這一層的任務是通過最小化模型對目標概念的響應來強化模型的魯棒性，同時確保模型生成質(zhì)量不受影響。我們通過更新模型參數(shù) θ，優(yōu)化概念擦除的效果，同時保留模型的正常生成能力。
其中
為目標遺忘內(nèi)容，
為遺忘目標函數(shù)，
為生成攻擊的目標函數(shù)。
保留效用的正則化（Utility-Retaining Regularization）
一個重大挑戰(zhàn)是，直接引入對抗性訓練可能會導致圖像生成質(zhì)量下降。這是因為模型在面對對抗性提示時，往往會過于專注于抵御不當輸入，而忽視了對其他正常內(nèi)容的生成能力。為此，我們提出了保留效用的正則化方法。
該方法的核心思想是引入一個 “保留集” （
），這是一組與目標概念無關的文本提示。這些提示用于確保模型在訓練過程中保留其生成正常內(nèi)容的能力。例如，在處理 “裸露” 相關概念時，保留集中的提示可能涉及描述風景、動物、建筑等非敏感內(nèi)容。
保留集的選擇尤為重要。我們用 COCO 和 ImageNet 等數(shù)據(jù)集中的非敏感的物體類別制作成文本提示作為保留集，并通過大模型進行篩選，確保保留集中的提示與待擦除概念無關。這些提示集通過以下正則化項加入優(yōu)化目標：
其中
為未經(jīng)過機器遺忘的初始模型參數(shù)。通過該機制，AdvUnlearn 框架能夠有效解決對抗性訓練帶來的質(zhì)量下降問題，確保生成圖像的多樣性和精度。
最佳優(yōu)化模塊的選擇 （Modularity Exploriation）: 文本編碼器 vs. UNet
擴散模型通常由多個核心模塊組成，其中文本編碼器（Text Encoder）和 UNet 是兩個主要組件。UNet 通常用于處理圖像的逐步去噪過程，而文本編碼器則將文本提示轉換為適合生成圖像的表示。
在傳統(tǒng)的概念擦除方法中，UNet 被視為模型生成圖像的核心，因此大多數(shù)現(xiàn)有技術都集中在 UNet 上進行優(yōu)化。然而，在我們的研究中，我們發(fā)現(xiàn)文本編碼器在應對對抗性提示攻擊時表現(xiàn)更佳。文本編碼器在處理輸入提示的過程中，負責將文本信息轉化為圖像生成的關鍵特征。因此，對文本編碼器進行優(yōu)化，可以更好地控制模型對輸入中細微變化的反應，從而提升模型在面對惡意輸入時的魯棒性。
具體來說，我們發(fā)現(xiàn)，優(yōu)化文本編碼器能夠顯著降低對抗性攻擊的成功率，同時保留生成高質(zhì)量圖像的能力。此外，經(jīng)過對抗性訓練優(yōu)化的文本編碼器可以作為一個 “即插即用” 的模塊，適用于不同類型的擴散模型，從而提高了框架的擴展性和通用性。
實證結果：AdvUnlearn 在多個任務中的表現(xiàn)
為了驗證 AdvUnlearn 框架的有效性，我們進行了廣泛的實驗證明。實驗涵蓋了多個概念擦除任務，包括裸露概念擦除、藝術風格擦除和對象擦除。UnlearnDiffAtk [4] 作為默認的對抗性提示攻擊來檢查遺忘模型的魯棒性。
裸露概念擦除
在處理敏感的 “裸露” 概念時，AdvUnlearn 展現(xiàn)了卓越的性能。相比傳統(tǒng)的概念擦除方法，AdvUnlearn 顯著降低了對抗性提示攻擊的成功率（ASR）。例如，在常見的對抗性提示攻擊場景下，ASR 從原始模型的 100% 下降至 21.13%，表明模型能夠有效抵御惡意輸入。同時，通過引入保留集，模型在生成無關內(nèi)容（如風景或物體圖像）時保持了高質(zhì)量，F(xiàn)ID 僅比標準模型略微上升。
藝術風格擦除
我們還對 “梵高風格” 的擦除進行了實驗。在這一任務中，AdvUnlearn 同樣表現(xiàn)出了優(yōu)異的性能。經(jīng)過對抗性訓練后的模型，能夠有效抵御對抗性提示攻擊，成功率降至 2%。相比之下，未經(jīng)優(yōu)化的模型在攻擊下幾乎無法抵御，生成的圖像仍保留了梵高的風格。此外，AdvUnlearn 在這一任務中成功保持了高質(zhì)量的圖像生成能力。
對象概念擦除
在對象擦除任務中，我們對 “教堂” 這一具體對象進行了擦除實驗。AdvUnlearn 不僅能夠有效抹去生成 “教堂” 圖像的能力，還能在對抗性提示攻擊下維持圖像生成的質(zhì)量。與其他擦除方法相比，AdvUnlearn 框架能夠在保證生成質(zhì)量的前提下，顯著提升魯棒性，使得模型在面對惡意輸入時更加可靠。
總結與展望
AdvUnlearn 框架是首個結合對抗性訓練和概念擦除技術的系統(tǒng)性方法，旨在提升擴散模型在概念擦除任務中的魯棒性。通過對文本編碼器進行優(yōu)化，以及引入保留效用的正則化機制，AdvUnlearn 成功平衡了魯棒性和模型效用。我們的實驗結果顯示，AdvUnlearn 能夠在多種場景下有效抵御對抗性提示攻擊，并保持高質(zhì)量的圖像生成。
隨著生成式 AI 的應用范圍不斷擴大，如何確保模型的安全性和可靠性成為一個關鍵問題。AdvUnlearn 框架為解決這一問題提供了新的路徑，并為未來的生成模型技術提供了重要的參考。如果你想深入了解 AdvUnlearn 框架的技術細節(jié)或實驗結果，歡迎訪問 GitHub 項目頁面（https://github.com/OPTML-Group/AdvUnlearn）。
