Jak Sophos szkoli potężny, lekki wykrywacz złośliwego oprogramowania PDF na ultra skalę za pomocą Amazon SageMaker

Ten post jest współautorami Salmy Taoufiq i Harini Kannan z Sophos.

Jako lider cyberbezpieczeństwa nowej generacji, Sophos stara się chronić ponad 500,000 150 organizacji i miliony klientów w ponad XNUMX krajach przed ewoluującymi zagrożeniami. Oparty na analizie zagrożeń, uczeniu maszynowym (ML) i sztucznej inteligencji Sophos X-Ops, Sophos dostarcza szeroką i zróżnicowaną ofertę zaawansowanych produktów i usług do zabezpieczania i ochrony użytkowników, sieci i punktów końcowych przed phishingiem, oprogramowaniem ransomware, złośliwym oprogramowaniem i szeroki zakres cyberataków.

Połączenia Grupa Sophos Artificial Intelligence (AI) (SophosAI) nadzoruje rozwój i utrzymanie głównej technologii bezpieczeństwa Sophos ML. Bezpieczeństwo to problem dużych zbiorów danych. Aby uniknąć wykrycia, cyberprzestępcy nieustannie tworzą nowe ataki. Przekłada się to na kolosalne zbiory danych o zagrożeniach, z którymi grupa musi współpracować, aby jak najlepiej chronić klientów. Jednym z godnych uwagi przykładów jest wykrywanie i eliminowanie plików, które sprytnie zawierały złośliwe oprogramowanie, gdzie zbiory danych są w terabajtach.

W tym poście skupiamy się na systemie wykrywania złośliwego oprogramowania firmy Sophos w szczególności dla formatu plików PDF. Pokazujemy, jak używa SophosAI Amazon Sage Maker rozproszone szkolenie z terabajtami danych w celu trenowania wydajnego, lekkiego modelu XGBoost (Extreme Gradient Boosting). Dzięki temu ich zespół może szybciej iterować po dużych danych treningowych dzięki automatycznemu dostrajaniu hiperparametrów i bez zarządzania podstawową infrastrukturą treningową.

Rozwiązanie jest obecnie bezproblemowo zintegrowane z produkcyjnym potoku szkoleniowym, a model jest wdrażany na milionach punktów końcowych użytkowników za pośrednictwem Usługa punktów końcowych Sophos.

Kontekst przypadku użycia

Niezależnie od tego, czy chcesz podzielić się ważną umową, czy zachować fantazyjny projekt swojego CV, najczęściej wybierany jest format PDF. Jego powszechne stosowanie i powszechne przekonanie, że takie dokumenty są hermetyczne i statyczne, uśpiły użytkowników w fałszywym poczuciu bezpieczeństwa. PDF stał się zatem preferowanym wektorem infekcji w arsenale napastników. Złośliwe działania przy użyciu plików PDF są najczęściej osiągane poprzez osadzanie ładunku JavaScript uruchamianego przez czytnik PDF w celu pobrania wirusa z identyfikatora URI, sabotowania komputera użytkownika lub kradzieży poufnych informacji.

Sophos wykrywa złośliwe pliki PDF w różnych punktach ataku za pomocą zestawu modeli deterministycznych i ML. Jedno z takich podejść ilustruje poniższy diagram, na którym złośliwy plik PDF jest dostarczany za pośrednictwem poczty e-mail. Gdy tylko zostanie podjęta próba pobrania, uruchamia złośliwy skrypt wykonywalny, aby połączyć się z serwerem dowodzenia i kontroli atakującego. Detektor PDF SophosAI blokuje próbę pobrania po wykryciu, że jest ona złośliwa.

Inne sposoby obejmują blokowanie plików PDF w punkcie końcowym, wysyłanie złośliwych plików do piaskownicy (gdzie są oceniane przy użyciu wielu modeli), przesyłanie złośliwego pliku do infrastruktury oceniania i generowanie raportu bezpieczeństwa i tak dalej.

Motywacja

