Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs

Resumo Técnico: Classificação de CVE para CWE com RoBERTa Fine-Tuned

1. O Problema

A mapeamento de entradas de CVE (Common Vulnerabilities and Exposures) para categorias CWE (Common Weakness Enumeration) é fundamental para a triagem, priorização e remediação de vulnerabilidades. No entanto, o processo atual enfrenta desafios significativos:

• Ruído nos Dados: As etiquetas CWE atribuídas pelo National Vulnerability Database (NVD) são frequentemente inconsistentes, genéricas demais (ex: uso excessivo de "CWE-20" como catch-all) ou imprecisas.
• Limitações dos Modelos Atuais: Modelos de Linguagem de Grande Porte (LLMs) generalistas, como o GPT-4, alcançam acurácias moderadas (72%), enquanto modelos menores (ex: LLaMA3-8B) performam significativamente pior (44,7%). Além disso, soluções proprietárias de alto desempenho (como o Sec-Gemini do Google) não são abertas e suas métricas exatas são difíceis de auditar.
• Ineficiência Paramétrica: Existe uma lacuna entre modelos pequenos (eficientes) e grandes LLMs, onde modelos menores geralmente falham em tarefas complexas de segurança cibernética.

2. Metodologia

A. Construção do Dataset (290K Amostras)

• Fonte: Coleta de 336.777 entradas CVE (1999–2026) do NVD.
• Refinamento de Rótulos (AI-Assisted): Para superar o ruído do NVD, os autores utilizaram o Claude Sonnet 4.6 (via Anthropic Batch API) para reetiquetar todas as descrições.
  • O modelo foi instruído a classificar cada CVE em uma das 205 classes CWE alvo.
  • Consenso: 73,1% das etiquetas do NVD coincidiam exatamente com as do Claude; ao considerar equivalências hierárquicas (pai/filho), a concordância sobe para 84,5%.
  • Estratégia de Split:
    • Treino (234.770): Inclui amostras com concordância, discordância (usando o rótulo do AI) e CVEs anteriormente não rotulados.
    • Teste/Validação (Agreement-Filtered): Apenas amostras onde NVD e AI concordam, criando um conjunto de alta confiança para avaliação interna.
• Decontaminação: Todas as 2.000 amostras do benchmark externo CTI-Bench foram removidas para garantir avaliação externa limpa.
• Anotação Adicional: Mapeamento para técnicas MITRE ATT&CK (361 IDs únicos) incluído no dataset liberado.

B. Arquitetura e Treinamento

• Modelo Base: RoBERTa-base (125M parâmetros), um encoder de linguagem pré-treinado.
• Estratégia de Fine-Tuning em Duas Fases:
  1. Fase 1 (Warm-up): Congelamento das primeiras 8 camadas do transformer e da camada de embeddings. Apenas as 4 camadas superiores e a cabeça de classificação são treinadas (4 épocas). Isso estabiliza a cabeça de classificação sem esquecer representações pré-treinadas.
  2. Fase 2 (Full Fine-tuning): Descongelamento de todos os parâmetros para ajuste fino completo (9 épocas).
• Hiperparâmetros: Batch size 32, sequência máxima de 384 tokens, label smoothing 0.1, otimizador com cosine schedule. Treinado em uma única NVIDIA RTX 5080 por ~4 horas.

3. Contribuições Principais

1. Dataset de Grande Escala: Liberação pública de um dataset de 290K amostras com rótulos refinados por IA e anotações MITRE ATT&CK.
2. Abordagem de Treinamento: Validação de uma estratégia de two-phase fine-tuning que maximiza o desempenho de modelos pequenos.
3. Eficiência Paramétrica: Demonstração de que um modelo de 125M parâmetros pode competir com LLMs de 8B parâmetros (64x menores) em tarefas específicas de segurança.
4. Análise de Granularidade: Identificação da discrepância de granularidade na hierarquia CWE como um fator de confusão crítico nos benchmarks atuais.

4. Resultados

A. Avaliação Interna (Conjunto de Teste Filtrado)

• Acurácia Top-1: 87,4% (vs. 84,9% de um baseline TF-IDF simples).
• Macro F1: 60,7% (ganho de +15,5 pontos percentuais sobre o baseline).
• Observação: O ganho no Macro F1 indica que o modelo RoBERTa é superior na classificação de classes de CWE raras, onde o baseline baseado em palavras-chave falha.

B. Benchmark Externo (CTI-Bench RCM)

• Desempenho: O modelo alcançou 75,6% de acurácia estrita (IC 95%: 72,8–78,2%).
• Comparação:
  • Cisco Foundation-Sec-8B-Reasoning (8B parâmetros): 75,3%.
  • GPT-4: 72,0%.
  • LLaMA3-8B: 44,7%.
• Conclusão: A diferença de 0,3% entre o RoBERTa (125M) e o Cisco (8B) não é estatisticamente significativa. O modelo pequeno é 64 vezes mais eficiente em parâmetros, mantendo desempenho competitivo.

C. Análise de Erros e Hierarquia CWE

• Uma grande parte das "erros" no CTI-Bench deve-se a uma mismatch de granularidade: o modelo prevê filhos específicos (ex: CWE-121 Stack Buffer Overflow), enquanto o ground truth do NVD usa pais genéricos (ex: CWE-119 Buffer Overflow).
• Ao aplicar uma avaliação hierarquia-consciente (onde pai/filho conta como correto), a acurácia no CTI-Bench sobe para 86,5%, fechando quase toda a lacuna com o desempenho interno.
• Validação Manual: Revisão de 100 casos de discordância mostrou que o Claude Sonnet foi mais preciso em 72% dos casos, frequentemente corrigindo rótulos genéricos ou errados do NVD.

5. Significado e Impacto

• Viabilidade de Modelos Pequenos: O trabalho prova que não é necessário um LLM massivo para tarefas de classificação de vulnerabilidades; um encoder específico e bem treinado é suficiente e mais eficiente.
• Qualidade de Dados > Tamanho do Modelo: A qualidade dos rótulos (refinados por IA) e a estratégia de treinamento foram mais determinantes para o sucesso do que o tamanho do modelo.
• Crítica aos Benchmarks: O artigo alerta que benchmarks atuais penalizam injustamente modelos que preveem classificações mais específicas e úteis (filhos da hierarquia CWE) quando o ground truth é genérico.
• Recursos Abertos: Todos os dados, código e o modelo pré-treinado foram liberados publicamente no Hugging Face, permitindo reprodutibilidade e avanço na comunidade de segurança cibernética.

Em suma, este estudo estabelece um novo padrão para a classificação CVE-to-CWE, demonstrando que a combinação de dados de alta qualidade (AI-refined) e fine-tuning estratégico em modelos compactos pode superar ou igualar soluções muito maiores e proprietárias.