# Enhancing Cell Tracking with a Time-Symmetric Deep Learning Approach

Gergely Szabó, Paolo Bonaiuti, Andrea Ciliberto, András Horváth
ITK, PPCU, Práter st. 50/A, Budapest, 1083, Hungary.
IFOM, Via Adamello, 16, Milan, 20139, Italy.
szabo.gergely, horvath.andras @itk.ppke.hu
paolo.bonaiuti, andrea.ciliberto @ifom.eu

## Abstract

The accurate tracking of live cells using video microscopy recordings remains
a challenging task for popular state-of-the-art image processing based object
tracking methods. In recent years, several existing and new applications have
attempted to integrate deep-learning based frameworks for this task, but most
of them still heavily rely on consecutive frame based tracking embedded in
their architecture or other premises that hinder generalized learning. To
address this issue, we aimed to develop a new deep-learning based tracking
method that relies solely on the assumption that cells can be tracked based on
their spatio-temporal neighborhood, without restricting it to consecutive
frames. The proposed method has the additional benefit that the motion
patterns of the cells can be learned completely by the predictor without any
prior assumptions, and it has the potential to handle a large number of video
frames with heavy artifacts. The efficacy of the proposed method is
demonstrated through biologically motivated validation strategies and compared
against multiple state-of-the-art cell tracking methods.

##  1 Introduction

For several decades, tracking of solid objects in video recordings has been a
partially solved problem. Classical image processing-based pattern matching
approaches, such as Scale-Invariant Feature Transform (SIFT) [1] combined with
Random Sample Consensus (RANSAC) [2], have the potential to track objects with
fixed shapes invariantly to shift, scale, rotation, and partially to lighting
conditions. Furthermore the extension of such methods with variants of the
Kalman filter [3][4] can introduce object velocity or other degrees of freedom
into the estimation, which can substantially reduce the position estimation
noise, particularly when the introduced additional parameters tend to be
continuous. However, these methods base their estimation on prior assumptions,
which greatly reduces their capabilities when the assumptions do not hold
true. These limitations are especially apparent in live cell tracking on
videomicroscopy recordings, as cells tend to change shape and morphological
description rapidly, unpredictably, and in a highly nonlinear manner.
Additionally, cells can grow and divide with time, have no real velocity, and
often exhibit unpredictable behavior as a living, non-rigid object.
Furthermore, cells on the same recordings tend to be very similar, making the
task even more challenging. Therefore, the most promising solution to address
these issues is using machine learning, where the feature importances and
other estimates are learned by the model and depend mainly on the distribution
of the training dataset.

In recent years, several cell tracking software and libraries have switched to
machine-learning-based solutions, mainly utilizing convolutional neural
networks (CNNs), as they generally perform well on image processing tasks.
Some of these tools, such as CellTracker [5], Usiigaci [6], and YeaZ [7],
separate segmentation and tracking and use deep learning only for high-quality
segmentation. However, they use classical methods, such as the Kalman filter,
bipartite graph matching, or other methods for tracking which ultimately rely
on a rigid metric system. In contrast, other tools, such as CellTrackR [8],
and Ilastik [9], adopt deep learning for tracking either in an end-to-end
manner or as a separate segmentation followed by tracking. This approach shows
advantages both theoretically and empirically, as demonstrated by the often
superior performance of machine learning based and other data-driven methods
when sufficient amount of training data is available. [10][11] However, to the
best of our knowledge, all of these tools rely on a frame-by-frame identity
matching approach, utilizing only consecutive frames for tracking. This
approach is a major simplification and disregards a large amount of useful
information that could be present in the temporal neighborhood of the given
video frames.

Therefore, we developed a complete tracking pipeline that separates
segmentation and tracking because end-to-end training with our method was
unfeasible. Our method uses deep learning for both stages utilizing the
temporal neighborhood of consecutive frames for prediction, makes skipping of
frames possible up to the size of the temporal neighborhood making the
tracking substantially more robust and uses metrics only between different
predictions for the same cell on the same frame minimizing the introduction of
methodical assumptions via the choice of the similarity metric. The most
similar method in the literature that we know of is DeepTrack [12], which was
recently developed for tracking cars. However, while this method utilizes the
local temporal neighborhood via encoding objects using Temporal Convolutional
Networks (TCNs), the network architecture and the metric system for object
matching is substantially different, making our solution architecturally
novel.

###  1.1 Data Source

The training and validation data used to generate the results in this paper
are videos of budding yeast cells dividing in a microfluidic device. Cells are
_S. Cerevisiae_ wild-type-like, with W303 genetic background. Yeast is a
unicellular eukaryote, which shares many essential genes and phenotypes with
multi-cellular eukaryotes, such as mammals. On the other hand, its relatively
small genome ($\sim 12\times 10^{6}$ base pairs), short doubling time and easy
genetic tools allow to perform experiments in a short time and with limited
costs. Such experiments would be very costly or even impossible to do in
higher eukaryotes. Not surprisingly, many discoveries performed in higher
eukaryotes were inspired by original studies performed in yeast.

In the experiments, cells were grown overnight at 30 ∘C in complete YPD medium
and synchronized by the use of the pheromone $\alpha$-factor. Imaging started
when they were released from this stimulus. Cells were free to grow and
duplicate, trapped in the microfluidic device to prevent them from moving in
the field of view. Images were acquired every 3 minutes for 3.5 hours using a
DeltaVision Elite imaging system equipped with a phase-contrast objective. The
field of view is a square of 111.1 $\mu$m size, resulting in a 512x512 pixels
image. The duration of the cell-cycle of the analysed cells is $\sim 1.5$
hours, while their size is $\sim 5\mu$m.

The segmentation and tracking models were trained on a dataset comprising of
314 movies, and for validation purposes during model design, 35 additional
movies were used. These movies were acquired using the same microscope,
objective, and image size, but with varying durations and time-lapses. On
average, each movie contained $35.8$ frames and $20.2$ individual cells to be
tracked. The training and validation sample numbers of the segmentation model
were independent of model parameters, while for the tracking model, the model
parameter tracking range ($TR$) substantially influenced the number of
available samples, as shown in Table 1. Further details on the model parameter
$TR$ will be discussed in Section 2.2. For final testing and evaluation of the
tools and model parameters, 5 additional independent recordings were used,
containing a total of 4629 total number of identified cell instances.

| Tracking
---
Range
| Valid Local
---
Tracks
Training Samples | Validation Samples
1 | 99019 | 88993 | 10026
2 | 64659 | 58059 | 6600
3 | 47940 | 43018 | 4922
4 | 37347 | 33487 | 3860
5 | 29888 | 26769 | 3119
6 | 24373 | 21803 | 2570
7 | 19931 | 17819 | 2112
8 | 16278 | 14533 | 1745
9 | 13301 | 11847 | 1454
10 | 10915 | 9709 | 1206

Table 1: Number of valid local tracking samples for model training in function
of tracking range ($TR$), and the disribution of training and validation
samples.

###  1.2 Segmentation Introduction

In case of tracking pipelines not trained in an end-to-end manner,
segmentation is commonly employed as a preliminary stage before tracking. The
reason being that the detection of the objects is mandatory for tracking, and
detection via segmentation can provide essential information about the objects
present, simplifying and improving the manageability of the tracking task. In
the context of cell tracking, high-quality segmentation of cells is also
important from a biological perspective as it offers valuable information
regarding the size and shape of cells, as well as serving as a quality check
for the tracking process.

From both the target and solution architecture perspectives, segmentation can
be classified into semantic segmentation and instance segmentation. In the
realm of CNN-based image processing, semantic segmentation is typically
considered the easier task and can be accomplished using relatively simpler
network architectures like SegNet [13], U-Net [14], and others. However,
accurately separating objects of the same category after semantic segmentation
can be highly challenging and may necessitate the utilization of classical
image processing methods, which we previously discussed as being less
desirable. On the other hand, recent advancements in deep learning,
particularly the introduction of Mask R-CNN-like architectures [15], have made
it possible to achieve instance segmentation solely based on deep learning
techniques. Although these architectures tend to be more complex and rigid due
to the combination of convolutional, fully connected, and potentially other
network connection types, they are highly valuable when it is essential to
separate instances of objects belonging to the same category.

In our approach, we employ a distinct step of Mask R-CNN-based instance
segmentation prior to tracking, grounded in the previously described
arguments. This step serves two purposes: to provide accurate segmentation
results and to serve as initialization for the tracker. Additionally, the
tracking method itself partially relies on segmentation too. However, in this
context, the objective is to separate a single cell from the background and
other cells, rendering semantic segmentation adequate for the task. In
subsequent sections we give a detailed architectural description for of both
components of our pipeline.

###  1.3 Tracking Introduction

While deep-learning based segmentation has gained popularity in modern cell
tracking software, tracking often relies on comparing consecutive frames using
metrics like Euclidean distance or Intersection over Union (IOU). In contrast,
more advanced but still rigid methods, such as variants of the Kalman filter,
consider motility patterns along with position and shape for tracking.
Following such frame-to-frame similarity measures, various techniques,
including variants of the Hungarian method [16] can be employed for optimal
assignment without repetition.

