# Neural Architecture Search on Efficient Transformers and Beyond

Zexiang Liu1, Dong Li1, Kaiyue Lu1, Zhen Qin1, Weixuan Sun3,1, Jiacheng Xu1,
Yiran Zhong1,2
1SenseTime Research 2Shanghai AI Lab 3Australian National University  Zexiang
Liu, Dong Li, and Kaiyue Lu contributed equally to this research.
Corresponding author: Yiran Zhong<EMAIL_ADDRESS>
###### Abstract

Recently, numerous efficient Transformers have been proposed to reduce the
quadratic computational complexity of standard Transformers caused by the
Softmax attention. However, most of them simply swap Softmax with an efficient
attention mechanism without considering the customized architectures specially
for the efficient attention. In this paper, we argue that the handcrafted
vanilla Transformer architectures for Softmax attention may not be suitable
for efficient Transformers. To address this issue, we propose a new framework
to find optimal architectures for efficient Transformers with the neural
architecture search (NAS) technique. The proposed method is validated on
popular machine translation and image classification tasks. We observe that
the optimal architecture of the efficient Transformer has the reduced
computation compared with that of the standard Transformer, but the general
accuracy is less comparable. It indicates that the Softmax attention and
efficient attention have their own distinctions but neither of them can
simultaneously balance the accuracy and efficiency well. This motivates us to
mix the two types of attention to reduce the performance imbalance. Besides
the search spaces that commonly used in existing NAS Transformer approaches,
we propose a new search space that allows the NAS algorithm to automatically
search the attention variants along with architectures. Extensive experiments
on WMT’14 En-De and CIFAR-10 demonstrate that our searched architecture
maintains comparable accuracy to the standard Transformer with notably
improved computational efficiency.

##  1 Introduction

Efficient Transformers [20, 17, 25, 31] have achieved remarkable advances in
recent years. They reduce the quadratic computational complexity of the
standard Transformer [35] by spasifying or approximating Softmax attention in
a more efficient fashion. Currently, the configuration of the efficient
network, _e.g_., the number of heads and the input embedding size, is directly
copied from the Transformer, which may not be suitable for efficient
Transformers [25]. The customized network structure specifically for efficient
Transformers has not been well studied. However, the manual design always
involves expensive engineering work and could be sub-optimal due to the human
bias [14]. Hence, in this paper, we aim to explore how to automatically find
an appropriate architecture for efficient Transformers.

Figure 1: Performance comparison in accuracy and efficiency. Models marked
with “*” are the optimal architectures that found by NAS. The size of the
circle represents the relative parameter scale on each task. The proposed
mixed use of Softmax and efficient attention, i.e., mFormer, is beneficial for
balancing accuracy and efficiency. Notice: (a) The cosFormer presents very low
BLEU on the translation task due to the inappropriate convergence. (b) The
vanilla Transformer and cosFormer in CIFAR-10 have a smaller FLOPS because
they only have 6 layers, while the searched optimal models (including ours)
have more than 10 layers. With similar FLOPS, ours has a much higher accuracy.

To achieve the goal, we need to (1) specify a suitable efficient Transformer
as the target, and (2) find an approach to automating the network design. For
the first point, we give a brief overview of existing efficient Transformers.
In terms of how to treat the Softmax attention, they can be roughly classified
into pattern-based and kernel-based [25]. Pattern-based methods sparsify the
attention matrix with predefined or learnable patterns, _e.g_., chunking input
sequences into fixed blocks [26], calculating attention at fixed intervals
[7], or using axial attention [15]. Although the sparsity generated by
specific patterns is beneficial for simplifying attention calculation, it
still suffers from the quadratic complexity with respect to the input length
$N$ [33] and is complicated to implement. Differently, kernel-based methods
aim at reducing the quadratic complexity into linear (i.e., $O(N)$ in both
time and space complexity) [25, 17]. They reformulate the self-attention
mechanism to avoid explicitly computing the $N\times N$ matrix. Moreover, they
are easier to reproduce in practice. Considering the intriguing properties of
kernel methods, we choose the cosFormer [25], a new kernel-based model with
state-of-the-art performance among efficient Transformers, as our target. We
attempt to seek for its customized and optimal architecture.

Figure 2: The proposed search space for attention types. It contains the
Softmax attention and linear attention (i.e., cosFormer attention [25] in this
paper). We let the model automatically determine which attention to use with
NAS.

To automate the network design, we take advantage of neural architecture
search (NAS) [11, 23]. It has been widely employed in searching standard
Transformer architectures in Natural Language Processing [29, 14, 38, 16, 36]
and Computer Vision [5, 13, 10, 6]. These studies mainly focus on refining
search space and/or improving search algorithms. For example, AutoAttend [14]
explores the effecitve connections between the Query, Key, and Value by
creating primitive operations to the input data, _e.g_., $1\times 1$
convolution and $1\times 3$ max pooling. The evolved Transformer [29] applies
tournament selection [27] to NAS to generate candidate models more robustly.
However, these methods suffer from long training and large search costs
because all the candidates need to be optimized, evaluated, and ranked. For
the purpose of lowering these costs, we utilize RankNAS [16], a new efficient
NAS framework for searching the standard Transformer [35]. It can
significantly speed up the search procedure through pairwise ranking, search
space pruning, and the hardware-aware constraint.

Given both the efficient Transformer (i.e., cosFormer [25]) and the search
algorithm (i.e., RankNAS [16]), we conduct a preliminary study on pure
efficient Transformer-based NAS. Specifically, we replace Softmax with the
linear attention introduced in the cosFormer, and then search it with RankNAS.
To comprehensively investigate the optimal architecture, we perform the search
on two representative tasks in NLP and CV fields, i.e., machine translation
(WMT’14 En-De [4]) and image classification (CIFAR-10 [19]). Generally, we
observe that the optimal structure of the cosFormer has fewer computational
costs, _e.g_., smaller FLOPS as illustrated in Fig. 1. However, the general
accuracy is less comparable to that of the standard Transformer (see the
vertical axis in Fig. 1). This performance imbalance between accuracy and
efficiency has also been revealed in other efficient Transformers [32, 17, 18,
37, 8].

Considering that the vanilla Softmax attention and linear attention have their
own distinctions regarding the performance, we propose to use Softmax and
linear attention in a mixed way for the benefit of a better balance between
accuracy and efficiency (named as “mFormer”). Moreover, we expect the model to
determine which type of attention to use automatically. To this end, we
introduce a new search space specially for attention and incorporate it into
the NAS framework. After re-searching the optimal architecture, we find that
the combination of the two attention types achieves comparable performance to
the Transformer with notably improved efficiency on both tasks (see Fig. 1).

In summary, we make three primary contributions:

  * •

To the best of our knowledge, it is the first work to search the optimal
architecture of efficient Transformers. We employ NAS for searching the
cosFormer, a representative efficient model. The searched results give new
insights to the community, i.e., how to design customized, effective, and
efficient network structures for efficient Transformers.

  * •

We propose a novel usage of attention, i.e., mixing Softmax attention and
linear attention in the Transformer, and define a new search space for
attention search in the NAS framework. This pushes existing Transformer search
further by enabling automatic selection of appropriate attention types.

  * •

The proposed mixed attention achieves better balance between accuracy and
efficiency, i.e., having comparable performance to the standard Transformer
while maintaining good efficiency. This is validated on both machine
translation and image classification tasks.

Figure 3: Illustration of attention. (a) Softmax calculates
$\mathbf{Q}\mathbf{K}^{\mathrm{T}}$ first, leading to the quadratic complexity
with respect to the sequence length $N$. (b) Linear attention decomposes the
similarity function with a kernel function, and multiplies $\mathbf{K}$ and
$\mathbf{V}$ first. This reduces the quadratic complexity to linear.

##  2 Related Work

###  2.1 Efficient transformers

The standard Transformer [35] suffers from the quadratic time and space
complexity caused by the Softmax attention. To reduce computational costs,
efficient Transformers either sparsify or approximate Softmax in a more
efficient fashion. Sparsifying the attention is adopted by the pattern-based
methods, where the attention matrix is sparsified with predefined or learnable
patterns. [26] chunks the input sequence into fixed blocks, so that the
complexity is reduced from $N^{2}$ to $B^{2}$ ($B$ is the block size smaller
than $N$). Alternatively, [22] downsamples the input into fixed length.
Instead of adjusting the sequence length, [7, 3] make use of strided/dilated
patterns to compute attention at fixed intervals. Compared with these methods
with one fixed pattern, the combination of multiple patterns can diversify the
coverage of attention access [33]. For example, [15] calculates the attention
along every axis of the input, and [7] aggregates the attention from both
strided and local patterns. To further improve the quality of patterns,
learning patterns in a data-driven manner is becoming a new trend [18, 28,
32]. The benefit over fixed patterns is that learnable ones are able to more
precisely cluster input tokens based on token relevance while still
maintaining the efficiency [33]. In summary, pattern-based methods can enhance
the efficiency by sparsifying the Softmax attention. However, they still have
the quadratic complexity, and it will increase with the input length becoming
larger. In addition to that, they are relatively complicated to reproduce in
practice.

Another general category of efficient Transformers is based on kernels (or
kernel functions), which aims to reduce the complexity from quadratic to
linear. By this means, Softmax can be re-written with other forms to avoid the
$N\times N$ attention matrix [33]. To achieve this, [37] assumes a low-rank
prior within the $N\times N$ structure, and transforms the Key and Value into
a lower dimension with extra projection layers. [24] approximates Softmax with
the production of a series of Gaussian kernels, and alternatively, [9] uses
Haar kernels instead. The PerFormer [8] employs orthogonal random features to
generate random kernels. The Linear Transformer [17] makes use of the
associative property of matrix products, and reformulates the similarity
function with kernel decomposition. The cosFormer [25] follows this
kernelization strategy, and utilizes ReLU [1] as the kernel function with an
additional cosine reweighting mechanism. Generally, kernel-based methods can
linearize the complexity and effectively enhance the efficiency. Moreover,
they are easier to implement than pattern-based approaches.

###  2.2 Neural architecture search (NAS)

NAS [11, 23] aims to automatically find the most appropriate network
architecture, and has been extensively applied to Computer Vision [5, 13, 10,
6] and Natural Language Processing [29, 14, 38, 16, 36]. The core of NAS is to
design a suitable search space, rank all candidate architectures generated
from the space, and find the optimal structure. For Transformers in NLP, the
Evolved Transformer (ET) [29] is the pioneering work that applies NAS into
Transformer architecture search. It defines a large-scale search space, which
covers the components in the input, normalization, layers, the output
dimension, activations, _etc_. An evolution algorithm [27] is adopted to make
architecture selection more stable and robust, i.e., consistently selecting
the most promising architecture based on the fitness. Although ET can find a
better and more efficient structure, its computational costs are still very
large. This is because the algorithm has to cover all the search features.
Besides, training the evolution process is also time-consuming.

Subsequent works focus on improving the search algorithm and/or pruning the
search space. [41] employs the differentiable architecture search (DARTS) [21]
that includes all the operations into a node (i.e., constituting a supernet)
and leverages Softmax for the specific choice. This method relaxes the need of
training each candidate network separately. [34] decomposes the Transformer
structure into smaller components, and introduces an one-shot search algorithm
based on sampling. HAT [36] introduces a hardware-aware constraint for
accelerating the search process. RankNAS [16] treats NAS as a pairwise ranking
problem, which significantly speeds up the search process.

Pruning the search space, i.e., only retaining the most important search
features, is another approach to reducing the search costs. TextNAS [38]
specifies a tailored search space for text representation, which consists of
convolutional, recurrent, pooling, and self-attention layers. The Primer [30]
modifies the search space with squaring ReLU activations and a depth-wise
convolutional layers after the Query, Key, and Value. AutoAttend [14]
specially searches the connections between QKV. RankNAS [16] proposes a
feature selection algorithm that measures the importance of each feature and
sorts out the most useful ones for further search. In summary, search space
pruning can effectively improve the search efficiency and maintain good
performance.

##  3 NAS on Efficient Transformers

In this section, we first give a brief review of the preliminary knowledge of
the cosFormer [25] and RankNAS [16]. After that, we specify the main steps of
applying NAS to the efficient Transformer search and discuss the general
results. Finally, we introduce the proposed idea of using mixed attention.

###  3.1 Preliminary knowledge

####  3.1.1 cosFormer

Given the Query $\mathbf{Q}\in\mathbb{R}^{N\times d}$, Key
$\mathbf{K}\in\mathbb{R}^{N\times d}$, and Value
$\mathbf{V}\in\mathbb{R}^{N\times d}$ ($d$ is the feature dimension of each
head), the standard Transformer [35] produces the output
$\mathbf{O}\in\mathbb{R}^{N\times d}$ by multiplying the normalized
$\mathbf{Q}\mathbf{K}$ attention to $\mathbf{V}$, i.e.,

$\mathbf{O}=\rho(\mathbf{Q}\mathbf{K}^{\mathrm{T}})\mathbf{V},$ (1)

where $\rho$ is the similarity function, _e.g_., Softmax in the Transformer.
The calculation of $\mathbf{Q}\mathbf{K}^{\mathrm{T}}\in\mathbb{R}^{N\times
N}$ leads to the quadratic complexity with respect to the sequence length $N$
(see Fig. 3(a)).

Figure 4: Search space in this paper. The blue blocks are primitive features
in Transformer search, and the gray ones denote the proposed attention type
search.

Linearizing attention can reduce the complexity from quadratic to linear [17,
25]. This is achieved by decomposing the similarity function $\rho$ and
multiplying $\mathbf{K}$ and $\mathbf{V}$ first, i.e.,

$\rho(\mathbf{Q}\mathbf{K}^{\mathrm{T}})\mathbf{V}=(\varphi(\mathbf{Q})\varphi(\mathbf{K})^{\mathrm{T}})\mathbf{V}=\varphi(\mathbf{Q})(\varphi(\mathbf{K})^{\mathrm{T}}\mathbf{V}),$
(2)

where $\varphi$ is a kernel function that transforms $\mathbf{Q}$ and
$\mathbf{K}$ into hidden representations. In this case, the complexity
$O(N^{2}d)$ lowers to $O(Nd^{2})$. Generally we always have $N\gg d$, so
$O(N^{2}d)\approx O(N^{2})$ and $O(Nd^{2})\approx O(N)$, yielding the linear
complexity.

In essence, it is recognized that the similarity function should be non-
negative and non-linear [25, 17]. The non-negativity strengthens the
positively correlated features, and [25] achieves this by directly utilizing
the ReLU function [1], i.e., $\mathrm{ReLU}(x)=\mathrm{max}(0,x)$. The non-
linearity concentrates the most relevant features to be aggregated and
stabilizes training [12, 2]. To this end, [25] further introduces a cosine-
based reweighting mechanism such that nearby tokens are encouraged to have
higher attention connection. The entire formulation is

$\mathbf{O}_{i}=\sum_{j=1}^{N}(\mathrm{ReLU}(\mathbf{Q}_{i})\mathrm{ReLU}(\mathbf{K}_{j})^{\mathrm{T}}cos(\frac{\pi}{2}\cdot\frac{i-j}{N}))\mathbf{V}_{j}.$
(3)

Here for simplicity, we ignore the normalization term. Eq. 3 can be re-written
as the linear form in Eq. 2, and we refer readers to the original paper [25]
for more details.

####  3.1.2 RankNAS

Conventional NAS methods [29, 14, 38] have to evaluate the performance of
every candidate network for ranking, which is time-consuming especially when
the search space is large. The recently-proposed RankNAS [16] can effectively
reduce training costs with the following three distinctions:

  1. 1.

Pairwise ranking. Instead of sorting all candidate architectures, RankNAS
treats NAS as a pairwise ranking problem. That is, a binary classification is
conducted to each candidate pair, and the label is simplified into correctly
ordered or incorrectly ordered as per their estimated performance.

  2. 2.

Search space pruning. RankNAS prunes the search space by only containing the
most important features that have large impact to the performance.

  3. 3.

Hardware-aware search. In addition to the standard search based on losses,
RankNAS also leverages a hardware constraint, i.e., estimating the latency and
discarding the fastest and slowest 10% models.

Considering the good efficiency of RankNAS, we take it as our NAS framework.
More technical details can be found in [16].

###  3.2 RankNAS on cosFormer

####  3.2.1 Search space

Primitive features (from RankNAS [16]) | Enc Layer Num | 6
---|---|---
Enc Emb Dim | 768, 1024
Enc FFN Dim | 2048, 3072, 4096, 5120
Enc Head Num | 4, 8, 16
Dec Layer Num | 1, 2, 3, 4, 5, 6
Dec Emb Dim | 768, 1024
Dec FFN Dim | 2048, 3072, 4096, 5120
Dec Head Num | 4, 8, 16
En-De Head Num | 4, 8, 16
En-De Connect | 1, 2, 3
Proposed feature | Enc Self-Attn | Softmax, Linear
Dec Self-Attn | Softmax, Linear
En-De Cross-Attn | Softmax, Linear
Table 1: Search space for WMT’14 En-De translation.

Our primitive search space is adopted from RankNAS, which contains fundamental
features including the embedding size, the number of encoder/decoder layers,
the number of heads, and the dimension of the feed-forward network (see Fig.
4). We also make corresponding modification to the space as per each task.
Detailed definition is introduced in Section LABEL:sec:setting.

####  3.2.2 Main steps

We replace all Softmax with the cosFormer attention, and follow RankNAS [16]
for the search process. A supernet containing all possible architectures is
firstly trained. Afterwards, the loss and latency data are separately
collected and ranked. Based on the loss and latency constraints, the most
important features are selected, i.e., search space pruning. The evolution
algorithm is performed on the refined search space, and the optimal
architecture is sorted out. Finally, we re-train the optimal network from
scratch.

####  3.2.3 Discussion

For comparison, we also repeat the above steps to the standard Transformer
[35] only with the Softmax attention. For simplicity, we denote the optimal
architecture of the cosFormer as “cosFormer$*$”, and that of the Transformer
as “Transformer$*$”. From Fig. 1, we find that on both tasks, the cosFormer$*$
presents better efficiency than the Transformer$*$ (smaller FlOPS) with a
comparable parameter scale, but the general accuracy is less competitive. We
name this phenomenon as the imbalance between accuracy and efficiency, which
has also been observed in other efficient Transformers [32, 17, 18, 37, 8].

In terms of the attention itself, the advantage of Softmax in accuracy is
associated with its ability in “imposing a categorical distribution constraint
on the query-context relevance scores” [40]. This constraint has two essential
properties: (1) dense coverage, i.e., attending every feature to the query
[35]; and (2) non-linear concentration, i.e., larger weights are only assigned
to more relevant features [12, 2, 25]. In spite of the useful properties, we
should also notice that the primary limitation of Softmax is its high
computational complexity, as analyzed in Section 3.1.

Figure 5: Detailed searched features on WMT’14 En-De. Layers with linear
attention are suffixed with “_m”. Dashed lines reflect the tendency of
changes.

Head number. We plot head numbers in Fig. 3.2.1(b)(d)(e). In general, the
cosFormer$*$ requires a smaller number of heads.

En-De Connect. From Fig. 3.2.1(f), we observe that decoders in both models
only attend the last layer of the encoder, i.e., obtaining the cross
information only from the last encoder block.

Parameters. The cosFormer$*$ has 12.6M fewer parameters than the
Transformer$*$.

####  4.2.2 Accuracy

On the WMT’14 En-De test set, the cosFormer$*$ presents extremely poor
BLEU†††It has the inappropriate convergence, and we are communicating with the
authors of [25] for dealing with this issue.. It reflects the limitation of
the linear attention in achieving comparable accuracy to the Transformer$*$.

We also investigate the importance of each feature to the accuracy, which is
measured by RankNAS [16]. The top-5 ranked features are:

  * •

cosFormer$*$: En-De Connect, Dec Layer Num, En-De Head Num, Dec FFN Dim, Dec
Head Num.

  * •

Transformer$*$: Dec Layer Num, En-De Head Num, Dec FFN Dim, Dec Head Num, Enc
FFN Dim.

Hence, when designing the architecture for the efficient Transformer, we need
to focus more on the encoder-decoder interaction, _e.g_., the number of
encoder layers to attend and the cross-attention heads. Both the cosFormer$*$
and Transformer$*$ are also sensitive to the decoder layer number, but the
best performance does not necessarily need the largest number of layers.

|  |  Transformer$*$ |  cosFormer$*$ |  mFormer (ours)
---|---|---|---|---
Searched architectures |  Enc Layer Num |  14 |  12 |  12
Enc Emb Dim |  384 |  448 |  384
Enc FFN Dim AVG |  1328 |  1419 |  1232
Enc Head Num AVG |  6.86 |  6.25 |  6.33
Searched results |  #Parameters (M) |  24.26 |  24.31 |  19.31
FLOPS (G) |  10.90 |  9.55 |  8.39
Accuracy (%) |  95.10 |  88.40 |  93.59
Table 4: Optimal architecture and performance comparison on CIFAR-10 image
classification. “AVG” means the average value. Figure 6: Detailed searched
features on image classification. Layers with linear attention are suffixed
with “_m”. Dashed lines reflect the tendency of changes.

####  4.2.3 Efficiency

Clearly, the cosFormer$*$ is more efficient, i.e., its FLOPS is 8.32G while
that of the Transformer$*$ is 9.86G. This difference is mainly due to the more
complex optimal structure of the Transformer. The Softmax attention also
increases the computational costs, which will be detailed in Section 4.5. The
top-5 ranked features regarding the efficiency are:

  * •

cosFormer$*$: Dec Layer Num, En-De Head Num, En-De Connect, Enc FFN Dim, Dec
Head Num.

  * •

Transformer$*$: Dec Layer Num, Dec Head Num, Enc Head Num, Enc FFN Dim, Dec
FFN Dim.

It is obvious that the number of decoder layers has the largest impact to the
efficiency. For the efficient Transformer, the encoder-decoder interaction
also plays an essential role in determining the computational costs. By
contrast, the Transformer$*$’s efficiency is less sensitive to the cross
interaction.

|  | #Parameters (M) | FLOPS (G) | BLEU/Accuracy (%)
---|---|---|---|---
Machine translation | mFormer | 75.59 | 8.085 | 25.79
mFormer (all linear) | 75.59 | 8.081 | 3.53
mFormer (all Softmax) | 75.59 | 8.09 | 27.95
Vanilla cosFormer [25] | 213.00 | 13.39 | Null
cosFormer$*$ | 79.78 | 8.32 | 3.05
cosFormer$*$ (all Softmax) | 79.78 | 8.34 | 27.83
Vanilla Transformer [35] | 213.00 | 12.7 | 28.4
Transformer$*$ | 92.38 | 9.86 | 28.34
Transformer$*$ (all linear) | 92.38 | 9.84 | 3.47
Image classification | mFormer | 19.31 | 8.39 | 93.59
mFormer (all linear) | 19.23 | 7.55 | 83.49
mFormer (all Softmax) | 19.33 | 8.67 | 94.35
Vanilla cosFormer [25] | 19.34 | 7.60 | 87.13
cosFormer$*$ | 24.31 | 9.55 | 88.40
cosFormer$*$ (all Softmax) | 24.42 | 10.66 | 94.80
Vanilla Transformer [35] | 19.31 | 8.07 | 88.70
Transformer$*$ | 24.26 | 10.90 | 95.10
Transformer$*$ (all linear) | 24.14 | 9.48 | 81.82
Table 5: Ablation study on attention types.

####  4.2.4 Summary

We summarize empirical hints to design an appropriate architecture for the
efficient Transformer on machine translation: (1) Fewer decoder layers; (2)
Focusing more on the encoder-decoder interaction. The encoder layers to attend
and the number of encoder-decoder heads are not necessarily too large; (3) The
encoder FFN dimension is more important to the efficiency while the decoder
FFN dimension is more important to the accuracy.

###  4.3 Optimal architecture comparison on image classification

The optimal structures of the cosFormer and Transformer significantly
outperform their vanilla versions in accuracy (see Table 5, which further
verifies the necessity and usefulness of NAS. The FLOPS is larger because
vanilla models only have 6 encoder layers. We compare optimal structures in
the following.

####  4.3.1 Network structure

Table 4 shows the optimal architectures of the cosFormer [25] and Transformer
[35]. Clearly, the Transformer$*$ has more layers, a smaller embedding size,
and a smaller averaged FFN dimension than the cosFormer$*$. From a general
view, the cosFormer prefers a “shallow and wide” structure while the optimal
architecture of the Transformer tends to be “deep and thin”.

FFN dimension. Fig. 3.2.1(a) displays the FFN dimensions in different layers.
From the tendency curve, we observe that the cosFormer$*$ has a relatively
smaller FFN dimension in intermediate layers. Differently, the dimension is
smaller in early layers in the Transformer$*$.

Head number. The head numbers in encoder layers are plotted in Fig. 3.2.1(b).
The cosFormer$*$ generally has more heads in early and top layers. By
contrast, the head numbers in the Transformer$*$ do not vary too much and
their values are large.

Parameters. The cosFormer$*$ and the Transformer$*$ have a similar parameter
scale, i.e., $\sim$24M.

####  4.3.2 Accuracy

On the CIFAR-10 test set, the cosFormer$*$ achieves the 88.4% accuracy, which
is surpassed by the Transformer$*$ (95.10%) by around 7.6%.

According to the importance to the accuracy, the features ranked by RankNAS
[16] are:

  * •

cosFormer$*$: Layer Num, FFN Dim, Emb Dim, Head Num.

  * •

Transformer$*$: FFN Dim, Head Num, Emb Dim, Layer Num.

Clearly, the number of layers has the largest impact to the accuracy of the
cosFormer. Similar to the case in machine translation, the cosFormer$*$ does
not select the greatest layer number.

####  4.3.3 Efficiency

The FLOPS of the cosFormer$*$ is 7.01G, which is significantly better than the
Transformer$*$ (10.9G) by 36%. It reveals the advantage of the cosFormer’s
linear attention in efficiency. The important features regarding the
efficiency are:

  * •

cosFormer$*$: Layer Num, Emb Dim, Head Num, FFN Dim.

  * •

Transformer$*$: Layer Num, FFN Dim, Head Num, Emb Dim.

Similar to the conclusion in machine translation, the layer number in image
classification is also the most related to the efficiency. The efficient model
has the preference of using fewer layers, which further reduces the
computational burden.

####  4.3.4 Summary

We give a brief summary of how to appropriately design the efficient
Transformer in image classification: (1) Fewer layers; (2) A smaller FFN
dimension in intermediate layers with a slightly larger embedding size; (3) A
smaller head number in intermediate layers.

###  4.4 Our results using mixed attention

####  4.4.1 Machine translation

Table LABEL:tab:mt_search shows the optimal structure and the performance of
our mFormer. In general, the mFormer is “shallow and thin”, which has the
fewest parameters and smallest FLOPS among the three optimal models. The
linear attention exists in Layer 2 and 4 in the encoder, and all the other
encoder/decoder layers retain the Softmax attention. In mFormer, the FFN
dimensions in the encoder undergo an “up-and-down” change, i.e., early layers
rely on smaller FFN sizes. The number of heads in the encoder tends to be
smaller in intermediate layers. The primary features in the decoder of the
mFormer consistently have the no larger quantity than the other two, except
that the head number in the first decoder layer is larger than that of the
cosFormer$*$.

In terms of the performance, the mFormer achieves comparable BLEU to the
Transformer$*$ with 18% fewer parameters and 18% smaller FLOPS. It also
significantly outperforms the cosFormer$*$ in accuracy with slightly better
efficiency. The results indicate that the mixed use of attention can
effectively facilitate the reduction of the imbalance between accuracy and
efficiency.

####  4.4.2 Image classification

We display the optimal architecture of the proposed mFormer in Table 4.
Intuitively, our structure is ”shallow and thin”, which is consistent with the
observation in machine translation. This is also reflected by the parameter
scale, i.e., we have around 20% fewer parameters than the other two. The
linear attention exists in Layer 3, 7 and 11 (accounting for 25%), and the
remaining attention types in other layers are all Softmax. The FFN dimensions
are large in early layers, and has a “down-up-down” tendency in subsequent
layers. The changes of head numbers seem to be inversely proportional to the
FFN dimension, i.e., larger FFN dimensions normally correspond to smaller head
numbers.

Remarkably, the mFormer outperforms the cosFormer$*$ to a large margin in
accuracy (by 5.19 in percentage) while maintaining good efficiency, i.e., it
only has 1.38G more FLOPS than the cosFormer$*$. Moreover, it achieves
comparable performance to the Transformer$*$ in accuracy (the latter surpasses
ours only by 1.51 in percentage) but is significantly more efficient (our
FLOPS is smaller than the Transformer$*$ by 2.51G, which is $\sim$23%). It
demonstrates that the use of mixed attention can achieve a better balance
between accuracy and efficiency on the image classification task.

###  4.5 Ablation study on attention types

Since the attention search is newly introduced in the NAS framework, we
specially study the influence of each attention type in optimal architectures.
This is to validate (1) the effectiveness of the proposed mixed attention in
balancing accuracy and efficiency, and (2) the inherent difference between
Softmax and linear attention in performance. The results on machine
translation and image classification are reported in Table 5.

####  4.5.1 All Softmax attention

We first uniform the mixed attention and linear attention in the optimal
architectures of the mFormer and cosFormer$*$ with Softmax. For our mFormer,
the accuracy is slightly improved but the FLOPS becomes larger, indicating
that the better accuracy is always along with the sacrifice in efficiency.
This performance imbalance is more obvious when we replace all the linear
attention in the cosFormer$*$ with Softmax, where the accuracy is
significantly enhanced with much degraded computational efficiency. Notably,
our model is able to achieve comparable accuracy to the performance with all
Softmax and keep satisfactory efficiency without dropping too much.

####  4.5.2 All linear attention

In this step, we replace the mixed attention in the mFormer and Softmax
attention in the Transformer$*$ uniformly with the linear attention. We notice
a significant performance drop in accuracy on both tasks. It demonstrates the
less competitive ability of the linear attention in yielding accurate results.
However, after the replacement, the efficiency gain is remarkable, which
further indicates the advantage of the efficient Transformer in reducing
computational costs. The mixed use of attention is beneficial for achieving
both comparable accuracy and efficiency.

##  5 Conclusion

In this paper, we utilize NAS to find the optimal architecture of efficient
Transformers (i.e., the cosFormer [25]). The searched architectures reveal
that the optimal structures of the efficient Transformer are relatively
lighter than those of the standard Transformer, _e.g_., the reduced parameter
scale and improved FLOPS. This provides useful insights to the community for
the appropriate design of efficient Transformers. However, the general
accuracy of efficient models is less competitive. Based on this observation,
we propose a novel usage of attention, i.e., employing the linear attention
and Softmax in a mixed manner in each layer. The searched optimal architecture
presents comparable accuracy to the Transformer and maintains as good
efficiency as the efficient Transformer. The proposed method supplies a new
direction in the Transformer study, i.e., taking advantage of both Softmax and
linear attention. In our future work, we will study the mixed attention on
large-scale pretrain models as well as other downstream tasks.

## References

  * [1] Abien Fred Agarap.  Deep learning using rectified linear units (relu).  arXiv preprint arXiv:1803.08375, 2018.
  * [2] Titsias RC AUEB et al.  One-vs-each approximation to softmax for scalable estimation of probabilities.  Advances in Neural Information Processing Systems, 29, 2016.
  * [3] Iz Beltagy, Matthew E Peters, and Arman Cohan.  Longformer: The long-document transformer.  arXiv preprint arXiv:2004.05150, 2020.
  * [4] Ondřej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, et al.  Findings of the 2014 workshop on statistical machine translation.  In Proceedings of the ninth workshop on statistical machine translation, pages 12–58, 2014.
  * [5] Boyu Chen, Peixia Li, Chuming Li, Baopu Li, Lei Bai, Chen Lin, Ming Sun, Junjie Yan, and Wanli Ouyang.  Glit: Neural architecture search for global and local image transformer.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12–21, 2021.
  * [6] Minghao Chen, Houwen Peng, Jianlong Fu, and Haibin Ling.  Autoformer: Searching transformers for visual recognition.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12270–12280, 2021.
  * [7] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers.  arXiv preprint arXiv:1904.10509, 2019.
  * [8] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, David Belanger, Lucy Colwell, et al.  Masked language modeling for proteins via linearly scalable long-context transformers.  arXiv preprint arXiv:2006.03555, 2020.
  * [9] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al.  Rethinking attention with performers.  arXiv preprint arXiv:2009.14794, 2020.
  * [10] Mingyu Ding, Xiaochen Lian, Linjie Yang, Peng Wang, Xiaojie Jin, Zhiwu Lu, and Ping Luo.  Hr-nas: searching efficient high-resolution neural architectures with lightweight transformers.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2982–2992, 2021.
  * [11] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter.  Neural architecture search: A survey.  The Journal of Machine Learning Research, 20(1):1997–2017, 2019\.
  * [12] Bolin Gao and Lacra Pavel.  On the properties of the softmax function with application in game theory and reinforcement learning.  arXiv preprint arXiv:1704.00805, 2017.
  * [13] Chengyue Gong, Dilin Wang, Meng Li, Xinlei Chen, Zhicheng Yan, Yuandong Tian, Vikas Chandra, et al.  Nasvit: Neural architecture search for efficient vision transformers with gradient conflict aware supernet training.  In International Conference on Learning Representations, 2021.
  * [14] Chaoyu Guan, Xin Wang, and Wenwu Zhu.  Autoattend: Automated attention representation search.  In International Conference on Machine Learning, pages 3864–3874. PMLR, 2021.
  * [15] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans.  Axial attention in multidimensional transformers.  arXiv preprint arXiv:1912.12180, 2019.
  * [16] Chi Hu, Chenglong Wang, Xiangnan Ma, Xia Meng, Yinqiao Li, Tong Xiao, Jingbo Zhu, and Changliang Li.  RankNAS: Efficient neural architecture search by pairwise ranking.  In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 2469–2480, Online and Punta Cana, Dominican Republic, Nov. 2021. Association for Computational Linguistics.
  * [17] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In International Conference on Machine Learning, pages 5156–5165. PMLR, 2020.
  * [18] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  In International Conference on Learning Representations, 2020.
  * [19] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [20] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh.  Set transformer: A framework for attention-based permutation-invariant neural networks.  In International Conference on Machine Learning, pages 3744–3753. PMLR, 2019.
  * [21] Hanxiao Liu, Karen Simonyan, and Yiming Yang.  Darts: Differentiable architecture search.  arXiv preprint arXiv:1806.09055, 2018.
  * [22] Peter J Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer.  Generating wikipedia by summarizing long sequences.  arXiv preprint arXiv:1801.10198, 2018.
  * [23] Yuqiao Liu, Yanan Sun, Bing Xue, Mengjie Zhang, Gary G Yen, and Kay Chen Tan.  A survey on evolutionary neural architecture search.  IEEE transactions on neural networks and learning systems, 2021\.
  * [24] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong.  Random feature attention.  In International Conference on Learning Representations, 2021.
  * [25] Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong.  cosformer: Rethinking softmax in attention.  In International Conference on Learning Representations, 2022.
  * [26] Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang.  Blockwise self-attention for long document understanding.  arXiv preprint arXiv:1911.02972, 2019.
  * [27] Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V Le.  Regularized evolution for image classifier architecture search.  In Proceedings of the aaai conference on artificial intelligence, volume 33, pages 4780–4789, 2019.
  * [28] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier.  Efficient content-based sparse attention with routing transformers.  Transactions of the Association for Computational Linguistics, 9:53–68, 2021.
  * [29] David So, Quoc Le, and Chen Liang.  The evolved transformer.  In International Conference on Machine Learning, pages 5877–5886. PMLR, 2019.
  * [30] David R So, Wojciech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V Le.  Primer: Searching for efficient transformers for language modeling.  arXiv preprint arXiv:2109.08668, 2021.
  * [31] Weixuan Sun, Zhen Qin, Hui Deng, Jianyuan Wang, Yi Zhang, Kaihao Zhang, Nick Barnes, Stan Birchfield, Lingpeng Kong, and Yiran Zhong.  Vicinity vision transformer.  arXiv preprint arXiv:2206.10552, 2022.
  * [32] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan.  Sparse sinkhorn attention.  In International Conference on Machine Learning, pages 9438–9447. PMLR, 2020.
  * [33] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler.  Efficient transformers: A survey.  arXiv preprint arXiv:2009.06732, 2020.
  * [34] Henry Tsai, Jayden Ooi, Chun-Sung Ferng, Hyung Won Chung, and Jason Riesa.  Finding fast transformers: One-shot neural architecture search by component composition.  arXiv preprint arXiv:2008.06808, 2020.
  * [35] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [36] Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han.  Hat: Hardware-aware transformers for efficient natural language processing.  In Annual Conference of the Association for Computational Linguistics, 2020.
  * [37] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  arXiv preprint arXiv:2006.04768, 2020.
  * [38] Yujing Wang, Yaming Yang, Yiren Chen, Jing Bai, Ce Zhang, Guinan Su, Xiaoyu Kou, Yunhai Tong, Mao Yang, and Lidong Zhou.  Textnas: A neural architecture search space tailored for text representation.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 9242–9249, 2020.
  * [39] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al.  Big bird: Transformers for longer sequences.  Advances in Neural Information Processing Systems, 33:17283–17297, 2020.
  * [40] Biao Zhang, Ivan Titov, and Rico Sennrich.  Sparse attention with linear units.  arXiv preprint arXiv:2104.07012, 2021.
  * [41] Yuekai Zhao, Li Dong, Yelong Shen, Zhihua Zhang, Furu Wei, and Weizhu Chen.  Memory-efficient differentiable transformer architecture search.  arXiv preprint arXiv:2105.14669, 2021.

Supplementary Material

##  6 Training Details

We specify training details of machine translation and image classification
tasks below. For those not listed, we use the default setting in RankNAS [16].

###  6.1 Machine translation

  * •

max relative length: 8

  * •

optimizer: adam

  * •

adam betas: (0.9, 0.998)

  * •

weight decay: 0.1

  * •

max tokens: 4400

  * •

criterion: label smoothed cross entropy

  * •

label smoothing: 0.1

  * •

min learning rate: $10^{-9}$

  * •

max update: 200000

  * •

warmup updates: 4000

  * •

lr scheduler: inverse_sqrt

  * •

lr: 0.0007

  * •

warmup init lr: $10^{-7}$

  * •

max lr: $1$

###  6.2 Image classification

  * •

max relative length: 14

  * •

optimizer: adam

  * •

adam betas: (0.9, 0.998)

  * •

weight decay: 0.05

  * •

max tokens: 100000

  * •

criterion: label smoothed cross entropy

  * •

label smoothing: 0.1

  * •

min learning rate: $10^{-7}$

  * •

max update: 140000

  * •

max sentences: 64

  * •

warmup updates: 976

  * •

lr scheduler: cosine

  * •

lr: $10^{-6}$

  * •

warmup init lr: $10^{-6}$

  * •

max lr: $10^{-3}$

##  7 Feature Importance in mFormer

With RankNAS [16], we can know the importance of each feature when searching
the proposed mFormer.

###  7.1 Machine translation

The top 5 features regarding the accuracy are: Dec Layer Num, En-De Head Num,
En-De Head Num, Enc Attn Type, and Dec FFN Dim. Similar to the Transformer$*$,
the most important feature is the number of decoder layers. In our optimal
structure, this number is 2, which indicates that the best performance does
not necessarily need many layers. Besides, the encoder-decoder interaction
also has large impact to the accuracy, which requires careful design. The
proposed attention type is important as well, further validating the
effectiveness of attention search.

The top 5 features regarding the efficiency are: Dec Layer Num, Dec Head Num,
En-De Connect, En-De Head Num, and Enc FFN Dim. Again, the number of decoder
layers has the largest impact to the efficiency. Note that the proposed
attention type is not contained here, which is because it only exists in
encoders with 6 layers and the efficiency gain is thus not very large.

###  7.2 Image classification

The top 5 features regarding the accuracy are: Enc Attn Type, Enc Layer Num,
Enc FFN Dim, Enc Head Num, and Enc Embed Dim. The top 5 features regarding the
efficiency are: Enc Layer Num, Enc Attn Type, Enc Embed Dim, Enc Head Num, and
Enc FFN Dim.

Clearly, on this task, our attention type plays a crucial role in both
accuracy and efficiency, which well verifies the effectiveness of the mixed
use of attention.

##  8 Linear Attention in Decoders

On the machine translation task, when we put the linear attention in decoder
layers, the BLEU values are all 0 (corresponding to three cases: 1) only
decoder self attention, 2) only encoder-decoder attention, and 3) both). We
also find that the output sentences have the “repeater” problem, i.e.,
repeating a single word for many times. We have not found the specific reason
for this issue, and we are communicating it with the authors in [25]. We infer
that the potential reason is the locality in the cosFormer attention [25] may
be more suitable for feature fusion in encoders.


