Hoe Sophos een krachtige, lichtgewicht PDF-malwaredetector op ultraschaal traint met Amazon SageMaker

Dit bericht is mede geschreven door Salma Taoufiq en Harini Kannan van Sophos.

Als leider in cyberbeveiliging van de volgende generatie, Sophos streeft ernaar meer dan 500,000 organisaties en miljoenen klanten in meer dan 150 landen te beschermen tegen zich ontwikkelende bedreigingen. Aangedreven door bedreigingsinformatie, machine learning (ML) en kunstmatige intelligentie van Sophos X-Ops, levert Sophos een breed en gevarieerd portfolio van geavanceerde producten en diensten om gebruikers, netwerken en eindpunten te beveiligen en te beschermen tegen phishing, ransomware, malware en het brede scala aan cyberaanvallen dat er is.

De Sophos Artificial Intelligence (AI)-groep (SophosAI) houdt toezicht op de ontwikkeling en het onderhoud van Sophos' belangrijkste ML-beveiligingstechnologie. Beveiliging is een big data-probleem. Om detectie te ontwijken, bedenken cybercriminelen voortdurend nieuwe aanvallen. Dit vertaalt zich in kolossale dreigingsdatasets waarmee de groep moet werken om klanten zo goed mogelijk te verdedigen. Een opmerkelijk voorbeeld is de detectie en verwijdering van bestanden die sluw met malware waren doorspekt, waarbij de datasets in terabytes zijn.

In dit bericht richten we ons specifiek op het malwaredetectiesysteem van Sophos voor het PDF-bestandsformaat. We laten zien hoe SophosAI gebruikt Amazon Sage Maker gedistribueerde training met terabytes aan gegevens om een ​​krachtig lichtgewicht XGBoost-model (Extreme Gradient Boosting) te trainen. Hierdoor kan hun team sneller grote trainingsgegevens herhalen met automatische afstemming van hyperparameters en zonder de onderliggende trainingsinfrastructuur te beheren.

De oplossing is momenteel naadloos geïntegreerd in de productietrainingspijplijn en het model dat wordt geïmplementeerd op miljoenen gebruikerseindpunten via de Sophos-eindpuntservice.

Gebruik case-context

Of u nu een belangrijk contract wilt delen of het fraaie ontwerp van uw cv wilt behouden, het PDF-formaat is de meest gebruikelijke keuze. Het wijdverbreide gebruik en de algemene perceptie dat dergelijke documenten luchtdicht en statisch zijn, hebben gebruikers in een vals gevoel van veiligheid gesust. PDF is daarom een ​​infectievector bij uitstek geworden in het arsenaal van aanvallers. Kwaadaardige acties met behulp van PDF's worden meestal bereikt door een JavaScript-payload in te sluiten die door de PDF-lezer wordt uitgevoerd om een ​​virus van een URI te downloaden, de computer van de gebruiker te saboteren of gevoelige informatie te stelen.

Sophos detecteert kwaadaardige PDF-bestanden op verschillende punten van een aanval met behulp van een ensemble van deterministische en ML-modellen. Een dergelijke benadering wordt geïllustreerd in het volgende diagram, waar het schadelijke PDF-bestand via e-mail wordt afgeleverd. Zodra er een downloadpoging wordt gedaan, wordt het schadelijke uitvoerbare script geactiveerd om verbinding te maken met de Command and Control-server van de aanvaller. De PDF-detector van SophosAI blokkeert de downloadpoging nadat hij heeft gedetecteerd dat deze kwaadaardig is.

Andere manieren zijn onder meer het blokkeren van de PDF-bestanden in het eindpunt, het verzenden van de schadelijke bestanden naar een sandbox (waar het wordt gescoord met behulp van meerdere modellen), het indienen van het schadelijke bestand bij een score-infrastructuur en het genereren van een beveiligingsrapport, enzovoort.

Motivatie