However, these methods are sensitive to several hyper-parameters, artifacts,
and unusual cases. Even with perfect segmentation and correct assignment,
frame-to-frame metrics are not expected to yield a perfect match unless the
object moves in a completely predictable way, with estimable process and
measurement noises. [4] This sensitivity arises due to the movement, speed,
shape changes, and other descriptors of the object, as well as environmental
factors such as lighting conditions. Consequently, manual fine-tuning of
parameters on a per-recording basis, or even adopting different parameters
within a single record, is often necessary.

Machine-learning based techniques offer an alternative to address these
challenges, as they can learn the movement patterns of the tracked object and
adapt to environmental changes when provided with diverse training data.
Artificial data augmentation can also be used to enhance data variety or
simulate specific anomalies. However, in popular cell tracking software, these
solutions are usually implemented on a frame-to-frame level, potentially
leading to sensitivity to unexpected temporally local anomalies, such as
changes in lighting conditions or administration of drugs that alter cell
movement patterns or morphology.

Instead we propose a multi-frame based assignment method that predicts the
position of a marked cell on several consecutive frames using a state-of-the-
art segmentation network. The assignment of local tracks is then carried out
using the Hungarian method, permitting multiple frame skipping on a
hierarchical level that decreases with increasing temporal distance.

##  2 Architectural Design and Features

###  2.1 Instance Segmentation

The emergence of Mask R-CNNs and similar architectures, coupled with the
availability of libraries and software packages utilizing it, has greatly
simplified instance segmentation for various applications. This holds true for
segmenting living cells as well, although quality can greatly vary depending
on the cell type and the training data at hand. Therefore, instead of focusing
on developing a novel architecture, we trained an instance segmentation model
using the Detectron2 [17] environment. To improve the performance and
robustness of the model, we employed extensive data augmentation techniques
during the training process.

The specific model we chose was a Mask R-CNN architecture with ResNet-50
feature pyramid backbone pretrained on the COCO instance segmentation dataset
[18] with 128 ROI heads for the single ”cell” object type. The model was
trained through 120,000 iterations with a base learningrate of 0.0025,
minibatch size of 4 and otherwise the deafult parameters of the Detectron2
environment. For artificial data augmentation, we employed various
transformation options provided by the Detectron2 environment and the
Albumentations library [19]. The chosen transformations and their respective
parameters are summarized in Table 2.

Detectron2 transformations
---
Transformation | Parameters
Random Brightness | Intensity Min = 0.7, Intensity Max = 1.3, P = 0.2
Random Contrast | Intensity Min = 0.7, Intensity Max = 1.3, P = 0.2
Random Flip | P = 0.1
Random Extent | Intensity Min = 0.7, Intensity Max = 1.3, No Shift, P = 0.2

Albumentations transformations
---
Transformation | Parameters
Grid Distortion | Num Steps = 10, Distort Limit = 0.2, P = 0.2
Elastic Transformation | Alpha = 991, Sigma = 8, Alpha Affine = 50, P = 0.2

Table 2: Data augmentation transformation parameters for Detectron2 and
Albumentations-based instance segmentation training

The trained instance segmentation model showed exceptional performance in both
object detection and shape recognition. Therefore we decided to not only use
its results as inputs for the subsequent tracking method in the pipeline but
also utilize them as the final segmentation output. This decision was made in
contrast to the option of using segmentation estimates of the tracking method.
Examples of such instance segmentation results can be seen in Figure 1. This
approach effectively combines the strengths of both stages of the pipeline.

Figure 1: A showcase of robust segmentation results, depicting accurate cell
detection and segmentation under varying lighting conditions, while
effectively avoiding detection of obviously dead cells.

###  2.2 Tracking

Our method is based on the concept that, although the direction of time is
important in cell tracking, certain morphological changes such as cell growth
generally occur in the positive temporal direction. Therefore, having
information about the future position of the cell can be advantageous for
tracking. As a result, we have developed a time-symmetric tracking method that
utilizes the local temporal neighborhood in both the negative and positive
temporal directions. Based on this architecture, the model learns to capture
the direction of time and temporally unidirectional changes from the data,
instead of relying on a temporally biased architecture. Our approach draws
notable inspiration from the 2019 study of X. Wang [20]. However, instead of
using cycle consistency losses, we sought to create an architecture that
remains invariant to the direction of time itself. After the predictions are
generated by this architecture, the resulting local tracks can be optimally
matched and linked together based on prediction information within their
temporal overlaps.

#### Local Tracking

The local tracking around a single frame is performed on all cells
individually distinguished by the previous instance segmentation step via a
neural network designed for semantic segmentation. The channel parameters of
this network can be defined the tracking range ($TR$) parameter. The
segmentator model input for a single cell instance consists of $(TR+1+TR)+1$
images as channels of the single input where $TR+1+TR$ images are raw video
frames centered around the frame on which the given cell is to be tracked,
while on the additional $1$ channel a single white dot marks the cell to be
tracked at its centroid and the rest of the image is black. The determination
of the centroid coordinates $(x_{c},y_{c})$ is based on the segmentation
instance for the given cell using image moments. The equations are as follows:

$x_{c}=\frac{\sum\sum x\cdot I(x,y)}{\sum\sum I(x,y)}$ (1)
$y_{c}=\frac{\sum\sum y\cdot I(x,y)}{\sum\sum I(x,y)}$ (2)

The target outputs for the local tracker consist of $n+1+n$ frames, containing
the segmentation of the cell marked on the last channel of the input. An
illustration of this local tracking architecture is presented in Figure 2.

The reason for using the centroid of the segmentation and not the segmentation
itself in the input is that according to our experiments if the segmentation
information would be present there, the local tracker would not learn how to
perform semantic segmentation. Instead it would copy the input segmentation to
all outputs, providing false segmentations to all outputs except the middle
one. However via using the centroid as a marker, the local tracker is forced
to learn how to perform segmentation for the single marked cell on all
temporal instances present, and thus it will be able to perform tracking via
segmentation. Furthermore, if a cell instance is temporally close to the
beginning or end of the recording, such that it is within n frames of the
boundary, direct forward or backward tracking using a kernel size of n+1+n is
not feasible. Nevertheless, the number of channels, which determines the local
tracking distance, is a fixed parameter of the network architecture. To
address this issue, we incorporate temporal padding by repeating the first or
last frame. Additionally, we ensured that such instances are included during
training to make the network capable of handling such edge cases and maintain
its temporal invariance.

Based on empirical evidence, we found that among the tested semantic
segmentation models, variants of the DeepLabV3+ [21] architecture from Pytorch
Segmentation Models [22] yielded the best performance for this task. The
models were trained through 20 epochs with a minibatch size of 10 in all
experiments using stochastic gradient descent optimizer with cosine annealing
learning rate scheduling having warm restarts during the first 15 epochs, and
a gradual cooldown during the last 5 epochs. The model output had sigmoid
activation with binary crossentropy loss.

Figure 2: Depiction of input and output information structure of local
tracking via segmentation with local tracking range ($TR$) = 2. Notice that
while the left to right positive temporal direction of the data is
recognizable from cell growth, and the model learns this during training, the
architecture itself has no built in directional preference, resulting in a
completely data-driven estimation for local tracking.

Artificial data augmentation of the local tracking samples was performed using
the torchvision transformations library [23]. For positional transformations
random horizontal flip [p=0.5] and random affine [degrees=(-40, 40),
scale=(0.7, 1.3), p=0.4] transformations were uniformly applied to the input
and ground truth data. For color transformations color jitter [brightness=0.5,
contrast=0.5, p=0.4] transformation was applied for each input frame
individually to improve model robustness for unexpected lighting artefacts.

#### Global Tracking

As the output of local tracking, a single cell instance marked by the video
frame and its position is tracked forwards and backwards through $TR$ frames.
This local tracking output can be obtained for each cell instance detected by
the instance segmentation step, serving as the input for global tracking.
Therefore the aim of global tracking is to match local tracks based on the
semantic segmentation predictions in their overlapping areas, and chain them
together based on ID matching creating full cell tracks based on a globally
optimal consensus.

In practice, at first, we obtain this global consensus on a $t$ to $t+1$
level, as usually the cells can be tracked from one frame to the next.
However, if a cell instance did not receive any matches above the minimal
requirement threshold, the matching can be repeated for larger temporal
distances. This secondary matching is performed after all $t$ to $t+1$
matching is done, and only between the remaining candidates. Secondary
matching is also performed hierarchically in terms of temporal difference
meaning that temporally closer instances get matched first and only the
remaining ones can get matched later with larger temporal differences. The
maximal possible matching distance is $2TR$, but matches over $TR$ are
unreliable as they are based on relatively few segmentation predictions, and
only indirectly contain information about the segmentation of the central cell
instances to be matched. Therefore we believe, using $TR$ as the maximal
temporal distance between local track matching is logically the most sound
choice, but in practice other choices could have minor benefits depending on
the dataset.

