Hvordan Sophos trener en kraftig, lett PDF-skadevaredetektor i ultraskala med Amazon SageMaker

Dette innlegget er medforfatter av Salma Taoufiq og Harini Kannan fra Sophos.

Som leder innen neste generasjons cybersikkerhet, Sophos streber etter å beskytte mer enn 500,000 150 organisasjoner og millioner av kunder i over XNUMX land mot stadige trusler. Drevet av trusselintelligens, maskinlæring (ML) og kunstig intelligens fra Sophos X-Ops, leverer Sophos en bred og variert portefølje av avanserte produkter og tjenester for å sikre og forsvare brukere, nettverk og endepunkter mot phishing, løsepengeprogramvare, skadelig programvare og det store spekteret av nettangrep der ute.

De Sophos Artificial Intelligence (AI) gruppe (SophosAI) overvåker utviklingen og vedlikeholdet av Sophos' viktigste ML-sikkerhetsteknologi. Sikkerhet er et stort dataproblem. For å unngå oppdagelse lager nettkriminelle stadig nye angrep. Dette oversetter seg til kolossale trusseldatasett som gruppen må jobbe med for å forsvare kundene best mulig. Et bemerkelsesverdig eksempel er gjenkjenning og eliminering av filer som var utspekulert med skadelig programvare, der datasettene er i terabyte.

I dette innlegget fokuserer vi på Sophos sitt malware-deteksjonssystem for PDF-filformatet spesifikt. Vi viser frem hvordan SophosAI bruker Amazon SageMaker distribuert trening med terabyte med data for å trene en kraftig, lett XGBoost (Extreme Gradient Boosting) modell. Dette gjør at teamet deres kan iterere over store treningsdata raskere med automatisk hyperparameterinnstilling og uten å administrere den underliggende treningsinfrastrukturen.

Løsningen er for øyeblikket sømløst integrert i produksjonsopplæringspipelinen og modellen distribuert på millioner av brukerendepunkter via Sophos endepunkttjeneste.

Bruk case kontekst

Enten du vil dele en viktig kontrakt eller bevare den fancy designen til CV-en din, er PDF-formatet det vanligste valget. Den utbredte bruken og den generelle oppfatningen om at slike dokumenter er lufttette og statiske har lurt brukerne inn i en falsk følelse av sikkerhet. PDF har derfor blitt en infeksjonsvektor i angripernes arsenal. Ondsinnede handlinger ved bruk av PDF-er oppnås oftest ved å bygge inn en JavaScript-nyttelast som kjøres av PDF-leseren for å laste ned et virus fra en URI, sabotere brukerens maskin eller stjele sensitiv informasjon.

Sophos oppdager ondsinnede PDF-filer på forskjellige punkter i et angrep ved å bruke et ensemble av deterministiske og ML-modeller. En slik tilnærming er illustrert i følgende diagram, der den skadelige PDF-filen leveres via e-post. Så snart et nedlastingsforsøk blir gjort, utløser det det ondsinnede kjørbare skriptet for å koble til angriperens kommando- og kontrollserver. SophosAIs PDF-detektor blokkerer nedlastingsforsøket etter å ha oppdaget at det er skadelig.

Andre måter inkluderer blokkering av PDF-filene i endepunktet, sending av de skadelige filene til en sandkasse (hvor den scores ved bruk av flere modeller), sende inn den skadelige filen til en scoringsinfrastruktur og generere en sikkerhetsrapport, og så videre.

Motivasjon

For å bygge en trebasert detektor som kan dømme ondsinnede PDF-filer med høy selvtillit, samtidig som det tillater lavt strømforbruk for endepunktdatabehandling og raske slutningssvar, fant SophosAI-teamet at XGBoost-algoritmen var en perfekt kandidat for oppgaven. Slike forskningsveier er viktige for Sophos av to grunner. Å ha kraftige, men små modeller utplassert på nivå med kundeendepunkter har stor innvirkning på selskapets produktanmeldelser av analytikere. Det gir også, og enda viktigere, en bedre brukeropplevelse totalt sett.

Teknisk utfordring

Fordi målet var å ha en modell med et mindre minnefotavtrykk enn deres eksisterende PDF-malwaredetektorer (både på disk og i minne), snudde SophosAI XGBoost, en klassifiseringsalgoritme som har dokumentert å produsere drastisk mindre modeller enn nevrale nettverk og samtidig oppnå imponerende ytelse på tabelldata. Før du begir deg ut i modellering av XGBoost-eksperimenter, var en viktig vurdering selve størrelsen på datasettet. Sophos' kjernedatasett av PDF-filer er faktisk i terabyte.

