Como a Sophos treina um detector de malware de PDF poderoso e leve em ultraescala com o Amazon SageMaker

Este post é co-autoria de Salma Taoufiq e Harini Kannan da Sophos.

Como líder em segurança cibernética de última geração, Sophos se esforça para proteger mais de 500,000 organizações e milhões de clientes em mais de 150 países contra ameaças em evolução. Alimentado por inteligência de ameaças, aprendizado de máquina (ML) e inteligência artificial da Sophos X-Ops, a Sophos oferece um portfólio amplo e variado de produtos e serviços avançados para proteger e defender usuários, redes e endpoints contra phishing, ransomware, malware e a ampla gama de ataques cibernéticos que existem.

A Grupo Sophos Inteligência Artificial (IA) (SophosAI) supervisiona o desenvolvimento e a manutenção da principal tecnologia de segurança de ML da Sophos. A segurança é um problema de big data. Para evitar a detecção, os cibercriminosos estão constantemente criando novos ataques. Isso se traduz em conjuntos de dados de ameaças colossais com os quais o grupo deve trabalhar para melhor defender os clientes. Um exemplo notável é a detecção e eliminação de arquivos que foram habilmente atados com malware, onde os conjuntos de dados estão em terabytes.

Neste post, focamos especificamente no sistema de detecção de malware da Sophos para o formato de arquivo PDF. Mostramos como a SophosAI usa Amazon Sage Maker treinamento distribuído com terabytes de dados para treinar um modelo XGBoost (Extreme Gradient Boosting) leve e poderoso. Isso permite que sua equipe itere grandes dados de treinamento mais rapidamente com ajuste automático de hiperparâmetros e sem gerenciar a infraestrutura de treinamento subjacente.

Atualmente, a solução está perfeitamente integrada ao pipeline de treinamento de produção e ao modelo implantado em milhões de endpoints de usuários por meio do Serviço de endpoint Sophos.

Contexto do caso de uso

Se você deseja compartilhar um contrato importante ou preservar o design sofisticado do seu CV, o formato PDF é a escolha mais comum. Seu uso generalizado e a percepção geral de que tais documentos são herméticos e estáticos levaram os usuários a uma falsa sensação de segurança. O PDF tornou-se, portanto, um vetor de infecção preferido no arsenal dos invasores. Ações maliciosas usando PDFs são mais frequentemente alcançadas por meio da incorporação de uma carga JavaScript que é executada pelo leitor de PDF para baixar um vírus de um URI, sabotar a máquina do usuário ou roubar informações confidenciais.

A Sophos detecta arquivos PDF maliciosos em vários pontos de um ataque usando um conjunto de modelos determinísticos e de ML. Uma dessas abordagens é ilustrada no diagrama a seguir, onde o arquivo PDF malicioso é entregue por e-mail. Assim que uma tentativa de download é feita, ele aciona o script executável malicioso para se conectar ao servidor de Comando e Controle do invasor. O detector de PDF da SophosAI bloqueia a tentativa de download depois de detectar que é malicioso.

Outras formas incluem bloquear os arquivos PDF no endpoint, enviar os arquivos maliciosos para uma sandbox (onde são pontuados usando vários modelos), enviar o arquivo malicioso para uma infraestrutura de pontuação e gerar um relatório de segurança e assim por diante.

Motivação

Para construir um detector baseado em árvore que pode condenar PDFs maliciosos com alta confiança, permitindo baixo consumo de energia de computação de endpoint e respostas rápidas de inferência, a equipe da SophosAI descobriu que o algoritmo XGBoost é um candidato perfeito para a tarefa. Esses caminhos de pesquisa são importantes para a Sophos por duas razões. Ter modelos poderosos, mas pequenos, implantados no nível de endpoints do cliente tem um alto impacto nas análises de produtos da empresa pelos analistas. Ele também, e mais importante, fornece uma melhor experiência geral ao usuário.

Desafio técnico

Como o objetivo era ter um modelo com uma pegada de memória menor do que seus detectores de malware PDF existentes (tanto no disco quanto na memória), a SophosAI transformou o XGBoost, um algoritmo de classificação com um histórico comprovado de produzir modelos drasticamente menores do que as redes neurais, obtendo resultados impressionantes desempenho em dados tabulares. Antes de se aventurar na modelagem de experimentos XGBoost, uma consideração importante foi o tamanho do conjunto de dados. De fato, o conjunto de dados principal de arquivos PDF da Sophos está em terabytes.

