Comment Sophos forme un détecteur de malware PDF puissant et léger à très grande échelle avec Amazon SageMaker

Cet article est co-écrit par Salma Taoufiq et Harini Kannan de Sophos.

En tant que leader de la cybersécurité de nouvelle génération, Sophos s'efforce de protéger plus de 500,000 150 organisations et des millions de clients dans plus de XNUMX pays contre l'évolution des menaces. S'appuyant sur l'intelligence des menaces, l'apprentissage automatique (ML) et l'intelligence artificielle de Sophos X-Ops, Sophos propose un portefeuille large et varié de produits et services avancés pour sécuriser et défendre les utilisateurs, les réseaux et les terminaux contre le phishing, les ransomwares, les malwares et le large éventail de cyberattaques qui existent.

La Groupe Sophos Intelligence Artificielle (IA) (SophosAI) supervise le développement et la maintenance de la principale technologie de sécurité ML de Sophos. La sécurité est un problème de mégadonnées. Pour échapper à la détection, les cybercriminels élaborent constamment de nouvelles attaques. Cela se traduit par des ensembles de données colossaux sur les menaces avec lesquels le groupe doit travailler pour mieux défendre ses clients. Un exemple notable est la détection et l'élimination de fichiers astucieusement mêlés à des logiciels malveillants, où les ensembles de données sont en téraoctets.

Dans cet article, nous nous concentrons sur le système de détection des logiciels malveillants de Sophos pour le format de fichier PDF en particulier. Nous montrons comment SophosAI utilise Amazon Sage Maker formation distribuée avec des téraoctets de données pour former un puissant modèle léger XGBoost (Extreme Gradient Boosting). Cela permet à leur équipe d'itérer plus rapidement sur de grandes données de formation avec un réglage automatique des hyperparamètres et sans gérer l'infrastructure de formation sous-jacente.

La solution est actuellement intégrée de manière transparente dans le pipeline de formation en production et le modèle déployé sur des millions de terminaux utilisateurs via le Service de terminaux Sophos.

Contexte du cas d'utilisation

Que vous souhaitiez partager un contrat important ou préserver le design sophistiqué de votre CV, le format PDF est le choix le plus courant. Son utilisation généralisée et la perception générale que ces documents sont hermétiques et statiques ont endormi les utilisateurs dans un faux sentiment de sécurité. Le PDF est donc devenu un vecteur d'infection de choix dans l'arsenal des attaquants. Les actions malveillantes utilisant des fichiers PDF sont le plus souvent réalisées via l'intégration d'une charge utile JavaScript qui est exécutée par le lecteur PDF pour télécharger un virus à partir d'un URI, saboter la machine de l'utilisateur ou voler des informations sensibles.

Sophos détecte les fichiers PDF malveillants à différents stades d'une attaque à l'aide d'un ensemble de modèles déterministes et ML. Une de ces approches est illustrée dans le diagramme suivant, où le fichier PDF malveillant est envoyé par e-mail. Dès qu'une tentative de téléchargement est effectuée, il déclenche le script exécutable malveillant pour se connecter au serveur de commande et de contrôle de l'attaquant. Le détecteur PDF de SophosAI bloque la tentative de téléchargement après avoir détecté qu'il est malveillant.

