{'fraction_non_alphanumeric': 0.03016127046657639, 'fraction_numerical': 0.010217899790717714, 'mean_word_length': 5.080973175296319, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Transformers for End-to-End InfoSec Tasks: A Feasibility Study

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.


INTRODUCTION
The abundance of labeled data sources has been a major driver in the success of ML in various applications such as image recognition [20], audio recognition [17], natural langugage processing (NLP) [13], and InfoSec [24,28,32,34,44]. For applications which involve image, natural language text, or audio data, a recent dominating trend has been to feed raw data with increasingly minimal preprocessing directly to an ML model in favor of fitting a model on hand-engineered features. For many applications, this approach of learning from raw data works well. Applications of ML in InfoSec, however, require dealing with unique data formats which must be approached in different ways, and this is one of the reasons why InfoSec applications of ML still rely heavily on relatively simple models fit on hand-engineered features.
This may be in part because properly learning features based on raw data is less trivial for some of the data formats inherent to security problems, for example, in [6], a convolutional neural network (CNN) which operates on raw PE files and is closely analogous to CV or NLP CNN [32] under-performs a gradient boosted decision tree model trained with default parameters on hand-engineered features.
From a scalability perspective, hand-engineered features for In-foSec tasks are non-trivial because creating them often requires expensive and time-consuming interaction with subject matter experts (SMEs). Moreover, ongoing attacker innovations mean that feature extractors may need to be continuously updated in order to keep pace. Models that learn features directly from raw data, on the other hand, require no such updates to feature extraction processors and do not require as much SME interaction. Thus, there is substantial motivation for research into more effective ML models for InfoSec tasks that operate on raw data.
With respect to other domains where ML models operating on raw data have advanced the state of the art, NLP is arguably most analogous to InfoSec. Some InfoSec tasks can even be framed as NLP problems, including email spam detection [11], website content categorization [35], social media phishing [43], and certain areas of data leak prevention (DLP) [3][4][5]. When framing InfoSec tasks in NLP terms, the raw data does not necessarily conform to typical prosaic sentence and paragraph structures. Some InfoSec tasks, e.g., source code attribution [2], contain little to no "natural language" yet may possess their own linguistic structure. InfoSec problems also deviate from NLP problems in terms of how their classifiers are typically trained. InfoSec classifiers in the industry are typically fit on millions to hundreds of millions of weakly labeled samples, where supervisory signals are e.g., aggregated over threat feeds or externally derived vendor scores. NLP classifiers, particularly those which rely on transformers, typically use large datasets for self-supervised pre-training of a base architecture, but relatively small datasets for fine tuning final layers for a given classification task. We explore the ramifications of this within this paper.
InfoSec data also more commonly contains long-ranging sequential dependencies in which one or more tokens far-removed from a given token may strongly influence the probability of that token. Given the close analogy to NLP, and because transformer architectures have revolutionized performance in the NLP domain, it seems fruitful to explore the feasibility of applying transformers to InfoSec tasks, especially since they address the long range dependency issue using an attention mechanism.
However, InfoSec is a broad field, and sequential dependency ranges vary dramatically based on the data type. In NLP literature, "long-range" dependencies are typically considered to span tens to tens of thousands of tokens -often backwards -within a sequence. Similar ranges may apply for certain types of InfoSec tasks, but not for others. For example, binary classification, where function calls and declarations can be almost arbitrarily separated may involve relevant dependencies spanning forward or backward millions or even billions of bytes.
While research has been conducted on scaling transformer context window sizes, additional research is required to develop transformer architectures that work well for a broad array of InfoSec tasks, specifically those with their own data formats that deviate from conventional applications of NLP. In this paper, we conduct a feasibility study using two different data formats: URLs and PE files. Other works [10,32,42] have applied deep learning based on raw data from these formats and explored the learnt representations. Applying transformers towards these tasks is in many ways a natural extension of this line of research. For each data format we perform a malicious/benign detection task. For URLs, we aim to detect those linking to malicious content, while for PEs we aim to detect files containing malware. We chose to focus on these two data formats because they are not natural language but have distinctly different characteristics under which we can explore the viability of transformer models: URLs are relatively short sequences of characters, can fit nicely into a canonical transformer's context window, and allow for relatively efficient training on a GPU cluster. This allows us to explore the effects of different training strategies and loss functions under feasible iteration times. In contrast, PE files span long sequences of bytes, with large quantities of content that potentially provides little to no malicious or benign indications (e.g., padding, images, or encrypted payloads). Thus, modifications for longer sequences must be made to the canonical transformer to feasibly operate on raw bytes (e.g., hidden state caching [12], sparse attention patterns [9], locality sensitive hashing [19], etc.).
The contributions of our feasibility study are as follows: • A performance comparison between our novel transformer approaches and other more conventional approaches to URL classification. Using a transformer model, we are able to achieve performance on par with our top benchmark models. • A performance comparison of multiple URL transformer training regimes, including with and without auto-regressive pre-training. We demonstrate that in contrast to NLP applications which rely heavily on self-supervised pre-training, this strategy does not readily improve performance on the URL classification task. • A novel balanced mixed objective transformer loss function which balances a classification objective with an auxiliary auto-regressive objective during training. This is designed to encourage the transformer to fuse both sequential context with class information in its hidden states. • A performance comparison between our own modification of an adaptive attention span transformer [45] and other bytebased and feature-based models on truncated PE files from the EMBER dataset. This includes featurizing the EMBER PE files after truncation to offer a fair comparison, which has typically not been performed in the literature when comparing raw byte models to featurized models. • An exploration of various ways to encode bytes with little to no information loss for raw-byte models, including transformers. We investigate whether or not these encodings may be able to improve performance by extending effective context window size at the cost of increased vocabulary size.

BACKGROUND
Transformers, introduced by Vaswani et al. [47] have revolutionized NLP and a variety of other discrete sequence modeling domains. This is partly due to their ability to directly incorporate long-term dependencies in a sequence and partly because they are easily trained in parallel. Transformers were originally formulated for sequence transduction tasks, e.g., neural machine translation (NMT) [7], wherein a source sequence is encoded via an encoder stack of attention and feed-forward layers, then decoded via a similar decoder stack, using the input sequence as a source context for decoding to the target sequence. Later models, which constitute massively pre-trained representations, meant for fine-tuning for non-transductive tasks abandon the encoder/decoder setup entirely, opting in favor of either an encoder or a decoder. Depending on the problem, the flow of information from the input sequence may be either bidirectional, e.g., [13], or left to right (L-R) [31]. While transformers can be implemented in a variety of ways, a typical implementation of a transformer attention stack consists of the following components: (1) Embedding: The input is projected into an embedding space.
(2) Positional Encoding: This is a signal added to each of the embedding vectors to imbue each embedding with a positional order in the sequence (which would otherwise not be tracked by the attention mechanism). This can optionally be  Figure 1: For a given context window size, transformers utilize padding/masking operations to perform inference over different sequence lengths. Let be the length of the sequence and be the size of the context window, where ≤ −1. Let 1 through be the embeddings corresponding to the sequence elements, in this case, characters within a URL. Let CLS be the embedding corresponding to the CLS token (position + 1), and let 1 through − −1 be embeddings of padding tokens. Let 1 through be the final hidden states derived from the transformer. A standard next character prediction task is shown in a, where the hidden states are each passed through a classifier , and the softmax outputs correspond to probabilities of the subsequent tokens. b We can also use the output of +1 , passed through a binary classifier to derive a probability that the input sequence is malicious or benign. c During training, we can use all token outputs to fit a malicious/benign predictor and a next-character prediction loss simultaneously. Note that hidden states corresponding to padding are ignored. For our implementation, left-to-right masking operations prevent prior hidden states from attending to any padding-related sequence information.
added to hidden states within the transformer, but we do not do this in our implementation. typically added to inputs shorter than the context window of the transformer. In order to avoid "attending" to the padding, masking is used. Masking is also used to enforce sequential dependencies, e.g., for our implementation, we enforce an L-R sequential dependence.
We will elaborate on the architecture of our specific implementation later on in this paper. However, we present a high-level schematic of how multiple prediction tasks can be performed using a transformer in Fig. 1. Fig. 1a depicts a standard next-character prediction task from the literature. In this case, the outputs are softmax probabilities over the vocabulary. Fig. 1b depicts using the transformer as a binary malicious/benign classifier -the baseline "decode-to-label" approach presented in Sec. 3.1. Fig. 1c depicts the mixed objective approach presented in Sec. 3.3, where a loss over the next character prediction output is used as an auxiliary loss in conjunction with the main classification task loss during training.
A downside of this architecture is that the transformer's multiheaded self-attention mechanism scales quadratically with respect to input sequence length. This is particularly problematic in the case of raw bytes of a binary file, whose dependencies can be spread over large swaths of the executable. Many different approaches to addressing these computational and memory inefficiencies have been proposed, ranging expanding the context window by ignoring tokens [9,45], reducing the representation in terms of memory and parameter size [19,22], sequence-level recurrence [12], and even dropping attention altogether in favor of pre-trained convolutions [46].

Transformers for InfoSec
ML for InfoSec has been researched for decades [40], but widespread industry adoption has occurred only over the past several years. Of the adopted models, the majority of them utilize some form of handcrafted features (e.g., [6,15,21,33,37,39]), with only a few operating on raw data. Sequence models, e.g., recurrent neural networks (RNNs) have been applied with some success [28], but only to niche problems. Training RNNs for most problems is fundamentally not scalable to long sequences, both in terms of the hefty latency and memory requirements, and in terms of lack of the general loss of information from long term dependencies.
Such long term dependencies are common in NLP tasks, where it has become common to pre-train transformers in an auto-regressive manner on large unlabeled sequences of text, via next token prediction for L-R models and masked language modeling (MLM) for bidirectional models, then fine tune on smaller quantities of labeled data. This is often necessary, when potentially few labeled examples are available for the task of interest. However, many InfoSec applications where ML works well have millions to hundreds of millions of labeled samples. These labels are typically derived from an aggregation of multiple weak labeling sources [16,36,39].
We are not the first to apply transformers to InfoSec tasks. Li et al. applied transformers to malware detection in [24]. However, contrary to their hierarchical approach with multiple transformers and neural networks trained with different training/fine-tuning regimes,our approach is built on a single end-to-end transformer. Moreover, their approach also involves disassembly which adds intrinsic context for sequential modeling, but is expensive in terms of added processing time. Our approach by contrast, operates on raw bytes. While their work demonstrates superior cross-validated performance on assembly code representations, their evaluation is severely limited in scale (∼ 10k samples) and further evaluation is needed to assess its efficacy in realistic scenarios.
Another approach by Pei et al. [30] uses a hierarchical transformer trained on dynamic micro-traces of a number of different functions in an MLM regime. The learnt representation is then fine-tuned to yield state-of-the-art results on a variety of semantic similarity tasks (e.g., determining semantic similarity of two functions across different architectures, compiler optimizations, etc.). The authors additionally provide a dataset for semantic similarity benchmarking. However, it differs from ours insofar as their hierarchical transformer uses micro-traces as its input modality whereas ours uses raw bytes.
While transformers excel at sequence modeling, in this work, our predominant application is binary detection. In order to encourage the transformer to fuse sequential context with class information, in this work we combine the tasks of next character prediction and malicious/benign classification into a common loss function. Our approach is a natural extension of other research [14,18,37,38] on multi-objective/multi-task training, which demonstrates that incorporating auxiliary losses can improve the performance of a classifier on the main task. In contrast to these previous approaches, however, our approach is applied to transformers and uses a novel technique for dynamically re-weighting the per-task loss, such that no auxiliary loss term ever dominates regardless of the loss magnitude.
We apply transformers to two InfoSec tasks: malicious URL detection and malicious PE file detection.
Our work is not the first to address malicious URL detectionthere are several approaches in the literature that use convolutional architectures to classify malicious/benign URLs and domain generation algorithm ( DGA)-generated domains including [23,41,42,48]. However, they use different benchmark data sets, and are thus not directly comparable to our work. We benchmark against similar, internally developed convolutional and non-convolutional architectures in this paper, Note that the primary focus of this work is not creating the optimal URL classifier, but exploring how to train transformers for a typical non-NLP InfoSec machine learning task and contrasting with transformer training regimes for NLP problems.
With respect to detecting malicious PE files, several works have tackled this problem, e.g., [6,10,14,21,32,37]. While many apply classifiers fit over raw features (e.g., string hashes, byte entropy histograms, etc.), the most meaningful comparisons are models which operate on raw bytes, e.g., [10,32]. However, there is an additional challenge when applying transformers to long sequences -specifically, the size of the sequence context window that can fit in memory is substantially limited. There have been several methods proposed in the literature to ameliorate this, including modeling using sparse attention patterns partial span lengths, or utilizing approximations to avoid quadratic-complexity, e.g., [8,12,19,26,45]. We apply our own variation of the adaptive span method proposed by Sukhbaatar et al. [45] in Section 4.

URL TRANSFORMER APPROACH
Let be a dataset of URLs with binary labels . Let be a generic URL from with label ∈ {0, 1}, 0 corresponding to benign and 1 corresponding to malicious.

Baseline: Decode-to-Label
We refer to our baseline transformer approach as "decode-to-label", where the transformer ingests a URL, and only the hidden state corresponding to the a special classification token (⟨CLS⟩) is used to predict the label. Contrary to sequence transduction tasks that leverage an encoder transformer and a decoder transformer to encode a sequence into another sequence (e.g., language translation), we can think of this approach as using a transformer strictly as a decoder from a source sequence to a malicious/benign label. Our "decode-to-label" approach uses a left-to-right (L-R) decoder, similar to OpenAI's GPT [31] approaches, with a ⟨CLS⟩ token placed at the end of the sequence. Sequence information is propagated through the attention layer states in a L-R manner, meaning that hidden state in the th attention layer is fed information corresponding to hidden states 1, 2, . . . , − 1, in the ( − 1)th layer. The final classification is made via a feed-forward neural network (FFNN), which is fed the top layer's hidden state corresponding to the ⟨CLS⟩ token. The final dense layer of the FFNN projects the output to a 1D value. This value is then passed through a Sigmoid activation function to assume a final prediction ℎ( ) ∈ [0, 1]. Binary cross entropy between the prediction ℎ( ) and label is evaluated during training and the associated gradients are backpropagated. The associated Classification Loss is then: In contradistinction to most transformer literature, this approach does not utilize any explicit loss over hidden states corresponding to tokens within the sequence. Our rationale for applying the "decodeto-label" approach is to provide a benchmark against which to assess any gains and losses that have been introduced by explicitly optimizing sequential information into the transformer.

Next Character Prediction Pre-Training and Fine-Tuning
Next character prediction tasks have been applied throughout the transformer literature. For this regime, no labels are used during pre-training. Instead, we have a pre-training setˆwhich may or may not have associated malicious/benign labels. Each subsequent character in the URL serves to label each previous character. For example, given input URL =https://www.xyzsecurity.com/, the "label" sequence would be ttps://www.xyzsecurity.com/⟨CLS⟩. Note that the ⟨CLS⟩ token is omitted from the input sequence. We encode each character with its respective ASCII byte value ranging from 0 to 255. Note that in practice only a subset of these byte values are manifest in our data. We use the value 256 to represent our ⟨CLS⟩ token, yielding 257 distinct input embeddings.
Next character prediction is performed over each hidden state of the transformer output, corresponding to the embedded vector of the sequence up until that point. A feed-forward neural network (FFNN), which takes the corresponding transformer hidden state as input is used as a predictor of the next character. The output of the FFNN is a 257-element softmax, with the first 256 output probabilities corresponding to the probabilities of specific byte values as the next character and the last output probability corresponding to the probability of the ⟨CLS⟩ token (i.e., the end of the sequence). Note that for the next character prediction task, the input sequenceî gnores the ⟨CLS⟩ token at input; this token is only used as a "label" for the last character of the URL.
Loss is evaluated as the categorical cross entropy over the entire sequence, normalized by the sequence length . Let (·) be an indicator function which evaluates to 1 if the argument is true and 0 otherwise. The next character loss function over the URLb ecomes: Backpropagation of is used to train the underlying transformer representation. During fine-tuning, the pre-trained representation is loaded, potentially with lower attention layers frozen. With a decreased learning rate, and are fed to the transformer with used to train a malicious/benign predictor and update the unfrozen parameters of the transformer.

Balanced Mixed Objective Training
This approach aims to jointly optimize for both next character prediction and malicious/benign classification across dataset with labels . Note that this does not preclude next character prediction pre-training over another dataset. The rationale behind this approach is built on prior research which suggests that optimizing over multiple (correlated) tasks simultaneously leads to a better performing classifier with more stable convergence characteristics [37].
Following this rationale, we apply a mixed objective optimization approach, which balances main malicious/benign determination task with an auxiliary next character prediction loss. Contrary to previous research, which uses ad-hoc fixed weights on main and auxiliary task losses, our novel approach employs an adaptive balancing scheme, which ensures that no single loss term dominates, regardless of the loss value. Our loss-weighting strategy is as follows: At iteration of training, corresponding to one mini-batch, values and are balancing multipliers computed for each mini-batch, and are assumed constant when computing the gradient of the loss function. They are used to weight the classification and next character prediction loss components. Respectively, ensures that accounts for + of and ensures that accounts for + of . Values and are hyperparameters which we fix during training. Note that for simplification we can normalize such that + := 1 and say that and are the respective loss fractions themselves. For our experiments, we set := := 0.5 unless specified otherwise. Values and are computed for each minibatch, according to: Note that in this work, we apply only two loss types. However, we could trivially extend our approach to different loss types as follows. Given generic multiplier for the th loss term at iteration , , and desired loss contribution fraction / =1 , we compute as follows:

BYTE TRANSFORMER APPROACH
As generic byte sequences tend to be much longer than URLs, the stock transformer architecture utilized in Sec. 3 quickly becomes infeasible. There are a multitude of approaches for extending context window sequence lengths in the literature (see Sec. 2). Of these approaches, the notion of adaptive attention span, introduced in [45], is intriguing. Moreover, the implementation by [45] is relatively mature and less stringently designed around specific NLP tasks. Thus, we chose this method as the basis for a feasibility study of a Byte Transformer. Note, however, that any of the plethora of other techniques discussed in Section 2 could be employed to this end. Adaptive span transformers conserve memory by applying the attention operation only over a prior sub-span of the context window, similar to Sparse Transformers. Theoretically, over multiple layers of the Attention + FFNN stack, information from the entire sequence should percolate into the final attention layer, despite the reduction in context window size. The Adaptive Attention Span paper also incorporated a parameter that tunes the respective attention span of each layer in the Attention + FFNN stack. A simplified schematic, aimed to capture the intuition behind Adaptive Span Transformers is presented in Fig. 2. Here, four sampled (not necessarily subsequent) blocks of Attention + FFNN are depicted in two separate stacks; one from an Adaptive Span transformer and one from a canonical "Full Attention" transformer. For the Adaptive Span transformer, the overlap in span from hidden states between prior and subsequent blocks notionally propagates sequence information over the full context of the input sequence as Attention + FFNN blocks are stacked. If partial span information can be successfully fused over the full stack of Attention + FFNN blocks, then equivalent performance to a canonical "Full Attention" transformer can be obtained at significantly reduced memory requirements. In [45], Sukhbaatar et al. validate this approach for canonical NLP problems, and we employ their approach as well as a spinoff of their implementation for our Byte Transformer.
We modified the original code from Sukhbaatar et al. by removing the log softmax multinomial output (for next character prediction tasks) and placing a dense layer with a sigmoid output atop the final hidden state. While we kept the hidden state caching functionality from Dai et al., which allows for tackling even longer sequences [12], we did not use this functionality during our experiments on EMBER (see Sec. 6.1), as this 1) makes training on batches of data more difficult for varying length sequences and 2) takes longer to train. We note, however, that on our infrastructure, combining state caching with reduced attention spans (see Sec. 6.1), we were able to feasibly train on sequences of up to 20kB.
During training, we used generators which enforced balancing of malware and goodware samples within each minibatch, with a minibatch size of 32. For our optimizer, we used SGD, with a learning rate of 0.005, momentum of 0.9, and a decay rate of 0.0001. For our Byte Transformer experiments, we utilized only the classification loss without jointly optimizing next-token prediction (see Fig. 1).

