In che modo Sophos addestra un rilevatore di malware PDF potente e leggero su scala ultra con Amazon SageMaker

Questo post è co-autore di Salma Taoufiq e Harini Kannan di Sophos.

In qualità di leader nella sicurezza informatica di nuova generazione, Sophos si impegna a proteggere più di 500,000 organizzazioni e milioni di clienti in oltre 150 paesi dalle minacce in evoluzione. Basato su intelligence sulle minacce, machine learning (ML) e intelligenza artificiale di Sophos X-Ops, Sophos offre un portafoglio ampio e variegato di prodotti e servizi avanzati per proteggere e difendere utenti, reti ed endpoint da phishing, ransomware, malware e la vasta gamma di attacchi informatici là fuori.

I Gruppo Sophos Intelligenza Artificiale (AI). (SophosAI) supervisiona lo sviluppo e la manutenzione della principale tecnologia di sicurezza ML di Sophos. La sicurezza è un problema di big data. Per eludere il rilevamento, i criminali informatici creano costantemente nuovi attacchi. Ciò si traduce in enormi set di dati sulle minacce con cui il gruppo deve lavorare per difendere al meglio i clienti. Un esempio degno di nota è il rilevamento e l'eliminazione di file che sono stati abilmente intrecciati con malware, in cui i set di dati sono in terabyte.

In questo post, ci concentreremo specificamente sul sistema di rilevamento del malware di Sophos per il formato di file PDF. Mostriamo come utilizza SophosAI Amazon Sage Maker training distribuito con terabyte di dati per addestrare un modello XGBoost (Extreme Gradient Boosting) leggero e potente. Ciò consente al proprio team di eseguire l'iterazione su dati di addestramento di grandi dimensioni più velocemente con l'ottimizzazione automatica degli iperparametri e senza gestire l'infrastruttura di addestramento sottostante.

La soluzione è attualmente perfettamente integrata nella pipeline di formazione della produzione e il modello è distribuito su milioni di endpoint utente tramite il Servizio endpoint Sophos.

Usa il contesto del caso

Sia che tu voglia condividere un contratto importante o preservare il design elegante del tuo CV, il formato PDF è la scelta più comune. Il suo uso diffuso e la percezione generale che tali documenti siano ermetici e statici hanno cullato gli utenti in un falso senso di sicurezza. Il PDF è quindi diventato un vettore di infezione preferito nell'arsenale degli aggressori. Le azioni dannose che utilizzano i PDF vengono spesso raggiunte incorporando un payload JavaScript eseguito dal lettore PDF per scaricare un virus da un URI, sabotare la macchina dell'utente o rubare informazioni riservate.

Sophos rileva i file PDF dannosi in vari punti di un attacco utilizzando un insieme di modelli deterministici e ML. Uno di questi approcci è illustrato nel diagramma seguente, in cui il file PDF dannoso viene recapitato tramite e-mail. Non appena viene effettuato un tentativo di download, viene attivato lo script eseguibile dannoso per la connessione al server Command and Control dell'attaccante. Il rilevatore PDF di SophosAI blocca il tentativo di download dopo aver rilevato che è dannoso.

Altri modi includono il blocco dei file PDF nell'endpoint, l'invio dei file dannosi a una sandbox (dove viene valutato utilizzando più modelli), l'invio del file dannoso a un'infrastruttura di punteggio e la generazione di un rapporto sulla sicurezza e così via.

Motivazione

Per creare un rilevatore ad albero in grado di condannare PDF dannosi con elevata sicurezza, consentendo al contempo un basso consumo energetico di elaborazione degli endpoint e risposte rapide di inferenza, il team di SophosAI ha riscontrato che l'algoritmo XGBoost è un candidato perfetto per l'attività. Tali vie di ricerca sono importanti per Sophos per due ragioni. Avere modelli potenti ma piccoli implementati a livello di endpoint dei clienti ha un forte impatto sulle recensioni dei prodotti dell'azienda da parte degli analisti. Inoltre, e soprattutto, fornisce una migliore esperienza utente in generale.

Sfida tecnica

Poiché l'obiettivo era quello di avere un modello con un footprint di memoria inferiore rispetto ai rilevatori di malware PDF esistenti (sia su disco che in memoria), SophosAI ha trasformato XGBoost, un algoritmo di classificazione con una comprovata esperienza nella produzione di modelli drasticamente più piccoli rispetto alle reti neurali ottenendo risultati impressionanti performance su dati tabulari. Prima di avventurarsi nella modellazione degli esperimenti XGBoost, una considerazione importante era l'enorme dimensione del set di dati. In effetti, il dataset di base dei file PDF di Sophos è in terabyte.