Based on this methodology, the global consensus of ID matching can be turned
into a one-to-one assignment problem for $N\times M$ candidates, where $N$ is
the number of candidates in frame $t$ and $M$ is the number of candidates in
frame $t+\Delta t$ where $1<\Delta t\leq 2TR$. Such assignment problems can be
optimally solved by variants of the Hungarian method in polynomial time even
for non-square matrices, but this requires a sound metric choice to measure
the goodness of the candidate matches. In our case these candidate matches are
the $2TR+1-\Delta t$ segmentation pairs in the overlapping temporal region of
the local tracks to be compared. Therefore in an ideal scenario where perfect
segmentations are achieved in local tracking, any metrics that solely compare
segments within the same frame and calculate the average of these comparisons
could be employed. This is because segmentation instances belonging to the
same cell should result in a perfect match. However in practice, the mean of
Intersection over Union (IOU) turned out to be the best metric choice as it
takes both the positional and morphological differences of the segmentations
into account, and gives 0 similarity for segmentations with no overlap
regardless of the distance. This process of calculating the metric similarity
measure based on local tracks is depicted in Figure 4.

After computing the selected metric between all global tracks of frame $t$ and
$t+\Delta t$, a threshold can be employed to establish a minimum required
similarity based on the metric results. This step eliminates candidates that
do not meet the desired level of similarity. Subsequently, the Hungarian
method can be applied to determine the optimal global consensus for ID
matching between frame $t$ and $t+\Delta t$, as depicted in Figure 5. Via
performing these steps in the hierarchical approach as described earlier,
complete global tracks are generated based on a global consensus of
similarity, which solely relies on the prediction capabilities of the local
tracker. An example of successful global tracking using this pipeline is
presented in Figure 3.

Figure 3: A side by side display of tracking results, demonstrating successful
tracking of all cells in the same recording with a temporal difference of 8
frames. New cells were appropriately assigned new IDs while maintaining
consistent tracking of existing cells. Figure 4: Schematic structure of the
metric similarity measurement step between $2TR+1$ long local tracks of
different cell instances on frames with a temporal distance of $\Delta t$. The
solid lines indicate the central segmentation instances with a $\Delta t$
temporal distance to be matched, while the arrows indicate the similarity
metric between the segmentation estimates for each time frame. Subsequently,
the individual metric results are averaged to obtain a single measure of
similarity. Figure 5: Illustration of the metric similarity based assignment
step using the Hungarian method for non-square matrices. Newly unassigned
cells, such as cell number 4 receive a new ID, while the previously assigned
cells retain the ID of their corresponding previous instances.

#### Skipped Instance Interpolation

By employing the aforementioned methodology for global tracking, we can
effectively address inconsistency errors in segmentation that commonly arise
in realistic scenarios when $\Delta t>1$. In such cases, although the
resulting tracks will be complete, they may not be continuous, and
segmentations for missed frames will be absent. To tackle this issue, we use a
positional linear interpolation method, which is described by the following
equations:

$\Delta c(x,y)=\frac{\begin{aligned}
(t-t_{\text{last}})c_{\text{next}}(x,y)+\\\
(t_{\text{next}}-t)c_{\text{last}}(x,y)\end{aligned}}{t_{\text{next}}-t_{\text{last}}}$
(3) $S(x,y)=S_{\text{last}}(x,y)+c(x,y)$ (4)

Here, $S(x,y)$ and $t$ represent the interpolated segmentation and its
corresponding time, respectively. $S_{\text{last}}(x,y)$,
$c_{\text{last}}(x,y)$, and $t_{\text{last}}$ denote the segmentation,
centroid, and time of the last occurrence of the cell with the given ID, while
$c_{\text{next}}(x,y)$ and $t_{\text{next}}$ represent the centroid and time
of the next occurrence of the cell with the given ID. Although this method
only shifts the segmentation from the last occurrence to the linearly
interpolated position without altering its shape, in practice the results have
proven satisfactory. Still, in the future, incorporating a method that
linearly interpolates the shape of the segmentation by considering both the
last and next occurrences could potentially yield minor improvements.

##  3 Evaluation

For evaluation purposes, we opted to calculate F-scores based on binary
decisions of correctness for segmentation and tracking, instead of utilizing
continuous regression metrics. This approach offers a clear numerical
evaluation of the capabilities of the compared tools and tool versions.
F-scores for segmentation were computed using IOU, with a minimum similarity
threshold of $0.5$. We selected this threshold as it is strict enough to
assess the accuracy of correct cell detection while disregarding minor
segmentation discrepancies that hold minimal biological significance. In the
case of tracking evaluation, F-scores were also calculated based on the
segmentation IOU, with the additional inclusion of ID matching. True positive
values, also referred to as ”links” were registered if the segmentations
matched on the same frame, and the IDs matched consecutively for both the
prediction and the ground truth. This tracking evaluation method of link
matching is described in the work of M. Primet (2011) [24].

###  3.1 Comparative Evaluation

For comparative evaluation, we selected Phylocell [25] [26] and YeaZ [27] as
the two other cell tracking tools for comparison. The reason behind this
choice was that the ground truth training and testing data were generated
using Phylocell and were later manually corrected by experts. On the other
hand YeaZ provides a fair comparison to our solution as it employs a similar
segmentation pipeline and a somewhat comparable tracking pipeline, utilizing
the Hungarian method. As the ground truth data excludes cells that eventually
leave the microscope’s field of view throughout the recording (from now on
referred to as ”border tracks”), we conducted two evaluation scenarios for the
tools. In one scenario, border tracks were removed in post-processing from all
predictions to ensure an unbiased comparison, while in the other scenario, the
prediction results were left unaltered.

Figure 6: Comparative evaluation of segmentation and tracking performances of
Phylocell, YeaZ and our pipeline based on F-measures of segmentation IOU and
tracking link matches with enabled and disabled border track predictions.
Disabled border tracks present an unbiased comparison between the tools, while
enabled border tracks show results with the unaltered outputs.

The segmentation and tracking F-measure results are depicted in Figure 6. It
is evident that our solution substantially surpasses both other tools in terms
of segmentation and tracking quality, particularly in the more equitable
scenario where the border tracks were eliminated. In the unaltered prediction
case, the performance of Phylocell remained comparable to its prior results
since the tool already automatically removes most border tracks. Conversely,
YeaZ and our solution exhibited notably lower performance due to domain
discrepancy between the predictions and ground truth data. Nevertheless, even
in this case, our solution outperformed Phylocell in terms of expected F-score
values, although with a substantially higher degree of variance on the lower
side.

###  3.2 Hyperparameter Dependencies

The parameters we considered to have a major impact on tracking performance
for both local and global tracking are the local tracking range, the
complexity of the local tracker model backbone, and the metric used for global
consensus.

Figure 7: Comparative evaluation of tracking F-scores in function of local
tracking ranges, backbone model complexities and global consensus metrics with
enabled and disabled border track predictions. Disabled border tracks present
an unbiased comparison between the parameters, while enabled border tracks
show results with the unaltered outputs.

The differences in tracking F-scores resulting from these parameters are
presented in Figure 7. Somewhat surprisingly, we observed minimal differences
in performance based on the local tracking range and model complexity.
Furthermore, smaller local tracking ranges exhibited marginally better
performance, suggesting that the benefits of larger tracking kernels were
outweighed by the increased model complexity, which slightly hindered
training. On the other hand, these results also indicate that lightweight
models with smaller local tracking ranges and simpler backbones are suitable
for tracking predictions, resulting in substantially shorter inference times.
As an additional benefit, while performing inference without GPU acceleration
remains slower by orders of magnitude, it is more reasonable with such
lightweight models, making initial pipeline testing more accessible. As
anticipated, the use of IOU as the metric for global consensus substantially
outperformed Euclidean distance, underscoring the importance of incorporating
cell morphology in short-term tracking.

##  4 Conclusion

In this paper, we have presented a comprehensive pipeline for cell
segmentation and tracking of budding yeast cells, introducing a novel
architecture for cell tracking based on local time-symmetric tracking and
global optimal assignment. Our approach incorporates re-trainable neural
networks for both the segmentation and tracking stages, enabling the
architecture’s potential for generalization across different cell types,
provided labeled training data is available. The extensive evaluation of our
method, utilizing an F-score based metric system separately for segmentation
and tracking, highlights its effectiveness. Empirical results from rigorous
testing demonstrate that our pipeline achieves state-of-the-art performance,
outperforming the compared tools YeaZ and Phylocell in both segmentation and
tracking tasks. This indicates the robustness and potential applicability of
our method in cell tracking scenarios. By combining the advantages of deep
learning-based segmentation and time-symmetric local tracking with linear
optimal assignment based on the Hungarian method, we have developed a
promising pipeline that has the potential to substantially contribute to cell
tracking tasks in biological research.

## Author Contributions

GS designed and developed the primary pipeline, encompassing model design,
training, and user interaction functionalities. GS played a major role in
writing the manuscript. PB conceptualized the evaluation methods, generated
comparative results for the assessed tools, contributed to architectural
design, and participated in manuscript preparation. AC provided valuable
theoretical insights and constructive feedback from a biological standpoint,
contributing substantially to the manuscript. AH offered essential theoretical
insights and constructive critiques from a machine learning perspective, made
substantial contributions to pipeline design, and participated in manuscript
preparation. All authors have reviewed and approved the final manuscript.

## Acknowledgements

PB was supported by an Umberto Veronesi fellowship. The AC lab is supported by
the italian association for cancer research (AIRC IG-21556) and the Hungarian
National Research Development and Innovation Office, through grant
TKP2021-EGA-42. The funder played no role in pipeline design, result
evaluation, or the writing of this manuscript.

