Cum Sophos antrenează un detector de malware PDF puternic și ușor la scară ultra cu Amazon SageMaker

Această postare este în colaborare cu Salma Taoufiq și Harini Kannan de la Sophos.

În calitate de lider în securitatea cibernetică de ultimă generație, Sophos se străduiește să protejeze peste 500,000 de organizații și milioane de clienți din peste 150 de țări împotriva amenințărilor în evoluție. Având la bază inteligența amenințărilor, învățarea automată (ML) și inteligența artificială de la Sophos X-Ops, Sophos oferă un portofoliu larg și variat de produse și servicii avansate pentru a securiza și apăra utilizatorii, rețelele și punctele finale împotriva phishing, ransomware, malware și gama largă de atacuri cibernetice de acolo.

Grupul Sophos de inteligență artificială (AI). (SophosAI) supraveghează dezvoltarea și întreținerea tehnologiei majore de securitate ML Sophos. Securitatea este o problemă de date mari. Pentru a evita detectarea, infractorii cibernetici creează în mod constant atacuri noi. Acest lucru se traduce în seturi de date colosale de amenințări cu care grupul trebuie să lucreze pentru a apăra cel mai bine clienții. Un exemplu notabil este detectarea și eliminarea fișierelor care au fost împletite cu viclenie cu malware, unde seturile de date sunt în terabytes.

În această postare, ne concentrăm asupra sistemului de detectare a malware-ului Sophos în special pentru formatul de fișier PDF. Vă prezentăm cum utilizează SophosAI Amazon SageMaker antrenament distribuit cu terabytes de date pentru a antrena un model XGBoost (Extreme Gradient Boosting) ușor și puternic. Acest lucru permite echipei lor să itereze mai rapid date mari de antrenament, cu reglarea automată a hiperparametrilor și fără a gestiona infrastructura de antrenament de bază.

Soluția este în prezent integrată perfect în conducta de instruire de producție și modelul este implementat pe milioane de puncte finale ale utilizatorilor prin intermediul Serviciul terminal Sophos.

Contextul cazului de utilizare

Indiferent dacă doriți să împărtășiți un contract important sau să păstrați designul elegant al CV-ului dvs., formatul PDF este cea mai comună alegere. Utilizarea sa pe scară largă și percepția generală că astfel de documente sunt etanșe și statice i-au amânat pe utilizatori într-un fals sentiment de securitate. Prin urmare, PDF a devenit un vector de infecție de alegere în arsenalul atacatorilor. Acțiunile rău intenționate care utilizează PDF-uri sunt realizate cel mai adesea prin încorporarea unei încărcături utile JavaScript care este rulată de cititorul PDF pentru a descărca un virus dintr-un URI, a sabota computerul utilizatorului sau a fura informații sensibile.

Sophos detectează fișiere PDF rău intenționate în diferite puncte ale unui atac folosind un ansamblu de modele deterministe și ML. O astfel de abordare este ilustrată în diagrama următoare, în care fișierul PDF rău intenționat este livrat prin e-mail. De îndată ce este făcută o încercare de descărcare, declanșează scriptul executabil rău intenționat să se conecteze la serverul de comandă și control al atacatorului. Detectorul PDF al SophosAI blochează încercarea de descărcare după ce detectează că este rău intenționat.

Alte modalități includ blocarea fișierelor PDF din punctul final, trimiterea fișierelor rău intenționate într-un sandbox (unde este punctat folosind mai multe modele), trimiterea fișierului rău intenționat la o infrastructură de scor și generarea unui raport de securitate și așa mai departe.

motivaţia

Pentru a construi un detector bazat pe arbore care poate condamna PDF-urile rău intenționate cu mare încredere, permițând în același timp un consum redus de energie de calcul și răspunsuri rapide de inferență, echipa SophosAI a descoperit că algoritmul XGBoost este candidatul perfect pentru această sarcină. Astfel de căi de cercetare sunt importante pentru Sophos din două motive. Având modele puternice, dar mici, implementate la nivelul punctelor finale ale clienților, are un impact mare asupra recenziilor de produse ale companiei de către analiști. De asemenea, și mai important, oferă o experiență de utilizare mai bună în general.

Provocare tehnică

Deoarece scopul era acela de a avea un model cu o amprentă de memorie mai mică decât detectoarele lor de malware PDF existente (atât pe disc, cât și în memorie), SophosAI a transformat XGBoost, un algoritm de clasificare cu un record dovedit de a produce modele drastic mai mici decât rețelele neuronale, obținând în același timp rezultate impresionante. performanța pe datele tabelare. Înainte de a te aventura în modelarea experimentelor XGBoost, o considerație importantă a fost dimensiunea mare a setului de date. Într-adevăr, setul de date de bază al fișierelor PDF al Sophos este în terabyți.

