Sådan træner Sophos en kraftfuld, let PDF-malwaredetektor i ultraskala med Amazon SageMaker

Dette indlæg er medforfattet af Salma Taoufiq og Harini Kannan fra Sophos.

Som førende inden for næste generations cybersikkerhed, Sophos stræber efter at beskytte mere end 500,000 organisationer og millioner af kunder på tværs af over 150 lande mod nye trusler. Drevet af trusselsintelligens, machine learning (ML) og kunstig intelligens fra Sophos X-Ops leverer Sophos en bred og varieret portefølje af avancerede produkter og tjenester til at sikre og forsvare brugere, netværk og slutpunkter mod phishing, ransomware, malware og den brede vifte af cyberangreb derude.

Sophos Artificial Intelligence (AI) gruppe (SophosAI) fører tilsyn med udviklingen og vedligeholdelsen af ​​Sophos' store ML-sikkerhedsteknologi. Sikkerhed er et big-data-problem. For at undgå opdagelse laver cyberkriminelle konstant nye angreb. Dette udmønter sig i kolossale trusselsdatasæt, som gruppen skal arbejde med for bedst muligt at forsvare kunderne. Et bemærkelsesværdigt eksempel er påvisning og eliminering af filer, der på snedigt vis var fyldt med malware, hvor datasættene er i terabyte.

I dette indlæg fokuserer vi på Sophos' malware-detektionssystem til PDF-filformatet specifikt. Vi viser, hvordan SophosAI bruger Amazon SageMaker distribueret træning med terabyte data for at træne en kraftfuld letvægts XGBoost-model (Extreme Gradient Boosting). Dette giver deres team mulighed for at iterere over store træningsdata hurtigere med automatisk hyperparameterjustering og uden at administrere den underliggende træningsinfrastruktur.

Løsningen er i øjeblikket problemfrit integreret i produktionstræningspipeline og modellen implementeret på millioner af brugerendepunkter via Sophos slutpunktstjeneste.

Brug case kontekst

Uanset om du vil dele en vigtig kontrakt eller bevare det smarte design af dit CV, er PDF-formatet det mest almindelige valg. Dens udbredte brug og den generelle opfattelse af, at sådanne dokumenter er lufttætte og statiske, har lullet brugerne ind i en falsk følelse af sikkerhed. PDF er derfor blevet en infektionsvektor i angriberens arsenal. Ondsindede handlinger ved hjælp af PDF'er opnås oftest ved at indlejre en JavaScript-nyttelast, der køres af PDF-læseren for at downloade en virus fra en URI, sabotere brugerens maskine eller stjæle følsomme oplysninger.

Sophos registrerer ondsindede PDF-filer på forskellige punkter i et angreb ved hjælp af en ensemble af deterministiske og ML-modeller. En sådan fremgangsmåde er illustreret i det følgende diagram, hvor den ondsindede PDF-fil leveres via e-mail. Så snart der er lavet et downloadforsøg, udløser det det ondsindede eksekverbare script til at oprette forbindelse til angriberens kommando- og kontrolserver. SophosAIs PDF-detektor blokerer downloadforsøget efter at have opdaget, at det er ondsindet.

Andre måder omfatter blokering af PDF-filerne i slutpunktet, afsendelse af de ondsindede filer til en sandkasse (hvor den scores ved brug af flere modeller), indsendelse af den ondsindede fil til en scoringsinfrastruktur og generering af en sikkerhedsrapport, og så videre.

Motivation

For at bygge en træ-baseret detektor, der kan dømme ondsindede PDF-filer med høj tillid, samtidig med at det giver mulighed for lavt strømforbrug til endpointcomputere og hurtige inferenssvar, fandt SophosAI-teamet, at XGBoost-algoritmen var en perfekt kandidat til opgaven. Sådanne forskningsmuligheder er vigtige for Sophos af to grunde. At have kraftfulde, men små modeller implementeret på niveau med kundeendepunkter, har stor indflydelse på virksomhedens produktanmeldelser fra analytikere. Det giver også, og endnu vigtigere, en bedre brugeroplevelse generelt.

Teknisk udfordring

