# Space-Time Attention with Shifted Non-Local Search

Kent Gauen & Stanley Chan
Department of Electrical and Computer Engineering
Purdue University
<EMAIL_ADDRESS>
###### Abstract

Efficiently computing attention maps for videos is challenging due to the
motion of objects between frames. While a standard non-local search is high-
quality for a window surrounding each query point, the window’s small size
cannot accommodate motion. Methods for long-range motion use an auxiliary
network to predict the most similar key coordinates as offsets from each query
location. However, accurately predicting this flow field of offsets remains
challenging, even for large-scale networks. Small spatial inaccuracies
significantly impact the attention module’s quality. This paper proposes a
search strategy that combines the quality of a non-local search with the range
of predicted offsets. The method, named Shifted Non-Local Search, executes a
small grid search surrounding the predicted offsets to correct small spatial
errors. Our method’s in-place computation consumes 10 times less memory and is
over 3 times faster than previous work. Experimentally, correcting the small
spatial errors improves the video frame alignment quality by over 3 dB PSNR.
Our search upgrades existing space-time attention modules, which improves
video denoising results by 0.30 dB PSNR for a 7.5% increase in overall
runtime. We integrate our space-time attention module into a UNet-like
architecture to achieve state-of-the-art results on video denoising.

##  1 Introduction

Figure 1: Comparing the Search Space of Attention Modules. (From left to
right) ViT uses an exhaustive, global grid search which is computationally
costly (Dosovitskiy et al., 2020). A non-local search can be implemented
efficiently but does not shift the search space according to the motion
between frames (Hassani et al., 2023). The predicted offsets used in Guided
Deformable Attention allow for long-range dependencies, but the flow fields
contain small spatial inaccuracies (Liang et al., 2022b). Our method, the
Shifted Non-Local Search, combines the quality of a non-local search with the
range of predicted offsets. It executes a small grid search surrounding the
predicted offsets to correct the small spatial error.

Attention modules form data-dependent receptive fields to aggregate related
features from arbitrary coordinates. This functionality is considered to be
central to the success of large-scale networks (Dosovitskiy et al., 2020;
Hassani et al., 2023; Tian et al., 2020; Liang et al., 2022b). Recent efforts
aggregate features across frames of a video, enabling deep networks to learn
temporal representations of a scene. For images, the receptive fields are
often bounded by a window surrounding the query location to reduce computation
and the risk of overfitting. However, across frames of a video, this window
must shift to data-dependent locations according to the motion. Long-range
offsets are required, such as optical flow or nearest neighbors field (Barnes
et al., 2010; Ranjan & Black, 2017a).

Non-local search strategies, such as NATTEN, provide excellent short-range
receptive fields (Hassani et al., 2023). However, this category of method does
not offset the search window, so it cannot handle the motion inherent to a
space-time search. Alternative methods, such as Guided Deformable Attention,
predict long-range offsets (e.g. optical flow) using an auxiliary network to
accommodate motion (Liang et al., 2022b). However, accurately predicting flow
fields remains an open challenge, even for large-scale networks (Butler et
al., 2012).

This paper combines the quality of the non-local search with the range of
predicted offsets. Our method, named Shifted Non-Local Search (Shifted-NLS),
executes a small windowed grid search surrounding the predicted offset. For a
marginal increase in wall-clock runtime, our search method acts as a
correction step to the predicted offsets. In addition, our grid search is
differentiable, which allows networks to learn long-range offsets. Our method
works for attention because, unlike optical flow’s goal of estimating apparent
motion, standard attention modules are defined through a grid search. We show
our search method improves video alignment, upgrades existing space-time
attention modules, and enables a state-of-the-art architecture for video
denoising.

Critically, this paper also offers a practical means to compute the Shifted
Non-Local Search. An important related work, named N3Net, already offers a
similar method (Plötz & Roth, 2018). However, their method is not presented in
the context of attention modules and requires integer-spaced indexing, which
does not allow backpropagation through the grid indices. Additionally, the
N3Net search’s forward runtime is 3-7x slower than our implementation and
requires over a 10-25x spike in GPU memory. We believe these exorbitant
computational demands explain why the module has not been adopted in more
recent works on space-time attention, and our Pytorch-friendly module offers a
practical alternative (Paszke et al., 2019).

In summary, our contributions are: (i) We propose the shifted non-local search
module for space-time attention. The module corrects spatial errors of
predicted offsets using a high-fidelity windowed grid search. (ii) Our
implementation uses in-place computation to reduce computational demands
compared to previous work, using 10 times less memory and executing 3 times
faster than N3Net (Plötz & Roth, 2018). While our code is not explicitly
optimized for speed, our search’s runtime is only 1 - 2.5 times slower than a
highly-optimized space-only non-local search (Hassani et al., 2023). (iii) We
show our search method improves video denoising quality. Specifically, we
replace Guided Deformable Attention’s auxiliary network with our space-time
search to improve denoising quality by $0.30$ dB PSRN (Liang et al., 2022b).
We replace a single image denoiser’s global search with our space-time
attention module to improve the denoised quality by over $1$ dB PSNR. We also
present an architecture using our space-time attention block which achieves
state-of-the-art results.

##  2 Related Works

Space-Only Attention: Attention modules often use a modified search space to
be computationally efficient, and most of them search only space (Dosovitskiy
et al., 2020; Mou et al., 2021; Liu et al., 2021). SWIN uses an asymmetric
search space but can be efficiently computed using native Pytorch code.
Hassani et al. (2023) offers an efficient non-local search but cannot
accommodate long-range offsets. Xia et al. (2022) applies predicted offsets
for single images but suffers from the inaccuracy of using a network to
predict the flow field.

Space-Time Attention: Recent works propose temporal attention modules using
predicted offsets learned with auxiliary networks which are inspired by
Deformable Convolution (Dai et al., 2017). Examples of these methods include
the temporal mutual self-attention module (TSMA), the temporal deformed
alignment module (TDAN), and the guided deformable attention module (GDA)
(Liang et al., 2022a; Tian et al., 2020; Liang et al., 2022b). Each method
predicts pixel-level offsets and warps an adjacent frame to match a reference
frame. These methods all suffer from the limitation of learning these offsets.
Plötz & Roth (2018) proposed N3Net which does execute a shifted grid search,
but their implementation is not connected to attention modules and their
method also does not propagate gradients through its grid search. The
computational cost of their method is also significant; their method runs 3
times slower and requires 10 times more memory than ours. Video Non-Local
Bayes is a classical method that can be formulated to be an attention module
(Arias & Morel, 2018). Figure 1 compares the search space of related works on
a single frame.

Restoration Architectures: Presented concurrently with new attention modules,
authors often present new architecture designs for video restoration tasks.
TDAN is used for video super-resolution, and RVRT is applied to video super-
resolution, deblurring, and denoising (Tian et al., 2020; Liang et al.,
2022b). Their attention blocks apply to only frame pairs, while our attention
module can search multiple frames in parallel.

##  3 Method

Figure 2: The Shifted Non-Local Search for Space-Time Attention. This figure
depicts a space-time attention module using the Shifted Non-Local Search. The
query points are deformed using the predicted offsets. Next, a grid search is
executed surrounding the predicted offsets, and then the most similar
locations are chosen from the search window. These locations are aggregated
using a module such as Guided Deformable Attention.

###  3.1 Problem Setup

Preparation. An input video, ${\bm{V}}_{\text{in}}$, is projected with a
$1\times 1$ convolution, ${\bm{Q}}_{V}=\text{conv}_{1\times
1}({\bm{V}}_{\text{in}})$. Let ${\bm{Q}}_{V}$, ${\bm{K}}_{V}$, and
${\bm{V}}_{V}$ be the query, key, and value videos with shape $T\times H\times
W\times F$ denoting frames, height, width, and features. The total number of
pixels is $N_{P}$, and we denote $N_{Q}=\left[T\times\left(H\times
W\right)/S_{Q}^{2}\right]$ as the number of reference locations with $S_{Q}$
defined as the reference stride. We similarly define $N_{K}$ and $S_{K}$ for
the number of search locations and search stride. A video can be represented
as a database of patches with patch size $P$ using unfold,
${\bm{K}}_{P}=\text{unfold}({\bm{K}}_{V},P)$. The unfolded queries and keys,
${\bm{Q}}_{P},{\bm{K}}_{P}$, take shape $N_{Q}\times\left[F\times
P^{2}\right]$ and $N_{K}\times\left[F\times P^{2}\right]$.

Attention. Attention consists of two steps: search and aggregate. Searching
returns the top-$K$ most similar values and indices between the queries and
keys written as ${\bm{D}},{\bm{I}}=\text{search}({\bm{Q}}_{V},{\bm{K}}_{V},K)$
with shapes $N_{Q}\times K$ for $K$ neighbors. Aggregation combines the
patches indexed with ${\bm{I}}$, often as a sum weighted with values
$\exp(-\lambda{\bm{D}}[i,k])$, $\lambda>0$, or learned weights ${\bm{W}}[k]$.
Together it is written as
${\bm{V}}_{\text{out}}=\text{aggregate}({\bm{K}}_{V},{\bm{D}},{\bm{I}})$. A
single output patch can be written as
${\bm{V}}_{\text{out}}[{\bm{i}}]=\sum_{k=1}^{K}\exp(-\lambda{\bm{D}}[i,k]){\bm{V}}_{V}[{\bm{I}}[i,k]]$,
where ${\bm{i}}$ is there rasterized $i^{th}$ query location.

Searching. The search step computes the similarity between reference patches
from ${\bm{Q}}_{V}$ and patches from a search space within ${\bm{K}}_{V}$. The
search space is written as a tensor $\hat{{\bm{I}}}$ with size $N_{Q}\times M$
taking values $\hat{{\bm{I}}}[i,j]\in\\{1,\ldots,N_{K}\\}$ which is the
$j^{\text{th}}$ search patch of the $i^{\text{th}}$ reference patch. The
similarity function outputs a number given two input patches,
$s:\mathbb{R}^{F\cdot P^{2}}\times\mathbb{R}^{F\cdot P^{2}}\to\mathbb{R}$. The
similarity values are written as a tensor $\hat{{\bm{D}}}$ with shape
$N_{Q}\times M$ taking values
$\hat{{\bm{D}}}[i,j]=s\left({\bm{Q}}_{V}[{\bm{i}}],{\bm{K}}_{V}[\hat{{\bm{I}}}[i,j]]\right)$
for patches centered at $i$ and $\hat{{\bm{I}}}[i,j]$, respectively. The top-K
most similar patches, according to $\hat{{\bm{D}}}$, are returned from the
search, denoted ${\bm{D}},{\bm{I}}$, with shapes $N_{Q}\times K$. For example,
the searched keys of a Non-Local Search are defined using a spatial window
($W_{s}$) surrounding the query location. Say index ${\bm{i}}$ corresponds to
spatial coordinates $(a,b)$, then the $i^{th}$ row of the search space,
denoted $\hat{{\bm{I}}}[i,\cdot]$, contains indices
$\\{(a+l,b+m)\\}_{l,m=-W_{s}/2,-W_{s}/2}^{W_{s}/2,W_{s}/2}$.