Zespół SophosAI uznał, że algorytm XGBoost jest idealnym kandydatem do tego zadania, aby zbudować detektor oparty na drzewie, który może skazywać złośliwe pliki PDF z dużą pewnością, jednocześnie pozwalając na niski pobór mocy obliczeniowej punktu końcowego i szybkie reakcje wnioskowania. Takie ścieżki badawcze są ważne dla Sophos z dwóch powodów. Posiadanie potężnych, ale niewielkich modeli wdrożonych na poziomie punktów końcowych klientów ma duży wpływ na oceny produktów firmy przez analityków. Ponadto, co ważniejsze, zapewnia ogólnie lepsze wrażenia użytkownika.

Wyzwanie techniczne

Ponieważ celem było posiadanie modelu o mniejszym zużyciu pamięci niż istniejące wykrywacze złośliwego oprogramowania PDF (zarówno na dysku, jak i w pamięci), SophosAI przekształciło XGBoost, algorytm klasyfikacji o udowodnionym doświadczeniu w tworzeniu drastycznie mniejszych modeli niż sieci neuronowe, przy imponującym wydajność na danych tabelarycznych. Przed przystąpieniem do modelowania eksperymentów XGBoost ważnym czynnikiem był sam rozmiar zestawu danych. Rzeczywiście, podstawowy zestaw danych plików PDF Sophos jest w terabajtach.

Dlatego głównym wyzwaniem było trenowanie modelu za pomocą dużego zestawu danych bez konieczności próbkowania w dół. Ponieważ dla wykrywacza kluczowe jest nauczenie się wykrywania wszelkich ataków opartych na plikach PDF — nawet ataków typu „igła w stogu siana” i całkowicie nowatorskich, aby lepiej chronić klientów Sophos — niezwykle ważne jest korzystanie ze wszystkich dostępnych różnorodnych zestawów danych.

W przeciwieństwie do sieci neuronowych, w których można trenować w partiach, w przypadku XGBoost potrzebujemy całego zestawu danych treningowych w pamięci. Największy zestaw danych szkoleniowych dla tego projektu ma ponad 1 TB i nie ma możliwości trenowania na taką skalę bez wykorzystania metodologii rozproszonej struktury szkoleniowej.

Omówienie rozwiązania

SageMaker to w pełni zarządzana usługa ML oferująca różne narzędzia do budowania, szkolenia, optymalizacji i wdrażania modeli ML. The Wbudowane biblioteki algorytmów SageMaker składa się z 21 popularnych algorytmów ML, w tym XGBoost. (Aby uzyskać więcej informacji, zobacz Uprość uczenie maszynowe dzięki XGBoost i Amazon SageMaker.) Dzięki wbudowanemu algorytmowi XGBoost możesz skorzystać z open-source SageMaker XGBoost Pojemnik określając wersję platformy większą niż 1.0-1, która ma ulepszoną elastyczność, skalowalność, rozszerzalność i funkcję Managed Spot Training oraz obsługuje formaty wejściowe, takie jak Parquet, który jest formatem używanym dla zestawu danych PDF.

Głównym powodem, dla którego SophosAI wybrał SageMaker, jest możliwość czerpania korzyści z w pełni zarządzanego szkolenia rozproszonego na wielowęzłowych instancjach procesora poprzez proste określenie więcej niż jednej instancji. SageMaker automatycznie dzieli dane na węzły, agreguje wyniki w węzłach równorzędnych i generuje jeden model. Instancje mogą być instancjami Spot, co znacznie obniża koszty szkolenia. Z wbudowany algorytm dla XGBoost, możesz to zrobić bez dodatkowego niestandardowego skryptu. Rozproszone wersje XGBoost istnieją również jako open source, takie jak XGBoost-Promień i XGBoost4J-Spark, ale ich użycie wymaga budowania, zabezpieczania, dostrajania i samozarządzających się rozproszonych klastrów obliczeniowych, co stanowi znaczny wysiłek w uzupełnieniu do rozwoju naukowego.

