\section{Introduction}
Named Entity Recognition (NER) is a core task of information extraction which enables many downstream tasks such as question answering \cite{molla2006named, lee-etal-2020-answering}, machine translation \cite{babych2003improving, ugawa2018neural}, entity linking \cite{martins-etal-2019-joint}, and event extraction \cite{wadden2019entity}.
In recent years, with a proliferation of pre-trained language models \cite{devlin2019bert, liu2019roberta, lee2020biobert}, the performance of benchmark datasets has reached a plateau.

Growing interests in the current performance raise several questions about generalization behavior such as: breaking down a holistic performance into fine-grained categories \cite{fu-etal-2020-interpretable, fu2020rethinking}, splitting benchmark datasets based on seen and unseen mentions \cite{lin2020rigourous}, and partitioning unseen mentions into synonym and concept in the biomedical domain \cite{kim2021your}.
Furthermore, it has been proved that pre-trained language models have been highly overestimated by focusing on dataset biases to solve the particular benchmark datasets~\cite{jia2017adversarial, gururangan2018annotation, nie2020adversarial}.

However, exploiting biases without learning the underlying task hinders the generalization capacity of pre-trained language models, thus numerous studies have suggested debiasing methods to prevent overfitting on superficial cues in training dataset \cite{gururangan2018annotation, he2019unlearn, clark2019don, utama2020towards, utama2020mind, ko2020look, kim2021your, du2021interpreting}.
Despite such commendable efforts, there is an issue while designing a bias-only model.
First, \citet{clark2019don} introduced a \textit{produce-of-expert} style to scale down the gradients of biased examples, but it ignored a class prior due to the difficulty of modeling and class of targeted dataset is uniformly distributed.
Furthermore, in Biomedical Named Entity Recognition (BioNER), \citet{kim2021your} enhanced generalizability via this debiasing framework that replaces bias-only model as a class distribution given word frequency.
Nonetheless, it is impossible to apply informative statistics when the train and test dataset is dissimilar.

In this paper, we argue that the bias-only model should reflect informative statistics and design a class prior.
To this end, we substitute the bias-only model with a statistic of Pointwise Mutual Information (PMI).
When neural models fail to reflect informative statistics because the train and test dataset are dissimilar in a word unit~\cite{leevy2018survey, li2020dice}, we apply the PMI based on a subword unit (Section~\ref{sec:subword}).
This helps because of an issue called the \textit{Out-Of-Vocabulary} (OOV) problem: by default, a word typically does not observe in the training dataset but exists in the evaluation dataset.
Furthermore, PMI was used to measure a correlation score between word and class in a particular dataset~\cite{gururangan2018annotation}.
By leveraging this, we can prevent a situation that ignores a class prior to a class-imbalance task such as NER (Section~\ref{sec:pmi}).
We refer to the class imbalance problem as the number of negative examples is overwhelmingly larger than that of positive examples.
Finally, we observe that attributes of entity length are highly correlated with the models' predictions.
Thus, we modify a class distribution to focus on a complex structure and long-named entity that contains adjectives or conjunctions within the entity~\cite{wei2020multichannel, cho2020combinatorial}.
In other words, our method adaptively smooth the class distribution of the bias-only model based on its entity length (Section~\ref{sec:scaling}).
The key intuition of our approach, depicted in Figure~\ref{fig:model architecture}, is to explicitly incorporate dataset biases into the target model to be more robust against the biases.

Our contributions are as follows:
(1) We propose a model-agnostic debiasing method that a bias-only model replaced with statistics of  PMI.
(2) Our method alleviates the class imbalance inherent for the task of NER and utilizes a subword frequency to tackle the OOV problem.
(3) Our method adaptively smooths a class distribution of bias-only model on long-named and complex structure entity.
(4) Our experiments demonstrate that our method shows overall improvement of generalization ability and in-domain performance in both general and biomedical domain NER datasets.

\section{Related work}
\paragraph{Named Entity Recognition.} Named Entity Recognition (NER) has long been studied and received attention as a core task of information extraction.
Initially, researchers of NER focused on constructing a model using handcraft features \cite{zhou2002named, bender2003maximum, settles2004biomedical}.
Since it was labor-intensive, deep learning approaches automatically extract features and label a sequence \cite{dos2015boosting, lample2016neural, habibi2017deep}.
Recently, with the advent of pre-training on a large corpus, contextualized language models performed significantly on the benchmark datasets and reached a plateau on performances \cite{devlin2019bert, liu2019roberta, beltagy2019scibert, lee2020biobert}.

However, these raise questions about the generalization behavior of pre-trained language models.
Specifically, in NER tasks, \citet{fu2020rethinking} has divided the set of test entities into different subsets and broken down overall performance into interpretable categories.
To explore the generalization ability, \citet{lin2020rigourous} partitioned the benchmark datasets as seen and unseen mention sets to make the situations of erasing name regularity, mention coverage, and context diversity.
In BioNER, \citet{kim2021your} split unseen mention sets into synonym and concept based on Concept Unique ID (CUI) overlaps between training and test datasets.
The authors adopted a debiasing method to enhance generalization behavior.
Motivated by debiasing framework with a prior distribution of task-related statistics \cite{ko2020look, kim2021your}, we introduce another statistic that fits well for the task of class imbalance.

\begin{figure}[]
 \centering
 \includegraphics[width=200px]{figure/ablation_length_of_entities.PNG}
 \caption{Biomedical entities are long and complex compared to general domain.}
 \label{fig:ablation length of entities}
\end{figure}

In terms of an entity length, there are several error cases that long and complex structure biomedical terms that are not well predicted.
\citet{hong2020dtranner} shows that conjunction or special character "/" within entity cause models hard to predict in types of disease and chemical.
\citet{wei2020multichannel} demonstrated three main error reasons as follows: complex and long-named entities, annotation inconsistency on corpora, and abbreviations.
On evaluation metric, \citet{fu-etal-2020-interpretable} observed the influence of entity length varies on structural biases of different decoders such as MLP or CRF.
Since the bias-only model helps to not over-trust on dataset bias, we hypothesize that relative smoothing on the bias-only models' predictions can have better robustness when dealing with long entities.

\begin{figure*}[t]
 \centering
 \includegraphics[width=460px]{figure/model_figure_modify.png}
 \caption{(a) We compute a PMI that captures a correlation between subword and its class. The PMI reflects solutions of \textit{Out-Of-Vocabulary} (OOV) issue and class imbalance. For example, we depict the OOV issue with a word \textit{'Latent'} that does not appear in the training dataset but occurs in the test dataset. We also design a class prior to reflect class imbalance inherent in the NER task. (b) A red line indicates that computing the PMI is only added during the training time. We use a temperature scaling~\cite{guo2017calibration} to smooth a class distribution of the normalized PMI depend on the entity length, i.e., $T_i := 1 + \lambda \cdot L_i$ if $L_i > 1$.}
 \label{fig:model architecture}
\end{figure*}

\paragraph{Dataset Bias and Debiasing Methods.}
In NLP tasks, dataset bias is a superficial cue or surface form of text associated with specific labels.
Recently, pre-trained language models~\cite{devlin2019bert,liu2019roberta,lee2020biobert} tend to exploit the dataset biases to predict rather than understanding the relationship between words and labels~\cite{gururangan2018annotation, mccoy2019right, geirhos2020shortcut, bender-koller-2020-climbing, utama2020towards, du2021interpreting}.
This eventually leads to the advent of debiasing methods to enhance the generalization ability.
\citet{clark2019don} introduced a \textit{product-of-expert} style as a model-agnostic approach of debiasing method to scale down the gradients on the biased examples.
\citet{mahabadi2020end} proposed a bias-only model to focus on learning hard examples by reducing the relative importance of biased examples using focal loss.
\citet{schuster-etal-2019-towards} suggested a regularization method to down-weigh the give-away phrases that cause the bias.
Since previous debiasing methods are vulnerable to preserve in-domain performance while improving out-of-domain performance, a \textit{confidence regularization}~\cite{utama2020mind} method addressed this trade-off by connecting robustness against dataset biases and overconfident problem in neural models~\cite{papernot2016distillation, feng2018pathologies}.

\vspace{-1mm}

\section{Backgrounds}
\paragraph{Task Formulation.}
A NER is a task of identifying a word or phrase that corresponding to instances such as a person, location, organization, and miscellaneous.
A key challenge is extracting and classifying named entities in a corpus with pre-defined tags of entities.
We use BIO tagging scheme \cite{ramshaw1999text}, where B and I tag is concatenated with all target entity types and O tag indicates that a token belongs to no chunk, e.g., B-PER, B-LOC.
Formally, a sentence consists of an input sequence \textbf{X} $=$ \textbraceleft$x_1, x_2, \ldots, x_N$\textbraceright with length $N$, its ground truth labels of input sequence \textbf{Y} $=$ \textbraceleft$y_1, y_2, \ldots, y_N$\textbraceright, and its predicted labels of output sequence $\hat{\textbf{Y}}$ $=$ $\{\hat{y}_1, \hat{y}_2, \ldots, \hat{y}_N\}$ where $y_i$ $\in$ \textbraceleft$1, 2, \ldots, C$\textbraceright~ and $C$ are the number of classes.
There are two widely used decoding strategies: tag independent decoding (i.e., softmax) and tag dependent decoding (i.e., Conditional Random Field~\cite{lafferty2001conditional}).
The goal of NER is to predict mentions by assigning output label $\hat{y}_{i}$ for each token $x_{i}$ using these decoding strategies: $P(\hat{y}_i | \textbf{X})$ or $P(\hat{y}_i | \textbf{X}, \hat{y}_1, \hat{y}_2, \ldots, \hat{y}_{i-1})$.

\paragraph{Partition Benchmark Dataset.}
\label{sec:partition}
Given annotated entities in the training dataset, we make an entity dictionary to split the evaluation dataset.
In other words, using this dictionary generates the two mention sets: mentions that are seen and unseen during training.
In particular BioNER task, a raw dataset suggests the CUI of each entity which signifies the same concept (or meaning) of different entities to form a concept dictionary.
We define our evaluation dataset to be split based on these dictionaries as follows,

\begin{enumerate}[label=(\roman*)]
    \item $\textsc{Mem} := \{ M | M \in E_{tr}(M)\}$
    \item $\textsc{Unseen} := \{M|M \notin {E_{tr}}(M) \}$
    \item $\textsc{Syn} := \{ M | M \notin E_{tr}(M), M \in C_{tr}(M) \}$
    \item $\textsc{Con} := \{ M | M \notin E_{tr}(M), M \notin C_{tr}(M) \}$,