###  3.2 The Shifted Non-Local Search

A shifted non-local search returns the similarities and indices for a search
space and input offsets,
${\bm{D}},{\bm{I}}=\text{search}({\bm{Q}}_{V},{\bm{K}}_{V},{\bm{F}})$. The
predicted offsets are denoted ${\bm{F}}$ and shift the center of the square
search window across time. The distances are computed as follows:
${\bm{D}}[i,k]=s\left({\bm{Q}}_{V}[{\bm{r}}[i]],{\bm{K}}_{V}[{\bm{r}}[i]+{\bm{g}}[k]+{\bm{F}}[{\bm{r}}[i]]]\right)$
where ${\bm{r}}$ is the list of $N_{Q}$ query locations and ${\bm{g}}$ is a
fixed grid of $N_{K}$ key locations. Figure 2 shows the grid search offset
across time (${\bm{r}}[i]+{\bm{g}}[k]+{\bm{F}}[{\bm{r}}[i]]$). Since
${\bm{F}}[\cdot]$ is generally a floating-point number, the value of the key
is computed with bilinear interpolation. Namely,
${\bm{K}}_{V}[{\bm{F}}[\cdot]]=\sum_{a,b}G(a,{\bm{F}}[\cdot,0])G(b,{\bm{F}}[\cdot,1]){\bm{K}}_{V}[a,b]$
which is efficiently computable for both a forward and backward passes (Jeon &
Kim, 2017). This step allows for gradients to backpropagate through the grid-
search window to learn offset values, ${\bm{F}}$.

While the idea is conceptually straightforward, it remains unclear whether
this grid search meaningfully improves search quality and how it should be
computed. First, we will convince the reader that optical flow only needs
small spatial corrections, which a grid search can approximate. Second, we
discuss why the Shifted Non-Local Search for videos should be executed in-
place, since standard tiling strategies for images do not apply. We note the
need for an in-place operator suggests a gap between the hardware design for
images and videos.

###  3.3 Why are predicted offsets not enough?

Figure 3: Predicted offsets are only a few pixels away from their optimal
location. This figure shows query points in the reference frame (top; black
points), and their counterparts in the adjacent frame shifted with optical
flow (bottom; blue points). The optical flow points are then corrected by a
grid search of size $41\times 41$ (bottom; yellow points). The two sets of
colored points in the bottom figure show the corrections are spatially similar
to the initial estimate from optical flow. The right subfigure plots the
distribution of these corrections. The peak value is positioned at the center,
indicating no correction is necessary for 3.5% of all cases. The two ellipses
form the 68% and 90% confidence intervals.

We investigate the limitation of predicting offset locations for attention by
reviewing results in the closely related research area of optical flow.

Millions of parameters for a 6-pixel error. The best methods for optical flow
today, according to the Sintel-Clean benchmark, report an average end-point
error of about 1 pixel (Butler et al., 2012). Meanwhile, the classical
pyramid-based method of 2014 reports an error of 6.73 pixels (Sun et al.,
2014). Although the average improvement of about 6 pixels is impressive, this
gap is closed using sophisticated training methods and network architectures
with millions of parameters. Some applications claim to hugely benefit from
the subpixel accuracy of these methods. However, it seems unlikely that _each
instance_ of an attention module will require its own auxiliary network with
millions of parameters to simply predict regions of high affinity.

Assessing the Error of Optical Flow for Attention. While the end-point-error
is compared against an optical flow groundtruth, we qualitatively find the
error to be similar when optical flow is used to estimate locations for
attention. Figure 3 qualitatively shows that optical flow’s errors are
concentrated in a small region surrounding the initial estimate, despite a
large search grid of size $41\times 41$. This supports our idea to execute a
small windowed grid search to correct of the predicted offsets’ errors.

###  3.4 An Inplace Computation

Figure 4: Video Dynamics Challenge Existing Computational Approaches.
Searching across time is computationally challenging because _spatially
adjacent patches in a reference frame have data-dependent spatial locations in
adjacent frames_. This figure shows two neighboring locations in one frame
(the chick and the kitten) move to separate spatial locations in the next
frame. The benefit of NATTEN’s tiling is lost because the search windows no
longer overlap (Hassani et al., 2023). The rightmost subfigure plots the
number of global memory reads, showing overlapping search spaces read a factor
of 3-5 fewer times.

Our In-Place Computation. Our in-place computation of the Shifted Non-Local
Search executes each query-key pair’s similarity within its own computational
unit, or CUDA thread (NVIDIA et al., 2020). Unlike other methods, our search
does not require storing additional data related to the video. However, the
global memory access pattern of our spatial window is undesirable for current
hardware, which increases the method’s runtime. Section 4.4 shows despite this
issue, our method is 3 - 7x faster than N3Net. Under some conditions, our
search is even faster than NATTEN’s optimized non-local search.

Limitations of NATTEN. NATTEN is designed to execute a non-local search with a
small runtime (Hassani et al., 2023). Their core efficiency comes from
reducing the number of reads from global memory by sharing global reads across
the threads of a block. This tiling principle cannot be extended to space-time
because the search windows shift to data-dependent locations, as depicted in
Figure 4. In general, these shifted windows will not overlap. Let $Q=3$ be the
tiled size and $W_{s}$ as the window size; then overlapping windows require
only $Q+W_{s}-1$ global reads while non-overlapping windows require $Q\cdot
W_{s}^{2}$ global reads. The far-right subfigure in Figure 4 plots these two
quantities, showing a significant disparity between the two cases.

Limitations of N3Net. A video can be represented as a database of patches of
size $P$ using unfold, ${\bm{K}}_{P}=\text{unfold}({\bm{K}}_{V},P)$. The
queries and keys, ${\bm{Q}}_{P},{\bm{K}}_{P}$, take shape
$N_{Q}\times\left[F\times P^{2}\right]$ and $N_{K}\times\left[F\times
P^{2}\right]$. Normally, operators can batch across the $N_{Q}$ and $N_{K}$
dimensions to control memory consumption. However, the data-dependent indexing
across space-time makes batching across the keys impossible. The entire key
database must be simultaneously represented in memory since each query patch
may access any key patch. If all queries are searched in parallel, the patch
database increases memory consumption by
$P^{2}\times\left(1/S_{Q}^{2}+1/S_{K}^{2}\right)$. For example, if $P=3$ and
$S_{Q}=S_{K}=1$, the memory consumption of the video increases by about
$18\times$.

###  3.5 Aggregating Space-Time Features

Once the top-K features are identified from a search, the features must be
aggregated. Our search is compatible with existing aggregation methods, such
as Guided Deformable Attention (GDA). However, GDA is limited to aggregating
features from a single frame. For our STAN architecture, we would like a
module to aggregate features across multiple frames, and we would like to
learn weights across the top-$K$ patches, similar to PacNet (Vaksman et al.,
2021). Specifically, if similarity is measured using the sum of squared
differences, then a patch from the output can be written as follows:
${\bm{V}}_{\text{out}}[{\bm{i}}]=\sum_{k=1}^{K}{\bm{W}}[k]\exp(-\lambda{\bm{D}}[i,k]){\bm{V}}_{V}[{\bm{I}}[i,k]]$,
where ${\bm{W}}$ are the learned convolution weights. To implement this logic,
we create a module to stack $K$ patches and apply 3D convolution to reduce the
stack across $K$. Stacking is slower than specialized code, but the intention
of this module is to foster future, creative development of aggregation
schemes. More details are in Supplemental Section 7.

##  4 Experiments

###  4.1 Video Frame Alignment

The Shifted Non-Local Search (Shifted-NLS) corrects the small spatial errors
of predicted offsets such as optical flow. This section illustrates the
significant impact of these small spatial errors by quantifying each search
module’s search quality on video alignment. Alignment quality quantifies how
well a search method can identify similar patches across space-time. The
Shifted-NLS method outperforms both the Non-Local Search (NLS) and the
predicted offsets by 3 dB PSNR.

This experiment uses the first 10 frames from the DAVIS training dataset
(Pont-Tuset et al., 2017). When searching and computing the TV-L1 optical
flow, we add a small amount of Gaussian noise ($\sigma^{2}=15$) to simulate
the uncertainty of the trained query and key values of an attention module
within a network during training (Zach et al., 2007). We compute alignment
quality using the PSNR between the noise-free aligned and reference images.
Both the Shifted-NLS and the NLS methods use our implementation since NATTEN’s
implementation of NLS requires specialized configurations.

Figure 5: The Shifted-NLS Improves Alignment Quality. The Shifted-NLS achieves
a better alignment with fewer search pixels than the NLS. Along each line the
search window increases, $\\{1,3,9,15,21,27,33\\}$. The size of each point
indicates the peak allocated memory when executing the search. A larger search
grid is slower and consumes more memory.

Figure 6: Video Motion Impacts the Alignment Quality. The y-axis plots the
difference in the alignment quality between the Shifted-NLS and the NLS. The
x-axis is an estimate of the video’s motion using the average absolute value
of the optical flow. Each point represents a single video sequence. The dotted
line is drawn at $0$ dB PSNR: no difference.

Figure 6 compares the alignment quality and runtime of the Shifted-NLS and the
NLS as the search space expands. Each point is associated with a spatial
window size, $W_{s}\in\\{1,3,9,15,21,27,33\\}$. A spatial window of size 1
indicates no search. As the window size increases, both the quality and
runtime also increase. The increase in PSNR plateaus when the window size is
about $W_{s}=9$ for the Shifted-NLS, while it plateaus around $W_{s}=21$ for
the NLS. This matches our intuition that optical flow contains small spatial
errors, which our grid search corrects. By using a small search window, say
$9$ pixels, the Shifted-NLS improves the alignment of optical flow by over 3
dB PSNR (marked by a solid-lined circle).

Figure 6 shows our Shifted-NLS method’s improvement depends on video motion.
Each point is the difference in PSNR between the Shifted-NLS and the NLS for
each video in the DAVIS training dataset. When motion is larger than about 3
pixels, Shifted-NLS improves the alignment quality by more than 3 dB PSNR.
When the average motion is less than 1 pixel, the Shifted-NLS actually
degrades the search quality. In the case of small motion, the offset values
act as added noise inside the Shifted-NLS.

Figure 7: Comparing alignment quality for different search methods. This
figure uses the indicated row’s search method to align the adjacent frame
(bottom) to the reference frame (second from bottom). The Non-Local Search and
the Shifted Non-Local Search use a search space of $11\times 11$. The bottom-
right corner in each aligned frame reports the PSNR between the aligned and
reference images; higher is better.

Figure 7 shows qualitative examples of the aligned images. For both the NLS
and the Shifted-NLS methods, the spatial search window is fixed to $11\times
11$. The NLS method creates a doubling effect because the search radius can
not escape the motion shifts. For example, the first number of the speed limit
sign (leftmost column) reads “5” rather than “2”. The “2” is dropped entirely
because pixels in the adjacent frame contain only dust pixels, not sign
pixels. The Shifted-NLS largely removes the doubling effect, but not entirely.
When the optical flow is inaccurate, a doubling effect still appears. For
example, in the third column, a face appears where only a fence should be
visible.

The errors from Predicted Offsets create a warping effect similar to
psychedelic art or the melting clocks of Salvador Dalí. The Shifted-NLS method
visually removes the warping effect. The wavy edges in many of the examples
are replaced with sharp ones. Instead, erroneous locations from the Shifted-
NLS fill regions with incorrect values. The football in the second column
fills a yellow stripe in the center grey pixels.

###  4.2 Upgrading Space-Time Attention

We upgrade Guided Deformable Attention (GDA) with our Shifted Non-Local Search
(Shifted-NLS) module to show the value of correcting the errors of predicted
offsets for video denoising (Liang et al., 2022b). GDA requires 9 offsets for
each pixel in the image. In the original network, 9 offsets are output from a
small convolution network whose input includes optical flow. Our method omits
the small network and searches the local region surrounding the optical flow.
In this experiment, our spatial window is $9\times 9$ and the temporal window
is fixed to $1$ by architecture design. The most similar 9 locations are
selected to replace the offsets from the network. Table 1 shows the denoising
quality improves when using our search method compared to using predicted
offsets. The improvement is roughly $0.30$ dB across all noise levels, which
is a quantitative improvement often attributed to an entirely new
architecture.

Table 1: Improving Video Denoising. [PSNR$\uparrow$/SSIM$\uparrow$/ST-RRED$\downarrow$] This table reports denoising results on the DAVIS dataset augmented using our Shifted Non-Local Search module. The original network aggregates features from 9 locations. Instead of identifying the locations with a network, our module executes the Shifted-NLS modules surrounding the optical flow offsets. $\sigma$ | Original | Shifted-NLS
---|---|---
10 | 38.69/0.966/0.004 | 38.90/0.967/0.004
20 | 35.32/0.933/0.013 | 35.58/0.936/0.012
30 | 33.39/0.902/0.026 | 33.68/0.907/0.024
40 | 32.02/0.873/0.042 | 32.35/0.881/0.040
50 | 30.93/0.844/0.062 | 31.30/0.854/0.058
Time (sec) | 23.86 | 25.56
Mem (GB) | 10.06 | 10.06

###  4.3 Space-Time Attention Network (STAN)

We integrate the Shifted Non-Local Search into our Space-Time Attention
Network (STAN). The architecture is a simple mixture of the UNet and RVRT
networks (Ronneberger et al., 2015). We train the network for video denoising
on the DAVIS train-val dataset (Pont-Tuset et al., 2017). We test the network
on the DAVIS testing dataset and the Set8 dataset (Tassano et al., 2020). Due
to space, we regulate details to Supplemental Section 6.4.

Table 2 shows our network achieves state-of-the-art results on video
denoising. We note the original RVRT network reports better results, but we
are unable to reproduce them (see Supplemental Section 6.3). We copy RVRT’s
training procedure from their paper since no training code is provided and use
the same procedure for both RVRT and STAN. Our method outperforms all other
published video denoising methods, which supports the hypothesis that the
Shifted Non-Local search is a useful module for space-time attention (Arias &
Morel, 2018; Tassano et al., 2020; Vaksman et al., 2021).

Table 2: State-of-the-art Video Denoising. [PSNR$\uparrow$/SSIM$\uparrow$/ST-
RRED$\downarrow$] This table reports state-of-the-art results on video
denoising. *RVRT and STAN explicitly use space-time attention. We report
reproduced RVRT results with further details in Supplemental Section 6.3.

| $\sigma^{2}$ | VNLB | FastDVDNet | PaCNet | RVRT (Reproduced)* | STAN*
---|---|---|---|---|---|---
Set8 | 10 | 37.26 | 36.44 | 37.06 | 36.66/0.955/0.003 | 37.19/0.960/0.002
20 | 33.72 | 33.43 | 33.94 | 33.47/0.918/0.011 | 34.27/0.931/0.007
30 | 31.74 | 31.68 | 32.05 | 31.65/0.885/0.022 | 32.58/0.905/0.013
40 | 30.39 | 30.46 | 30.70 | 30.38/0.855/0.035 | 31.39/0.880/0.021
50 | 29.24 | 29.53 | 29.66 | 29.41/0.829/0.052 | 30.46/0.856/0.030
DAVIS | 10 | 38.85 | 38.71 | 39.97 | 39.29/0.970/0.003 | 40.22/0.976/0.002
20 | 35.68 | 35.77 | 36.82 | 36.00/0.942/0.010 | 37.30/0.956/0.007
30 | 33.73 | 34.04 | 34.79 | 34.12/0.915/0.021 | 35.54/0.937/0.012
40 | 32.32 | 32.82 | 33.34 | 32.80/0.891/0.034 | 34.26/0.918/0.020
50 | 31.13 | 31.86 | 32.20 | 31.78/0.868/0.050 | 33.26/0.901/0.029
Time (sec) | 497.93 | 0.11 | 182.34 | 1.63 | 3.26
GPU Memory (GB) | 0.0 | 0.37 | 12.35 | 4.25 | 10.75
Parameters ($10^{6}$) | N/A | 2.4 | 2.9 | 12.8 | 12.1

Figure 8: Qualitatively Comparing Denoised Outputs. [PSNR$\uparrow$] RVRT and
STAN use space-time attention and are trained using the same procedure. STAN
recovers more small details than RVRT.

###  4.4 Computational Benchmarking

This section compares the computation for three non-local search strategies.
The Shifted-NLS and N3Net methods execute a shifted search, and NATTEN
executes a highly optimized non-shifted search. Figures 10 and 10 report
benchmark results of each search method executed on a $5$ frame video with
resolution $160\times 160$. The input resolution is limited by N3Net’s
enormous memory consumption. Due to NATTEN’s tiling, its reference stride is
fixed at $1$. The other two search methods vary the reference stride to $1$ or
$2$ as indicated by the dotted and solid lines, respectively.

Figure 9: The Shifted-NLS Consumes Less Memory. Our Shifted Non-Local Search
module consumes less than $20\times$ the memory of N3Net’s search method and
less than $9\times$ the memory of NATTEN. N3Net expands the video into a
database of patches, which requires an absurd consumption of GPU memory.
NATTEN’s tiling first reads an image chunk into shared CUDA memory to reduce
the number of global reads, growing the GPU consumption.

Figure 10: The Runtime of Shifted-NLS is Competitive. NATTEN can not be used
for a space-time search but acts as a highly efficient baseline. Importantly,
NATTEN’s efficient requires using a reference stride of $1$, while both other
search methods can use any reference stride. The Shifted-NLS method is $2.4$
times slower than NATTEN when the reference stride is $1$, but it is $2.5$
times faster when the reference stride is $2$.

Figure 10 reports memory consumption for an input video with $32$ features
across $4$ heads. N3Net requires dramatically more memory than the Shifted-NLS
module since it explicitly constructs a database of patches using the Unfold
operator. When the spatial window size is 9, N3Net consumes 10.27 GB of
memory, while the Shifted-NLS consumes 0.27 GB of memory. NATTEN constructs a
more sophisticated database of patches within the shared memory of a CUDA
kernel, which reduces memory consumption to only $1.28$ GB. NATTEN searches
pairs of frames, so parallel searching across space-time requires stacking
frames of the temporal search window along the batch dimension111We note this
temporal-window-stacking is not measured in NATTEN’s runtime. Our in-place
search does not copy any video data to reduce memory consumption.

Figure 10 reports runtimes using $2$ heads and a fixed search window of size
$9$. The Shifted-NLS module is slower than NATTEN when the reference grid is
fixed to $1$. This is expected as NATTEN is specifically designed for fast
execution under this condition. As the number of features grows, the runtimes
of NATTEN and Shifted-NLS converge. A factor of NATTEN’s increased runtime
stems from slower reads of the growing shared memory buffer, expanding to
accommodate the increasing number of features. NATTEN also must read from
memory (at least) twice: once from global memory and once from shared memory.
For a tile of $Q$ queries, the total number of global reads is reduced from
$F\left(Q\cdot W_{S}^{2}\right)$ to $F\left(Q+W_{S}-1\right)$. But NATTEN must
also read from shared memory a total of $F\left(Q\cdot W_{S}^{2}\right)$
times. As the number of features ($F$) grows, the growing number of total
reads increases the total runtime. N3Net is far slower than both methods;
N3Net’s runtime for a reference stride of $1$ is too slow to plot clearly. We
attribute this slow runtime to the construction of the patch database, which
requires using an enormous amount of GPU memory.

Notably, the Shifted-NLS is faster than NATTEN when the reference stride can
be set to $2$. Future iterations of NATTEN may allow for a reference stride
greater than $1$, but the premise of tiling requires overlapping search
spaces. So even if NATTEN is modified to allow for a flexible stride, the
impact on runtime is unclear.

##  5 Conclusion

This paper presents a Shifted Non-Local Search module for space-time
attention. We first observe the errors of offsets predicted from auxiliary
networks require only small spatial corrections. Rather than train a large-
scale network with millions of parameters, we propose using a small grid
search to correct these errors. Our in-place implementation of the Shifted
Non-Local Search avoids absurd memory spikes with a competitive runtime.
Correcting the small spatial errors corresponds to over a 3 dB improvement
when aligning adjacent frames. The module upgrades existing space-time
attention modules, such as Guided Deformable Attention, to improve video
denoising quality by over $0.30$ dB PSNR. We also integrate the module into a
simple architecture for video denoising and achieve state-of-the-art results.
As this module is designed for learning temporal representations, future work
can apply this method to additional computer vision tasks such as instance
segmentation and video synthesis.

## References

  * Arias & Morel (2018) Pablo Arias and Jean-Michel Morel.  Video denoising via empirical bayesian estimation of space-time patches.  _Journal of Mathematical Imaging and Vision_ , 60(1):70–93, 2018.
  * Barnes et al. (2010) Connelly Barnes, Eli Shechtman, Dan B Goldman, and Adam Finkelstein.  The generalized patchmatch correspondence algorithm.  In _Proceedings of the European Conference on Computer Vision_ , pp. 29–43. Springer, 2010.
  * Butler et al. (2012) D. J. Butler, J. Wulff, G. B. Stanley, and M. J. Black.  A naturalistic open source movie for optical flow evaluation.  In A. Fitzgibbon et al. (Eds.) (ed.), _Proceedings of the European Conference on Computer Vision_ , Part IV, LNCS 7577, pp. 611–625. Springer-Verlag, October 2012.
  * Dai et al. (2017) Jifeng Dai, Haozhi Qi, Yuwen Xiong, Yi Li, Guodong Zhang, Han Hu, and Yichen Wei.  Deformable convolutional networks.  In _Proceedings of the IEEE International Conference on Computer Vision_ , pp. 764–773, 2017.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Falcon et al. (2019) William Falcon et al.  Pytorch lightning.  _GitHub. Note: https://github.com/PyTorchLightning/pytorch-lightning_ , 3, 2019.
  * Hassani et al. (2023) Ali Hassani, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi.  Neighborhood attention transformer.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 6185–6194, 2023.
  * Hu et al. (2018) Jie Hu, Li Shen, and Gang Sun.  Squeeze-and-excitation networks.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 7132–7141, 2018.
  * Jeon & Kim (2017) Yunho Jeon and Junmo Kim.  Active convolution: Learning the shape of convolution for image classification.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 4201–4209, 2017.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Liang et al. (2021) Jingyun Liang, Jiezhang Cao, Guolei Sun, Kai Zhang, Luc Van Gool, and Radu Timofte.  Swinir: Image restoration using swin transformer.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 1833–1844, 2021.
  * Liang et al. (2022a) Jingyun Liang, Jiezhang Cao, Yuchen Fan, Kai Zhang, Rakesh Ranjan, Yawei Li, Radu Timofte, and Luc Van Gool.  Vrt: A video restoration transformer.  _arXiv preprint arXiv:2201.12288_ , 2022a.
  * Liang et al. (2022b) Jingyun Liang, Yuchen Fan, Xiaoyu Xiang, Rakesh Ranjan, Eddy Ilg, Simon Green, Jiezhang Cao, Kai Zhang, Radu Timofte, and Luc V Gool.  Recurrent video restoration transformer with guided deformable attention.  _Advances in Neural Information Processing Systems_ , 35:378–393, 2022b.
  * Liu et al. (2021) Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 10012–10022, 2021.
  * Loshchilov & Hutter (2016) Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  _arXiv preprint arXiv:1608.03983_ , 2016.
  * Mou et al. (2021) Chong Mou, Jian Zhang, Xiaopeng Fan, Hangfan Liu, and Ronggang Wang.  Cola-net: Collaborative attention network for image restoration.  _IEEE Transactions on Multimedia_ , 24:1366–1377, 2021\.
  * Niklaus (2018) Simon Niklaus.  A reimplementation of SPyNet using PyTorch.  https://github.com/sniklaus/pytorch-spynet, 2018.
  * NVIDIA et al. (2020) NVIDIA, Péter Vingelmann, and Frank H.P. Fitzek.  Cuda, release: 10.2.89, 2020.  URL https://developer.nvidia.com/cuda-toolkit.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems 32_ , pp. 8024–8035. Curran Associates, Inc., 2019.
  * Plötz & Roth (2018) Tobias Plötz and Stefan Roth.  Neural nearest neighbors networks.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Pont-Tuset et al. (2017) Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alexander Sorkine-Hornung, and Luc Van Gool.  The 2017 davis challenge on video object segmentation.  _arXiv:1704.00675_ , 2017.
  * Ranjan & Black (2017a) Anurag Ranjan and Michael J Black.  Optical flow estimation using a spatial pyramid network.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 4161–4170, 2017a.
  * Ranjan & Black (2017b) Anurag Ranjan and Michael J. Black.  Optical flow estimation using a spatial pyramid network.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , 2017b.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention_ , pp. 234–241, 2015.
  * Sun et al. (2014) Deqing Sun, Stefan Roth, and Michael J Black.  A quantitative analysis of current practices in optical flow estimation and the principles behind them.  _International Journal of Computer Vision_ , 106:115–137, 2014.
  * Tassano et al. (2020) Matias Tassano, Julie Delon, and Thomas Veit.  Fastdvdnet: Towards real-time deep video denoising without flow estimation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 1354–1363, 2020.
  * Tian et al. (2020) Yapeng Tian, Yulun Zhang, Yun Fu, and Chenliang Xu.  Tdan: Temporally-deformable alignment network for video super-resolution.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 3360–3369, 2020.
  * Vaksman et al. (2021) Gregory Vaksman, Michael Elad, and Peyman Milanfar.  Patch craft: Video denoising by deep modeling and patch matching.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 2157–2166, 2021.
  * Xia et al. (2022) Zhuofan Xia, Xuran Pan, Shiji Song, Li Erran Li, and Gao Huang.  Vision transformer with deformable attention.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 4794–4803, 2022.
  * Zach et al. (2007) Christopher Zach, Thomas Pock, and Horst Bischof.  A duality based approach for realtime tv-l1 optical flow.  In _Joint Pattern Recognition Symposium_ , pp. 214–223, 2007.

##  6 Additional Experimental Details

###  6.1 Dataset Summary

Our experiments use the DAVIS and Set8 datasets datasets (Pont-Tuset et al.,
2017; Tassano et al., 2020). The DAVIS train-val dataset consists of 90
sequences and the number of frames in a sequence varies from $40$ \- $104$
with resolution $480\times 854$. The DAVIS testing dataset consists of $30$
sequences. Set8 consists of four sequences from a GoPro camera and four
sequences from the DAVIS test set. For Set8, the number of frames in a
sequence varies from $35$ \- $85$ with resolution $540\times 960$.

###  6.2 Upgrading Space-Only Attention

Is Losing Spatial Resolution Worth Space-Time Features? In practice, space-
time attention might replace space-only attention. Is the loss of spatial
information worth the temporal information? Theoretically assessing the impact
of losing spatial information on model quality is difficult, so we instead
seek experimental evidence. Using the COLA-Net architecture for video
denoising, we find that space-time attention is significantly better than
global or non-local space-only attention.

Experimental Details. We finetune the COLA-Net denoising network provided by
the original authors for roughly 30 epochs with 400 randomly selected sample
sequences from the DAVIS training dataset (Pont-Tuset et al., 2017). Each
sequence consists of $5$ frames with resolution $128\times 128$. We train on
two Titan RTX GPUs with a batch size of $8$ using PyTorch Lightning (Falcon et
al., 2019). The optimizer is the Adam with a learning rate of $5\cdot 10^{-5}$
(Kingma & Ba, 2014) with a step learning rate with a step size of $5$ epochs
and decay rate $\gamma=0.1$. Our search step keeps $30$ neighbors. As a
control, we finetune the original model using the authors’ original
configuration and observe no meaningful change in denoising quality.

Space-Time Features are Better than Space-Only Features for Denoising. Table 3
shows a COLA-Net architecture with various types of attention is tested on the
DAVIS dataset with varying Gaussian noise intensities ($\sigma$) (Pont-Tuset
et al., 2017; Mou et al., 2021). COLA-Net originally uses cross-scale
attention, which is a global, strided search. We replace only the attention
module with attention indicated by the column, and finetune the network for
$1/2$ an epoch on the DAVIS training dataset with a learning rate of $1e-4$
using the Adam optimizer (Kingma & Ba, 2014). The results show space-time non-
local attention with optical flow outperforms the original network by over $1$
dB PSNR and it is the best of all the attention modules. The space-time search
and the space-time attention module use a spatial window of size $15\times
15$, a patch size of $7$, and a reference stride of $4$. The space-time search
uses a temporal window of size $2$. We aggregate the results using our
weighted sum of patches.

Space-Time Attention is Faster Than Global Attention. Table 3’s runtime and
memory consumption is measured using a $5$ frame video of resolution
$230\times 230$. This is the largest input video which fits on our 24 GB
NVIDIA RTX 3090 GPU. The global attention module is over 3.5 times slower than
our space-time search and consumes almost 20 times the memory. The space-only
and space-time searches use our Shifted Non-Local Search module with a
temporal window size of $0$ and $4$, respectively.

Table 3: Upgrading Search-Only Search. [PSNR$\uparrow$/SSIM$\uparrow$/ST-RRED$\downarrow$] This table reports denoising results on the DAVIS dataset using various attention modules within COLA-Net. The origianl attention module uses global, strided search named cross-scale attention. We report results using the noise intensities ($\sigma^{2}$) to match the original paper.  Dataset | $\sigma$ | Original (Global) | Non-Local | Space-Time
---|---|---|---|---
DAVIS | 15 | 34.26/0.915/1.3 | 34.14/0.912/1.4 | 35.78/0.936/0.8
30 | 30.93/0.848/5.2 | 30.80/0.843/5.3 | 32.41/0.878/2.6
50 | 28.67/0.784/13.6 | 28.51/0.776/13.5 | 29.73/0.810/7.2
Runtime (seconds) |  | 1.703 | 0.242 | 0.439
Memory (GB) |  | 14.449 | 0.730 | 0.736
Figure 11: Qualitatively Comparing Denoised Outputs. [PSNR$\uparrow$] RVRT and
STAN use space-time attention and are trained using the same procedure, so we
qualitatively compare the two networks. Small details are recovered using
STAN’s improved space-time attention module.

###  6.3 Upgrading Space-Time Attention: Additional Details