Dodatkowo, Automatyczne dostrajanie modelu SageMaker, znany również jako dostrajanie hiperparametrów, znajduje najlepszą wersję modelu, uruchamiając wiele zadań szkoleniowych z określonymi zakresami hiperparametrów. Następnie wybiera wartości hiperparametrów, które skutkują modelem, który działa najlepiej, mierzonym przez metrykę dla danego zadania ML.

Poniższy schemat ilustruje architekturę rozwiązania.

Warto zauważyć, że gdy SophosAI rozpoczął eksperymenty XGBoost przed przejściem do SageMaker, podjęto próby użycia dużej pamięci Elastyczna chmura obliczeniowa Amazon (Amazon EC2) instancje (na przykład r5a.24xlarge i x1.32xlarge), aby trenować model na jak największej próbce danych. Jednak te próby trwały średnio ponad 10 godzin i zwykle kończyły się niepowodzeniem z powodu braku pamięci.

W przeciwieństwie do tego, korzystając z algorytmu SageMaker XGBoost i bezproblemowego mechanizmu rozproszonego szkolenia, SophosAI może trenować model wzmacniający na dużą skalę na kolosalnym zestawie danych szkoleniowych PDF w ciągu 20 minut. Zespół musiał tylko przechowywać dane na Usługa Amazon Simple Storage (Amazon S3) jako pliki Parquet o podobnej wielkości i wybierz typ wystąpienia EC2 i żądaną liczbę wystąpień, a SageMaker zarządzał podstawową infrastrukturą klastra obliczeniowego i rozproszonym treningiem między wieloma węzłami klastra. Pod maską, SageMaker dzieli dane na węzły, używając ShardedByS3Key, aby równomiernie rozdzielić obiekty plików między każdą instancją i używa implementacji XGBoost Protokół królika (niezawodny interfejs AllReduce i broadcast) do uruchamiania przetwarzania rozproszonego i komunikacji między węzłami podstawowymi i równorzędnymi. (Aby uzyskać więcej informacji na temat agregacji histogramu i rozgłaszania przez węzły, zobacz XGBoost: skalowalny system wzmacniania drzew.)

Poza szkoleniem jednego modelu, dzięki SageMaker, XGBoost strojenie hiperparametrów został również szybki i łatwy dzięki możliwości jednoczesnego przeprowadzania różnych eksperymentów w celu dostrojenia najlepszej kombinacji hiperparametrów. Dostrajalne hiperparametry obejmują zarówno hiperparametry specyficzne dla dawki przypominającej, jak i hiperparametry specyficzne dla funkcji obiektywnej. Dwie strategie wyszukiwania oferowane są: losowe lub bayesowskie. Strategia wyszukiwania bayesowskiego okazała się cenna, ponieważ pomaga znaleźć lepsze hiperparametry niż zwykłe wyszukiwanie losowe, w mniejszej liczbie iteracji eksperymentalnych.

Informacje o zbiorze danych

Modelowanie wykrywania złośliwego oprogramowania PDF SophosAI opiera się na różnych funkcjach, takich jak histogramy n-gramowe i funkcje entropii bajtów (Aby uzyskać więcej informacji, zobacz MEADE: w kierunku mechanizmu wykrywania złośliwych załączników do wiadomości e-mail). Metadane i funkcje wyodrębnione z zebranych plików PDF są przechowywane w rozproszonej hurtowni danych. Następnie obliczany jest zestaw danych obejmujący ponad 3,500 funkcji, następnie dzielony na podstawie czasu na zestawy treningowe i testowe oraz przechowywany w partiach jako pliki Parquet w Amazon S3, aby SageMaker był łatwo dostępny do zadań szkoleniowych.

Poniższa tabela zawiera informacje o danych treningowych i testowych.

Dataset	Liczba przykładów	Liczba pilników do parkietu	Całkowity rozmiar
Trening	70,391,634	5,500	~ 1010 GB
Testowanie	1,242,283	98	~ 18 GB

Rozmiary danych zostały obliczone według wzoru:

Rozmiar danych = N × (n_F + nie_L) × 4

Formuła ma następujące parametry:

• N to liczba próbek w zbiorze danych
• n_F to liczba funkcji, z n_F = 3585
• n_L to liczba podstawowych etykiet prawdy, z n_L = 1
• 4 to liczba bajtów potrzebnych dla typu danych obiektów: float32

Ponadto poniższe wykresy kołowe przedstawiają rozkład etykiet zarówno zestawów treningowych, jak i testowych, wywołując nierównowagę klas napotykaną w zadaniu wykrywania złośliwego oprogramowania PDF.

Rozkład przesuwa się ze zbioru uczącego do zbioru testowego Jeden miesiąc. W celu zasymulowania rzeczywistego scenariusza wdrożenia i uniknięcia tymczasowego podsłuchiwania stosuje się podział zbioru danych na podstawie czasu na szkolenie i testowanie. Ta strategia umożliwiła również SophosAI ocenę prawdziwych możliwości uogólniania modelu w obliczu niespotykanych wcześniej, na przykład, zupełnie nowych ataków na pliki PDF.

Eksperymenty i wyniki

Aby rozpocząć eksperymenty, zespół SophosAI przeszkolił podstawowy model XGBoost z domyślnymi parametrami. Następnie rozpoczęli dostrajanie hiperparametrów w programie SageMaker przy użyciu strategii Bayesa, która jest tak prosta, jak określenie hiperparametry do dostrojenia i żądany zakres wartości, metrykę oceny (w tym przypadku ROC (Receiver Operating Characteristic) AUC) oraz zestawy uczące i walidacyjne. W przypadku wykrywacza złośliwego oprogramowania PDF SophosAI nadał priorytet hiperparametrom, w tym liczbie rund wzmacniających (num_round), maksymalna głębokość drzewa (max_depth), wskaźnik uczenia się (eta) oraz współczynnik próbkowania kolumn podczas budowania drzew (colsample_bytree). Ostatecznie uzyskano najlepsze hiperparametry i wykorzystano je do trenowania modelu na pełnym zestawie danych, a następnie oceniono na zestawie testowym wstrzymania.

Poniższy wykres przedstawia metrykę celu (ROC AUC) w porównaniu do 15 zadań szkoleniowych uruchomionych w ramach zadania dostrajania. Najlepsze hiperparametry to te odpowiadające dziewiątej pracy szkoleniowej.

Na początku eksperymentów SophosAI na SageMakerze szczególnie ważnym pytaniem, na które należy odpowiedzieć, było: jakiego rodzaju instancje i ile z nich jest potrzebnych do trenowania XGBoost na dostępnych danych? Ma to kluczowe znaczenie, ponieważ użycie niewłaściwej liczby lub typu instancji może być stratą czasu i pieniędzy; szkolenie na pewno się nie powiedzie z powodu braku pamięci lub, jeśli używa się zbyt wielu zbyt dużych instancji, może to stać się niepotrzebnie kosztowne.

XGBoost to algorytm związany z pamięcią (w przeciwieństwie do algorytmu obliczeniowego). Tak więc wystąpienie obliczeniowe ogólnego przeznaczenia (na przykład M5) jest lepszym wyborem niż wystąpienie zoptymalizowane pod kątem obliczeń (na przykład C4). Aby podjąć świadomą decyzję, istnieje prosta wskazówka SageMaker dotycząca wyboru liczby instancji wymaganych do uruchomienia szkolenia na pełnym zbiorze danych:

Całkowity rozmiar danych treningowych × współczynnik bezpieczeństwa^(*) < Liczba instancji × Całkowita pamięć typu instancji

W tym przypadku: całkowity rozmiar danych treningowych × współczynnik bezpieczeństwa (12) = 12120 GB

W poniższej tabeli podsumowano wymagania, gdy wybrany typ wystąpienia to ml.m5.24xlarge.

Rozmiar szkolenia × Współczynnik bezpieczeństwa (12)	Pamięć instancji ml.m5.24xlarge	Minimalna liczba wystąpień wymagana do szkolenia
12120 GB	384 GB	32

^*Ze względu na charakter uczenia rozproszonego XGBoost, który wymaga załadowania całego uczącego zestawu danych do obiektu DMatrix przed uczeniem oraz dodatkowej wolnej pamięci, zalecany jest współczynnik bezpieczeństwa 10–12.