\end{enumerate}
which $E_{tr}(M)$ and $C_{tr}(M)$ refers to the entity and concept dictionary of training dataset respectively, and $M$ is a mention indicating a surface form of the entity.
Memorization (\textsc{Mem}) is composed of mentions that appear in the training dataset.
On the contrary, \textsc{Unseen} contains the mentions that did not occur in the training dataset.
More specifically, in BioNER dataset, we can further partition \textsc{Unseen} mention set as Synonym Generalization (\textsc{Syn}) and Concept Generalization (\textsc{Con})~\cite{kim2021your}.

\paragraph{Evaluation Metrics.}
Evaluation of NER is divided into two parts: boundary detection and type identification \cite{li2020survey}.
General domain NER is to identify and classify text spans into human-annotated categories \cite{sang2003introduction}.
On the other hand, most of BioNER is to compute an entity-level boundary detection.
The common evaluation metric of NER and BioNER is to compute entity-level metric using micro-averaged F1 score:
$F1 = \frac{2*P*R}{P+R}$, which P and R are Precision and Recall respectively.
In spite of using the common evaluation metric as an in-domain performance, we follow the evaluation metric of \cite{kim2021your} because predicted entities are likely not to exist in the annotated entities.
Thus, we only report the recall score of the partitioned mention sets (e.g., \textsc{Mem}, \textsc{Unseen}, \textsc{Syn}, and \textsc{Con}).

\paragraph{Subword Frequency.}
\label{sec:subword}
We depict an \textit{Out-Of-Vocabulary} (OOV) issue in Figure~\ref{fig:model architecture}(a).
If a class distribution of a particular word is different in training and test datasets then the models fail to reflect informative statistics~\cite{leevy2018survey, li2020dice}.
For example, a word \textit{'Latent'} does not appear in the training dataset but occurs as an entity in the evaluation dataset as \textit{'Latent BHRF1 transcripts'}.
In this situation, it is difficult to apply a meaningful statistic based on the word unit.
However, if subword statistics are applied, \textit{'Latent'} can split into subwords of \textit{'Late'} and \textit{'\#\#nt'}.
Since the word \textit{'Late'} occurs with only 'O' label in the training dataset, debiasing to predict without 'O' label enables to capture the entity of \textit{'Latent BHRF1 transcripts'}.
The example is fully detailed in Section~\ref{sec:informative}.

\paragraph{Modeling a Class Prior.}
A joint modeling of $f_\theta(\hat{\textbf{Y}}|\textbf{X})$ and $\frac{P(\textbf{Y}|\textbf{X})}{P(\textbf{Y})}$ have been used by the current debiasing methods~\cite{clark2019don}.
However, these methods use a bias-only model with learned parameters and ignore $P(\textbf{Y})$ factor while designing.
Because it is hard to reflect a class prior with a pointer network or a situation that classes of a particular dataset are uniformly distributed.
In contrast, a class imbalance is severe in the task of NER which needs to address the dominating influence of easy-negative examples~\cite{meng2019dsreg, li2020dice, mahabadi2020end}.
In Table~\ref{tab:all dataset statistics}, we could easily see that ratio of negative examples to positive examples are extremely high in the NER tasks.
We thus argue that a class prior needs to be considered to reduce the influence of easy-negative examples.

\begin{table}[]
\centering
\begin{adjustbox}{max width=0.5\textwidth}
\begin{tabular}{ l l c c c c }
\toprule
\multicolumn{1}{l}{Dataset} & \multicolumn{1}{l}{Entity Type} & \multicolumn{1}{c}{\begin{tabular}[c]{@{}c@{}}\# of\\ Sentences\end{tabular}} & \begin{tabular}[c]{@{}c@{}}Positive\\ Examples\end{tabular} & \begin{tabular}[c]{@{}c@{}}Negative\\ Examples\end{tabular} & Ratio \\
\midrule
NCBI-disease & Disease & 8,230 & 13K & 203K & 15.47 \\
BC5CDR-disease & Disease & 12,058 & 14K & 313K & 22.29 \\
BC5CDR-chem & Drug/Chem. & 12,126 & 14K & 313K & 22.08 \\
BC4CHEMD & Drug/Chem. & 89,664 & 129K & 2.3M & 17.83 \\
BC2GM & Gene/Protein & 21,370 & 44K & 535K & 11.96 \\
JNLPBA & Gene/Protein & 27,051 & 122K & 637K & 5.20 \\
LINNAEUS & Species & 21,459 & 4.3K & 478K & 110.71 \\
Species-800 & Species & 8,723 & 6.7K & 220K & 32.79 \\
CoNLL-2003 & General & 14,986 & 34K & 170K & 4.98 \\
WNUT2017 & General & 3,394 & 15K & 74K & 4.79 \\
\bottomrule
\end{tabular}
\end{adjustbox}
\caption{Data Statistics of the NER tasks. All datasets are listed with details: entity type, number of sentences, positive examples ('B' and 'I' tag), negative examples ('O' tag), and the ratio of negative examples to positive examples (negative examples / positive examples).}
\label{tab:all dataset statistics}
\end{table}

\section{Methods}
We start by introducing our debiasing method which utilizes PMI as a bias-only model to alleviate the OOV issue and class imbalance.
For long entities, we provide an explanation of smoothing the class distribution via temperature scaling~\cite{guo2017calibration}.

\subsection{Debiasing Framework}
We illustrate an overall framework of our debiasing method in Figure~\ref{fig:model architecture}(b).
Given the input sequence of tokens \textbf{X} and its ground truth labels \textbf{Y}, a model has to predict its labels of output sequence $\hat{\textbf{Y}}$.
As a simple ensemble method, we use a Bias Product that trains a product of experts \cite{hinton2002training, smith2005logarithmic} and adds the log probabilities of bias-only model and target model $f_\theta$ \cite{clark2019don, he2019unlearn}.
It is intuitive that the target model can prevent focusing on dataset biases during training.
We compute a final probability as follows,

\begin{equation}
    \hat{p_{i}} = softmax(log(p_i) + log(b_i)),
\end{equation}
where $\hat{p_i} \propto p_i \cdot b_i$, a $log(p_i)$ is log probability of the target model $f_\theta$ which we want to not over-trust on the surface form, and $log(b_i)$ is log probability of the bias-only model.

Besides, Bias Product can be extended to Learned-Mixin and Learned-Mixin+H~\cite{clark2019don}.
The Learned-Mixin method allows the target model to determine how much to trust dataset bias given input sequence.
On the contrary, the Learned-Mixin+H method used to prevent the target model from neglecting bias probability $log(b_i)$ by adding an entropy penalty.
Detailed explanations are described in Appendix~\ref{app:learned-mixin}.

Overall, our method can operate on these ensemble frameworks without introducing additional parameters.
We define a loss function of our model as below:
\vspace{-0.2cm}
\begin{equation}

    L = -\frac{1}{N}\sum_{i=1}^{N} y_{i} \cdot softmax(log(p_i) + log(b_i)),
\end{equation}
where $N$ is the length of an input sequence.
It minimizes the negative log-likelihood of the distribution of final probability $\hat{p_i}$ and its ground-truth label $y_i$.
Note that $log(b_i)$ is only added during the training process while it does not apply at the inference time.
This process is depicted with a red arrow in Figure~\ref{fig:model architecture}.

\subsection{Pointwise Mutual Information}
\label{sec:pmi}
To cope with the class-distribution shift of word between training and test dataset~\footnote{PMI was used as a correlation score between the word and class in a particular dataset \cite{gururangan2018annotation}.
}, we compute the PMI between subword and class in the training dataset:

\vspace{-0.2cm}
\begin{equation}
\hspace*{-0.2cm}
PMI(subword, class) = log\frac{P(subword,class)}{P(subword) \cdot P(class)}. \nonumber
\label{eq:PMI}
\end{equation}
For a numerical issue, we apply add-K smoothing to each subword frequency~\footnote{K is a hyperparameter and we use K=100 same as \cite{gururangan2018annotation} to emphasize correlations of word and class.}.

Overall, PMI captures a correlation between subwords and their corresponding class.
We further suggest a way to understand how each mathematical expression operates.
A joint distribution $P(subword, class)$ assigns low training signals to subword of which class distributions are dominant to a particular label.
However, $P(subword, class)$ focuses on frequent subwords independent of class such as \textit{'disease'} or \textit{'deficiency'}.
To handle this problem, a $P(subword)$ alleviates a situation in which subword appears frequently regardless of the class.
Finally, a $P(class)$ enhances regularization when the class imbalance between positive examples (i.e., B and I tag) and negative examples (i.e., O tag) is severe.
Since there are many easy-negative examples, we want to prevent the training signals are ruled by these examples.

Because PMI is much larger than the logit value of the target model $f_\theta$, we normalize the PMI to a probability distribution.
We then compute the normalized PMI (i.e., $PMI_{ic}$) and final probability $\hat{p_i}$ as below,

\begin{equation}
    PMI_{ic} = \frac{exp(PMI_{ic})}{\sum_{c'=1}^{C}exp(PMI_{ic'})}
\end{equation}
\begin{equation}
    \hat{p_{i}} = softmax(log(p_i) + log(PMI_i))
\end{equation}
where $C$ is the number of classes in a particular dataset, and $PMI_{ic}$ indicates the probability of the PMI corresponding to each class $c$.
$PMI_i = \{ PMI_{i1}, PMI_{i2} \ldots, PMI_{iC} \}$ is a class distribution of the PMI for each token $x_i$ in the training dataset.

