Improving Generalization on Cybersecurity Tasks with Multi-Modal Contrastive Learning

TL;DR

This paper introduces a multi-modal contrastive learning framework that leverages textual descriptions to improve payload classification in cybersecurity, reducing shortcut learning and enhancing generalization.

Contribution

It proposes a novel two-stage contrastive learning approach that transfers knowledge from text to payloads, addressing generalization issues in cybersecurity ML models.

Findings

Reduces shortcut learning compared to baselines

Improves performance on large-scale private dataset

Effective transfer of knowledge from text to payloads

Abstract

The use of ML in cybersecurity has long been impaired by generalization issues: Models that work well in controlled scenarios fail to maintain performance in production. The root cause often lies in ML algorithms learning superficial patterns (shortcuts) rather than underlying cybersecurity concepts. We investigate contrastive multi-modal learning as a first step towards improving ML performance in cybersecurity tasks. We aim at transferring knowledge from data-rich modalities, such as text, to data-scarce modalities, such as payloads. We set up a case study on threat classification and propose a two-stage multi-modal contrastive learning framework that uses textual vulnerability descriptions to guide payload classification. First, we construct a semantically meaningful embedding space using contrastive learning on descriptions. Then, we align payloads to this space, transferring knowledge from text to payloads. We evaluate the approach on a large-scale private dataset and a synthetic benchmark built from public CVE descriptions and LLM-generated payloads. The methodology appears to reduce shortcut learning over baselines on both benchmarks. We release our synthetic benchmark and source code as open source.