Portanto, o principal desafio foi treinar o modelo com um grande conjunto de dados sem precisar fazer downsample. Como é crucial para o detector aprender a identificar quaisquer ataques baseados em PDF - mesmo agulhas no palheiro e completamente novos para melhor defender os clientes da Sophos - é de extrema importância usar todos os diversos conjuntos de dados disponíveis.

Ao contrário das redes neurais, onde você pode treinar em lotes, para o XGBoost, precisamos de todo o conjunto de dados de treinamento na memória. O maior conjunto de dados de treinamento para este projeto tem mais de 1 TB e não há como treinar em tal escala sem utilizar as metodologias de uma estrutura de treinamento distribuída.

Visão geral da solução

O SageMaker é um serviço de ML totalmente gerenciado que fornece várias ferramentas para criar, treinar, otimizar e implantar modelos de ML. o Bibliotecas de algoritmos integradas do SageMaker consistem em 21 algoritmos de ML populares, incluindo XGBoost. (Para mais informações, veja Simplifique o aprendizado de máquina com XGBoost e Amazon SageMaker.) Com o algoritmo integrado XGBoost, você pode tirar proveito do código aberto Contêiner SageMaker XGBoost especificando uma versão de estrutura superior a 1.0-1, que melhorou a flexibilidade, a escalabilidade, a extensibilidade e o Managed Spot Training e oferece suporte a formatos de entrada como Parquet, que é o formato usado para o conjunto de dados PDF.

A principal razão pela qual a SophosAI escolheu o SageMaker é a capacidade de se beneficiar do treinamento distribuído totalmente gerenciado em instâncias de CPU de vários nós simplesmente especificando mais de uma instância. O SageMaker divide automaticamente os dados entre nós, agrega os resultados entre nós pares e gera um único modelo. As instâncias podem ser instâncias spot, reduzindo significativamente os custos de treinamento. Com o algoritmo embutido para XGBoost, você pode fazer isso sem nenhum script personalizado adicional. Versões distribuídas do XGBoost também existem como código aberto, como XGBoost-RayName e XGBoost4J-Spark, mas seu uso requer a construção, proteção, ajuste e autogerenciamento de clusters de computação distribuída, o que representa um esforço significativo adicional ao desenvolvimento científico.

Além disso, SageMaker ajuste automático do modelo, também conhecido como ajuste de hiperparâmetros, encontra a melhor versão de um modelo executando muitos trabalhos de treinamento com intervalos de hiperparâmetros que você especifica. Em seguida, ele escolhe os valores de hiperparâmetros que resultam em um modelo com o melhor desempenho, conforme medido por uma métrica para a tarefa de ML fornecida.

O diagrama a seguir ilustra a arquitetura da solução.

Vale a pena notar que, quando a SophosAI iniciou os experimentos XGBoost antes de recorrer ao SageMaker, foram feitas tentativas de usar memória grande Amazon Elastic Compute Nuvem (Amazon EC2) (por exemplo, r5a.24xlarge e x1.32xlarge) para treinar o modelo em uma amostra de dados tão grande quanto possível. No entanto, essas tentativas levaram mais de 10 horas em média e geralmente falhavam devido à falta de memória.

Por outro lado, usando o algoritmo SageMaker XGBoost e um mecanismo de treinamento distribuído sem complicações, a SophosAI pode treinar um modelo de reforço em escala no colossal conjunto de dados de treinamento PDF em questão de 20 minutos. A equipe só teve que armazenar os dados em Serviço de armazenamento simples da Amazon (Amazon S3) como arquivos Parquet de tamanho semelhante e escolha um tipo de instância EC2 e o número desejado de instâncias, e o SageMaker gerenciou a infraestrutura de cluster de computação subjacente e distribuiu o treinamento entre vários nós do cluster. Sob o capô, o SageMaker divide os dados entre os nós usando ShardedByS3Key para distribuir os objetos de arquivo igualmente entre cada instância e usa a implementação XGBoost do Protocolo de coelho (AllReduce confiável e interface de transmissão) para iniciar o processamento distribuído e a comunicação entre nós primários e de mesmo nível. (Para obter mais detalhes sobre a agregação de histograma e transmissão entre nós, consulte XGBoost: um sistema de aumento de árvore escalável.)