\begin{table*}[t]
\begin{adjustbox}{width=\textwidth}
\begin{tabular}{ l c c c c c c c | c c c c | c c c c }
\toprule
\multicolumn{1}{c}{\multirow{3}{*}{Method}} & \multicolumn{1}{c}{\multirow{3}{*}{\textsc{Subword}~}} &
\multicolumn{1}{c}{\multirow{3}{*}{\textsc{Class}~}} &
\multicolumn{1}{c}{\multirow{3}{*}{\textsc{Temp}~}} & \multicolumn{4}{c}{NCBI-disease} & \multicolumn{4}{c}{BC5CDR-disease} & \multicolumn{4}{c}{BC5CDR-chem} \\ \cmidrule{5-16}
\multicolumn{1}{c}{} & \multicolumn{1}{c}{} & & \multicolumn{1}{c}{} & \begin{tabular}[c]{@{}c@{}}\textsc{Mem}\\ (R)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}\textsc{Syn}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Con}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}Total\\ (P/R/F) \end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Mem}\\ (R)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}\textsc{Syn}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Con}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}Total\\ (P/R/F) \end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Mem}\\ (R)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}\textsc{Syn}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Con}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}Total\\ (P/R/F) \end{tabular}  \\ \midrule
$BiLSTMCRF$ & \cmark & \xmark & \xmark & 93.5 & 78.8 & 79.3 & 85.2 / 88.7 / 86.9 & 89.7 & 74.5 & 71.2 & \textbf{78.5} / 82.9 / 80.6 & 94.8 & 80.2 & 84.3 & 90.5 / 91.2 / 90.8 \\
$BiLSTMCRF_{BP}$ & \xmark & \xmark & \xmark & 91.8 & 79.5 & 80.1 & 84.8 / 88.8 / 86.7 & 89.1 & 75.0 & 72.3 & 78.0 / 83.1 / 80.5 & 94.4 & 81.7 & 85.1 & 91.0 / 91.7 / 91.3 \\
$BiLSTMCRF_{BP}$ & \cmark & \cmark & \cmark & 93.8 & \textbf{81.7} & \textbf{82.1} & \textbf{85.5} / \textbf{89.4} / \textbf{87.4} & \textbf{90.2} & \textbf{76.2} & 72.9 & 78.4 / \textbf{83.9} / \textbf{81.1} & 94.9 & 81.9 & 84.7 & \textbf{91.2} / 91.8 / \textbf{91.5} \\
$BiLSTMCRF_{LM}$ & \cmark & \cmark & \cmark & \textbf{94.0} & 80.8 & 81.5 & \textbf{85.5} / 89.0 / 87.2 & 90.0 & 75.8 & \textbf{73.4} & 78.1 / 83.8 / 80.9 & \textbf{95.2} & \textbf{82.4} & \textbf{85.2} & 90.8 / \textbf{92.2} / \textbf{91.5} \\ \midrule
$BioBERT$  & \cmark & \xmark & \xmark & 95.5 & 80.8 & 83.5 & 86.3 / 90.5 / 88.3 & 91.3 & 76.7 & 73.4 & 79.8 / 85.5 / 82.5 & 97.2 & 82.9 & 86.9 & 91.4 / 92.8 / 92.1 \\
$BioBERT_{BP}$ & \xmark & \xmark & \xmark & 93.9 & 81.1 & 84.1 & 85.7 / 90.7 / 88.1 & 90.9 & 77.1 & 74.3 & 79.2 / 86.1 / 82.5 & 96.8 & 83.1 & 88.0 & 91.2 / 92.9 / 92.0 \\
$BioBERT_{BP}$ & \cmark & \cmark & \cmark & 95.7 & \textbf{83.6} & \textbf{86.5} & \textbf{86.8} / \textbf{91.5} / \textbf{89.0} & 92.0 & 78.5 & 75.1 & \textbf{80.8} / 86.0 / \textbf{83.3} & 97.1 & 83.8 & \textbf{88.2} & \textbf{91.5} / 93.1 / 92.3 \\
$BioBERT_{LM}$ & \cmark & \cmark & \cmark & \textbf{96.0} & 82.1 & 84.0 & 86.7 / 91.0 / 88.8 & \textbf{92.2} & \textbf{78.8} & \textbf{75.2} & 79.8 / \textbf{86.2} / 82.9 & \textbf{97.2} & \textbf{84.4} & 87.5 & 91.4 / \textbf{93.3} / \textbf{92.4} \\
\bottomrule
\end{tabular}
\end{adjustbox}
\caption{Performance of debiasing method on the biomedical domain NER datasets. Each dataset is partitioned as Memorization (\textsc{Mem}), Synonym Generalization (\textsc{Syn}), and Concept Generalization (\textsc{Con}).
We use \cmark and \xmark ~to show what components are used. Best performances are in bold.}
\label{tab:maintable 1}
\end{table*}

\subsection{Bias-weighted Scaling on Long-Named and Complex Entity}
\label{sec:scaling}
As illustrated in Figure~\ref{fig:model architecture}, our method adaptively involves smoothing a class distribution of PMI using temperature scaling~\cite{guo2017calibration}.
Biomedical entities are long named and have a complex structure that contains conjunctions or adjectives within entities~\cite{wei2020multichannel, cho2020combinatorial}.
This can be easily observed through our Figure~\ref{fig:ablation length of entities}.
Thus, we define a temperature $T_i$ that depends on the length $L_i$ of the entity to smooth the class distribution of PMI.
Given the PMI probability $PMI_{ic}$, the new statistic is defined as below,

\begin{equation}
    T_i := \begin{cases}
        1 + \lambda \cdot L_i, & \text{if $L_i$ > 1} \\
        1, & \text{otherwise}
    \end{cases}
\end{equation}
\begin{equation}
    PMI_{ic} = \frac{exp(PMI_{ic} / T_i)}{\sum_{c'=1}^{C}exp(PMI_{ic'} / T_i)}.
\end{equation}
$T_i$ is a temperature~\cite{guo2017calibration} of token $x_i$ to smooth a class distribution of $PMI_{ic}$ depend on the entity length $L_i$, and $\lambda$ is a hyperparamter.
As $T \text{→} ~\infty$, the probability $PMI_{ic}$ approaches $1 / C$, which represents maximum uncertainty.
Conversely, as $T \text{→} ~ 0$, the probability collapses to a point mass.

It is worth noticing that since PMI requires task-related information, it does not introduce an additional memory.
We can also choose other statistics such as random value or word frequency corresponding to its class.
However, our experiments show that using PMI achieves the best results.
Moreover, improving generalizability with our method seems to be beneficial in other class-imbalance tasks.

\section{Experiments}
In this section, we evaluate our debiasing method on general domain and biomedical domain datasets of the NER task.
We address three main questions in the experiments as follows:
(1) How much effective of our method as a bias-only model using subword frequency, class prior, and length dependent scaling? (2) Does bias-weighted scaling predict long-named and complex-structure entities? (3) What component can highly support the capacity of generalization?

\begin{table*}[t]
\begin{adjustbox}{width=\textwidth}
\begin{tabular}{l c c c c c c | c c}
\toprule
 \multicolumn{1}{c}{\multirow{3}{*}{Method}} & \multicolumn{1}{c}{\multirow{3}{*}{\textsc{Subword}~}} &
\multicolumn{1}{c}{\multirow{3}{*}{\textsc{Class}~}} &
\multicolumn{1}{c}{\multirow{3}{*}{\textsc{Temp}~}} & \multicolumn{3}{c}{CoNLL 2003} & \multicolumn{2}{c}{WNUT 2017}      \\ \cmidrule{5-9}
& & & & \textsc{Mem} (R) &  \textsc{Unseen} (R) &  Total (P/R/F) & Total (P/R/F) & Surface (P/R/F) \\ \midrule
$BiLSTMCRF$  & \cmark & \xmark & \xmark & 93.2 & 87.1  & 89.8 / 90.1 / 89.9 & 57.8 / 31.2 / 40.5 & 55.6 / 29.1 / 38.2 \\
$BiLSTMCRF_{BP}$  & \xmark & \xmark & \cmark & 92.8 & 87.5  & 89.5 / 90.5 / 90.0 & 49.8 / 34.5 / 40.7 & 48.1 / 32.3 / 38.6 \\
$BiLSTMCRF_{BP}$  & \cmark & \cmark & \cmark & \textbf{93.5} & \textbf{87.7}  & \textbf{90.1} / \textbf{90.7} / \textbf{90.4} & 51.2 / \textbf{36.0} / \textbf{42.3} & 49.9 / \textbf{34.1} / \textbf{40.5} \\
$BiLSTMCRF_{LM}$  & \cmark & \cmark & \cmark & 93.3 & \textbf{87.7}  & 90.0 / 90.6 / 90.3 & \textbf{58.2} / 30.9 / 40.3 & \textbf{57.1} / 33.2 / 40.4 \\ \midrule
$BERT$ & \cmark & \xmark & \xmark & 94.7 & 88.0  & 90.7 / 91.3 / 91.0 & 58.6 / 36.7 / 45.1 & 57.5 / 34.7 / 43.1 \\
$BERT_{BP}$ & \xmark & \xmark & \xmark & 93.8 & 88.4  & 90.1 / 91.7 / 90.9 & 56.0 / 36.3 / 44.3 & 54.6 / 34.6 / 42.3 \\
$BERT_{BP}$ & \cmark & \cmark & \cmark & \textbf{95.2} & \textbf{89.1}  & \textbf{90.9} / \textbf{92.0} / \textbf{91.5} & 55.3 / \textbf{42.1} / \textbf{47.8} & 54.3 / \textbf{40.9} / \textbf{46.6} \\
$BERT_{LM}$ & \cmark & \cmark & \cmark & 94.8 & 88.8 & 90.8 / 91.8 / 91.3 & \textbf{58.9} / 36.5 / 45.0 & \textbf{60.0} / 35.1 / 44.3\\
\bottomrule
\end{tabular}
\end{adjustbox}
\caption{Performance of debiasing method on the general domain NER and biomedical domain NER datasets. Best performances are in bold.}
\label{tab:maintable 2}
\end{table*}

\subsection{Experimental Settings}