URL EXPERIMENTS
We collected a dataset of URLs, down-sampled over 2-3 months from a threat intel feed in late 2019. Each URL had multiple weak malicious/benign labels. We derived a single malicious/benign label for each URL using Snorkel [36]. The training set consisted of 1,007,451 labeled URLs with 180,052 malicious and 826,333 benign. The validation set consisted of 111,930 URLs with 20,013 malicious and 91,813 benign. The test set consisted of 279,874 URLs with 50,029 malicious and 229,604 benign respectively. We additionally collected a dataset of 20 million unlabeled URLs for pre-training experiments, disjoint from train, test, and validation sets. Note that we did not perform any cleanup of these URLs, meaning that our dataset includes lengthy/obscure URLs as well as raw IPs; thus the reported performance numbers are not comparable with those from other authors with proprietary datasets (e.g., [42]).

Transformer Base Topology
We implemented a transformer with 20 layers of attention, a context window size of 256, a model hidden state size of 64, a feed-forward dimension of 128, and 4 attention heads per-layer. We also employed dropout, with a dropout ratio of 0.1. For classification, we apply a feed forward neural network (FFNN) atop the output corresponding to the ⟨CLS⟩ token. This FFNN consists of a 64-dimensional input followed by Layer Normalization, a hidden layer of 32-dimensions, Exponential Linear Unit activation, another hidden layer, reducing dimensionality from 32 to 1, and a sigmoid output. For next character prediction, we apply a similar architecture, but with a softmax output of 257 dimensions, predicting bytes 0 through 255 and the ⟨CLS⟩ token (which we ascribe label 256).