## Code and Data availability

We have made our segmentation and tracking pipeline openly accessible through
a dedicated GitHub repository, which can be found at the following URL:
https://github.com/SzaboGergely0419/PPCU_IFOM_YeastTracker. This repository
not only hosts the main pipeline but also links to supplementary resources
including various versions of the tracker model and sample test data.
Furthermore, we have created a comprehensive Google Colab environment within
the repository that enables the execution of the entire pipeline. Currently
the training codes are not included in the public repository. However, we will
make them available for peer review and plan to release them on the
aforementioned repository in the foreseeable future.

## References

  * [1] D. G. Lowe, “Object recognition from local scale-invariant features,” in Proceedings of the seventh IEEE international conference on computer vision, vol. 2, pp. 1150–1157, Ieee, 1999.
  * [2] M. A. Fischler and R. C. Bolles, “Random sample consensus: a paradigm for model fitting with applications to image analysis and automated cartography,” Communications of the ACM, vol. 24, no. 6, pp. 381–395, 1981\.
  * [3] R. E. Kálmán, “A new approach to linear filtering and prediction problems,” Journal of Basic Engineering, vol. 82, no. 1, pp. 35–45, 1960.
  * [4] S. J. Julier and J. K. Uhlmann, “New extension of the kalman filter to nonlinear systems,” in Signal processing, sensor fusion, and target recognition VI, vol. 3068, pp. 182–193, Spie, 1997.
  * [5] T. Hu, S. Xu, L. Wei, X. Zhang, and X. Wang, “Celltracker: an automated toolbox for single-cell segmentation and tracking of time-lapse microscopy images,” Bioinformatics, vol. 37, no. 2, pp. 285–287, 2021.
  * [6] H.-F. Tsai, J. Gajda, T. F. Sloan, A. Rares, and A. Q. Shen, “Usiigaci: Instance-aware cell tracking in stain-free phase contrast microscopy enabled by machine learning,” SoftwareX, vol. 9, pp. 230–237, 2019.
  * [7] N. Dietler, M. Minder, V. Gligorovski, A. M. Economou, D. A. H. L. Joly, A. Sadeghi, C. H. M. Chan, M. Koziński, M. Weigert, A.-F. Bitbol, et al., “A convolutional neural network segments yeast microscopy images with high accuracy,” Nature communications, vol. 11, no. 1, p. 5723, 2020\.
  * [8] Y. Chen, Y. Song, C. Zhang, F. Zhang, L. O’Donnell, W. Chrzanowski, and W. Cai, “Celltrack r-cnn: A novel end-to-end deep neural network for cell segmentation and tracking in microscopy images,” in 2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI), pp. 779–782, IEEE, 2021\.
  * [9] S. Berg, D. Kutra, T. Kroeger, C. N. Straehle, B. X. Kausler, C. Haubold, M. Schiegg, J. Ales, T. Beier, M. Rudy, et al., “Ilastik: interactive machine learning for (bio) image analysis,” Nature methods, vol. 16, no. 12, pp. 1226–1232, 2019.
  * [10] D. Vlah, A. Kastrin, J. Povh, and N. Vukašinović, “Data-driven engineering design: A systematic review using scientometric approach,” Advanced Engineering Informatics, vol. 54, p. 101774, 2022.
  * [11] C. Sun, A. Shrivastava, S. Singh, and A. Gupta, “Revisiting unreasonable effectiveness of data in deep learning era,” in Proceedings of the IEEE international conference on computer vision, pp. 843–852, 2017.
  * [12] V. Katariya, M. Baharani, N. Morris, O. Shoghli, and H. Tabkhi, “Deeptrack: Lightweight deep learning for vehicle trajectory prediction in highways,” IEEE Transactions on Intelligent Transportation Systems, vol. 23, no. 10, pp. 18927–18936, 2022.
  * [13] V. Badrinarayanan, A. Kendall, and R. Cipolla, “Segnet: A deep convolutional encoder-decoder architecture for image segmentation,” IEEE transactions on pattern analysis and machine intelligence, vol. 39, no. 12, pp. 2481–2495, 2017.
  * [14] O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,” Medical Image Computing and Computer-Assisted Intervention (MICCAI), vol. 9351, pp. 234–241, 2015.
  * [15] K. He, G. Gkioxari, P. Dollár, and R. Girshick, “Mask r-cnn,” in Proceedings of the IEEE international conference on computer vision, pp. 2961–2969, 2017.
  * [16] H. W. Kuhn, “The hungarian method for the assignment problem,” Naval research logistics quarterly, vol. 2, no. 1-2, pp. 83–97, 1955.
  * [17] Y. Wu, A. Kirillov, F. Massa, W.-Y. Lo, and R. Girshick, “Detectron2.” https://github.com/facebookresearch/detectron2, 2019.
  * [18] T.-Y. Lin, M. Maire, S. Belongie, J. Hays, P. Perona, D. Ramanan, P. Dollár, and C. L. Zitnick, “Microsoft coco: Common objects in context,” in Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740–755, Springer, 2014.
  * [19] A. Buslaev, V. I. Iglovikov, E. Khvedchenya, A. Parinov, M. Druzhinin, and A. A. Kalinin, “Albumentations: fast and flexible image augmentations,” Information, vol. 11, no. 2, p. 125, 2020.
  * [20] X. Wang, A. Jabri, and A. A. Efros, “Learning correspondence from the cycle-consistency of time,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2566–2576, 2019.
  * [21] L.-C. Chen, Y. Zhu, G. Papandreou, F. Schroff, and H. Adam, “Encoder-decoder with atrous separable convolution for semantic image segmentation,” in ECCV, 2018.
  * [22] P. Iakubovskii, “Segmentation models pytorch.” https://github.com/qubvel/segmentation_models.pytorch, 2019.
  * [23] PyTorch Contributors, “torchvision.transforms.” https://pytorch.org/docs/stable/torchvision/transforms.html, 2023.
  * [24] M. Primet, Probabilistic methods for point tracking and biological image analysis.  PhD thesis, Université René Descartes-Paris V, 2011.
  * [25] G. Charvin, “Phylocell,” 2021.
  * [26] S. Fehrmann, C. Paoletti, Y. Goulev, A. Ungureanu, H. Aguilaniu, and G. Charvin, “Aging yeast cells undergo a sharp entry into senescence unrelated to the loss of mitochondrial membrane potential,” Cell reports, vol. 5, no. 6, pp. 1589–1599, 2013.
  * [27] F. Padovani, B. Mairhörmann, P. Falter-Braun, J. Lengefeld, and K. M. Schmoller, “Segmentation, tracking and cell cycle analysis of live-cell imaging data with cell-acdc,” BMC biology, vol. 20, no. 1, p. 174, 2022\.


IQshort=IQ,long=in-phase and quadrature PDUsshort=PDUs,long=Protocol Data
Units RRCshort=RRC,long=Radio Resource Control ZMQshort=ZMQ,long=ZeroMQ
LSTMshort=LSTM,long=Long Short-Term Memory LALshort=LAL,long=Listen-and-Learn
TSshort=TS,long=Technical Specifications NRshort=NR,long=New Radio
eMBBshort=eMBB,long=enhancing Mobile Broadband mMTCshort=mMTC,long=massive
Machine Type Communications O-RANshort=O-RAN,long=Open Radio Access Network
NASshort=NAS,long=Non-Access Stratum PHYshort=PHY,long=physical
ASshort=AS,long=Access Stratum UEshort=UE,long=User Equipment
MITMshort=MITM,long=man-in-the-middle OTAshort=OTA,long=over-the-air
MACshort=MAC,long=medium access control CNshort=CN,long=core network
gNBshort=gNB,long=gNodeB DoSshort=DoS,long=Deny of Service
BSshort=BS,long=Base Station NSAshort=NSA,long=Non Standard-Alone
SAshort=SA,long=Standard-Alone EPSshort=EPS,long=Evolved Packet System
eNBshort=eNB,long=evolved NodeBs AKAshort=AKA,long=authentication and key
agreement EC-AKAshort=EC-AKA,long=Ensured confidentiality Authentication and
Key agreement ASMEshort=ASME,long=Access Security Management Entity
UPshort=UP,long=User Plane KDFshort=KDF,long=key derivation function
IMSIshort=IMSI,long=international mobile subscriber identity
S-TMSIshort=S-TMSI,long=SAE Temporary Mobile Subscriber Identity
OAIshort=OAI,long=Open Air Interface NFVshort=NFV,long=Network Function
Virtualization SDRshort=SDR,long=Software Defined Radios
SUTshort=SUT,long=System Under Test

# From Ambiguity to Explicitness: NLP-Assisted 5G Specification Abstraction
for Formal Analysis