\paragraph{Dataset.}
We evaluate our debiasing method on ten benchmark NER datasets including eight biomedical domains and two general domains.
For biomedical domain, we use eight datasets: disease entities (NCBI-disease \cite{dougan2014ncbi}, BC5CDR-disease \cite{li2016biocreative}), drug and chemical entities (BC5CDR-chemical \cite{li2016biocreative}, BC4CHEMD \cite{krallinger2015chemdner}), gene and protein entities (BC2GM \cite{smith2008overview}, JNLPBA \cite{kim2004introduction}), and species entities (LINNAEUS \cite{gerner2010linnaeus}, Species-800 \cite{pafilis2013species}).
We follow the pre-processed version of previous work \cite{crichton2017neural}~\footnote{https://github.com/cambridgeltl/MTL-Bioinformatics-2016}.
For the general domain, we use two English datasets: CoNLL-2003 shared benchmark dataset \cite{sang2003introduction} and WNUT-2017 which focuses on identifying unusual, previously-unseen entities \cite{derczynski2017results}.
WNUT-2017 dataset has its own evaluation metric which measures how good systems are at correctly recognizing a diverse range of entities, rather than just the very frequent surface forms.
We focus on improving both in-domain performance and recall of partitioned datasets (e.g., \textsc{Mem}, \textsc{Syn}, \textsc{Con}, and \textsc{Unseen}).
Descriptions of data statistics are suggested in Appendix~\ref{app:dataset statistic}.

\paragraph{Baselines.}
Our target models are composed of Bi-LSTM and BERT architectures which have been widely used in many works for NER.
We use BERT \cite{devlin2019bert} in general domain and BioBERT \cite{lee2020biobert} in biomedical domain.
We leverage BERT and BioBERT token embeddings to train Bi-LSTM~\cite{huang2015bidirectional} with CRF decoder ($BiLSTMCRF$).
On the bias-only model, We compare PMI with a class distribution of word frequency~\cite{kim2021your}.
Rather than trying PMI solely on the Bias-Product ($BP$), we extend our experiments to show how the PMI can operate with other debiasing methods such as Learned-Mixin ($LM$) and Learned-Mixin+H~\cite{clark2019don}.
We denote \cmark and \xmark ~when using subword frequency (\textsc{Subword}), class prior (\textsc{Class}), and bias-weighted scaling (\textsc{Temp}).
Description of training details and hyperparameters are suggested in Appendix~\ref{app:hyperparameter}.

\begin{table*}[t]
\begin{adjustbox}{width=\textwidth}
\begin{tabular}{ l c c c c c c c | c c c c | c c c c }
\toprule
\multicolumn{1}{c}{\multirow{3}{*}{Method}} & \multicolumn{1}{c}{\multirow{3}{*}{\textsc{Subword}~}} & \multicolumn{1}{c}{\multirow{3}{*}{\textsc{Class}~}} &
\multicolumn{1}{c}{\multirow{3}{*}{\textsc{Temp}~}} & \multicolumn{4}{c}{NCBI-disease} & \multicolumn{4}{c}{BC5CDR-disease} & \multicolumn{4}{c}{BC5CDR-chem} \\ \cmidrule{5-16}
\multicolumn{1}{c}{} & \multicolumn{1}{c}{} & \multicolumn{1}{c}{} & \multicolumn{1}{c}{} & \begin{tabular}[c]{@{}c@{}}\textsc{Mem}\\ (R)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}\textsc{Syn}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Con}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}Total\\ (P/R/F) \end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Mem}\\ (R)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}\textsc{Syn}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Con}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}Total\\ (P/R/F) \end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Mem}\\ (R)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}\textsc{Syn}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}\textsc{Con}\\ (R)\end{tabular} & \begin{tabular}[c]{@{}c@{}}Total\\ (P/R/F) \end{tabular} \\ \midrule
$BioBERT_{BP}$ & \xmark & \xmark & \xmark & 93.9 & 81.1 & 84.1 & 85.7 / 90.7 / 88.1 & 90.9 & 77.1 & 74.3 & 79.2 / 86.1 / 82.5 & 96.8 & 83.1 & 88.0 & 91.2 / 92.9 / 92.0 \\
$BioBERT_{BP}$ & \cmark & \xmark & \xmark & 94.2 & 80.8 & 83.6 & 86.7 / 90.5 / 88.6 & 91.8 & 76.0 & 73.2 & 80.5 / 85.6 / 83.0 & 97.4 & 81.4 & 86.7 & 91.7 / 92.7 / 92.2 \\
$BioBERT_{BP}$ & \xmark & \cmark & \xmark & 95.5 & 82.4 & 83.8 & 86.4 / 91.0 / 88.6 & 91.8 & 77.5 & 73.8 & 80.3 / 86.0 / 83.1 & 97.4 & 83.5 & 87.5 & 91.7 / 92.6 / 92.2 \\
$BioBERT_{BP}$ & \xmark & \xmark & \cmark & 95.5 & 81.3 & 82.5 & 86.7 / 90.8 / 88.7 & \textbf{92.9} & 75.8 & 71.5 & 80.3 / 85.9 / 83.0 & 97.7 & 81.4 & 86.8 & 91.7 / 92.9 / 92.3 \\
$BioBERT_{BP}$ & \cmark & \cmark & \xmark & 93.8 & 83.5 & \textbf{87.1} & 86.4 / 90.7 / 88.5 & 91.6 & 77.3 & 74.7 & 80.1 / \textbf{86.3} / 83.1 & 97.3 & 83.3 & \textbf{88.4} & 91.6 / 92.8 / 92.2 \\
$BioBERT_{BP}$ & \cmark & \xmark & \cmark & 95.6 & 82.4 & 80.3 & 86.5 / 90.5 / 88.4 & 92.5 & 76.4 & 73.2 & 81.2 / 86.1 / 83.6 & 97.7 & 83.3 & 87.4 & \textbf{91.8} / \textbf{93.2} / \textbf{92.5} \\
$BioBERT_{BP}$ & \xmark & \cmark & \cmark & 95.4 & \textbf{83.9} & 84.2 & \textbf{87.0} / 90.6 / 88.8 & 92.3 & 73.2 & 70.4 & \textbf{82.7} / 84.8 / \textbf{83.7} & \textbf{97.8} & 82.3 & 87.3 & 91.6 / \textbf{93.2} / 92.4 \\
$BioBERT_{BP}$ & \cmark & \cmark & \cmark & \textbf{95.7} & 83.6 & 86.5 & 86.8 / \textbf{91.5} / \textbf{89.0} & 92.0 & \textbf{78.5} & \textbf{75.1} & 80.8 / 86.0 / 83.3 & 97.1 & \textbf{83.8} & 88.2 & 91.5 / 93.1 / 92.3 \\
\bottomrule
\end{tabular}
\end{adjustbox}
\caption{Ablation study of three components on biomedical domain: subword frequency~(\textsc{Subword}), class prior~(\textsc{Class}), and length dependent scaling~(\textsc{Temp}). Best performances are in bold.}
\label{tab:ablation table 1}
\end{table*}

\begin{table*}[t]
\centering
\begin{adjustbox}{width = \textwidth}
\begin{tabular}{l l l l l l l l l l l l }
\toprule
\multicolumn{12}{ c }{\textbf{JNLPBA (Out-Of-Vocabulary 'Latent')}} \\ \midrule
\multicolumn{12}{ l }{\begin{tabular}[c]{@{}l@{}}
Sentence: \textbf{Latent BHRF1 transcripts} encoding \textbf{bcl - 2 homologue} and \textbf{BCRF1 transcripts} encoding \textbf{viral interleukin ( vIL )} \\
~~~~~~~~~~~~~~~~~ \textbf{- 10} were detected in one and two of eight patients , respectively .\\
$BioBERT$: \mb{BHRF1 transcripts}, bcl - 2, BCRF1 transcripts, viral interleukin ( vIL ) - 10 \\
$BioBERT_{BP}$ (+ None): \mb{BHRF1 transcripts}, bcl - 2 homologue, BCRF1 transcripts, viral interleukin ( vIL ) - 10 \\
$BioBERT_{BP}$ (+ \textsc{Subword}): \ch{Latent BHRF1 transcripts}, bcl - 2 homologue, BCRF1 transcripts, viral interleukin ( vIL ) - 10\end{tabular}} \\
\midrule
\multicolumn{12}{ c }{\textbf{BC2GM (Long-Named Entity)}} \\ \midrule
\multicolumn{12}{ l }{\begin{tabular}[c]{@{}l@{}}
Sentence: \textbf{Rhinovirus 2A protease} and \textbf{foot - and - mouth - disease virus L protease} were used to analyze the association of \\
~~~~~~~~~~~~~~~~ \textbf{eIF4G} with \textbf{eIF4A} , \textbf{eIF4E} , and \textbf{eIF3} . \\
$BioBERT$: Rhinovirus 2A protease, \mb{mouth - disease virus L protease}, eIF4G, eIF4A, eIF4E, eIF3 \\
$BioBERT_{BP}$ (+ None): Rhinovirus 2A protease, \mb{mouth disease virus L protease}, eIF4G, eIF4A, eIF4E, eIF3 \\
$BioBERT_{BP}$ (+ \textsc{Temp}): Rhinovirus 2A protease, \ch{foot - and - mouth - disease virus L protease}, eIF4G, eIF4A, eIF4E, eIF3 \end{tabular}} \\
\bottomrule
\end{tabular}
\end{adjustbox}
\caption{Example of dealing with OOV problem and long-named entity. Ground truth labels are in bold, and wrong answer is in red, and correct answer is in blue.}
\label{tab:qualitative example}
\end{table*}

\subsection{Effectiveness of Our Debiasing Method}
Table~\ref{tab:maintable 1} and \ref{tab:maintable 2} show our results on both general and biomedical domain.
In Table~\ref{tab:maintable 1}, using three components, namely \textsc{Subword}, \textsc{Class}, and \textsc{Temp}, show significant improvements of recall on \textsc{Syn} and \textsc{Con}.
On \textit{$BiLSTMCRF_{BP}$} architecture, the average enhancements of our method (row 3) in \textsc{Syn} and \textsc{Con} are about 1.2\% and 0.74\% compared to the baseline (row 2)~\cite{kim2021your} respectively.
On \textit{$BioBERT_{BP}$} architecture, improvements of our method (row 7) are about 1.53\% in \textsc{Syn} and 1.13\% in \textsc{Con} compared to the current debiasing method respectively.
In spite of these improvements, our method preserves or even improves recall of \textsc{Mem}, and enhances in-domain performance (i.e., F1 score) about 0.5\%.
In the general domain, in Table~\ref{tab:maintable 2}, the CoNLL-2003 dataset shows a similar trend but not much high value.
However, our method shows a significant increase in F1 score of in-domain performance and surface form about 3.5\% and 4.3\% respectively.

Not surprisingly, a Learned-Mixin framework (row 4 and 8) that prevents the target model from reversing the bias decreases the generalizability while increases memorization ability.
We also observe that the Learned-Mixin+H method shows unstable results because two-regularization functions (i.e., entropy penalty and temperature scaling) conflict with each other to a bias-only model.
In other words, the entropy penalty encourages this bias-only model to be non-uniform, but the temperature scaling forces the model to be uniform.
Due to the low and unstable results, we remove the Learned-Mixin+H method in our results.
We conduct 5 runs with different seeds for all experiments.
The t-tests indicate that p $<$ 0.05.
We report other datasets for the space limit in Appendix~\ref{app:experimental results}.