Comparison of Different Training Regimes
We performed comparisons of our four training techniques, ensuring that performance converged on the validation set. Baseline (DecodeToLabel): As a baseline, we compute loss only over the binary malicious/benign prediction, performing neither pre-training nor next character prediction. We trained the model for 15 epochs with minibatch size 512. For this training regime we used PyTorch's default Adam optimizer.
Auto-regressive Pre-training and Fine Tuning -Training Set (FineTune): For this experiment, we first performed pre-training over the entire training set for the next-character prediction task for 15 epochs. We then performed fine-tuning, freezing the first 16 attention layers and trained the decode-to-label task for 15 epochs at a reduced learning rate. For Pre-Training, we used PyTorch's default Adam optimizer. For fine-tuning, we used SGD optimizer with a learning rate of 1e-4. Auto-regressive Pre-training and Fine Tuning -Pre-Training Set (FineTune 20M): For this experiment, we first performed pretraining over the 20 million URL pre-training dataset using Py-Torch's default Adam optimizer. Only 2 epochs were required for convergence using a dataset of this magnitude. We then performed fine-tuning, freezing the first 16 attention layers and trained the decode-to-label task for 15 epochs at a reduced learning rate. For fine-tuning, we used SGD optimizer with a learning rate of 1e-4.

URL Transformer Results
We report results in terms of Receiver Operating Characteristics (ROC) curves and area under the ROC curve (AUC).
Surprisingly, fine-tuning from a pre-trained initialization resulted in poor performance compared to training from scratch from a Xavier initialization (see Fig. 3). We additionally attempted fine-tuning, freezing all but the last 4 layers of the transformer, but witnessed similar performance to our original fine-tuning regime. This is despite a consistent decrease and general convergence in loss for both of the pre-trained representations. This suggests that there is less immediate task transfer between next-character prediction for URLs than there is for next-character prediction in NLP contexts.
Despite the failure of pre-training with an auto-regressive loss to deliver performance gains, we did find that we were able to achieve marginal performance improvements by introducing an auxiliary next-character prediction loss. This is consistent with the findings of Rudd et al. [37].
While both and from Sec. 3 tended to converge with training, we noticed was consistently an order of magnitude greater than . This is not surprising, as the classification output predicts a single binary malicious/benign score from the ⟨CLS⟩ token embedding, while the next character prediction loss is summed over the full length of the URL. While the re-scaling in the joint loss term in Eq. 3 prevents either or from having an undue influence over the optimization process, this may affect the rate at which and converge but does not directly affect the magnitude of or .

Comparison Models
As a viability comparison, we benchmarked our transformer against several other models which we have developed for malicious/benign URL detection. Substantial development and testing effort went into these models. We summarize them in this section.
Random Forest on SME-Derived Features: For this model, proprietary features were derived with the help of subject matter experts (SMEs). These feature vectors consist of binary values/counts derived from parsing the URL and checking if specific parsed values from the URL string reside in various lists corresponding to likely indicators of malicious or benign content. The derived feature vectors therefore characterize the content of the URL. We then fit a random forest on these extracted feature vectors. The random forest classifier consisted of 30 trees, each with a maximum depth of 20. Nodes were split based on information gain. For all other parameters, we used Scikit-learn's [29] defaults.
LSTM on Raw URLs: For this model, we fit a long short-term memory neural network (LSTM) over the URLs, using embeddings of size 50 and an LSTM hidden size of 100. We performed optimization for a max of 100 epochs with early stopping based on validation performance. During optimization, we used an Adam optimizer with Keras's default parameters and a minibatch size of 128.
1D CNN on Raw URLs: This model first embeds each input token into a 40-dimensional vector, then follows up with a dropout layer (p=0.2) and a stack of two 1D convolutions with ReLU activations. The convolutional layers have 256 and 100 filters respectively with kernel sizes 5 and 3. The output of the convolutional layer stack is then globally max-pooled and passed through a hidden size layer of size 256, a dropout layer (p=0.2), and a ReLU activation. This is then transformed to 1D via a final fully-connected layer and passed through a Sigmoid output.
Our mixed objective transformer model (MixedObjective in Fig.  4) outperforms all but one of the comparison models, substantially outperforming the random forest feature-based model (RF) and the LSTM and performing on-par with the CNN model.
As previously stated, our URL results are not directly comparable to other works, as they are derived from a different dataset. Our URL dataset was collected with the use case of checking embedded URLs in suspicious emails in mind, and using the results as indicators to flag an email as suspicious or warn around a specific link contained within. The URLs chosen were unfiltered and intentionally representative of more challenging edge cases in a production pipeline, where ML detection is one of many tools in the defensive arsenal. Thus, the ROC curves presented herein should not be read as performance of a standalone commercial system. With respect to URL detection rates, this is a more general and more difficult use-case than checking URLs strictly meant to appear as human-readable content, e.g., a spoofed payment processing page. This is also a different and potentially more difficult problem than simply detecting URLs/domains generated by domains generation algorithms (DGAs), as the content on a website is incidental to how the domain was generated, i.e., DGA-generated domains are not necessarily malicious and non-DGA-generated domains are not necessarily benign. However, our results are consistent with other literature in terms of rank order of similar classifiers. Saxe and Berlin [42] found that a 1D CNN on character embeddings (similar to ours) outperformed a baseline trained on extracted features, while Yu et al. [48] found that this same architecture yielded state of the art for DGA detection, slightly outperforming LSTM models and significantly outperforming random forests over lexical features.

PE EXPERIMENTS
In Sec. 6.1 we examine the viability of our Byte Transformer via experiments on the EMBER 2018 PE dataset [6]. We then investigate encoding schemes to potentially improve the effectiveness of bytebased classifiers in Sec. 6.2 and Sec. 6.3.

Performance Comparison on EMBER
Due to GPU memory imposed constraints on transformer context windows, even for adaptive span transformers, we truncated each of the EMBER PE files to 4096 bytes. Thus, performance comparisons should be read in terms of relative rank of each benchmark, rather than compared to numbers in the original EMBER dataset. This was the only change that we made to the EMBER benchmark; for consistency, we maintained the original temporal dataset splits and followed the training and evaluation protocol from [6]. We conducted training using a Cirrascale server with 8 Tesla M40 24GB GPUs.
Note that we were able to effectively expand the context window size to 20kB using the TransformerXL caching mechanism [12], but opted not to employ this in our experiments as it significantly increased training time, and added additional levels of complexity.
We trained our Byte Transformer on the EMBER training set using the following selection of hyperparameters. Note that this was one of a few hyperparameter choices that actually attained convergence: • For the classifier, we used a logistic regressor fit atop the 32dimensional hidden state corresponding to a ⟨CLS⟩ token. Logistic regression weights and transformer parameters were jointly optimized during training. We trained the Byte Transformer for 15 epochs after which validation convergence was attained.
To assess the relative performance of our Byte Transformer model, we fit three benchmark classifiers. Two of these classifiers were variations of the MalConv and LightGBM models from the original EMBER paper, with minor modifications for a fair comparison. For MalConv, we limited the context to 4096 bytes, and for the GBM model, we modified the EMBER feature extractors to only extract from the first 4096 bytes of each file.
Note that 4096 bytes is the approximate size of a typical PE header. In practical adversarial settings, an attacker could bypass this model by ensuring that malicious byte sequences manifest beyond the extent of the file's initial 4096 bytes. However, real-world industry malware detection approaches commonly employ defensein-depth strategies in which this model would be one of many complementary detection methods to help counteract such an evasion. Nevertheless, future work towards extending this truncation window using other efficient Transformer architectures [9,12,19] represents a promising research path forward from this feasibility study.
We additionally employed a deep convolutional neural network that combines a 10-dimensional, learnable embedding layer with a series of five interleaved convolutional and max-pooling layers arranged hierarchically so that the original input size is reduced by one quarter (1/4) after each layer [10]. We refer to this as the MG-CNN.
In other papers, byte-based models often truncate inputs beyond a maximum length, while feature-based models use features extracted from the entire file irrespective of length, potentially presenting an overly optimistic view of feature-based classifiers. Our comparison, by contrast, is fair in the sense that inputs to each respective model are extracted using the same context window length. Results of this comparison are shown in Fig. 5. We see that our Byte Transformer and MalConv models perform comparably, but under-perform the MG-CNN and lightGBM models, which also perform comparably. This suggests that there is a motivation for the extra layers added in the MG-CNN over MalConv, and that much of the information learned during training may overlap with content extracted from EMBER features. This finding is consistent with [27], where the authors found that a convolutional model trained on EMBER largely uses portions of the header when discriminating betweeen malicious/benign.