Prin urmare, principala provocare a fost antrenarea modelului cu un set mare de date fără a fi nevoie să eșantioneze. Deoarece este esențial ca detectorul să învețe să detecteze orice atacuri bazate pe PDF - chiar și cele cu ace în carul de fân și cele complet noi pentru a apăra mai bine clienții Sophos - este de cea mai mare importanță să folosiți toate seturile de date disponibile.

Spre deosebire de rețelele neuronale, în care vă puteți antrena în loturi, pentru XGBoost, avem nevoie de întregul set de date de antrenament în memorie. Cel mai mare set de date de instruire pentru acest proiect este de peste 1 TB și nu există nicio modalitate de a se antrena la o asemenea scară fără a utiliza metodologiile unui cadru de instruire distribuit.

Prezentare generală a soluțiilor

SageMaker este un serviciu ML gestionat complet, care oferă diverse instrumente pentru a construi, antrena, optimiza și implementa modele ML. The Biblioteci de algoritmi încorporate SageMaker constau din 21 de algoritmi ML populari, inclusiv XGBoost. (Pentru mai multe informații, vezi Simplificați învățarea automată cu XGBoost și Amazon SageMaker.) Cu algoritmul încorporat XGBoost, puteți profita de open-source Container SageMaker XGBoost prin specificarea unei versiuni de cadru mai mare decât 1.0-1, care a îmbunătățit flexibilitatea, scalabilitatea, extensibilitatea și Managed Spot Training și acceptă formate de intrare precum Parquet, care este formatul utilizat pentru setul de date PDF.

Principalul motiv pentru care SophosAI a ales SageMaker este capacitatea de a beneficia de instruirea distribuită complet gestionată pe instanțele CPU cu mai multe noduri, prin simpla specificare a mai multor instanțe. SageMaker împarte automat datele în noduri, agregează rezultatele între noduri egale și generează un singur model. Instanțele pot fi Instanțe Spot, reducând astfel semnificativ costurile de instruire. Cu algoritm încorporat pentru XGBoost, puteți face acest lucru fără niciun script personalizat suplimentar. Versiunile distribuite de XGBoost există și ca sursă deschisă, cum ar fi XGBoost-Ray și XGBoost4J-Spark, dar utilizarea lor necesită construirea, securizarea, reglarea și autogestionarea clusterelor de calcul distribuite, ceea ce reprezintă un efort semnificativ suplimentar dezvoltării științifice.

În plus, Ajustarea automată a modelului SageMaker, cunoscut și sub denumirea de ajustare a hiperparametrilor, găsește cea mai bună versiune a unui model rulând multe joburi de antrenament cu intervale de hiperparametri pe care le specificați. Apoi alege valorile hiperparametrului care rezultă într-un model care are cele mai bune performanțe, măsurate printr-o metrică pentru sarcina ML dată.

Următoarea diagramă ilustrează arhitectura soluției.

Este demn de remarcat faptul că, atunci când SophosAI a început experimentele XGBoost înainte de a apela la SageMaker, s-au încercat să folosească memoria mare. Cloud Elastic de calcul Amazon Instanțe (Amazon EC2) (de exemplu, r5a.24xlarge și x1.32xlarge) pentru a antrena modelul pe un eșantion cât mai mare de date. Cu toate acestea, aceste încercări au durat în medie mai mult de 10 ore și, de obicei, au eșuat din cauza lipsei de memorie.

În schimb, utilizând algoritmul SageMaker XGBoost și un mecanism de antrenament distribuit fără probleme, SophosAI ar putea antrena un model de amplificare la scară pe setul colosal de date de antrenament PDF în doar 20 de minute. Echipa nu trebuia decât să stocheze datele pe Serviciul Amazon de stocare simplă (Amazon S3) ca fișiere Parquet de dimensiuni similare și alegeți un tip de instanță EC2 și numărul dorit de instanțe, iar SageMaker a gestionat infrastructura clusterului de calcul de bază și a distribuit instruirea între mai multe noduri ale clusterului. Sub capotă, SageMaker împarte datele în noduri folosind ShardedByS3Key pentru a distribui obiectele fișierului în mod egal între fiecare instanță și utilizează implementarea XGBoost a sistemului. Protocolul iepurelui (interfață AllReduce și broadcast de încredere) pentru a lansa procesarea distribuită și a comunica între nodurile primare și peer. (Pentru mai multe detalii despre agregarea histogramei și difuzarea între noduri, consultați XGBoost: Un sistem scalabil de creștere a arborilor.)