\subsection{Quantitative Analysis}
Overall, adopting three components not only boosts the performance of memorization ability under all cases but also improves the generalizability by preventing dataset bias.
In Table~\ref{tab:ablation table 1}, we point out that the class prior (\textsc{Class}) and subword frequency (\textsc{Subword}) improves a memorization ability (\textsc{Mem}) while decreasing the capacity of generalization (\textsc{Syn} and \textsc{Con}).
However, unexpectedly, a combination of \textsc{Class} and \textsc{Subword} gives a decrease in the \textsc{Mem} but high increase in \textsc{Syn} and \textsc{Con}.
Furthermore, we observe that applying a temperature scaling (\textsc{Temp}) dependent on entity length enhances the ability of memorization.
Since we suggest the \textsc{Temp} to predict well on the long-named entity, it increases an in-domain performance.
Although the performance improvement itself seems insignificant, we will show that the improvement is large for long entities in the next paragraph.
To sum up, it seems that the three components keep an equilibrium to make a more robust model with a high capacity for generalization and memorization.

\begin{figure}[]
 \centering
 \includegraphics[width=200px]{figure/length_ablation_2.PNG}
 \caption{Bias-weighted scaling on long-named and complex entities. L(E) indicates an entity length.}
 \label{fig:length ablation}
\end{figure}

\paragraph{Adjusting $\lambda$ for Temperature Scaling.}
\label{exp:adjust lambda}
As shown in Figure~\ref{fig:length ablation}, for BC2GM~\cite{smith2008overview} dataset, we conduct an ablation study by only converting a $\lambda$ during training.
Following previous work~\cite{fu-etal-2020-interpretable}, we adopt a fine-grained evaluation metric to partition the mention set based on an entity length.
In Figure~\ref{fig:ablation length of entities}, since biomedical entities are much longer than the general entities, we suggest the plot to select entity length in the range from 1 to 8.
As we can see, there is little change in the performance of the mention sets with a small length of entities.
On the other hand, in the long-named entities (E(L) $\geqq$ 8), we investigate that the F1 score improvement is up to 2\%.
Detail numerical values are reported in Appendix~\ref{app:lambda results}.

\subsection{Qualitative Analysis}
\label{sec:informative}
In this section, we give two intuitive examples to better understand our approach.
In Table~\ref{tab:qualitative example}, we show two examples of dealing with \textit{Out-Of-Vocabulary} issue and predicting long-named entity.
For example, the word \textit{'Latent'} does not appear in the training dataset but occurs in the evaluation dataset.
We thus cannot apply the class distribution of word frequency based on the word unit.
However, \textit{'Latent'} can be split into \textit{'Late'} and \textit{'\#\#nt'} using a subword tokenizer.
Considering the word \textit{'Late'} highly co-occurs with 'O' label, the bias-only model force the target model to predict without 'O' label.
It enables to capture the entity of \textit{'Latent BHRF1 transcripts'}.
By adopting bias-weighted scaling, our approach can capture long and complex structure entities.
It enables to predict \textit{foot - and - mouth - disease virus L protease} which includes conjunction (i.e., and) and special symbol (i.e., -).
In BioNER tasks, there are especially much more long entities exist but we only show a simpler case to give an intuition.

\section{Conclusion}
In this paper, we propose a bias-only model based on task-related subword statistics to alleviate the \textit{Out-Of-Vocabulary} issue for class imbalance tasks such as NER.
To this end, we adopt PMI by preventing solely focusing on dataset biases.
We argue that the bias-only model should reflect informative statistics via subword frequency and design a class prior.
Based on previous studies and our observations, we address long-named and complex structure entities with bias-weighted scaling.
Our experimental results also demonstrate the effectiveness and generalizability of the PMI.
To sum up, training a target model with three components can make a more robust model with a high capacity for generalization and memorization.


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

A major goal in the study of artificial and natural intelligence is to understand how language can scaffold more general problem-solving skills \citep[e.g.][]{spelke2017core}, and how these skills in turn shape language itself \citep[e.g.][]{gibson2017color}.
In NLP and machine learning, \textbf{latent language policies} \citep[LLPs;][]{andreas2017l3} provide a standard framework for studying these questions.
An LLP consists of \emph{instructor} and \emph{executor} subpolicies: the instructor generates natural language messages (e.g.\ high-level commands or subgoals), and the executor maps these messages to sequences of low-level actions (\cref{fig:teaser}).
LLPs have been used to construct interactive agents capable of complex reasoning (e.g.\ programming by demonstration) and planning over long horizons \citep[e.g.\ in strategy games;][]{minirts}.
They promise an effective and interpretable interface between
planning and control.

\begin{figure}[t!]
    \footnotesize
    \includegraphics[width=\columnwidth]{figs/teaser}
    \\[-0.5em]
    \strut \hspace{21pt} (a) \hspace{36pt} (b) \hspace{37pt} (c) \hspace{46pt} (d)
    \vspace{-.5em}
    \caption{In latent language policies,
    instructor agents (a) send natural-language commands (b) to executor agents (c), which execute them in an interactive environment (d).
    Jointly trained instructor--executor pairs learn to use messages in ways inconsistent with their natural language meanings (top shows a real message--action pair from a model described in \cref{sec:minirts}). We show that \emph{multitask} training with a population of task-specific instructors stabilizes message semantics and in some cases improves model performance.}
    \label{fig:teaser}
    \vspace{-1em}
\end{figure}

However, they present a number of challenges for training. %
As LLPs employ a human-specified space of high-level commands, they must be initialized with human supervision, typically obtained by pretraining the executor.
On its own, this training paradigm restricts the quality of the learned executor policy to that exhibited in (possibly suboptimal) human supervision. For tasks like the real-time strategy game depicted in \cref{fig:teaser}, we would like to study LLPs trained via \textbf{reinforcement learning (RL)}, \emph{jointly} learning from a downstream reward signal, and optimizing both instructors and executors for task success rather than fidelity to human teachers.

Training LLPs via RL has proven difficult. Past work has identified two main
challenges: primarily, the LLP-specific problem of \textbf{semantic drift}, in which agents come to deploy messages in ways inconsistent with their original (natural language) meanings \cite{lewis2017deal,lee2019countering}; secondarily, the general problem of \textbf{sample inefficiency} in RL algorithms \citep{kakade2003sample, brunskill2013sample}. Model-free deep RL is particularly notorious for requiring enormous amounts of interaction with the environment \citep{munos2016safe, atari}. For LLPs to meet their promise as flexible, controllable, and understandable tools for deep learning, better approaches are needed to limit semantic drift and perhaps improve sample efficiency.

While semantic change is a constant and well-documented feature of human languages \cite{mcmahon1994understanding}, (human) word meanings are on the whole remarkably stable relative to the rate of change in the tasks for which words are deployed \cite{karjus2020communicative}.
In particular, disappearance of lexical items is mitigated by increased population size \citep{bromham2015rate} and increased frequency of use \cite{pagel2007frequency}.
Drawing on these facts about stabilizing factors in human language,
we hypothesize that training of machine learning models with latent language variables can be made more robust by incorporating a population of instructors with diverse communicative needs that exercise different parts of the lexicon.

We describe a multitask LLP training scheme in which task-specific instructors communicate with a shared executor.
We show that complex long-horizon LLPs can be effectively tuned via \emph{joint reinforcement learning} of instructors and executors using \emph{multitask training}:

\begin{itemize}
    \item
        \cref{sec:signaling} presents a formal analysis of LLP training as an iterated Lewis signalling game \citep{lewis1969convention}. By modeling learning in this game as a dynamical system, we completely characterize a class of simple policies that are subject to semantic drift. We show that a particular multitask training scheme eliminates the set of initializations that undergo semantic drift.

    \item
        \cref{sec:minirts} evaluates the empirical effectiveness of multitask learning in a real-time strategy game featuring rich language, complex complex dynamics, and LLPs implemented with deep neural networks.
        Again, we show that multitask training reduces semantic drift (and improves sample efficiency) of LLPs in multiple game variants.
\end{itemize}
Together, these results show that diverse shared goals and communicative needs can facilitate (and specifically stabilize) learning of communication strategies.

\begin{figure*}
\centering
\includegraphics[height=1in,clip,trim=0 5.5in 0 0.2in]{figs/signaling}
\hfill
\includegraphics[height=1in,clip,trim=0.3in 4in 1.0in 0]{figs/signaling_payoffs}
\caption{A signaling game with two possible observations, two possible messages, and two possible actions. The instructor observes either a triangle or a square, then sends a message to a executor, who pushes either the red or blue buttons. The players' reward depends on the observation and the action but not the message. Two possible reward functions, $R$ and $R'$, are shown at right.}
\label{fig:signaling}
\end{figure*}

\section{Background and Related Work}
\label{sec:related}

Deep reinforcement learning (DRL) has recently made impressive progress on many challenging domains such as games \citep{mnih2013, silver2016}, locomotion \citep{schulman2015} and dexterous manipulation tasks \citep{gu2016,rajeswaran2017}. However, even state-of-the-art approaches to reinforcement struggle with tasks involving complex goals, sparse rewards, and long time horizons. A variety of models and algorithms for \emph{hierarchical} reinforcement learning have been proposed to address this challenge \citep{dayan1993feudal,dietterich2000hierarchical,sutton1999,bacon2017option} via supervised or unsupervised training of a fixed, discrete set of sub-policies.

Language can express arbitrary goals, and has compositional structure that allows generalization across commands. %
Building on this intuition, several recent papers have explored hierarchical RL in which natural language is used to parameterize the space of high-level actions \citep{oh2017zero,andreas2017,shu2017hierarchical,jiang2019language,minirts}. While there are minor implementation differences between all these approaches, we will refer to them collectively as \textbf{latent language policies (LLPs)}. Like other hierarchical agents, an LLP consists of a pair of subpolicies: an instructor $\speaker(\msg \mid \obs)$ and an executor $\listener(\action \mid \msg, \obs)$. An LLP takes actions by first sampling a string-valued message $\msg \sim \speaker$ from the instructor, and then an action $\action \sim \listener$ from the executor. For these messages to correspond to \emph{natural language}, rather than arbitrary strings, policies need some source of information about what human language users think they mean. This is typically accomplished by pretraining executors via human demonstrations or reinforcement learning; here we focus on the ingredients of effective joint RL of instructors and executors.

