Det här inlägget är medförfattare av Salma Taoufiq och Harini Kannan från Sophos.

Som ledare inom nästa generations cybersäkerhet, Sophos strävar efter att skydda mer än 500,000 150 organisationer och miljontals kunder i över XNUMX länder mot föränderliga hot. Drivs av hotintelligens, maskininlärning (ML) och artificiell intelligens från Sophos X-Ops, levererar Sophos en bred och varierad portfölj av avancerade produkter och tjänster för att säkra och försvara användare, nätverk och slutpunkter mot nätfiske, ransomware, skadlig programvara och det breda utbudet av cyberattacker där ute.

Smakämnen Sophos Artificiell Intelligens (AI) grupp (SophosAI) övervakar utvecklingen och underhållet av Sophos stora ML-säkerhetsteknologi. Säkerhet är ett stort dataproblem. För att undvika upptäckt, skapar cyberbrottslingar ständigt nya attacker. Detta leder till kolossala hotdatauppsättningar som gruppen måste arbeta med för att försvara kunderna på bästa sätt. Ett anmärkningsvärt exempel är upptäckt och eliminering av filer som var listigt spetsade med skadlig programvara, där datamängderna är i terabyte.

I det här inlägget fokuserar vi på Sophos system för upptäckt av skadlig programvara för PDF-filformatet specifikt. Vi visar upp hur SophosAI använder Amazon SageMaker distribuerad träning med terabyte data för att träna en kraftfull lätt XGBoost-modell (Extreme Gradient Boosting). Detta gör att deras team kan iterera över stora träningsdata snabbare med automatisk hyperparameterjustering och utan att hantera den underliggande träningsinfrastrukturen.

Lösningen är för närvarande sömlöst integrerad i produktionsträningspipelinen och modellen distribueras på miljontals användarslutpunkter via Sophos slutpunktstjänst.

Använd fallsammanhang

Oavsett om du vill dela ett viktigt kontrakt eller bevara den snygga designen av ditt CV, är PDF-formatet det vanligaste valet. Dess utbredda användning och den allmänna uppfattningen att sådana dokument är lufttäta och statiska har invaggat användarna i en falsk känsla av säkerhet. PDF har därför blivit en infektionsvektor i angriparnas arsenal. Skadliga handlingar med PDF-filer uppnås oftast genom att bädda in en JavaScript-nyttolast som körs av PDF-läsaren för att ladda ner ett virus från en URI, sabotera användarens dator eller stjäla känslig information.

Sophos upptäcker skadliga PDF-filer vid olika punkter under en attack med hjälp av en ensemble av deterministiska och ML-modeller. Ett sådant tillvägagångssätt illustreras i följande diagram, där den skadliga PDF-filen levereras via e-post. Så snart ett nedladdningsförsök görs utlöser det det skadliga körbara skriptet att ansluta till angriparens kommando- och kontrollserver. SophosAI:s PDF-detektor blockerar nedladdningsförsöket efter att ha upptäckt att det är skadligt.

Andra sätt inkluderar att blockera PDF-filerna i slutpunkten, skicka de skadliga filerna till en sandlåda (där den får poäng med flera modeller), skicka in den skadliga filen till en poänginfrastruktur och generera en säkerhetsrapport och så vidare.

Motivation

För att bygga en trädbaserad detektor som kan döma skadliga PDF-filer med hög tillförsikt, samtidigt som den tillåter låg effektförbrukning för slutpunktsdatorer och snabba slutledningssvar, fann SophosAI-teamet att XGBoost-algoritmen var en perfekt kandidat för uppgiften. Sådana forskningsvägar är viktiga för Sophos av två skäl. Att ha kraftfulla men ändå små modeller utplacerade på nivån av kundslutpunkter har stor inverkan på företagets produktrecensioner av analytiker. Det ger också, och ännu viktigare, en bättre användarupplevelse totalt sett.

Teknisk utmaning

