乾貨!基於GAN的稀有樣本生成

aitime論道 發佈 2022-07-05T04:58:17.205677+00:00

例如,在安全領域,系統設計人員希望獲得具有攻擊功能的數據包來修補漏洞;論文題目:RareGAN: Generating Samples for Rare Classes。

在實際應用中,我們經常需要稀有類別的樣本。例如,在安全領域,系統設計人員希望獲得具有攻擊功能的數據包來修補漏洞;在系統和網絡領域,維護人員希望獲得能使系統性能下降的數據包來測試系統;在機器學習領域,稀有類別的圖像可以用於調試機器學習算法。


在本工作中,我們提出了RareGAN,一個基於生成對抗網絡(Generative Adversarial Networks, GANs)的稀有樣本生成算法。


實驗表明,RareGAN具有很好的通用性:它能被廣泛地用於上述安全、系統、網絡、機器學習的應用中。


同時,RareGAN具有很好的效果:相比之前最先進的針對DNS放大攻擊的漏洞挖掘算法,RareGAN在花費更少的DNS請求的情況下能達到更好的攻擊樣本覆蓋度。


本期AI TIME PhD直播間,我們邀請到美國卡內基梅隆大學博士生——林梓楠,為我們帶來報告分享《基於GAN的稀有樣本生成》。


林梓楠:美國卡內基梅隆大學博士生,本科畢業於清華大學(連續三年獲得國家獎學金)。主要研究方向為生成模型的基礎理論,及生成模型在網絡、系統、安全、隱私上的應用。相關工作發表於NeurIPS(Spotlight獎),ICML,AISTATS,AAAI,IMC(最佳論文候選),SIGCOMM等。主頁:http://zinanlin.me/


1

動機


近幾年,生成模型在學術界是一個很火的領域,同時在工業界也得到了很多應用。



在生成模型領域,被廣泛研究的一個方向是無條件生成:比如訓練一個模型來生成很多的動物圖片。另一個方向是有條件生成:不僅能生成動物圖片,還能控制生成的圖片是何種動物。


我們研究了一個不一樣的問題。第一,我們只想生成稀有類別的樣本。這裡稀有類別指的是比例特別小的類別,比如說熊貓的圖片;第二,初始的數據集是沒有類別標記的。第三,我們在訓練中可以請求某些樣本的類別標籤,不過請求的數量存在上限。



這個問題在系統、網絡、安全領域有很多的應用。


例子1:DNS放大攻擊



DNS是一個我們日常常用的網絡協議。在DNS放大攻擊中,黑客把給DNS伺服器發送的請求中自己的IP位址替換成了其想要攻擊的IP位址。


DNS伺服器返回的數據包通常比DNS請求中的數據包大很多,所以攻擊者用很少的網絡帶寬就能耗盡被攻擊者的網絡帶寬。因為攻擊者利用DNS伺服器放大了請求中的數據包大小,因此這個網絡攻擊被稱作DNS放大攻擊。



作為DNS服務的運營商,我們希望知道哪些DNS請求數據包具有高放大率以便提前阻止DNS伺服器被黑客利用。我們可以將該問題抽象為我們之前提到的問題。


1.我們希望生成高放大率的DNS請求數據包。這些數據包在所有合法的DNS請求中非常稀疏,因此可以被看做是稀有類。


2.初始情況下,這些樣本是沒有類別標記的,因為我們不知道哪些數據包具有高放大率。


3.我們可以在訓練過程中將數據包發送到DNS伺服器得到它們是否為高放大率的標籤。但這個過程需要一定的時間。為了保證系統的效率,我們希望限制請求的標籤的數量。


例子2:性能壓力測試



在系統和網絡領域中,性能壓力測試是一個很常見的任務。給定一個系統,性能壓力的目標是獲得能使得系統資源使用率(如CPU、內存占用率)過高的數據包,並測試系統在這些數據包下的性能。獲得能使得系統資源使用率過高的數據包的這個問題也可以被抽象為我們之前提到的問題。


1.資源使用率過高的數據包在所有數據包中非常稀疏,因此可以被視為稀有類。


2.初始情況下,這些樣本是沒有類別標記的,因為我們不知道哪些數據包具有高資源使用率。


3. 訓練過程中我們可以測試數據包的資源占用以獲得類別標籤。由於測試過程可能非常耗時,我們需要對請求的標籤數量設置一個上限。


之前工作中的問題


針對這些問題,網絡、系統、安全領域的研究人員提出過許多的解決方案。然而大多數解決方案都需要設備、伺服器的原始碼。在現實場景中,由於設備的生產者和使用者常常來自於不同的機構,設備的使用者通常沒有接觸原始碼的條件。因此,我們亟需一個不依賴系統原始碼的解決方案。



這幾年,生成模型尤其是GAN取得了特別大的進展。我們的想法是:我們能否訓練一個GAN來生成這些稀有類別的數據包?


2

方法


難點


我們發現,簡單地使用GAN或者Conditional GAN並不能達到很好的效果。我們以DNS放大攻擊為例來解釋。


GAN