Reinforcement learning has been widely used to improve supervised language generation policies, particularly for dialogue \cite{li2016deep,lewis2017deal}, translation \cite{ranzato2015sequence,wu2016google} and summarization \cite{stiennon2020learning}.
Here, we instead focus on models where language is a latent variable as part of a hierarchical policy for a non-linguistic task.

As noted in \cref{sec:intro}, an observed shortcoming of reinforcement learning in all these settings is its susceptibility to \textbf{semantic drift}.
In the literature on human language change \citep{blank1999new}, semantic drift refers to a variety of phenomena, including specific terms becoming more general, general terms becoming specific, and parts coming to refer to wholes.
In machine learning, it refers broadly to the use of messages inconsistent with their natural language meanings in language-generation policies \citep{lazaridou2020multi}.

\citet{lee2019countering}
mitigate semantic drift in pivot-based machine translation by using visual grounding, whereas \citet{lu2020countering} periodically update a student model on data generated by an RL teacher.
Work in emergent communication has found that reinforcement learning tends not to learn policies with natural language-like properties \cite{kottur2017natural}, although population-based training has been found to be helpful \citep{gupta2019seeded}. Most relatedly to our work, \citet{lazaridou2020multi} train speaker-listener agents jointly in a visual referential communication task and introduce auxiliary loss functions for stabilizing training. Our work focuses on a more general setting where the interactions are temporally extended, have large action spaces and is partially observable.
\citet{agarwal2018community} use populations of agents to reduce semantic drift in visual dialogue. We view the current paper's analysis of multitask learning as complementary to these approaches from the emergent communication literature; future work might consider ways of combining the two.

A great deal of recent work in both RL \citep[e.g.][]{jaderberg2016reinforcement,shelhamer2016loss} and language processing \citep[e.g.][]{clark2019bam,gururangan2020don} has observed that carefully designed training objectives can serve as a source of model-agnostic inductive bias. Our results bring these two lines of work together: multitask training improves the faithfulness and adaptability of learned language understanding models, even when optimizing for a downstream reward.

\section{Multitask Communication in Theory: Lewis Signaling Games}
\label{sec:signaling}

We begin our analysis with the simple \textbf{signaling game} depicted in
\cref{fig:signaling}. In this game, one agent receives an observation, then
sends a message to another agent, which then performs an action.
Signaling games like this one are widely studied in NLP as models of reference
resolution and language generation \citep{frank2012predicting}. The instructor--executor pair may together be viewed as the simplest LLP of the kind described in \cref{sec:related}.

Formally, a (2-observation, 2-message) \textbf{Lewis signalling game} is defined by:
\begin{itemize}
    \item a set of \textbf{observations} $\obss = \{\obs_1, \obs_2\}$
    \item a set of \textbf{messages} $\msgs = \{\msg_1, \msg_2\}$
    \item a set of \textbf{actions} $\actions = \{\action_1, \action_2\}$
    \item a \textbf{reward function} $\reward : \obss \times \actions \to \reals$
\end{itemize}
The game is played between two agents: a \textbf{instructor} (with parameters $\sparam$),
which receives an observation and samples an observation-specific message from a distribution $\speaker(\msg \mid \obs; \sparam)$;
and a \textbf{executor} (with parameters $\lparam$), which receives the instructor's message and uses it to sample an action from a distribution $\listener(\action \mid \msg; \lparam)$. The agents then receive a reward $R(\obs, \action)$ that depends on the observation and action but not on the message sent.
 This policy's expected reward is given by:
\begin{equation}\label{eq:reward}
    \sum_{\substack{
      \obs \in \obss \\
      \msg \in \msgs \\
      \action \in \actions
    }}
    p(\obs)
    \speaker(\msg \mid \obs; \sparam)
    \listener(\action \mid \msg; \lparam)
    \reward(\obs, \action)
    ~ .
\end{equation}
Gradient ascent on \cref{eq:reward} with respect to $\sparam$ and $\lparam$ (e.g.\ using a policy gradient algorithm; \citeauthor{williams1992simple}, \citeyear{williams1992simple}) can be used to improve the expected reward obtained by an LLP.

As an example,
the right portion of \cref{fig:signaling} shows two reward functions $\reward$ and $\reward'$. In both, each observation is paired with a single action, and the executor must take the action corresponding to the observation to receive a positive reward. For $\reward$, two strategies obtain the optimal expected reward of 1: one in which $\speaker(\msg_1 \mid \blacktriangle) = 1$ and $\listener(\texttt{red} \mid \msg_1) = 1$, and one in which $\speaker(\msg_1 \mid \blacksquare) = 1$ and $\listener(\texttt{blue} \mid \msg_1) = 1$.
Almost every initialization of $\lparam$ and $\sparam$ (excluding a set of \emph{pooling equilibria}; see e.g.\ \citeauthor{huttegger2010evolutionary}, \citeyear{huttegger2010evolutionary}) converges to one of these two strategies when agents are jointly trained to optimize \cref{eq:reward}. %

\paragraph{Semantic drift}
Suppose, as shown in \cref{fig:signaling}, the messages $\msg_1$ and $\msg_2$ are not arbitrary symbols, but correspond to the natural language expressions $\msg_1 = $ \emph{push red} and $\msg_2 =$ \emph{push blue}.
In this case, only of the policies described above corresponds to the semantics of natural language---namely, the one in which $\listener(\action_1 \mid \msg_1) = \listener(\texttt{red} \mid \textit{push red}) = 1$. What is needed to ensure that a pair of agents playing converge to the natural language strategy?

In the analysis that follows, we will consider instructor and executor agents (each with a single scalar parameter $\in [0, 1]$):
\begin{align}
    \label{eq:speaker}
    \speaker(\msg_j \mid \obs_i; \sparam) = \begin{cases}
        \sparam & i = j \\
        1 - \sparam & \textrm{otherwise}
    \end{cases} \\
    \listener(\action_j \mid \msg_i) = \begin{cases}
        \lparam & i = j \\
        1 - \lparam & \textrm{otherwise}
    \end{cases}
\end{align}
For the game depicted in \cref{fig:signaling}, we would like to avoid any outcome in which, after training, $\lparam = \listener(\texttt{red} \mid \textit{push red})  < \half$. More generally, let us assume that we have an initial set of executor parameters that are possibly suboptimal but correspond to natural language semantics in the sense that $\listener(a \mid \msg_i; \lparam\att{0}) > \half$ if and only if the meaning of $\msg$ is \emph{do $a$}. In this case, we will say that a parameter initialization
$(\sparam\att{0}, \lparam\att{0})$ undergoes \textbf{executor semantic drift} if, after training, any such $\listener(\action \mid \msg_i; \lparam) = \lparam < \half$.

To analyze semantic drift in this game, we consider the final values of the parameters $(\sparam, \lparam)$ when optimized from an initialization $(\sparam\att{0}, \lparam\att{0})$.
For the reward function $R$ depicted in \cref{fig:signaling}, we can perform gradient ascent on \cref{eq:reward} with respect to $\sparam$ and $\lparam$ in this model by observing that:
\begin{align}
    \label{eq:gradients}
    \diffp{J}{{\sparam}} &= \lparam - \half  &
    \diffp{J}{{\lparam}} &= \sparam - \half
\end{align}

By considering the limiting behavior of gradient ascent as step size goes to zero (a \textbf{gradient flow}; see \cref{sec:appendixA}), it is possible to give a closed-form expression for the value of these parameters as a function of time:

\begin{prop}
\label{prop:single}
    Suppose $\lparam\att{0} + \sparam\att{0} < 1$. Then two agents optimizing \cref{eq:reward} via \cref{eq:gradients} undergo semantic drift (converging to $\lparam = 0$).
\end{prop}

Proof is given in \cref{sec:appendixA}. %
Note in particular that semantic drift will occur whenever $\sparam\att{0} < 1 - \lparam\att{0}$, which can occur even assuming a well-initialized executor with $\lparam > \half$. \cref{fig:flow} in the appendix provides a visualization of learning dynamics and these drift-susceptible initializations.
However, we will next show that this drift can be eliminated via multitask training.

\paragraph{Multitask signaling games}
Consider a \textbf{multitask} version of this game with the \textbf{two} reward functions $R$ and $R'$ depicted in \cref{fig:signaling}. As discussed in the introduction and depicted in \cref{fig:teaser}, our approach to multitask training focuses on sharing a single executor $\listener(\action \mid \msg; \lparam)$ between multiple task-specific instructors, here $\speaker(\msg \mid \obs; \soneparam)$ and $\speaker(\msg \mid \obs; \stwoparam)$, both parameterized as in \cref{eq:speaker}. As above, we train $(\soneparam, \stwoparam, \lparam)$ jointly to optimize:
\begin{align}\label{eq:joint-reward}
    \hspace{-.1em}&\sum_{\obs, \msg, \action}
    p(\obs)
    \speaker(\msg | \obs; \soneparam)
    \listener(\action | \msg; \lparam)
    \reward(\obs, \action) \nonumber\\
    \hspace{-.1em}&+
    \sum_{\obs, \msg, \action}
    p(\obs)
    \speaker(\msg | \obs; \stwoparam)
    \listener(\action | \msg; \lparam)
    \reward'(\obs, \action)  \hspace{-.1em}
\end{align}
We assume that the two instructors share the same initialization, with $\soneparam\att{0} = \stwoparam\att{0} = \sparam\att{0}$.
In this case, the following is true:
\begin{prop}
    \label{prop:multi}
    Suppose $\lparam\att{0} > \half$ and $\soneparam\att{0} = \stwoparam\att{0}$. Then three agents optimizing \cref{eq:joint-reward} via its gradient flow \emph{do not} undergo semantic drift. In fact, the eventual executor parameter $\lparam\att{t}$ is \emph{independent} of the initial speaker parameters $\soneparam\att{0}$ and $\stwoparam\att{0}$.
\end{prop}
Proof is again given in \cref{sec:appendixA}.
It is important to emphasize that these results concern the simplest possible policies for the signaling games considered here: agents with a single parameter which already ``bake in'' the assumption that different signals should trigger different behaviors.
We leave generalization of this formal analysis to general signaling games with more complex agents and message spaces for future work, noting that---at least in this simple case---we have succeeded in constructing a concrete multitask objective that reduces (indeed eliminates) the set of initial model parameters subject to semantic drift.

\section{Multitask Communication in Practice: The MiniRTS Environment}
\label{sec:minirts}

