Cómo Sophos entrena un potente y ligero detector de malware PDF a gran escala con Amazon SageMaker

Esta publicación es coautora de Salma Taoufiq y Harini Kannan de Sophos.

Como líder en ciberseguridad de última generación, Sophos se esfuerza por proteger a más de 500,000 150 organizaciones y millones de clientes en más de XNUMX países contra las amenazas en evolución. Con tecnología de inteligencia de amenazas, aprendizaje automático (ML) e inteligencia artificial de Sophos X-Ops, Sophos ofrece una cartera amplia y variada de productos y servicios avanzados para proteger y defender a los usuarios, las redes y los puntos finales contra phishing, ransomware, malware y la amplia gama de ciberataques que existen.

El Grupo de inteligencia artificial (IA) de Sophos (SophosAI) supervisa el desarrollo y mantenimiento de la principal tecnología de seguridad de ML de Sophos. La seguridad es un problema de big data. Para evadir la detección, los ciberdelincuentes crean constantemente nuevos ataques. Esto se traduce en conjuntos de datos de amenazas colosales con los que el grupo debe trabajar para defender mejor a los clientes. Un ejemplo notable es la detección y eliminación de archivos que fueron astutamente enlazados con malware, donde los conjuntos de datos están en terabytes.

En esta publicación, nos enfocamos específicamente en el sistema de detección de malware de Sophos para el formato de archivo PDF. Mostramos cómo SophosAI usa Amazon SageMaker entrenamiento distribuido con terabytes de datos para entrenar un modelo XGBoost (Extreme Gradient Boosting) ligero y potente. Esto permite a su equipo iterar sobre grandes datos de entrenamiento más rápido con el ajuste automático de hiperparámetros y sin administrar la infraestructura de entrenamiento subyacente.

Actualmente, la solución se integra a la perfección en la canalización de capacitación de producción y el modelo se implementa en millones de puntos finales de usuarios a través de la Servicio de punto final de Sophos.

Contexto de caso de uso

Ya sea que desee compartir un contrato importante o conservar el elegante diseño de su CV, el formato PDF es la opción más común. Su uso generalizado y la percepción general de que dichos documentos son herméticos y estáticos han llevado a los usuarios a una falsa sensación de seguridad. Por lo tanto, PDF se ha convertido en el vector de infección elegido en el arsenal de los atacantes. Las acciones maliciosas que usan archivos PDF se logran con mayor frecuencia mediante la incrustación de una carga útil de JavaScript que ejecuta el lector de PDF para descargar un virus de un URI, sabotear la máquina del usuario o robar información confidencial.

Sophos detecta archivos PDF maliciosos en varios puntos de un ataque mediante un conjunto de modelos deterministas y de aprendizaje automático. Uno de estos enfoques se ilustra en el siguiente diagrama, donde el archivo PDF malicioso se envía por correo electrónico. Tan pronto como se realiza un intento de descarga, activa el script ejecutable malicioso para conectarse al servidor de comando y control del atacante. El detector de PDF de SophosAI bloquea el intento de descarga después de detectar que es malicioso.

Otras formas incluyen bloquear los archivos PDF en el endpoint, enviar los archivos maliciosos a un espacio aislado (donde se puntúa utilizando varios modelos), enviar el archivo malicioso a una infraestructura de puntuación y generar un informe de seguridad, etc.

Motivación

Para construir un detector basado en árbol que pueda condenar archivos PDF maliciosos con alta confianza, al mismo tiempo que permite un bajo consumo de energía de cómputo de punto final y respuestas de inferencia rápidas, el equipo de SophosAI encontró que el algoritmo XGBoost era un candidato perfecto para la tarea. Estas vías de investigación son importantes para Sophos por dos razones. Tener modelos potentes pero pequeños desplegados a nivel de puntos finales de los clientes tiene un gran impacto en las revisiones de productos de la empresa por parte de los analistas. También, y lo que es más importante, proporciona una mejor experiencia de usuario en general.

Desafío técnico

Debido a que el objetivo era tener un modelo con una huella de memoria más pequeña que sus detectores de malware PDF existentes (tanto en el disco como en la memoria), SophosAI convirtió XGBoost, un algoritmo de clasificación con un historial comprobado de producir modelos drásticamente más pequeños que las redes neuronales mientras logra impresionantes rendimiento en datos tabulares. Antes de aventurarse en el modelado de experimentos XGBoost, una consideración importante era el tamaño total del conjunto de datos. De hecho, el conjunto de datos principal de archivos PDF de Sophos está en terabytes.