Derfor var hovedutfordringen å trene modellen med et stort datasett uten å måtte nedsample. Fordi det er avgjørende for detektoren å lære å oppdage PDF-baserte angrep – selv nål-i-høystakken og helt nye for å bedre forsvare Sophos-kunder – er det av ytterste viktighet å bruke alle tilgjengelige forskjellige datasett.

I motsetning til nevrale nettverk, hvor du kan trene i grupper, for XGBoost, trenger vi hele treningsdatasettet i minnet. Det største opplæringsdatasettet for dette prosjektet er over 1 TB, og det er ingen måte å trene i en slik skala uten å bruke metodikkene til et distribuert treningsrammeverk.

Løsningsoversikt

SageMaker er en fullstendig administrert ML-tjeneste som tilbyr ulike verktøy for å bygge, trene, optimalisere og distribuere ML-modeller. De SageMaker innebygde biblioteker av algoritmer består av 21 populære ML-algoritmer, inkludert XGBoost. (For mer informasjon, se Forenkle maskinlæring med XGBoost og Amazon SageMaker.) Med den innebygde XGBoost-algoritmen kan du dra nytte av åpen kildekode SageMaker XGBoost-beholder ved å spesifisere en rammeversjon større enn 1.0-1, som har forbedret fleksibilitet, skalerbarhet, utvidbarhet og Managed Spot Training, og støtter inndataformater som Parquet, som er formatet som brukes for PDF-datasettet.

Hovedgrunnen til at SophosAI valgte SageMaker er muligheten til å dra nytte av den fullstendig administrerte distribuerte opplæringen på multi-node CPU-forekomster ved ganske enkelt å spesifisere mer enn én forekomst. SageMaker deler automatisk dataene på tvers av noder, samler resultatene på tvers av peer-noder og genererer en enkelt modell. Forekomstene kan være Spot-forekomster, og dermed redusere treningskostnadene betydelig. Med innebygd algoritme for XGBoost, kan du gjøre dette uten ekstra tilpasset skript. Distribuerte versjoner av XGBoost finnes også som åpen kildekode, som f.eks XGBoost-Ray og XGBoost4J-Spark, men bruken av dem krever bygging, sikring, innstilling og selvadministrerende distribuerte databehandlingsklynger, noe som representerer betydelig innsats i tillegg til vitenskapelig utvikling.

I tillegg SageMaker automatisk modellinnstilling, også kjent som hyperparameterinnstilling, finner den beste versjonen av en modell ved å kjøre mange treningsjobber med intervaller av hyperparametere som du angir. Den velger deretter hyperparameterverdiene som resulterer i en modell som yter best, målt med en metrikk for den gitte ML-oppgaven.

Følgende diagram illustrerer løsningsarkitekturen.

Det er verdt å merke seg at da SophosAI startet XGBoost-eksperimenter før han vendte seg til SageMaker, ble det gjort forsøk på å bruke stort minne Amazon Elastic Compute Cloud (Amazon EC2)-forekomster (for eksempel r5a.24xlarge og x1.32xlarge) for å trene modellen på et så stort utvalg av dataene som mulig. Disse forsøkene tok imidlertid mer enn 10 timer i gjennomsnitt og mislyktes vanligvis på grunn av tom minne.

I motsetning til dette, ved å bruke SageMaker XGBoost-algoritmen og en problemfri distribuert treningsmekanisme, kunne SophosAI trene en boostermodell i skala på det kolossale PDF-treningsdatasettet i løpet av 20 minutter. Teamet måtte bare lagre dataene på Amazon enkel lagringstjeneste (Amazon S3) som Parkett-filer av lignende størrelse, og velg en EC2-instanstype og ønsket antall instanser, og SageMaker administrerte den underliggende dataklyngeinfrastrukturen og distribuerte opplæring mellom flere noder i klyngen. Under panseret deler SageMaker dataene på tvers av noder ved å bruke ShardedByS3Key for å fordele filobjektene likt mellom hver instans og bruker XGBoost-implementering av Rabit protokoll (pålitelig AllReduce og kringkastingsgrensesnitt) for å starte distribuert prosessering og kommunisere mellom primær- og noder. (For mer informasjon om histogramaggregering og kringkasting på tvers av noder, se XGBoost: Et skalerbart treforsterkningssystem.)