We next verify whether this result extends to the complex LLP-learning tasks discussed in \cref{sec:related}.
Our focus in this section is the \minirts environment of \citet{minirts} (depicted in \cref{fig:teaser}), in which agents must build and control an army of units like archers, spearmen, swordman, cavalry, and dragons, each with specialized abilities, with the goal of destroying the opponent's town center. Using this game, \citet{minirts} crowdsourced a dataset of high-level instructions (like \emph{attack with dragon} and \emph{send idle peasant to mine}) paired with low-level action sequences (\cref{fig:teaser}). They showed that an LLP trained on this supervised data via behavior cloning significantly outperformed a flat policy trained with imitation learning directly on low-level action sequences.

Here we investigate (1) whether these policy-cloned LLP can be further improved via reinforcement learning directly on a sparse win--loss signal from the game, (2) whether we can improve sample efficiency during reinforcement learning by jointly training executor models on multiple game variants simultaneously through multitask learning, and (3) whether semantic drift can be avoided during multi-task training. Below, \cref{sec:minirts-task}, \cref{sec:minirts-model} and \cref{sec:minirts-training} provide more detail about the task, model, and training procedure. \cref{sec:minirts-experiments} reports experimental results.

\subsection{Task and Training Data}
\label{sec:minirts-task}

\label{section:dataset}

\minirts is a partially-observable real-time strategy game environment, in which the actions of a large number of units must be coordinated on long time scales to defeat an opposing player. In a typical episode, a player must use its initial units to gather resources, use resources to build specialized structures for producing other units, and finally deploy these units to attack the opposing player's base. This involves challenging problems in both low-level tactics (controlling the placement of individual units for resource-gathering and combat) and high-level strategy (deciding which unit types to build, and when to deploy them).

\minirts additionally features a dataset collected from pairs of humans playing collaboratively against rule-based opponents. One human, the \emph{instructor}, designs high-level strategies and describes them in natural language. The other human, the \emph{executor} observes the environment state as well as the natural language strategy descriptions from the instructor and selects appropriate low-level actions.  The dataset consists of 5,392 games, with a total of 76,045 (instruction, execution) pairs.

\begin{figure*}
\centering
\includegraphics[scale=0.4]{imgs/obs_encoder.pdf}
\vspace{-.8em}

\caption{
State representations for the \minirts environment.
The model encodes the spatial observations, non-spatial observations, instructions, and auxiliary information about environment state at each timestep. These features are used by both the \emph{instructor} and \emph{executor} models. Reproduced with permission from \citet{minirts}.}
\label{fig:model}
\end{figure*}

\subsection{Model}
\label{sec:minirts-model}

\citet{minirts} use the labeled data to train an LLP for the \minirts environment. Our experiments use the same model architecture (\cref{fig:model}), which we briefly review here; see the original for details.

\paragraph{Observation encoder}
\label{section:encoder}
The \emph{instructor} and \emph{executor} models condition on a fixed-sized
  representation of the current game state which are constructed using different
  encoders for various aspects of the game state (\cref{fig:model}):
\begin{itemize}
    \item \emph{Spatial input encoder}: The spatial information of the map is encoded using a convolutional neural network.

    \item \emph{Non-spatial input encoder}:
    The non-spatial attributes and internal state of game objects are encoded using a simple MLP. These include attributes like the number of enemy units, the agent's units, and resource locations.

    \item \emph{Instruction encoder}:
    The current instruction is encoded with a recurrent neural network.

    \item \emph{Auxiliary encoder}:
    Global variables, such as the total number of resources collected, are additionally encoded with an MLP.
\end{itemize}

\paragraph{Instructor model}

The \emph{instructor} takes in the game state from the observation encoder and produces instructions. The 500 instructions appearing most frequently in the training set are encoded with an RNN into a fixed-sized vector. The score for each instruction is proportional to its dot product with the game state encoding. This instructor model achieved the best performance on several metrics in the original work \citep{minirts}. By restricting the instructor to the most frequent 500 well-formed natural language strings, we are able to focus our attention on \textbf{semantic} drift. A generative model free to generate arbitrary strings might also be subject to \textbf{syntactic} drift.

\paragraph{Executor model}
\label{section:executor}

The \emph{executor} predicts an action for every unit controlled by the agent based on of the current observation as encoded by the various encoders. The executor then predicts an action based on these features. In particular, for each unit, it predicts one of the 7 action types (\textsc{Idle, Continue, Gather, Attack, Train Unit, Build Building, Move}), an action target location (for the \textsc{Move, Attack, Gather} and \textsc{Build Building} actions) and a unit type (for the \textsc{Train Unit} and \textsc{Build Building} actions). %
Taking the product of the factorized action and location arguments across all units, the action space for the executor can be enormous, with as many as $10^{24}$ distinct actions available on a single turn.

\subsection{Training}
\label{sec:minirts-training}

As mentioned above, the original work of \citet{minirts} (and other work on learning LLPs) focused on behavior cloning or independent supervision of the \emph{instructor} and \emph{executor}. In the current paper, we are interested in the the dynamics of joint reinforcement learning of LLPs in both single- and multitask settings. Experiments in \cref{sec:minirts-experiments} make use of models trained with all three strategies.
\paragraph{Rule-based opponent pool}
Agents are trained against a similar pool of rule-based bots \citep[see][]{minirts} used to collect the human data. These bots follows a randomly selected, unit-specific strategy, building a fixed number of \textsc{swordmen, spearmen, cavalry, archers} or \textsc{dragons} and attacking as soon as they are constructed.

\paragraph{Behavior cloning}
Behavior-cloned models are trained using the supervised \minirts dataset. Given a collection of game observations $\obs$, each annotated with a high-level action $\msg$ and a low-level action $\action$, we maximize:
\begin{align}
    \max_{\sparam, \lparam} ~ &\sum_{\obs, \msg, \action} \big[ \log \speaker(\msg \mid \obs; \sparam) \nonumber \\
    \label{eq:behavior-cloning}
    & \qquad + \log \listener(\action \mid \msg, \obs; \lparam) \big] ~ .
\end{align}
During training, one frame is taken from every $K$ frames to form the supervised learning dataset.
To preserve unit level actions for the \emph{executor} training, all actions that happen in $[tK, (t+1)K)$ frames are stacked onto the $tK$th frame.

\paragraph{Reinforcement learning}

To train agents via reinforcement learning, we initialize them with the behavior cloning objective in \cref{eq:behavior-cloning} to provide an initial, human-meaningful grounding of message semantics (analogous to the initialization of the executor parameter $\lparam\att{0}$ in \cref{sec:signaling}). We then fine-tune them on game success, providing a sparse reward of 1 when agents win the game, -1 when they lose or draw the game.

As in \cref{sec:signaling}, learned agents are trained on the game reward, using a \textbf{proximal policy optimization} \textbf{(PPO)} \cite{ppo} objective to optimize the expected reward: %
\begin{equation}
    \expectation_{(s, a) \sim (\speaker, \listener)} R(s, a) ~ .
\end{equation}

\paragraph{Multi-task RL}
The original game of \citet{minirts} is defined by a set of \textbf{attack multipliers}: the aforementioned rock-paper-scissors dynamic arises because spearmen are especially effective against cavalry, archers against dragons, etc. To create alternative ``tasks'' in the MiniRTS environment, we create alternative versions of the game featuring different multipliers: e.g.\ making dragons invulnerable to archers or cavalry extra-effective against swordsmen. \cref{tab:original_rule} shows these multipliers for the original rule, and a set of game \textbf{variants} with different multipliers are described in \cref{sec:appendixB}. These variants are labeled B--J in the experiments that follow. Multiplier changes have significant effects on the optimal high-level strategy, affecting both which units are most effective overall, and how players should respond to opponents' choices.

