Wie Sophos mit Amazon SageMaker einen leistungsstarken, leichten PDF-Malware-Erkenner im Ultramaßstab trainiert

Dieser Beitrag wurde gemeinsam von Salma Taoufiq und Harini Kannan von Sophos verfasst.

Als führendes Unternehmen für Cybersicherheit der nächsten Generation Sophos ist bestrebt, mehr als 500,000 Organisationen und Millionen von Kunden in über 150 Ländern vor neuen Bedrohungen zu schützen. Basierend auf Threat Intelligence, maschinellem Lernen (ML) und künstlicher Intelligenz von Sophos X-Ops bietet Sophos ein breites und vielfältiges Portfolio an fortschrittlichen Produkten und Services, um Benutzer, Netzwerke und Endpunkte vor Phishing, Ransomware, Malware und mehr zu schützen und zu verteidigen die breite Palette von Cyberangriffen da draußen.

Das Sophos Gruppe für künstliche Intelligenz (KI). (SophosAI) überwacht die Entwicklung und Wartung der wichtigsten ML-Sicherheitstechnologie von Sophos. Sicherheit ist ein Big-Data-Problem. Um der Entdeckung zu entgehen, entwickeln Cyberkriminelle ständig neuartige Angriffe. Dies führt zu kolossalen Bedrohungsdatensätzen, mit denen die Gruppe arbeiten muss, um Kunden bestmöglich zu schützen. Ein bemerkenswertes Beispiel ist die Erkennung und Beseitigung von Dateien, die auf listige Weise mit Malware versehen wurden, wobei die Datensätze in Terabyte liegen.

In diesem Beitrag konzentrieren wir uns speziell auf das Malware-Erkennungssystem von Sophos für das PDF-Dateiformat. Wir zeigen, wie SophosAI verwendet Amazon Sage Maker verteiltes Training mit Terabytes an Daten, um ein leistungsstarkes, leichtes XGBoost-Modell (Extreme Gradient Boosting) zu trainieren. Dadurch kann ihr Team mit automatischer Hyperparameter-Abstimmung und ohne Verwaltung der zugrunde liegenden Trainingsinfrastruktur schneller über große Trainingsdaten iterieren.

Die Lösung ist derzeit nahtlos in die Produktionsschulungspipeline integriert und das Modell wird über das auf Millionen von Benutzerendpunkten bereitgestellt Sophos Endpoint-Dienst.

Anwendungsfallkontext

Egal, ob Sie einen wichtigen Vertrag teilen oder das ausgefallene Design Ihres Lebenslaufs bewahren möchten, das PDF-Format ist die häufigste Wahl. Seine weit verbreitete Verwendung und die allgemeine Wahrnehmung, dass solche Dokumente luftdicht und statisch sind, haben die Benutzer in ein falsches Sicherheitsgefühl gewiegt. PDF ist daher zu einem bevorzugten Infektionsvektor im Arsenal von Angreifern geworden. Böswillige Aktionen mit PDFs werden meistens durch das Einbetten einer JavaScript-Nutzlast erreicht, die vom PDF-Reader ausgeführt wird, um einen Virus von einer URI herunterzuladen, den Computer des Benutzers zu sabotieren oder vertrauliche Informationen zu stehlen.

Sophos erkennt schädliche PDF-Dateien an verschiedenen Angriffspunkten mithilfe eines Ensembles aus deterministischen und ML-Modellen. Ein solcher Ansatz wird im folgenden Diagramm veranschaulicht, in dem die schädliche PDF-Datei per E-Mail zugestellt wird. Sobald ein Download-Versuch unternommen wird, löst dies das bösartige ausführbare Skript aus, um eine Verbindung zum Command-and-Control-Server des Angreifers herzustellen. Der PDF-Detektor von SophosAI blockiert den Download-Versuch, nachdem er erkannt hat, dass er schädlich ist.

Andere Möglichkeiten umfassen das Blockieren der PDF-Dateien im Endpunkt, das Senden der schädlichen Dateien an eine Sandbox (wo sie mit mehreren Modellen bewertet werden), das Senden der schädlichen Datei an eine Scoring-Infrastruktur und das Generieren eines Sicherheitsberichts und so weiter.

Motivation