Por lo tanto, el principal desafío fue entrenar el modelo con un gran conjunto de datos sin tener que reducir la muestra. Debido a que es crucial que el detector aprenda a detectar cualquier ataque basado en PDF, incluso los ataques de aguja en el pajar y los completamente nuevos para defender mejor a los clientes de Sophos, es de suma importancia utilizar todos los diversos conjuntos de datos disponibles.

A diferencia de las redes neuronales, en las que se puede entrenar por lotes, para XGBoost necesitamos todo el conjunto de datos de entrenamiento en la memoria. El conjunto de datos de capacitación más grande para este proyecto tiene más de 1 TB, y no hay forma de capacitar a tal escala sin utilizar las metodologías de un marco de capacitación distribuido.

Resumen de la solución

SageMaker es un servicio de ML completamente administrado que proporciona varias herramientas para crear, entrenar, optimizar e implementar modelos de ML. los Bibliotecas de algoritmos integradas de SageMaker consta de 21 algoritmos ML populares, incluido XGBoost. (Para más información, ver Simplifique el aprendizaje automático con XGBoost y Amazon SageMaker.) Con el algoritmo integrado XGBoost, puede aprovechar el código abierto Contenedor SageMaker XGBoost especificando una versión de marco superior a 1.0-1, que ha mejorado la flexibilidad, escalabilidad, extensibilidad y Managed Spot Training, y admite formatos de entrada como Parquet, que es el formato utilizado para el conjunto de datos PDF.

La razón principal por la que SophosAI eligió SageMaker es la capacidad de beneficiarse del entrenamiento distribuido completamente administrado en instancias de CPU de múltiples nodos simplemente especificando más de una instancia. SageMaker divide automáticamente los datos entre los nodos, agrega los resultados entre los nodos del mismo nivel y genera un solo modelo. Las instancias pueden ser instancias puntuales, lo que reduce significativamente los costos de capacitación. Con el algoritmo incorporado para XGBoost, puede hacerlo sin ningún script personalizado adicional. También existen versiones distribuidas de XGBoost como código abierto, como Rayo XGBoost y XGBoost4J-Chispa, pero su uso requiere construir, asegurar, ajustar y autogestionar clústeres de computación distribuida, lo que representa un importante esfuerzo adicional al desarrollo científico.

Además, Ajuste automático del modelo SageMaker, también conocido como ajuste de hiperparámetros, encuentra la mejor versión de un modelo ejecutando muchos trabajos de entrenamiento con rangos de hiperparámetros que usted especifica. Luego elige los valores de hiperparámetros que dan como resultado un modelo que funciona mejor, según lo medido por una métrica para la tarea de ML dada.

El siguiente diagrama ilustra la arquitectura de la solución.

Vale la pena señalar que, cuando SophosAI comenzó los experimentos con XGBoost antes de recurrir a SageMaker, se intentaron utilizar memorias de gran tamaño. Nube informática elástica de Amazon (Amazon EC2) (por ejemplo, r5a.24xlarge y x1.32xlarge) para entrenar el modelo en una muestra de datos tan grande como sea posible. Sin embargo, estos intentos demoraron más de 10 horas en promedio y generalmente fallaron debido a que se quedaron sin memoria.

Por el contrario, al usar el algoritmo SageMaker XGBoost y un mecanismo de entrenamiento distribuido sin problemas, SophosAI pudo entrenar un modelo de refuerzo a escala en el colosal conjunto de datos de entrenamiento de PDF en cuestión de 20 minutos. El equipo solo tuvo que almacenar los datos en Servicio de almacenamiento simple de Amazon (Amazon S3) como archivos Parquet de tamaño similar y elija un tipo de instancia EC2 y la cantidad deseada de instancias, y SageMaker administró la infraestructura del clúster de cómputo subyacente y distribuyó la capacitación entre múltiples nodos del clúster. Bajo el capó, SageMaker divide los datos entre nodos mediante ShardedByS3Key para distribuir los objetos de archivo por igual entre cada instancia y utiliza la implementación XGBoost de la Protocolo de conejo (AllReduce confiable e interfaz de transmisión) para iniciar el procesamiento distribuido y comunicarse entre los nodos principales y los pares. (Para obtener más detalles sobre la agregación de histogramas y la transmisión entre nodos, consulte XGBoost: un sistema escalable de refuerzo de árboles.)