Shiyu Yuan1, Jingda Yang1, Sudhanshu Arya1, Carlo Lipizzi1 and Ying Wang1
1Stevens Institute of Technology, Hoboken, NJ
Email: 1{syuan14, jyang76, sarya, clipizzi<EMAIL_ADDRESS>
###### Abstract

Formal method-based analysis of the 5G Wireless Communication Protocol is
crucial for identifying logical vulnerabilities and facilitating an all-
encompassing security assessment, especially in the design phase. Natural
Language Processing (NLP) assisted techniques and most of the tools are not
widely adopted by the industry and research community. Traditional formal
verification through a mathematics approach heavily relied on manual logical
abstraction prone to being time-consuming, and error-prone. The reason that
the NLP-assisted method did not apply in industrial research may be due to the
ambiguity in the natural language of the protocol designs nature is
controversial to the explicitness of formal verification. To address the
challenge of adopting the formal methods in protocol designs, targeting (3GPP)
protocols that are written in natural language, in this study, we propose a
hybrid approach to streamline the analysis of protocols. We introduce a two-
step pipeline that first uses NLP tools to construct data and then uses
constructed data to extract identifiers and formal properties by using the NLP
model. The identifiers and formal properties are further used for formal
analysis. We implemented three models that take different dependencies between
identifiers and formal properties as criteria. Our results of the optimal
model reach valid accuracy of 39% for identifier extraction and 42% for formal
properties predictions. Our work is proof of concept for an efficient
procedure in performing formal analysis for large-scale complicate
specification and protocol analysis, especially for 5G and nextG
communications.

###### Index Terms:

5G, formal dependency table, NLP, identifier, formal property, language model

##  I Introduction

The fifth-generation (5G) technology introduces more flexibility than previous
generations of cellular networks for novel functionalities and higher
performance. It also enables novel customization when applied to various
verticals. Meanwhile, the flexibility and customization, the design of more
complicated specifications and protocols, along with amendment agreements
bring forth the need for more sophisticated protocol examinations and
verification to detect functional and nonfunctional vulnerabilities inherited
from previous generations of networks and newly introduced. To ensure the
security of these specifications and protocols, fuzz testing, and formal
verification are employed to detect vulnerabilities in the system and verify
the accuracy of implemented protocol behaviors.

To expedite the assessment of protocol security, formal verification is
employed to mathematically model identifiers and properties, using formal
methods and fuzz testing to reason about and establish their security [1, 2,
3].Fuzz testing is an effective method that utilizes a plethora of random or
mutated inputs to uncover system vulnerabilities. This automated approach
significantly reduces labor costs associated with manual testing efforts.
However, its effectiveness relies on the quality and diversity of test inputs
generated [1, 4].

In the realm of 5G Wireless Communication Protocol, formal verification plays
a pivotal role in ensuring the reliability, security, interoperability, and
performance of 5G networks, including protocol correctness, security measures,
and overall system robustness [5]. Unlike automated physical simulation
testing approaches, formal verification is a labor-intensive process that
necessitates manual expertise in extracting protocols’ procedures and
relationships. Consequently, there is a pressing need for an automated system
capable of analyzing protocols, extracting identifiers, and establishing their
relationships.

Recently, an NLP based 5G vulnerability detection approach has been proposed
[6]. The cross-layer 5G vulnerability detection is conducted via fuzzing-
generated run-time profiling with piloting on srsRAN. To characterize the
performance, the authors leverage various machine learning classification
algorithms, including K-Nearest Neighbors, Random Forest, and Logistic
Regression with performance reaching as high as $\approx 95.9\%$. Although the
approach offers high accuracy, it relies on generating random or semi-random
inputs to the target system, thereby may not be able to explore all possible
code paths. This can result in missing certain vulnerabilities that require
specific sequences of actions to trigger. Moreover, due to the complex 5G
network, fuzzing can lead to a massive number of test cases, making it
challenging to handle the state-space explosion problem effectively.

##  II Related Work

NextG Wireless Communication Protocol and Formal Verification: With the rapid
development of wireless communication technologies and the ongoing efforts to
address vulnerabilities in previously defined protocols, experts regularly
introduce new protocols that modify high-risk rules and introduce additional
constraints. However, the comprehension of these newly released protocols by
human experts remains a significant challenge in efficiently transferring them
into formal properties for cybersecurity purposes. Therefore, there is a
growing need for an automatic framework that can extract identifiers and
analyze formal relationships, ensuring the security of protocols while
alleviating the burden of manual analysis. Formal verification has emerged as
a valuable tool in numerous research studies for detecting logical
vulnerabilities and ensuring error-free systems with propriety [5]. In
contrast to traditional fuzz testing, formal verification leverages
mathematical proofs to rigorously verify the logical correctness of
procedures. By focusing on the properties and relationships of each
identifier, formal verification offers a systematic approach to ensure the
integrity and reliability of system behavior. However, due to the stringent
nature of mathematical proofs, formal verification typically requires the
expertise of security professionals for manual analysis. The extensive
protocol analysis and recurring context updates necessitate the development of
efficient protocol analyzer tools to streamline the extraction of identifiers
and their properties in the field of cybersecurity [7].

Identifier and Formal Property Labeling: Identifiers in communication
protocols often represent specific protocol components, each with a unique
role in ensuring efficient and effective communication. The network
identifiers in 5G system are divided into subscriber identifiers and UE
identifiers [8]. These identifiers could be protocol names (e.g., ‘5G’, ‘LTE’,
‘TCP’), specific algorithm names (e.g., ‘RSA’, ’AES’), or other technical
terminologies unique to the communication domain (e.g., ‘RCCRequest’,
‘RCCConfiguration’). Simultaneously, formal properties associated with these
identifiers refer to the mathematically definable characteristics that
describe the behavior of the protocols. These might include attributes like
‘Authentication’, ‘Accounting’,‘Integrity’,‘Confidentiality’ [1]. Together,
the identifier and its formal properties provide a detailed specification for
a protocol component, outlining its function, behavior, and performance within
the communication system. The extraction of formal verification information
provides a foundational step towards a deeper understanding and analysis of
complex communication protocols such as 5G.

Information Extraction in Natural Language Processing: Information Extraction
(IE) is a task in NLP that involves automatically extracting structured
information from unstructured text data [9]. This typically involves
processing text and identifying relevant pieces of information, such as
entities, relationships between entities, and specific attributes of entities
or relationships. IE has three main tasks, named entity recognition (NER),
semantic role labeling (SRL), and relation extraction (RE) [10] with two main
approaches on these three tasks, a heuristic approach based on the transition-
based model and supervised learning through a neural network and annotation
data [11]. Despite the availability of generic-domain datasets for IE, domain-
specific IE resources, such as annotated datasets for NER or RE, are scarce.
This scarcity presents a significant challenge when executing IE tasks within
a specific domain. Several studies provide NER pipeline such as in work [12],
but their training corpus are more general cybersecurity content instead of
communication protocols.

####  II-1 NER

NER is to identify named entities or semantic roles in a sequence of tokens
[9]. In this study, our focus is on domain-specific terminology in 5G
communication protocol, also referred to as ’identifiers’. The focus of this
study is to undertake a systematic analysis of protocol documents, with the
specific aim of locating and extracting the pertinent identifiers contained
within them and predicting the formal properties that delineate the
relationships between the extracted identifiers.

####  II-2 RE

RE is a task in information extraction that involves identifying and
classifying semantic relationships between entities in text. These entities
are typically pre-identified, and the goal of relation extraction is to
ascertain how these entities are connected [13]. However, in our task, the
entities are not pre-identified. Instead, the entities need to be located and
extracted during the training phase, which adds a layer of complexity to the
task. In the context of our study, relations are considered as formal
properties between identifiers.

To address the research gap, streamline the analysis pipeline, and mitigate
the potential risks of experimental errors, we propose a unique solution
leveraging natural language processing techniques to analyze the 5G protocol.
The protocol, primarily written in natural language, serves as the basis for
our solution. Our approach specifically integrates two different natural
language processing methodologies for the analysis. Initially, we employ a
traditional heuristic technique utilizing information extraction methods to
identify named entities and semantic roles. Subsequently, we implement a pre-
trained language model to extract the entities or identifiers and the formal
property between the identifiers within these sentences.

Contributions:  To this end, we make the following contributions:

  * •

This work presents a first-of-its-kind and a new approach to doing formal
analysis employing the NLP technique. The proposed approach aims to overcome
the limitations imposed by the use of the conventional methods, thereby
allowing wider state-space explosion and coverage. We elicit the set of
properties from the informal standardization documents and encode them as
analyzable queries in the formal model.

  * •

The presented method eliminates the need for resource-intensive physical
experiments by utilizing NLP techniques. This approach promotes an automated
and efficient protocol analysis pipeline, reducing reliance on specialized
setup and equipment while minimizing the risk of error. Consequently, our
method significantly improves the accessibility and scalability of the
analysis.

  * •

In particular, a novel approach to doing formal analysis employing two
complementary NLP solutions is presented. The first, a traditional heuristic,
uses information extraction techniques to identify named entities
(identifiers) and semantic roles within the protocol text. The second
leverages a pre-trained language model to extract corresponding entities
(identifiers) and formal properties within the sentences.

In summary, our proposed method offers a novel, automated, and efficient
approach to 5G Wireless Communication Protocol analysis. By leveraging NLP
techniques, we can analyze the protocol more efficiently, and intuitively,
addressing several of the limitations inherent in current approaches.

The remainder of this paper is structured as follows. Section 3 provides a
comprehensive overview of the methodologies and the algorithms we developed
for this study. The experimental setup and the evaluation metrics used for the
study are detailed in Section 4. The results and analysis of these experiments
are presented in Section 5. We provide a discussion and interpretation of
these results in Section 6. Finally, the conclusions are drawn in Section 7.

##  III Methodology

Traditional IE tasks typically rely on annotated datasets for NER, with
entities for RE tasks being pre-defined. However, our research encounters two
major obstacles: the scarcity of datasets and the unavailability of pre-
defined entities for RE tasks. To address these challenges, we propose a novel
method.

(1) Regarding the annotated dataset, domain experts have noted that tables
within these protocols are rich in domain-specific terminologies, most of
which are the identifiers we aim to extract. We have leveraged the inherent
structure of these tables embedded in the protocols as a natural means to
extract these identifiers.

(2) Addressing the issue of unavailable pre-defined entities in the RE task,
we proposed a pipeline capable of concurrently extracting identifiers and
relations (also considered as formal properties) during the training process.

The whole system is illustrated in Figure. 1 which includes two components of
Information Extraction and Formal Analysis .

Figure 1: System Diagram

Based on the constructed dataset, we proposed a hybrid joint training model
based on pretrained language model. Specifically, we first locate and extract
the identifiers and use the extracted identifiers and the context sentence as
latent representation to classify the relations between the identifiers.

###  III-A Heuristic Information Extraction

The 5G Wireless Communication Protocol comprises both structured data (such as
tables) and unstructured data (such as text). For structured data in the PDF
version of the protocol, we employ Tabula [14] to extract tables embedded
within the file, storing the extracted tables in a list of dataframes. In
protocols like 5G, tables are integral data repositories, housing critical
operational and management information, and encapsulating domain-specific
terminologies. These terminologies, encapsulating complex protocol features
and representing the unique language of this field, are structured within
tables, presenting details on operational parameters, network configurations,
routing information, device details, resource management, and security
measures. To leverage the rich terminology within tables and maximize the use
of information from the extracted dataframes, we collect all the words in the
dataframes and apply a hierarchical filter to exclude unwanted words. The
first filter is a general one, using stopwords to filter out generic words.
The second is a domain filter, utilizing a term list provided by domain
experts to filter out unrelated domain terms.

For unstructured data, we leverage Stanford CoreNLP [15] to extract sources,
targets, and relations. Our primary focus is on identifying identifiers and
their formal properties, where the source and target represent the
identifiers, and the relation symbolizes the formal properties between these
identifiers. We utilize CoreNLP’s Part-of-Speech (POS) tagging and dependency
parsing pipelines in this process. POS tagging assigns suitable part-of-speech
tags (such as noun, verb, adjective, etc.) to each token in a given text.
Concurrently, dependency parsing evaluates the grammatical structure of a
sentence based on the dependencies between words.

We apply the filtered domain entity list to the extracted source-relation-
target (SRT) triples from CoreNLP, filtering the source and target tokens in
the SRT triple with our entities of interest. Following the same method, we
filter out unwanted predicates to specify the relations between identifiers as
shown in Figure 2. Terminology extraction through Tabula and SRT extraction
through coreNLP is illustrated in Information Extraction in Figure. 1 and
detailed depicted in Figure.2

Figure 2: Information Extraction Pipeline

###  III-B Identifier and Formal Property

Because our study is a pivot and preliminary research to analyze formal
protocol utilizing natural language techniques, the main objective of this
research is to propose an analysis framework and make a proof concept of the
whole pipeline. Consequently, we do not strictly define the identifiers used
in this study. Instead, we adopt domain-specific terminologies from
cybersecurity as our identifiers.

In the context of formal verification, for four main formal properties
‘Authentication’, ‘Accounting’,‘Integrity’,‘Confidentiality’ [1], they may
take the shape of state transition diagrams, preconditions, postconditions,
invariants, and more abstract concepts used to mathematically model the
protocol’s operation. In communication protocols, these formats are written in
natural language and may be classified into formal properties. For instance,
in the sentence ‘The RRCReconfiguration includes the masterCellGroup
containing the reportUplinkTxDirectCurrentTwoCarrier’, the identifiers are
‘RRCReconfiguration’ and ‘masterCellGroup’. The formal property connecting
these two identifiers is ‘Integrity’, but the word ‘includes’ in the sentence
represents ‘Integrity’. In this study, we do not train the model to classify
formal properties directly. Instead, we predict the specific words
representing formal properties. The representative words, provided by domain
experts, are shown in Table II.

###  III-C Hybrid Joint Sequence Labeling and Text Classification Task

The main goal of our project is to extract the identifiers and their relations
from the sentences in the communication protocol. The traditional NLP approach
is to locate and extract the annotated tokens from sentences which is purely a
sequence labeling task. Yet, we can construct the annotated dataset from
communication protocol by using specific NLP tools such as coreNLP, the
annotated dataset includes the identifiers and relations and their
corresponding position information in the sentences. Considering the project
goal of formal analysis based on communication protocol, we want to extract
the identifiers and we want to know the relation or the formal property
between the identifiers. So we convert the pure sequence labeling task into
two sub-tasks: token classification task for identifiers extraction, and text
classification for relation or formal proprieties prediction. For the token
classification task of identifiers extraction, we follow the finetune-on-pre-
trained models’ pipeline. Pretrained models like BERT, RoBERTa, and GPT have
been trained on massive amounts of text data. These models learn general
features from the language like syntax, semantics, and word relationships. By
fine-tuning these models on a specific task, we can leverage these previously
learned features, allowing the model to perform well even with a smaller
amount of task-specific data. Specifically, we used a pre-trained language
model (PLM) and build a multi-layer perceptron (MLP) classifier on top of the
PLM. For the text classification task of formal property prediction, we
consider the information dependency of formal properties that depend on the
identifiers also the context information for here which is the sentence. We
concatenate the hidden state of the sentence and identifiers and pass the
integrated latent representation to the formal property classifiers to predict
the relation between the two identifiers under the context(sentence).

We designed three models to investigate the dependency of formal property on
the identifiers and context. The first one is a disjoint training model in
that we build two classifiers to separately extract identifiers and formal
properties. The second one is jointly training the identifiers and formal
properties, so we only consider the dependency between identifiers and formal
properties. The third one is based on our proposed joint method which
considers the dependency of formal properties.

(a) Disjoint Model

(b) Joint1 Model

(c) Joint2 Model

Figure 3: Model Architectures

####  III-C1 Dataset

In this study, we use the extracted and filtered SRT triples from core NLP and
corresponding sentences as our annotation dataset. Our goal is to extract the
identifiers and their corresponding formal properties. Therefore, we designate
the source as ’identifier0’, the target as ’identifier1’, and the relation as
the ’formal property’. Consequently, we have two entity types: ’identifier0’
and ’identifier1’. As for relation types, we adhere to the list of interested
relations provided by domain experts as shown in Table. II. In our case, we
have a total of 23 relations, thus yielding 23 relation types. Given the
influence of sequence length on NER task performance highlighted in prior
research [11], we prepend the corresponding sentence to the identifier
triples, as shorter sequences have been shown to yield better results in named
entity extraction.

####  III-C2 Model Selection

In this study, we used transformer-based [16] Masked Language Model (MLM),
since the bidirectionality property enables such a model to use context from
both left and right to predict the masked word, which is crucial for
accurately predicting each token’s label in sequence labeling tasks. Also, we
need the sentence information as another dependency factor of formal
properties, the bidirectional understanding of MLM can provide a more holistic
view of a sentence, capturing relationships and dependencies between all parts
of the input and providing better sentence information. Combined the two
factors, we used BERT as our base model, BERT [17] is a pre-trained MLM, that
has shown state-of-the-art performance on a wide array of NLP tasks, including
sequence labeling tasks like NER [18] and text classification task which is
beneficial for identifiers extraction and formal property classification in
our project.

####  III-C3 Training Strategy

Our study aims to extract the identifiers and the formal property between the
extracted identifiers under a specific context. Given their interdependent
nature, to examine the inter-dependency of identifiers and formal properties,
we propose three training strategies.

The first approach involves training classifiers for identifiers and formal
properties separately. Specifically, we pass the hidden state of the sentence
separately to the identifier classifier and the formal property classifier. We
then train both classifiers in parallel to obtain their respective logits.
This process is illustrated in Figure. 3(a). The second one is jointly
training the identifiers and formal properties, that we first pass the
identifiers to the classifier and use the logits of identifiers as input and
pass the logits into formal property classifiers to get the formal property
logits. This method did not consider the context information influence on the
formal properties. The third one is based on our proposed method that builds
the hybrid joint training model to extract the identifiers first. Then use the
identifiers logits and sentence hidden states to predict the formal properties
of the identifiers under the context (sentences). The algorithm is explained
in Algorithm 1 and model architecture is shown in 3(c). Since formal
properties depend on the identifiers, even if two identifiers concurrently
appear, under different contexts (sentences), the formal properties between
the identifiers may change. Our hypothesis is that the third model that not
only considers identifiers information but also takes context information to
the formal properties prediction may provide the most representative
information among the three models. Representative in this research is the
prediction accuracy of formal properties.

Input Sequence of tokens $X=\\{x_{1},x_{2},...,x_{n}\\}$
Output Identifiers label $I=\\{0,0,1,1,...\\}$,
Formal property label $P=\\{0,0,0,10,...\\}$

1:procedure Joint_Protocol_Extraction($X$)

2: $E=embedding(X)$ $\triangleright$ Embedding from pretrained model: BERT

3: $H=pretrained\\_model(E)$ $\triangleright$ Get hidden state $H$ from
pretrained model like BERT

4: $I^{\prime}=identitfier\\_classifer(H)$ $\triangleright$ Predict the
identifiers label from hidden state

5: $I=argmax(softmax(I^{\prime}))$

6: $H^{\prime}=concatenate(H,I^{\prime})$

7: $P^{\prime}=fp\\_classifer(H^{\prime})$ $\triangleright$ Predict the formal
property label from hidden state

8: $P=argmax(softmax(P^{\prime}))$

9:return $I=\\{0,0,1,1,...\\},P=\\{0,0,0,10,...\\}$

Algorithm 1 Algorithm of Joint Identifier and Formal_Property Extraction

##  IV Experiment Setup

###  IV-A Model Configuration

We utilize ’bert-base-cased’ as our selected BERT model, with the number of
layers (L) set to 12, and a hidden dimension size of 768. Given our goal to
leverage pre-trained knowledge and transfer it to our domain-specific task of
in-domain token classification, we do not freeze the pre-trained layer. On top
of the PLM, we add a classification layer with two MLPs for identifiers and
formal properties, as depicted in Figures. 3(a) and 3(c). The model
configuration and hyper-parameter settings are shown in Table I.

TABLE I: Model Config and Hyperparameters Setting Model Config |  Hyperparameters
---|---
No of layers |  12
Hidden Dimension |  768
Learning rate |  1e-4
num_epochs |  20 (40 only for Joint1)
Optimizer |  AdamW
weight_decay |  0.1
no_regularization layer |  ”bias”, ”ln_1”, ’ln_2’

###  IV-B Raw Input Documents to Dataset Construction

The source document that we used to analyze in this research is ‘5G NR Radio
Resource Control (RRC) Protocol specification’ [8]. This document outlines the
procedures and command formats associated with 5G communication at the RRC
layer, serving as a crucial source for understanding and exploring 5G
communication protocols.

####  IV-B1 Identifier Categorization

For the filtered source and target elements in the SRT triples derived from
CoreNLP (as shown in Figure 2), we label the source as ‘identifier0’ and the
target as ‘identifier1’ in each sentence. Tokens not classified as identifiers
are categorized as ‘O’. Therefore, we have two types of identifiers, and the
total number of token labels is three, including the non-identifier tokens.

####  IV-B2 Formal Property Categorization

We use the formal property list provided by domain experts for categorization,
as shown in Table II. There are 23 formal property words, meaning we have 24
types of labels, including non-formal property tokens represented by ‘O’.

TABLE II: Predicate Convert Formal_Property Formal Property Word List |  Formal Property Category
---|---
‘decode’,‘encode’ |  Confidentiality
‘verify’ |  Integrity
‘access’, ‘reestablish’ |  Authentication
‘count’ |  Accounting
‘build’,‘complete’, ‘append’,‘belong’,‘store’,‘contain’, ‘include’,‘combine’ |  Belong
‘imply’,‘establish’,‘modify’, ‘denote’,‘utilize’,‘set’, ‘change’,‘define’,‘=’ |  Generation

###  IV-C Training and Evaluation

We employ a hierarchical regularization strategy to avoid over-penalizing
during training, which might lead to an under-fitting issue. For all bias
modules, MLP layers, and linear layers in the multihead_attnetion component of
the transformer block, we refrain from regularization during training. The
details are elaborated in Table I. All results are presented in the ensuing
results section.

We adhere to the standard evaluation metrics used in sequence labeling tasks,
which involve computing the ratio of correctly extracted identifiers and non-
identifier tokens to the total number of tokens in the sequence:
$acc\\_idf=(n_{(correct\\_idf)}\quad+\quad\
n_{(incorrect\\_idf)})/n_{tokens}$. This measurement provides an overview of
the model’s performance in terms of correctly identifying both the entities
(identifiers in our case) and other tokens. The $n_{(correct\\_idf)}$ denotes
the total number of correctly extracted identifiers in a batch,
$n_{(incorrect\\_idf)}$ denotes the total number of correctly extracted non-
identifiers in a batch, $n_{tokens}$ is the total number of tokens in a batch.

For formal properties, we adopt a slightly different evaluation metric. Given
the complexity and potential multi-label nature of formal properties, we
evaluate the performance by $acc\\_fprop=n_{(correct\\_fprop)}/batch\\_size$,
which is the ratio of correctly predicted properties to the total batch size,
and the$n_{(correct\\_fprop)}$ is the correct predicted formal properties in a
batch, ‘batch_size’ is the number of samples in a batch.

##  V Result Analysis

Heuristic Information Extraction: In the information extraction task, we
specifically concentrated on two sub-tasks: extracting terminologies and
performing SRT extraction. The extracted terminologies are used as a
dictionary to further filter the SRT triples. The filtered SRT triple will
only contain the in-domain terminologies in the source and target.

Terminology Extraction:  For this phase, we used a 5G communication protocol
spanning $1197$ pages as our source document. Utilizing tabula, we were
successful in extracting $696$ tables from this document. These tables were
then serialized into lists of words, which subsequently allowed us to extract
a total of $10905$ terminologies.

####  V-1 Identifier and Formal Property Extraction

By using coreNLP, we extract $19586$ source-relation-target-sentence groups
from the documents. We incorporated the terminologies derived from the
previous step, adhered to a hierarchical filtering strategy, and applied
lemmatization. This process yielded $16244$ groups.

The formal property list provided by the domain expert is as in table II, so
the total formal_property type is 24 including ’O’ which stands for not
formal_property. By filtering for the formal_property type of relation within
the original SRT groups, we retained $2495$ SRT groups. Since the majority of
the sentences within these groups contained fewer than 150 tokens, we removed
sentences exceeding 200 tokens from the groups and set a maximum sequence
length of $256$ in BERT. Consequently, we ended up with $1486$ source-
relation-target-sentence groups.

Hybrid Training Disjoint Result: : In Disjoint training, we separately train
the identifier classifier and the formal property classifier in parallel,
which means the loss in formal properties did not inherit logits from
identifiers. From the loss result, we see that the identifiers’ loss
converged, but the formal property loss did not converge as showed in
Figure.4(b). For accuracy, identifiers extraction accuracy is almost 60% in
Figure. 4(a). We have 3 labels for identifiers, the random guess is 30% for
each label. But for formal properties prediction, the result is in line with
the loss which means overfit for FP training which indicates an overfitting
problem.

Hybrid Training Joint Result:  In the joint training process, we implemented
two joint strategies to investigate the sentence influence on formal property
prediction. The first model (Joint1) did not take the sentence hidden state to
predict formal property, we only pass the identifier logits after the
identifier classifier to the formal property classifier to get the formal
property logits and make the final formal property prediction. In the second
joint model (Joint2), we concatenate the sentence hidden state and identifier
logits and pass the merged hidden information to the formal property
classifier. On comparing, we found the formal property prediction accuracy is
much lower in Joint1 than in Joint2, but after around 12 epochs, we found the
formal property prediction in Joint1 reached a similar performance as Joint2
which indicates the sentence influence may accumulate during training and make
contribute to the formal property prediction through the identifiers’ logits.
Another thing that needs to be noticed is in Joint1 the accuracy is still in
an increasing trend, but in Joint2, the performance for both identifiers and
formal properties seems static. So we extend training sessions from 20 to 40
of the Joint1 strategy, and the result is shown in Figure. 5(b). Results show
that both the loss and accuracy improved. But the prediction performance of
formal properties is with limitation. Partially because the formal properties
training depends on the identifiers training, which still faces the
overfitting problem. So the limited improved identifier training also
restricts the formal property training.

TABLE III: Identifier & Formal Property Training Accuracy Train_Strategy |  Train_Idf |  Train_Property |  Valid_Idf |  Valid_Property
---|---|---|---|---
Disjoint |  0.3339 |  0.9536 |  0.3993 |  0.4351
Joint1 |  0.3773 |  0.6170 |  0.5343 |  0.2799
jpint1_40 |  0.4029 |  0.9050 |  0.3927 |  0.4253
Joint2 |  0.3356 |  0.9523 |  0.3127 |  0.4459

(a) Disjoint Training Acc

(b) Disjoint Training Loss

Figure 4: Disjoint Training Accuracy

(a) Joint1 20 Epoch Training Accuracy

(b) Joint1 40 Epoch Training Accuracy

(c) Joint2 20 Epoch Training Loss

Figure 5: Joint Training Result

Further, we investigate the incorrect prediction distribution of the Joint1
model. We extract the incorrect perdition result from the last training epoch.
We noticed that for the identifier incorrect prediction, most of the incorrect
prediction is non-identifiers in which the ground truth label should be ‘1’ or
‘2’, but incorrectly predicted to non-identifier label ‘O’ as shown in Figure.
6(b). This result may indicate an unbalanced label in the identifiers because
most of the token labels in sentences are labeled ‘O’, the non-identifier,this
may cause very unbalanced label distribution in the training. In formal
property incorrect predictions, we found that most of the wrongly predicted
labels are ‘3’ which predicted to ‘4’ or ‘17’ as shown in Figure. 6(a). Based
on partial formal property distribution in Table. IV, we noticed the incorrect
prediction of formal property is not due to unbalanced data, this may result
from the identifier dependency, which if the identifier prediction is wrong,
the formal property will also predict incorrectly. We will discuss this in the
following section.

(a) Joint1 Incorrect Formal Property Prediction

(b) Joint1 Incorrect Identifier Prediction

Figure 6: Joint1 Incorrect Prediction TABLE IV: Formal_property_Stats (partial) formal property |  formal property IDs |  formal property Counts
---|---|---
include |  3 |  620
contain |  4 |  327
utilize |  17 |  10

##  VI Discussion and Limitation

###  VI-A Heuristic Terminology and Semantic Roles Extraction

In the preparation of the dataset, we employed Stanford’s CoreNLP to extract
SRT triples, utilizing its deterministic PTBTokenizer for tokenization. For
the construction of the model, the BERT architecture was selected as the
framework. Both require the processing of text into individual tokens.
However, they differ significantly in their tokenization methodologies. While
CoreNLP employs PTBTokenizer, which operates on a deterministic basis, BERT
makes use of a word-piece tokenizer that subdivides words into subword units.
This discrepancy in tokenization between PTB and the word piece can result in
differing token representations for the same lexical item, adding a layer of
complexity to the alignment and integration of these two components within our
research workflow.

###  VI-B Training Result and Analysis

The results from both the disjoint and joint training approaches revealed that
the joint training strategy outperforms the disjoint method in both an
identifier and formal property tasks. The reason lies in the sequential
architecture of the joint training model. In this model, the gradient from the
final loss first passes through the formal property classifiers and then both
the identifier classifier and the pre-trained BERT layers during back-
propagation. In contrast, even though the identifier and formal properties
classifiers share the same final loss in disjoint training, the gradient
passes equally to both classifiers. The formal properties loss enhances the
final loss that generates the gradient, but in disjoint training, the
identifier and formal property classifiers don’t share sequential back-
propagation. Consequently, the identifier classifier learns more through the
formal properties loss, which influences the identifier classifier’s learning
behavior. For our extended training of the Joint1 model, we noticed some
improvement compare to the Disjoint model and Joint2 model. But the
improvement is not significant, the main reason behind is an overfitting
problem. We tried several training strategies to alleviate this problem, such
as increasing the dropout rate, using higher regularization weight in the
optimizer, increasing model complexity, etc, but did not see significant
improvement. This indicates the overfitting problem is not just due to model
complexity or experiment setup, whereas dataset size may be the most crucial
reason since we only have around $1400$ training samples. Based on the result,
in the future, we will construct more data and implement the same training
strategy to investigate the performance. Another aspect that may influence the
overall performance is unrestricted identifiers and formal properties, this
problem may be trivial since our goal is to provide a proof of concept of the
pipeline. We plan to employ a more fine-grained dataset to refine this
procedure in future studies.

For the incorrect predictions, we found unbalanced label distribution present
in the identifier and consequently influence formal properties predictions. In
the future, we will mask out the non-identifier tokens during back-propagation
which will provide more guided loss on the identifiers and may also improve
the prediction of formal property tokens.

###  VI-C Disjoint and Joint Training

We used the averaged result to represent the performance of each model in
Table III. The experimental results underscore the superior performance of the
Disjoint model, which trains identifiers and formal properties in parallel.
However, it is crucial to recognize the inherent complexity of our project.
While both identifiers and formal properties are dependent on the sentence
structure, the latter also relies heavily on the former. Neglecting the
influence of identifiers on formal properties could lead to significant design
flaws in the experiment. Such oversight may create a deceptive impression of
improved training performance, masking underlying issues that could undermine
the integrity of the model. In future research, we will delve into the model
architecture and find the optimal solution to balance sentence and identifier
influence on formal properties’ prediction.

##  VII Conclusion and Future Research

In this paper, we introduce an innovative approach to performing formal
verification analysis using NLP-assissed techniques. Our method involves an
automated process pipeline for extracting terminologies and semantic roles
from domain-specific documents, which subsequently form the foundation of our
annotation training dataset. Moreover, we employ a transformer-based language
model as the backbone for our classifiers, further enhancing the efficacy of
our approach. Upon comparing the results from disjoint and joint training
strategies, we find that only passing the identifier logits to extract formal
properties slightly outperforms other methods. This proposed pipeline
illustrates the potential of our novel approach in the realm of formal
verification analysis.

This study serves as a fundamental exploration of the application of NLP for
formal verification analysis, showcasing our preliminary experimental results.
As we continue to develop our approach, we aim to fine-tune our methodologies,
address identified challenges, and enhance the performance and versatility of
our model in subsequent research pursuits. Our future efforts will focus on
refining the dataset with ground-truth identifiers and verified formal
properties. With regards to the model, we intend to persist with the joint
approach, while seeking solutions to further improve the accuracy of
identifier extraction and FP prediction, thus ensuring that our methodology
remains robust and adaptable to diverse formal verification analysis tasks.

## Acknowledgment

This effort was sponsored by the Defense Advanced Research Project Agency
(DARPA) under grant no. D22AP00144. The views and conclusions contained herein
are those of the authors and should not be interpreted as necessarily
representing the official policies or endorsements, either expressed or
implied, of DARPA or the U.S. Government.

## References

  * [1] J. Yang and Y. Wang, “Formal and fuzzing amplification: Targeting vulnerability detection in 5g and beyond,” _arXiv preprint arXiv:2307.05758_ , 2023.
  * [2] J. Yang, S. Arya, and Y. Wang, “Formal-guided fuzz testing: Targeting security assurance from specification to implementation for 5g and beyond,” _arXiv preprint arXiv:2307.11247_ , 2023.
  * [3] D. Basin, J. Dreier, L. Hirschi, S. Radomirovic, R. Sasse, and V. Stettler, “A formal analysis of 5g authentication,” in _Proceedings of the 2018 ACM SIGSAC conference on computer and communications security_ , 2018, pp. 1383–1396.
  * [4] J. Yang, Y. Wang, T. X. Tran, and Y. Pan, “5g rrc protocol and stack vulnerabilities detection via listen-and-learn,” in _2023 IEEE 20th Consumer Communications & Networking Conference (CCNC)_. IEEE, 2023, pp. 236–241.
  * [5] S. R. Hussain, M. Echeverria, I. Karim, O. Chowdhury, and E. Bertino, “5greasoner: A property-directed security and privacy analysis framework for 5g cellular network protocol,” in _Proceedings of the 2019 ACM SIGSAC Conference on Computer and Communications Security_ , 2019, pp. 669–684.
  * [6] Z. Wang and Y. Wang, “Nlp-based cross-layer 5g vulnerabilities detection via fuzzing generated run-time profiling,” _arXiv preprint arXiv:2305.08226_ , 2023.
  * [7] J. Zhang, Q. Wang, L. Yang, and T. Feng, “Formal verification of 5g-eap-tls authentication protocol,” in _2019 IEEE fourth international conference on data science in cyberspace (DSC)_. IEEE, 2019, pp. 503–509.
  * [8] 3GPP, “5G; NR; Radio Resource Control (RRC); Protocol specification,” 3rd Generation Partnership Project (3GPP), Technical Specification (TS) 38.331, 07 2020, version 16.1.0.
  * [9] D. Jurafsky and J. H. Martin, “Speech and language processing: An introduction to natural language processing, computational linguistics, and speech recognition.”
  * [10] V. Khetan, E. Wetherley, E. Eneva, S. Sengupta, A. E. Fano _et al._ , “Knowledge graph anchored information-extraction for domain-specific insights,” _arXiv preprint arXiv:2104.08936_ , 2021.
  * [11] S. Yuan and C. Lipizzi, “Information extraction in domain and generic documents: Findings from heuristic-based and data-driven approaches,” _arXiv preprint arXiv:2307.00130_ , 2023.
  * [12] M. T. Alam, D. Bhusal, Y. Park, and N. Rastogi, “Cyner: A python library for cybersecurity named entity recognition,” _arXiv preprint arXiv:2204.05754_ , 2022.
  * [13] M. Banko and O. Etzioni, “The tradeoffs between open and traditional relation extraction,” in _Proceedings of ACL-08: HLT_ , 2008, pp. 28–36.
  * [14] J. Mendes and J. Saraiva, “Tabula: A language to model spreadsheet tables,” _arXiv preprint arXiv:1707.02833_ , 2017.
  * [15] C. D. Manning, M. Surdeanu, J. Bauer, J. R. Finkel, S. Bethard, and D. McClosky, “The stanford corenlp natural language processing toolkit,” in _Proceedings of 52nd annual meeting of the association for computational linguistics: system demonstrations_ , 2014, pp. 55–60.
  * [16] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [17] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” _arXiv preprint arXiv:1810.04805_ , 2018.
  * [18] P. Liu, W. Yuan, J. Fu, Z. Jiang, H. Hayashi, and G. Neubig, “Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing,” _ACM Computing Surveys_ , vol. 55, no. 9, pp. 1–35, 2023.