Training Details. The upgraded RVRT networks are trained for 180,000
iterations with a batch size of 8 using clip lengths of 10 with resolution
$256\times 256$. Each batch is corrupted with Gaussian noise using a random
noise parameter, $\sigma\sim\text{Uniform}[0,50]$. Weight updates use the Adam
optimizer with an initial learning rate of $4\times 10^{-4}$ and decrease with
the Cosine Annealing scheme (Kingma & Ba, 2014; Loshchilov & Hutter, 2016).
RVRT’s internal SpyNet model is initialized with pre-trained weights which are
fixed for the first 30,000 iterations and learns with a learning rate reduced
by 75%. (Ranjan & Black, 2017b; Niklaus, 2018). STAN uses TV-L1 optical flow,
instead of SpyNet (Zach et al., 2007). Notably, we execute only 180,000
iterations instead of 600,000 due to limited computing resources.

Upgrading Guided Deformable Attention. A Guided Deformable Attention (GDA)
module requires offsets, denoted ${\bm{F}}$, with shape $H\times W\times
9\times 2$ where ${\bm{F}}[h_{i},w_{i},k]$ the $(x,y)$ shift from
$(h_{i},w_{i})$ between a reference and adjacent frame. These offsets are
output from a network whose input includes a single optical-flow-like offset,
denoted $\widetilde{{\bm{F}}}$ with shape $H\times W\times 2$. In summary, for
reference and adjacent frames, denoted ${\bm{V}}_{\text{ref}}$ and
${\bm{V}}_{\text{adj}}$, a GDA module uses
${\bm{F}}=\text{net}({\bm{V}}_{\text{ref}},{\bm{V}}_{\text{adj}},\widetilde{{\bm{F}}})$.
Our search step replaces the network and outputs the top-$K$ ($K=9$) offsets;
${\bm{F}}=\text{search}({\bm{V}}_{\text{ref}},{\bm{V}}_{\text{adj}},\widetilde{{\bm{F}}})$.

Ablation Study. Our Shifted Non-Local Search uses hyperparameters not learned
during network training. To better understand their impact on video denoising
quality, we report results for various settings. The search stride, $S_{K}$,
denotes the spacing between two points in the grid search. When $S_{K}=0.5$ or
$S_{K}=1$, then each grid point is spaced $1/2$ or $1$ pixel apart,
respectively. Table 4 shows the results for a various number of search
parameters. The best settings used a search stride of $0.5$ and a spatial
window of $9$, outperforming the same window size with a search stride of $1$.
This suggests subpixel corrections to the predicted offsets may be more
beneficial than a larger search radius. A search radius that is too large or
too small also decreases the network quality. We hypothesis a large search
radius will overfit to noise and a small radius does not have enough radius to
properly correct the errors from the predicted offsets.

Table 4: Ablation Experiments for Shifted Non-Local Search.[PSNR$\uparrow$/SSIM$\uparrow$/ST-RRED$\downarrow$] A small spatial window ($W_{s}=3$) is unable to correct errors and a large spatial window ($W_{s}=15$) overfits to noise. A fractional search stride allows for subpixel correction. $\sigma$ | $W_{s}=9,S_{K}=0.5$ | $W_{s}=9,S_{K}=1$ | $W_{s}=15,S_{K}=1$ | $W_{s}=3,S_{K}=1$
---|---|---|---|---
10 | 38.90/0.967/0.004 | 38.77/0.967/0.004 | 38.68/0.966/0.004 | 38.61/0.965/0.004
20 | 35.58/0.936/0.012 | 35.40/0.934/0.013 | 35.27/0.933/0.013 | 35.23/0.932/0.013
30 | 33.68/0.907/0.024 | 33.44/0.904/0.025 | 33.29/0.901/0.026 | 33.27/0.900/0.026
40 | 32.35/0.881/0.040 | 32.07/0.875/0.042 | 31.89/0.872/0.043 | 31.87/0.871/0.044
50 | 31.30/0.880/0.027 | 30.97/0.847/0.061 | 30.78/0.843/0.063 | 30.76/0.841/0.064

###  6.4 Space-Time Attention Network (STAN) for Video Denoising

Figure 12: Space-Time Attention Network (STAN) Architecture. The STAN macro-
level architecture design is inspired from UNet (Ronneberger et al., 2015).
Each block’s design is inspired by the RVRT network (Liang et al., 2022b).

Architecture. The macro-level architecture is designed to match the multi-
scale architecture of UNet (Ronneberger et al., 2015). Experiments adding more
depth resulted in significantly worse quality models and downscaling is needed
to allow the increased number of parameters. Each block contains layer
normalization, our space-time attention module, channel attention, and lastly
Residual Swin Transformer Blocks (Hu et al., 2018; Liang et al., 2021). The
block structure is inspired by the RVRT network (Liang et al., 2022b). The
space-time attention layer is described in Section 3.1, and we use the
aggregation scheme described Section 3.5. Notably, no positional embeddings
are used in the space-time attention module of STAN. Our networks are set to
$N_{1}=1$, $N_{2}=2$, and $N_{3}=4$ with 1, 4, and 12 heads. STAN uses TV-L1
for optical flow (Zach et al., 2007).

Experimental Details. The RVRT and STAN networks are trained for 450,000
iterations with a batch size of 8 using clip lengths of 16 with resolution
$256\times 256$. Each batch is corrupted with Gaussian noise using a random
noise parameter, $\sigma\sim\text{Uniform}[0,50]$. Weight updates use the Adam
optimizer with an initial learning rate of $4\times 10^{-4}$ and decrease with
the Cosine Annealing scheme (Kingma & Ba, 2014; Loshchilov & Hutter, 2016).
The internal SpyNet model is initialized with pre-trained weights which are
fixed for the first 30,000 iterations and uses a learning rate reduced by 75%.
(Ranjan & Black, 2017b; Niklaus, 2018). Notably, we execute only 450,000
iterations instead of 600,000 due to limited computing resources. We only
notice a 0.18 dB increase across the final $120,000$ iterations, which does
not close the gap of over 1 dB difference remaining between reported RVRT
results and our reproduced RVRT results.

Valid Conclusions Despite the Non-Reproducible Related Work. While we cannot
reproduce the results from RVRT, we believe this does not change conclusions
regarding our Shifted Non-Local Search module. Since the training procedure is
copied directly from the RVRT paper and the procedure is identical to all RVRT
models and STAN, our conclusions remain scientifically supported.

Reproducibility Description. While the authors of RVRT do not provide training
code, a procedure is described in their paper. However, when we follow their
training procedure, the network’s quality is significantly worse than reported
by the original paper. An open issue on RVRT’s GitHub describes being unable
to reproduce their trained network’s results. Since training a single RVRT
network takes approximately 10 days which costs about $$560$ on
AWS222https://aws.amazon.com/ec2/instance-types/g4/, we do not believe it is
our role to debug the author’s method.

Ablation Experiment. We include an ablation study to assess the impact of
space-time attention’s temporal window size in Table 5. Each column’s
configuration is trained from scratch for $180,000$ epochs, following the
training procedure previously described. The table each network on the Set8
dataset.

Table 5: The Impact of Temporal Features.[PSNR$\uparrow$/SSIM$\uparrow$/ST-RRED$\downarrow$] Temporal features improve the denoising quality of the STAN architecture. The temporal window size is the number of frames searched for each query point from each frame. When $W_{t}=1$, no adjacent frames are searched. $\sigma$ | $W_{t}=5$ | $W_{t}=3$ | $W_{t}=1$
---|---|---|---
10 | 36.71/0.957/0.003 | 36.59/0.956/0.003 | 36.42/0.953/0.003
20 | 33.94/0.927/0.008 | 33.77/0.925/0.008 | 33.51/0.920/0.009
30 | 32.29/0.900/0.014 | 32.10/0.896/0.015 | 31.82/0.890/0.016
40 | 31.11/0.874/0.023 | 30.92/0.870/0.024 | 30.74/0.862/0.026
50 | 30.19/0.850/0.033 | 30.00/0.845/0.035 | 29.72/0.838/0.038

##  7 Pytorch Example of Space-Time Attention

We provide example PyTorch code for our space-time attention using the Shifted
Non-Local Search module named _NonLocalSearch_ within the code snippet.

⬇

# – imports –

import torch as th

conv2d = torch.nn.functional.conv2d

conv3d = torch.nn.functional.conv3d

from einops import rearrange

import stnls

\par# – data –

B = 1 # batch size

T = 5 # number of frames

F = 32 # number of features

H = 256 # video resolution height

W = 256 # ideo resolution width

device = ”cuda:0”

vid_in = th.randn((B,T,F,H,W),device=device)

fflow = th.randn((B,T,2,H,W),device=device)

bflow = th.randn((B,T,2,H,W),device=device)

\par# – queries, keys, and values –

proj_weights = th.randn((F,F,1,1),device=device)

q_vid = conv2d(vid,proj_weights)

k_vid = conv2d(vid,proj_weights)

v_vid = conv2d(vid,proj_weights)

\par# – searching –

ws = 9 # spatial window (W_s)

wt = 1 # temporal window; searches 2*wt+1 (aka W_t) frames

ps = 3 # patch size (P)

K = 5 # number of neighbors

HD = 2 # number of heads

stride0 = 1 # reference stride (S_Q)

stride1 = 1 # search stride (S_K)

search = stnls.NonLocalSearch(ws,wt,ps,K,nheads=HD,

stride0=stride0,stride1=stride1)

dists,inds = search(q_vid,k_vid,fflow,bflow)

# dists.shape = (B,HD,T*nH*nW), nH = ((H-1)//stride0+1)

# inds.shape = (B,HD,T*nH*nW,3)

\par# – aggregate –

ps = 5 # patch size; can change from search

stacking = stnls.NonLocalStack(ps,stride0)

stacked = stacking(v_vid,dists,inds)

# stacked.shape = (B,HD,K,T,F’,H,W) where F’ = F/HD

stacked = rearrange(stacked,’b hd k t f h w -> b t k (hd f) h w’)

agg_weights = th.randn((F,F,K,1,1),device=device)

vid_out = conv3d(stacked,agg_weights,stride=(K,1,1))

# vid_out.shape = (B,T,F,H,W)


[True anomaly $({}^{\circ})$]trueanomalyΘ Kazuaki Kashihara, Kuntal Kumar Pal,
Chitta Baral, Robert P Trevino 11institutetext: School of Computing and
Augmented Intelligence, Arizona State University, USA
11email<EMAIL_ADDRESS>22institutetext: Design Pickle, USA
22email<EMAIL_ADDRESS>
# Prompt-Based Learning for
Thread Structure Prediction
in Cybersecurity Forums

Kazuaki Kashihara 11 Kuntal Kumar Pal 11 Chitta Baral 11 Robert P Trevino 22

###### Abstract