Eftersom målet var att ha en modell med ett mindre minnesfotavtryck än deras befintliga PDF-malwaredetektorer (både på disken och i minnet), förvandlade SophosAI XGBoost, en klassificeringsalgoritm som har bevisat att de producerar drastiskt mindre modeller än neurala nätverk samtidigt som den uppnådde imponerande prestanda på tabelldata. Innan man vågade sig på att modellera XGBoost-experiment var ett viktigt övervägande själva storleken på datamängden. Sophos kärndatauppsättning av PDF-filer är faktiskt i terabyte.

Därför var den största utmaningen att träna modellen med en stor datamängd utan att behöva nedsampla. Eftersom det är avgörande för detektorn att lära sig att upptäcka alla PDF-baserade attacker – även nål-i-höstacken och helt nya sådana för att bättre försvara Sophos-kunder – är det av yttersta vikt att använda alla tillgängliga olika datauppsättningar.

Till skillnad från neurala nätverk, där du kan träna i omgångar, behöver vi för XGBoost hela träningsdataset i minnet. Den största utbildningsdatauppsättningen för detta projekt är över 1 TB, och det finns inget sätt att träna i en sådan skala utan att använda metoderna i ett distribuerat utbildningsramverk.

Lösningsöversikt

SageMaker är en helt hanterad ML-tjänst som tillhandahåller olika verktyg för att bygga, träna, optimera och distribuera ML-modeller. De SageMaker inbyggda bibliotek av algoritmer består av 21 populära ML-algoritmer, inklusive XGBoost. (För mer information, se Förenkla maskininlärningen med XGBoost och Amazon SageMaker.) Med den inbyggda XGBoost-algoritmen kan du dra nytta av öppen källkod SageMaker XGBoost Container genom att ange en ramversion större än 1.0-1, som har förbättrad flexibilitet, skalbarhet, utökningsbarhet och Managed Spot Training, och stöder inmatningsformat som Parquet, vilket är formatet som används för PDF-datauppsättningen.

Den främsta anledningen till att SophosAI valde SageMaker är möjligheten att dra nytta av den fullt hanterade distribuerade utbildningen på CPU-instanser med flera noder genom att helt enkelt specificera mer än en instans. SageMaker delar automatiskt upp data mellan noder, aggregerar resultaten över peer-noder och genererar en enda modell. Förekomsterna kan vara Spot-instanser, vilket avsevärt minskar utbildningskostnaderna. Med inbyggd algoritm för XGBoost, kan du göra detta utan något extra anpassat skript. Distribuerade versioner av XGBoost finns också som öppen källkod, som t.ex XGBoost-Ray och XGBoost4J-Spark, men deras användning kräver att bygga, säkra, ställa in och självhantera distribuerade datorkluster, vilket representerar betydande ansträngningar utöver den vetenskapliga utvecklingen.

Dessutom används SageMaker automatisk modelljustering, även känd som hyperparameterjustering, hittar den bästa versionen av en modell genom att köra många träningsjobb med intervall av hyperparametrar som du anger. Den väljer sedan hyperparametervärdena som resulterar i en modell som presterar bäst, mätt med ett mått för den givna ML-uppgiften.

Följande diagram illustrerar lösningsarkitekturen.

Det är värt att notera att när SophosAI startade XGBoost-experiment innan han vände sig till SageMaker, gjordes försök att använda stort minne Amazon Elastic Compute Cloud (Amazon EC2)-instanser (till exempel r5a.24xlarge och x1.32xlarge) för att träna modellen på ett så stort urval av data som möjligt. Dessa försök tog dock mer än 10 timmar i genomsnitt och misslyckades vanligtvis på grund av att minnet tog slut.

Däremot, genom att använda SageMaker XGBoost-algoritmen och en problemfri distribuerad träningsmekanism, kunde SophosAI träna en boostermodell i skala på den kolossala PDF-träningsdatauppsättningen på 20 minuter. Teamet behövde bara lagra data på Amazon enkel lagringstjänst (Amazon S3) som Parkettfiler av liknande storlek, och välj en EC2-instanstyp och önskat antal instanser, och SageMaker hanterade den underliggande klusterinfrastrukturen och distribuerade utbildning mellan flera noder i klustret. Under huven delar SageMaker upp data över noder med ShardedByS3Key för att fördela filobjekten lika mellan varje instans och använder XGBoost-implementering av Rabit protokoll (pålitligt AllReduce och broadcast-gränssnitt) för att starta distribuerad bearbetning och kommunicera mellan primära och peer noder. (För mer information om histogramaggregation och sändning över noder, se XGBoost: Ett skalbart trädförstärkningssystem.)