Effects of Byte Encoding Schemes
Note that both our Byte Transformer and the MG-CNN share two common characteristics: First, their convergence is extremely sensitive to hyperparameter selection. Second, their practical application is somewhat hindered by limited context window sizes imposed by memory constraints (though the MG-CNN can feasibly incorporate far longer sequences). Since successful applications of transformers on raw data in the NLP space leverage specific encoding strategies, e.g., RoBERTa models leverage a byte-pair encoding [25], we explore whether similar encoding schemes can potentially benefit InfoSec byte-based models.
To shorten training times and perform this investigation using more realistic context windows, we performed the following comparisons on the full EMBER dataset using the MG-CNN model with each sample truncated at 102400 characters and different encoding algorithms: • Baseline (Baseline): MG-CNN model fit on raw bytes.
• Byte Pair Encoding (Byte Pair): For this, we performed 10 iterations, greedily adding 10 characters per iteration. This lead to a total vocab size of 396. We first truncated our input at 102400 bytes and then performed the byte pair encoding, which compressed each input at the expense of increased vocab size. • Removal of Common Padding Sequences (No Pad): We removed padding bytes for 0xFF and 0x00 for sequences longer than 3 subsequent padding bytes of the same type and replaced them with a sequence of length 3. As with the byte pair encoding, we first truncated at 100kB and then performed the compression. Surprisingly, none of our compression schemes improved performance beyond the baseline. This is specifically surprising, as the byte pair extra model introduces no information loss, yet incorporates longer contexts from the original sequence.