Utover å bare trene én modell, med SageMaker, XGBoost hyperparameterinnstilling ble også gjort raskt og enkelt med muligheten til å kjøre forskjellige eksperimenter samtidig for å finjustere den beste kombinasjonen av hyperparametre. De justerbare hyperparametrene inkluderer både boosterspesifikke og objektive funksjonsspesifikke hyperparametre. To søkestrategier tilbys: tilfeldig eller bayesisk. Den Bayesianske søkestrategien har vist seg å være verdifull fordi den hjelper til med å finne bedre hyperparametre enn bare et tilfeldig søk, i færre eksperimentelle iterasjoner.

Datasettinformasjon

SophosAIs PDF-malwaredeteksjonsmodellering er avhengig av en rekke funksjoner som n-gram-histogrammer og byte-entropifunksjoner (For mer informasjon, se MEADE: Mot en motor for ondsinnet e-postvedlegg). Metadata og funksjoner hentet fra innsamlede PDF-filer lagres i et distribuert datavarehus. Et datasett med over 3,500 funksjoner blir deretter beregnet, videre delt basert på tid i trenings- og testsett og lagret i batcher som parkettfiler i Amazon S3 for å være lett tilgjengelig for SageMaker for treningsjobber.

Tabellen nedenfor gir informasjon om trenings- og testdata.

datasett	Antall prøver	Antall parkettfiler	Total størrelse
Kurs	70,391,634	5,500	~ 1010 GB
Test	1,242,283	98	~ 18 GB

Datastørrelsene er beregnet etter formelen:

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

Formelen har følgende parametere:

• N er antall prøver i datasettet
• n_F er antall funksjoner, med n_F = 3585
• n_L er antall grunnsannhetsetiketter, med n_L = 1
• 4 er antall byte som trengs for funksjonenes datatype: float32

I tillegg gir de følgende kakediagrammene etikettfordelingen for både trenings- og testsettene, og fremkaller klasseubalansen i PDF-malwareoppdagingsoppgaven.

Fordelingen skifter fra treningssettet til XNUMX-måneders testsett. En tidsbasert oppdeling av datasettet i trening og testing brukes for å simulere det virkelige utrullingsscenarioet og unngå tidsmessig snoking. Denne strategien gjorde det også mulig for SophosAI å evaluere modellens sanne generaliseringsevner når de ble møtt med tidligere usett splitter nye PDF-angrep, for eksempel.

Eksperimenter og resultater

For å kickstarte eksperimenter trente SophosAI-teamet en baseline XGBoost-modell med standardparametere. Så begynte de å utføre finjustering av hyperparameter med SageMaker ved å bruke den Bayesianske strategien, som er så enkel som å spesifisere hyperparametere som skal stilles inn og ønsket verdiområde, evalueringsmetrikken (ROC (Receiver Operating Characteristic) AUC i dette tilfellet) og trenings- og valideringssettene. For PDF-malwaredetektoren prioriterte SophosAI hyperparametere inkludert antall forsterkingsrunder (num_round), den maksimale tredybden (max_depth), læringsraten (eta), og kolonnenes prøvetakingsforhold når du bygger trær (colsample_bytree). Til slutt ble de beste hyperparametrene oppnådd og brukt til å trene en modell på hele datasettet, og til slutt evaluert på holdout-testsettet.

Følgende plot viser objektiv metrikk (ROC AUC) vs. de 15 treningsjobbene som kjøres innenfor tuningjobben. De beste hyperparametrene er de som tilsvarer den niende treningsjobben.

I begynnelsen av SophosAIs eksperimenter på SageMaker var et spesielt viktig spørsmål å svare på: hvilken type forekomster og hvor mange av dem trengs for å trene XGBoost på dataene som er tilgjengelig? Dette er avgjørende fordi å bruke feil nummer eller type forekomst kan være bortkastet tid og penger; treningen vil garantert mislykkes på grunn av å gå tom for minne, eller, hvis du bruker for mange for store forekomster, kan dette bli unødvendig dyrt.