With recent trends indicating cyber crimes increasing in both frequency and
cost, it is imperative to develop new methods that leverage data-rich hacker
forums to assist in combating ever evolving cyber threats. Defining
interactions within these forums is critical as it facilitates identifying
highly skilled users, which can improve prediction of novel threats and future
cyber attacks. We propose a method called Next Paragraph Prediction with
Instructional Prompting (NPP-IP) to predict thread structures while grounded
on the context around posts. This is the first time to apply an instructional
prompting approach to the cybersecurity domain. We evaluate our NPP-IP with
the Reddit dataset and Hacker Forums dataset that has posts and thread
structures of real hacker forums’ threads, and compare our method’s
performance with existing methods. The experimental evaluation shows that our
proposed method can predict the thread structure significantly better than
existing methods allowing for better social network prediction based on forum
interactions.

###### keywords:

Instructional prompts, Thread structure prediction, Thread Structure, Social
network, Unstructured forums, Cybersecurity

##  1 Introduction

Cybercrimes cost trillions of dollars in damages worldwide each year,
impacting different sectors of society ranging from national defense to
private industry [6]. Current trends indicate a considerable rise in
cybercrimes in the next several years as hacker tools become more
sophisticated and ubiquitous [18]. This is, in part, due to the advent of the
dark web, which has given hackers the opportunity to interact, profit, and
exchange information on dark web forums [8]. Identifying key user interactions
within these dark forums can assist in identifying prominent hackers with
knowledge of novel threats as well as predicting potential cyber attacks.
Thus, the thread structure of a forum becomes important in generating social
networks based on user interactions as shown in Fig. 1 [7].

Unfortunately, most of the hacker forums are unstructured making it difficult
to identify user interactions through post replies in an automated manner.
Moreover, although many of these dark forums have a rich source of text
information in different threads that discuss specific topics, such as
malware, virus, illegal items, and other illegal activities, the recent
reports indicate that $90\%$ of posts on popular dark web hacking forums are
made by those looking to solicit hacker services instead of the hackers
themselves [27, 2]. Traditional methods used to define social networks on
unstructured forums such as Creator-oriented Network and Last Reply-oriented
Network [13] are based upon temporal interaction assumptions that do not
consider the full context of the user interactions based on the content of the
posts. Kashihara et al. previously introduced a powerful deep learning method
called Next Paragraph prediction (NPP) designed to define social networks
using posts from the Reddit forums [12]. The NPP method outperformed
traditional methods as well as BERT’s Next Sentence Prediction (NSP) [3] when
defining social networks from posts.

Figure 1: Example of a thread from an unstructured Hacker forum in the darkweb
(left), and the predicted thread structure (right).

Building upon the NPP method, we propose the Next Paragraph Prediction with
Instructional Prompting (NPP-IP) that leverages cutting-edge Prompt-based
learning to assist in social network construction from posts. Using the Reddit
dataset [12] consisting of over $105$ threads and $1,648$ posts, we train and
evaluate the model against both traditional methods as well as the original
NPP method. In addition, we test the model using real unstructured hacker
forums data, where $20$ threads are manually annotated by human experts to
identify interactions based on posts. Since the hacker forums data is hard to
extract and manual annotation takes time and labor tense, we create the above
limited dataset for the hacker forums in this work. The results show that NPP-
IP performs $2.68$ to $4.70$ percentage points better than the other existing
methods.

Contributions: Our research contributions in this research are listed as
follows:

  * •

Prompt-based learning (instructional prompting) is introduced into a deep
learning method called Next Paragraph Prediction for social network
construction of forum data.

  * •

We apply Prompt-based learning to Cybersecurity domain for the first time.

  * •

The evaluation results show that our method can predict thread structures
better than the existing methods.

  * •

The results indicate that the proposed method is robust enough to train using
data from one cyber-related forum and apply to another cyber-related forum.

##  2 Related Work

In this section, we give a brief overview of previous work on reply
relationships identification, thread structure prediction, and instructional
prompting.

###  2.1 Reply Relationships Identification

In general, reply relationships are explicitly defined in several social
platforms such as Reddit, Twitter, and Facebook. However, unstructured forums
and some of the online communication tools such as Telegram and other instant
messaging platforms have a challenge to identify the reply relationships since
users usually do not use explicit reply marks but directly post related
messages or posts to communicate. Thus, several methods are proposed to
identify the reply relationships in various platforms.

The first step of the conversation disentanglement task is usually reply
relationship identification [4, 15, 10]. There are several classification
models that have been developed to detect reply relationships between pairs of
messages. A linear binary classifier considering conversation features and
content features is proposed by Elsner et al. [4]. A random forest classifier
with 250 trees that uses a feature vector in order to describe the
relationship between the two messages is developed by Mehri et al. [15].

Neural models are used in recent methods to represent utterances. The joint
model of reply relationship and pairwise relationship under pointer network
model is proposed by Yu et al. [29]. Another approach is that the problem of
conversation structure modeling is defined as identifying the parent
utterances to which each utterance in the conversation responds to by Zhu et
al. [31]. They designed a novel masking mechanism using masked hierarchical
transformer. Zhang et al. [30] proposed a method based on multi-features to
identify reply relationships from Telegram groups using BERT model to learn
the textual representation of messages and utilize the user’s contextual
features that contain richer information to overcome the limitation of short
messages.

###  2.2 Thread Structure Prediction

In order to build social networks from forums, member interactions must be
correctly identified via posts on threads. There are two network
representations introduced [13] for building the social network in forums:
Creator-oriented Network and Last Reply-oriented Network. The Last Reply-
oriented Network is widely used for the social network analysis in the recent
works [21, 1, 14, 23, 20, 11]. Fig. 2 shows the sample structure of Creator-
oriented Network and Last Reply-oriented Network. Since these two traditional
network conversion approaches are based on limited information and
considerable assumptions on interactions between users, the social structures
of the networks are likely not accurate representations. Other recent works
have predicted helpful posts in the forums [9] using a neural network based
model that determines whether the post is useful or not. However, the
importance of a post has very little utility when predicting interactions and
thus social networks. More recently, Kashihara et al. [12] proposed the Next
Paragraph Prediction (NPP) method which extended BERT’s Next Sentence
Prediction to predict the response post from the previous post. This method
allows for the Reconstruction of social networks using thread structure
prediction.

Figure 2: Example of Creator-oriented Network and Last Reply-oriented Network.

###  2.3 Instructional Prompting

Building effective discrete prompts for language models (LM) to perform NLP
tasks is an active area of research [25, 24, 26, 22]. Such prompts are often
extremely short and may not include a complete definition of complex tasks. In
contrast, the recent works [16, 17] give instructions encode detailed
instructions as they were used to collect the dataset. Driven many empirical
analysis by [16], the framing instructional prompting has demonstrated
considerable improvements across LMs.

##  3 Model Description

Figure 3: The original NPP model (Left) combines a pair of posts to predict
whether one post is a response to the other. Our NPP-IP model (Right)
incorporates instruction prompt information into the NPP structure allowing
for task information to be leveraged.

Our proposed NPP-IP model is based on infusing the original dataset with
specific task instructions using an instruction prompting function. Formally,
the instruction prompting function $f_{prompt}(\cdot)$ is defined as

$\displaystyle f_{prompt}(x)=I||x,$ (1)

where —— represents concatenation of instruction prompt $I$ with training
sample $x$. Instruction prompt $I$ is formally defined as:

Task Description:
You are given two posts and you need to generate True if they are the direct
reply relation, otherwise generate False.
Positive Example:
post1: Windows Defender Gets a New Name: Microsoft Defender
post2: Bring back MSE and its ui even logo looks cool…
output: True
Negative Example:
post1: Windows Defender Gets a New Name: Microsoft Defender
post2: Title says it
output: False”

Training sample $x$ is formally defined as

$\displaystyle x=\text{Post $k$ || [sep] || Post $k+i$},$ (2)

which represents a pair of concatenated posts at index $k$ and $k+i$ with a
separation key $[sep]$, such that $i\neq 0$.

The NPP-IP model leverages five framing techniques defined in [16] for framing
the instruction prompting information $I$. (i) First, the Use Low Level
Patterns technique is accomplished by providing a simple task descriptor to
correctly output a value of True or False if a reply relationship exists
between posts without including any cybersecurity jargon. (ii and iii) Second,
Itemized Instructions are provided via the positive and negative examples with
the corresponding output in bulleted list format for thread structure
prediction. The positive and negative examples also fulfill the Break It Down
technique by defining simpler sub-tasks corresponding to identifying negative
and positive examples. This is also where cybersecurity information is
introduced into the instructional prompt. (iv) Next, Enforce Constraints is
accomplished by constraining the examples to their respective outputs of True
or False. (v) Lastly, the Specialize Instructions technique is accomplished by
specifically stating the expected output in both task description and
examples.

Figure 4: The data structure for NPP model and NPP-IP model.

Fig. 3 shows the BERT-based neural network structure used by an NPP model as
well as the resultant NPP-IP model after introducing instructional prompting
information. The original dataset gives two posts as its input, where the
label space is defined as {True, False}, defining whether posts share a direct
response relation or not. Including instructional prompting provides critical
task information for both positive and negative cases, which are then used in
the embedding and subsequent prediction task during training. Fig. 4 shows the
data structure of the model for original NPP model and NPP-IP model.

##  4 Evaluation and Results

###  4.1 Datasets

A curated Reddit dataset [12] was used to train and evaluate our proposed
model. The Reddit dataset is ideally suited for thread structure analysis
given its tree-like structure within different threads. The Reddit dataset has
the threads from the following ten topics from “cybersecurity” field and
extracted the threads of these topics: “cyber_security”, “AskNetsec”,
“ComputerSecurity”, “cyberpunk”, “cybersecurity”, “Hacking”, “Malware”,
“Malwarebytes”, and “security”. Our proposed model was also evaluated using 20
hacker forum threads from three English hacker forums annotated by human
experts, which is referred to as the “Hacker Forums” dataset. The forum thread
data is from CYR3CON (Cyber Security Works)111https://www.cyr3con.ai. The
average posts per thread is $15.4$. Four cybersecurity experts checked posts
in each thread, and annotated a relation of two posts in a thread which the
two posts are direct response relations or not. The site names and usernames
are anonymized. Table 1 and Table 2 show the basic statistics of Reddit and
Hacker Forums dataset respectively. In the Reddit dataset, we create all pairs
of post combinations in a thread. The positive pair of posts is that the
second post is the direct reply of the first post. There are 14744 pairs of
posts including both positive and negative. Then, we split the pairs into
train, dev, and test sets and the split ratio is $60\%$, $10\%$, and $30\%$
respectively.