GAN包含生成器和鑑別器。下圖的下半部分展示了GAN的基本結構。



一個簡單的利用GAN的方法是,我們從未標記的DNS請求數據包中隨機挑選B個,其中B代表請求標籤數量的上限。然後我們把這B個數據包發送給DNS伺服器,獲得他們的放大率。然後我們只保留高放大率的數據包(圖中的紅色圓圈),並用它們來訓練一個GAN。


理想情況下,這個GAN能學習到高放大率數據包的特徵,並且生成更多的高放大率的數據包。然而,由於高放大數據包非常稀有,通過這個過程我們只能得到很少的訓練樣本。這會導致 GAN 過擬合。


Conditional GAN


我們注意到上面方法中被丟棄的大量的低放大率的數據包(圖中灰色圓圈)實際上對學習過程很有幫助,因為高放大率和低放大率數據包的特徵是互補的。我們能否將低放大率的數據包也納入到訓練過程中呢?Conditional GAN就能做到這點。


Conditional GAN是一種特殊的GAN:不僅能夠生成新的樣本,還能夠控制樣本的類別。ACGAN是一種被廣泛使用的Conditional GAN。下圖中的下半部分展示了ACGAN的基本結構。相比原始的GAN,ACGAN多了一個分類器用於預測樣本的類別。



一個簡單的利用ACGAN的方法是,我們用上述過程中獲得的高放大率和低放大率的數據包一起訓練一個ACGAN。之後我們可以指定生成類別為稀有類,來讓ACGAN生成更多稀有類別的數據包。然而,該方法效果還是不好。原因在於稀有類樣本還是太少。


在本工作中我們提出了幾個方法來改進ACGAN中鑑別器分類器在少樣本情況下的效果。



方法一:利用未標記樣本提高鑑別器效果



我們注意到ACGAN中鑑別器的損失函數實際上並不需要標籤,於是我們改變了ACGAN的訓練方法,使用所有未標記的樣本來訓練鑑別器。這大大增加了鑑別器訓練樣本的數量。


方法二:利用主動學習提高分類器效果



為了提高ACGAN中分類器的效果,我們借鑑了主動學習的方法。在新方法中,我們將訓練過程分為幾輪。


第一輪中,我們只標註B0個樣本來訓練一個初始的ACGAN;之後將未標記的樣本輸入給ACGAN的分類器並保留分類器對結果最不確信的B1個樣本輸入到DNS伺服器中獲得它們真實的標籤,並將它們加入到訓練數據集中,以便分類器可以學習改進。我們將多次循環這個過程,直到花光所有B的標籤預算。


3

結果


圖片


儘管我們研究的問題源於網絡、系統、安全領域,我們發現,我們提出的方法再圖片數據集上也有很好的效果。



Red: generated image

Green: the nearest training image

Yellow: their overlap


上圖展示了MNIST數據及上的效果,我們將數字0視為稀有類,其中刪除了一些數字0使得其占比為1%。第一行來自真實數據集。第二行和第三行的結果證實了我們之前的猜測:由於稀有類占比過少,直接用GAN或者ACGAN無法取得特別好的效果。GAN發生了嚴重的過擬合;ACGAN生成的樣本質量很差。


上圖中最後一行是我們提出的RareGAN方法。我們可以看到RareGAN生成的樣本具有很好的多樣性和質量,而且RareGAN能從訓練集中推廣以生成訓練集中不存在的樣本。


下圖展示了CIFAR10數據集上的效果。



For each method:

• First row: generated image

• Second row: the nearest training image


我們同樣可以看到,相對於之前的方法而言,RareGAN生成的樣本同樣具有很好的多樣性和質量。


網絡、系統、安全中的應用


回到先前提到的性能壓力測試。下圖中X軸指的是生成數據包的質量,即生成數據包和目標的高資源占用率的數據包這兩個分布之間的吻合程度,越小越好;Y軸是樣本多樣性,越大越好。我們看到,RareGAN同時具有很好的多樣性和質量。



我們還對之前提到的DNS放大攻擊做了相同的實驗。同樣,RareGAN生成的數據包同時具有很好的多樣性和質量。



我們還與之前安全研究人員提出的專門為DNS放大攻擊設計的方法AmpMAP(發表於2021年security頂會 USENIX Security)進行了對比。



我們可以看到,我們提出的 RareGAN 可以達到更好的樣本質量和多樣性,同時花費更小的標籤預算。



4

總結


網絡、系統和安全領域中的許多實際問題都可以抽象為從稀有類中生成更多的樣本(比如數據包)。先前研究者在網絡、系統、安全領域中提出的方法通常需要系統的原始碼,但系統的維護人員很難拿到原始碼。


在我們的這項工作中,我們將這個問題抽象為一個新的生成模型的問題。我們提出了一種基於 GAN 的方法RareGAN來解決這個問題。RareGAN特別的通用,並且比現有方法具有更好的樣本質量和多樣性。我們相信 RareGAN 在網絡、系統和安全領域將會有很多的應用。


論文題目:

RareGAN: Generating Samples for Rare Classes


論文連結:

https://arxiv.org/abs/2203.10674

關鍵字: