abstract: Lane detection is to detect lanes on the road and provide the accurate location and shape of each lane. It severs as one of the key techniques to enable modern assisted and autonomous driving systems. However, several unique properties of lanes challenge the detection methods. The lack of distinctive features makes lane detection algorithms tend to be confused by other objects with similar local appearance. Moreover, the inconsistent number of lanes on a road as well as diverse lane line patterns, e.g. solid, broken, single, double, merging, and splitting lines further hamper the performance. In this paper, we propose a deep neural network based method, named LaneNet, to break down the lane detection into two stages: lane edge proposal and lane line localization. Stage one uses a lane edge proposal network for pixel-wise lane edge classification, and the lane line localization network in stage two then detects lane lines based on lane edge proposals. Please note that the goal of our LaneNet is built to detect lane line only, which introduces more difficulties on suppressing the false detections on the similar lane marks on the road like arrows and characters. Despite all the difficulties, our lane detection is shown to be robust to both highway and urban road scenarios method without relying on any assumptions on the lane number or the lane line patterns. The high running speed and low computational cost endow our LaneNet the capability of being deployed on vehicle-based systems. Experiments validate that our LaneNet consistently delivers outstanding performances on real world traffic scenarios.
author: Ze Wang[^1]
Department of Electrical and Computer Engineering, Duke University
`email@example.com`; Weiqiang Ren
Horizon Robotics, Inc.
`email@example.com`; Qiang Qiu
Department of Electrical and Computer Engineering, Duke University
`firstname.lastname@example.com`
bibliography: egbib.bib
title: LaneNet: Real-Time Lane Detection Networks for Autonomous Driving

# Introduction

Optical image based lane detection method is a key component of modern driving assistant systems. However, the detection of lanes remains challenging for several reasons. First, the appearance of a lane is typically extremely simple, which provides no complicated or distinctive features for lane detection, and dramatically increases the risk of false positive detections. Furthermore, diverse lane patterns, such as solid, broken, splitting, and merging lanes make independent lane modelling difficult. Algorithms based on hand-crafted features can only solve the lane detection in limited scenarios. And most existing methods also require strict assumptions on lanes, e.g. lanes are parallel , lanes are straight or close to straight , which are not always valid especially in urban situations. In recent year, some methods have been proposed to address lane detection under few assumptions to the lanes, yet there is still a large space for further progress on the robustness to diverse real-world scenarios. Methods based on deep neural networks , especially convolutional neural networks stimulate a promising research direction and also inspire the idea of our LaneNet. Moreover, considering that the lane detection runs on vehicle-based systems, where computation resources are severely limited, the computational cost of a lane detection method should also be considered as a key indicator of the overall performance.

In this paper, to strive for a generalized, low computational cost, and real-time vehicle-based solution, we propose a lane detection method called LaneNet. The proposed LaneNet breaks down the lane detection task into two stages, i.e. lane edge proposal and lane line localization, respectively; and each involves an independent deep neural network. In the lane edge proposal stage, the proposal network runs binary classification on every pixel of an input image for generating lane edge proposals, which are served as the input to the lane line localization network in the second stage.

The neural networks in both stages are designed for high accuracy, low computational cost, and high running speed. Specifically, a light-weight encoder-decoder architecture is adopted for lane edge proposal, where stacked depthwise separable convolution and $1 \times 1$ convolution layers are used for fast feature encoding, and non-parametric decoding layers for fast feature resolution recovery. The obtained proposal map is then transformed to lane edge coordinates and fed to the second stage, where, a high-speed lane line localization network, consisting of a point feature encoder and a LSTM decoder, localize the lane lines robustly under various scenarios.

Such two-stage design of LaneNet brings additional desirable properties. Firstly, the lane edge map produced by the proposal network serves as interpretable intermediate features, which to some extent alleviates the impact of the black-box property of neural network based method, and makes the detection failures more trackable. The two-stage process allows the parameters of lane line localization network to be refined in a weakly-supervised manner which alleviate the strong demand for well-annotated training samples. Additional, an efficient dimensionality reduction is performed when transforming the lane edge proposal map to the lane edge coordinates between the two stages, which further reduces the complexity and the network scale of the lane line localization network, and speed up the detection without any compromise to the accuracy. Last but not the least, the function of lane edge proposal network can be integrated into semantic segmentation network and further reduces the overall computational cost of the driving assistant systems.

All these advantages mentioned above contribute to a robust and reliable lane detection for autonomous driving systems and driving assistant systems. Extensive experiments as well as comparisons are conducted to validate our LaneNet as a lane detection method with remarkable accuracy and fast speed.

The remaining of this paper is structured as follows: In Section  we concisely describe the architecture of our model. We provide the details regarding both the training data and the training process in Section . Experiments and results are shown in Section . Section  provides a brief overview and summary on other lane detection methods. We conclude the paper with Section .

# Lane proposal and localization networks

In this section, we first concisely describe two neural networks designed to together accomplish lane detection in diverse scenarios. The first proposal network detects the edges of lane marks, and produce a pixel-wise lane edge proposal map. The second localization network determines the localization and shape of each lane based on lane edge proposals. Follows the work of , our LaneNet only takes the Inverse Perspective Mapping (IPM) of an image as input, which helps get rid of the perspective effect in the image, and so lanes that appear to converge at the horizon line are now vertical and parallel. However, our LaneNet does not rely on the assumption that the lanes are parallel (or close to parallel) so that the robustness is guaranteed. The two proposed networks play different roles in lane detection, yet they share several common design principles. Firstly, different from those algorithms running on servers or even computation centers, lane detection runs on vehicles so that the computational cost should be reduced as possible while the inference should be fast enough for real-time decision making. Additionally, the prediction accuracy and the robustness to diverse scenarios are crucial. The details of our LaneNet and how to meet the aforementioned principles are discussed in this section.

## Lane edge proposal network

To filter irrelevant information contained in optical images, we first process an image by a lane edge proposal network to generate a binary lane edge proposal map indicating the position of the pixels likely lie on the edge of lane segments. The reason we adopt a deep neural network instead of using hand-craft features to proposal maps is that it is essentially difficult to find one or a set of robust hand-craft features for filtering lane edges in an input image. As an illustration, Fig.  shows a typical IPM image of the front view of a vehicle. As marked by red circles, road marks, characters on the roads, guardrails along the roads, and even the light reflected by other vehicles share similar local appearances with lanes, thus are easy to be detected as false positives. To derive accurate lane edge proposal on each pixel, contextual information as well as global appearance should all be taken into consideration to suppress false positive detections caused by similar local appearances. Deep convolutional neural network thus become a desirable tool thanks to its powerful feature extraction capacity.

Encoder-decoder architectures are widely used in dense prediction tasks like semantic segmentation , which typically utilize convolutional layers and deconvolutional layers (transpose convolution layers) for feature encoding and decoding. For a higher efficiency, our lane edge proposal network adopts a light-weight encoder-decoder architecture. The encoder takes an IPM image of the front view of a vehicle as the input, and hierarchically extracts the features. The decoder progressively recovers the resolution of the feature map and produce a pixel-wise lane edge proposal map.

#### Encoder

The encoder of the proposal network replaces the standard convolution operations by the combination of depthwise separable convolutions and pointwise convolutions ($1 \times 1$ convolutions) to significantly lower the computational cost as in . Specifically, depthwise separable convolutional layers with a kernel size of 3 are stacked for progressive feature extraction. Each depthwise separable convolution layer is followed by a $1 \times 1$ convolution layer for channel-wise information aggregation. As mentioned above, there exists a lot of objects that share similar local appearance with lanes in the input images, to prevent false positive detections, the context information should be properly extracted and preserved in the encoding stage. In order to enlarge the reception field of the encoder while guarantee the low computational cost, we adopt dilated convolutional kernels in the encoder layers. Specifically, as plotted in Fig.  (a), three depthwise convolution layers that followed by a $1 \times 1$ convolution layer are included to form a convolution block that is used for feature extraction on one particular feature resolution. Among the three depthwise separable convolutional layers, the first one has a dilation rate of 1, which corresponds to a standard separable convolution, while the following two layers use dilation rate of 2 and 4, respectively, to enlarge the reception field. Introducing dilated kernels successfully improves the valid reception field without any additional parameters or computational cost, which properly balances the efficiency and effectiveness.

#### Decoder

In order to recover the feature resolution and produce the pixel-wise lane edge map of an input image, we design a decoder architecture that follows the encoder. Deconvolution layers (transpose convolution layers) are widely used to enlarge the intermediate features in modern deep neural networks. However, the computational cost and the training difficulty make deconvolution layers an inappropriate choice in our model. Inspired by the success of adopting sub-pixel convolution layers in image super-resolution , we use sub-pixel convolution layers here in our model to progressively recover the feature resolution. It brings the beneficial properties including completely parameter-free and no computational cost, which are much desirable in our case of lane detection. Skip connections are used to provide high resolution feature for accurate lane edge localization. The final high resolution feature map is then used to generate the pixel-wise lane edge probability map and the value of each pixel indicates the confidence that the network decides this pixel lies on the edge of a lane segment in the input image.

## Lane line localization network

The lane edge proposal network projects an image to the corresponding binary lane edge proposal map where structure and appearance are far simpler, while detecting the lane lines from a lane edge proposal map remains sophisticated. The diverse patterns of lanes, e.g., splitting lanes, merging lanes, intersection, and uncertain total number of lanes challenge the detection. We complete our lane line detection method with a light-weight, ultra fast, yet powerful neural network referred to lane line localization network. Our lane line localization network also adopts an encoding-decoder structure. It takes the coordinates of the lane edges as input, and apply a series of 1D convolution and pooling operations to encode an input to a low dimensional holistic representation. A long short-term memory network (LSTM) based decoder is used to progressively decode the representation to the parameters of each lane in the image.

Specifically, given a binary lane edge map, we use coordinates of lane edge points instead of directly inputting the lane edge map to the lane line localization network. Such a design brings us two major benefits: First is that the size of the input is reduced from $w \times h$ to $n \times 2$, where n is the number of the lane edge points in the binary lane edge map. This reduction allows the compact architecture of our model and the ultra fast prediction. Secondly, such coordinate from of inputs enable a weakly supervised training method of the lane line localization network, which will be discussed in section .

Apart from the benefits, using coordinates, instead of lane edge map as model inputs comes with an additional requirement to the model, which is the invariance to the input order. Since the lane edge points are sampled from 2D images, it's hard to find a proper and reasonable rule to sort the points. And considering that in most cases, the positions of the lane edge points are sufficient to determine the number and the shape of lanes on the road, the prediction of our model should be invariant to the input point order so that the prediction remains consistent regardless any shuffle to the order of the input points. Inspired by the recent progress on novel deep neural network architecture PointNet which is specifically designed for analysing point cloud data, an input order insensitive encoder is designed in our lane line localization network to encode the input points to a low dimensional holistic representation of the lane edges. Specifically, to achieve invariance to the input order, shared fully-connected layers are utilized to perform non-linear projection on the feature of each point simultaneously. This operation can be simply implemented by a series of 1D convolutional layers whose kernel size is identical to the size of the feature of each point. Max pooling layer is adopted to aggregate the feature of all the points to achieve a global information fusion which is also insensitive to the input order. Different from the model proposed in PointNet, we do not use any transformation to the input or the features, and the max pooling layer is used for multiple times to encourage the information communication among all the points under multiple semantic levels. Specifically, we refer a series of non-linear feature extraction to each point and a pooling based feature aggregation to a stage in encoder network as shown in Fig.  (b), and multiple stages are used in our network for effective information communication among all the points and the final representation is drawn by applying average pooling to the final feature.

After the input coordinates of the points are encoded to a low-dimensional representation, we adopt LSTM, which serves as a solution to the uncertain number of lanes, to decode the parameters of each lane. Specifically, we fit each lane in the image by a quadratic function, and use LSTM to progressively predict the parameters depicting each lane along with a confidence score from the left side of the image to the right side of the image one by one. The prediction is terminated when the confidence score is lower than a predefined threshold which indicates the network believes there is no more lane in the image.

# Training

## Lane edge proposal network

The training of the lane edge proposal network is straightforward. In our case, the model takes an image of the front view of a vehicle as input, and outputs a lane edge probability map of the same size as the input image. All the parameters are trained in a completely supervised manner. For each training image, we provide an annotation map where '1' indicates this pixel is a positive point laying on the edge of one lane segment, and '0' elsewhere. The entire network is end-to-end optimized by stochastic gradient descent. Note that in every annotation map, the number of positive points are always far smaller than number of negative points. So in the training process, we dynamical weight the loss of positive points and negative points according to their relative amounts, and the loss function is written as: $$\label{}
%\small
l_{cls}= \sum_{i}^{B} \sum_{n}^{k} y_{n}^{i} \times \log P_{n}^{i} + \beta (1 - y_{n}^{i}) \times \log(1 - P_{n}^{i})$$ where $B$ is the batch size of one iteration in training process and $P^{i}$ denotes the prediction of the $i$-th sample in this mini-batch. $P_{n}^{i}$ indicating the $n$-th pixel in the image, and $k$ is total number of pixel in the image which equals to $w \times h$. $\beta$ is the ratio of the number of positive points to the number of negative points, which acts as a dynamic weight for balancing the losses from negative points and positive point. It could be computed by $$\label{}
\beta = \frac{\sum y^{i}}{w \times h - \sum y^{i}}$$ for simplicity.

After training, the network directly produces a lane edge proposal map of the input image where the value of each pixel indicates the confidence of this pixel lies on the edge of one lane segment.

## Lane line localization network

In our data, each lane line is represented as a quadratic function, i.e., the shape of each lane is annotated by three continuous numbers. In practice, however, training the network to predict the parameters of each quadratic function works poorly as each parameter is usually of significant different orders of magnitude, e.g. the quadratic term of a lane is nearly zero yet the constant term is thousands times larger, which makes the construction of a loss function difficult. Simply normalize the parameters results in disappointing performance as well. Therefore, we solve the problem by training the network to predict the point locations where a lane intersects with the top, middle and the bottom line of an image, respectively. Fig.  shows the way we construct the training objectives for each lane. In our case, the input image has a size of $h \times w$, so the values we train the network to predict are the three $X$ coordination values of the points which lies on the intersections of the lane with three horizontal lines $Y=0$, $Y=h/2$, and $Y=h$. We refer these three values to the keys values of a lane. By transferring the training objective from the parameters of a quadratic function to the key values, the predicted values are of relatively similar orders of magnitude so that the training becomes stable, and converging at bad local minimum is prevented. A simple matrix multiplication can directly transform the parameters of a quadratic function to the corresponding key values:

$$\label{}
%\small
\left [ K_{1}, K_{2}, K_{3} \right ]=\left [ P_{2}, P_{1}, P_{0} \right ] \begin{bmatrix} 0& \left ( h/2 \right ) ^{2} &h^{2}  \\0 & h/2 & h \\0&0&0\end{bmatrix} ,
%\vspace{-1mm}$$ where we use $\left [K_{1}, K_{2}, K_{3} \right ]$ to denote the key values of a lane, and $\left [P_{2}, P_{1}, P_{0} \right ]$ to indicate the list of quadratic term, linear term and constant term.

While training the lane line localization network in a completely supervised way, we can use the average of the $l_2$ distance between predicted key values and the real key values of each lane as the loss function and, minimize the distance using stochastic gradient descent.

Furthermore, using lane edge coordinates as network input brings us an additional important benefit. Considering the large cost of annotating each lane in an image using a quadratic function, and any inaccurate annotations affecting the training result, we design a weakly supervised training method for the lane line localization network, which only requires the number of lanes in a image as ground truth, we call this form of loss the min-distance loss. When using min-distance loss, the loss function is computed by the sum of the distance of each input point to its nearest estimated line.

Specifically, when training the network using min-distance loss, the output of the networks is still the key values of the lanes. After we get the predicted key values, we transform them back to corresponding quadratic functions representing the estimated lanes: $$\label{}
%\small
[P_{2}^{pr}, P_{1}^{pr}, P_{0}^{pr}]=[K_{1}^{pr}, K_{2}^{pr}, K_{3}^{pr}]\begin{bmatrix} 0&\left ( h/2 \right ) ^{2}&h^{2}\\0&h/2&h\\0&0&0\end{bmatrix} ^{-1},
%\vspace{-2mm}$$ which gives us a set of estimated quadratic functions of all lanes. As mentioned above, the input of the network are the coordinates of all the lane edge points. Now we take all the $X$ coordination values of the points to each estimated quadratic functions, and this will give us the projected $Y$ coordination values of the points to the lanes. The absolute values of the distance between the real $Y$ coordination value and the projected $Y$ coordination values are the horizontal distances of the points to the estimated lanes. Ideally, we take the sum of the horizontal distances of all the points to its nearest estimated lane line as the loss, and minimizing this loss pushes the predicted lane lines to the exact centers of ground truth lane lines.