* [33] Liu, S.-Q., Wang, Z., and Zhang, Y.  Variational bihamiltonian cohomologies and integrable hierarchies I: foundations.  Communications in Mathematical Physics. (2023)
  * [34] Liu, S.-Q., and Zhang, Y.  Jacobi structures of evolutionary partial differential equations.  Adv. Math. 227, 73–130 (2011)
  * [35] Liu, S.-Q., and Zhang, Y.  Bihamiltonian cohomologies and integrable hierarchies I: a special case.  Comm. Math. Phys. 324, 897–935 (2013)
  * [36] Lorenzoni, P., Shadrin, S., and Vitolo, R.  Miura-reciprocal transformations and localizable Poisson pencils.  arXiv preprint arXiv:2301.04475 (2023).
  * [37] Lorenzoni, P., and Vitolo, R.  Weakly nonlocal Poisson brackets, Schouten brackets and supermanifolds.  J. Geom. Phys. 149, 103573 (2020)
  * [38] Oevel, W., Fuchssteiner, B., Zhang, H., and Ragnisco, O.  Mastersymmetries, angle variables, and recursion operator of the relativistic toda lattice.  J. Math. Phys. 30, 2664–2670 (1989)
  * [39] Pavlov, M.  Conservation of the ”forms” of the Hamiltonian structures upon linear substitution for independent variables.  Math. Notes 57, 489–495 (1995)
  * [40] Rogers, C., and Shadwick, W. F.  Bäcklund transformations and their applications, vol. 161.  Academic press New York (1982)
  * [41] Shabat, A., and Zakharov, V.  Exact theory of two-dimensional self-focusing and one-dimensional self-modulation of waves in nonlinear media.  Sov. Phys. JETP 34, 62–70 (1972)
  * [42] Strachan, I. A. B., and Stedman, R.  Generalized Legendre transformations and symmetries of the WDVV equations.  J. Phys. A 50, 095202 (2017)
  * [43] Suris, Y. B.  The problem of integrable discretization: Hamiltonian approach, vol. 219.  Birkhäuser (2012)
  * [44] Tsarev, M.  The geometry of Hamiltonian systems of hydrodynamic type: The generalized hodograph method.  Math. USSR Izv. 37, 397–419 (1991)
  * [45] Verosky, J. M.  Negative powers of Olver recursion operators.  J. Math. Phys. 32, 1733–1736 (1991)
  * [46] Xue, T., and Zhang, Y.  Bihamiltonian systems of hydrodynamic type and reciprocal transformations.  Lett. Math. Phys. 75, 79–92 (2006)