Effects of Kilo-Gram Encoding Schemes
In this section, we examine to what extent we can effectively extend the length of a context window by either removing common kilogram [33] sequences and keeping all other bytes, i.e., blacklisting kilograms, or by keeping only common kilo-gram sequences and removing all other bytes, i.e., whitelisting kilo-grams.
For these experiments, we first flagged kilograms for each sample up to 12288 bytes, performed whitelisting and blacklisting by keeping only or removing only these kilograms respectively, then truncated everything in each sequence beyond 4096 bytes. We chose 6-grams because these seemed to work well both in the kilograms paper and as features for some of our unpublished detection models.
Performance comparisons of baseline, whitelist, and blacklist models are shown in the Fig. 7. Each approach utilized an MG-CNN model.
Neither blacklisting nor whitelisting improved classification performance, despite the fact that broader byte information from larger context window spans get packed into the input for the MG-CNN. Since whitelisting yielded a better classifier than blacklisting, this potentially suggests that 1) broader context delivered by kilogram whitelisting is important and 2) kilo-grams associated bytes have greater signal content than non-kilogram associated bytes. This is consistent with the finding that kilograms serve as a useful feature representation for a variety of InfoSec problems.

DISCUSSION
We have demonstrated that transformers can achieve performance comparable to or better than that of other top-performing models for URL classification. We have also found that, contrary to NLP domains, wherein auto-regressive pre-training substantially enhances performance in a fine-tuned regime, for our URL data, auto-regressive pre-training on a large corpus yields no apparent gains for the classification task and makes it substantially more difficult to fit a performant model. This suggests that the next character prediction task has too little apparent correlation with the task of malicious/benign prediction for effective/stable transfer. Interestingly, utilizing next character prediction as an auxiliary loss function [37] yields improvements over training solely to predict the label. Note that this occurs even with a relatively large portion of the overall loss term (50 %) devoted to the auxiliary loss. This suggests that mixed objective optimization is more effective at correlating across heterogeneous loss terms than fine-tuning.
Note that in this work we did not perform rigorous hyperparameter searches for either the URL transformer or the byte transformer. Given a motivation to use transformers for these or other problems, frameworks like Optuna [1] could be utilized to arrive at a better selection of hyperparameters which could very likely improve performance for both tasks studied herein. However, such a hyperparameter search would require considerable compute time/resources, and would be, in many respects, beyond the scope of this paper. This paper aims to explore the viability of transformers for endto-end infosec tasks, not arrive at the some "optimal" transformer topology. This is one of the reasons why our architectural choices were strongly guided by well-studied approaches in the literature. We would also note that there are myriad other architectures in the literature that could be applied in our experimental settings.