Topic Name | TH | Posts
---|---|---
cyber_security | 8 | 48
AskNetsec | 14 | 338
ComputerSecurity | 12 | 110
cyberpunk | 11 | 176
cybersecurity | 11 | 158
Hacking | 12 | 370
Hacking_Tutorial | 12 | 110
Malware | 9 | 82
Malwarebytes | 8 | 72
security | 8 | 184
Table 1: The Reddit dataset consisted of ten cybersecurity topics. “TH” is defined as the number of threads in each topic while “Posts” is defined as the number of Posts across the different threads. Forum # | TH | Posts
---|---|---
Forum1 | 7 | 169
Forum2 | 7 | 80
Forum3 | 6 | 58
Table 2: The Hacker Forums dataset consisted of 20 threads from three hacker forums. “TH” is defined as the number of threads in each topic while “Posts” is defined as the number of Posts across the different threads. Method | Model | P | R | F1
---|---|---|---|---
CO | - | 0.00 | 1.00 | 0.01
LR | - | 0.72 | 0.12 | 0.20
NPP | BE-B | 0.42 | 0.46 | 0.44
BE-L | 0.36 | 0.51 | 0.42
RB-B | 0.59 | 0.33 | 0.43
RB-L | 0.41 | 0.58 | 0.48
NPP-IP | BE-B | 0.48 | 0.46 | 0.47
BE-L | 0.64 | 0.41 | 0.50
RB-B | 0.62 | 0.43 | 0.51
RB-L | 0.39 | 0.56 | 0.46
Table 3: Results from the Reddit test data show that the NPP-IP method outperformed all other methods for thread structure prediction across all but one of the different BERT language models analyzed. |  | Forum1 | Forum2 | Forum3
---|---|---|---|---
Method | Model | P | R | F1 | P | R | F1 | P | R | F1
CO | - | 0.31 | 1.00 | 0.47 | 0.27 | 1.00 | 0.43 | 0.12 | 1.00 | 0.21
LR | - | 0.50 | 0.00 | 0.01 | 0.50 | 0.09 | 0.16 | 0.50 | 0.13 | 0.21
NPP | BE-B | 0.40 | 0.37 | 0.39 | 0.37 | 0.61 | 0.46 | 0.33 | 0.65 | 0.44
BE-L | 0.94 | 0.27 | 0.41 | 0.50 | 0.34 | 0.41 | 0.50 | 0.33 | 0.40
RB-B | 0.59 | 0.38 | 0.46 | 0.29 | 0.50 | 0.37 | 0.48 | 0.43 | 0.41
RB-L | 0.54 | 0.55 | 0.54 | 0.55 | 0.58 | 0.54 | 0.45 | 0.40 | 0.41
NPP-IP | BE-B | 0.55 | 0.39 | 0.45 | 0.71 | 0.63 | 0.67 | 0.61 | 0.56 | 0.58
BE-L | 0.70 | 0.43 | 0.53 | 0.85 | 0.31 | 0.45 | 0.61 | 0.60 | 0.57
RB-B | 0.50 | 0.37 | 0.42 | 0.52 | 0.58 | 0.48 | 0.53 | 0.84 | 0.46
RB-L | 0.50 | 0.87 | 0.43 | 0.50 | 0.34 | 0.41 | 0.50 | 0.33 | 0.40
Table 4: Results from each of the anonymous hacker forums demonstrated that
the NPP-IP model outperformed all other models. The NPP and NPP-IP models were
both trained with Reddit data further demonstrating NPP-IPs inference
performance robustness on unrelated cyber forums.

###  4.2 Metrics and Task

Our proposed NPP-IP method was evaluated against several different methods for
thread structure prediction using cybersecurity related posts. Two language
models, BERT (BE) and RoBERTa (RB), were explored when training the NPP and
proposed NPP-IP models, where -B and -L represent base and large models for
each LM respectively. As shown in Fig. 5, our NPP-IP method outperformed the
original NPP method based on the F1 score using the Reddit data across all but
one of the LMs. We compared performance with well known methods, Creator-
Oriented Network (CO) and Last Reply-Oriented Network (LR), using Precision
(P), Recall (R), and F1 score (F1) metrics.

###  4.3 Libraries and Hyperparameters

In order to build, train, and evaluate both NPP and NPP-IP methods, we use
publicly available libraries and set hyperparameters, and they are described
as follows: torchtext 0.8.0 and PyTorch 1.7.1 [19], pytorch-lightning 1.2.2
[5], and transformers 3.4 [28] on Google Colab (Nvidia K80 12 GB GPU) were
used. We use the hidden dropout probability as 0.15. The batch size was set to
8 and the learning rate was set to 5e-6. The model was trained with $>10$
epochs. Convergence was observed around 3 epochs with limited over-fitting and
the maximum sequence length was set to 250.

Figure 5: An F1 score comparison of the Reddit Forum data using different
BERT-based language models indicates that our proposed NPP-IP model (Red)
outperforms NPP (Blue) in all but one language model, achieving the highest
score using RoBERTa LM.

###  4.4 Results

The results are reported in Tables 3 and 4 for Reddit and Hacker Forums
datasets, respectively. Both tables show a clear improvement from our NPP-IP
method compared to most methods across both datasets. There were two cases
where NPP outperformed NPP-IP with both using the large RoBERTa (RB-L)
language model. As shown in Table 4, in two of the three hacker forums, our
proposed NPP-IP method with BERT-B LM reached the highest F1 score while NPP-
IP with BERT-L LM recorded the highest F1 score for the third forum.

##  5 Analysis and Discussion

In this section, we analysis the results of Reddit and Hacker Forums datasets,
especially NPP-IP performance, Precision and Recall Implications, and various
Error cases.

###  5.1 NPP-IP Performance

As far as the authors are aware, this is the first time that instructional
prompts have been applied to text-based cybersecurity data. As the results of
Reddit and Hacker Forums datasets show in Table 3 and Table 4, NPP-IP performs
better than NPP in most of the cases. The improvement on the Reddit dataset
from NPP to NPP-IP across different LMs, ranged between $3-8\%$ across F1
scores. Similar improvements were observed using the real world Hackers Forums
dataset, ranging between $3-11\%$ difference across F1 scores. This is despite
the fact that the models were trained on Reddit only data. As shown, in Table
1, the Reddit dataset is comprised of cybersecurity related topics across the
different threads. This evidence is consistent with NPP-IP’s ability to better
detect and leverage cybersecurity related information compared to other well
known methods for social network construction based on thread structure
prediction. Moreover, the framing of the instructional prompt using cyber
related information may also have improved its performance across different
forums. This is a significant discovery since annotating new datasets,
especially in the cyber realm, is costly, requiring considerable human
experts’ efforts to collect a decent size of data for training and testing.
More research needs to be conducted to determine the extent to which framing
cyber related instruction prompts can make text-based analysis more robust
across different cyber forums and datasets.

###  5.2 Precision and Recall Implications

As precision and recall scores in Tables 3 and 4 show, both are considerably
low in Reddit dataset, with recall scores much lower than precision scores in
the Hacker Forums dataset. One possible explanation for this observed behavior
is that publicly available pre-trained LMs were used. These LMs are pre-
trained by a wide range of topics across a massive size of data. However, the
cybersecurity field is in a constant state of flux - changing the meaning of
words and adding new words quite frequently. We suspect that LMs could not
understand many of the cybersecurity keywords in posts predicting thread
structures that were consistent with actual social interaction. Thus, re-
training LMs with cybersecurity data should be explored to improve the
performance.

###  5.3 Error Cases

There are several error cases that are not easy to solve. Fig. 6 shows an
interesting case we found in the Reddit dataset. For the question post “Can
you send a link to the tutorial?”, a user responded “*URL*” and “in the
comment”. In the ground truth, “in the comment” is the response of “*URL*”,
however, both of our models predicted “in the comment” is the response of “Can
you send a link to the tutorial?”. We think that “in the comment” reinforces
the post “*URL*” and also answers the original question. Since the ground
truth is based on the thread tree structure, it only has one interaction even
if it can interact with multiple posts or users. However, due to the tree
structure in Reddit, the ground truth from the subreddit structure is assigned
to only one of them. We found some cases that our methods predicted a post
replied to multiple posts, and only one of them is correct as we mentioned
before. Thus, these cases may decrease the performance of our methods.

Figure 6: An interesting case in Reddit dataset. Since the actual URL is
harmful site URL, we replaced it as *URL*.

In the Hacker Forums dataset, two hacker forums have a feature to quote the
referencing post in the same thread. However, the quote feature catches not
only the referencing post but also the referencing post’s referencing post if
it has. Fig. 7 shows the example of this case. User C replied (referenced)
User B post, however, the User B referenced User A post as well. Then, both
NPP and NPP-IP models predicted the pair of User C post and User A post are
direct reply relationship since User C post contains User A post content
through referencing User B post. We observed many this false positive cases,
and this type of error effected some performance in the Hacker Forums dataset.
Some pre-process to remove reference’s reference post content will be needed
to solve this issue.

Figure 7: An interesting case in Hacker Forums dataset.

##  6 Conclusion

Predicting thread structures within cybersecurity forums is a crucial
component in defining key social networks used to identify prominent users who
provide useful information. Identifying these users can facilitate prediction
and prevention of future cyber incidents and attacks.

A prompt-based learning model called Next Paragraph Prediction with
Instructional Prompting (NPP-IP) for predicting thread structures across
different cybersecruity topics was introduced. The method was evaluated using
two different datasets and compared against several well known methods. The
results show that the NPP-IP method had considerable improvement over existing
methods, achieving the highest F1 score across different real world hacker
forum datasets.

In the future work, we plan to train the language models with cybersecurity
related data for adapting cybersecurity context to the models to improve the
performance. In addition, we will apply these methods to the social network
analysis to replace from the assumption methods’ networks to our context
considered thread structures for comparing the advantage of the improved
quality of the networks.

## Ethical Considerations

In this research, we use one dataset from the other work with the agreement of
using the dataset for this research only. We created the Hacker Forum dataset
where the data from CYR3CON (Cyber Security Works), and they already
anonymized the site names and usernames. We have an agreement with CYR3CON to
use the data for this research only, and not sharing the data in public. The
Hacker Forums dataset has randomly picked 20 threads that have average 15.4
posts per thread from three English hacker forums. The dataset was annotated
by four cybersecurity experts (employees of CYR3CON) in a week as a part of
their jobs. Our goal is to construct thread structure from unstructured
forums. Further precautions taken include not identifying individuals
(including not publishing usernames), and presenting results objectively. In
addition, we use well-known publicly released language models, BERT and
RoBERTa, for our experiments.