Utöver att bara träna en modell, med SageMaker, XGBoost hyperparameterjustering gjordes också snabbt och enkelt med möjligheten att köra olika experiment samtidigt för att finjustera den bästa kombinationen av hyperparametrar. De avstämbara hyperparametrarna inkluderar både boosterspecifika och objektiva funktionsspecifika hyperparametrar. Två sökstrategier erbjuds: slumpmässigt eller Bayesian. Den Bayesianska sökstrategin har visat sig vara värdefull eftersom den hjälper till att hitta bättre hyperparametrar än en ren slumpmässig sökning, i färre experimentella iterationer.

Datauppsättningsinformation

SophosAI:s PDF-malodetekteringsmodellering förlitar sig på en mängd funktioner som n-gram histogram och byte entropi funktioner (För mer information, se MEADE: Mot en motor för upptäckt av skadliga e-postbilagor). Metadata och funktioner som extraherats från insamlade PDF-filer lagras i ett distribuerat datalager. En datauppsättning med över 3,500 3 funktioner beräknas sedan, delas upp baserat på tid i tränings- och testset och lagras i omgångar som parkettfiler i Amazon SXNUMX för att vara lättillgängliga för SageMaker för utbildningsjobb.

Följande tabell ger information om tränings- och testdata.

Datastorlekarna har beräknats enligt formeln:

Datastorlek = N × (n_F + nej_L) x 4

Formeln har följande parametrar:

• N är antalet prover i datamängden
• n_F är antalet funktioner, med n_F = 3585
• n_L är antalet marksanningsetiketter, med n_L = 1
• 4 är antalet byte som behövs för funktionernas datatyp: float32

Dessutom tillhandahåller följande cirkeldiagram etikettfördelningen för både tränings- och testuppsättningarna, vilket framkallar klassobalansen i uppgiften att upptäcka skadlig PDF.

Fördelningen skiftar från träningssetet till Enmånadstestset. En tidsbaserad uppdelning av datamängden i utbildning och testning tillämpas för att simulera det verkliga implementeringsscenariot och undvika tidsmässig snooping. Den här strategin gjorde det också möjligt för SophosAI att utvärdera modellens verkliga generaliseringsförmåga när man ställs inför tidigare osynliga helt nya PDF-attacker, till exempel.

Experiment och resultat

För att kickstarta experiment tränade SophosAI-teamet en baslinje XGBoost-modell med standardparametrar. Sedan började de utföra hyperparameterfinjustering med SageMaker med den Bayesianska strategin, vilket är så enkelt som att specificera hyperparametrar som ska ställas in och det önskade intervallet av värden, utvärderingsmåttet (ROC (Receiver Operating Characteristic) AUC i detta fall) och tränings- och valideringsuppsättningarna. För detektorn för PDF-skadlig programvara prioriterade SophosAI hyperparametrar inklusive antalet boostronder (num_round), det maximala träddjupet (max_depth), inlärningshastigheten (eta), och kolumnernas provtagningsförhållande när man bygger träd (colsample_bytree). Så småningom erhölls de bästa hyperparametrarna och användes för att träna en modell på den fullständiga datamängden, och slutligen utvärderades på holdout-testsetet.

Följande plot visar objektivmåttet (ROC AUC) kontra de 15 träningsjobben som körs inom trimningsjobbet. De bästa hyperparametrarna är de som motsvarar det nionde träningsjobbet.

I början av SophosAI:s experiment på SageMaker var en särskilt viktig fråga att besvara: vilken typ av instanser och hur många av dem behövs för att träna XGBoost på data till hands? Detta är avgörande eftersom att använda fel nummer eller typ av instans kan vara ett slöseri med tid och pengar; träningen kommer sannolikt att misslyckas på grund av att minnet tar slut, eller om du använder för många för stora instanser kan detta bli onödigt dyrt.