Um einen baumbasierten Detektor zu entwickeln, der bösartige PDFs mit hoher Zuverlässigkeit erkennen kann und gleichzeitig einen geringen Energieverbrauch der Endpoint-Rechenleistung und schnelle Inferenzreaktionen ermöglicht, fand das SophosAI-Team den XGBoost-Algorithmus als perfekten Kandidaten für diese Aufgabe. Solche Forschungswege sind für Sophos aus zwei Gründen wichtig. Die Bereitstellung leistungsstarker, aber kleiner Modelle auf der Ebene der Kundenendpunkte hat einen großen Einfluss auf die Produktbewertungen des Unternehmens durch Analysten. Es bietet auch, und was noch wichtiger ist, insgesamt eine bessere Benutzererfahrung.

Technische Herausforderung

Da das Ziel darin bestand, ein Modell mit einem kleineren Speicherbedarf als die bestehenden PDF-Malware-Detektoren (sowohl auf der Festplatte als auch im Speicher) zu haben, wandte sich SophosAI an XGBoost, einen Klassifizierungsalgorithmus, der nachweislich drastisch kleinere Modelle als neuronale Netze produziert und dabei beeindruckende Ergebnisse erzielt Leistung auf Tabellendaten. Bevor man sich an die Modellierung von XGBoost-Experimenten wagte, war die schiere Größe des Datensatzes eine wichtige Überlegung. Sophos Kerndatensatz von PDF-Dateien liegt in der Tat in Terabyte.

Daher bestand die größte Herausforderung darin, das Modell mit einem großen Datensatz ohne Downsampling zu trainieren. Da es für den Detektor entscheidend ist, alle PDF-basierten Angriffe zu erkennen – selbst Nadel-im-Heuhaufen und völlig neuartige, um Sophos-Kunden besser zu schützen – ist es von größter Bedeutung, alle verfügbaren unterschiedlichen Datensätze zu verwenden.

Im Gegensatz zu neuronalen Netzwerken, bei denen Sie in Stapeln trainieren können, benötigen wir für XGBoost den gesamten Trainingsdatensatz im Speicher. Der größte Trainingsdatensatz für dieses Projekt umfasst über 1 TB, und es gibt keine Möglichkeit, in einem solchen Umfang zu trainieren, ohne die Methoden eines verteilten Trainingsframeworks zu nutzen.

Lösungsüberblick

SageMaker ist ein vollständig verwalteter ML-Dienst, der verschiedene Tools zum Erstellen, Trainieren, Optimieren und Bereitstellen von ML-Modellen bereitstellt. Das SageMaker integrierte Bibliotheken von Algorithmen bestehen aus 21 beliebten ML-Algorithmen, einschließlich XGBoost. (Weitere Informationen finden Sie unter Vereinfachen Sie maschinelles Lernen mit XGBoost und Amazon SageMaker.) Mit dem integrierten XGBoost-Algorithmus können Sie die Vorteile von Open Source nutzen SageMaker XGBoost-Container durch Angabe einer Framework-Version größer als 1.0-1, die über verbesserte Flexibilität, Skalierbarkeit, Erweiterbarkeit und Managed Spot Training verfügt und Eingabeformate wie Parquet unterstützt, das das Format ist, das für den PDF-Datensatz verwendet wird.

Der Hauptgrund, warum sich SophosAI für SageMaker entschieden hat, ist die Möglichkeit, vom vollständig verwalteten verteilten Training auf CPU-Instanzen mit mehreren Knoten zu profitieren, indem einfach mehr als eine Instanz angegeben wird. SageMaker teilt die Daten automatisch auf Knoten auf, aggregiert die Ergebnisse über Peer-Knoten hinweg und generiert ein einziges Modell. Die Instanzen können Spot-Instanzen sein, wodurch die Schulungskosten erheblich reduziert werden. Mit dem eingebauter Algorithmus für XGBoost, können Sie dies ohne zusätzliches benutzerdefiniertes Skript tun. Verteilte Versionen von XGBoost existieren auch als Open Source, wie z XGBoost-Ray und XGBoost4J-Spark, aber ihre Verwendung erfordert das Erstellen, Sichern, Optimieren und Selbstverwalten verteilter Rechencluster, was einen erheblichen Aufwand zusätzlich zur wissenschaftlichen Entwicklung darstellt.