CONCLUSION
We have demonstrated first steps for training transformers for InfoSec tasks from scratch and applying them to two heterogeneous InfoSec datasets: malicious/benign URLs and goodware/malware PE binaries.
While these datasets are not representative of all data in the ML for InfoSec space, they do reflect the common property of having a multitude of samples with labels derived from weak labeling sources, differing from common NLP classification tasks.
For URLs, we have shown that unlike in NLP domains, there is little benefit to pre-training, but substantial benefit from mixed objective optimization. To this end, we have introduced a novel loss function which dynamically re-balances gradients of auxiliary losses with the main task loss at each training step, potentially improving training stability. Utilizing this loss function, our URL transformer model performs on-par with our top-scoring benchmark model. Performance could potentially be further improved using a bi-directional information flow trained under an MLM-like regime.
For PEs, we have shown that we can successfully train adaptive span transformers from scratch on raw byte sequences, and for a limited context length perform on par with a Malconv model. Certain methods for extending this context length are trivial, e.g., via caching hidden states, and incorporating other scaling tricks could be feasible with some alterations (e.g., via a longformer-like locality-sensitive hashing approximation [8]), though these may lengthen training times and introduce additional complexity.
While we have laid out first principles for training transformers for InfoSec tasks, there is still much work required to make these models widely viable: First, transformers are extremely computeintensive compared to other models that operate on raw data. For the URL problem, a transformer would need to substantially outperform top-performing convolutional models in order to justify extra compute costs during training and deployment. While this is also the case for PE binaries or other byte sequence tasks, for these tasks, context length is also a problem that is difficult to fully address even with special tricks (sparse alterations to attention, adaptive span, hidden state caching, etc.).
We are presently working to extend our models to other data formats, including ways to incorporate even longer context and different training regimes that may improve classification performance.