Más allá de entrenar un modelo, con SageMaker, Ajuste de hiperparámetro XGBoost también se hizo rápido y fácil con la capacidad de ejecutar diferentes experimentos simultáneamente para ajustar la mejor combinación de hiperparámetros. Los hiperparámetros sintonizables incluyen tanto hiperparámetros específicos del refuerzo como específicos de la función objetiva. Dos estrategias de búsqueda se ofrecen: aleatorio o bayesiano. La estrategia de búsqueda bayesiana ha demostrado ser valiosa porque ayuda a encontrar mejores hiperparámetros que una mera búsqueda aleatoria, en menos iteraciones experimentales.

Información del conjunto de datos

El modelo de detección de malware en PDF de SophosAI se basa en una variedad de funciones, como histogramas de n-gramas y funciones de entropía de bytes (para obtener más información, consulte MEADE: Hacia un motor de detección de archivos adjuntos de correo electrónico malicioso). Los metadatos y las funciones extraídas de los archivos PDF recopilados se almacenan en un almacén de datos distribuido. A continuación, se calcula un conjunto de datos de más de 3,500 características, se divide aún más según el tiempo en conjuntos de entrenamiento y prueba y se almacena en lotes como archivos Parquet en Amazon S3 para que SageMaker pueda acceder fácilmente a ellos para trabajos de entrenamiento.

La siguiente tabla proporciona información sobre los datos de entrenamiento y prueba.

Conjunto de datos	Número de muestras	Número de archivos de parquet	Tamaño total
Formación	70,391,634	5,500	~ 1010 GB
Probar	1,242,283	98	~ 18 GB

Los tamaños de los datos se han calculado siguiendo la fórmula:

Tamaño de datos = N × (n_F + norte_L) × 4

La fórmula tiene los siguientes parámetros:

• N es el número de muestras en el conjunto de datos
• n_F es el número de características, con n_F = 3585
• n_L es el número de etiquetas de verdad del terreno, con n_L = 1
• 4 es el número de bytes necesarios para el tipo de datos de las funciones: float32

Además, los siguientes gráficos circulares proporcionan la distribución de etiquetas de los conjuntos de entrenamiento y prueba, lo que revela el desequilibrio de clases que se enfrenta en la tarea de detección de malware en PDF.

La distribución cambia del conjunto de entrenamiento al conjunto de prueba de un mes. Se aplica una división basada en el tiempo del conjunto de datos en entrenamiento y prueba para simular el escenario de implementación de la vida real y evitar la intromisión temporal. Esta estrategia también permitió a SophosAI evaluar las verdaderas capacidades de generalización del modelo cuando se enfrentó a nuevos ataques de PDF nunca antes vistos, por ejemplo.

Experimentos y resultados

Para iniciar los experimentos, el equipo de SophosAI entrenó un modelo XGBoost de referencia con parámetros predeterminados. Luego, comenzaron a realizar un ajuste fino de hiperparámetros con SageMaker utilizando la estrategia bayesiana, que es tan simple como especificar el hiperparámetros a sintonizar y el rango de valores deseado, la métrica de evaluación (ROC (Receiver Operating Characteristic) AUC en este caso) y los conjuntos de entrenamiento y validación. Para el detector de malware PDF, SophosAI priorizó los hiperparámetros, incluido el número de rondas de refuerzo (num_round), la profundidad máxima del árbol (max_depth), la tasa de aprendizaje (eta), y la razón de muestreo de las columnas cuando se construyen árboles (colsample_bytree). Finalmente, se obtuvieron los mejores hiperparámetros y se usaron para entrenar un modelo en el conjunto de datos completo y, finalmente, se evaluaron en el conjunto de prueba de reserva.

El siguiente gráfico muestra la métrica objetiva (ROC AUC) frente a los 15 trabajos de entrenamiento que se ejecutan dentro del trabajo de ajuste. Los mejores hiperparámetros son los correspondientes al noveno trabajo de entrenamiento.

Al comienzo de los experimentos de SophosAI en SageMaker, una pregunta especialmente importante que debía responderse era: ¿qué tipo de instancias y cuántas se necesitan para entrenar XGBoost con los datos disponibles? Esto es crucial porque usar el número o tipo de instancia incorrecto puede ser una pérdida de tiempo y dinero; el entrenamiento está destinado a fallar debido a que se queda sin memoria o, si se usan demasiadas instancias demasiado grandes, esto puede volverse innecesariamente costoso.

XGBoost es un algoritmo vinculado a la memoria (a diferencia del vinculado a la computación). Por lo tanto, una instancia informática de propósito general (por ejemplo, M5) es una mejor opción que una instancia optimizada para computación (por ejemplo, C4). Para tomar una decisión informada, existe una guía simple de SageMaker para elegir la cantidad de instancias requeridas para ejecutar el entrenamiento en el conjunto de datos completo:

Tamaño total de datos de entrenamiento × factor de seguridad^(*) < Recuento de instancias × Memoria total del tipo de instancia

En este caso: Tamaño total de datos de entrenamiento × Factor de seguridad (12) = 12120 GB

La siguiente tabla resume los requisitos cuando el tipo de instancia elegido es ml.m5.24xlarge.

Tamaño del entrenamiento × factor de seguridad (12)	Memoria de instancia ml.m5.24xlarge	Recuento mínimo de instancias requerido para el entrenamiento
12120 GB	384 GB	32

^*Debido a la naturaleza del entrenamiento distribuido de XGBoost, que requiere que todo el conjunto de datos de entrenamiento se cargue en un objeto DMatrix antes del entrenamiento y memoria libre adicional, se recomienda un factor de seguridad de 10 a 12.

Para observar más de cerca la utilización de la memoria para un entrenamiento completo de SageMaker de XGBoost en el conjunto de datos proporcionado, proporcionamos el gráfico correspondiente obtenido del entrenamiento Reloj en la nube de Amazon vigilancia. Para este trabajo de entrenamiento, se usaron 40 instancias ml.m5.24xlarge y la utilización máxima de la memoria alcanzó alrededor del 62 %.

El costo de ingeniería ahorrado al integrar un servicio de ML administrado como SageMaker en la canalización de datos es de alrededor del 50 %. La opción de usar instancias de spot para trabajos de entrenamiento y ajuste de hiperparámetros redujo los costos en un 63 % adicional.

Conclusión

Con SageMaker, el equipo de SophosAI pudo resolver con éxito un proyecto complejo de alta prioridad mediante la creación de un modelo XGBoost ligero de detección de malware en PDF que es mucho más pequeño en el disco (hasta 25 veces más pequeño) y en la memoria (hasta 5 veces más pequeño) que su antecesor del detector. Es un detector de malware pequeño pero poderoso con ~0.99 AUC y una tasa de verdaderos positivos de 0.99 y una tasa de falsos positivos de . Este modelo se puede volver a entrenar rápidamente y su rendimiento se puede monitorear fácilmente a lo largo del tiempo, ya que lleva menos de 20 minutos entrenarlo con más de 1 TB de datos.

Puede aprovechar el algoritmo integrado de SageMaker XGBoost para construir modelos con sus datos tabulares a escala. Además, también puede probar los nuevos algoritmos integrados de Amazon SageMaker LightGBM, CatBoost, AutoGluon-Tabular y Tab Transformer como se describe en este blog.

---

Sobre los autores

Salma Taoufiq es científico de datos sénior en Sophos y trabaja en la intersección del aprendizaje automático y la ciberseguridad. Con una formación universitaria en ciencias de la computación, se graduó de la Universidad de Europa Central con una maestría. en Matemáticas y sus Aplicaciones. Cuando no está desarrollando un detector de malware, Salma es una ávida senderista, viajera y consumidora de novelas de suspenso.

harini kannan es científico de datos en SophosAI. Ella ha estado en ciencia de datos de seguridad por ~4 años. Anteriormente fue científica principal de datos en Capsule8, que fue adquirida por Sophos. Ha dado charlas en CAMLIS, BlackHat (EE. UU.), Open Data Science Conference (Este), Data Science Salon, PyData (Boston) y Data Connectors. Sus áreas de investigación incluyen la detección de ataques basados ​​en hardware mediante contadores de rendimiento, análisis del comportamiento del usuario, aprendizaje automático interpretable y detección de anomalías no supervisadas.

Hasan Poonawala es un arquitecto senior de soluciones especializado en IA / ML en AWS, con sede en Londres, Reino Unido. Hasan ayuda a los clientes a diseñar e implementar aplicaciones de aprendizaje automático en producción en AWS. Tiene más de 12 años de experiencia laboral como científico de datos, practicante de aprendizaje automático y desarrollador de software. En su tiempo libre, a Hasan le encanta explorar la naturaleza y pasar tiempo con amigos y familiares.

Patel Digno es líder de soporte empresarial en AWS. Trabaja con los clientes para diseñar, implementar y operar en la nube a escala. Sus áreas de interés son las prácticas de MLOps y DevOps y cómo pueden ayudar a los clientes en su viaje a la nube. Fuera del trabajo, disfruta de la fotografía, jugar voleibol y pasar tiempo con amigos y familiares.