Zusätzlich Automatische Modellabstimmung von SageMaker, auch bekannt als Hyperparameter-Tuning, findet die beste Version eines Modells, indem viele Trainingsjobs mit von Ihnen angegebenen Bereichen von Hyperparametern ausgeführt werden. Es wählt dann die Hyperparameterwerte aus, die zu einem Modell führen, das die beste Leistung liefert, gemessen an einer Metrik für die gegebene ML-Aufgabe.

Das folgende Diagramm zeigt die Lösungsarchitektur.

Es ist erwähnenswert, dass, als SophosAI XGBoost-Experimente startete, bevor es sich an SageMaker wandte, Versuche unternommen wurden, großen Speicher zu verwenden Amazon Elastic Compute-Cloud (Amazon EC2)-Instances (z. B. r5a.24xlarge und x1.32xlarge), um das Modell mit einer möglichst großen Datenstichprobe zu trainieren. Diese Versuche dauerten jedoch im Durchschnitt mehr als 10 Stunden und scheiterten normalerweise, weil der Arbeitsspeicher knapp wurde.

Im Gegensatz dazu konnte SophosAI durch die Verwendung des SageMaker XGBoost-Algorithmus und eines problemlosen verteilten Trainingsmechanismus ein Booster-Modell im großen Maßstab auf dem riesigen PDF-Trainingsdatensatz in nur 20 Minuten trainieren. Das Team musste die Daten nur noch speichern Amazon Simple Storage-Service (Amazon S3) als Parquet-Dateien ähnlicher Größe und wählen Sie einen EC2-Instance-Typ und die gewünschte Anzahl von Instances aus, und SageMaker verwaltete die zugrunde liegende Compute-Cluster-Infrastruktur und verteilte das Training auf mehrere Knoten des Clusters. Unter der Haube teilt SageMaker die Daten mithilfe von ShardedByS3Key auf Knoten auf, um die Dateiobjekte gleichmäßig auf alle Instanzen zu verteilen, und verwendet die XGBoost-Implementierung von Kaninchen-Protokoll (zuverlässige AllReduce- und Broadcast-Schnittstelle), um eine verteilte Verarbeitung zu starten und zwischen primären und Peer-Knoten zu kommunizieren. (Weitere Einzelheiten zur Aggregation und Übertragung von Histogrammen über Knoten hinweg finden Sie unter XGBoost: Ein skalierbares Tree-Boosting-System.)

Mit SageMaker trainieren Sie nicht nur ein Modell, XGBoost-Hyperparameter-Tuning wurde auch durch die Möglichkeit, verschiedene Experimente gleichzeitig auszuführen, schnell und einfach gemacht, um die beste Kombination von Hyperparametern zu optimieren. Die abstimmbaren Hyperparameter umfassen sowohl Booster-spezifische als auch Zielfunktions-spezifische Hyperparameter. Zwei Suchstrategien angeboten: zufällig oder bayesianisch. Die Bayes'sche Suchstrategie hat sich als wertvoll erwiesen, da sie hilft, bessere Hyperparameter als eine bloße zufällige Suche in weniger experimentellen Iterationen zu finden.

Datensatzinformationen

Die PDF-Malware-Erkennungsmodellierung von SophosAI stützt sich auf eine Vielzahl von Funktionen wie N-Gramm-Histogramme und Byte-Entropie-Funktionen (Weitere Informationen finden Sie unter MEADE: Auf dem Weg zu einer Engine zur Erkennung bösartiger E-Mail-Anhänge). Aus gesammelten PDF-Dateien extrahierte Metadaten und Merkmale werden in einem verteilten Data Warehouse gespeichert. Anschließend wird ein Datensatz mit über 3,500 Merkmalen berechnet, nach Zeit weiter in Trainings- und Testsätze aufgeteilt und stapelweise als Parquet-Dateien in Amazon S3 gespeichert, damit SageMaker für Trainingsaufträge leicht darauf zugreifen kann.

Die folgende Tabelle enthält Informationen zu den Trainings- und Testdaten.

Datensatz	Anzahl von Beispielen	Anzahl Parkettdateien	Gesamtgröße
Ausbildung	70,391,634	5,500	~ 1010 GB
Test	1,242,283	98	~ 18 GB

Die Datengrößen wurden nach folgender Formel berechnet:

Datengröße = N × (n_F + n_L) × 4

Die Formel hat die folgenden Parameter:

• N ist die Anzahl der Proben im Datensatz
• n_F ist die Anzahl der Features, mit n_F = 3585
• n_L die Anzahl der Ground-Truth-Labels ist, wobei n ist_L = 1
• 4 ist die Anzahl der Bytes, die für den Datentyp der Features benötigt werden: float32

Darüber hinaus stellen die folgenden Tortendiagramme die Bezeichnungsverteilung sowohl der Trainings- als auch der Testsätze bereit, wodurch das Klassenungleichgewicht herausgearbeitet wird, mit dem die PDF-Malware-Erkennungsaufgabe konfrontiert ist.

Die Verteilung verschiebt sich vom Trainingsset zum Einmonatigen Testset. Eine zeitbasierte Aufteilung des Datensatzes in Training und Test wird angewendet, um das reale Einsatzszenario zu simulieren und zeitliches Schnüffeln zu vermeiden. Diese Strategie ermöglichte es SophosAI auch, die wahren Verallgemeinerungsfähigkeiten des Modells zu evaluieren, wenn es zum Beispiel mit zuvor unbekannten brandneuen PDF-Angriffen konfrontiert wurde.

Experimente und Ergebnisse

Um Experimente anzukurbeln, trainierte das SophosAI-Team ein XGBoost-Basismodell mit Standardparametern. Dann begannen sie mit der Hyperparameter-Feinabstimmung mit SageMaker unter Verwendung der Bayes'schen Strategie, die so einfach ist wie die Angabe von Hyperparameter abzustimmenden und gewünschten Wertebereich, die Bewertungsmetrik (in diesem Fall ROC (Receiver Operating Characteristic) AUC) und die Trainings- und Validierungssätze. Für den PDF-Malware-Detektor hat SophosAI Hyperparameter priorisiert, darunter die Anzahl der Boosting-Runden (num_round), die maximale Baumtiefe (max_depth), die Lernrate (eta) und das Stichprobenverhältnis der Spalten beim Erstellen von Bäumen (colsample_bytree). Schließlich wurden die besten Hyperparameter ermittelt und verwendet, um ein Modell mit dem vollständigen Datensatz zu trainieren und schließlich mit dem Holdout-Testsatz zu bewerten.

Das folgende Diagramm zeigt die objektive Metrik (ROC AUC) im Vergleich zu den 15 Trainingsjobs, die innerhalb des Optimierungsjobs ausgeführt werden. Die besten Hyperparameter sind diejenigen, die dem neunten Trainingsjob entsprechen.

Zu Beginn der Experimente von SophosAI mit SageMaker war eine besonders wichtige Frage zu beantworten: Welche Art von Instanzen und wie viele davon werden benötigt, um XGBoost mit den vorliegenden Daten zu trainieren? Dies ist von entscheidender Bedeutung, da die Verwendung der falschen Nummer oder des falschen Instanztyps Zeit- und Geldverschwendung sein kann. Das Training wird zwangsläufig aufgrund von Speichermangel fehlschlagen oder, wenn zu viele zu große Instanzen verwendet werden, kann dies unnötig teuer werden.

XGBoost ist ein speichergebundener (im Gegensatz zu einem rechengebundenen) Algorithmus. Daher ist eine Compute-Instance für allgemeine Zwecke (z. B. M5) eine bessere Wahl als eine rechenoptimierte Instance (z. B. C4). Um eine fundierte Entscheidung zu treffen, gibt es eine einfache SageMaker-Richtlinie zur Auswahl der Anzahl der Instanzen, die erforderlich sind, um das Training auf dem vollständigen Datensatz auszuführen:

Gesamtgröße der Trainingsdaten × Sicherheitsfaktor^(*) < Instanzanzahl × Gesamtspeicher des Instanztyps

In diesem Fall: Gesamtgröße der Trainingsdaten × Sicherheitsfaktor (12) = 12120 GB

Die folgende Tabelle fasst die Anforderungen zusammen, wenn der ausgewählte Instance-Typ ml.m5.24xlarge ist.

Trainingsgröße × Sicherheitsfaktor (12)	Instanzspeicher ml.m5.24xlarge	Erforderliche Mindestanzahl von Instanzen für das Training
12120 GB	384 GB	32