Het SophosAI-team vond het XGBoost-algoritme een perfecte kandidaat voor de taak om een ​​op bomen gebaseerde detector te bouwen die kwaadaardige pdf's met veel vertrouwen kan veroordelen, terwijl het een laag eindpuntverbruik van rekenkracht en snelle inferentiereacties mogelijk maakt. Dergelijke onderzoekskanalen zijn om twee redenen belangrijk voor Sophos. De implementatie van krachtige maar kleine modellen op het niveau van de eindpunten van de klant heeft een grote impact op de productrecensies van het bedrijf door analisten. Het biedt ook, en nog belangrijker, een betere gebruikerservaring in het algemeen.

Technische uitdaging

Omdat het doel was om een ​​model te hebben met een kleinere geheugenvoetafdruk dan hun bestaande PDF-malwaredetectoren (zowel op schijf als in geheugen), maakte SophosAI XGBoost, een classificatie-algoritme met een bewezen staat van dienst bij het produceren van drastisch kleinere modellen dan neurale netwerken, terwijl het indrukwekkende resultaten behaalde. prestaties op tabelgegevens. Voordat we ons waagden aan het modelleren van XGBoost-experimenten, was een belangrijke overweging de enorme omvang van de dataset. Sophos' kerndataset van pdf-bestanden is inderdaad in terabytes.

Daarom was de grootste uitdaging het model te trainen met een grote dataset zonder te hoeven downsamplen. Omdat het van cruciaal belang is dat de detector alle op PDF gebaseerde aanvallen leert herkennen - zelfs naald-in-de-hooiberg en compleet nieuwe aanvallen om Sophos-klanten beter te verdedigen - is het van het grootste belang om alle beschikbare diverse datasets te gebruiken.

In tegenstelling tot neurale netwerken, waar je in batches kunt trainen, hebben we voor XGBoost de volledige trainingsdataset in het geheugen nodig. De grootste trainingsdataset voor dit project is meer dan 1 TB, en er is geen manier om op een dergelijke schaal te trainen zonder gebruik te maken van de methodologieën van een gedistribueerd trainingskader.

Overzicht oplossingen

SageMaker is een volledig beheerde ML-service die verschillende tools biedt voor het bouwen, trainen, optimaliseren en implementeren van ML-modellen. De Ingebouwde SageMaker-bibliotheken met algoritmen bestaat uit 21 populaire ML-algoritmen, waaronder XGBoost. (Voor meer informatie, zie Vereenvoudig machine learning met XGBoost en Amazon SageMaker.) Met het ingebouwde XGBoost-algoritme kunt u profiteren van de open-source SageMaker XGBoost-container door een framework-versie op te geven die groter is dan 1.0-1, met verbeterde flexibiliteit, schaalbaarheid, uitbreidbaarheid en Managed Spot Training, en ondersteunt invoerformaten zoals Parquet, het formaat dat wordt gebruikt voor de PDF-gegevensset.

De belangrijkste reden waarom SophosAI voor SageMaker heeft gekozen, is de mogelijkheid om te profiteren van de volledig beheerde gedistribueerde training op CPU-instanties met meerdere knooppunten door simpelweg meer dan één instantie te specificeren. SageMaker splitst de gegevens automatisch over knooppunten, aggregeert de resultaten over peer-knooppunten en genereert één enkel model. De instances kunnen Spot Instances zijn, waardoor de trainingskosten aanzienlijk worden verlaagd. Met de ingebouwd algoritme voor XGBoost, kunt u dit doen zonder extra aangepast script. Gedistribueerde versies van XGBoost bestaan ​​ook als open source, zoals: XGBoost-Ray en XGBoost4J-Spark, maar het gebruik ervan vereist het bouwen, beveiligen, afstemmen en zelf beheren van gedistribueerde computerclusters, wat een aanzienlijke inspanning betekent naast de wetenschappelijke ontwikkeling.