This weakly supervised training manner only needs the number of lanes as a supervision so that it significantly reduces the cost of annotation. In practice, only using min-distance loss to train the model from scratch leads to trapping in a poor local minimum thus we have to start with supervised training. We come up with two ways to adopt this min-distance loss to the training of our model. First one is that while training the network, we train the network to minimize a combination of the $l_2$ loss and the min-distance loss $$\label{}
\begin{aligned}
L &= L_{l_2} + \alpha L_{min-distance } \\
%&= l2dis(\sum_{i}^{m}(\theta^{pre}(La_{i}) - \theta^{gt}(La_{i})) + \alpha \sum_{j}^{n}mindis(p_{j})\\
%&= \sum_{i}^{m}(\theta^{pre}(La_{i}) - \theta^{gt}(La_{i}))^2 + \alpha \sum_{j}^{n} d^{j}
\end{aligned}$$ where $\alpha$ is a hyperparameter for balancing two loss functions. From our observation, minimizing the two loss functions simultaneously successfully overcomes the annotation inaccuracy so that the estimated lines fit exactly to the lanes in the image regardless some minor misalignments between the annotations and the real lanes, and results in an improvement on generalization.

After the network is trained by well-annotated samples, we can further refine the network on some weakly annotated samples which only has the annotation of number of lanes. Since the network has already been optimized, only using min-distance loss for refinement does not collapse the learned model, and can further refines the network parameters.

# Experiments

We test our LaneNet on the real world traffic data. In this section, we compare our method with other lane detection method, and evaluate the performance. Our data consist of more than 5,000 annotated front view images taken on both highways and urban roads. For each image in the dataset, we annotate all the edges of lane segments, and fit each lane line by a quadratic function. We sample 600 samples from the dataset to be the test data and use the remaining data to train our LaneNet. Specifically, the test set is split into two sub-sets according to the difficulty.

We follow the widely used evaluation metric and adopt two criteria: the true positive rate (TPR) which is calculated as TPR = (the number of detect lanes) / (the number of target lanes), and false positive rate (FPR) which is calculated as FPR = (the number of false positives) / (the number of target lanes). Each lane should be only detected once, so that both over estimation and under estimation of the total number of lanes are undesirable, e.g. detecting a dashed lane for multiple times results in a high FPR, and detecting splitting lanes to be one lane will lower the TPR. We select the lane detection method proposed in to compare the performance since it also relies on few assumptions, and the code is publicly available [^2]. Table  shows the results. The performance of our LaneNet is reported under the following training setting in terms of supervisions: lane edge proposal network is trained using the entire training dataset with full supervision. Weak supervision is introduced in the training of lane line localization network. Additional weakly labelled training samples (about 1,000 samples annotated by only the number of lane lines in the images) are used for fine tuning the parameters of lane line localization network.

```latex
\begin{table*}
    \begin{center}
        %\small
        \caption{Experiment results.}
        \label{tab:exp}
        \begin{tabular}{r|ccc|ccc}
            \toprule
            Difficulty & \multicolumn{3}{c|}{Easy (1170 lanes)} & \multicolumn{3}{c}{Hard (1035 lanes)} \\
            & Detected & TPR & FPR  & Detected & TPR & FPR \\
            \midrule

            Junhwa \textit{et al.} \cite{Hur2013Multi}  & 916 & 78.2\% & 9.5\% & 801 & 77.4\% & 15.9\% \\

            \textbf{LaneNet} & 1146 & 97.9\% & 2.7\% & 1001 & 96.7\% & 3.9\% \\
            \bottomrule
        \end{tabular}
    \end{center}
    %\vspace{-5mm}
\end{table*}
```

We present some illustrative results in Fig. . All the predictions are transformed back to the original images from the IPM images We transform all the predictions on IPM images back to original images for clear illustrations. The result of lane edge proposal network is robust to other objects on the roads, e.g. the lane marks and the vehicles, which we attribute to the contextual information extracting ability of deep neural network based lane edge proposal network, e.g., for every example, there is no false positive detection on the arrows on the roadways, which indicates that our deep neural network based detection method makes decision base on large contextual area, so that is robust to objects with similar local appearances. The localizations of the lanes on the roads are accurate regardless the shapes and the number of the lanes, which demonstrate the effectiveness of our lane localization network to diverse real-world scenarios. The dashed lane lines are properly estimated, which we believe it's attributed to the ability of our lane localization network to evaluate the lane edge proposals from a global view instead of simply clustering the nearby points.

We also conduct a additional experiment on validating the performance improvements by applying the weakly supervised loss function and additional weakly labelled samples. Results are presented in Table . We can see that the weak supervision loss consistently improves the detection on both easy and hard sub-datasets. Further fine tuning the network using weakly labelled data improves the performance on samples with hard difficulties significantly. The fine tuning using weakly labelled data actually enables us to improve the network performance on extreme cases with low cost. We believe that more carefully collected hard samples will further improve the network to a even better stage.

```latex
\begin{table*}

    \begin{center}
        %\small
        \caption{Performance improvements by the weakly supervised loss function and additional weakly labelled samples.}
        \label{tab:add}
        \begin{tabular}{r|ccc|ccc}
            \toprule
            Difficulty & \multicolumn{3}{c|}{Easy (1170 lanes)} & \multicolumn{3}{c}{Hard (1035 lanes)} \\
            & Detected & TPR & FPR  & Detected & TPR & FPR \\
            \midrule

            Baseline   & 996 & 85.1\% & 3.4\% & 904 & 87.3\% & 7.7\% \\

            Weak supervision & 1129 & 96.4\% & 2.8\% & 932 & 90.0\% & 5.1\% \\
            Weak supervision + Tuning on weakly labelled data & 1146    & 97.9\% & 2.7\% & 1001 & 96.7\% & 3.9\% \\
            \bottomrule
        \end{tabular}
    \end{center}
    %\vspace{-5mm}
\end{table*}
```

Additionally, when using a NVIDIA Titan Xp GPU, our lane edge proposal network runs at the speed of 330 frame per second (FPS), and the lane line localization network is even 4 times faster, which enable our entire LaneNet to process images at the speed of 250 FPS. When running on an embedded GPU platform, e.g. NVIDIA Jetson TX1, the speed turns to 26 FPS without specific modifications, which is fast enough for real-time detection. The model size of the entire LaneNet is restricted to less than 1GB. Both the Compact model size and the high running speed further enable the deployment of our LaneNet on vehicles.

# Related work

As a key component of autonomous driving system, lane detection have been widely studied for years. Despite the intensive demands of real world applications, lane detection remains challenging. One of the key reasons is the lack of distinct features. To solve this, many works seek to leverage the structural feature of lanes. Vanishing point is considered to be an important feature for the detection of lanes and is used in . As for local features, gradient on the lane edges is a popular tool for locating the lanes, e.g. in , two filters which detect left and right edge, respectively, are adopted for finding the strong lane edges in the images. However, gradient based detection methods are vulnerable to scenarios with complex irrelevant objects. From our observation, the performances of gradient based methods drop dramatically when large amount of other vehicles appear in the images, which are unavoidable in real world scenarios.

In recent years, with the fast development of deep learning and the hardware support, a lot of lane detection methods based on deep neural networks have emerged and show excellent performances since the training of deep neural network enables detection methods to learn features that are far more robust than the hand-crafted features. Alexandru Gurghian *et al.* propose DeepLanes which use deep neural network to estimate lane position on the both side of the vehicle. Jun Li *et al.* use a multitask convolutional neural network to simultaneously detect the presences and the geometric attributes of the lane marks in a patch of the image, however, the patch based detection makes it hard to infer the global structure of the lane lines on the road, thus is incapable of offering decisive information for navigation, so they propose to further combine the convolutional neural networks with recurrent neural networks to introduce context information and make structural predictions. Similar to our method, Xue Li *et al.* propose using neural networks to detect the edge of lane lines, and use Hough transform to detect different lanes which the robustness might be weak when encountering with complicated scenarios like extremely curved lane lines.

A more comprehensive review on recent advances in lane detection can be found in .

# Conclusion

In this paper, we proposed a lane detection method consists of two deep neural networks. The lane edge proposal network takes an IPM image of the front view of a vehicle as input and produces a lane edge proposal map. Given the lane edge map, the lane line localization network is then in charge of inferring the location of each lane in the image. Our lane detection method does not rely on any assumptions and can be applied to various situations. Using deep neural network endows our method with great robustness and the two stages detection pipeline reduces the computational cost and allows our lane line localization network to be trained in a manner which combines supervised and weekly supervised learning, this remarkably reduces the cost of labelling training data. Extensive experiments demonstrate the speed, accuracy, and robustness of our LaneNet to diverse scenarios. The future work will focus on introducing lane tracking to our LaneNet for a more stable video based lane line detection method.

[^1]: The majority of the work was done while interning at Horizon Robotics, Inc..

[^2]: We use the public version of the code which is not exactly identical with the original implementation as noted in the code.


abstract: Ransomware has appeared as one of the major global threats in recent days. The alarming increasing rate of ransomware attacks and new ransomware variants intrigue the researchers to constantly examine the distinguishing traits of ransomware and refine their detection strategies. Application Programming Interface (API) is a way for one program to collaborate with another; API calls are the medium by which they communicate. Ransomware uses this strategy to interact with the OS and makes a significantly higher number of calls in different sequences to ask for taking action. This research work utilizes the frequencies of different API calls to detect and classify ransomware families. First, a Web-Crawler is developed to automate collecting the Windows Portable Executable (PE) files of 15 different ransomware families. By extracting different frequencies of 68 API calls, we develop our dataset in the first phase of the two-phase feature engineering process. After selecting the most significant features in the second phase of the feature engineering process, we deploy six Supervised Machine Learning models: Naïve Bayes, Logistic Regression, Random Forest, Stochastic Gradient Descent, K-Nearest Neighbor, and Support Vector Machine. Then, the performances of all the classifiers are compared to select the best model. The results reveal that Logistic Regression can efficiently classify ransomware into their corresponding families securing 99.15% overall accuracy. Finally, instead of relying on the 'Black box' characteristic of the Machine Learning models, we present the post-hoc analysis of our best-performing model using 'SHapley Additive exPlanations' or SHAP values to ascertain the transparency and trustworthiness of the model's prediction.
author:
;
;
title: Application of Explainable Machine Learning in Detecting and Classifying Ransomware Families Based on API Call Analysis
  [^1]

# Introduction

Recently, ransomware has become one of the biggest global challenges that are agitating peoples' normal lives. Being harmful software, it applies symmetric and asymmetric cryptography to inscribe user information and poses a Denial-of-Service (DoS) attack on the intended user \[1\]. The unique functional process of ransomware attacks makes it more harmful than any malware attacks and causes irreversible losses. According to \[2\], Fig. 1 illustrates the number of publicized ransomware attacks in 2021, with inflation of 25% than the same time in the previous year. Although the report does not include the number of supply chain attacks, it is creating a big interference in providing healthcare, purchasing groceries, and even loading fuel in vehicles. Examples of these attacks are the Kaseya attack, the colonial pipeline attack, etc. In addition, in the first six months of 2021, the FBI's Internet Crime Complaint Center documented 2,084 ransomware attacks \[3\], and the U.S. Treasury's Financial Crimes Enforcement Network (FinCEN) recorded the cost of around \$590 million related to ransomware activities during that period \[4\]. Moreover, distinct ransomware variants are being detected regularly and more than 130 different ransomware variants have been identified from 2020 till this year causing an inevitable disturbance in day-to-day lives. \[5\].

Due to the increasing number of ransomware variants and ransomware attacks, researchers have been earnestly involving themselves to look for efficient ways to improve the scenarios. While some researchers are analyzing the distinctive behaviors of ransomware by executing it in a secure environment called Dynamic Analysis \[6\]-\[11\], some researchers are analyzing the ransomware without any execution, referred to as Static Analysis \[12\], \[13\]. However, a good number of researchers are combining these two approaches and adopting a Hybrid Analysis Approach \[14\], \[15\]. In this research, we have opted for the dynamic analysis approach for its ability to detect and classify ransomware based on behavioral patterns regardless of the code obfuscation techniques deployed by the ransomware programmers \[16\], \[17\]. The main contributions of this paper are:

-   Develop a Web-Crawler, 'GetRansomware' to automate collecting the Windows Portable Executable (PE) files of 15 different ransomware families from the ransomware repository. The Web-Crawler is essential to automate searching and downloading the samples and to cut down the manual workload, but no prior works targeted this scenario.

-   Develop our dataset and conduct feature selection through a two-phase feature engineering process that includes- 'Feature Extraction' from the sample binaries, and 'Feature Selection' to select the most important features for each ML classifier.

-   Develop, evaluate and compare the performance of six State-of-the-art Supervised Machine Learning models. Our approach includes utilizing Recursive Feature Elimination with Cross-Validation (RFECV) for selecting the significant features and RandomSearchCV for selecting the optimum hyperparameter values for each ML classifier. Thereby we attempt to optimize each model's performance before the comparison is made.

-   Present the post-hoc analysis of the best-performing model using 'SHapley Additive exPlanations' or SHAP values to ascertain the transparency and trustworthiness of the model's prediction. This insight presents a better idea about which features are more dominant in detecting and classifying the ransomware families. While explainability has been widely presented in malware detection scenarios, to the best of the authors' knowledge, till today, no prior works presented their models' explainability that considered only the ransomware families.

The rest of this paper is structured as follows: Section discusses the related works. Section presents our proposed method. The experimental results and discussion are made in Section . Section presents our model's explainability. Section concludes the paper with the direction for future works.

# Related Works

Most researchers prefer the dynamic analysis approach because it can delineate the behaviors of the ransomware in a more explicit manner. Maniath et al. \[6\] analyzed the API call sequence of 157 ransomware and presented an LSTM-based ransomware detection method. Despite securing 96.67% accuracy, this work lacks complete information about the ransomware families/variants, and the number of benign software used for the experiment. VinayaKumar et al. \[7\] proposed an MLP-based ransomware detection method focusing on API call frequency and secured 100%, and 98% accuracy for binary and multi-class classification respectively. However, they deployed a simple MLP network that failed to distinguish CryptoWall and Cryptolocker ransomware. Z. Chen et al. \[8\] used the API Call Flow Graph (CFG) generated from the extracted API sequence of 83 ransomware and 83 benign software. Regardless of securing 98.2% exactness using the Logistic Regression model, the work is based on a smaller dataset that includes only four ransomware families. Also, graph-similarity analysis requires higher computational power that some systems may fail to provide. Takeuchi et al. \[9\] used API call sequences extracted from 276 ransomware, and 312 benign files to identify zero-day ransomware attacks. Although the work secured 97.48% accuracy by deploying the Support Vector Machine, the accuracy of this work decreases while using standardized vector representation because of the less diverse dataset. Using the Intel Pin Tool, Bae et al. \[10\] extracted the API call sequences from 1000 ransomware, 900 malware, and 300 benign files. Their sequential process includes generating an n-gram sequence, input vector, and Class Frequency Non-Class Frequency (CF-NCF) for every sample before fitting their model. Regardless of obtaining 98.65% accuracy using the Random Forest classifier, the model's performance can be improved with the help of deception-based techniques. Hwang et al. \[11\] analyzed the API call sequence of 2507 ransomware and 3886 benign files. They used two Markov chains, one for ransomware and another for benign software to capture the API call sequence patterns. By using Random Forest, they compensate Markov Chains and control FPR and FNR to achieve better performance. Despite securing 97.3% accuracy, their model produces high FPR that can be improved with the help of signature-based techniques.

A good number of researchers chose the static analysis approach to detect ransomware. Baldwin and Dehghantanha \[12\] analyzed the opcode characteristics of 5 crypto-ransomware families and 350 benign samples. Their experiment involved the WEKA AI toolset, and the experimental results showed an accuracy of 96.5% while recognizing five crypto-ransomware families and benign software by using the Support Vector Machine classifier. However, their work could be improved by extending the dataset and extracting those groups of opcodes identified during the evaluation of attribute selection. Zhang et al. \[13\] analyzed the opcode-based characteristics of 1787 ransomware of 8 different ransomware families and 100 benign software. Their technique included moving opcode groupings to the N-gram sequence and afterward Term Frequency Inverse Document Frequency (TF-IDF). Five ML classifiers were used with 10-fold cross-validation among which the Random Forest classifier achieved the highest 91.43% exactness. However, their model could not distinguish Reveton, CryptoWall, and Locky.

Some researchers adopted a hybrid analysis approach that combines the features extracted from the dynamic and static analyses. Subedi et al. \[14\] used both dynamic and static analysis on the library, assembly, and function calls. Moreover, they came up with a new analysis tool, namely, CRSTATIC which was deployed to build signatures that could classify ransomware families with the help of reverse engineering. However, they analyzed only 450 samples of ransomware for their experiment and their experimental results showed some important correlations among various code components instead of detecting ransomware. Ferrante et al. \[15\] proposed a crossover way to deal with android ransomware recognition. While the dynamic detection strategy considered memory utilization, system call insights, CPU, and network utilization, the static identification technique utilized the recurrence of opcodes. They analyzed 2386 benign android applications and 672 applications containing ransomware for their experiment and used 3 ML classifiers- J48, Naïve Bayes and Logistic Regression. Although their experimental result showed 100% accuracy, their test set for the hybrid model included only 9 ransomware applications that were not detected in the static analysis phase but were detected in the dynamic analysis.

Regardless of the relentless effort of the prior research works, we have found some common gaps and challenges, such as a) manual data collection process, b) not having a large or diverse dataset that focuses on both Crypto and Locker types of ransomware, c) using the traditional sandbox environment that often falls short providing in-depth analysis report (i.e., Cuckoo Sandbox), d) modeling binary classification models to detect ransomware from malware or benign software, and e) not presenting the models' explainability. Therefore, the goal of our research work lies in a) automating the data collection process, b) constructing a balanced dataset consisting of the most potential features extracted through a methodical dynamic analysis of both Crypto and Locker types of ransomware, c) utilizing an efficient sandbox environment that conducts large scale integration with other antivirus engines and provides in-depth analysis reports (i.e., Falcon Sandbox), d) developing efficient and faster Supervised Machine Learning model that can attain a higher degree of accuracy in case of multi-class classification, and e) presenting the explainability of our model to affirm the trustworthiness of the model's prediction.

