A Context-Aware Approach for Textual Adversarial Attack through Probability Difference Guided Beam Search

Abstract: Textual adversarial attacks expose the vulnerabilities of text classifiers and can be used to improve their robustness. Existing context-aware methods solely consider the gold label probability and use the greedy search when searching an attack path, often limiting the attack efficiency. To tackle these issues, we propose PDBS, a context-aware textual adversarial attack model using Probability Difference guided Beam Search. The probability difference is an overall consideration of all class label probabilities, and PDBS uses it to guide the selection of attack paths. In addition, PDBS uses the beam search to find a successful attack path, thus avoiding suffering from limited search space. Extensive experiments and human evaluation demonstrate that PDBS outperforms previous best models in a series of evaluation metrics, especially bringing up to a +19.5% attack success rate. Ablation studies and qualitative analyses further confirm the efficiency of PDBS.

\section{Introduction}

In this paper, we explore Pre-trained Language Model (PLM) based textual adversarial attack, and focus on the score-based and non-targeted black-box setting. Previous work \cite{papernot2016, TextFooler} demonstrate that applying tiny perturbations to texts may fool a text classifier (a.k.a. target model) while these perturbations are usually human imperceptible, raising concerns on safety in reality. However, comprehensive research on textual adversarial attack helps to improve the robustness of target models \cite{Wallace2019}, thus attracting much research attention.

The most recent work introduces various PLMs, e.g., BERT \cite{BERT} and RoBERTa \cite{RoBERTa}, into the textual adversarial attack, since PLMs enable us to generate more contextual and fluent adversarial examples. Specifically, to generate adversarial examples, BERT-Attack \cite{BERTAttack} first finds the vulnerable tokens of the input text sequence through a Replace action and then uses BERT to generate substitutes for these tokens.  BAE \cite{BAE} uses Replace and Insert operations to iteratively mask vulnerable tokens of the input and then uses BERT to generate substitutes for the masked tokens. CLARE \cite{CLARE} uses Replace, Insert and Merge actions to mask tokens of the input text sequence and infills the masked tokens using substitutes generated by RoBERTa.
The score-based black-box setting assumes class label probabilities of a classifier are accessible. For example, the three label probabilities for input sequence $\textrm{x}^{(0)}$ are 0.16, 0.64, and 0.20 (see Figure \ref{model2}).
Based on these scores, the above models formulate the task as a path search problem, and they propose to use the greedy search by selecting the minimum gold label probability in each attack iteration, such as the $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,1)} \rightarrow \textrm{x}^{(2,1)}$.
However, we find that in multi-class ($\geq$ 3) classifications, solely considering the gold probability may be insufficient. For example, the $\textrm{x}^{(1,2)}$ is intuitively easier to attack successfully than the $\textrm{x}^{(1,1)}$, since the probability difference of the former (0.09) is less than the latter (0.22).
But the greedy search selects $\textrm{x}^{(1,1)}$ when only considering the gold label probability.

In addition, we demonstrate that the greedy search limits the search space, leading to attack failure in some cases. For example, in Figure \ref{model2} the greedy search fails anyway, whether using the gold class probability or the probability difference. However,  there is a successful attack path, namely $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,3)} \rightarrow \textrm{x}^{(2,7)}$.

\begin{figure*}[t]
\centering
\includegraphics[width=0.85\textwidth]{fig/fig11}
\caption{A textual attack example regarding the triple-class text classification, where $\{\textrm{y}_1, \textrm{y}_2$, $\textrm{y}_3\}$ are class labels. For clarity, we assume there are only three perturbations applied to one text sequence. $\textrm{x}^{(0)}$ is the input text sequence with gold label $\textrm{y}_2$ and Iter. 1 \& 2 are two attack iterations. The values in rectangles are probabilities of related class labels, which are predicted by the target model.  $\textrm{x}^{(i,j)}$ is a candidate example, where $i$ is the iteration number and $j$ denotes the $j$-$th$ candidate. The red value around $\textrm{x}^{(i,j)}$ is the \textbf{probability difference}, which is defined as the gold label ($\textrm{y}_2$) probability minus the maximum probability of the other labels (\{$\textrm{y}_1,\textrm{y}_3$\}). The blue path $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,1)} \rightarrow \textrm{x}^{(2,1)}$ is the greedy search path when selecting the minimum gold label probability in each iteration. The green path $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,2)} \rightarrow \textrm{x}^{(2,4)}$ is the greedy search path when selecting the minimum probability difference in each iteration.
When using the probability difference guided beam search (set beam size to 2), we obtain a successful attack path -- the red path $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,3)} \rightarrow \textrm{x}^{(2,7)}$, where $\textrm{x}^{(2,7)}$ fools the target model. However, using the gold label probability to guide the beam search (set size to 2) still leads to failure.
}
\label{model2}
\end{figure*}

To tackle the above issues, we propose \textbf{PDBS}, a context-aware approach for textual adversarial attack through \textbf{P}robability \textbf{D}ifference guided \textbf{B}eam \textbf{S}earch.
Following CLARE \cite{CLARE}, PDBS uses three RoBERTa-based actions to obtain candidate adversarial examples, but it is different from CLARE in path search.
Specifically, we use the probability difference to guide the path search, where the probability difference is a trade-off between the probabilities of the gold label and the other labels.
In addition, we propose to use the beam search for the path search. Compared to the greedy search, the beam search provides multiple search channels, thus improving the attack success rate.
For example, in Figure \ref{model2} using the probability difference guided beam search enables us to obtain the successful path $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,3)} \rightarrow \textrm{x}^{(2,7)}$.

We conduct experiments on Yelp, AG, MNLI, and QNLI to evaluate PDBS. Experimental results show that PDBS consistently outperforms previous state-of-the-art models in terms of A-rate, Mod, Sim, PPL, and overall GErr. Specifically, PDBS brings up to +19.5\% A-rate gains. In addition, the ablation studies ($\S$\ref{5.1}) validate the effectiveness of the probability difference and the beam search.

In sum, we summarize the contributions as follows: (1) we propose the first textual attack model using probability difference to guide the attack path search; (2) we are the first to use the beam search in non-targeted textual adversarial attack; (3) our model achieves new state-of-the-art results across a set of benchmark datasets.

\section{Related Work}
Textual Adversarial attack has been widely studied.
Previous work can be divided into two types: white-box \cite{papernot2016,Ebrahimi2018,Wallace2019,Meng2020} and black-box. The white-box work needs structures and parameters of target models, which are inaccessible in reality. By contrast, the black-box work only requires the easily accessible class probabilities, thus attracting much research attention.

\subsection{Heuristic Rule-based Textual Adversarial Attack under Black-box Setting}
Early text adversarial attack models mainly use heuristic rule-based methods to generate candidate adversarial examples. Liang et al. \shortcite{Liang2018} use three actions to generate examples, but these examples are out of fluency and lose semantic information. To tackle this, recent work proposes more effective methods. Alzantot et al. \shortcite{Alzantot2018} use the genetic algorithm to find synonyms for word replacement. Ren et al. \shortcite{Ren2019} first obtain the word saliency and then greedily replace words with synonyms derived from WordNet \cite{WordNet}. Zang et al. \shortcite{Zang2020} use the particle swarm optimization algorithm to find synonyms having the same sememe of original words. Jin et al. \shortcite{TextFooler} first rank words and then greedily replace them with synonyms obtained by semantically enhanced embeddings \cite{mrkvsic2016}. Although these models achieve good results, they are context-unware, disabling them from generating contextual and fluent adversarial examples. By contrast, our model is based on the PLM, enabling us to avoid the above problems.

\subsection{PLM-based Textual Adversarial Attack under Black-box Setting}
To generate candidate adversarial examples in a context-aware manner, recent work proposes to use PLMs.
Li et al. \shortcite{BERTAttack} first use a Replace operation to find the vulnerable words based on the gold label probability, then it iteratively replaces the vulnerable words with synonyms generated by BERT \cite{BERT}.
Garg and Ramakrishnan \shortcite{BAE} propose a model similar to Li et al. \shortcite{BERTAttack}, but they add an Insert operation when finding vulnerable words.
More recently, Li et al. \shortcite{CLARE} propose CLARE, which uses a mask-then-infill strategy to obtain candidate adversarial examples with RoBERTa \cite{RoBERTa}.
The above three PLM-based models generate much more contextual and fluent adversarial examples than rule-based models.
However, we demonstrate there are two common flaws in them: (1) they only use the gold label probability to guide the attack path search, and (2) they use the greedy search, which limits the search space. Compared to them, our model considers all label probabilities (probability difference) and uses the beam search while also based on the PLM.

\section{Approach}
\subsection{Problem Definition}

In this paper, we focus on generating context-aware textual adversarial examples under the black-box and non-targeted settings. Textual adversarial example generation centers around a text classification model $f(\cdot)$, referred to as the target model. The black-box setting enables us to access probability outputs of $f(\cdot)$.
We first define a text classification dataset $\mathcal{D}=\{(\textrm{x}^{(1)},\textrm{y}^{(1)}),(\textrm{x}^{(2)},\textrm{y}^{(2)}),…,(\textrm{x}^{(d)},\textrm{y}^{(d)})\}$ and the class label set $ \mathcal{Y} = \{y_1,y_2,…, y_\ell\}$, where $\textrm{x}^{(i)}$ is a text sequence, $\textrm{y}^{(i)}$ is the gold class label of $\textrm{x}^{(i)}$ and $\textrm{y}^{(i)} \in \mathcal{Y}$. Given a text-label pair $(\textrm{x},\textrm{y}) \in \mathcal{D}$ (assume $f(\textrm{x}) = \textrm{y}$), an adversarial example $\textrm{x}'$ is supposed to modify $\textrm{x}$ to fool the target model: $f(\textrm{x}') \neq \textrm{y}$ (non-targeted). In addition,  modifications made to $\textrm{x}'$ should be minimal, such that $\textrm{x}'$ should be close to $\textrm{x}$ and a human cannot percesive the modifications \cite{CLARE}. We achieve above constraints by restricting the similarity between $\textrm{x}'$ and $\textrm{x}$, namely $\textrm{sim}(\textrm{x}',\textrm{x}) \geq \epsilon$, where $\textrm{sim}(\cdot,\cdot)$ calculates the similarity and $\epsilon$ is a pre-set threshold. Following previous work \cite{CLARE,TextFooler}, we use neural networks as the $\textrm{sim}(\cdot,\cdot)$ to calculate the cosine similarity of text pairs in the embedding space.

\subsection{Mask-then-Infill Strategy}\label{3.2}
Li et al. \shortcite{CLARE} demonstrate that Pre-trained Language Model (PLM) helps to produce adversarial examples that are more fluent and closer to the original text, and they propose CLARE, which contains a Mask-then-Infill strategy. The strategy consists of three perturbation operations, i.e., Replace (R), Insert (I) and Merge (M). For a fair comparison, we also use the strategy in our model. But we propose a more effective way to select substitute tokens (see Eq.\ref{8}) and a more effective scoring function to calculate action scores (see Eq.\ref{9}).

Given the text-label pair $\textrm{(x,y)}$ and a position $i$, where $\textrm{x} = x_1x_2…x_n$ and $1\leq i \leq n$, the above three operations apply masks at the given position of $\textrm{x}$ using operation-specific methods.

The Replace obtains a masked sequence by replacing the token $x_i$ with the [mask] token.
\begin{equation}
\textrm{R}(i, \textrm{[mask]}) = x_1...x_{i-1}\textrm{[mask]}x_{i+1}...x_n
\end{equation}

The Insert obtains a masked sequence by inserting the [MASK] token between the token ${x}_i$ and ${x}_{i+1}$, increasing the sequence length by 1.
\begin{equation}
\textrm{I}(i, \textrm{[mask]}) = x_1...x_i\textrm{[mask]}x_{i+1}...x_n
\end{equation}

The Merge obtains a masked sequence by replacing the bigram ${x}_i{x}_{i+1}$ with the [MASK] token, decreasing the sequence length by 1.\footnote{For the last position (i.e., $n$), the operation just replaces the token ${x}_n$ with the [MASK] token.}
\begin{equation}
\textrm{M}(i, \textrm{[mask]}) = x_1...x_{i-1}\textrm{[mask]}x_{i+2}...x_n
\end{equation}

A total of $3n$ masked sequences can be obtained for \textrm{x}: 3 sequences for each position. For clarity, we use $\hat{\textrm{x}}$ to denote any one of the $3n$ sequences.
\begin{equation}
{\hat{\textrm{x}} \in \{ \textrm{R}(i, \textrm{[mask]}), \textrm{I}(i, \textrm{[mask]}), \textrm{M}(i, \textrm{[mask]})\}_{i=1}^{n}}
\end{equation}

The \textbf{infilling action} first use a MLM to generate a substitute token set $\mathcal{Z}$ for the mask positon in $\hat{\textrm{x}}$, where $\mathcal{Z}$ is the vocabulary list $\mathcal{V}$ of the MLM initially. Then it selects a $z \in \mathcal{Z}$ to infill the [mask] position in $\hat{\textrm{x}}$.
\begin{equation}
\textrm{infill}(\hat{\textrm{x}},\textrm{[mask]},z) = x_1…x_{j-1}zx_{j+1}…x_m
\end{equation}
where we assume the [mask] token is at the position $j$ and the length of $\hat{\textrm{x}}$ is $m$.\footnote{The Insert and Merge change sequence lengths, so we use a variable $m$ to denote the length of the masked sequence $\hat{\textrm{x}}$.}
We refer to the $x_1…x_{j-1}zx_{j+1}…x_m$ as a \textbf{candidate adversarial example}, which is denoted as  $\hat{\textrm{x}}_z$ for clarity.

To obtain high-quality adversarial examples, the infilling action only selects $z$ when satisfying the following constraints.
\begin{compactitem}
\item The probability of $z$ predicted by the MLM is larger than a threshold: $p_{\scriptsize{\textrm{MLM}}}(z|\hat{\textrm{x}}) > \alpha$
\item The similarity between $\hat{\textrm{x}}_z$ and $\textrm{x}$ is larger than a threshold: $\textrm{sim}(\hat{\textrm{x}}_z, $\textrm{x}$) > \beta$
\end{compactitem}

We use the $\alpha$ and $\beta$ values reported in {CLARE}: $\alpha = 5e^{-3}$, $\beta=0.7$.
We then obtain the substitute token set $\mathcal{Z}$ as follows.
\begin{equation}
{ \mathcal{Z}= \{z \in \mathcal{V}\}  \ s.t.   \
p_{\scriptsize{\textrm{MLM}}}(z|\hat{\textrm{x}}) > \alpha, \textrm{sim}(\hat{\textrm{x}}_{z}, \textrm{x}) > \beta}
\end{equation}

Previous work \cite{BERTAttack,BAE,CLARE} selects the $z \in \mathcal{Z}$ that most confuses the $f(\cdot)$ as the final substitute token for the infilling action, where
\begin{equation}
z = \underset{z \in \mathcal{Z}}{\arg  \min} \  p_{{f}(\cdot)}(\textrm{y}|\hat{\textrm{x}}_z)
\label{7}
\end{equation}

We observe that the above $z$ solely considers the probability of the gold class label $\textrm{y}$.
However, we demonstrate that for multi-class classification tasks (\# class $\geq$ 3), taking the probabilities of all class labels into consideration is more effective for generating adversarial examples. Thus we select the $z \in \mathcal{Z}$ according to the probability difference, which is calculated by gold label probability minus the maximum probability of the other class labels.
\begin{equation}
z = \underset{z \in \mathcal{Z}, \ \textrm{y}' \in \mathcal{Y}\setminus \{\textrm{y}\}}{\arg  \min} \  [p_{{f}(\cdot)}(\textrm{y}|\hat{\textrm{x}}_z) - p_{{f}(\cdot)}(\textrm{y}'|\hat{\textrm{x}}_z) ]
\label{8}
\end{equation}
where $\mathcal{Y}\setminus \{\textrm{y}\}$ denotes removing $\textrm{y}$ from $\mathcal{Y}$. For binary classifications, Eq.\ref{7} and Eq.\ref{8} make no difference.

In addition, the Eq.\ref{7} uses the greedy search since it solely selects the token that minimizes the gold label probability, which limits the search space. To tackle this, we propose to use the beam search. Specifically, we select the $z$ for $\mathcal{K}$ times and search an attack path through $\mathcal{K}$ channels (see $\S\ref{beam}$), where $\mathcal{K}$ is the beam size.
For the $k\mbox{-}th$ of the $\mathcal{K}$ selections, we set the constraint of Eq.\ref{8} to  ($z \in \mathcal{Z}\setminus\{z^{(1)},z^{(2)},...,z^{(k-1)}\}, \textrm{y}' \in \mathcal{Y}\setminus \{\textrm{y}\}$), where $\{z^{(1)},z^{(2)},...,z^{(k-1)}\}$ is the set of selected tokens. Finally, we obtain $\mathcal{K}$ infilling actions for the masked sequence $\hat{\textrm{x}}$, where each action infills the [mask] position with a token from
$\{z^{(1)},z^{(2)},...,z^{(\mathcal{K})}\}$. Thus for the text sequence $\textrm{x}$ which has $3n$ masked sequences, we obtain a total of $(3*n*\mathcal{K})$ actions. For each action, we use the probability difference as its score.
\begin{equation}
s_{(\textrm{x, y})}(a) = \underset{\textrm{y}' \in  \mathcal{Y}\setminus \{\textrm{y}\}}\min [p_{f(\cdot)}(\textrm{y}|a(\textrm{x})) - p_{f(\cdot)}(\textrm{y}'|a(\textrm{x}))]
\label{9}
\end{equation}
where $a$ denotes the action. $a(\textrm{x})$ denotes applying $a$ to x, generating a candidate adversarial example. And the lower the score, the better the action.

\subsection{ Probability Difference Guided Beam Search}\label{beam}
We formulate the textual attacking as a path search problem and use probability difference to guide the beam search to find an attack path, where the destination example in the path can fool the target model, as the red path $\textrm{x}^{(0)} \rightarrow \textrm{x}^{(1,3)} \rightarrow \textrm{x}^{(2,7)}$ in Figure \ref{model2} shows.

Algorithm 1 shows the path search procedure using PDBS. We divide the total $T$ search iterations (Line 6-16) into two cases: the first iteration and the subsequent iterations. In the first iteration (i.e., $t=1$, Line 8-10), we start the beam search from the sequence $\textrm{x}^{(0)}$, and first obtain its top $\mathcal{K}$ best actions according to action scores (Line 9). Then we sort these actions (Line 17) and sequentially apply them to $\textrm{x}^{(0)}$ (Line 18-19). Each action generates a candidate adversarial example (Line 20), and if the example can fool $f(\cdot)$, it attacks successfully and we return it (Line 21-22). Else we regard the $\mathcal{K}$ examples generated by the $\mathcal{K}$ actions as perturbed text sequences and pass them to the next iteration.

If we fail to attack in the first iteration, we continue the beam search in subsequent iterations (i.e., $t \geq 2$, Line 11-16). In the $t$-$th$ iteration, we conduct path search on the $\mathcal{K}$ perturbed text sequences generated in the $(t-1)$ iteration. For each of the $\mathcal{K}$ sequences, we first obtain its top $\mathcal{K}$ best actions (Line 13). Thus we can obtain a total of $\mathcal{K}^2$ actions. Next, we first sort the $\mathcal{K}^2$ actions (Line 17), and then we select the top $\mathcal{K}$ best actions and sequentially apply them to their texts (Line 18-20).
Each action generates a candidate adversarial example (Line 20), and if the example can fool $f(\cdot)$, it attacks successfully, and we return it (Line 21-22). Else we regard the $\mathcal{K}$ examples generated by the $\mathcal{K}$ actions as perturbed text sequences and pass them to the $(t+1)$ iteration.