## Limitations

In this study, we have presented a new approach of predicting thread structure
in cybersecurity forums. However, there are several limitations. First, there
is a challenge of accessibility of datasets, especially (public) sensitive.
Many of the previous works used hacker forums as their datasets, however, most
of them are not published due to copyright and other restrictions. Kashihara
et al. [12] provided the Reddit dataset for this study only. In addition,
CYR3CON provided us with raw hacker forums data for this research, and this
data will not be published. Second, there is a language limitation. To compare
our approach with the related work [12] on the same dataset, we need to stick
with English since they used Reddit threads in English for their dataset.
There are many hacker forums in not only English but also other languages,
however, we selected 20 hacker forum threads from three English hacker forums
for our evaluation since we fine-tuned English BERT and RoBERTa models for our
evaluation. Lastly, another limitation is the size of ground truth. The
original Reddit dataset by [12] used the Reddit tree structure that shows the
reply relationship of each pair of posts that is the direct response as the
ground truth. However, many hacker forums we used for evaluation do not have
such a structure, and human experts needed to check every combination of post
pairs under a thread to determine the relationship of each pair. Thus, we use
the 20 threads from three English hacker forums annotated by human experts for
the ground truth.

## References

  * [1] Almukaynizi, M., Grimm, A., Nunes, E., Shakarian, J., Shakarian, P.: Predicting cyber threats through hacker social networks in darkweb and deepweb forums. In: Proceedings of the 2017 International Conference of The Computational Social Science Society of the Americas. p. 12. ACM (2017)
  * [2] Culafi, A.: Ninety percent of dark web hacking forum posts come from buyers. shorturl.at/fEGW4 (2021), accessed: 2022-04-01
  * [3] Devlin, J., Chang, M., Lee, K., Toutanova, K.: BERT: pre-training of deep bidirectional transformers for language understanding. In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers). pp. 4171–4186 (2019), https://aclweb.org/anthology/papers/N/N19/N19-1423/
  * [4] Elsner, M., Charniak, E.: You talking to me? A corpus and algorithm for conversation disentanglement. In: McKeown, K.R., Moore, J.D., Teufel, S., Allan, J., Furui, S. (eds.) ACL 2008, Proceedings of the 46th Annual Meeting of the Association for Computational Linguistics, June 15-20, 2008, Columbus, Ohio, USA. pp. 834–842. The Association for Computer Linguistics (2008), https://aclanthology.org/P08-1095/
  * [5] Falcon, W., Cho, K.: A framework for contrastive self-supervised learning and designing a new approach. arXiv preprint arXiv:2009.00104 (2020)
  * [6] Fox, J.: Cybersecurity statistics 2021. https://www.cobalt.io/blog/cybersecurity-statistics-2021 (2021), accessed: 2022-04-01
  * [7] Fu, T., Abbasi, A., Chen, H.: Interaction coherence analysis for dark web forums. In: IEEE International Conference on Intelligence and Security Informatics, ISI 2007, New Brunswick, New Jersey, USA, May 23-24, 2007, Proceedings. pp. 342–349 (2007), https://doi.org/10.1109/ISI.2007.379495
  * [8] Goel, S.: Cyberwarfare: connecting the dots in cyber intelligence. Communications of the ACM 54(8), 132–140 (2011)
  * [9] Halder, K., Kan, M., Sugiyama, K.: Predicting helpful posts in open-ended discussion forums: A neural architecture. In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers). pp. 3148–3157 (2019), https://www.aclweb.org/anthology/N19-1318/
  * [10] Jiang, J., Chen, F., Chen, Y., Wang, W.: Learning to disentangle interleaved conversational threads with a siamese hierarchical network and similarity ranking. In: Walker, M.A., Ji, H., Stent, A. (eds.) Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2018, New Orleans, Louisiana, USA, June 1-6, 2018, Volume 1 (Long Papers). pp. 1812–1822. Association for Computational Linguistics (2018)
  * [11] Johnsen, J.W., Franke, K.: Identifying proficient cybercriminals through text and network analysis. In: IEEE International Conference on Intelligence and Security Informatics, ISI 2020, Arlington, VA, USA, November 9-10, 2020. pp. 1–7. IEEE (2020), https://doi.org/10.1109/ISI49825.2020.9280523
  * [12] Kashihara, K., Shakarian, J., Baral, C.: Social structure construction from the forums using interaction coherence. In: Proceedings of the Future Technologies Conference. pp. 830–843. Springer (2020)
  * [13] L’Huillier, G., Álvarez, H., Ríos, S.A., Aguilera, F.: Topic-based social network analysis for virtual communities of interests in the dark web. SIGKDD Explorations 12(2), 66–73 (2010), https://doi.org/10.1145/1964897.1964917
  * [14] Marin, E., Shakarian, J., Shakarian, P.: Mining key-hackers on darkweb forums. In: 1st International Conference on Data Intelligence and Security, ICDIS 2018, South Padre Island, TX, USA, April 8-10, 2018. pp. 73–80 (2018), https://doi.org/10.1109/ICDIS.2018.00018
  * [15] Mehri, S., Carenini, G.: Chat disentanglement: Identifying semantic reply relationships with random forests and recurrent neural networks. In: Kondrak, G., Watanabe, T. (eds.) Proceedings of the Eighth International Joint Conference on Natural Language Processing, IJCNLP 2017, Taipei, Taiwan, November 27 - December 1, 2017 - Volume 1: Long Papers. pp. 615–623. Asian Federation of Natural Language Processing (2017), https://aclanthology.org/I17-1062/
  * [16] Mishra, S., Khashabi, D., Baral, C., Choi, Y., Hajishirzi, H.: Reframing instructional prompts to gptk’s language. CoRR abs/2109.07830 (2021), https://arxiv.org/abs/2109.07830
  * [17] Mishra, S., Khashabi, D., Baral, C., Hajishirzi, H.: Natural instructions: Benchmarking generalization to new tasks from natural language instructions. CoRR abs/2104.08773 (2021), https://arxiv.org/abs/2104.08773
  * [18] Morgan, S.: Hackerpocalypse cybercrime report 2016. https://cybersecurityventures.com/hackerpocalypse-cybercrime-report-2016/ (2016), accessed: 2022-04-01
  * [19] Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., Lerer, A.: Automatic differentiation in pytorch. In: NIPS-W (2017)
  * [20] Pete, I., Hughes, J., Chua, Y.T., Bada, M.: A social network analysis and comparison of six dark web forums. In: IEEE European Symposium on Security and Privacy Workshops, EuroS&P Workshops 2020, Genoa, Italy, September 7-11, 2020. pp. 484–493. IEEE (2020), https://doi.org/10.1109/EuroSPW51379.2020.00071
  * [21] Phillips, E., Nurse, J.R., Goldsmith, M., Creese, S.: Extracting social structure from darkweb forums (2015)
  * [22] Reynolds, L., McDonell, K.: Prompt programming for large language models: Beyond the few-shot paradigm. In: Kitamura, Y., Quigley, A., Isbister, K., Igarashi, T. (eds.) CHI ’21: CHI Conference on Human Factors in Computing Systems, Virtual Event / Yokohama Japan, May 8-13, 2021, Extended Abstracts. pp. 314:1–314:7. ACM (2021), https://doi.org/10.1145/3411763.3451760
  * [23] Sarkar, S., Almukaynizi, M., Shakarian, J., Shakarian, P.: Predicting enterprise cyber incidents using social network analysis on the darkweb hacker forums. In: 2018 International Conference on Cyber Conflict, CyCon U.S. 2018, Washington, DC, USA, November 14-15, 2018. pp. 1–7 (2018), https://cyberdefensereview.army.mil/Portals/6/Documents/CyConUS18“%20Conference“%20Papers/Session2-Paper2.pdf
  * [24] Scao, T.L., Rush, A.M.: How many data points is a prompt worth? In: Toutanova, K., Rumshisky, A., Zettlemoyer, L., Hakkani-Tür, D., Beltagy, I., Bethard, S., Cotterell, R., Chakraborty, T., Zhou, Y. (eds.) Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021. pp. 2627–2636. Association for Computational Linguistics (2021), https://doi.org/10.18653/v1/2021.naacl-main.208
  * [25] Schick, T., Schütze, H.: Few-shot text generation with natural language instructions. In: Moens, M., Huang, X., Specia, L., Yih, S.W. (eds.) Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021. pp. 390–402. Association for Computational Linguistics (2021), https://doi.org/10.18653/v1/2021.emnlp-main.32
  * [26] Tam, D., Menon, R.R., Bansal, M., Srivastava, S., Raffel, C.: Improving and simplifying pattern exploiting training. In: Moens, M., Huang, X., Specia, L., Yih, S.W. (eds.) Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021. pp. 4980–4991. Association for Computational Linguistics (2021), https://doi.org/10.18653/v1/2021.emnlp-main.407
  * [27] Technologies, P.: Custom hacking services. https://www.ptsecurity.com/ww-en/analytics/custom-hacking-services/ (2021), accessed: 2022-04-01
  * [28] Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T.L., Gugger, S., Drame, M., Lhoest, Q., Rush, A.M.: Transformers: State-of-the-art natural language processing. In: Liu, Q., Schlangen, D. (eds.) Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, EMNLP 2020 - Demos, Online, November 16-20, 2020. pp. 38–45. Association for Computational Linguistics (2020), https://doi.org/10.18653/v1/2020.emnlp-demos.6
  * [29] Yu, T., Joty, S.R.: Online conversation disentanglement with pointer networks. In: Webber, B., Cohn, T., He, Y., Liu, Y. (eds.) Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020. pp. 6321–6330. Association for Computational Linguistics (2020), https://doi.org/10.18653/v1/2020.emnlp-main.512
  * [30] Zhang, P., Qi, Y., Li, Y., Ya, J., Wang, X., Liu, T., Shi, J.: Identifying reply relationships from telegram groups using multi-features fusion. In: 2021 IEEE Sixth International Conference on Data Science in Cyberspace (DSC). pp. 321–327. IEEE (2021)
  * [31] Zhu, H., Nan, F., Wang, Z., Nallapati, R., Xiang, B.: Who did they respond to? conversation structure modeling using masked hierarchical transformer. In: The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020\. pp. 9741–9748. AAAI Press (2020), https://ojs.aaai.org/index.php/AAAI/article/view/6524