Dincolo de doar antrenamentul unui model, cu SageMaker, Ajustarea hiperparametrului XGBoost a fost, de asemenea, făcut rapid și ușor cu capacitatea de a rula diferite experimente simultan pentru a regla cea mai bună combinație de hiperparametri. Hiperparametrii reglabili includ atât hiperparametri specifici amplificatorului, cât și hiperparametri specifici funcției obiective. Două strategii de căutare sunt oferite: aleatoriu sau bayesian. Strategia de căutare bayesiană s-a dovedit a fi valoroasă deoarece ajută la găsirea unor hiperparametri mai buni decât o simplă căutare aleatorie, în mai puține iterații experimentale.

Informații despre setul de date

Modelarea de detectare a malware-ului PDF de la SophosAI se bazează pe o varietate de caracteristici, cum ar fi histogramele n-grame și caracteristicile de entropie de octeți (Pentru mai multe informații, consultați MEADE: Către un motor de detectare a atașamentelor de e-mail rău intenționate). Metadatele și caracteristicile extrase din fișierele PDF colectate sunt stocate într-un depozit de date distribuit. Un set de date de peste 3,500 de caracteristici este apoi calculat, împărțit în continuare în funcție de timp în seturi de instruire și testare și stocat în loturi ca fișiere Parquet în Amazon S3 pentru a fi ușor accesibil de către SageMaker pentru joburi de instruire.

Următorul tabel oferă informații despre datele de instruire și test.

Setul de date	Număr de probe	Numărul de dosare de parchet	Marimea totala
Pregătire	70,391,634	5,500	~ 1010 GB
Test	1,242,283	98	~ 18 GB

Dimensiunile datelor au fost calculate după formula:

Dimensiunea datelor = N × (n_F + nu_L) × 4

Formula are următorii parametri:

• N este numărul de mostre din setul de date
• n_F este numărul de caracteristici, cu n_F = 3585
• n_L este numărul de etichete de adevăr de bază, cu n_L = 1
• 4 este numărul de octeți necesari pentru tipul de date al caracteristicilor: float32

În plus, următoarele diagrame circulare oferă distribuția etichetelor atât pentru seturile de instruire, cât și pentru cele de testare, provocând dezechilibrul de clasă cu care se confruntă sarcina de detectare a malware PDF.

Distribuția se schimbă de la setul de antrenament la setul de test de o lună. Se aplică o împărțire în timp a setului de date în instruire și testare pentru a simula scenariul de implementare din viața reală și pentru a evita istoria temporală. Această strategie a permis, de asemenea, SophosAI să evalueze adevăratele capacități de generalizare ale modelului atunci când se confruntă cu atacuri PDF nevăzute anterior, de exemplu.

Experimente și rezultate

Pentru a începe experimentele, echipa SophosAI a antrenat un model de bază XGBoost cu parametri impliciti. Apoi au început să efectueze reglarea fină a hiperparametrului cu SageMaker folosind strategia Bayesiană, care este la fel de simplă ca și specificarea hiperparametri de reglat și intervalul dorit de valori, metrica de evaluare (ROC (Receiver Operating Characteristic) AUC în acest caz) și seturile de antrenament și validare. Pentru detectorul de malware PDF, SophosAI a prioritizat hiperparametrii, inclusiv numărul de runde de amplificare (num_round), adâncimea maximă a copacului (max_depth), rata de învățare (eta), și raportul de eșantionare a coloanelor la construirea arborilor (colsample_bytree). În cele din urmă, cei mai buni hiperparametri au fost obținuți și utilizați pentru a antrena un model pe setul de date complet și, în cele din urmă, au fost evaluați pe setul de testare holdout.

Următorul grafic arată metrica obiectivă (ROC AUC) față de cele 15 joburi de antrenament executate în cadrul jobului de reglare. Cei mai buni hiperparametri sunt cei corespunzători celui de-al nouălea job de antrenament.

La începutul experimentelor SophosAI pe SageMaker, o întrebare deosebit de importantă la care trebuie să răspundem a fost: ce tip de instanțe și câte dintre ele sunt necesare pentru a instrui XGBoost cu privire la datele disponibile? Acest lucru este crucial deoarece utilizarea unui număr sau tip greșit de instanță poate fi o pierdere de timp și bani; antrenamentul este obligat să eșueze din cauza lipsei de memorie sau, dacă utilizați prea multe instanțe prea mari, acest lucru poate deveni inutil de costisitor.