D'autres moyens incluent le blocage des fichiers PDF dans le point de terminaison, l'envoi des fichiers malveillants dans un bac à sable (où ils sont notés à l'aide de plusieurs modèles), la soumission du fichier malveillant à une infrastructure de notation et la génération d'un rapport de sécurité, etc.

motivation

Pour créer un détecteur basé sur une arborescence capable de condamner les PDF malveillants avec une grande confiance, tout en permettant une faible consommation d'énergie informatique et des réponses d'inférence rapides, l'équipe SophosAI a trouvé que l'algorithme XGBoost était un candidat parfait pour la tâche. De telles pistes de recherche sont importantes pour Sophos pour deux raisons. Le déploiement de modèles puissants mais petits au niveau des points de terminaison des clients a un impact important sur les évaluations des produits de l'entreprise par les analystes. Il offre également, et plus important encore, une meilleure expérience utilisateur globale.

Défi technique

Parce que l'objectif était d'avoir un modèle avec une empreinte mémoire plus petite que leurs détecteurs de logiciels malveillants PDF existants (à la fois sur disque et en mémoire), SophosAI a transformé XGBoost, un algorithme de classification ayant fait ses preuves dans la production de modèles considérablement plus petits que les réseaux de neurones tout en réalisant des performances impressionnantes. performances sur des données tabulaires. Avant de s'aventurer dans la modélisation des expériences XGBoost, une considération importante était la taille même de l'ensemble de données. En effet, l'ensemble de données de base des fichiers PDF de Sophos est en téraoctets.

Par conséquent, le principal défi consistait à former le modèle avec un grand ensemble de données sans avoir à sous-échantillonner. Parce qu'il est crucial que le détecteur apprenne à repérer toutes les attaques basées sur PDF - même les attaques les plus inopinées et les plus nouvelles pour mieux défendre les clients Sophos - il est de la plus haute importance d'utiliser tous les divers ensembles de données disponibles.

Contrairement aux réseaux de neurones, où vous pouvez vous entraîner par lots, pour XGBoost, nous avons besoin de l'ensemble des données d'entraînement en mémoire. Le plus grand ensemble de données de formation pour ce projet dépasse 1 To, et il n'y a aucun moyen de former à une telle échelle sans utiliser les méthodologies d'un cadre de formation distribué.

Vue d'ensemble de la solution

SageMaker est un service ML entièrement géré fournissant divers outils pour créer, former, optimiser et déployer des modèles ML. La Bibliothèques d'algorithmes intégrées SageMaker se composent de 21 algorithmes ML populaires, y compris XGBoost. (Pour plus d'informations, voir Simplifiez l'apprentissage automatique avec XGBoost et Amazon SageMaker.) Avec l'algorithme intégré XGBoost, vous pouvez profiter de l'open-source Conteneur SageMaker XGBoost en spécifiant une version de framework supérieure à 1.0-1, qui a amélioré la flexibilité, l'évolutivité, l'extensibilité et la formation Spot gérée, et prend en charge les formats d'entrée comme Parquet, qui est le format utilisé pour le jeu de données PDF.

La principale raison pour laquelle SophosAI a choisi SageMaker est la possibilité de bénéficier de la formation distribuée entièrement gérée sur les instances de CPU multi-nœuds en spécifiant simplement plus d'une instance. SageMaker divise automatiquement les données entre les nœuds, agrège les résultats entre les nœuds homologues et génère un modèle unique. Les instances peuvent être des instances ponctuelles, réduisant ainsi considérablement les coûts de formation. Avec le algorithme intégré pour XGBoost, vous pouvez le faire sans aucun script personnalisé supplémentaire. Des versions distribuées de XGBoost existent également en open source, telles que XGBoost-Ray ainsi que XGBoost4J-Spark, mais leur utilisation nécessite la construction, la sécurisation, le réglage et l'autogestion de clusters de calcul distribué, ce qui représente un effort important en plus du développement scientifique.

En outre, Réglage automatique du modèle SageMaker, également connu sous le nom de réglage d'hyperparamètres, trouve la meilleure version d'un modèle en exécutant de nombreuses tâches d'entraînement avec des plages d'hyperparamètres que vous spécifiez. Il choisit ensuite les valeurs d'hyperparamètres qui donnent un modèle qui fonctionne le mieux, tel que mesuré par une métrique pour la tâche ML donnée.

Le diagramme suivant illustre l'architecture de la solution.

Il convient de noter que, lorsque SophosAI a commencé les expériences XGBoost avant de se tourner vers SageMaker, des tentatives ont été faites pour utiliser une grande mémoire Cloud de calcul élastique Amazon (Amazon EC2) (par exemple, r5a.24xlarge et x1.32xlarge) pour entraîner le modèle sur un échantillon de données aussi large que possible. Cependant, ces tentatives prenaient plus de 10 heures en moyenne et échouaient généralement en raison d'un manque de mémoire.

En revanche, en utilisant l'algorithme SageMaker XGBoost et un mécanisme de formation distribuée sans tracas, SophosAI pourrait former un modèle booster à grande échelle sur le colossal ensemble de données de formation PDF en l'espace de 20 minutes. L'équipe n'avait qu'à stocker les données sur Service de stockage simple Amazon (Amazon S3) en tant que fichiers Parquet de taille similaire, et choisissez un type d'instance EC2 et le nombre d'instances souhaité, et SageMaker gère l'infrastructure de cluster de calcul sous-jacente et distribue la formation entre plusieurs nœuds du cluster. Sous le capot, SageMaker divise les données entre les nœuds à l'aide de ShardedByS3Key pour répartir les objets de fichier de manière égale entre chaque instance et utilise l'implémentation XGBoost du Protocole Rabit (interface AllReduce et broadcast fiable) pour lancer le traitement distribué et communiquer entre les nœuds primaires et pairs. (Pour plus de détails sur l'agrégation d'histogrammes et la diffusion entre les nœuds, reportez-vous à XGBoost : un système de boost d'arborescence évolutif.)

Au-delà de la simple formation d'un modèle, avec SageMaker, Réglage des hyperparamètres XGBoost a également été rendu rapide et facile avec la possibilité d'exécuter différentes expériences simultanément pour affiner la meilleure combinaison d'hyperparamètres. Les hyperparamètres accordables comprennent à la fois des hyperparamètres spécifiques au rappel et spécifiques à la fonction objective. Deux stratégies de recherche sont proposés : aléatoire ou bayésien. La stratégie de recherche bayésienne s'est avérée utile car elle permet de trouver de meilleurs hyperparamètres qu'une simple recherche aléatoire, en moins d'itérations expérimentales.

Informations sur l'ensemble de données

La modélisation de la détection des logiciels malveillants PDF de SophosAI repose sur une variété de fonctionnalités telles que les histogrammes n-grammes et les fonctionnalités d'entropie d'octets (pour plus d'informations, reportez-vous à MEADE : Vers un moteur de détection de pièces jointes malveillantes). Les métadonnées et les fonctionnalités extraites des fichiers PDF collectés sont stockées dans un entrepôt de données distribué. Un ensemble de données de plus de 3,500 3 fonctionnalités est ensuite calculé, divisé en fonction du temps en ensembles de formation et de test et stocké par lots sous forme de fichiers Parquet dans Amazon SXNUMX pour être facilement accessible par SageMaker pour les tâches de formation.

Le tableau suivant fournit des informations sur les données de formation et de test.

Ensemble de données	Nombre d'échantillons	Nombre de limes parquet	Taille totale
Formation	70,391,634	5,500	~ 1010 Go
Teste	1,242,283	98	~ 18 Go

Les tailles de données ont été calculées selon la formule :

Taille des données = N × (n_F + n_L)×4

La formule a les paramètres suivants :

• N est le nombre d'échantillons dans l'ensemble de données
• n_F est le nombre de caractéristiques, avec n_F = 3585
• n_L est le nombre d'étiquettes de vérité terrain, avec n_L = 1
• 4 est le nombre d'octets nécessaires pour le type de données des fonctionnalités : float32

De plus, les diagrammes circulaires suivants fournissent la distribution des étiquettes des ensembles d'apprentissage et de test, évoquant le déséquilibre des classes rencontré dans la tâche de détection des logiciels malveillants PDF.

La distribution passe de l'ensemble d'apprentissage à l'ensemble de test d'un mois. Une division temporelle de l'ensemble de données en formation et en test est appliquée afin de simuler le scénario de déploiement réel et d'éviter l'espionnage temporel. Cette stratégie a également permis à SophosAI d'évaluer les véritables capacités de généralisation du modèle face à des attaques PDF inédites, par exemple.

Expériences et résultats

Pour lancer les expériences, l'équipe SophosAI a formé un modèle XGBoost de base avec des paramètres par défaut. Ensuite, ils ont commencé à effectuer un réglage fin des hyperparamètres avec SageMaker en utilisant la stratégie bayésienne, qui est aussi simple que de spécifier le hyperparamètres à régler et la plage de valeurs souhaitée, la métrique d'évaluation (ROC (Receiver Operating Characteristic) AUC dans ce cas) et les ensembles d'apprentissage et de validation. Pour le détecteur de logiciels malveillants PDF, SophosAI a donné la priorité aux hyperparamètres, y compris le nombre de tours de boosting (num_round), la profondeur maximale de l'arbre (max_depth), le taux d'apprentissage (eta), et le taux d'échantillonnage des colonnes lors de la construction des arbres (colsample_bytree). Finalement, les meilleurs hyperparamètres ont été obtenus et utilisés pour former un modèle sur l'ensemble de données complet, et finalement évalués sur l'ensemble de test d'exclusion.

Le graphique suivant montre la métrique objective (ROC AUC) par rapport aux 15 tâches de formation exécutées dans la tâche de réglage. Les meilleurs hyperparamètres sont ceux correspondant au neuvième travail d'apprentissage.

Au début des expérimentations de SophosAI sur SageMaker, une question particulièrement importante à laquelle répondre était : quel type d'instances et combien d'entre elles sont nécessaires pour entraîner XGBoost sur les données disponibles ? Ceci est crucial car l'utilisation d'un mauvais numéro ou d'un mauvais type d'instance peut être une perte de temps et d'argent ; la formation est vouée à l'échec en raison d'un manque de mémoire ou, si vous utilisez trop d'instances trop volumineuses, cela peut devenir inutilement coûteux.

XGBoost est un algorithme lié à la mémoire (par opposition à lié au calcul). Ainsi, une instance de calcul à usage général (par exemple, M5) est un meilleur choix qu'une instance optimisée pour le calcul (par exemple, C4). Pour prendre une décision éclairée, il existe une simple directive SageMaker pour choisir le nombre d'instances nécessaires pour exécuter la formation sur l'ensemble de données complet :

Taille totale des données d'entraînement × Facteur de sécurité^(*) < Nombre d'instances × Mémoire totale du type d'instance

Dans ce cas : Taille totale des données d'entraînement × Facteur de sécurité (12) = 12120 XNUMX Go

Le tableau suivant résume les exigences lorsque le type d'instance choisi est ml.m5.24xlarge.

Taille de l'entraînement × Facteur de sécurité (12)	Mémoire d'instance ml.m5.24xlarge	Nombre d'instances minimum requis pour la formation
12120 GB	384 GB	32

^*En raison de la nature de la formation distribuée XGBoost, qui nécessite que l'intégralité de l'ensemble de données de formation soit chargée dans un objet DMatrix avant la formation et de la mémoire libre supplémentaire, un facteur de sécurité de 10 à 12 est recommandé.

Pour examiner de plus près l'utilisation de la mémoire pour une formation SageMaker complète de XGBoost sur l'ensemble de données fourni, nous fournissons le graphique correspondant obtenu à partir de la formation Amazon Cloud Watch surveillance. Pour cette tâche d'entraînement, 40 instances ml.m5.24xlarge ont été utilisées et l'utilisation maximale de la mémoire a atteint environ 62 %.

Le coût d'ingénierie économisé en intégrant un service ML géré comme SageMaker dans le pipeline de données est d'environ 50 %. La possibilité d'utiliser des instances ponctuelles pour les tâches de formation et de réglage d'hyperparamètres a permis de réduire les coûts de 63 % supplémentaires.

Conclusion

Avec SageMaker, l'équipe SophosAI a pu résoudre avec succès un projet complexe hautement prioritaire en créant un modèle léger XGBoost de détection de logiciels malveillants PDF qui est beaucoup plus petit sur disque (jusqu'à 25 fois plus petit) et en mémoire (jusqu'à 5 fois plus petit) que son prédécesseur du détecteur. C'est un petit mais puissant détecteur de logiciels malveillants avec ~ 0.99 AUC et un taux de vrais positifs de 0.99 et un taux de faux positifs de . Ce modèle peut être rapidement réentrainé, et ses performances peuvent être facilement surveillées dans le temps, car il faut moins de 20 minutes pour l'entraîner sur plus de 1 To de données.

Vous pouvez tirer parti de l'algorithme intégré de SageMaker XGBoost pour créer des modèles avec vos données tabulaires à grande échelle. De plus, vous pouvez également essayer les nouveaux algorithmes Amazon SageMaker intégrés LightGBM, CatBoost, AutoGluon-Tabular et Tab Transformer, comme décrit dans ce blogue.

---

À propos des auteurs

Salma Taoufik est Senior Data Scientist chez Sophos, travaillant à l'intersection de l'apprentissage automatique et de la cybersécurité. Avec une formation de premier cycle en informatique, elle est diplômée de l'Université d'Europe centrale avec un MSc. en mathématiques et ses applications. Lorsqu'elle ne développe pas de détecteur de logiciels malveillants, Salma est une randonneuse passionnée, une voyageuse et une consommatrice de thrillers.

Harini Kannan est Data Scientist chez SophosAI. Elle travaille dans le domaine de la science des données de sécurité depuis environ 4 ans. Elle était auparavant principale scientifique des données chez Capsule8, qui a été rachetée par Sophos. Elle a donné des conférences à CAMLIS, BlackHat (USA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) et Data Connectors. Ses domaines de recherche incluent la détection d'attaques basées sur le matériel à l'aide de compteurs de performances, l'analyse du comportement des utilisateurs, le ML interprétable et la détection d'anomalies non supervisée.

Hassan Poonawala est un architecte de solutions spécialiste senior en IA/ML chez AWS, basé à Londres, au Royaume-Uni. Hasan aide ses clients à concevoir et à déployer des applications de machine learning en production sur AWS. Il a plus de 12 ans d'expérience professionnelle en tant que scientifique des données, praticien de l'apprentissage automatique et développeur de logiciels. Dans ses temps libres, Hasan aime explorer la nature et passer du temps avec ses amis et sa famille.

Digant Patel est responsable de l'assistance aux entreprises chez AWS. Il travaille avec les clients pour concevoir, déployer et exploiter le cloud à grande échelle. Ses domaines d'intérêt sont les pratiques MLOps et DevOps et comment cela peut aider les clients dans leur parcours vers le cloud. En dehors du travail, il aime la photographie, jouer au volleyball et passer du temps avec ses amis et sa famille.