\begin{table}[b!]
\centering
\footnotesize
\begin{adjustbox}{width=\columnwidth}
\begin{tabular}{|l||c|c|c|c|c|}
\hline
& \multicolumn{5}{|c|}{\textbf{Attack multiplier}} \\
\hline
\hline
\textbf{Unit name} & \textbf{Swordman} & \textbf{Spearman} & \textbf{Cavalry} & \textbf{Archer} & \textbf{Dragon} \\
\hline
\textsc{swordman} & 1.0 & 1.5 & 0.5 & 1.0 & 0.0 \\
\hline
\textsc{spearman} & 0.5 & 1.0 & 1.5 & 1.0 & 0.0 \\
\hline
\textsc{cavalry} & 1.5 & 0.5 & 1.0 & 1.0 & 0.0 \\
\hline
\textsc{archer} & 0.5 & 0.5 & 0.5 & 0.5 & 2.0 \\
\hline
\textsc{dragon} & 1.0 & 1.0 & 1.0 & 0.5 & 1.0 \\
\hline
\end{tabular}
\end{adjustbox}
\caption{Attack multipliers for the \textbf{original} game rules. For example, cavalry are extra-effective against swordsmen (1.5 in Swordsman col.); only archers and dragons can attack dragons (nonzero entries in Dragon col.). See \cref{sec:appendixB} for other game variants' multipliers.}
\label{tab:original_rule}
\end{table}

As in \cref{sec:signaling}, we perform multitask LLP training in the \minirts environment by jointly optimizing expected reward across multiple game variants at once, assigning each variant its own set of \emph{instructor} parameters $\sparam$ (initialized to the same value) but sharing a single set of executor parameters $\lparam$ across all contexts. The training pseudo-code can be found in \cref{sec:appendixPseudoCode}.

\subsection{Experiments}
\label{sec:minirts-experiments}

 Unlike in the signaling game considered in \cref{sec:signaling}, \minirts is complex, and we cannot take for granted that reinforcement learning of LLPs (with either ordinary or multitask objectives) will converge to an improved good solution at all.
 We thus begin with an analysis of policy performance and sample efficiency, then conclude this section with an analysis of semantic drift. (Model training details can be found in \cref{sec:appendixC}.)

\subsubsection{Evaluating performance and sample efficiency}
\begin{table}[b!]
    \centering
    \begin{adjustbox}{width=\columnwidth}
    \footnotesize

    \begin{tabular}{lccc}

    \toprule
    \textbf{Training} & \textbf{Evaluation} & \multicolumn{2}{c}{\textbf{Win rate}}  \\
    \textbf{strategy} & \textbf{environment} &  {\scriptsize (standard)} & {\scriptsize (3$\times$ training)} \\
    \midrule
    BC & \multirow{3}{*}{original} & 30.3  & -\\
    RL\ssjoint[orig.] &  & 65.7 & 86.9\\
    RL\ssmulti[orig., B, C] &  & \bf 76.5 & 90.6\\
    \midrule
    BC & \multirow{3}{*}{variant G} & 11.6 & -\\
    RL\ssjoint[G] & & 73.0 & 74.1\\
    RL\ssmulti[G, H, J] & & \bf 75.7 & \bf 77.6\\
    \midrule
    BC & \multirow{3}{*}{variant H} & 26.2 & -\\
    RL\ssjoint[H] & & 82.2 & \bf 91.4 \\
    RL\ssmulti[G, H, J] & &  79.4 & 83.5\\
    \midrule
    BC &  & 14.6 & -\\
    RL\ssjoint[J] & variant J & 87.2 & 93.0 \\
    RL\ssmulti[G, H, J] &  &  \bf 91.2 & 93.7\\
    \bottomrule
    \end{tabular}
    \end{adjustbox}
    \caption{Evaluation of policy quality in \minirts. Policies are evaluated against a rule-based opponent pool in four environments: \minirts with original rules, and three rule variants described in \cref{sec:appendixB}. We compare the original behavior-cloned LLP of \citet{minirts} (BC) with one fine-tuned directly on the evaluation environment (RL\ssjoint[env]) and one with multitask tuning on the evalution environment and two others (RL\ssmulti[env1, env2, env3]). Both RL fine-tuning strategies significantly outperform their behavior-cloned initializer. When using the same number of game episodes per training environment, RL\ssmulti is generally best but when RL\ssjoint is provided additional budget, it sometimes beats RL\ssmulti. Differences between models in environments original, G and J are significant ($p < 0.05$ under a permutation test).}
    \label{tab:performance}
\end{table}
To evaluate policy quality and sample complexity, we compare the final win rate (against the fixed pool of rule-based agents) for the policy-cloned (BC), RL-tuned (RL\ssjoint), and multitask-RL-tuned (RL\ssmulti) agents described above. We perform this evaluation for multiple game configurations: \textbf{original}, with the same rules used by \citet{minirts} for human data collection and evaluation, and 3 alternative variants (\textbf{variant G}, \textbf{variant H}, \textbf{variant J}) , in which the relative strengths of various units has been modified (see \cref{sec:appendixB}). We train 4 separate (RL\ssjoint) agents  corresponding to each of the environments and 2 (RL\ssmulti) agents. Following \citet{d2019sharing}, we provide both training strategies with a fixed budget of training experience \emph{across environments}: both RL\ssjoint and RL\ssmulti have been trained on the same number of game episodes per training environment. We also present the win rates for RL\ssjoint and RL\ssmulti, when trained on $3 \times$ more episodes per environment.

Results are shown in \cref{tab:performance}. Both RL fine-tuning strategies allow the policy to significantly improve over the behavior-cloned initializer, showing that effective reinforcement learning of LLPs is possible in \minirts. In most environments, performance of the model fine-tuned with the multi-task training is higher than ordinary joint RL training. When RL\ssjoint{} is provided extra training budget, it sometimes surpasses the performance of RL\ssmulti{} model with standard number of episodes. However, when RL\ssmulti{} is also given the extra training budget, it performs better in all but one environment.
At a high level, these results indicate that multitask training of LLPs can be applied at small (and in some cases no) cost in accuracy and significantly less per-environment training cost. %

\subsubsection{Evaluating semantic drift}
\begin{table}[b!]
    \centering
    \begin{adjustbox}{width=\columnwidth}
    \begin{tabular}{llccc}
    \toprule
     \multirow{2}{*}{\textbf{Instructor}} & \multirow{2}{*}{\textbf{Executor}} & \textbf{Eval.} & \multicolumn{2}{c}{\textbf{Win rate}} \\
                                 &                           & \textbf{env.} &         {\scriptsize (standard)}              & {\scriptsize (3$\times$ training)}\\
    \midrule
    \multirow{2}{*}{BC} & RL\ssjoint[orig.] & \multirow{2}{*}{original} & 48.4 & 59.2\\
                        & RL\ssmulti[orig., B, C] &  & \bf 60.7 & \bf 67.2\\
    \midrule
    \multirow{2}{*}{RL\ssspeaker[D]} & RL\ssjoint[orig.] & \multirow{2}{*}{variant D} & 74.8 & 85.9 \\
                                     & RL\ssmulti[orig., B, C] & & \bf 88.3 & 89.1\\
    \midrule
    \multirow{2}{*}{RL\ssspeaker[E]} & RL\ssjoint[orig.] & \multirow{2}{*}{variant E} & 57.5 & 68.7 \\
                                     & RL\ssmulti[orig., B, C] &  & \bf 72.9 & \bf 76.6\\
    \midrule
    \multirow{2}{*}{RL\ssspeaker[F]} & RL\ssjoint[orig.] & \multirow{2}{*}{variant F} & 73.2 & 83.8 \\
                                     & RL\ssmulti[orig., B, C] & & \bf 87.3 & \bf 92.3\\
    \bottomrule
    \end{tabular}
    \end{adjustbox}
    \caption{Evaluation of semantic drift in \minirts. Here, reinforcement-learned executor models are paired with instructors different from those they are trained with: either the original behavior-cloned instructor, or a instructor fine-tuned in an entirely different environment. The multitask executor RL\ssmulti{} performs better than RL\ssjoint{} when paired with new instructors, even when RL\ssjoint{} is given additional training budget. Differences in all environments are significant ($p < 0.05$ under a permutation test).}
    \label{tab:drift}
\end{table}
Next, we consider the second question from the introduction: outside of performance effects, does multitask training with populations of instructors reduce semantic drift in executors? We present two different quantitative evaluations that provide different perspectives on the answer to this question.
\paragraph{Semantic drift}
In \minirts, executor semantic drift occurs when the executor performs actions that are not consistent with the instruction produced by the instructor. ( i.e., \emph{create spearman} instruction produced by the instructor leads to the executor producing \emph{swordman} instead). In particular, this occurs in RL\ssjoint because the instructor and executor can co-adapt to new semantics during exploration as they are only trained to win the game.
\paragraph{Agent interoperability}
First, we evaluate the robustness of executor policies to alternative choices of instructors. Specifically, we pair each RL-trained executor with an instructor trained either via behavior cloning (and thus guaranteed to implement the human annotators' semantics) or fine-tuning (RL\ssspeaker) on a different game variant from the executor (and thus not co-adapted with it). Intuitively, to succeed at this task, executors must follow messages produced by the instructors trained in that domain. Executors that have undergone less semantic drift should perform better when paired with these different instructors.
Results are shown in \cref{tab:drift}; here, it can be seen that multitask learning matches or exceeds the performance of single-task training on this evaluation of semantic drift in all rule variants studied, even when RL\ssjoint\ is provided additional training budget. As evidence that performance comes from instructor--executor pairs, rather than executors alone, using a random coach paired with RL\ssmulti[orig., B, C] on variant D gives 33.2\% accuracy. Additionally, when RL\ssmulti[orig., B, C] is paired with a coach from a different variant, we get an accuracy of just 41\% on variant D.
\paragraph{Low-level action semantics}
As an alternative means of gaining insight into learned behaviors, we can directly inspect the correspondence between instructor messages and executor actions. We do this by uniformly sampling messages from a random instructor, then feeding them to the RL\ssmulti{} and RL\ssjoint{} executors and observing their choice of low-level actions $\action$. We then restrict these these $(\msg, \action)$ pairs to those in which (1) the text of $\msg$ includes one of the words \emph{create, build, train} or \emph{make} and the name of a unit (\emph{peasant}, \emph{spearman}, etc.) and (2) $\action$ is a \textsc{Train Unit} action for \emph{any} unit.
We then compute the empirical probability $P(\textrm{unit}_1 \in \action \mid \textrm{unit}_2 \in \msg )$ as shown in \cref{fig:action-drift}. If there is semantic drift, we expect to observe non-zero probability on the off-diagonal entries (the executor is building units different from those it is instructed to build). RL\ssmulti{} places less probability mass on the off-diagonal entries compared to RL\ssjoint{}, consistent with less semantic drift. In \cref{fig:action-drift}, one can also note that some word meanings change more than others. We hypothesize that, this is because like in natural languages, environmental pressures cause the meanings of some words to change at a greater rate than others. In this case, the dynamics of the game makes the \emph{spearman} unit slightly stronger than the \emph{swordman} unit overall. This results in unexpectedly good performance for players who accidentally misinterpret \emph{swordman} as \emph{spearman}. Therefore, this creates pressure for the conventional meaning of \emph{swordman} to shift more than other units.
\begin{figure}
    \centering
    \includegraphics[width=\columnwidth]{figs/heatmap-Original.jpg}
    \vspace{-1.8em}
    \caption{ Message--action drift in \minirts. Unit X on the x-axis indicates that a message $\msg$ of the form \emph{build unit X} was sampled. Unit Y on the y-axis shows the low-level \textsc{Train Unit Y} action $\action$ sampled by the executor. Matrix entries show the empirical probability $P(\action \mid \msg )$ computed using the relative frequencies of the sampled instructor messages and the corresponding executor actions. %
    The total sum of off-diagonal entries is 2.98 for RL\ssjoint[orig.] and 1.98 for RL\ssmulti, indicating less semantic drift for RL\ssmulti[orig., B, C].
    }
    \label{fig:action-drift}
\end{figure}

Taken together, these two evaluation results show that, when fine-tuning a policy initialized via imitation learning on the same objective, ordinary RL can be quite effective: the resulting executor model performs well even when paired with other instructors. But as in \cref{sec:signaling}, multitask training is even more helpful, especially by reducing semantic drift in both familiar and new environments.

\section{Conclusions}
We have presented a theoretical and empirical analysis of semantic drift and sample efficiency in multitask reinforcement learning of latent language policies (LLPs).
In a Lewis signaling game, we proved that multitask training can completely eliminate semantic drift. In a two-player real-time strategy game, we showed that multitask training is effective at mitigating semantic drift, improves the quality of learned policies and is sample efficient. Future work might integrate these results with other forms of population-based training (like those proposed by \citet{gupta2019seeded} for reference games) and explore other environmental factors affecting dynamics of language change in populations of learned agents.

\section*{Acknowledgments}
We thank Hengyuan Hu for assistance in reproducing the original work and MIT Supercloud for compute resources. We also thank Eric Chu, Alana Marzoev, Ekin Akyürek and Debby Clymer for feedback and proof-reading.

\bibliographystyle{acl_natbib}