Fordi målet var at have en model med et mindre hukommelsesfodaftryk end deres eksisterende PDF-malwaredetektorer (både på disk og i hukommelsen), forvandlede SophosAI XGBoost, en klassifikationsalgoritme med en dokumenteret rekord for at producere drastisk mindre modeller end neurale netværk og samtidig opnå imponerende ydeevne på tabeldata. Før man begav sig ud i at modellere XGBoost-eksperimenter, var en vigtig overvejelse selve størrelsen af ​​datasættet. Faktisk er Sophos' kernedatasæt af PDF-filer i terabyte.

Derfor var den største udfordring at træne modellen med et stort datasæt uden at skulle downsample. Fordi det er afgørende for detektoren at lære at spotte alle PDF-baserede angreb - selv nål-i-høstakken og helt nye for bedre at forsvare Sophos-kunder - er det af yderste vigtighed at bruge alle tilgængelige forskellige datasæt.

I modsætning til neurale netværk, hvor du kan træne i batches, skal vi til XGBoost bruge hele træningsdatasættet i hukommelsen. Det største træningsdatasæt for dette projekt er på over 1 TB, og der er ingen måde at træne i en sådan skala uden at bruge metoderne i en distribueret træningsramme.

Løsningsoversigt

SageMaker er en fuldt administreret ML-tjeneste, der tilbyder forskellige værktøjer til at bygge, træne, optimere og implementere ML-modeller. Det SageMaker indbyggede biblioteker af algoritmer består af 21 populære ML-algoritmer, inklusive XGBoost. (For mere information, se Forenkle maskinlæring med XGBoost og Amazon SageMaker.) Med den indbyggede XGBoost-algoritme kan du drage fordel af open source SageMaker XGBoost Container ved at specificere en rammeversion større end 1.0-1, som har forbedret fleksibilitet, skalerbarhed, udvidelsesmuligheder og Managed Spot Training og understøtter inputformater som Parquet, som er det format, der bruges til PDF-datasættet.

Hovedårsagen til, at SophosAI valgte SageMaker, er evnen til at drage fordel af den fuldt administrerede distribuerede træning på multi-node CPU-instanser ved blot at angive mere end én instans. SageMaker opdeler automatisk data på tværs af noder, aggregerer resultaterne på tværs af peer noder og genererer en enkelt model. Forekomsterne kan være Spot Forekomster, hvorved uddannelsesomkostningerne reduceres markant. Med indbygget algoritme til XGBoost, kan du gøre dette uden yderligere brugerdefineret script. Distribuerede versioner af XGBoost findes også som open source, som f.eks XGBoost-Ray , XGBoost4J-Spark, men deres brug kræver opbygning, sikring, tuning og selvadministrerende distribuerede computerklynger, hvilket repræsenterer en betydelig indsats ud over den videnskabelige udvikling.

Desuden SageMaker automatisk model tuning, også kendt som hyperparameter tuning, finder den bedste version af en model ved at køre mange træningsjob med intervaller af hyperparametre, som du angiver. Den vælger derefter de hyperparameterværdier, der resulterer i en model, der yder bedst, målt med en metrik for den givne ML-opgave.

Følgende diagram illustrerer løsningsarkitekturen.

Det er værd at bemærke, at da SophosAI startede XGBoost-eksperimenter før han henvendte sig til SageMaker, blev der gjort forsøg på at bruge stor hukommelse Amazon Elastic Compute Cloud (Amazon EC2)-forekomster (f.eks. r5a.24xlarge og x1.32xlarge) for at træne modellen på et så stort udsnit af dataene som muligt. Disse forsøg tog dog mere end 10 timer i gennemsnit og mislykkedes normalt, fordi de løb tør for hukommelse.

I modsætning hertil kunne SophosAI ved at bruge SageMaker XGBoost-algoritmen og en problemfri distribueret træningsmekanisme træne en boostermodel i skala på det kolossale PDF-træningsdatasæt i løbet af 20 minutter. Holdet skulle kun gemme dataene på Amazon Simple Storage Service (Amazon S3) som Parket-filer af lignende størrelse, og vælg en EC2-instanstype og det ønskede antal instanser, og SageMaker administrerede den underliggende computerklyngeinfrastruktur og distribuerede træning mellem flere noder i klyngen. Under hætten deler SageMaker dataene på tværs af noder ved hjælp af ShardedByS3Key til at fordele filobjekterne ligeligt mellem hver instans og bruger XGBoost-implementering af Rabit protokol (pålidelig AllReduce og broadcast interface) til at starte distribueret behandling og kommunikere mellem primære og peer noder. (For flere detaljer om histogramsammenlægning og udsendelse på tværs af noder, se XGBoost: Et skalerbart træforstærkningssystem.)