^*Aufgrund der Art des verteilten XGBoost-Trainings, bei dem der gesamte Trainingsdatensatz vor dem Training in ein DMatrix-Objekt geladen werden muss, und zusätzlicher freier Speicher, wird ein Sicherheitsfaktor von 10–12 empfohlen.

Um die Speicherauslastung für ein vollständiges SageMaker-Training von XGBoost auf dem bereitgestellten Datensatz genauer zu betrachten, stellen wir das entsprechende Diagramm bereit, das aus den Trainings gewonnen wurde Amazon CloudWatch Überwachung. Für diesen Trainingsjob wurden 40 ml.m5.24xlarge-Instanzen verwendet und die maximale Speicherauslastung erreichte etwa 62 %.

Die eingesparten Engineering-Kosten durch die Integration eines verwalteten ML-Dienstes wie SageMaker in die Datenpipeline betragen etwa 50 %. Die Option, Spot-Instances für Trainings- und Hyperparameter-Tuning-Jobs zu verwenden, senkt die Kosten um weitere 63 %.

Zusammenfassung

Mit SageMaker konnte das SophosAI-Team ein komplexes Projekt mit hoher Priorität erfolgreich lösen, indem es ein leichtgewichtiges XGBoost-Modell zur Erkennung von PDF-Malware erstellte, das viel kleiner auf der Festplatte (bis zu 25-mal kleiner) und im Arbeitsspeicher (bis zu 5-mal kleiner) ist als sein eigenes Detektor Vorgänger. Es ist ein kleiner, aber mächtiger Malware-Detektor mit ~0.99 AUC und einer True-Positive-Rate von 0.99 und einer False-Positive-Rate von . Dieses Modell kann schnell neu trainiert werden, und seine Leistung kann im Laufe der Zeit einfach überwacht werden, da es weniger als 20 Minuten dauert, um es mit mehr als 1 TB Daten zu trainieren.

Sie können den integrierten Algorithmus von SageMaker nutzen XGBoost zum Erstellen von Modellen mit Ihren tabellarischen Daten im Maßstab. Darüber hinaus können Sie auch die neuen integrierten Amazon SageMaker-Algorithmen LightGBM, CatBoost, AutoGluon-Tabular und Tab Transformer ausprobieren, wie hier beschrieben Blog.

---

Über die Autoren

Salma Taoufiq ist Senior Data Scientist bei Sophos und arbeitet an der Schnittstelle zwischen maschinellem Lernen und Cybersicherheit. Sie hat einen Bachelor-Hintergrund in Informatik und schloss ihr Studium an der Central European University mit einem MSc ab. in Mathematik und ihren Anwendungen. Wenn sie nicht gerade einen Malware-Detektor entwickelt, ist Salma eine begeisterte Wanderin, Reisende und Konsumentin von Thrillern.

Harini Kannan ist Data Scientist bei SophosAI. Sie ist seit ~4 Jahren in der Sicherheitsdatenwissenschaft tätig. Zuvor war sie Principal Data Scientist bei Capsule8, das von Sophos übernommen wurde. Sie hat Vorträge bei CAMLIS, BlackHat (USA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) und Data Connectors gehalten. Zu ihren Forschungsgebieten gehören die Erkennung hardwarebasierter Angriffe mithilfe von Leistungszählern, die Analyse des Benutzerverhaltens, interpretierbares ML und die unüberwachte Erkennung von Anomalien.

Hasan Poonawala ist Senior AI/ML Specialist Solutions Architect bei AWS mit Sitz in London, Großbritannien. Hasan unterstützt Kunden bei der Entwicklung und Bereitstellung von Machine-Learning-Anwendungen in der Produktion auf AWS. Er hat über 12 Jahre Berufserfahrung als Data Scientist, Machine Learning Praktiker und Softwareentwickler. In seiner Freizeit liebt Hasan es, die Natur zu erkunden und Zeit mit Freunden und Familie zu verbringen.

Digant Patel ist Enterprise Support Lead bei AWS. Er arbeitet mit Kunden zusammen, um die Cloud in großem Umfang zu entwerfen, bereitzustellen und zu betreiben. Seine Interessengebiete sind MLOps- und DevOps-Praktiken und wie sie Kunden bei ihrer Cloud-Reise helfen können. Außerhalb der Arbeit fotografiert er gerne, spielt Volleyball und verbringt Zeit mit Freunden und Familie.