Si-Qi Liu,

Department of Mathematical Sciences, Tsinghua University
Beijing 100084, P.R. China
<EMAIL_ADDRESS>
Zhe Wang,

Department of Mathematical Sciences, Tsinghua University
Beijing 100084, P.R. China
<EMAIL_ADDRESS>
Youjin Zhang,

Department of Mathematical Sciences, Tsinghua University
Beijing 100084, P.R. China
<EMAIL_ADDRESS>


remarkRemark hypothesisHypothesis claimClaim Leveraging Sharing Communities to
Achieve …F. W. Bentrem, M. A. Corsello, and J. J. Palm SIAM2021_supplement

# Leveraging Sharing Communities to Achieve Federated Learning for
Cybersecurity††thanks: Submitted to the editors March 5, 2021; accepted for
publication March 25, 2021; revised April 15, 2021. This work was funded by
the Defense Advanced Research Projects Agency under contract no.
W911NF-18-C-0019 through subcontract under the University of Virginia.

Frank W. Bentrem Commonwealth Computer Research, Inc., Charlottesville, VA (,
, , https://ccri.com<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>Michael A. Corsello22footnotemark: 2 Joshua J.
Palm22footnotemark: 2

###### Abstract

Automated cyber threat detection in computer networks is a major challenge in
cybersecurity. The cyber domain has inherent challenges that make traditional
machine learning techniques problematic, specifically the need to learn
continually evolving attacks through global collaboration while maintaining
data privacy, and the varying resources available to network owners. We
present a scheme to mitigate these difficulties through an architectural
approach using community model sharing with a streaming analytic pipeline. Our
streaming approach trains models incrementally as each log record is
processed, thereby adjusting to concept drift resulting from changing attacks.
Further, we designed a community sharing approach which federates learning
through merging models without the need to share sensitive cyber-log data.
Finally, by standardizing data and Machine Learning processes in a modular
way, we provide network security operators the ability to manage cyber threat
events and model sensitivity through community member and analytic method
weighting in ways that are best suited for their available resources and data.

###### keywords:

cybersecurity, federated learning, streaming machine learning

68Q85, 68T05, 68U01

##  1 Introduction

Computers have been subject to malicious attacks for many years. However, the
sophistication and impacts of those attacks have grown dramatically.
Successful, timely detection of such attacks is essential to enable defenders
to prevent or respond to an attack to mitigate or minimize damage. The
consequences from such attacks include substantial financial losses, leaked
personal data, and the release of proprietary information or state secrets.
The immense human and computational resources required to analyze network logs
is prohibitive and often still misses novel attacks. As a result, there is a
growing need to incorporate advanced techniques such as machine learning in
the cyber defender’s toolbox. However, the cyber domain presents key
challenges to traditional batch machine learning techniques [1, 2].

Problem 1.1 The continuous evolution of threat vectors requires continual
learning.

Problem 1.2 The volume and sensitivity of network log data prevents long-term
retention and sharing across organizations.

Problem 1.3 Barriers-to-entry prevent small or more vulnerable organizations
from adopting complex or expensive technology.

Problem 1.4 Feedback from network security operators is important to identify
novel attacks and mitigate as quickly as possible.

In this paper, we present solutions to these problems through an architecture
that (1) is streaming, (2) is federated, and (3) manages feedback. The paper
is organized in the following way. Our targeted constraints are in Section 2,
our system architecture is in Section 3, the implemented machine learning
algorithms are in Section 4, and the summary follows in Section 5.

##  2 Targeted Constraints

Individual organizations that train models to recognize cyber threats are
severely data limited, however, privacy concerns prevent them from directly
sharing raw network data logs. Even anonymized data may collectively reveal
sensitive information. We therefore propose a federated learning system [3]
where trained standardized models are shared across organizations through
their parameters only without sharing any raw data. The models are then merged
in a way to improve performance compared to individually trained models.
Details on the sharing and merging of trained models are provided in Section
4.

In order to entice a large group of organizations to participate in a
collaborative learning enterprise, we lower barriers-to-entry by minimizing
the computational requirements regarding platform, bandwidth, processing
power, and storage. Our design choices were guided by this constraint, where
we preferred small algorithms, minimal sharing, and streaming processing. Some
modifications to the machine learning algorithms designed for traditional
batch processing are needed to enable streaming learning. (For a similar
approach see[4].)

##  3 System Architecture

To achieve the simultaneous goals of supporting real-time analytics of
streaming log data and ad hoc offline analytics, we propose a simple directed
acyclic graph (DAG) computational model. This model assumes a well-defined
data model to facilitate “pluggable” analytics and multi-organizational data
federation. The compute flow is based upon a small number of autonomous boxes
that are connected in a specific, simple way as illustrated in Fig. 1.

Figure 1: Core computation DAG.

The flow of data through the DAG ensures data at each stage is consistent and
provides guarantees of how, when and where data is both shared and integrated.
Sharing occurs at “stores” where sharable data is integrated across
organizations, as shown in Fig. 2, in the live stream using well-known model
combiners specific to each analytic shared. In this way, the development of a
new analytic involves the creation of a local training processor, a model
combining routine, and a prediction processor.

Figure 2: Computation DAG extended with feedback stores where sharing occurs.

Extending the sharing through the use of communities will enable data to be
aggregated across community members and then re-shared as the “community
consensus” for further integration at member sites.

##  4 Machine Learning Algorithms

Finally, any number of prediction algorithms can be plugged into the streaming
framework along with methods to merge the separately trained models. We built
a prototype for our described system using three example classes of machine
learning algorithms, Example 4.1 (neural network), Example 4.2 (naive Bayes),
and Example 4.3 (Random Forest), which illustrate how a large number of
different classification algorithms might be implemented. The input data for
the machine learning algorithms were obtained using existing raw HTTP data
logs that were preprocessed to obtain 81-dimensional vectors $\mathbf{x}$
(features), most of which were processed according to the method described by
Oprea, et al.[5] Limited malicious labels were derived from VirusTotal and
FireEye™ where available, and the host traffic ranking was used as a proxy for
benign labels. The streaming federated algorithm with scheduled model sharing
is described in Algorithm 1.

Algorithm 1 Streaming federated algorithm

while not past time for scheduled sharing do

Receive HTTP record $R$

Preprocess to get input vector $\mathbf{x}$

if $R$ contains label information then

Extract labels $\mathbf{y}$

Perform training iteration with $\mathbf{x}\xrightarrow{}\mathbf{y}$

else

Perform prediction on $\mathbf{x}$

end if

end while

Exchange model parameters with collaborators

Merge models

Repeat from top

Example 4.1 (Neural Network) To maintain a small model size, for our neural
network model, we avoided deep learning and implemented a multilayer
perceptron with five hidden layers of size [64, 32, 16, 8, 4] and output size
two for the malicious and benign scores. To merge neural network models we
perform a weighted average for the weights and biases of each node in the
network. So for each weight $w_{i}$ and bias $b_{i}$, where $i$ is the model
index, we obtain the merged weight $w^{\prime}$ and bias $b^{\prime}$ through
the weighted averages

(1) $w^{\prime}=\langle\mathbf{a},\mathbf{w}\rangle\hskip
10.0pt\text{and}\hskip 10.0ptb^{\prime}=\langle\mathbf{a},\mathbf{b}\rangle,$

where $\mathbf{a}$ is the (averaging) weights for the shared models with
$||\mathbf{a}||_{1}=1$, and $\mathbf{w}$ and $\mathbf{b}$ are the the weights
and biases for all shared models at a given node.

Example 4.2 (Naive Bayes) For the Naive Bayes algorithm, we generate two
histograms $\mathbf{h}_{k,i}$ for each feature $i$, one histogram for the
benign labels and one for the malicious labels, where $k\in\\{\text{benign,
malicious}\\}$. Let $N_{k}=||\mathbf{h}_{k,i}||_{1}$ represent the cumulative
numbers of benign and malicious records in the histograms. Then the benign and
malicious likelihoods $\mathcal{L}_{k}$ and evidence $\mathcal{E}$ for each
feature $i$ are given by

(2) $\mathcal{L}_{k,i}=\frac{\mathbf{h}_{k,i}}{N_{k}}\hskip
10.0pt\text{and}\hskip
10.0pt\mathcal{E}_{i}=\frac{\sum_{k}\mathbf{h}_{k,i}}{\sum_{k}N_{k}},$

respectively. The probabilities $p(k\mid\mathbf{x})$ for a log record with
feature vector $\mathbf{x}=(x_{1},x_{2},\ldots,x_{n})$ being benign
($k=\text{benign}$) or malicious ($k=\text{malicious}$) are given by

(3)
$p(k\mid\mathbf{x})=\frac{N_{k}}{\sum_{k}N_{k}}\prod_{i=1}^{n}\frac{\mathcal{L}_{k,i}(x_{i})}{\mathcal{E}_{i}(x_{i})},$

where $\mathcal{L}_{k,i}(x_{i})$ are the likelihoods and
$\mathcal{E}_{i}(x_{i})$ is the evidence for the histogram bin corresponding
to the value $x_{i}$ of feature $i$. The proof for Eq. 3 is given in
LABEL:sec:nb_proof. Model merging is accomplished by simply summing the
histograms from the shared models.

Example 4.3 (Random Forest) As an example of an ensemble model, the Random
Forest algorithm implements some number $m$ of constituent models that are
combined for prediction. The Random Forest algorithm is trained in the normal
way. To merge the shared ensembles, we select a sampling of constituent models
from each shared ensemble so that we maintain $m$ models in the merged
ensemble.

##  5 Summary

We presented an automated cyber threat detection system that leverages sharing
communities for collaborative, federated learning with a streaming
architecture. The computation DAG addresses the need for data privacy and low
barriers-to-entry. We completed a prototype with streaming machine learning
algorithms with the capability to merge shared models. The results from our
internal testing demonstrate the feasibility of our approach and the
effectiveness of sharing and merging shared models, however a detailed
performance analysis is contingent on more robust labeling.

## Acknowledgments

We would like to acknowledge the University of Virginia PCORE-CHASE team for
their support.

## References

  * [1] G. Apruzzese, M. Colajanni, L. Ferretti, A. Guido, and M. Marchetti, On the effectiveness of machine and deep learning for cyber security, in 2018 10th international conference on cyber Conflict (CyCon), IEEE, 2018, pp. 371–390.
  * [2] M. Conti, T. Dargahi, and A. Dehghantanha, Cyber threat intelligence: challenges and opportunities, in Cyber Threat Intelligence, Springer, 2018, pp. 1–6.
  * [3] J. Konečnỳ, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon, Federated learning: Strategies for improving communication efficiency, arXiv preprint arXiv:1610.05492, (2016).
  * [4] J. Montiel, M. Halford, S. M. Mastelini, G. Bolmier, R. Sourty, R. Vaysse, A. Zouitine, H. M. Gomes, J. Read, T. Abdessalem, and A. Bifet, River: machine learning for streaming data in python, 2020, https://arxiv.org/abs/2012.04740.
  * [5] A. Oprea, Z. Li, R. Norris, and K. Bowers, Made: Security analytics for enterprise threat detection, in Proceedings of the 34th Annual Computer Security Applications Conference, 2018, pp. 124–136.

See pages - of CCRi_SIAM2021Supplement.pdf