# Proposed Method

Fig. 2 illustrates the process overview of our proposed method which consists of three steps: Data Collection, Feature Engineering, and Classification.

## Data Collection

We have developed a Web-Crawler, namely, 'GetRansomware' to automate searching and downloading Windows Portable Executable (PE) files of ransomware from VirusShare \[18\]. We have also shared our Web-Crawler on our GitHub repository for public access \[19\]. About 95% of our sample ransomware was collected from VirusShare using GetRansomware. The rest of the ransomware samples were collected from theZoo \[20\] and Hybrid-Analysis.com \[21\]. Our samples included both Crypto and Locker types of ransomware and every sample was downloaded as a password-protected compressed file from the repositories. In addition, we have collected benign software from the PORTABLE FREEWARE collection \[28\] which includes portable executables of benign media players, file explorers, file compression tools, file managers, image editors, etc. Table presents the description of our dataset.

|  **Ransomware**   | **Number of samples** |
|:-----------------:|:---------------------:|
|    Cerber(c1)     |          95           |
| CryptoLocker (c2) |          93           |
|  CryptoWall (c3)  |          97           |
|     Eris (c4)     |          98           |
|     Hive (c5)     |          100          |
|    Jigsaw (c6)    |          93           |
|    Locky (c7)     |          95           |
|     Maze (c8)     |          100          |
|     Mole (c9)     |          100          |
|    Sage (c10)     |          100          |
|    Satan (c11)    |          100          |
|    Shade (c12)    |          97           |
| TeslaCrypt (c13)  |          96           |
|   Virlock (c14)   |          93           |
|  WannaCry (c15)   |          93           |
|    Benign (c0)    |          100          |
|       Total       |         1550          |

Description of Dataset

## Feature Engineering

The scarcity of the ransomware dataset is one of the major challenges that hinder the research work in this area \[22\]. Therefore, we aim to construct our dataset through a thorough feature engineering process to overcome this challenge. The feature engineering step is composed of two phases. The phases are:

1.  Phase 1: Feature Extraction

2.  Phase 2: Feature Selection

### Phase 1: Feature Extraction

From the wide range of distinct behavioral features, we have considered utilizing API call frequencies for our classifiers. Application Programming Interface (API) calls are made by the application or program running at a user level to request system services or gain access to the system resources, such as network connection, files, registry, processes, etc. The OS performs the requested services by issuing these calls, and the outcomes are returned to the caller user applications. Thus, API calls made by the ransomware program allow the attackers to obtain the control of the system and perform malicious activities. Since API calls provide a complete picture of a particular program, analyzing API call behavior leads the researchers to better understand the program's behavior \[23\], \[24\]. Therefore, we have considered extracting the API call frequencies by executing the Windows PE files of the ransomware and benign software in a secure virtual environment and utilizing them as distinguishing features for our classification models.

We have analyzed the PE files with the help of Hybrid-Analysis.com \[21\], powered by the CrowdStrike Falcon Sandbox \[25\]. To automate submitting malicious binaries, pull the analysis report after the analysis, and perform advanced or required search queries on the database, Falcon Sandbox provides a free, convenient, and efficient API key that one can obtain from an authorized user account. For analysis, we have used our API key and Falcon Sandbox Python API Connector- VxAPI wrapper \[26\] to automatically submit the binaries from our system. After submission, Falcon Sandbox runs the binaries in a Virtual Machine (VM) and captures the run-time behaviors as illustrated in Fig. 3. Later, it shows the analysis results on the web interface.

Contrary to the prior research works where the analysis tasks were done using the Cuckoo Sandbox \[6\], \[7\], \[9\], \[11\], we have analyzed the PE files using the Falcon Sandbox that used a VM (Windows 7 64-bit) to execute the PE files. Falcon Sandbox can analyze a wide variety of files with or without dependencies, such as PE (i.e., .com, .exe, .dll, etc.), Office (i.e., .doc, .docx, .ppt, .pptx, etc.), and PDF, etc. Also, Falcon Sandbox incorporates many other services, such as VirusTotal, Thug honeyclient, OPSWAT Metadefender, TOR, NSRL (Whitelist), Phantom, and a large number of antivirus engines to provide an integrated and accurate analysis report compared to other sandboxes. While executing the binaries, we have set run-time to the maximum available duration in the Falcon Sandbox to deal with the delayed execution techniques deployed by the attackers. The total time for the analysis was (1550 PE files \* 7 minutes) = 181 hours = 8 days approximately. Next, we obtained the analysis report by using our API key. As we aim to focus on the API call behaviors, we have only sorted and computed the frequency of each API call. At the end of the Feature Extraction phase, we obtained our dataset 'RansomClass' consisting of the frequencies of 68 distinct API calls commonly called by the 15 different ransomware families and benign software as presented in Table .

|     **Name of the API Calls**     |                                |
|:---------------------------------:|:------------------------------:|
|         1\. FindWindowExW         |      2\. LdrGetDllHandle       |
|    3\. NtAdjustPrivilegesToken    |       4\. NtAlertThread        |
|    5\. NtAllocateVirtualMemory    | 6\. NtAlpcSendWaitReceivePort  |
|         7\. NtConnectPort         |       8\. NtCreateEvent        |
|         9\. NtCreateFile          |        10\. NtCreateKey        |
|        11\. NtCreateKeyEx         |      12\. NtCreateMutant       |
|       13\. NtCreateSection        |     14\. NtCreateThreadEx      |
|     15\. NtCreateUserProcess      |     16\. NtDelayExecution      |
|       17\. NtDeleteValueKey       |   18\. NtDeviceIoControlFile   |
|        19\. NtEnumerateKey        |    20\. NtEnumerateValueKey    |
|       21\. NtFsControlFile        |    22\. NtGetContextThread     |
|      23\. NtMapViewOfSection      |     24\. NtNotifyChangeKey     |
|    25\. NtOpenDirectoryObject     |        26\. NtOpenEvent        |
|          27\. NtOpenFile          |         28\. NtOpenKey         |
|         29\. NtOpenKeyEx          |       30\. NtOpenMutant        |
|        31\. NtOpenProcess         |    32\. NtOpenProcessToken     |
|        33\. NtOpenSection         |     34\. NtOpenThreadToken     |
|    35\. NtProtectVirtualMemory    |   36\. NtQueryAttributesFile   |
|     37\. NtQueryDefaultLocale     |   38\. NtQueryDirectoryFile    |
|    39\. NtQueryInformationFile    | 40\. NtQueryInformationProcess |
|   41\. NtQueryInformationToken    |        42\. NtQueryKey         |
|        43\. NtQueryObject         | 44\. NtQuerySystemInformation  |
|       45\. NtQueryValueKey        |   46\. NtQueryVirtualMemory    |
| 47\. NtQueryVolumeInformationFile |        48\. NtReadFile         |
|     49\. NtReadVirtualMemory      |  50\. NtRequestWaitReplyPort   |
|        51\. NtResumeThread        |    52\. NtSetContextThread     |
|     53\. NtSetInformationFile     |    54\. NtSetInformationKey    |
|   55\. NtSetInformationProcess    |  56\. NtSetInformationThread   |
|     57\. NtSetSecurityObject      |       58\. NtSetValueKey       |
|      59\. NtTerminateProcess      |     60\. NtTerminateThread     |
|     61\. NtUnmapViewOfSection     | 62\. NtWaitForMultipleObjects  |
|         63\. NtWriteFile          |   64\. NtWriteVirtualMemory    |
|       65\. NtYieldExecution       |       66\. OpenSCManager       |
|         67\. OpenServiceW         |     68\. SetWindowsHookEx      |

List of API Calls

### Phase 2: Feature Selection

Feature selection is one of the potential steps toward building an efficient ML model because it reduces the probability of overfitting by removing redundant data, improves the model's accuracy by eliminating irrelevant features, and thereby reduces training time. At the beginning of the Feature Selection phase, we have evenly divided (stratified train-test split) our dataset into train data (75%) and test data (25%) to avoid data leakage. Next, we have applied Recursive Feature Elimination with Cross-Validation (RFECV) to our train data for selecting the most significant features for each ML classifier. RFECV is a wrapper-style feature selection method that wraps a given ML model and selects the optimal number of features for that model by recursively eliminating $0-n$ features in each loop. Next, it selects the best performing subset of features based on the accuracy or the score of cross-validation. Besides, RFECV also removes the dependencies and collinearities existing in the model. By using RFECV, we have selected 6 distinct subsets of features for 6 ML classifiers. Table presents the list of features that have been selected by RFECV for each ML classifier. To represent the names of the selected features, the serial numbers from Table are used in Table . These features have been selected by setting 'min_features_to_select' as 34 (half of the features), cv=5, and 'scoring'= 'accuracy' so that RFECV would select at least half of the features based on the optimum accuracy over the 5-fold cross-validation.

|                  Classifier                  | Selected Features | Total |
|:--------------------------------------------:|:-----------------:|:-----:|
|                      LR                      |                   |       |
| 37-39,42,43,45,46,48,53,54,56,59,61,63,64,67 |        41         |       |
|                     SGD                      |                   |       |
|       45-48,51,53,54,56,59,61,63,64,67       |        44         |       |
|                     KNN                      |                   |       |
|       43-48,50,51,53,54,56,59,61,63,64       |        42         |       |
|                      NB                      |                   |       |
|  41-44,48,49,51,53,54,56,58,59,61-65,67,68   |        44         |       |
|                      RF                      |                   |       |
|             50,51,54-56,59,62,63             |        44         |       |
|                     SVM                      |                   |       |
|       45-48,50,51,53,54,56,59,61,63,64       |        44         |       |

List of features selected by RFECV for each ML classifier

## Classification

We have employed Supervised Machine Learning algorithms to detect and classify 15 ransomware families into corresponding categories. Supervised Machine Learning algorithms are trained on the labeled dataset to make a decision in response to the unseen test dataset. As our work focuses on detecting and classifying 15 ransomware families, we have employed the following 6 state-of-the-art Supervised Machine Learning algorithms that are widely used for both binary and multi-class classification problems as per requirement:

Logistic Regression (LR): a type of statistical analysis that predicts the probability of a dependent variable from a set of independent variables using their linear combination.

Stochastic Gradient Descent (SGD): is an optimization algorithm to find the model parameters by updating them for each training data so that the best fit is reached between predicted and actual outputs.

K-Nearest Neighbor (KNN): estimates the likelihood of a new data point being a member of a specific group by measuring the distance between neighboring data points and the new data point.

Naïve Bayes (NB): is based on Bayes' theorem and predicts the probability of an instance belonging to a particular class.

Random Forest (RF): constructs multiple decision trees and finally determines the class selected by the maximum number of trees.

Support Vector Machine (SVM): takes one or more data points from different classes as inputs and generates hyperplanes as outputs that best distinguish the classes.

Since we aim to do multi-class classification and some classifiers are only designed for binary classification problems (i.e., Logistic Regression, Support Vector Machine, etc.), we cannot directly apply these classifiers for multi-class classification problems. Therefore, heuristic methods can be applied to divide a multi-class classification problem into several binary classification problems. There are two types of heuristic methods, such as One-vs-Rest (OvR) which splits the dataset into one class against all other classes each time, and One-vs-One (OvO) which splits the dataset into one class against every other class each time. We have applied the OvR method for our experiment to reduce the time and computational complexities. All these classifiers are built along with 'RandomSearchCV'- a hyperparameter optimization technique, to find the best combination of hyperparameters for maximizing the performance of the models' output in a reasonable time. Instead of exhaustively searching for the optimal values of the hyperparameters through a manually determined set of values (i.e., Grid Search), RandomSearchCV randomly searches the grid space and selects the best combination of hyperparameter values based on the accuracy or the score of cross-validation. Since we have used RFECV for feature selection and RandomSearchCV for hyperparameter optimization, the nested cross-validation technique has been applied to build our model.

# Experimental Results

We have evaluated our models in terms of Precision, Recall, F1-score, Accuracy, and Area Under the Receiver Operator Characteristic (AUROC). These performance metrics are measured as follows:

$$Precision = \frac{TP}{TP+FP} \label{eq:precision}$$

$$Recall = \frac{TP}{TP+FN} \label{eq:recall}$$

$$F1{\text -}score = \frac{2 \times Precision \times Recall}{Precision + Recall} \label{eq:f1}$$

$$Accuracy = \frac{TP + TN}{TP + TN + FP + FN} \times 100 \label{eq:acc}$$ where, TP = True Positives, FP = False Positives, TN = True Negative, FN = False Negatives.

Table describes the experimental results showing that LR outperforms other classifiers securing 99.15% overall accuracy. Although KNN achieves 85.65% accuracy, the processing time of this classifier is less than other classifiers as presented in Table . To examine the degree of separability of our best-performed machine learning model, we have constructed Area Under the Receiver Operator Characteristic (AUROC) curve of the LR classifier as shown in Fig. 4. We construct the AUROC curve by plotting the True Positive Rate or Sensitivity against the False Positive Rate or (1-Specificity). For threshold = 0.5, the average AUROC of LR is 99.55. This is an ideal situation for LR which indicates that the model has a 99.55% probability of predicting the correct class against the rest. Therefore, we can ascertain that LR is our best-performing classifier.

| $P_{avg}$ |  LR   |  SGD  |  KNN  |  NB   |  RF   |  SVM  |
|:---------:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|
| Accuracy  | 99.15 | 85.26 | 85.65 | 94.96 | 82.38 | 94.23 |
| Precision | 99.22 | 95.94 | 93.47 | 98.04 | 99.86 | 99.13 |
|  Recall   | 99.15 | 87.62 | 85.65 | 94.96 | 82.38 | 94.23 |
| F1-score  | 99.15 | 91.14 | 88.04 | 96.28 | 89.31 | 96.37 |
|   AUROC   | 99.55 | 93.67 | 92.55 | 97.41 | 91.19 | 97.09 |

Performance ($P_{avg}$) comparison of LR, SGD, KNN, NB, RF, and SVM Classifiers

| Classifier | Processing Time (in seconds) |
|:----------:|:----------------------------:|
|     LR     |            52.44             |
|    SGD     |            53.65             |
|    KNN     |            51.34             |
|     NB     |            53.77             |
|     RF     |            54.19             |

Classifier's processing time comparison

# Model Explainability

Although some predictive models may not require explanation because of their usage in a low-risk real-world environment, well-studied problem statements, and the availability of sufficient practical solutions, some models that deal with the real-world high-risk environment (i.e., ransomware detection/classification) need explanation. When it comes to the performance evaluation of any model, knowing both 'What' and 'Why' the models have taken these decisions is equally important. These insights lead us to dive deeper into the problems, explore the most important features that affect the model's prediction both locally and globally and thereby build the trustworthiness of the model's predictions. We present the explainability of our model's prediction with the help of 'SHapley Additive exPlanations' or SHAP values \[27\]. SHAP is based on the coalition game theory that measures each feature's individual contribution to the final output while conserving the sum of contributions being the same as the final result. Contrary to the other explanation techniques, such as permutation feature importance, LIME, etc., SHAP provides both global and local explainability. Besides, we can use SHAP values to explain a wide variety of models, such as DeepExplainer to explain Deep Neural Networks (i.e., Multi-Layer Perceptron, Convolutional Neural Network, etc.), TreeExplainer to explain tree-based models (i.e., Random Forest, XGBoost, etc.), and KernelExplainer to explain any model, etc.