Aby przyjrzeć się bliżej wykorzystaniu pamięci dla pełnego treningu SageMaker XGBoost na dostarczonym zbiorze danych, przedstawiamy odpowiedni wykres uzyskany z treningu Amazon Cloud Watch monitorowanie. Do tego zadania szkoleniowego użyto 40 instancji ml.m5.24xlarge, a maksymalne wykorzystanie pamięci osiągnęło około 62%.

Koszty inżynieryjne zaoszczędzone dzięki integracji zarządzanej usługi ML, takiej jak SageMaker, z potoku danych wynoszą około 50%. Możliwość korzystania z instancji Spot do zadań szkoleniowych i dostrajania hiperparametrów obniża koszty o dodatkowe 63%.

Wnioski

Dzięki SageMaker zespół SophosAI mógł z powodzeniem rozwiązać złożony projekt o wysokim priorytecie, budując lekki model XGBoost do wykrywania złośliwego oprogramowania do plików PDF, który jest znacznie mniejszy na dysku (do 25 razy mniejszy) i w pamięci (do 5 razy mniejszy) niż jego poprzednik detektora. Jest to mały, ale potężny wykrywacz złośliwego oprogramowania z ~0.99 AUC i prawdziwym pozytywnym wskaźnikiem 0.99 i fałszywym pozytywnym wskaźnikiem . Ten model można szybko przeszkolić, a jego wydajność można łatwo monitorować w czasie, ponieważ przeszkolenie go na więcej niż 20 TB danych zajmuje mniej niż 1 minut.

Możesz wykorzystać wbudowany algorytm SageMaker XGBoost do budowania modeli z danymi tabelarycznymi na dużą skalę. Dodatkowo możesz również wypróbować nowe wbudowane algorytmy Amazon SageMaker LightGBM, CatBoost, AutoGluon-Tabular i Tab Transformer, jak opisano w tym blog.

---

O autorach

Salma Taoufik jest starszym analitykiem danych w Sophos, pracującym na styku uczenia maszynowego i cyberbezpieczeństwa. Z wykształceniem licencjackim w dziedzinie informatyki, ukończyła Uniwersytet Środkowoeuropejski z tytułem mgr inż. w matematyce i jej zastosowaniach. Salma, nie rozwijając wykrywacza złośliwego oprogramowania, jest zapaloną turystką, podróżniczką i konsumentką thrillerów.

Hariniego Kannana jest analitykiem danych w SophosAI. Od około 4 lat zajmuje się nauką o bezpieczeństwie danych. Wcześniej pracowała jako Principal Data Scientist w Capsule8, która została przejęta przez Sophos. Wygłaszała prelekcje na CAMLIS, BlackHat (USA), Open Data Science Conference (Wschód), Data Science Salon, PyData (Boston) oraz Data Connectors. Jej obszary badań obejmują wykrywanie ataków sprzętowych za pomocą liczników wydajności, analizę zachowań użytkowników, interpretowalną ML i nienadzorowane wykrywanie anomalii.

Hasana Poonawala jest starszym architektem rozwiązań AI/ML w AWS z siedzibą w Londynie w Wielkiej Brytanii. Hasan pomaga klientom projektować i wdrażać aplikacje do uczenia maszynowego w środowisku produkcyjnym na AWS. Posiada ponad 12-letnie doświadczenie zawodowe jako data scientist, praktyk uczenia maszynowego i programista. W wolnym czasie Hasan uwielbia poznawać przyrodę i spędzać czas z przyjaciółmi i rodziną.

Diganta Patela jest kierownikiem ds. wsparcia korporacyjnego w AWS. Współpracuje z klientami, aby projektować, wdrażać i działać w chmurze na dużą skalę. Jego obszary zainteresowania to praktyki MLOps i DevOps oraz to, jak mogą one pomóc klientom w ich podróży do chmury. Poza pracą lubi fotografować, grać w siatkówkę i spędzać czas z przyjaciółmi i rodziną.