XGBoost är en minnesbunden (i motsats till datorbunden) algoritm. Så en allmän beräkningsinstans (till exempel M5) är ett bättre val än en beräkningsoptimerad instans (till exempel C4). För att fatta ett välgrundat beslut finns det en enkel SageMaker-riktlinje för att välja antalet instanser som krävs för att köra utbildning på hela datasetet:

Total träningsdatastorlek × säkerhetsfaktor^(*) < Antal instanser × instanstypens totala minne

I detta fall: Total träningsdatastorlek × säkerhetsfaktor (12) = 12120 GB

Följande tabell sammanfattar kraven när den valda instanstypen är ml.m5.24xlarge.

^*På grund av XGBoost-distribuerad tränings natur, som kräver att hela träningsdataset laddas in i ett DMatrix-objekt före träning och ytterligare ledigt minne, rekommenderas en säkerhetsfaktor på 10–12.

För att ta en närmare titt på minnesutnyttjandet för en fullständig SageMaker-utbildning av XGBoost på den tillhandahållna datamängden tillhandahåller vi motsvarande graf som erhålls från utbildningens amazoncloudwatch övervakning. För detta träningsjobb användes 40 ml.m5.24xlarge instanser och maximalt minnesutnyttjande nådde cirka 62 %.

Den tekniska kostnaden som sparas genom att integrera en hanterad ML-tjänst som SageMaker i datapipelinen är cirka 50 %. Möjligheten att använda Spot Instances för träning och hyperparameterjustering minskade kostnaderna med ytterligare 63 %.

Slutsats

Med SageMaker kunde SophosAI-teamet framgångsrikt lösa ett komplext högprioriterat projekt genom att bygga en lätt XGBoost-modell för upptäckt av skadlig programvara för PDF som är mycket mindre på disken (upp till 25 gånger mindre) och i minnet (upp till 5 gånger mindre) än dess detektorns föregångare. Det är en liten men mäktig detektor för skadlig programvara med ~0.99 AUC och en sann positiv frekvens på 0.99 och en falsk positiv frekvens på . Denna modell kan snabbt omskolas och dess prestanda kan enkelt övervakas över tid, eftersom det tar mindre än 20 minuter att träna den på mer än 1 TB data.

Du kan utnyttja SageMaker inbyggda algoritm XGBoost för att bygga modeller med dina tabelldata i stor skala. Dessutom kan du också prova de nya inbyggda Amazon SageMaker-algoritmerna LightGBM, CatBoost, AutoGluon-Tabular och Tab Transformer som beskrivs i detta blogg.

Om författarna

Salma Taoufiq är Senior Data Scientist på Sophos, som arbetar i skärningspunkten mellan maskininlärning och cybersäkerhet. Med en grundutbildning i datavetenskap tog hon examen från Central European University med en MSc. i matematik och dess tillämpningar. När hon inte utvecklar en detektor för skadlig programvara är Salma en ivrig vandrare, resenär och konsument av thrillers.

Harini Kannan är dataforskare på SophosAI. Hon har varit inom säkerhetsdatavetenskap i ~4 år. Hon var tidigare Principal Data Scientist på Capsule8, som förvärvades av Sophos. Hon har hållit föredrag på CAMLIS, BlackHat (USA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) och Data Connectors. Hennes forskningsområden inkluderar detektering av hårdvarubaserade attacker med hjälp av prestandaräknare, analys av användarbeteende, tolkbar ML och oövervakad upptäckt av anomalier.

Hasan Poonawala är Senior AI/ML Specialist Solutions Architect på AWS, baserad i London, Storbritannien. Hasan hjälper kunder att designa och distribuera maskininlärningsapplikationer i produktion på AWS. Han har över 12 års arbetslivserfarenhet som datavetare, maskininlärningsutövare och mjukvaruutvecklare. På sin fritid älskar Hasan att utforska naturen och umgås med vänner och familj.

Digant Patel är en Enterprise Support Lead på AWS. Han arbetar med kunder för att designa, distribuera och arbeta i moln i stor skala. Hans intresseområden är MLOps och DevOps praxis och hur det kan hjälpa kunder i deras molnresa. Utanför jobbet tycker han om att fotografera, spela volleyboll och umgås med vänner och familj.