To illustrate both local and global explainability, we have used KernerExplainer to explain our Logistic Regression model.

## Local Interpretation

For our classification model, the SHAP value is regarded as a 2-D array where the columns represent the features used in the model and the rows represent individual predictions predicted by the model. So, each SHAP value in this array indicates a specific feature's contribution to the output of that row's prediction as shown in Fig. 5. Here, a positive SHAP value specifies that a feature is positively pushing the base value or expected value to the model output. On the other hand, the negative SHAP value specifies that the feature is negatively pushing the base value to the model output. The base value or the mean model output is computed based on the train data. By using the Force plot, we can visualize this explanation for an individual prediction as shown in Fig. 6. For one instance, the features with higher SHAP values (red) are positively pushing the base value to the model output, and the features with lower SHAP values (blue) are negatively pushing the base value to the model output.

## Global Interpretation

Passing the array of SHAP values to a 'summary_plot' function creates a feature importance plot as shown in Fig. 7. It illustrates 20 highly contributing features to the model output. Here, the x-axis denotes the mean of the absolute SHAP value for each feature which indicates the total contribution of the feature to the model and the y-axis denotes the features used for the classification. The features are organized in descending order from top to bottom by how strongly they influence the model's decision. For our LR model, the frequency of NtAlpcSendWaitReceivePort is the most important feature with a higher mean absolute SHAP value followed by NtFsControlFile, NtWriteVirtualMemory, etc.

In addition to the summary plot, we have created a 'beeswarm' plot as presented in Fig. 8 which shows the distribution of the features' impact on the model output. The x-axis denotes the SHAP value and the y-axis denotes each feature used for the classification. Every point on this plot stands for one SHAP value for one feature and prediction. For our LR model, the lower values (blue) of NtAlpcSendWaitReceivePort have a negative impact on the model output, while the lower values of NtFsControlFile have a positive impact on the model output.

After obtaining these insights from the explanation of our LR model, we conducted another distinct experiment to reevaluate the performance of our LR model including only 20 highly contributing features as shown in Fig. 7 and Fig. 8. As illustrated in Table , our model could distinguish the ransomware families with 95.48% overall accuracy based on those features. This infers that the rest 21 features selected by RFECV for the LR classifier are also contributing to the model achieving its best 99.15% overall accuracy.

| Performance | LR (including only the top 20 features) |
|:-----------:|:---------------------------------------:|
|  Accuracy   |                  95.48                  |
|  Precision  |                  96.06                  |
|   Recall    |                  95.48                  |
|  F1-score   |                  95.45                  |
|    AUROC    |                  97.59                  |

Performance evaluation of LR model including 20 highly contributing features

To evaluate our model's performance concerning the existing classification methods, we have compared the accuracy of the previous API call-based ransomware detection and classification approaches with the proposed method. Both API call frequency and sequence-based detection methods were surveyed from the literature. Table illustrates the comparison of the proposed model with those previous works. It is evident that the proposed method provides higher accuracy while having the advantage of multi-class classification with model explainability. In addition, we have developed a Web-Crawler to automate collecting 15 different ransomware families unlike \[7\] where only 7 ransomware families were utilized for the experiment. Besides, we have conducted feature engineering and utilized hyperparameter optimization techniques to optimize each model's performance contrary to \[6\]-\[11\]. Therefore, our approach is unique and efficient in multi-class classification.

| Reference | Feature | Classifier | Multiclass | Accuracy (%) |
|:---|:---|:---|:---|:---|
| Maniath et al. \[6\] | API call | Long Short-Term Memory | No | 96.67 |
| VinayaKumar et al. \[7\] | API call | Multilayer Perceptron | Yes | 98.00 |
| Z.-G. Chen et al. \[8\] | API call | Logistic Regression | No | 98.20 |
| Takeuchi et al. \[9\] | API call | Support Vector Machine | No | 97.48 |
| Bae et al. \[10\] | API call | Random Forest | Yes | 98.65 |
| Hwang et al. \[11\] | API call | Markov Chain, Random Forest | No | 97.30 |
| Proposed Method | API call | Logistic Regression | Yes | 99.15 |

Performance comparison with existing ransomware classification models

# Conclusion

Ransomware has become the most popular weapon to cybercriminals recently for its easy way to earn money through untraceable cryptocurrency. In this research work, we have developed a Web-Crawler to automate searching and downloading both Crypto and Locker types of ransomware. By executing the binaries of ransomware and benign software using an advanced sandbox tool, we extracted API call frequencies, constructed our dataset, and selected features through a two-phase feature engineering process. Next, we have applied six state-of-the-art Supervised Machine Learning models to detect and classify ransomware families. Our approach includes utilizing Recursive Feature Elimination with Cross-Validation (RFECV) for selecting the significant features and RandomSearchCV for selecting the optimum hyperparameter values for each ML classifier. Thereby we attempt to optimize each model's performance. The results reveal that Logistic Regression can efficiently classify ransomware into their corresponding families securing 99.15% overall accuracy. Finally, instead of relying on the 'Black box' characteristic of the Machine Learning models, we also present the explainability of our best-performing model using SHAP values to ascertain the transparency and trustworthiness of the model's prediction. In the future, we aim to enlarge our dataset by including more samples and potential distinguishing features. In addition, we plan to apply state-of-the-art deep neural networks with explainability for ransomware classification.

[^1]: This work is funded by NetApp.


abstract: A legal document is usually long and dense requiring human effort to parse it. It also contains significant amounts of jargon which make deriving insights from it using existing models a poor approach. This paper presents the approaches undertaken to perform the task of rhetorical role labelling on Indian Court Judgements as part of SemEval Task 6: understanding legal texts, shared subtask A . We experiment with graph based approaches like Graph Convolutional Networks and Label Propagation Algorithm, and transformer-based approaches including variants of BERT to improve accuracy scores on text classification of complex legal documents.
author: **Anshika Gupta, Shaz Furniturewala, Vijay Kumari, Yashvardhan Sharma**
BITS Pilani,Pilani,Rajasthan
`(f20200111,f20200025)@pilani.bits-pilani.ac.in`
`(p20190065,yash)@pilani.bits-pilani.ac.in`
bibliography: anthology.bib; custom.bib
title: Steno AI at SemEval-2023 Task 6: Rhetorical Role Labeling of Legal Documents using Transformers and Graph Neural Networks

# Introduction

Rhetorical Role Labelling for Legal Documents refers to the task of classifying sentences from court judgements into various categories depending on their semantic function in the document. This task is important as it not only has direct applications in the legal industry but also has the ability to aid several other tasks on legal documents such as summarization and legal search. This task is still in it's early stages, with huge scope for improvement over the current state-of-the-art.

To facilitate automatic interpretation of legal documents by dividing them into topic coherent components, a rhetorical role corpus was created for Task 6, sub-task A of The International Workshop on Semantic Evaluation . Several applications of legal AI, including judgment summarizing, judgment outcome prediction, precedent search, etc., depend on this classification.

# Related Works with Comparison

The predominant technique used in Rhetorical Role Labeling over large datasets is based on the use of transformer-based models like LEGAL-BERT and ERNIE 2.0 , augmented by various heuristics or neural network models. The accuracy of these approaches has remained low over the years. The results are summarized in Table 1.

The dataset used to implement the above approaches is relatively small, consisting only of a few hundred annotated documents and 7 sentence classes.

| **Model**                 | **F1 score** |
|:--------------------------|:------------:|
| `LEGAL-BERT`              |    0.557     |
| `LEGAL-BERT + Neural Net` |    0.517     |
| `ERNIE 2.0`               |    0.505     |

Summary of related works on the task of rhetorical role labelling on legal text.

|     |     |
|:----|:---:|

Summary of related works on the task of rhetorical role labelling on legal text.

# Dataset

The dataset is made up of publicly available Indian Supreme Court Judgements. It consists of 244 train documents, 30 validation documents and 50 test documents making a total of 36023 sentences.
For every document, each sentence has been categorized into one of 13 semantic categories as follows:

1.  **PREAMBLE**: The initial sentences of a judgement mentioning the relevant parties

2.  **FAC**: Sentences that describe the events that led to the filing of the case

3.  **RLC**: Judgments given by the lower courts based on which the present appeal was made to the present court

4.  **ISSUE**: Key points mentioned by the court upon which the verdict needs to be delivered

5.  **ARG_PETITIONER**: Arguments made by the petitioner

6.  **ARG_RESPONDENT**: Arguments made by the respondent

7.  **ANALYSIS**: Court discussion of the facts, and evidence of the case

8.  **STA**: Relevant statute cited

9.  **PRE_RELIED**: Sentences where the precedent discussed is relied upon

10. **PRE_NOT_RELIED**: Sentences where the precedent discussed is not relied upon

11. **Ratio**: Sentences that denote the rationale/reasoning given by the Court for the final judgement

12. **RPC**: Sentences that denote the final decision given by the Court for the case

13. **None**: A sentence not belonging to any of the 12 categories

# Proposed Techniques and Algorithms

We try several different approaches for the task at hand. All our models use LEGAL-BERT as their base, and use various methods for further processing and refining of results.

The LEGAL-BERT family of models is a modified pretrained model based on the architecture of BERT . The variant used in this paper is LEGAL-BERT-BASE, a model with 12 layers, 768 hidden units, and 12 attention heads. It has a total of 110M parameters and is pretrained for 40 epochs on a corpus of 12 GB worth of legal texts.

This model was fine-tuned on the task dataset for 2 epochs with a learning rate of 1e-5 using the Adam optimizer and Cross entropy loss

## Direct Classification of CLS tokens

First, we used the default classifier of LEGAL-BERT to find the first set of predictions, to establish a baseline for our further experiments. Our next step used the CLS tokens extracted from the final hidden layer of this trained model.

Similar to the methodology of Gaoa et al. and Furniturewala et al. we utilised the CLS tokens from LEGAL-BERT for further classification models. This CLS token is a 768-dimensional semantic feature that represents BERT's understanding of the text input. It is a fixed embedding present as the first token in BERT's output to the classifier and contains all the useful extracted information present in the input text.

We tried directly applying various multi-layer neural networks to the extracted CLS tokens. These two models served as a baseline to assess the efficacy of our methods.

## Graph-Based Approaches

We implemented classificaton systems based on graph architectures. We modeled the data into a graph using cosine similarity on the CLS tokens generated by LEGAL-BERT. An edge was created between two sentences if and only if their CLS tokens had cosine similarity greater than 0.5, with the cosine similarity acting as edge weight. The threshold was included to minimize the presence of noise-heavy edges in the graph.

$$\cos ({\bf x},{\bf y}) = \frac{ \sum_{i=1}^{n}{{\bf x}_i{\bf y}_i} }{ \sqrt{\sum_{i=1}^{n}{({\bf x}_i)^2}} \sqrt{\sum_{i=1}^{n}{({\bf y}_i)^2}} }$$

The cosine similarity between two nodes, X and Y, is defined in equation (1), where x and y are the CLS tokens for nodes X and Y respectively, and n is the length of the CLS token, i.e. 768 in this case. The function for the final adjacency matrix is defined equation (2).

$$A_{XY} =
    \begin{cases}
      \cos ({\bf x},{\bf y}) & {if\hspace{2mm} \cos ({\bf x},{\bf y}) > 0.5}\\
      0 & {otherwise}
    \end{cases}$$

On this graph, we performed the label diffusion algorithm , to establish a graph-based baseline for our system. Random walk label diffusion assigns labels to an unlabeled node using the average of it's neighbours, weighted by their distance from the node.

$$\centering{F^{t+1} = \alpha \cdot P \cdot F^t + (1 - \alpha) * Y}$$ $$\centering{P = D^{-1/2}\cdot A\cdot D^{-1/2}}$$ $$\centering{F^* = (1 - \alpha) * (I - \alpha P)^{-1} \cdot Y}$$

To implement it, we combined the train and validation label array, one-hot encoded it and masked the validation labels. We then used equation (5) to generate predictions for each sentence. Here P is the normalised adjacency matrix, Y is the array of one-hot encoded labels, $\alpha$ is a hyper-parameter, D is the degree matrix, and Z is the array of predicted labels.

The matrix P is obtained via equation (4), normalizing the adjacency matrix A using the square root inverse of the degree matrix D. For our experimentation, we used $\alpha = 0.5$.

Furthermore, we used a two-layer Graph Convolution Network (GCN) to perform classifications on the data. Inspired by the methodology of BERTGCN , we used the LEGAL-BERT embeddings of each sentence as the node representation for our graph, and then performed graph convolutions on it.

The GCN architecture uses trainable weights to identify the optimal weightage that each neighbour of each node should have on its label. The use of two layers allows us to incorporate the context of one-hop neighbours into the label of a particular node.

$$\begin{aligned}
% \begin{equation}
Z &= f(X, A)\\ &= softmax(\hat{A} \cdot ReLU(\hat{A}XW^{(0)})W^{(1)})
% \end{equation}
\end{aligned}$$

We used equation (7) to predict the labels of the validation set. Here, Â represents the symmetrically normalized adjacency matrix, X is the feature vector which in this case is the LEGAL-BERT embeddings of the nodes, $W^i$ is the matrix of trainable weights in layer $i$.

The calculations required for this approach were extremely computationally expensive, so we were not able to train the model on the entire training set on a V100 server. We used half of the training documents for graph building and the prediction of labels. However, the LEGAL-BERT embeddings were generated by fine-tuning the model on all training documents.

## Context-Based LEGAL-BERT

Our final approach was a Context-Based LEGAL-BERT. We cleaned each sentence by removing all stopwords (such as 'a', 'an', 'the') present using the NLTK library. Then we created a 5 sentence input corresponding to any given input by concatenating its two preceeding sentences and its two succeeding sentences in order. These 5 sentences were separated using LEGAL-BERT's separater token \</s\>. Sentences at the beginning or end of a document were padded using a string of \<pad\> tokens.

These 5 sentence inputs were then tokenized using LEGAL-BERT's tokenizer and fed into the model using the baseline parameters. We used the default classifier to perform classification on these context-based inputs.

# Results

We trained the models and tested them on the validation set. The accuracy scores have been reported in Table 2.

We see that the performance of these models is significantly better than the previous attempts at this problem. The improvement of the results of previously studied models can be attributed to the increase in dataset size, along with other changes in the structure of the task.

However, our Context-based LEGAL-BERT approach outperforms the other frameworks by a significant margin. This exhibits that the context of each sentence is critically important in determining its label, and that we are successful in incorporating the context of each sentence into its representation.

We saw that graph-based approaches did not significantly improve performance compared to the current state-of-the-art models. However, it is important to note that we were unable to run the Graph Convolution Network using the entire train dataset due to compute constraints.

Despite such constraints, there might be other reasons for the mediocre performance of graph-based models. One possible reason is that the representation of the sentences used for building the model was not able to capture information necessary to make better predictions. This also explains how the Context-based LEGAL-BERT performed so much better - it improved the quality of sentence representation, successfully capturing a wider range of features pertaining to the task at hand.

| **Model**                  | **Accuracy** |
|:---------------------------|:------------:|
| `LEGAL-BERT`               |    65.56%    |
| `LEGAL-BERT + Classifier`  |    67.15%    |
| `Graph Label Diffusion`    |    66.34%    |
| `GCN`                      |    67.42%    |
| `Context-based LEGAL-BERT` |    71.02%    |

Summary of results obtained by the models on validation dataset

|     |     |
|:----|:---:|

Summary of results obtained by the models on validation dataset

# Conclusion and Future Work

In this paper, we tried several different techniques to perform a sentence classification task on legal documents. Through our experiments, we show that incorporating context into the CLS tokens of sentences offers a significant improvement of 5.5 percentage points over LEGAL-BERT.

Moreover, through our experiments on graph-based models, we show that improving the CLS tokens results in a better classification, compared to the regular CLS tokens used in a variety of different ways. The Context-based LEGAL-BERT model was not only more accurate but also less resource intensive.

For future improvements on these models, we could try the Graph Convolutional Network approach on the complete dataset. We could also try the various methods of classification, such as a custom neural network or label diffusion, on the context-based CLS tokens.

Moreover, we could further try to incorporate more sentences for context of each target sentence. This would require the use of a long-former model, since the total number of tokens passed into the model will increases.