XGBoost este un algoritm legat de memorie (spre deosebire de legat de calcul). Deci, o instanță de calcul cu scop general (de exemplu, M5) este o alegere mai bună decât o instanță optimizată pentru calcul (de exemplu, C4). Pentru a lua o decizie în cunoștință de cauză, există un ghid SageMaker simplu pentru alegerea numărului de instanțe necesare pentru a rula antrenamentul pe setul de date complet:

Dimensiunea totală a datelor de antrenament × factor de siguranță^(*) < Număr de instanțe × Memoria totală a tipului de instanță

În acest caz: Dimensiunea totală a datelor de antrenament × factor de siguranță (12) = 12120 GB

Următorul tabel rezumă cerințele când tipul de instanță ales este ml.m5.24xlarge.

Dimensiune antrenament × factor de siguranță (12)	Memorie de instanță ml.m5.24xlarge	Număr minim de instanțe necesare pentru instruire
12120 GB	384 GB	32

^*Datorită naturii instruirii distribuite XGBoost, care necesită încărcarea întregului set de date de antrenament într-un obiect DMatrix înainte de antrenament și memorie suplimentară liberă, se recomandă un factor de siguranță de 10-12.

Pentru a arunca o privire mai atentă asupra utilizării memoriei pentru un antrenament complet SageMaker a XGBoost pe setul de date furnizat, oferim graficul corespunzător obținut din programul de instruire. Amazon CloudWatch monitorizarea. Pentru această lucrare de antrenament, au fost utilizate 40 ml.m5.24xinstanțe mari și utilizarea maximă a memoriei a atins aproximativ 62 %.

Costul de inginerie economisit prin integrarea unui serviciu ML gestionat precum SageMaker în conducta de date este de aproximativ 50%. Opțiunea de a utiliza Instanțele Spot pentru lucrări de instruire și de reglare a hiperparametrilor reduce costurile cu încă 63%.

Concluzie

Cu SageMaker, echipa SophosAI ar putea rezolva cu succes un proiect complex de înaltă prioritate prin construirea unui model XGBoost de detectare a malware PDF ușor, care este mult mai mic pe disc (de până la 25 de ori mai mic) și în memorie (de până la 5 ori mai mic) decât acesta. predecesorul detectorului. Este un detector de malware mic, dar puternic, cu ~0.99 AUC și o rată pozitivă adevărată de 0.99 și o rată fals pozitivă de . Acest model poate fi reantrenat rapid, iar performanța lui poate fi monitorizată cu ușurință în timp, deoarece durează mai puțin de 20 de minute pentru a-l antrena pe mai mult de 1 TB de date.

Puteți utiliza algoritmul încorporat SageMaker XGBoost pentru a construi modele cu datele dvs. tabelare la scară. În plus, puteți încerca și noii algoritmi încorporați Amazon SageMaker LightGBM, CatBoost, AutoGluon-Tabular și Tab Transformer, așa cum este descris în acest articol. blogul.

---

Despre autori

Salma Taoufiq este Senior Data Scientist la Sophos, care lucrează la intersecția dintre învățarea automată și securitatea cibernetică. Cu studii de licență în informatică, ea a absolvit Universitatea Central Europeană cu un MSc. în Matematică și aplicațiile sale. Când nu dezvoltă un detector de malware, Salma este un drumeț pasionat, călător și consumator de thrillere.

Harini Kannan este Data Scientist la SophosAI. Ea lucrează în știința datelor de securitate de aproximativ 4 ani. Anterior, ea a fost principalul Data Scientist la Capsule8, care a fost achiziționată de Sophos. Ea a susținut discursuri la CAMLIS, BlackHat (SUA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) și Data Connectors. Domeniile ei de cercetare includ detectarea atacurilor bazate pe hardware folosind contoare de performanță, analiza comportamentului utilizatorului, ML interpretabil și detectarea nesupravegheată a anomaliilor.

Hasan Poonawala este Senior AI/ML Specialist Solutions Architect la AWS, cu sediul în Londra, Marea Britanie. Hasan îi ajută pe clienți să proiecteze și să implementeze aplicații de învățare automată în producție pe AWS. Are peste 12 ani de experiență de lucru ca om de știință a datelor, practician de învățare automată și dezvoltator de software. În timpul liber, lui Hasan îi place să exploreze natura și să petreacă timp cu prietenii și familia.

Digant Patel este lider de asistență pentru întreprinderi la AWS. El lucrează cu clienții pentru a proiecta, implementa și opera în cloud la scară. Domeniile sale de interes sunt practicile MLOps și DevOps și modul în care acestea pot ajuta clienții în călătoria lor în cloud. În afara serviciului, îi place să fotografieze, să joace volei și să petreacă timpul cu prietenii și familia.