Udover blot at træne en model, med SageMaker, XGBoost hyperparameter tuning blev også gjort hurtigt og nemt med muligheden for at køre forskellige eksperimenter samtidigt for at finjustere den bedste kombination af hyperparametre. De indstillelige hyperparametre inkluderer både booster-specifikke og objektive funktionsspecifikke hyperparametre. To søgestrategier tilbydes: tilfældig eller Bayesiansk. Den Bayesianske søgestrategi har vist sig at være værdifuld, fordi den hjælper med at finde bedre hyperparametre end en ren tilfældig søgning, i færre eksperimentelle iterationer.

Datasæt oplysninger

SophosAIs PDF-malwaredetektionsmodellering er afhængig af en række funktioner såsom n-gram-histogrammer og byte-entropifunktioner (For mere information, se MEADE: Mod en ondsindet registreringsmaskine for vedhæftede e-mails). Metadata og funktioner udtrukket fra indsamlede PDF-filer gemmes i et distribueret datavarehus. Et datasæt med over 3,500 funktioner bliver derefter beregnet, yderligere opdelt baseret på tid i trænings- og testsæt og gemt i batches som parketfiler i Amazon S3 for at være let tilgængelige for SageMaker til træningsjob.

Følgende tabel giver oplysninger om trænings- og testdata.

datasæt	Antal prøver	Antal parketfiler	Samlet størrelse
Kurser	70,391,634	5,500	~ 1010 GB
Test	1,242,283	98	~ 18 GB

Datastørrelserne er beregnet efter formlen:

Datastørrelse = N × (n_F + nej_L)×4

Formlen har følgende parametre:

• N er antallet af prøver i datasættet
• n_F er antallet af funktioner, med n_F = 3585
• n_L er antallet af grundsandhedsmærker, med n_L = 1
• 4 er antallet af bytes, der er nødvendige for funktionernes datatype: float32

Derudover giver de følgende cirkeldiagrammer etiketfordelingen af ​​både trænings- og testsæt, hvilket fremkalder klasseubalancen, som PDF-malwaredetekteringsopgaven står over for.

Fordelingen skifter fra træningssættet til XNUMX-måneders testsæt. En tidsbaseret opdeling af datasættet i træning og test anvendes for at simulere det virkelige implementeringsscenarie og undgå tidsmæssig snooping. Denne strategi gjorde det også muligt for SophosAI at evaluere modellens sande generaliseringsevner, når de f.eks. stod over for hidtil usete splinternye PDF-angreb.

Eksperimenter og resultater

For at kickstarte eksperimenter trænede SophosAI-teamet en baseline XGBoost-model med standardparametre. Derefter begyndte de at udføre hyperparameter finjustering med SageMaker ved hjælp af den Bayesianske strategi, hvilket er så simpelt som at specificere hyperparametre der skal indstilles og det ønskede interval af værdier, evalueringsmetrikken (ROC (Receiver Operating Characteristic) AUC i dette tilfælde) og trænings- og valideringssættene. For PDF-malwaredetektoren prioriterede SophosAI hyperparametre, herunder antallet af boost-runder (num_round), den maksimale trædybde (max_depth), indlæringshastigheden (eta), og kolonnernes prøveudtagningsforhold ved bygning af træer (colsample_bytree). Til sidst blev de bedste hyperparametre opnået og brugt til at træne en model på det fulde datasæt og til sidst evalueret på holdout-testsættet.

Følgende plot viser objektiv-metrikken (ROC AUC) vs. de 15 træningsjob, der køres inden for tuning-jobbet. De bedste hyperparametre er dem, der svarer til det niende træningsjob.

I begyndelsen af ​​SophosAI's eksperimenter på SageMaker var et særligt vigtigt spørgsmål at besvare: hvilken type instanser og hvor mange af dem er nødvendige for at træne XGBoost på de tilgængelige data? Dette er afgørende, fordi det kan være spild af tid og penge at bruge det forkerte antal eller den forkerte type instans; træningen er nødt til at mislykkes på grund af at løbe tør for hukommelse, eller hvis du bruger for mange for store instanser, kan dette blive unødvendigt dyrt.