Bovendien Automatische afstemming van SageMaker-modellen, ook wel hyperparameter-afstemming genoemd, vindt de beste versie van een model door veel trainingstaken uit te voeren met door u opgegeven bereiken van hyperparameters. Vervolgens kiest het de hyperparameterwaarden die resulteren in een model dat het beste presteert, gemeten aan de hand van een metriek voor de gegeven ML-taak.

Het volgende diagram illustreert de oplossingsarchitectuur.

Het is vermeldenswaard dat, toen SophosAI XGBoost-experimenten startte voordat hij zich tot SageMaker wendde, er pogingen werden ondernomen om een ​​groot geheugen te gebruiken Amazon Elastic Compute-cloud (Amazon EC2) instanties (bijvoorbeeld r5a.24xlarge en x1.32xlarge) om het model te trainen op een zo groot mogelijke steekproef van de gegevens. Deze pogingen duurden echter gemiddeld meer dan 10 uur en mislukten meestal omdat er onvoldoende geheugen was.

Door gebruik te maken van het SageMaker XGBoost-algoritme en een probleemloos gedistribueerd trainingsmechanisme, kon SophosAI in slechts 20 minuten een boostermodel op schaal trainen op de kolossale PDF-trainingsdataset. Het team hoefde de gegevens alleen op te slaan op Amazon eenvoudige opslagservice (Amazon S3) als Parquet-bestanden van vergelijkbare grootte, en kies een EC2-instantietype en het gewenste aantal instanties, en SageMaker beheerde de onderliggende rekenclusterinfrastructuur en gedistribueerde training tussen meerdere knooppunten van het cluster. Onder de motorkap splitst SageMaker de gegevens over knooppunten met behulp van ShardedByS3Key om de bestandsobjecten gelijkelijk over elke instantie te verdelen en gebruikt XGBoost-implementatie van de Konijnenprotocol (betrouwbare AllReduce en broadcast-interface) om gedistribueerde verwerking te starten en te communiceren tussen primaire en peer-knooppunten. (Voor meer details over de histogramaggregatie en uitzending over knooppunten, zie: XGBoost: een schaalbaar boomversterkingssysteem.)

Naast het trainen van één model, met SageMaker, XGBoost-hyperparameterafstemming werd ook snel en gemakkelijk gemaakt met de mogelijkheid om verschillende experimenten tegelijkertijd uit te voeren om de beste combinatie van hyperparameters te verfijnen. De afstembare hyperparameters omvatten zowel boosterspecifieke als objectieve functiespecifieke hyperparameters. Twee zoekstrategieën worden aangeboden: willekeurig of Bayesiaans. De Bayesiaanse zoekstrategie is waardevol gebleken omdat het helpt bij het vinden van betere hyperparameters dan een willekeurige zoekopdracht, in minder experimentele iteraties.

Gegevensset-informatie

SophosAI's PDF-malwaredetectiemodellering is gebaseerd op een verscheidenheid aan functies, zoals n-gramhistogrammen en byte-entropiefuncties (raadpleeg voor meer informatie MEADE: op weg naar een detectie-engine voor kwaadaardige e-mailbijlagen). Metadata en functies geëxtraheerd uit verzamelde PDF-bestanden worden opgeslagen in een gedistribueerd datawarehouse. Een dataset van meer dan 3,500 functies wordt vervolgens berekend, verder opgesplitst op basis van tijd in trainings- en testsets en in batches opgeslagen als Parquet-bestanden in Amazon S3 om gemakkelijk toegankelijk te zijn voor SageMaker voor trainingstaken.

De volgende tabel geeft informatie over de trainings- en testgegevens.

dataset	Aantal monsters	Aantal parketbestanden	Totale grootte
Trainingen	70,391,634	5,500	~ 1010 GB
test	1,242,283	98	~ 18 GB

De gegevensgroottes zijn berekend volgens de formule:

Gegevensgrootte = N × (n_F + n_L)×4

De formule heeft de volgende parameters:

• N is het aantal steekproeven in de dataset
• n_F is het aantal functies, met n_F = 3585
• n_L is het aantal grondwaarheidslabels, met n_L = 1
• 4 is het aantal bytes dat nodig is voor het gegevenstype van de functies: float32

Bovendien geven de volgende cirkeldiagrammen de labeldistributie van zowel de trainings- als testsets, waardoor de klassenonevenwichtigheid waarmee de PDF-malwaredetectietaak wordt geconfronteerd, wordt uitgelokt.

De verdeling verschuift van de trainingsset naar de Testset van een maand. Een op tijd gebaseerde splitsing van de dataset in training en testen wordt toegepast om het real-life implementatiescenario te simuleren en tijdelijke snooping te voorkomen. Deze strategie stelde SophosAI ook in staat om de echte generalisatiemogelijkheden van het model te evalueren wanneer ze bijvoorbeeld geconfronteerd werden met voorheen ongeziene gloednieuwe PDF-aanvallen.

Experimenten en resultaten

Om experimenten op gang te brengen, heeft het SophosAI-team een ​​baseline XGBoost-model met standaardparameters getraind. Daarna begonnen ze hyperparameter-fijnafstemming uit te voeren met SageMaker met behulp van de Bayesiaanse strategie, die zo eenvoudig is als het specificeren van de hyperparameters af te stemmen en het gewenste bereik van waarden, de evaluatiemetriek (ROC (Receiver Operating Characteristic) AUC in dit geval) en de trainings- en validatiesets. Voor de PDF-malwaredetector gaf SophosAI prioriteit aan hyperparameters, inclusief het aantal boosting-rondes (num_round), de maximale boomdiepte (max_depth), de leersnelheid (eta), en de bemonsteringsratio van kolommen bij het bouwen van bomen (colsample_bytree). Uiteindelijk werden de beste hyperparameters verkregen en gebruikt om een ​​model te trainen op de volledige dataset, en uiteindelijk geëvalueerd op de holdout-testset.

De volgende grafiek toont de objectieve metriek (ROC AUC) versus de 15 trainingstaken die binnen de afstemmingstaak worden uitgevoerd. De beste hyperparameters zijn die welke overeenkomen met de negende trainingstaak.

Aan het begin van SophosAI's experimenten met SageMaker was een bijzonder belangrijke vraag om te beantwoorden: wat voor soort instanties en hoeveel zijn er nodig om XGBoost te trainen op de beschikbare gegevens? Dit is cruciaal omdat het gebruik van het verkeerde nummer of type instantie een verspilling van tijd en geld kan zijn; de training is gedoemd te mislukken als gevolg van onvoldoende geheugen, of, als er te veel te grote instanties worden gebruikt, kan dit onnodig duur worden.

XGBoost is een geheugengebonden (in tegenstelling tot computergebonden) algoritme. Een rekeninstantie voor algemeen gebruik (bijvoorbeeld M5) is dus een betere keuze dan een voor rekenkracht geoptimaliseerde instantie (bijvoorbeeld C4). Om een ​​weloverwogen beslissing te nemen, is er een eenvoudige SageMaker-richtlijn voor het kiezen van het aantal instanties dat nodig is om training op de volledige dataset uit te voeren:

Totale grootte van trainingsgegevens × veiligheidsfactor^(*) < Aantal exemplaren × Totaal geheugen van exemplaartype

In dit geval: Totale grootte van trainingsgegevens × veiligheidsfactor (12) = 12120 GB

De volgende tabel geeft een overzicht van de vereisten wanneer het gekozen instantietype ml.m5.24xlarge is.

Trainingsgrootte × veiligheidsfactor (12)	Instantiegeheugen ml.m5.24xlarge	Minimaal aantal exemplaren vereist voor training
12120 GB	384 GB	32

