Transformers for End-to-End InfoSec Tasks: A Feasibility Study

TL;DR

This study explores the feasibility of using transformer models for end-to-end cybersecurity tasks involving URLs and PE files, highlighting the importance of tailored training strategies and architectural adaptations for optimal performance.

Contribution

The paper introduces novel end-to-end transformer architectures for InfoSec data formats and proposes a mixed objective training method, demonstrating competitive results with existing benchmarks.

Findings

Auto-regressive pre-training on URLs does not transfer well to classification tasks.

Auxiliary auto-regressive loss improves URL classification performance.

Adaptive span self-attention enables transformers to handle longer byte sequences in PE files.

Abstract

In this paper, we assess the viability of transformer models in end-to-end InfoSec settings, in which no intermediate feature representations or processing steps occur outside the model. We implement transformer models for two distinct InfoSec data formats - specifically URLs and PE files - in a novel end-to-end approach, and explore a variety of architectural designs, training regimes, and experimental settings to determine the ingredients necessary for performant detection models. We show that in contrast to conventional transformers trained on more standard NLP-related tasks, our URL transformer model requires a different training approach to reach high performance levels. Specifically, we show that 1) pre-training on a massive corpus of unlabeled URL data for an auto-regressive task does not readily transfer to binary classification of malicious or benign URLs, but 2) that using an auxiliary auto-regressive loss improves performance when training from scratch. We introduce a method for mixed objective optimization, which dynamically balances contributions from both loss terms so that neither one of them dominates. We show that this method yields quantitative evaluation metrics comparable to that of several top-performing benchmark classifiers. Unlike URLs, binary executables contain longer and more distributed sequences of information-rich bytes. To accommodate such lengthy byte sequences, we introduce additional context length into the transformer by providing its self-attention layers with an adaptive span similar to Sukhbaatar et al. We demonstrate that this approach performs comparably to well-established malware detection models on benchmark PE file datasets, but also point out the need for further exploration into model improvements in scalability and compute efficiency.