XGBoost er en minnebundet (i motsetning til beregningsbundet) algoritme. Så en generell beregningsforekomst (for eksempel M5) er et bedre valg enn en beregningsoptimalisert forekomst (for eksempel C4). For å ta en informert beslutning, er det en enkel SageMaker-retningslinje for å velge antall forekomster som kreves for å kjøre opplæring på hele datasettet:

Total treningsdatastørrelse × sikkerhetsfaktor^(*) < Forekomstantall × Forekomsttypens totale minne

I dette tilfellet: Total treningsdatastørrelse × sikkerhetsfaktor (12) = 12120 GB

Følgende tabell oppsummerer kravene når den valgte instanstypen er ml.m5.24xlarge.

Treningsstørrelse × sikkerhetsfaktor (12)	Forekomst Minne ml.m5.24xlarge	Minimum antall forekomster kreves for opplæring
12120 GB	384 GB	32

^*På grunn av XGBoost-distribuert trenings natur, som krever at hele treningsdatasettet lastes inn i et DMatrix-objekt før trening og ekstra ledig minne, anbefales en sikkerhetsfaktor på 10–12.

For å se nærmere på minneutnyttelsen for en full SageMaker-trening av XGBoost på det medfølgende datasettet, gir vi den tilsvarende grafen hentet fra opplæringens Amazon CloudWatch overvåkning. For denne treningsjobben ble det brukt 40 ml.m5.24xlarge forekomster og maksimal minneutnyttelse nådde rundt 62 %.

Ingeniørkostnadene som spares ved å integrere en administrert ML-tjeneste som SageMaker i datapipelinen er rundt 50 %. Alternativet til å bruke Spot Instances for trening og hyperparameterjusteringsjobber kuttet kostnadene med ytterligere 63 %.

konklusjonen

Med SageMaker kunne SophosAI-teamet lykkes med å løse et komplekst høyprioritert prosjekt ved å bygge en lett XGBoost-modell for deteksjon av skadelig programvare for PDF som er mye mindre på disken (opptil 25 ganger mindre) og i minnet (opptil 5 ganger mindre) enn dens detektorens forgjenger. Det er en liten, men mektig skadevaredetektor med ~0.99 AUC og en sann positiv rate på 0.99 og en falsk positiv rate på . Denne modellen kan raskt omskoleres, og ytelsen kan enkelt overvåkes over tid, fordi det tar mindre enn 20 minutter å trene den på mer enn 1 TB med data.

Du kan utnytte SageMaker innebygde algoritme Xgboost for å bygge modeller med tabelldataene dine i skala. I tillegg kan du også prøve de nye innebygde Amazon SageMaker-algoritmene LightGBM, CatBoost, AutoGluon-Tabular og Tab Transformer som beskrevet i denne blog.

---

Om forfatterne

Salma Taoufiq er senior dataforsker hos Sophos, og jobber i skjæringspunktet mellom maskinlæring og cybersikkerhet. Med en bachelorbakgrunn i informatikk, ble hun uteksaminert fra Central European University med en MSc. i matematikk og dens anvendelser. Når hun ikke utvikler en skadevaredetektor, er Salma en ivrig turgåer, reisende og forbruker av thrillere.

Harini Kannan er dataforsker ved SophosAI. Hun har vært i sikkerhetsdatavitenskap i ~4 år. Hun var tidligere hoveddataforsker ved Capsule8, som ble kjøpt opp av Sophos. Hun har holdt foredrag på CAMLIS, BlackHat (USA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) og Data Connectors. Forskningsområdene hennes inkluderer oppdagelse av maskinvarebaserte angrep ved hjelp av ytelsestellere, brukeratferdsanalyse, tolkbar ML og uovervåket anomalideteksjon.

Hasan Poonawala er en Senior AI/ML Specialist Solutions Architect ved AWS, basert i London, Storbritannia. Hasan hjelper kunder med å designe og distribuere maskinlæringsapplikasjoner i produksjon på AWS. Han har over 12 års arbeidserfaring som dataforsker, maskinlæringsutøver og programvareutvikler. På fritiden elsker Hasan å utforske naturen og tilbringe tid med venner og familie.

Digant Patel er en Enterprise Support Lead hos AWS. Han jobber med kunder for å designe, distribuere og operere i skyen i stor skala. Hans interesseområder er MLOps og DevOps-praksis og hvordan det kan hjelpe kunder i deres skyreise. Utenom jobben liker han å fotografere, spille volleyball og tilbringe tid med venner og familie.