Pertanto, la sfida principale è stata addestrare il modello con un set di dati di grandi dimensioni senza dover eseguire il downsampling. Poiché è fondamentale che il rilevatore impari a individuare eventuali attacchi basati su PDF, anche quelli più pericolosi e completamente nuovi per difendere al meglio i clienti Sophos, è della massima importanza utilizzare tutti i diversi set di dati disponibili.

A differenza delle reti neurali, in cui è possibile eseguire l'allenamento in batch, per XGBoost è necessario l'intero set di dati di addestramento in memoria. Il più grande set di dati di formazione per questo progetto è di oltre 1 TB e non c'è modo di allenarsi su tale scala senza utilizzare le metodologie di un framework di formazione distribuito.

Panoramica della soluzione

SageMaker è un servizio ML completamente gestito che fornisce vari strumenti per creare, addestrare, ottimizzare e distribuire modelli ML. Il Librerie di algoritmi integrate di SageMaker sono costituiti da 21 popolari algoritmi ML, incluso XGBoost. (Per ulteriori informazioni, vedere Semplifica l'apprendimento automatico con XGBoost e Amazon SageMaker.) Con l'algoritmo integrato XGBoost, puoi sfruttare l'open source Contenitore SageMaker XGBoost specificando una versione del framework superiore alla 1.0-1, che offre maggiore flessibilità, scalabilità, estensibilità e Managed Spot Training e supporta formati di input come Parquet, che è il formato utilizzato per il set di dati PDF.

Il motivo principale per cui SophosAI ha scelto SageMaker è la possibilità di beneficiare della formazione distribuita completamente gestita su istanze CPU multi-nodo semplicemente specificando più di un'istanza. SageMaker suddivide automaticamente i dati tra nodi, aggrega i risultati tra nodi peer e genera un unico modello. Le istanze possono essere istanze Spot, riducendo così notevolmente i costi di formazione. Con il algoritmo integrato per XGBoost, puoi farlo senza alcun script personalizzato aggiuntivo. Le versioni distribuite di XGBoost esistono anche come open source, ad esempio XGBoost-Ray ed XGBoost4J-Spark, ma il loro utilizzo richiede la creazione, la protezione, la messa a punto e l'autogestione di cluster di calcolo distribuito, il che rappresenta uno sforzo significativo aggiuntivo allo sviluppo scientifico.

Inoltre, Ottimizzazione automatica del modello di SageMaker, noto anche come ottimizzazione degli iperparametri, trova la versione migliore di un modello eseguendo molti processi di addestramento con intervalli di iperparametri specificati dall'utente. Quindi sceglie i valori dell'iperparametro che si traducono in un modello con le prestazioni migliori, come misurato da una metrica per l'attività ML specificata.

Il diagramma seguente illustra l'architettura della soluzione.

Vale la pena notare che, quando SophosAI ha avviato gli esperimenti XGBoost prima di passare a SageMaker, sono stati fatti tentativi per utilizzare la memoria di grandi dimensioni Cloud di calcolo elastico di Amazon (Amazon EC2) (ad esempio, r5a.24xlarge e x1.32xlarge) per addestrare il modello su un campione di dati il ​​più grande possibile. Tuttavia, questi tentativi hanno richiesto in media più di 10 ore e di solito non sono riusciti a causa dell'esaurimento della memoria.

Al contrario, utilizzando l'algoritmo SageMaker XGBoost e un meccanismo di addestramento distribuito senza problemi, SophosAI potrebbe addestrare un modello booster su larga scala sul colossale set di dati di addestramento PDF nel giro di 20 minuti. Il team doveva solo memorizzare i dati Servizio di archiviazione semplice Amazon (Amazon S3) come file Parquet di dimensioni simili e scegliere un tipo di istanza EC2 e il numero di istanze desiderato e SageMaker ha gestito l'infrastruttura del cluster di calcolo sottostante e ha distribuito la formazione tra più nodi del cluster. Sotto il cofano, SageMaker divide i dati tra i nodi utilizzando ShardedByS3Key per distribuire equamente gli oggetti file tra ciascuna istanza e utilizza l'implementazione XGBoost del Protocollo Rabit (affidabile interfaccia AllReduce e broadcast) per avviare l'elaborazione distribuita e comunicare tra nodi primari e peer. (Per maggiori dettagli sull'aggregazione dell'istogramma e sulla trasmissione tra i nodi, fare riferimento a XGBoost: un sistema di potenziamento dell'albero scalabile.)

Oltre alla semplice formazione di un modello, con SageMaker, Ottimizzazione dell'iperparametro XGBoost è stato inoltre reso semplice e veloce grazie alla possibilità di eseguire diversi esperimenti contemporaneamente per mettere a punto la migliore combinazione di iperparametri. Gli iperparametri sintonizzabili includono iperparametri specifici del booster e obiettivi specifici della funzione. Due strategie di ricerca sono offerti: casuali o bayesiani. La strategia di ricerca bayesiana si è rivelata preziosa perché aiuta a trovare iperparametri migliori rispetto a una semplice ricerca casuale, in un minor numero di iterazioni sperimentali.

Informazioni sul set di dati

La modellazione del rilevamento del malware PDF di SophosAI si basa su una varietà di funzionalità come gli istogrammi di n-grammi e le funzionalità di entropia dei byte (per ulteriori informazioni, fare riferimento a MEADE: Verso un motore di rilevamento di allegati e-mail dannosi). I metadati e le funzionalità estratti dai file PDF raccolti vengono archiviati in un data warehouse distribuito. Viene quindi calcolato un set di dati di oltre 3,500 funzionalità, ulteriormente suddiviso in base al tempo in set di addestramento e test e archiviato in batch come file Parquet in Amazon S3 per essere facilmente accessibile da SageMaker per i lavori di formazione.

La tabella seguente fornisce informazioni sui dati di addestramento e test.

dataset	Numero di campioni	Numero di file per parquet	Dimensione totale
Training	70,391,634	5,500	~ 1010 GB
Test	1,242,283	98	~ 18 GB

Le dimensioni dei dati sono state calcolate secondo la formula:

Dimensione dei dati = N × (n_F + n_L) × 4

La formula ha i seguenti parametri:

• N è il numero di campioni nel set di dati
• n_F è il numero di funzioni, con n_F = 3585
• n_L è il numero di etichette di verità di base, con n_L = 1
• 4 è il numero di byte necessari per il tipo di dati delle caratteristiche: float32

Inoltre, i seguenti grafici a torta forniscono la distribuzione delle etichette dei set di addestramento e test, provocando lo squilibrio di classi riscontrato nell'attività di rilevamento del malware PDF.

La distribuzione si sposta dal set di formazione al set di test di un mese. Viene applicata una suddivisione del set di dati basata sul tempo in formazione e test per simulare lo scenario di distribuzione della vita reale ed evitare lo snooping temporale. Questa strategia ha inoltre consentito a SophosAI di valutare le reali capacità di generalizzazione del modello, ad esempio di fronte a nuovi attacchi PDF mai visti prima.

Esperimenti e risultati

Per avviare gli esperimenti, il team di SophosAI ha addestrato un modello XGBoost di base con parametri predefiniti. Quindi hanno iniziato a eseguire la messa a punto degli iperparametri con SageMaker utilizzando la strategia bayesiana, che è semplice come specificare il iperparametri da regolare e l'intervallo di valori desiderato, la metrica di valutazione (in questo caso ROC (Receiver Operating Characteristic) AUC) e i set di training e validazione. Per il rilevatore di malware PDF, SophosAI ha dato priorità agli iperparametri, incluso il numero di round di potenziamento (num_round), la profondità massima dell'albero (max_depth), il tasso di apprendimento (eta) e il rapporto di campionamento delle colonne durante la costruzione di alberi (colsample_bytree). Alla fine, i migliori iperparametri sono stati ottenuti e utilizzati per addestrare un modello sull'intero set di dati e infine valutati sul set di test di tenuta.

Il grafico seguente mostra la metrica obiettivo (ROC AUC) rispetto ai 15 processi di formazione eseguiti all'interno del processo di ottimizzazione. I migliori iperparametri sono quelli corrispondenti al nono lavoro di addestramento.

All'inizio degli esperimenti di SophosAI su SageMaker, una domanda particolarmente importante a cui rispondere era: che tipo di istanze e quante di esse sono necessarie per addestrare XGBoost sui dati a disposizione? Questo è fondamentale perché l'utilizzo del numero o del tipo di istanza sbagliato può essere una perdita di tempo e denaro; l'addestramento è destinato a fallire a causa dell'esaurimento della memoria o, se si utilizzano troppe istanze troppo grandi, questo può diventare inutilmente costoso.

XGBoost è un algoritmo legato alla memoria (al contrario di quello legato al calcolo). Pertanto, un'istanza di calcolo per uso generico (ad esempio, M5) è una scelta migliore rispetto a un'istanza di calcolo ottimizzata (ad esempio, C4). Per prendere una decisione informata, esiste una semplice linea guida SageMaker per selezionare il numero di istanze necessarie per eseguire la formazione sull'intero set di dati:

Dimensione totale dei dati di allenamento × fattore di sicurezza^(*) < Conteggio istanza × Memoria totale del tipo di istanza

In questo caso: Dimensione totale dei dati di allenamento × Fattore di sicurezza (12) = 12120 GB

La tabella seguente riassume i requisiti quando il tipo di istanza scelto è ml.m5.24xlarge.

Dimensione dell'allenamento × Fattore di sicurezza (12)	Memoria di istanza ml.m5.24xlarge	Numero minimo di istanze richiesto per la formazione
12120 GB	384 GB	32

^*A causa della natura dell'addestramento distribuito XGBoost, che richiede il caricamento dell'intero set di dati di addestramento in un oggetto DMatrix prima dell'addestramento e della memoria libera aggiuntiva, si consiglia un fattore di sicurezza di 10–12.

Per dare un'occhiata più da vicino all'utilizzo della memoria per un addestramento completo di SageMaker di XGBoost sul set di dati fornito, forniamo il grafico corrispondente ottenuto dall'addestramento Amazon Cloud Watch monitoraggio. Per questo lavoro di formazione sono state utilizzate istanze di 40 ml.m5.24xlarge e l'utilizzo massimo della memoria ha raggiunto circa il 62%.

Il costo di progettazione risparmiato dall'integrazione di un servizio ML gestito come SageMaker nella pipeline di dati è di circa il 50%. L'opzione di utilizzare le istanze Spot per i lavori di addestramento e ottimizzazione degli iperparametri riduce i costi di un ulteriore 63%.

Conclusione

Con SageMaker, il team di SophosAI ha potuto risolvere con successo un progetto complesso ad alta priorità creando un modello XGBoost di rilevamento malware PDF leggero che è molto più piccolo su disco (fino a 25 volte più piccolo) e in memoria (fino a 5 volte più piccolo) del suo predecessore del rivelatore. È un piccolo ma potente rilevatore di malware con ~0.99 AUC e un tasso di veri positivi di 0.99 e un tasso di falsi positivi di . Questo modello può essere riqualificato rapidamente e le sue prestazioni possono essere facilmente monitorate nel tempo, poiché sono necessari meno di 20 minuti per addestrarlo su più di 1 TB di dati.

Puoi sfruttare l'algoritmo integrato di SageMaker XGBoost per costruire modelli con i tuoi dati tabulari in scala. Inoltre, puoi anche provare i nuovi algoritmi integrati di Amazon SageMaker LightGBM, CatBoost, AutoGluon-Tabular e Tab Transformer, come descritto in questo blog.

---

Circa gli autori

Salma Taufiq è un Senior Data Scientist presso Sophos, che lavora all'intersezione tra apprendimento automatico e sicurezza informatica. Con un background universitario in informatica, si è laureata presso la Central European University con un Master. in Matematica e sue applicazioni. Quando non sviluppa un rilevatore di malware, Salma è un'appassionata escursionista, viaggiatrice e consumatrice di thriller.

Harini Kannan è un Data Scientist presso SophosAI. Si occupa di sicurezza dei dati da circa 4 anni. In precedenza era la Principal Data Scientist di Capsule8, che è stata acquisita da Sophos. Ha tenuto conferenze a CAMLIS, BlackHat (USA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) e Data Connectors. Le sue aree di ricerca includono il rilevamento di attacchi basati su hardware utilizzando contatori di prestazioni, analisi del comportamento degli utenti, ML interpretabile e rilevamento di anomalie senza supervisione.

Hasan Poonawala è un Senior AI/ML Specialist Solutions Architect presso AWS, con sede a Londra, Regno Unito. Hasan aiuta i clienti a progettare e distribuire applicazioni di machine learning in produzione su AWS. Ha oltre 12 anni di esperienza lavorativa come data scientist, professionista dell'apprendimento automatico e sviluppatore di software. Nel suo tempo libero, Hasan ama esplorare la natura e trascorrere del tempo con amici e familiari.

Digan Patel è un Enterprise Support Lead presso AWS. Lavora con i clienti per progettare, distribuire e operare nel cloud su larga scala. Le sue aree di interesse sono le pratiche MLOps e DevOps e come possono aiutare i clienti nel loro percorso verso il cloud. Al di fuori del lavoro, gli piace fotografare, giocare a pallavolo e passare il tempo con gli amici e la famiglia.