^*Vanwege de aard van XGBoost-gedistribueerde training, waarbij de volledige trainingsgegevensset vóór de training in een DMatrix-object moet worden geladen en extra vrij geheugen, wordt een veiligheidsfactor van 10-12 aanbevolen.

Om het geheugengebruik voor een volledige SageMaker-training van XGBoost op de verstrekte dataset van naderbij te bekijken, bieden we de bijbehorende grafiek die is verkregen uit de training's Amazon Cloud Watch toezicht houden. Voor deze trainingstaak werden 40 ml.m5.24xgrote exemplaren gebruikt en het maximale geheugengebruik bereikte ongeveer 62%.

De engineeringkosten die worden bespaard door een beheerde ML-service zoals SageMaker in de datapijplijn te integreren, bedragen ongeveer 50%. De optie om Spot Instances te gebruiken voor training en hyperparameter-afstemmingstaken verlaagt de kosten met nog eens 63%.

Conclusie

Met SageMaker kon het SophosAI-team een ​​complex project met hoge prioriteit met succes oplossen door een lichtgewicht PDF-malwaredetectie XGBoost-model te bouwen dat veel kleiner is op schijf (tot 25 keer kleiner) en in het geheugen (tot 5 keer kleiner) dan zijn voorloper van de detector. Het is een kleine maar krachtige malwaredetector met een AUC van ~0.99 en een echt-positief percentage van 0.99 en een fout-positief percentage van . Dit model kan snel opnieuw worden getraind en de prestaties ervan kunnen in de loop van de tijd gemakkelijk worden gevolgd, omdat het minder dan 20 minuten duurt om het te trainen op meer dan 1 TB aan gegevens.

U kunt gebruikmaken van het ingebouwde algoritme van SageMaker XGBoost voor het bouwen van modellen met uw tabelgegevens op schaal. Daarnaast kunt u ook de nieuwe ingebouwde Amazon SageMaker-algoritmen LightGBM, CatBoost, AutoGluon-Tabular en Tab Transformer proberen, zoals beschreven in deze blog.

---

Over de auteurs

Salma Taoufiq is Senior Data Scientist bij Sophos, werkzaam op het snijvlak van machine learning en cybersecurity. Met een niet-gegradueerde achtergrond in computerwetenschappen, studeerde ze af aan de Central European University met een MSc. in wiskunde en haar toepassingen. Als hij geen malwaredetector ontwikkelt, is Salma een fervent wandelaar, reiziger en consument van thrillers.

Harini Kannan is een datawetenschapper bij SophosAI. Ze werkt al ongeveer 4 jaar in de wetenschap van beveiligingsgegevens. Daarvoor was ze Principal Data Scientist bij Capsule8, dat werd overgenomen door Sophos. Ze heeft lezingen gegeven op CAMLIS, BlackHat (VS), Open Data Science Conference (Oost), Data Science Salon, PyData (Boston) en Data Connectors. Haar onderzoeksgebieden omvatten het detecteren van op hardware gebaseerde aanvallen met behulp van prestatiemeters, analyse van gebruikersgedrag, interpreteerbare ML en detectie van niet-gecontroleerde anomalie.

Hassan Poonawala is een Senior AI/ML Specialist Solutions Architect bij AWS, gevestigd in Londen, VK. Hasan helpt klanten bij het ontwerpen en implementeren van machine learning-applicaties in productie op AWS. Hij heeft meer dan 12 jaar werkervaring als datawetenschapper, machine learning beoefenaar en softwareontwikkelaar. In zijn vrije tijd houdt Hasan ervan om de natuur te verkennen en tijd door te brengen met vrienden en familie.

Digante Patel is een Enterprise Support Lead bij AWS. Hij werkt samen met klanten om op grote schaal in de cloud te ontwerpen, implementeren en werken. Zijn interessegebieden zijn MLOps en DevOps-praktijken en hoe het klanten kan helpen bij hun cloudreis. Naast zijn werk houdt hij van fotografie, volleyballen en tijd doorbrengen met vrienden en familie.