Além de apenas treinar um modelo, com o SageMaker, Ajuste de hiperparâmetros XGBoost também foi rápido e fácil com a capacidade de executar diferentes experimentos simultaneamente para ajustar a melhor combinação de hiperparâmetros. Os hiperparâmetros ajustáveis ​​incluem hiperparâmetros específicos de reforço e específicos de função objetiva. Duas estratégias de pesquisa são oferecidos: aleatórios ou bayesianos. A estratégia de busca Bayesiana provou ser valiosa porque ajuda a encontrar melhores hiperparâmetros do que uma mera busca aleatória, em menos iterações experimentais.

Informações do conjunto de dados

A modelagem de detecção de malware em PDF da SophosAI conta com uma variedade de recursos, como histogramas de n-gram e recursos de entropia de byte (para obter mais informações, consulte MEADE: Rumo a um mecanismo de detecção de anexos de e-mail maliciosos). Metadados e recursos extraídos de arquivos PDF coletados são armazenados em um data warehouse distribuído. Um conjunto de dados de mais de 3,500 recursos é então calculado, dividido ainda mais com base no tempo em conjuntos de treinamento e teste e armazenado em lotes como arquivos Parquet no Amazon S3 para serem prontamente acessíveis pelo SageMaker para trabalhos de treinamento.

A tabela a seguir fornece informações sobre os dados de treinamento e teste.

Conjunto de dados	Número de amostras	Número de arquivos Parquet	Tamanho total
Training	70,391,634	5,500	~ 1010 GB
Test	1,242,283	98	~ 18 GB

Os tamanhos dos dados foram calculados seguindo a fórmula:

Tamanho dos dados = N × (n_F + n_L) × 4

A fórmula tem os seguintes parâmetros:

• N é o número de amostras no conjunto de dados
• n_F é o número de características, com n_F = 3585
• n_L é o número de rótulos de verdade, com n_L = 1
• 4 é o número de bytes necessários para o tipo de dados dos recursos: float32

Além disso, os gráficos de pizza a seguir fornecem a distribuição de rótulos dos conjuntos de treinamento e teste, provocando o desequilíbrio de classe enfrentado na tarefa de detecção de malware PDF.

A distribuição muda do conjunto de treinamento para o conjunto de teste de um mês. Uma divisão baseada em tempo do conjunto de dados em treinamento e teste é aplicada para simular o cenário de implantação da vida real e evitar a espionagem temporal. Essa estratégia também permitiu que a SophosAI avaliasse os verdadeiros recursos de generalização do modelo quando confrontado com novos ataques de PDF inéditos, por exemplo.

Experimentos e resultados

Para iniciar os experimentos, a equipe da SophosAI treinou um modelo XGBoost de linha de base com parâmetros padrão. Em seguida, eles começaram a realizar o ajuste fino de hiperparâmetros com o SageMaker usando a estratégia Bayesiana, que é tão simples quanto especificar o hiperparâmetros a ser ajustado e a faixa de valores desejada, a métrica de avaliação (ROC (Receiver Operating Characteristic) AUC neste caso) e os conjuntos de treinamento e validação. Para o detector de malware PDF, a SophosAI priorizou hiperparâmetros, incluindo o número de rodadas de reforço (num_round), a profundidade máxima da árvore (max_depth), a taxa de aprendizado (eta), e a taxa de amostragem das colunas ao construir árvores (colsample_bytree). Eventualmente, os melhores hiperparâmetros foram obtidos e usados ​​para treinar um modelo no conjunto de dados completo e, finalmente, avaliados no conjunto de teste de validação.

O gráfico a seguir mostra a métrica objetiva (ROC AUC) versus os 15 trabalhos de treinamento executados no trabalho de ajuste. Os melhores hiperparâmetros são aqueles correspondentes ao nono trabalho de treinamento.

No início dos experimentos da SophosAI no SageMaker, uma pergunta especialmente importante a ser respondida foi: que tipo de instâncias e quantas delas são necessárias para treinar o XGBoost nos dados disponíveis? Isso é crucial porque usar o número ou tipo errado de instância pode ser uma perda de tempo e dinheiro; o treinamento está fadado a falhar devido à falta de memória ou, se usar muitas instâncias grandes demais, isso pode se tornar desnecessariamente caro.

O XGBoost é um algoritmo vinculado à memória (em oposição ao vinculado à computação). Portanto, uma instância de computação de uso geral (por exemplo, M5) é uma escolha melhor do que uma instância otimizada para computação (por exemplo, C4). Para tomar uma decisão informada, há uma diretriz simples do SageMaker para escolher o número de instâncias necessárias para executar o treinamento no conjunto de dados completo:

Tamanho total dos dados de treinamento × fator de segurança^(*) < Contagem de Instâncias × Memória Total do Tipo de Instância

Neste caso: Tamanho total dos dados de treinamento × Fator de segurança (12) = 12120 GB

A tabela a seguir resume os requisitos quando o tipo de instância escolhido é ml.m5.24xlarge.

Tamanho do Treinamento × Fator de Segurança (12)	Memória da instância ml.m5.24xlarge	Contagem mínima de instâncias necessária para treinamento
12120 GB	384 GB	32

^*Devido à natureza do treinamento distribuído do XGBoost, que requer que todo o conjunto de dados de treinamento seja carregado em um objeto DMatrix antes do treinamento e memória livre adicional, um fator de segurança de 10 a 12 é recomendado.

Para dar uma olhada mais de perto na utilização de memória para um treinamento completo do SageMaker do XGBoost no conjunto de dados fornecido, fornecemos o gráfico correspondente obtido do treinamento Amazon CloudWatch monitoramento. Para este trabalho de treinamento, foram usadas instâncias de 40 ml.m5.24xlarge e a utilização máxima de memória atingiu cerca de 62%.

O custo de engenharia economizado ao integrar um serviço de ML gerenciado como o SageMaker ao pipeline de dados é de cerca de 50%. A opção de usar instâncias spot para trabalhos de treinamento e ajuste de hiperparâmetros reduz os custos em mais 63%.

Conclusão

Com o SageMaker, a equipe da SophosAI pode resolver com sucesso um projeto complexo de alta prioridade criando um modelo XGBoost de detecção de malware PDF leve que é muito menor em disco (até 25 vezes menor) e na memória (até 5 vezes menor) do que seu antecessor do detector. É um detector de malware pequeno, mas poderoso, com ~ 0.99 AUC e uma taxa de verdadeiro positivo de 0.99 e uma taxa de falso positivo de . Esse modelo pode ser retreinado rapidamente e seu desempenho pode ser facilmente monitorado ao longo do tempo, pois leva menos de 20 minutos para treiná-lo em mais de 1 TB de dados.

Você pode aproveitar o algoritmo integrado do SageMaker XGBoostName para construir modelos com seus dados tabulares em escala. Além disso, você também pode experimentar os novos algoritmos integrados do Amazon SageMaker LightGBM, CatBoost, AutoGluon-Tabular e Tab Transformer, conforme descrito neste blog.

---

Sobre os autores

Salma Taoufiq é Cientista de Dados Sênior na Sophos, trabalhando na interseção de aprendizado de máquina e segurança cibernética. Com formação em ciência da computação, ela se formou na Central European University com um mestrado. em Matemática e suas Aplicações. Quando não está desenvolvendo um detector de malware, Salma é um ávido caminhante, viajante e consumidor de thrillers.

Harini Kannan é Cientista de Dados na SophosAI. Ela está em ciência de dados de segurança há cerca de 4 anos. Ela foi anteriormente a Cientista de Dados Principal da Capsule8, que foi adquirida pela Sophos. Ela deu palestras no CAMLIS, BlackHat (EUA), Open Data Science Conference (East), Data Science Salon, PyData (Boston) e Data Connectors. Suas áreas de pesquisa incluem detecção de ataques baseados em hardware usando contadores de desempenho, análise de comportamento do usuário, ML interpretável e detecção de anomalias não supervisionadas.

Hasan Poonawala é arquiteto de soluções de especialista sênior em AI / ML na AWS, com sede em Londres, Reino Unido. Hasan ajuda os clientes a projetar e implantar aplicativos de aprendizado de máquina em produção na AWS. Ele tem mais de 12 anos de experiência de trabalho como cientista de dados, profissional de aprendizado de máquina e desenvolvedor de software. Em seu tempo livre, Hasan adora explorar a natureza e passar tempo com amigos e familiares.

Digante Patel é líder de suporte empresarial na AWS. Ele trabalha com os clientes para projetar, implantar e operar na nuvem em escala. Suas áreas de interesse são as práticas de MLOps e DevOps e como isso pode ajudar os clientes em sua jornada na nuvem. Fora do trabalho, ele gosta de fotografar, jogar vôlei e passar tempo com amigos e familiares.