XGBoost er en hukommelsesbundet (i modsætning til computerbundet) algoritme. Så en generel computerforekomst (for eksempel M5) er et bedre valg end en computeroptimeret forekomst (for eksempel C4). For at træffe en informeret beslutning er der en simpel SageMaker-retningslinje til at vælge antallet af tilfælde, der kræves for at køre træning på det fulde datasæt:

Samlet træningsdatastørrelse × sikkerhedsfaktor^(*) < Forekomstantal × Forekomsttypens samlede hukommelse

I dette tilfælde: Samlet træningsdatastørrelse × sikkerhedsfaktor (12) = 12120 GB

Følgende tabel opsummerer kravene, når den valgte instanstype er ml.m5.24xlarge.

Træningsstørrelse × sikkerhedsfaktor (12)	Instance Hukommelse ml.m5.24xlarge	Minimum antal forekomster påkrævet for træning
12120 DK	384 DK	32

^*På grund af karakteren af ​​distribueret XGBoost-træning, som kræver, at hele træningsdatasættet indlæses i et DMatrix-objekt før træning og yderligere ledig hukommelse, anbefales en sikkerhedsfaktor på 10-12.

For at se nærmere på hukommelsesudnyttelsen til en fuld SageMaker-træning af XGBoost på det medfølgende datasæt, giver vi den tilsvarende graf opnået fra træningens amazoncloudwatch overvågning. Til dette træningsjob blev der brugt 40 ml.m5.24xlarge instanser og maksimal hukommelsesudnyttelse nåede op på omkring 62 %.

De tekniske omkostninger, der spares ved at integrere en administreret ML-tjeneste som SageMaker i datapipelinen er omkring 50 %. Muligheden for at bruge Spot Instances til træning og tuning af hyperparametre reducerede omkostningerne med yderligere 63 %.

Konklusion

Med SageMaker kunne SophosAI-teamet med succes løse et komplekst højprioritetsprojekt ved at bygge en letvægts PDF-malwaredetektion XGBoost-model, der er meget mindre på disken (op til 25 gange mindre) og i hukommelsen (op til 5 gange mindre) end dens detektors forgænger. Det er en lille, men mægtig malwaredetektor med ~0.99 AUC og en sand positiv rate på 0.99 og en falsk positiv rate på . Denne model kan hurtigt genoptrænes, og dens ydeevne kan nemt overvåges over tid, fordi det tager mindre end 20 minutter at træne den på mere end 1 TB data.

Du kan udnytte SageMaker indbyggede algoritme XGBoost til at bygge modeller med dine tabeldata i skala. Derudover kan du også prøve de nye indbyggede Amazon SageMaker-algoritmer LightGBM, CatBoost, AutoGluon-Tabular og Tab Transformer som beskrevet i denne blog.

---

Om forfatterne

Salma Taoufiq er Senior Data Scientist hos Sophos, der arbejder i krydsfeltet mellem maskinlæring og cybersikkerhed. Med en bachelorbaggrund i datalogi dimitterede hun fra Central European University med en MSc. i matematik og dens anvendelser. Når Salma ikke udvikler en malware-detektor, er Salma en ivrig vandrer, rejsende og forbruger af thrillere.

Harini Kannan er dataforsker hos SophosAI. Hun har været i sikkerhedsdatavidenskab i ~4 år. Hun var tidligere Principal Data Scientist hos Capsule8, som blev opkøbt af Sophos. Hun har holdt foredrag på CAMLIS, BlackHat (USA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) og Data Connectors. Hendes forskningsområder omfatter detektering af hardwarebaserede angreb ved hjælp af ydeevnetællere, brugeradfærdsanalyse, fortolkelig ML og uovervåget anomalidetektion.

Hasan Poonawala er Senior AI/ML Specialist Solutions Architect hos AWS, baseret i London, Storbritannien. Hasan hjælper kunder med at designe og implementere maskinlæringsapplikationer i produktion på AWS. Han har over 12 års erhvervserfaring som data scientist, machine learning practitioner og softwareudvikler. I sin fritid elsker Hasan at udforske naturen og tilbringe tid med venner og familie.

Digant Patel er Enterprise Support Lead hos AWS. Han arbejder med kunder om at designe, implementere og operere i skyen i stor skala. Hans interesseområder er MLOps og DevOps-praksis, og hvordan det kan hjælpe kunder i deres cloud-rejse. Uden for arbejdet nyder han at fotografere, spille volleyball og tilbringe tid med venner og familie.