\begin{table}[H]
\centering
\begin{tabular}{l}
\toprule
\textbf{Algorithm 1} Adversarial Attack by PDBS                            \\ \midrule
\resizebox{0.950\linewidth}{!}
{
\begin{tabular}[c]{@{}l@{}}
1: \textbf{Input:} text sequence \textrm{x}, gold class label \textrm{y},
\\ \quad \quad \quad \quad \ target model $f(\cdot)$ \\
2: \textbf{Output:} An adversarial example \\
3: \textbf{Initialization}: $\textrm{x}^{(0)} \leftarrow \textrm{x}$ \\
4: $\mathcal{K} \leftarrow \ $beam size \\
5:  $  T \leftarrow \ $iteration times \\
\\
6: \textbf{for} $1\leq t \leq T$ \textbf{do} \\
7:\ \ \  \quad $\mathcal{P}_2 \leftarrow \varnothing$  \\
8:\ \ \  \quad  \textbf{if} \ $t == 1$ \ \textbf{do} \\
\quad \quad  \  \ \  \textcolor[RGB]{73,153,54}{// start beam search only from $\textrm{x}^{(0)}$} \\
9:\ \ \  \quad \quad $\mathcal{P}_2 \leftarrow $ \ GetBestActions($\textrm{x}^{(0)}$) \\
10: \quad   \textbf{end if} \\
11: \quad   \textbf{if \ $t  \geq 2$ \ do} \\
\quad \quad \ \ \ \textcolor[RGB]{73,153,54}{// continue beam search with $\mathcal{K}$ sequences} \\
12: \quad  \quad \textbf{for} $1 \leq k \leq \mathcal{K}$ \textbf{do}\\
13:\ \ \quad \quad   \quad $\mathcal{P}^{(k)} \leftarrow$\ \ GetBestActions($\textrm{x}^{(t-1,k)}$) \\
14:\ \ \quad \quad   \quad $\mathcal{P}_2 \leftarrow \mathcal{P}_2 \cup \mathcal{P}^{(k)}$ \\
15:\ \ \quad   \quad \textbf{end for} \\
\quad \quad \quad \quad \textcolor[RGB]{73,153,54}{// $\mathcal{P}_2$ contains $\mathcal{K}^2$ action-text pairs} \\
16: \quad   \textbf{end if} \\
\\
17:\ \ \quad Sort $\mathcal{P}_2$ based on action scores ascendingly \\
18:\ \ \quad \textbf{for} $1 \leq k \leq \mathcal{K}$ \textbf{do} \\
\quad \quad \quad \textcolor[RGB]{73,153,54}{// solely select the top $\mathcal{K}$ best pairs from $\mathcal{P}_2$} \\
19:\ \ \quad \quad ($a$, $\bar{\textrm{x}}$) $\leftarrow$ the $k\mbox{-}th$ best action-text pair in $\mathcal{P}_2$ \\
20:\ \ \quad \quad $\textrm{x}^{(t,k)} \leftarrow $ Apply $a$ to $\bar{\textrm{x}}$ \\
21:\ \ \quad \quad \textbf{if} $f(\textrm{x}^{(t,k)}) \neq \textrm{y}$ \textbf{then} \textbf{return} $\textrm{x}^{(t,k)}$ \\
\quad \quad \quad \quad \textcolor[RGB]{73,153,54}{// attack successfully} \\
22:\ \ \quad \quad \textbf{end if}\\
23:\ \ \quad \textbf{end for} \\
24: \textbf{end for}\\
25: \textbf{return} None \\
\\
26: \textbf{def}\ GetBestActions($\bar{\textrm{x}}$)\ \textbf{do} \\
\quad \  \ \  \textcolor[RGB]{73,153,54}{// return the top $\mathcal{K}$ best action-text pairs for $\bar{\textrm{x}}$} \\
27:\ \ \quad $\mathcal{P}_1 \leftarrow \varnothing$ \\
28:\ \ \quad $\mathcal{A} \leftarrow$ all infilling actions of $\bar{\textrm{x}}$\\
\quad \quad \  \ \  \ \textcolor[RGB]{73,153,54}{// a total of (3*$|\bar{\textrm{x}}|$*$\mathcal{K}$) actions} \\
29:\ \ \quad {Sort $\mathcal{A}$ based on action scores ascendingly} \\
\quad \quad \ \  \ \ \textcolor[RGB]{73,153,54}{// the lower the score, the better the action} \\
30: \ \ \  \ \  \textbf{for} $1 \leq k \leq \mathcal{K}$ \textbf{do} \\
31: \quad  \ \  \quad $a \leftarrow $ the $k\mbox{-}th$ action in $\mathcal{A}$ \\
\quad \quad \quad \  \ \ \ \ \textcolor[RGB]{73,153,54}{// select the ${k}\mbox{-}th$ best action from $\mathcal{A}$} \\
32: \quad \quad \ \  $\mathcal{P}_1 \leftarrow \mathcal{P}_1\cup \{(a,\bar{\textrm{x}})\}$ \\
33: \quad \ \  \textbf{end for} \\
34: \quad \ \  \textbf{return} $\mathcal{P}_1$ \\
35: \textbf{end def} \\
\bottomrule
\end{tabular}}
\end{tabular}
\label{alg}
\end{table}

\section{Experiment}
\subsection{Datasets and Baselines}
\textbf{Datasets.} To evaluate the proposed PDBS, we conduct experiments on the following benchmark datasets regarding text classification and Natural Language Inference (NLI).

\begin{compactitem}
\item \textbf{Yelp} \cite{Zhang2015}: a sentiment classification dataset from restaurant reviews, containing two classes: positive and negative.
\item \textbf{AG} \cite{Zhang2015}: a dataset classifying news articles to four classes: world, sports, business, and science/technology.
\item \textbf{MNLI} \cite{Williams2018broad}: a triple-class classification dataset regarding NLI. It is composed of premise-hypothesis pairs. Each pair is labeled with the relation between the premise and hypothesis. The relation set contains \textit{entailment}, \textit{contradiction} and \textit{neutral}.
\item \textbf{QNLI} \cite{Wang2018glue}: a binary classification dataset regarding NLI. Each instance is a question-answer pair, and is labeled whether the answer corresponds to the question.

\end{compactitem}

We summarize more dataset details in Table \ref{1111} \cite{CLARE}. Following previous work \cite{TextFooler,BERTAttack,CLARE}, we evaluate our model on a set of 1,000 examples.\footnote{For a fair comaprison, we actually use the exact 1,000 examples used in {CLARE} for each dataset. \url{https://github.com/cookielee77/CLARE}}
In addition, to comprehensively evaluate the PDBS, we conduct experiments on four additional datasets: DBpedia ontology dataset \cite{Zhang2015}, Stanford sentiment treebank \cite{Socher2013recursive}, Microsoft Research Paraphrase Corpus \cite{Dolan2005automatically}, and Quora Question Pairs \cite{Wang2018glue}. Limited by space, we report the results of these additional datasets in Appendix \ref{A}.

\begin{table}[htb] \small
\centering
\renewcommand\tabcolsep{4.5pt}
\begin{tabular}{lccrcc}
\toprule
 \textbf{Dataset}   & \multicolumn{1}{l}{\textbf{\# Class}} & \multicolumn{1}{l}{\textbf{Train}} & \multicolumn{1}{l}{\textbf{Test}} & \multicolumn{1}{l}{\textbf{Avg Len}} & \multicolumn{1}{l}{\textbf{Acc(\%)}} \\ \midrule
 {Yelp}      & 2                                       & 560K                               & 38.0K                               & 130                                  & 95.9                                 \\
{AG} & 4                                       & 120K                               & 7.6K                              &\ \ 46                                   & 95.0                                 \\ \midrule

 {MNLI}      & 3                                       & 392K                               & 9.8K                              & 23/11                                & 84.3                                 \\

 {QNLI}      & 2                                       & 105K                               & 5.4K                              & 11/31                                & 91.4                                 \\\bottomrule
\end{tabular}
\caption{\label{1111}Dataset details. The ``Acc(\%)'' denotes the target model’s accuracy on the original test set without adversarial attack.}
\end{table}

\begin{table*}[htb] \small
\centering
\renewcommand\tabcolsep{4.5pt}
\begin{tabular}{lccccccccccc}
\toprule
                        & \multicolumn{5}{c}{\textbf{Yelp}}                                                                                &                               & \multicolumn{5}{c}{\textbf{AG}}                                                                                \\ \midrule
\textbf{Model}          & \textbf{A-rate($\%$)\textcolor{red}{$\uparrow$}} & \textbf{Mod($\%$)\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{Sim\textcolor{red}{$\uparrow$}} & \textbf{PPL\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{GErr\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \multicolumn{1}{r}{\textbf{}} & \textbf{A-rate($\%$)\textcolor{red}{$\uparrow$}} & \textbf{Mod($\%$)\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{Sim\textcolor{red}{$\uparrow$}} & \textbf{PPL\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{GErr\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} \\ \midrule
TextFooler              & 77.0                             & 16.6                           & 0.70                   & 163.3                    & 1.23                      & \multicolumn{1}{r}{}          & 56.1                             & 23.3                           & 0.69                   & 331.3                    & 1.43                      \\
\multicolumn{1}{c}{+LM} & 34.0                             & 17.4                           & 0.73                   & 90.0                     & 1.21                      & \multicolumn{1}{r}{}          & 23.1                             & 21.9                           & 0.74                   & 144.6                    & 1.07                      \\
BERT-Attack             & 71.8                             & 10.7                           & 0.72                   & 90.8                     & 0.27                      & \multicolumn{1}{r}{}          & 63.4                             &\ \  7.9                            & 0.71                   & 90.6                     & 0.25                      \\
CLARE                   & 79.7                             & 10.3                           & 0.78                   & 83.5                     & 0.25                      & \multicolumn{1}{r}{}          & 79.1                             &\ \  6.1                            & 0.76                   & 86.0                     & 0.17                      \\ \midrule
PDBS                    & \textbf{99.2}                    & \textbf{\ \ 4.9}                   & \textbf{0.80}          & \textbf{64.2}            & \textbf{0.13}             & \multicolumn{1}{r}{\textbf{}} & \textbf{87.7}                    & \textbf{\ \ 6.0}                   & \textbf{0.78}          & \textbf{84.7}            & \textbf{0.07}             \\ \toprule
\textbf{}               & \multicolumn{5}{c}{\textbf{MNLI}}                                                                                & \textbf{}                     & \multicolumn{5}{c}{\textbf{QNLI}}                                                                                \\ \midrule
\textbf{Model}          & \textbf{A-rate($\%$)\textcolor{red}{$\uparrow$}} & \textbf{Mod($\%$)\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{Sim\textcolor{red}{$\uparrow$}} & \textbf{PPL\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{GErr\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \multicolumn{1}{r}{\textbf{}} & \textbf{A-rate($\%$)\textcolor{red}{$\uparrow$}} & \textbf{Mod($\%$)\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{Sim\textcolor{red}{$\uparrow$}} & \textbf{PPL\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} & \textbf{GErr\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}} \\ \midrule
TextFooler              & 59.8                             & 13.8                           & 0.73                   & 161.5                    & 0.63                      &                               & 57.8                             & 16.9                           & 0.72                   & 164.4                    & 0.62                      \\
\multicolumn{1}{c}{+LM} & 32.3                             & 12.4                           & 0.77                   & 91.9                     & 0.50                      &                               & 29.2                             & 17.3                           & 0.75                   & 85.0                     & 0.42                      \\
BERT-Attack             & 82.7                             &\ \  8.4                            & 0.77                   & 86.7                     & 0.04                      &                               & 76.7                             & 13.3                           & 0.73                   & 86.5                     & 0.03                      \\
CLARE                   & 88.1                             &\ \  7.5                            & 0.82                   & 82.7                     & \textbf{0.02}             &                               & 83.8                             & 11.8                           & 0.78                   & 76.7                     & \textbf{0.01}             \\ \midrule
PDBS                    & \textbf{98.3}                    & \textbf{\ \ 4.5}                   & \textbf{0.83}          & \textbf{73.9}            & 0.07                      &                               & \textbf{92.6}                    & \textbf{\ \ 8.3}                   & \textbf{0.79}          & \textbf{67.1}            & 0.07                      \\ \bottomrule
\end{tabular}

\caption{\label{2}Performance comparisons across the four datasets. We use fine-tuned BERT models as target models.
Bold font values denote the best results for each metric. \textcolor{red}{$\uparrow$}(\textcolor{red}{\textcolor{red}{\textcolor{red}{\textcolor{blue}{$\downarrow$}}}}) indicates the higher (lower) the better.}
\end{table*}

\noindent \textbf{Baselines.} We compare our model with previous state-of-the-art models regarding black-box textual adversarial attack, including TextFooler \cite{TextFooler}, TextFooler+LM \cite{CLARE}, BERT-Attack \cite{BERTAttack} and CLARE \cite{CLARE}, where the first two are heuristic rule-based and the last two are PLM-based methods.\footnote{Following CLARE, we do not list BAE \cite{BAE} since it has a similar performance as BERTAttack.}
We report more details of these baselines in Appendix \ref{B}.

\subsection{Implementation Details}
Following CLARE, we use a distilled version of RoBERT to infill the masked tokens. Across all the datasets, we set the iteration times $T$ to 10, the beam size $\mathcal{K}$ to 10. Moreover, we investigate model performance against various $T$ values in Appendix \ref{E}.
As for other hyper-parameters, we use just the same settings to CLARE.
 Following previous work \cite{BERTAttack,CLARE}, we use fine-tuned BERT models as the target models.\footnote{For a fair comparison, we actually use the fine-tuned BERT models used in CLARE.}

In addition, we also conduct experiments on two additional target models, namely Word-LSTM \cite{Hochreiter1997long} and ESIM \cite{Chen2017enhanced}, where the former is specific for text classification, and the latter is specific for NLI. We report experimental results on the two additional models in Appendix \ref{C}.

\subsection{Automatic Evaluation Metrics}
Following CLARE, we use the Attack success rate (\textbf{A-rate}), Modification Rate (\textbf{Mod}), Textual similarity (\textbf{Sim}), Perplexity (\textbf{PPL}) and Grammar error (\textbf{GErr}) to comprehensively evaluate the model performance, where the last four are used to evaluate the quality of adversarial examples. We report more metric details in Appendix \ref{D}.

\subsection{Main Results}\label{4.4}

We summarize the performance results in Table \ref{2}. Obviously, PDBS successfully triggers classification errors in the target models across the four datasets.

Specifically, PDBS achieves the current state-of-the-art performance in terms of A-rate, Mod, Sim, and PPL across the four datasets. To be more precise, (1) PDBS delivers +8.6\% to +19.5\% A-rate gains compared to the previous best model, i.e., CLARE. These gains explicitly validate the excellent attack ability of PDBS. (2) PDBS consistently decreases the Mod scores, especially up to -5.4\% on Yelp, implying that the adversarial examples are more humanly imperceptible. (3) The consistent optimal Sim scores indicate that our adversarial examples are closer to the original texts than those generated by previous models. (4) PDBS decreases the PPL by -9.75 on average, especially -19.3 on Yelp. These lower PPL scores confirm that our adversarial examples are more fluent than others.

In addition, PDBS is also superior to the listed baselines in terms of overall GErr scores (0.085 \textit{v.s.} 0.113, averaged scores), demonstrating that our adversarial examples are grammatically correct.

We attribute these exciting performance gains to the fact that (1) the probability difference takes all class probabilities into consideration and is more effective than solely using the gold label probability; (2) the beam search provides more search channels, enabling our model to avoid suffering from limited search space caused by the greedy search.

We present case studies of adversarial examples in Appendix \ref{F}.

\subsection{Human Evaluation}
To further evaluate the quality of adversarial examples, we set up human evaluations on three criteria: {similarity} (\textbf{Sim.}), {fluency \& grammaticality} (\textbf{Flu. \& Gra.}), and {label consistency} (\textbf{Con.}). We choose three independent volunteers who are native English speakers.
We randomly select 200 instances (original text, adversarial example, gold label) from AG and MNLI respectively, where the adversarial examples can fool our PDBS.

We first ask the volunteers to score (from 0-1) the similarity of original-adversarial text pairs. And the higher the score, the better the similarity. Then we evaluate the fluency and grammaticality by asking volunteers to score (from 1-5) the mixed adversarial examples and original texts. And the higher the score, the better the fluency and grammaticality. As for label consistency, we ask volunteers to annotate adversarial examples and compare their annotations with the gold labels. We finally average the scores marked by the three volunteers.

Table \ref{3} shows the evaluation results. We can observe that: (1) the {Sim.} scores are 0.86 and 0.89, indicating that the adversarial examples maintain the most semantics of the original texts, which preliminarily proves the validity of our model. (2) the {Flu. \& Gra.} scores of the adversarial examples are close to the original ones (3.6 \textit{v.s.} 4.1 and 3.7 \textit{v.s.} 3.7), implying that the adversarial examples have good fluency and satisfy the grammatical rules on the whole. (3) the {Con.} scores are 0.77 and 0.84, revealing that the majority of the adversarial examples are humanly imperceptible.

\begin{table}[thb] \small
\centering
\renewcommand\tabcolsep{5.5pt}
\begin{tabular}{llccc}
\toprule
\multicolumn{2}{c}{\textbf{Dataset}}            & \multicolumn{1}{l}{\textbf{Sim.\textcolor{red}{$\uparrow$}}} & \multicolumn{1}{l}{\textbf{Flu. \& Gra.\textcolor{red}{$\uparrow$}}} & \multicolumn{1}{l}{\textbf{Con.\textcolor{red}{$\uparrow$}}} \\ \midrule
\multirow{2}{*}{{AG}} & Original    & \multirow{2}{*}{0.86}                             & 4.1                                                              & \multirow{2}{*}{0.77}                                    \\
                                  & Adversarial &                                                   & 3.6                                                              &                                                          \\ \midrule
\multirow{2}{*}{{MNLI}}    & Original    & \multirow{2}{*}{0.89}                             & 4.0                                                              & \multirow{2}{*}{0.84}                                    \\
                                  & Adversarial &                                                   & 3.7                                                              &                                                          \\ \bottomrule
\end{tabular}
\caption{\label{3}Human evaluation results. }
\end{table}

In addition, the results on AG are worse than the ones on MNLI. We attribute it to the fact that the adversarial examples of AG have an overall higher modification rate than MNLI (6.2\% \textit{v.s.} 4.4\%).\footnote{Since we solely select a part of the two datasets, thus the two Mod scores are not the same to the ones in Table \ref{2}.}

\section{Analysis}
We first conduct ablation studies on the probability difference and the beam search ($\S5.1$). Then we
examine the transferability of adversarial examples ($\S5.2$). At last, we explore using PDBS to improve the robustness of the target models ($\S5.3$).

\subsection{Ablation Study} \label{5.1}
\textbf{Ablation on the Probability Difference.}
We report the ablation results on the probability difference (\textbf{P\&D}) in Table \ref{4}. We use the \texttt{-P\&D} to denotes ablating the P\&D from PDBS, instead we use the gold label probability to select substitutes (see Eq. \ref{7}, also select $\mathcal{K}$ times) and calculate the action score.
\begin{equation}
s_{(\textrm{x, y})}(a) = p_{f(\cdot)}(\textrm{y}|a(\textrm{x}))
\end{equation}

The ablation results show that the P\&D significantly promotes the A-rate, delivering +16.3\% and +15.6\% gains on MNLI and AG respectively, but the four quality evaluation metrics become worse in current scenario.\footnote{Since the probability difference and the gold label probability make no difference in binary classifications (see $\S$\ref{3.2}), we do not conduct ablations on Yelp and QNLI. }\footnote{In both datasets, the A-rate of \texttt{-P\&D} is inferior to CLARE, which we attribute to the fact that \texttt{-P\&D} and CLARE use different $T$ values. For a fair comparison, we conduct more experiments in Appendix \ref{G}.}
We attribute the worse results to the fact that the \texttt{-P\&D} generates much fewer adversarial examples, thus making it unfair for comparison. For a fair comparison, we select a part of the 1000 test instances of the two datasets, where if the \texttt{-P\&D} can attack an instance successfully, then the instance is selected. Finally, 785 and 672 instances are selected for MNLI and AG, which are referred to as \textbf{Reduced MNLI} and \textbf{Reduced AG}, respectively. We report extended ablation results on the two reduced datasets in Table \ref{5}. We observe that (1) the two A-rates of PDBS are also 100\%; (2) the P\&D consistently benefits PDBS in terms of the four quality evaluation metrics.

We conclude that using the probability difference is more effective than using the gold label probability, especially in terms of the A-rate.

\begin{table}[thb] \small
\centering
\renewcommand\tabcolsep{2pt}
\begin{tabular}{lccccc}
\toprule
                    \multicolumn{6}{c}{\textbf{MNLI}}                                                                                \\ \midrule
\textbf{Model}     & \multicolumn{1}{r}{\textbf{A-rate($\%$)\color{red}{$\uparrow$}}}  & \multicolumn{1}{r}{\textbf{Mod($\%$)\color{blue}{$\downarrow$}}} & \multicolumn{1}{r}{\textbf{Sim\color{red}{$\uparrow$}}} & \multicolumn{1}{r}{\textbf{PPL\color{blue}{$\downarrow$}}} & \multicolumn{1}{r}{\textbf{GErr\color{blue}{$\downarrow$}}} \\ \midrule
{PDBS}      & \textbf{98.3}                                        & 4.5                                                & 0.83                              & 73.9                                         & 0.07                                          \\
{\quad \texttt{-P\&D}} & 82.0                                                & \textbf{3.4}                                       & \textbf{0.84}                              & \textbf{71.7}                                & \textbf{0.06}                                 \\
\toprule
                    \multicolumn{6}{c}{\textbf{AG}}                                                                                \\ \midrule
\textbf{Model}     & \multicolumn{1}{r}{\textbf{A-rate($\%$)\color{red}{$\uparrow$}}} & \multicolumn{1}{r}{\textbf{Mod($\%$)\color{blue}{$\downarrow$}}} & \multicolumn{1}{r}{\textbf{Sim\color{red}{$\uparrow$}}} & \multicolumn{1}{r}{\textbf{PPL\color{blue}{$\downarrow$}}} & \multicolumn{1}{r}{\textbf{GErr\color{blue}{$\downarrow$}}} \\ \midrule
{PDBS}      & \textbf{87.7}                                        & 6.0                                                & 0.78                                       & 84.7                                         & 0.07                                          \\
{\quad \texttt{-P\&D}} & 72.1                                                 & \textbf{4.2}                                       & \textbf{0.79}                              & \textbf{79.2}                                & \textbf{0.05}                                 \\
 \bottomrule
\end{tabular}
\caption{\label{4}Ablation results on probability difference.}
\end{table}

\begin{table}[h] \small
\centering
\renewcommand\tabcolsep{2pt}
\begin{tabular}{lccccc}
\toprule
                    \multicolumn{6}{c}{\textbf{Reduced MNLI}}                                                                                \\ \midrule
\textbf{Model}     & \textbf{A-rate($\%$)\color{red}{$\uparrow$}}  & \textbf{Mod($\%$)\color{blue}{$\downarrow$}} & \textbf{Sim\color{red}{$\uparrow$}} & \textbf{PPL\color{blue}{$\downarrow$}} & \textbf{GErr\color{blue}{$\downarrow$}} \\ \midrule
{PDBS}      & \textbf{100.0}                  & \textbf{3.2}                   & \textbf{0.85}          & \textbf{69.9}            & \textbf{0.04}             \\
{\quad \texttt{-P\&D}} & \textbf{100.0}                  & 3.4                            & 0.84                   & 71.7                     & 0.06                      \\
   \toprule
                    \multicolumn{6}{c}{\textbf{Reduced AG}}                                                                                \\ \midrule
\textbf{Model}     & \textbf{A-rate($\%$)\color{red}{$\uparrow$}} & \textbf{Mod($\%$)\color{blue}{$\downarrow$}} & \textbf{Sim\color{red}{$\uparrow$}} & \textbf{PPL\color{blue}{$\downarrow$}} & \textbf{GErr\color{blue}{$\downarrow$}} \\ \midrule
{PDBS}      & \textbf{100.0}                  & \textbf{4.1}                   & \textbf{0.81}          & \textbf{73.6}            & \textbf{0.04}             \\
{\quad \texttt{-P\&D}} & \textbf{100.0}                  & 4.2                            & 0.79                   & 79.2                     & 0.05                      \\
 \bottomrule
\end{tabular}
\caption{\label{5}Extended ablation results on probability difference.}
\end{table}

\begin{figure}[t]
\centering
\includegraphics[width=0.38\textwidth]{fig/111}
\caption{Ablation results on the beam search.
}
\label{model3}
\end{figure}

\noindent \textbf{Ablation on the Beam Search.}
We report the ablation results on the beam search in Figure \ref{model3}. To be specific, we set the beam size $\mathcal{K}$ from 1 to 10.
When $\mathcal{K}$=1, the beam search is the greedy search. We observe that (1) when the $\mathcal{K}$ increases, the A-rate generally increases across the four datasets; (2) compared to the greedy search ($\mathcal{K}$=1), the beam search ($\mathcal{K}$>1) consistently brings A-rate gains. For example, the beam search with $\mathcal{K}$=10 brings up to +17.62\% A-rate scores on AG. We attribute these to the fact that the beam search provides more search channels, thus improving the performance.

Additionally, we report the performance of the four quality evaluation metrics in Table \ref{6}. We observe that the $\mathcal{K}$ value poses tiny effects on the four metrics, revealing that a bigger $\mathcal{K}$ value can bring a better A-rate score while ensuring the high-quality of the adversarial examples.

Limited by hardware capacity, we cannot run experiments with a bigger $\mathcal{K}$ due to GPU out of memory with 24GB of GPU memory. We believe that it is possible to obtain a better A-rate with a bigger $\mathcal{K}$, especially on MNLI and AG.

\begin{table}[htb] \small
\centering
\renewcommand\tabcolsep{3pt}
\begin{tabular}{lcccc}
\toprule
                 & \textbf{Mod($\%$)} & \textbf{Sim} & \textbf{PPL} & \textbf{GErr} \\ \midrule
\textbf{Yelp}    & 5.22$\pm$0.17                 &0.80$\pm$0.001           &64.5$\pm$0.3           &0.133$\pm$0.002         \\
\textbf{AG} &5.94$\pm$0.08                 &0.79$\pm$0.004           &84.7$\pm$0.2           & 0.078$\pm$0.006            \\
\textbf{MNLI}    & 4.50$\pm$0.11          & 0.83$\pm$0.004  & 74.1$\pm$0.2    & 0.070$\pm$0.003   \\
\textbf{QNLI}    & 8.28$\pm$0.03                 & 0.79$\pm$0.001          & 66.8$\pm$0.3          &0.078$\pm$0.007           \\ \bottomrule
\end{tabular}
\caption{\label{6}Ablation results on the beam search when the $\mathcal{K}$ varing from 1 to 10. We report the averaged value of each metric, as well as its change intervals.
}
\end{table}

\subsection{Transferability and Adversarial Training}
\textbf{Transferability.} The transferability is measured by if the adversarial examples generated with one target model can fool other target models directly \cite{Kurakin2016}.
For this purpose, we first collect the adversarial examples generated for Yelp and MNLI and then use other target models to predict their classes. Following previous work \cite{BERTAttack}, we select the Word-LSTM and ESIM as new target models for Yelp and MNLI, respectively. We report the results in Table \ref{7},  from with we can see that
there is a noticeable degree of transferability between models, and the transferability is higher in Yelp than in MNLI, which is consistent with the conclusion in \cite{TextFooler}.

\begin{table}[htb] \small
\centering
\begin{tabular}{llcc}
\toprule
\textbf{Dataset}               & \textbf{Model} & BERT & Word-LSTM \\ \midrule
\multirow{2}{*}{{Yelp}} & BERT           & \ \ 0.0  & 52.5      \\
                               & Word-LSTM      & 56.7 &\ \  0.0       \\ \toprule
\textbf{Dataset}               & \textbf{Model} & BERT & ESIM      \\ \midrule
\multirow{2}{*}{{MNLI}} & BERT           &\ \  0.0  & 65.1      \\
                               & ESIM           & 72.6 &\ \  0.0       \\ \bottomrule
\end{tabular}
\caption{\label{7}Transferability results. Row $i$ and column $j$ is the accuracy of adversarial examples generated for model $i$ evaluated on model $j$. The lower the score, the better the transferability.
}
\end{table}

\noindent \textbf{Adversarial Training.} We investigate if the adversarial examples can be used to improve the robustness of target models. For this purpose, we conduct a preliminary experiment on adversarial training by feeding the models both the original texts and the adversarial examples, where adversarial examples share the same labels as their original texts.

We collect the adversarial examples curated from the training sets of AG and MNLI and add them to the original training sets. We then use the mixed sets to fine-tune BERT from scratch and attack this adversarially fine-tuned model. The results in Table \ref{8} show that the A-rate and Mod after adversarial re-training become worse, indicating the greater difficulty of attacking. This reveals one of the potencies of our attack system: we can enhance the robustness of a model to future attacks by training it with the generated adversarial examples.

\begin{table}[htb]\small
\centering
\renewcommand\tabcolsep{2pt}
\begin{tabular}{llrrr}
\toprule
                                   &              & \textbf{Acc($\%$)\textcolor{red}{$\uparrow$}} & \textbf{A-rate($\%$)\textcolor{blue}{$\downarrow$}} & \textbf{Mod($\%$)\textcolor{red}{$\uparrow$}} \\ \midrule
\multirow{2}{*}{\textbf{AG}} & Original      & 95.0                         & 85.7                               & 6.0                          \\
                                  & +Adv.Training & 94.8(-0.2)                   & 66.1(-19.6)                        & 12.4(+6.4)                   \\ \midrule
\multirow{2}{*}{\textbf{MNLI}}    & Original      & 84.3                         & 98.3                               & 4.5                          \\
                                  & +Adv.Training & 84.0(-0.3)                   & 80.8(-17.5)                        & 10.2(+5.7)                   \\ \bottomrule
\end{tabular}
\caption{\label{8}Comparisons of original fine-tuning (``Original'') and adversarial fine-tuning (``+ Adv. Training'')  BERT models on AG and MNLI. }
\end{table}

\section{Conclusion}
In this work, we present PDBS, a PLM-based text adversarial attack model. PDBS uses probability difference to guide the beam search when seeking a successful attack path. Experimental results demonstrate that PDBS delivers +8.6\% to +19.5\% attack access rate, and is consistently superior to previous models in four quality evaluation metrics, indicating that the PDBS owns both efficient attack ability and the ability to generate high-quality adversarial examples. Qualitative analyses further show significant advantages of PDBS.


KGML-xDTD: A Knowledge Graph-based Machine Learning Framework for Drug Treatment Prediction and Mechanism Description

Abstract: Background: Computational drug repurposing is a cost- and time-efficient approach that aims to identify new therapeutic targets or diseases (indications) of existing drugs/compounds. It is especially critical for emerging and/or orphan diseases due to its cheaper investment and shorter research cycle compared with traditional wet-lab drug discovery approaches. However, the underlying mechanisms of action (MOAs) between repurposed drugs and their target diseases remain largely unknown, which is still a main obstacle for computational drug repurposing methods to be widely adopted in clinical settings.Results: In this work, we propose KGML-xDTD: a Knowledge Graph-based Machine Learning framework for explainably predicting Drugs Treating Diseases. It is a two-module framework that not only predicts the treatment probabilities between drugs/compounds and diseases but also biologically explains them via knowledge graph (KG) path-based, testable mechanisms of action (MOAs). We leverage knowledge-and-publication based information to extract biologically meaningful "demonstration paths" as the intermediate guidance in the Graph-based Reinforcement Learning (GRL) path-finding process. Comprehensive experiments and case study analyses show that the proposed framework can achieve state-of-the-art performance in both predictions of drug repurposing and recapitulation of human-curated drug MOA paths.Conclusions: KGML-xDTD is the first model framework that can offer KG-path explanations for drug repurposing predictions by leveraging the combination of prediction outcomes and existing biological knowledge and publications. We believe it can effectively reduce "black-box" concerns and increase prediction confidence for drug repurposing based on predicted path-based explanations, and further accelerate the process of drug discovery for emerging diseases.

\section{Introduction}\label{sec:intro}

Traditional drug development is a time-consuming process (from initial chemical identification to clinical trials and finally to FDA approval) that takes around 10-15 years and also comes along with billions-of-dollars investments and high failure rates \citep{Berdigaliyev_Nurken_2020}. Considering the rapid pace of novel disease evolution, it is urgent to find a more efficient and economical drug discovery method. Fortunately, it has been observed that a single drug can often be effective in treating multiple diseases. For example, thalidomide was originally used as an anti-anxiety medication \citep{Miller_M_T_1991}, and was later found to have the potential for the treatment of cancers \citep{Singhal_Seema_1999}. Hence, drug repurposing, also known as the identification of new uses for the approved or experimental drugs/compounds, might bring us the hope to address this urgent need with the advantage of a shorter research cycle, lower investments, and more pre-existing safety tests.

Existing drug repurposing approaches can roughly be categorized into three groups: experimental-based approaches (e.g., binding affinity assays, phenotypic screening), clinical-based approaches (e.g., off-label drug use analysis), and computational-based approaches (e.g., network-based approaches) \citep{Neha_Dhir_2020}. Due to the advancement of techniques, more and more publicly available biomedical data can be freely accessed in different databases such as DrugBank \citep{Wishart_David_S_2017DrugBank}, ChEMBL \citep{Gaulton_Anna_2012}, HMDB \citep{Wishart_David_S_2018HMDB}, which makes the computational approaches seem to be more cost-efficient, particularly when the goal is to prioritize repurposed targets for followup experimental investigation. One of the computational drug repurposing methods commonly used in recent years is to integrate existing biomedical relations from databases or literature into a so-called \textbf{biomedical knowledge graph} (BKG) where unknown drug-disease treatment relationships are predicted via different knowledge graph (KG)-based machine learning models \citep{Himmelstein_Daniel_Scott_2017, Vassilis_N_Ioannidis_2020, Zhang_Rui_2021, Zhang_Xiaolin_2021}. Although these KG-based models are demonstrated to have good predictive performance for drug repurposing, they struggle to explain why some drugs can be useful for treating a given disease in an intuitive and easy-to-understand fashion. To solve the "black-box" concern for drug repurposing prediction, some methods are proposed to leverage KG-based paths as explanations, as illustrated in Figure~\ref{fig:fig1}. However, these existing models cannot be efficiently applied to a large and general BKG without additional weighted edge information \citep{Sosa_Daniel_N_2020} or pre-defined meta-paths derived from domain experts or inefficient computational methods (e.g., degree-weighted path count -- "DWPC") \citep{Liu_Yushan_2021}.

\begin{wrapfigure}{hr}{0.5\textwidth}
  \centering
    \includegraphics[height=4.5cm]{./Figures/main_figure1}
  \caption{Drug repurposing prediction and path-based explanation.}
  \label{fig:fig1}
\end{wrapfigure}

In this study, we customize a large and standardized biomedical knowledge graph and propose a computational drug repurposing framework that predicts not only the treatment probabilities between drugs and diseases but also the KG-based \textbf{mechanism of action} (MOA) \citep{Davis_Ronald_L_2020} paths as their biomedical explanations based on the treatment predictions. For drug repurposing predictions, we first calculate attribute embedding as the initial feature of each node and employ the GraphSAGE model in an unsupervised manner to further capture the neighborhood information for each node, then a Random Forest model is utilized to predict the treatment probability of drug-disease pairs based on their embeddings.
To predict the MOA paths, we employ the ADversarial Actor-Critic (ADAC) reinforcement learning (RL) model \citep{Zhao_Kangzhi_2020} to perform path-finding on the knowledge graph. To encourage the RL model to find paths that are biologically reasonable, we amplify it with knowledge-and-publication-based "demonstration paths", paths that explain why a drug can treat a disease.
Although the underlying mechanisms of action between repurposed drugs and their target diseases largely remain vague, in this study, we assume that a repurposed drug follows similar molecular mechanisms as the known MOAs to treat different diseases. Based on this assumption, we define demonstration paths based on the known drug-target interactions from a curated drug database (e.g., DrugBank v5.1 \citep{Wishart_David_S_2017DrugBank}) and a chemical-knowledge-centric data provider (e.g., Molecular Data Provider v1.2 \footnote{\url{https://github.com/NCATSTranslator/Translator-All/wiki/Molecular-Data-Provider}}) as well as an adjusted PubMed-publication-based version of Normalized Google Distance (NGD) \citep{Cilibrasi_R_L_2007}.
In summary, the main contributions are summarized as follows:

\begin{itemize}
  \item We propose a novel computational model framework that both accurately predict how likely a drug can be used to treat a disease and also predict its corresponding knowledge graph-based mechanism of action path as the explanation of the predicted treatment.
  \item We are innovative in using a knowledge-based and publication-based method to extract demonstration paths from a BKG and leverage it to guide the RL model to identify biologically reasonable paths. Empirical results demonstrate the great effectiveness of it.
\end{itemize}

By comparing with the existing popular KG-based models and evaluating the predicted paths with an expert-curated path-based drug MOA database \emph{DrugMechDB} \citep{Mayers_Michael_2020}, we show that this proposed model framework outperforms the state-of-the-art baseline models on both the predictive performance of drug repurposing and the explanatory performance of recapitulating human-curated MOA paths provided by DrugMechDB. In further case studies, by comparing the model predictions with the real regulatory networks, we show that the proposed framework is effective in identifying biologically reasonable KG-based paths for real-world applications.

\section{Related Work}\label{sec:related_work}

\paragraph{Biomedical knowledge graphs}
A biomedical knowledge graph (BKG) is normally defined as a heterogeneous, semantic graph that integrates biomedical information from multiple data sources \citep{David_N_Nicholson_2020}. Many open-source BKGs have been developed for different biomedical research purposes. These existing BKGs can be categorized into three groups: database-based BKGs, literature-based BKGs, and mixed BKGs. The database-based BGKs (e.g., \emph{Hetionet} \citep{Himmelstein_Daniel_Scott_2017}, \emph{BioKG} \citep{Walsh_Brian_2020}, \emph{CBKH} \citep{Su_Chang_2021}) are constructed with biomedical data and their relations stored in existing biological databases. The literature-based BKGs are obtained by leveraging Natural Language Processing (NLP) techniques to extract the semantic relations from a large amount of available biomedical literature and electronic health record (EHR) data. They are more disease-specific if the text sources are only focused on specific diseases. For example, \citet{Zhang_Rui_2021} recently constructed a literature knowledge graph via the collections of Covid-19 literature from the Semantic MEDLINE Database (SemMedDB) \citep{Kilicoglu_Halil_2012} as well as  other Covid-19 datasets for Covid-19 drug predictions. The mixed BKGs (e.g., \emph{CKG} \citep{Santos_Alberto_2020}, \emph{DRKG} \citep{Ioannidis_Vassilis_N_2020}) are generated by combining the knowledge collected from both existing biological databases and available text-based data.

\paragraph{KG-based computational methods for drug repurposing}
KG-based computational approaches have been widely used in drug repurposing in recent years \citep{Zhang_Rui_2021,Al_Saleem_Jacob_2021,Yan_Vincent_K_C_2021}. These approaches treat drug repurposing as a link prediction task by applying different machine learning models to biomedical knowledge graphs. These models can mainly be classified into three groups: tensor-factorization-based models, translation-based models, and neural network-based models \citep{Zeng_Xiangxiang_2022}. The tensor-factorization-based models \citep{Nickel_Maximilian_2011,Yang_Bishan_2014} consider a KG as a cubic tensor in which three dimensions respectively represent head entity, tail entity, and their relations and then leverage tensor-factorization methods \citep{Rabanser_Stephan_2017} to recover the missing relations in KGs. The translation-based models \citep{Bordes_Antoine_2013,Yankai_Lin_2015,Sun_Zhiqing_2019} regard the semantic relations in KGs as a "translation" process, that is, given a triple (head entity, relation, tail entity) the head entity can be translated to the tail entity via the relation-specific "translation". The neural network-based models \citep{Vassilis_N_Ioannidis_2020,Dettmers_Tim_2017} leverage graph convolutions or 2D convolutions to learn the underlying relationships between two entities.

\paragraph{Explanation of drug repurposing}
Although drug repurposing prediction can accelerate the process of drug discovery, one of the greatest concerns in computational drug repurposing is the lack of biologically reasonable explanations, which hinders its wide adoption in clinical settings. Currently, there are few computational models designed for drug repurposing explanations. A common and intuitive explanation for drug repurposing leverages the semantic KG-based paths between given drug-disease pairs. To the best of our knowledge, existing models either utilize known statistics scores or use reinforcement learning models to find the most likely KG-based paths for explaining drug repurposing. For example, \citet{Sosa_Daniel_N_2020} applied a graph embedding model UKGE \citep{Chen_Xuelu_2019}, which utilizes the confidence scores of relation edges in a literature-based KG \emph{GNBR}, to identify new indications of drugs for rare diseases and then explain the results via the highest-ranking paths based on confidence scores. \citet{Liu_Yushan_2021} developed a Reinforcement Learning-based model “PoLo” that utilizes the biological meta-paths identified in \citet{Himmelstein_Daniel_Scott_2017} via a computationally inefficient method "DWPC" to supervise path searching for drug repurposing.

\section{Knowledge Graph}\label{sec:kg}

To accommodate biomedical-reasonable predictions of drug repurposing and mechanism of action, the ideal knowledge graph should accurately represent comprehensive and diverse interactions among known biological entities. Thus, we utilize the canonicalized version of the Reasoning Tool X Knowledge Graph 2 (\emph{RTX-KG2c}) \citep{Wood_E_C_2022}, one of the largest open-source biomedical knowledge graph (BKG) that has been widely used in the Biomedical Data Translator Project \citep{Consortium_2019,Biomedical_Data_Translator_Consortium_2019}. Compared to other commonly used open-source BKGs mentioned in Sec. \ref{sec:related_work}, \emph{RTX-KG2c} is biolink-model-based \footnote{Biolink model (\url{https://biolink.github.io/biolink-model}) is a standardized BKG ontology framework} standardized \citep{unni2022biolink} and regularly-updated.
The version 2.7.3 of \emph{RTX-KG2c} that we use aggregates data from 70 public knowledge sources into a large graph where all biological concepts (e.g., "ibuprofen") are represented as vertices and all concept-predicates-concepts (e.g., "ibuprofen - increases activity of - GP1BA gene") are presented as edges.
We customized it for drug repurposing purpose with four principles: 1). excluding the nodes whose categories are irrelevant to drug repurposing explanation (e.g., ”GeographicLocation” and ”Device”); 2). filtering out the low-quality edges based on our criteria; 3). removing the hierarchically redundant edges; 4). excluding all drug-disease edges (see more details in Appx. \ref{subsec:kg_preprocessing}). After this, 3,659,165 nodes with 33 distinct categories and 18,291,237 edges with 74 distinct types are left in our customized knowledge graph, which is used for downstream model training

\section{Proposed Method}\label{sec:method}

\begin{figure}[t]
		\centering
		\begin{tikzpicture}
		\draw (0,0 ) node[inner sep=0] {\includegraphics[width=1\columnwidth, trim={1.7cm 6.55cm 0.7cm 0.92cm}, clip]{./Figures/main_figure2.pdf}};
		\end{tikzpicture}
		\caption{Illustration of the drug repurposing prediction model (left, see Sec. \ref{subsec:gnn} for more details) and mechanism of action (MOA) prediction model (right, see Sec. \ref{subsec:moa_prediction} for more details).}
		\label{fig:figure2}
\end{figure}

In this section, we describe our proposed model framework (see Figure~\ref{fig:figure2}) that predicts drug repurposing and identifies the possible KG-based mechanism of action (MOA) paths as its biological explanations. The proposed framework consists of two components: drug repurposing prediction with GraphSAGE \citep{Hamilton_William_L_2017} plus a Random Forest (Sec. \ref{subsec:gnn}) and MOA prediction with an adversarial actor-critic reinforcement learning (RL) (Sec. \ref{subsec:moa_prediction}).
\subsection{Notation}

Let $\mathcal{G} = \{\mathcal{V}, \mathcal{E}\}$ be a directed knowledge graph, where each node $v \in \mathcal{V}$ represents a particular biomedical entity (e.g., a specific drug, disease, or gene, etc.) and each edge $e \in \mathcal{E}$ represents a biomedical relationship (e.g., \textit{interacts-with}). We use $\mathcal{V}^{\mathrm{drug}}$ to represent all the drug nodes and $\mathcal{V}^{\mathrm{disease}}$ to represent all the disease nodes (defined in item 4 in Appx. \ref{subsec:kg_preprocessing}).

\subsection{Drug Repurposing Prediction}\label{subsec:gnn}

Drug repurposing aims to identify new indications of approved or experimental drugs. We solve it as a link prediction problem on the knowledge graph $\mathcal{G}$. Specifically, given any drug-disease pair $(v_i, v_j)$ where $v_i \in \mathcal{V}^{\mathrm{drug}}$ and $v_j \in \mathcal{V}^{\mathrm{disease}}$, we predict the probability that drug $i$ can be used to treat disease $j$.
We first use GraphSAGE to calculate the embedding for each node. Ideally, the node embeddings should contain two kinds of information: node attributes and node neighborhoods. To capture the neighborhood information, we optimize GraphSAGE to encourage neighbor nodes to have similar embeddings and non-neighbor nodes to have distinct embeddings. Specifically, we perform random walks for each node to collect its neighborhood information and train the model to maximize a node's similarity with its neighbor nodes. For a node $u$, the loss is calculated as:
\begin{equation}
J_{\mathcal{G}}(\bm{z}_u) = - \log (\sigma(\bm{z}_u^{\top}\bm{z}_v)) - k \cdot \mathbb{E}_{v_n \sim P_n(v)} \log (\sigma(-\bm{z}_u^{\top}\bm{z}_{v_n}))
\end{equation}
where $\bm{z}_u, \bm{z}_v$ are respectively the embeddings of nodes $u,v$, $\sigma$ is the sigmoid function, $v$ is a node that co-occurs with $u$ in fixed-length random walks, $P_n$ represents negative sampling distribution, and $k$ indicates the number of negative samples (nodes not in $u$'s fixed-length neighborhood).

To capture the node attributes information, we utilize the PubMedBERT model \citep{Gu_Yu_2022}, a pre-trained language model designed for biomedical texts, to generate a node attribute embedding for each node based on the concatenation of the node's name and category. We further compress the embeddings to 100 dimensions with Principal Components Analysis (PCA) to reduce memory usage and use them as the initial node feature for GraphSAGE. In this way, the final GraphSAGE embedding of each node contains the information of both graph topology and node attributes.
We concatenate the GraphSAGE embeddings of drug-disease pairs and use them as input into a Random Forest model to classify each drug-disease pair into one of the "not treat", "treat" and "unknown" classes. We obtain "treat" and "not treat" drug-disease pairs from biomedical datasets (see Appx. \ref{subsubsec:training_data_collection}) and generate "unknown" drug-disease pairs through negative sampling \citep{Mikolov_Tomas_2013}.

\subsection{Mechanism of Action (MOA) Prediction}\label{subsec:moa_prediction}

When potential indications of a given drug are identified by the proposed drug repurposing model, a natural yet important question is: can we biologically explain the predictions?
We solve this by employing a reinforcement learning (RL) model to predict the KG-based MOA paths, which are essentially the paths on the knowledge graph from drug nodes to disease nodes. These KG-based MOA paths can semantically describe an abstract biological process of how a drug treats a disease.

\subsubsection{Demonstration Paths}\label{subsubsec:dp}

To encourage the RL agent to terminate the path searching at the expected diseases through a biologically reasonable path, we leverage so-called demonstration paths, a set of biologically likely paths (e.g., drug1-gene1-protein3-disease1), that explains the underlying reasons of why a drug can treat a disease. We extract the demonstration paths by using the known drug-target interactions collected from two curated biomedical data sources: DrugBank (v5.1) and Molecular Data Provider (v1.2) along with a literature-based source that utilizes the adjusted Normalized Google Distance (NGD) applied to concepts appearing in PubMed publication abstracts (see Appx. \ref{subsec:dp_extraction_preprocessing} for more details).

\subsubsection{ADAC-based Reinforcement Learning Model}\label{subsubsec:rl}
We formulate the MOA prediction as a path-finding problem and adapt the ADversarial Actor-Critic (ADAC) Reinforcement Learning model to solve it.
The reinforcement learning is defined as a Markov Decision Process (MDP) which contains:

\textbf{States}: Each state is defined as $s_{t}=\big(v_{drug}, v_{t}, (v_{t-1}, e_{t}), \dots, (v_{t-K}, e_{t-(K-1)}) \big)$ where $v_{drug} \in \mathcal{V}^{\mathrm{drug}}$ is a given starting drug node; $v_{t} \in \mathcal{V}$ represents the node where the agent locates at time $t$; the tuple $\big(v_{t-K}, e_{t-(K-1)}\big)$ represents the previous $K$th node and $(K-1)$th predicate. For the initial state $s_{0}$, the previous nodes and predicates are substituted by a special dummy node and predicate. We concatenate the embedding of all nodes and predicates of $s_{t}$ to get the state embedding, where the node embeddings are node attribute embeddings generated with the PubMedBERT model (see Sec. \ref{subsec:gnn}) and the predicate embeddings employ one-hot vectors.

\textbf{Actions}: The action space $A_{t}$ of each node $v_{t}$ includes a self-loop action and the actions to reach its outgoing neighbors in the graph $\mathcal{G}$. Due to memory limitation and extremely large outdegree of certain nodes in the knowledge graph, we prune the neighbor actions based on the PageRank scores if a node has neighbors more than a certain threshold. Specifically, we let $A_{t}=\mathrm{vcat}(\vec{a}_{self},\vec{a}_{1},\dots, \vec{a}_{k},\dots,\vec{a}_{n_{v_t}})$ where $\mathrm{vcat}$ indicates the vertical concatenation and $n_{v_t}$ is out-degree of node $v_t$, $v_t \in \mathcal{V}$. For each action $a_{t}=(v_{t},e_{t})$, we concatenate its node and predicate embeddings to obtain action embeddings. We learn two embedding matrices $E^{N_\mathrm{n} \times \mathrm{d}}$ and $E^{N_\mathrm{p} \times \mathrm{d}}$ respectively for nodes and predicates\footnote{Each sub-network uses separate embedding matrices.}, where $d$ represents the embedding dimension, $N_\mathrm{n}$ represents the number of nodes, and $N_\mathrm{p}$ represents the number of predicate categories.

\textbf{Rewards}: During the path searching process, the agent only receives a terminal reward $R_{e,T}$ from the environment ($R_{e,t}=0, \forall t<T$). Let $v_T$ be the last node of the path, and $\mathcal{N}_{drug}$ be the known diseases that drug $v_{drug}$ can treat, the terminal reward is calculated by using the proposed drug repurposing model (Sec. \ref{subsec:gnn}) via:
$$
  R_{e,T}=\left\{
    \begin{array}{ll}
      1, & \mbox{if $v_T \in \mathcal{N}_{drug}$}.\\
      p_{treat}, & \mbox{if $v_{T} \notin \mathcal{N}_{drug} ; v_{T} \in \mathcal{V}^{\mathrm{disease}}$ and $f(v_{drug}, v_{T})$ is predicted as "treat"}.\\
      0, & \mbox{if $v_{T} \notin \mathcal{N}_{drug} ; v_{T} \in \mathcal{V}^{\mathrm{disease}}$ and $f(v_{drug}, v_{T})$ is not predicted as "treat"}.\\
      -1, & \mbox{if $v_{T} \notin \mathcal{V}^{\mathrm{disease}}$}.
    \end{array}
  \right.
$$
where $p_{treat}$ is the "treat" class probability predicted by the drug repurposing model $f$.

The ADAC-based RL model consists of four sub-networks that share the same model architecture $\mathrm{MLP}^{i}$ (note that $i$ represents the id of each sub-network described later, such as $a$ for \textit{actor network}, $c$ for \textit{critic network}, etc.) but with different parameters:
\begin{equation}
\mathrm{MLP}^{i}(X) = \bm{f}(\bm{f}(XW^i_{1}+b^i_{1})W^i_{2}+b^i_{2})W^i_{3}+b^i_{3}
\end{equation}
where \{${W^i_{1}, W^i_{2}, W^i_{3}, b^i_{1}, b^i_{2}, b^i_{3}}$\} are the parameters and biases of linear transformations, $\bm{f}$ represents a batch normalization layer followed by an ELU activation function.

\textbf{Actor network}: The actor network learns a path-finding policy $\pi_{\theta}$ (note that $\theta$ represents all parameters of actor network) to guide the agent to choose an action $a_{t}$ from the action space $A_{t}$ based on current state $s_{t}$:
\begin{equation}
\pi_{\theta}(a_{t}|s_{t}, A_{t}) = \mathrm{softmax}(A_{t} \odot \mathrm{MLP}^{a}(s_t))
\end{equation}
where $\odot$ represents the dot product. Here, $\pi_{\theta}(a_{t}|s_{t}, A_{t})$ represents the probability of choosing action $a_{t}$ at time $t$ from the action space $A_{t}$ given the state $s_{t}$.

\textbf{Critic network}: The critic network \citep{Lillicrap_Timothy_P_2015} estimates the expected reward $Q_{\phi}(s_{t},a_{t})$ (note that $\phi$ represents all parameters of critic network) if the agent takes the action $a_{t}$ at state $s_{t}$ by:
\begin{equation}
Q_{\phi}(s_{t},a_{t}) =  \mathrm{MLP}^{c}(s_t) \odot a_{t}
\end{equation}
\textbf{Path discriminator network}: Since the RL agent only receives a terminal reward indicating whether it reaches an expected target, to encourage the agent to find biologically reasonable paths and provide intermediate rewards, we further guide it with demonstration paths.
This network is essentially a binary classifier that distinguishes whether a path segment $(s_{t}, a_{t})$ is from demonstration paths or generated by the actor network. We treat all the known demonstration path segments $(s_{t}^{D}, a_{t}^{D})$ as positive samples and all actor-generated non-demonstration path segments $(s_{t}, a_{t})$ as negative samples. The path discriminator network $D_{p}(s_{t},a_{t}) = \mathrm{sigmoid}(\mathrm{MLP}^{p}(s_{t} \oplus a_{t}))$, where $\oplus$ represents the concatenation operator, is optimized with:
\begin{equation}
L_{p} = - \mathbb{E}_{(s,a) \sim P_{D}}[\log(D_{p}(s,a))] - \mathbb{E}_{(s,a) \sim P_{A}}[\log(1-D_{p}(s,a))]
\end{equation}
where $P_{D}$ and $P_{A}$ respectively represent the demonstration path segment distribution and the actor-generated non-demonstration path segment distribution. Based on the probability $D_{p}(s_{t},a_{t})$, the path-discriminator-based intermediate reward $R_{p,t}$ is calculated as:
\begin{equation}
R_{p,t} = \log(D_{p}(s_{t},a_{t})) - \log(1-D_{p}(s_{t},a_{t})).
\end{equation}
\textbf{Meta-Path discriminator network}: Similar to the path discriminator, this network aims to judge whether the meta-path of the actor-generated paths is similar to that of demonstration paths. The meta-path is the path of node categories (e.g., ['Drug'→'Gene'→'BiologicalProcess'→'Disease']). Similarly, the meta-path discriminator $D_{m}(M) = \mathrm{sigmoid}(\mathrm{MLP}^{m}(M))$ is also a binary classifier where the meta-paths of demonstration paths are treated as positive samples while others are negative samples. We optimize it with the following loss:
\begin{equation}
L_{m}= - \mathbb{E}_{M \sim P^M_{D}}[\log(D_{m}(M))] - \mathbb{E}_{M \sim P^M_{A}}[\log(1-D_{m}(M))]
\end{equation}
where $M$ is a meta-path embedding vector defined as the concatenation of learned category embeddings of all nodes that appear in the path; $P^M_{D}$ and $P^M_{A}$ respectively represent the demonstration meta-path distribution and the actor-generated non-demonstration meta-path distribution.
The intermediate reward $R_{m,t}$ generated by the meta-path discriminator is calculated by:
\begin{equation}
R_{m,t} = \log(D_{m}(M)) - \log(1-D_{m}(M)).
\end{equation}
The final intermediate reward $R_{t}$ at time $t$ is then calculated as:
\begin{equation}
\label{func:reward}
R_{t} = {\alpha}_{p}R_{p,t}+{\alpha}_{m}R_{m,t}+(1-{\alpha}_{p}-{\alpha}_{m}){\gamma}^{T-t}R_{e,T}
\end{equation}
where ${\alpha}_{p} \in [0,1]$ and ${\alpha}_{m} \in [0, 1-{\alpha}_{p}]$ are hyperparameters,  $\gamma$ is the decay coefficient, and $R_{e,T}$ is defined in the "Rewards" section above.
To optimize the critic network, we minimize the Temporal Difference (TD) error \citep{Sutton_Richard_S_1988} with loss:
\begin{equation}
L_{c}= \mathrm{TD}^{2} = [(R_{t} + Q_{\phi}(s_{t+1},a_{t+1})) - Q_{\phi}(s_{t},a_{t})]^{2}.
\end{equation}
Since the goal of the actor network is to achieve the largest expected reward by learning an optimal actor policy, we optimize the actor network by maximizing $J({\theta}) = \mathbb{E}_{a \sim {\pi}_{\theta}}[Q_{\phi}(s_{t},a)]$. We use the REINFORCE algorithm \citep{Williams_Ronald_J_1992} to optimize the parameters. To encourage more diverse exploration in finding paths, we use the entropy of ${\pi}_{\theta}$ as a regularization term and optimize the actor network with the following stochastic gradient of the loss function $L_a$:
\begin{equation}
{\nabla}_{\theta}L_{a} = - {\nabla}_{\theta}J({\theta})= - \mathbb{E}_{{\pi}_{\theta}}[{\nabla}_{\theta}\mathrm{TD} \log{\pi}_{\theta}(a_{t}|s_{t})] - {\alpha}{\nabla}_{\theta}\mathrm{entropy}({\pi}_{\theta})
\end{equation}
where ${\pi}_{\theta}$ is the action probability distribution based on the actor policy and $\alpha$ is the entropy weight.

Following \citet{Zhao_Kangzhi_2020}, we train the ADAC-based RL model in a multi-stage way. First, we initialized the actor network using the behavior cloning method \citep{A_PomerleauDean_1991} in which the training set of demonstration paths is used to guide the sampling of the agent with Mean Square Error (MSE) loss. Then, in the first $z$ epochs, we freeze the parameters of the actor network and the critic network and respectively train the path discriminator network and meta-path discriminator network by minimizing $L_{p}$ and $L_{m}$. After $z$ epochs, we unfreeze the actor network and the critic network and optimize them together by minimizing a joint loss $L_{joint} = L_{a} + L_{c}$.

\section{Experiments}\label{sec:experiment}
In this section, we introduce training data used in this study (Sec. \ref{subsec:training_data}), evaluation setup (Sec. \ref{subsec:evaluation_setup}) as well as evaluation results (Sec. \ref{subsec:dr_prediction_evaluation} \& \ref{subsec:moa_prediction_evaluation}) respectively for drug repurposing prediction and MOA prediction. Implementation details of the proposed models are presented in Appx. \ref{sec:implementation}.

\subsection{Training Data}\label{subsec:training_data}

Training data of drug repurposing model consists of drug-disease pairs in three categories: "treat", "not treat", and "unknown". For the first two categories, we collect data from four different data sources (Appx. \ref{subsubsec:training_data_collection}). The raw data are pre-processed and separated into training, validation, and test sets with ratios of [0.8, 0.1, 0.1] for downstream model training (Appx. \ref{subsubsec:training_data_preprocessing}).

\subsection{Evaluation Metrics and Methods}\label{subsec:evaluation_setup}
We evaluate the proposed models on two types of tasks: \ul{predicting drug-disease "treat" probability} (i.e., drug repurposing prediction) and \ul{identifying biologically reasonable KG-based MOA paths from all candidates} (i.e., MOA prediction). These two tasks are evaluated based on classification accuracy-based metrics (e.g., accuracy, macro f1 score) and ranking-based metrics (e.g., mean percentile rank (MPR), mean reciprocal rank (MRR), and Hit@K) defined as follows:
\begin{equation}
\mbox{MPR} = \frac{1}{|PR|} \sum_{pr \in PR}{pr} \;\;\;\;\;\;\;\;
\mbox{MRR} = \frac{1}{|R|} \sum_{r \in R}{r^{-1}} \;\;\;\;\;\;\;\;
\mbox{Hit@K} = \frac{1}{|R|} \sum_{r \in R}{|r \leq k|},
\end{equation}
where $PR$ and $R$ are respectively a list of percentile ranks and a list of ranks of all true positive drug-disease pairs ("treat" category) in the evaluation dataset (combination of validation and test).

\textbf{Drug repurposing prediction}:
Due to the extremely long running time of certain baselines (e.g., GraphSAGE+SVM), we calculate the metrics MRR and Hit@K based on the rank of each true positive drug-disease pair among 1000 random drug-disease pairs (500 with drug id replacement and 500 with disease id replacement) and compare them across all models. We also show the comparison results with MRR and Hit@K based on "all nodes" replacement for the models that can be implemented in a reasonable time (Appx. \ref{subsec:gnn_eval_allnodes}). In addition, since our drug repurposing model does 3-class classification while other baselines do 2-class classification, for a fair comparison, we re-calculate "accuracy" and "macro f1 score" for our drug repurposing model by excluding the "unknown" class.

\textbf{MOA prediction}:
For the evaluation of MOA prediction (Sec. \ref{subsec:moa_prediction}), we use the DrugMechDB \citep{Mayers_Michael_2020}, an expert-curated path-based drug MOA database, to obtain the verified MOA paths as ground-truth data. We match each biological concepts in these verified MOA paths to the biological entities used in our customized knowledge graph (Sec. \ref{sec:kg}) and generate the corresponding KG-based MOA paths (Appx. \ref{subsec:DrugMechDB_path_mapping}). The matched KG-based MOA path is defined as a 3-hop KG-based path of which all four nodes appear in the verified MOA paths. For each drug-disease pair not in the training set, we calculate the path scores for all 3-hop KG paths between drug and disease with the path-finding policy learned from the ADAC-based RL model by using equation:
$\mbox{path score} = \sum_{i=1}^{k}{{\delta}^{i-1} \times \log(P_{i} \times N_{i})}$,
where $k$ is the number of hops in this path; $\delta$ is a decay coefficient (we set it to 0.9 in this study); $P_{i}$ represents the probability of choosing action $a_{i}$ in the $i$th hop following this path based on the trained ADAC-based RL model; $N_{i}$ is the number of possible actions in the $i$th hop. With these path scores, we obtain the ranks of the matched KG-based MOA paths and calculate the ranking-based metrics mentioned above. For those drug-disease pairs with multiple KG-based MOA paths, we use the highest ranks of their paths. We compare our ADAC-based RL model with the baseline models based on these metrics. In addition, we further perform some case studies to evaluate the effectiveness of the proposed method and present them in Appx. \ref{subsec:case_study}.

\subsection{Drug Repurposing Prediction Evaluation}\label{subsec:dr_prediction_evaluation}
\textbf{Baselines}:
We compare our proposed drug repurposing prediction model against several state-of-the-art (SOTA) KG-based  models and the variants of our proposed model.  \texttt{TransE} \citep{Bordes_Antoine_2013},  \texttt{TransR} \citep{Yankai_Lin_2015},  \texttt{RotatE} \citep{Sun_Zhiqing_2019} are the translation-distance-based models that regard a relation as a translation/rotation from a head entity to a tail entity.  \texttt{DistMult} \citep{Yang_Bishan_2014} is a bilinear model that measures the latent semantic similarity of a triple with a trilinear dot product. \texttt{ComplEx} \citep{Trouillon_Theo_2016} and \texttt{ANALOGY} \citep{Hanxiao_Liu_2017} are the extensions of \texttt{DistMult} that consider more complex relations. \texttt{SimpLE} \citep{Kazemi_Seyed_Mehran_2018} is a tensor-factorization-based model to learn the semantic relation of a triple. \texttt{GAT} \citep{velickovic_2018} is a popular graph attention model. Besides these SOTA baselines, we use the variants of our proposed model (e.g., pure GraphSAGE for link prediction \texttt{GraphSAGE-link}, the \texttt{GraphSAGE+logistic} model, the \texttt{GraphSAGE+SVM} model, and the \texttt{GraphSAGE+RF} model for 2-class classification\footnote{only considers true positive and true negative.}). Implementation details of these baselines are presented in Appx. \ref{sec:implementation}.

\textbf{Evaluation results}:
Table~\ref{tab:table1} shows the evaluation results of the proposed drug repurposing model and other baselines based on the metrics described in Sec. \ref{subsec:evaluation_setup}. As shown in the table, on the one hand, the proposed model outperforms almost all baselines in classification-based metrics even though it is a bit worse than the GAT model, which indicates its effectiveness in classifying known "treat" and "not treat" drug-disease pairs based on their attribute and neighborhood information on the knowledge graph (note that we conduct an ablation experiment to show the effectiveness of node attribute embeddings in Appx. \ref{subsec:abation_node_attribute_embs}). On the other hand, the best performance of our model in ranking-based metrics shows its capability in identifying new indications of existing drugs out of all the possible drug-disease pairs with relatively low false positives, which is of great importance for guiding clinical research.
Besides, comparing \texttt{2-class GraphSAGE+RF} with the vanilla GraphSAGE model (e.g., \texttt{GraphSAGE-link}), the results demonstrate the effectiveness of the Random Forest model over a neural network classifier in this task. Comparison between 2-class and 3-class GraphSAGE+RF models indicates the importance of using negative sampling to generate "unknown" drug-disease pairs for model training. With the "unknown"  drug-disease pairs, the \texttt{3-class GraphSAGE+RF} model achieves significant improvement in ranking-based metrics, which is essential when applying to real-world drug repurposing because it can reduce the false positives.
\begin{table}[t!]
    \caption{Comparing evaluation results between the proposed drug repurposing model with baselines. The values with $^{\ast}$ inside the parenthesis are the results after excluding the "unknown" category.}
	\label{tab:table1}
	\centering
	\begin{tabular}{lcccccc}
	    \toprule
		\textbf{Model} & \textbf{Accuracy} & \textbf{Macro F1 score} & \textbf{MRR} & \textbf{Hit@1} & \textbf{Hit@3} & \textbf{Hit@5} \\
		\midrule
		TransE & 0.706 & 0.706 & 0.275 & 0.111 & 0.295 & 0.446 \\
		TransR & 0.858 & 0.855 & 0.307 & 0.130 & 0.350 & 0.517 \\
		RotatE & 0.707 & 0.707 & 0.255 & 0.077 & 0.287 & 0.454 \\
		DistMult & 0.555 & 0.495 & 0.172 & 0.040 & 0.144 & 0.258 \\
		ComplEx & 0.624 & 0.456 & 0.133 & 0.023 & 0.106 & 0.196 \\
		ANALOGY & 0.597 & 0.467 & 0.179 & 0.045 & 0.147 & 0.272 \\
		SimplE & 0.600 & 0.475 & 0.161 & 0.037 & 0.135 & 0.236 \\
		GAT & \textbf{0.932} & \textbf{0.929} & 0.002 & 0 & 0 & 0.001 \\
		GraphSAGE-link & 0.919 & 0.915 & 0.002 & 0 & 0 & 0 \\
	    GraphSAGE+logistic & 0.787 & 0.779 & 0.002 & 0 & 0 & 0 \\
	    GraphSAGE+SVM & 0.809 & 0.795 & 0.002 & 0 & 0 & 0 \\
	    2-class GraphSAGE+RF & 0.929 & 0.925 & 0.271 & 0.176 & 0.311 & 0.386  \\
		\midrule
		3-class GraphSAGE+RF (ours) & 0.934 (0.929$^{\ast}$) & 0.922 (0.925$^{\ast}$)  & \textbf{0.360} & \textbf{0.211} & \textbf{0.410} & \textbf{0.524}  \\
		\bottomrule
	\end{tabular}
\end{table}
\subsection{Mechanism of Action (MOA) Prediction Evaluation}\label{subsec:moa_prediction_evaluation}
\textbf{Baselines}:
For a fair comparison, we choose the \texttt{MultiHop} model \citep{Lin_Xi_Victoria_2018} (Implementation details in Appx. \ref{sec:implementation}) as a baseline since it allows using a self-defined reward shaping strategy in its reward function as what we do in our proposed ADAC-based RL model. We don't compare with other existing models that either require pre-known edge importance information (see Sec. \ref{sec:related_work}) or cannot be trained within a reasonable time (e.g., within two weeks) on the customized knowledge graph. To show the importance of demonstration paths, we also compare with an ablated version of the proposed model that does not take advantage of the demonstration paths by setting $\alpha_p$ and $\alpha_m$ in Function \ref{func:reward} as $0$.

\textbf{Evaluation results}:
Table~\ref{tab:table2} shows the evaluation results. Although all the models receive the same terminal rewards from the environment, the proposed ADAC-based RL model achieves significantly better performance in identifying biologically reasonable KG-based MOA paths than other two baselines.
Comparison between the proposed model and the ADAD-based RL model without demonstration paths (i.e., ADAC RL w/o DP) further illustrates the great effectiveness in using demonstration paths to guide the path-finding process. Due to the massive searching space and sparse rewards, the RL agent often fails to find reasonable paths out of plenty of possible choices, while our model, with the intermediate guidance provided by the demonstration path, is able to identify the most biologically reasonable choices at each time step with a much higher probability. Moreover, we perform two specific case studies (see Appx. \ref{subsec:case_study}) where we further show the top-rank MOA paths predicted by our model are able to identify key molecules in real drug action regulatory networks.
\begin{table}[h!]
\centering
\caption{Evaluation results on MOA prediction on MultiHop model and ADAC-based RL model with and without demonstration paths (ADAC RL w/o DP).}\label{tab:table2}
{
	\begin{tabular}{lccccccc}
	    \toprule
		\textbf{Model} & \textbf{MPR} & \textbf{MRR} & \textbf{Hit@1} & \textbf{Hit@10} & \textbf{Hit@50} & \textbf{Hit@100} & \textbf{Hit@500} \\
		\midrule
		MultiHop & 61.541\% & 0.026 & 0.012 & 0.043 & 0.081 & 0.124 & 0.360 \\
		ADAC RL w/o DP & 73.281\% & 0.022 & 0.012 & 0.025 & 0.112 & 0.186 & 0.447 \\
		\midrule
		ADAC RL w/ DP (ours) & \textbf{ 94.410\%} & \textbf{0.123} & \textbf{0.062} & \textbf{0.242} & \textbf{0.509} & \textbf{0.640} & \textbf{0.857} \\
		\bottomrule
	\end{tabular}
}
\end{table}

\section{Conclusion}\label{sec:conclusion}
In this paper, we propose a computational drug repurposing model framework that predicts not only the treatment probabilities of drug-disease pairs but also the KG-based mechanism of action (MOA) paths as biological explanations for drug repurposing. We apply the proposed model framework to a large and standardized biomedical knowledge graph and conduct experiments to compare our framework with extensive state-of-the-art KG-based models. The results show that our proposed drug repurposing model achieves much better predictive performance in a comprehensive evaluation. We also show that by leveraging the proposed knowledge-and-publication-based demonstration paths to provide intermediate guidance during the path-finding process, our model can effectively find biologically reasonable paths out of a large amount of possible choices. Furthermore, two specific case studies are performed to illustrate the effectiveness of our models in identifying some key biological molecules in real drug action regulatory networks. We believe our proposed framework can effectively reduce the "black-box" concerns and increase prediction confidence for drug repurposing based on its predicted path-based explanations, which further accelerates the process of drug discovery for emerging diseases. Despite the comparably good performance of our models, some limitations remain to be solved. For example, our MOA prediction is based on an assumption that a repurposed drug follows similar molecular mechanisms as the known MOAs to treat different diseases, which needs further biologically experimental verification.

\section{Ethics Statement}\label{sec:ethics_discussion}
Our proposed method aims to provide a new model framework for drug repurposing prediction and explanation. The potential negative societal impacts might include the risk of misusing the drug repurposing capability (e.g., misleading patients to take drugs inappropriately). Although our model can provide certain suggestions for the new indications of drugs, it is mainly designed for assisting medical practitioners (e.g., doctors and licensed practitioners), who have professional training or knowledge to assess the accuracy of suggestions provided by our model. Therefore, the application of this method should be limited to drug development in pharmaceutical companies or pre-clinical drug research in relevant academic or medical organizations, but not be directly accessible to the public. Since drugs treat diseases in complex ways, there is still much unknown knowledge involved, and therefore any suggestions given by this proposed method should require further clinical testing but not be used to directly diagnose, prevent, or treat any diseases without FDA approval. In addition, although the data used in this study (including the data in the biomedical knowledge graph \emph{RTX-KG2c}) are from multiple data sources, they are all collected from publicly accessible databases or datasets or APIs with free licenses (except for DrugBank data that we have requested and obtained its non-commercial license). Note that the licenses of all data are shown in Table~\ref{tab:license_of_assets}. Therefore, they have excluded any personal or sensitive information. All code used in this study is originally created by ourselves or modified from the code of previous researchers (note that we illustrate this at the top of those scripts), and allowed to be disclosed with the permission of its original authors.

\begin{table}[ht!]
\centering
\caption{Summary of licenses of all data we used.}\label{tab:license_of_assets}
\scriptsize
{
	\begin{tabular}{lc}
	    \toprule
		\textbf{Existing assets} & \textbf{Licenses} \\
		\midrule
		RTX-KG2c (v2.7.3) \url{https://github.com/RTXteam/RTX-KG2} & MIT license, CC-BY 4.0 license \\
		DrugMechDB data (v1.1.0) \url{https://github.com/SuLab/DrugMechDB} & MIT license \\
        DrugBank data (v5.1) \url{https://go.drugbank.com} & DrugBank academic license \\
        Molecular Data Provider (v1.2) \url{https://github.com/broadinstitute/molecular-data-provider} & MIT license \\
        MyChem data (v1) \url{https://mychem.info} & Apache-2.0 license \\
        SemMedDB data (semmedVER43\_R) \url{https://lhncbc.nlm.nih.gov/ii/tools/SemRep_SemMedDB_SKR.html} & UMLS Metathesaurus license \\
        NDF-RT data (v2018AA) \url{https://bioportal.bioontology.org/ontologies/NDFRT} & UMLS Metathesaurus license \\
        RepoDB data \url{https://unmtid-shinyapps.net/shiny/repodb} & CC-BY 4.0 license \\
		\bottomrule
	\end{tabular}
}
\end{table}

\clearpage
\bibliographystyle{unsrtnat}

Quantitative Methods (q-bio.QM); Machine Learning (cs.LG)