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

擴散模型（Diffusion Models, DMs）已經(jīng)成為文本到圖像生成領(lǐng)域的核心技術(shù)之一。憑借其卓越的性能，這些模型可以生成高質(zhì)量的圖像，廣泛應(yīng)用于各類創(chuàng)作場景，如藝術(shù)設(shè)計、廣告生成等。然而，隨著擴散模型的日益普及，其帶來的安全問題也逐漸顯現(xiàn)。模型在處理開放式互聯(lián)網(wǎng)數(shù)據(jù)時，可能會在生成過程中輸出有害的、不適當(dāng)?shù)膬?nèi)容，例如裸露、暴力、侵犯版權(quán)的圖像等，這為其實際應(yīng)用帶來了道德和法律上的挑戰(zhàn)。

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

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

對抗性提示攻擊的關(guān)鍵在于通過改變文本輸入提示的細微細節(jié)，誘使模型生成不當(dāng)?shù)膬?nèi)容。這類攻擊可能會通過調(diào)整拼寫、加入特定符號、或是改變句子結(jié)構(gòu)來規(guī)避模型的防護。例如，即便模型已經(jīng)接受了擦除 “裸露” 相關(guān)內(nèi)容的訓(xùn)練，但通過輸入一些巧妙設(shè)計的提示詞，攻擊者依然可能讓模型重新生成被擦除的圖像。這種攻擊的成功率顯示了現(xiàn)有擦除技術(shù)在對抗惡意輸入時的薄弱性。

為了解決這一問題，我們提出了 AdvUnlearn 框架。這是一個將對抗性訓(xùn)練（Adversarial Training, AT）與概念擦除相結(jié)合的框架，旨在提升擴散模型在概念擦除任務(wù)中的魯棒性，防止模型在面對對抗性攻擊時再次生成被遺忘的內(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 框架：對抗性訓(xùn)練與概念擦除的融合

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

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

1. 下層優(yōu)化（Lower-level Optimization）：這一層的任務(wù)是生成對抗性提示，即通過對輸入文本提示的細微擾動，生成能夠誘使模型生成本應(yīng)被擦除的內(nèi)容的提示。這個過程的目標是找到這些對抗性提示，使得模型在面對這些提示時生成不合規(guī)內(nèi)容。
2. 上層優(yōu)化（Upper-level Optimization）：這一層的任務(wù)是通過最小化模型對目標概念的響應(yīng)來強化模型的魯棒性，同時確保模型生成質(zhì)量不受影響。我們通過更新模型參數(shù) θ，優(yōu)化概念擦除的效果，同時保留模型的正常生成能力。

其中為目標遺忘內(nèi)容，為遺忘目標函數(shù)，為生成攻擊的目標函數(shù)。

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

一個重大挑戰(zhàn)是，直接引入對抗性訓(xùn)練可能會導(dǎo)致圖像生成質(zhì)量下降。這是因為模型在面對對抗性提示時，往往會過于專注于抵御不當(dāng)輸入，而忽視了對其他正常內(nèi)容的生成能力。為此，我們提出了保留效用的正則化方法。

該方法的核心思想是引入一個 “保留集” （），這是一組與目標概念無關(guān)的文本提示。這些提示用于確保模型在訓(xùn)練過程中保留其生成正常內(nèi)容的能力。例如，在處理 “裸露” 相關(guān)概念時，保留集中的提示可能涉及描述風(fēng)景、動物、建筑等非敏感內(nèi)容。

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

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

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

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

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

具體來說，我們發(fā)現(xiàn)，優(yōu)化文本編碼器能夠顯著降低對抗性攻擊的成功率，同時保留生成高質(zhì)量圖像的能力。此外，經(jīng)過對抗性訓(xùn)練優(yōu)化的文本編碼器可以作為一個 “即插即用” 的模塊，適用于不同類型的擴散模型，從而提高了框架的擴展性和通用性。

實證結(jié)果：AdvUnlearn 在多個任務(wù)中的表現(xiàn)

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

裸露概念擦除

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

藝術(shù)風(fēng)格擦除

我們還對 “梵高風(fēng)格” 的擦除進行了實驗。在這一任務(wù)中，AdvUnlearn 同樣表現(xiàn)出了優(yōu)異的性能。經(jīng)過對抗性訓(xùn)練后的模型，能夠有效抵御對抗性提示攻擊，成功率降至 2%。相比之下，未經(jīng)優(yōu)化的模型在攻擊下幾乎無法抵御，生成的圖像仍保留了梵高的風(fēng)格。此外，AdvUnlearn 在這一任務(wù)中成功保持了高質(zhì)量的圖像生成能力。

對象概念擦除

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

總結(jié)與展望

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

隨著生成式 AI 的應(yīng)用范圍不斷擴大，如何確保模型的安全性和可靠性成為一個關(guān)鍵問題。AdvUnlearn 框架為解決這一問題提供了新的路徑，并為未來的生成模型技術(shù)提供了重要的參考。如果你想深入了解 AdvUnlearn 框架的技術(shù)細節(jié)或?qū)嶒灲Y(jié)果，歡迎訪問 GitHub 項目頁面（https://github.com/OPTML-Group/AdvUnlearn）。