Effective Unsupervised Domain Adaptation with
Adversarially Trained Language Models

The masking strategy learning problem described in eqn (2) is a minimax game of two players: the puzzle generator to select the subset resulting in the most challenging puzzle, and the MLM $\mathcal{B}_{\theta}$ to best solve the puzzle by reconstructing the masked tokens correctly. As optimising over the subsets is a hard combinatorial problem over the discrete space of $\mathcal{S}_{K}$, we are going to convert it to a continuous optimisation problem.

We establish a variational lower bound of the objective function over $S$ using the following inequality,

where $q(.)$ is the variational distribution provided by a neural network $\pi_{\phi}$. This variational distribution $q(S|\boldsymbol{x};\pi_{\phi})$ estimates the distribution over all subset of size $K$. It is straightforward to see that the weighted sum of negative log likelihood of all possible subsets is always less than the max value of them. Our minimax training objective is thus,

where $\mathcal{Z}$ is the partition function making sure the probability distribution sums to one,

The number of possible subsets is $|\mathcal{S}_{K}|=\binom{|\boldsymbol{x}|}{K}$, which grows exponentially with respect to $K$. In §4, we provide efficient dynamic programming algorithm for computing the partition function and sampling from this exponentially large combinatorial space. In the following, we present our model architecture and training algorithm for the puzzle generator $\phi$ and MLM $\theta$ parameters based on the variational training objective in eqn (5).

We learn the masking strategy through the puzzle generator network as shown in Figure 1. It is a feed-forward neural network assigning a selection probability $\pi_{\phi}(i|\boldsymbol{x})$ for each index $i$ given the original sentence $\boldsymbol{x}$, where $\phi$ denote the parameters. Inputs to the puzzle generator are the feature representations $\{\boldsymbol{h}_{i}\}_{i=1}^{n}$ of the original sequence $\{\boldsymbol{x}_{i}\}_{i=1}^{n}$. More specifically, they are output of the last hidden states of the MLM. The probability of perform masking at position $i$ is computed by applying sigmoid function over the feed-forward net output $\pi_{\phi}(i|\boldsymbol{x})=\sigma(\textrm{FFNN}(\boldsymbol{h}_{i}))$. From these probabilities, we can sample the masked positions in order to further train the underlying MLM $\mathcal{B}_{\theta}$.

We use an alternating optimisation algorithm to train the MLM $\mathcal{B}_{\theta}$ and the puzzle generator $\pi_{\phi}$ (Algorithm 1). The update frequency for $\pi_{\phi}$ is determined via a mixing hyperparameter $\beta$.

In order to sample from the variational distribution in eqn (6), we need to compute its partition function in eqn (7). Interestingly, the partition function can be computed using dynamic programming (DP).

Let us denote by $Z(j,k)$ the partition function of all subsets of size $k$ from the index set $\{j,..,|\boldsymbol{x}|\}$. Hence, the partition function of the $q$ distribution is $Z(1,K)$. The DP relationship can be written as,

The initial conditions are $Z(j,0)=1$ and

corresponding to two special terminal cases in selection process in which we have picked all $K$ indices, and we need to select all indices left to fulfil $K$.

This amounts to a DP algorithm with the time complexity $\mathcal{O}(K|\boldsymbol{x}|)$.

The DP in the previous section also gives rise to the sampling procedure. Given a partial random subset $S_{j-1}$ with elements chosen from the indices $\{1,..,j-1\}$, the probability of including the next index $j$, denoted by $q_{j}(\textrm{yes}|S_{j-1},\pi_{\phi})$, is

where $Z(j,k)$ values come from the DP table. Hence, the probability of not including the index $j$ is

In case the next index is chosen to be in the sample, then $S_{j+1}=S_{j}\cup\{j+1\}$; otherwise $S_{j+1}=S_{j}$.

The sampling process entails a sequence of binary decisions (Figure 1.b) in an underlying Markov Decision Process (MDP). It is an iterative process, which starts by considering the index one. At each decision point $j$, the sampler’s action space is to whether include (or not include) the index $j$ into the partial sample $S_{j}$ based on eqn (13). We terminate this process when the partially selected subset has $K$ elements.

The sampling procedure is described in Algorithm 2. In our MDP, we actually sample an index by generating Gumbel noise in each stage, and then select the choice (yes/no) with the maximum probability. This enables differentiation through the sampled subset, covered in the next section.

Once the sampling process is terminated, we then need to backpropagate through the parameters of $\pi_{\phi}$, when updating the parameters of the puzzle generator according to eqn (10).

More concretely, let us assume that we would like to sample a subset $S$. As mentioned in previous section, we need to decide about the inclusion of the next index $j$ given the partial sample so far $S_{j-1}$ based on the eqn (13). Instead of uniform sampling, we can equivalently choose one of these two outcomes as follows

where the random noise $\epsilon_{o_{j}}$ is distributed according to standard Gumbel distribution. Sampling a subset then amounts to a sequence of $\operatorname*{argmax}$ operations. To backpropagate through the sampling process, we replace the $\operatorname*{argmax}$ operators with softmax, as $\operatorname*{argmax}$ is not differentiable. That is,

The $\log$ product of the above probabilities for the decisions in a sampling path is returned as $l$ in Algorithm 2, which is then used for backpropagation.

We compare our adversarial learned masking strategy approach against random and various heuristic masking strategies which we propose:

• •

  Random. Masked tokens are sampled uniformly at random, which is the common strategy in the literature (Devlin et al., 2019; Liu et al., 2019).

• •

  POS-based strategy. Masked tokens are sampled according to a non-uniform distribution, where a token’s probability depends on its POS tag. The POS tags are obtained using spaCy.³³3https://spacy.io/ Content tokens such as verb (VERB), noun (N), adjective (ADJ), pronoun (PRON) and adverb (ADV) tags are assigned higher probability (80%) than other content-free tokens such as PREP, DET, PUNC (20%).

• •

  Uncertainty-based strategy. We select those tokens for which the current MLM is most uncertain for the reconstruction, where the uncertainty is measured by the entropy. That is, we aim to select those tokens with high $\textrm{Entropy}[Pr_{i}(.|\boldsymbol{x}_{\bar{S_{i}}};\mathcal{B}_{\theta})]$, where $\boldsymbol{x}_{\bar{S_{i}}}$ is the sentence $\boldsymbol{x}$ with the $i$th token masked out, and $Pr_{i}(.|\boldsymbol{x}_{\bar{S_{i}}};\mathcal{B}_{\theta})$ is the predictive distribution for the $i$th position in the sentence.

  Calculating the predictive distribution for each position requires one pass through the network. Hence, it is expensive to use the exact entropy, as it requires $|\boldsymbol{x}|$ passes. We mitigate this cost by using $Pr_{i}(.|\boldsymbol{x};\mathcal{B}_{\theta})$ instead, which conditions on the original unmasked sentence. This estimation only costs one pass through the MLM.

• •

  Adversarial learned strategy. The masking strategy is learned adversarially as in §3. The puzzle-generator update frequency $\beta$ (Algorithm 1) is set to 0.3 for all experiments.

These strategies only differ in how we choose the candidate tokens. The number of to-be-masked tokens is the same in all strategies (15%). Among them, 80% are replaced with [MASK], 10% are replaced with random words, the rest are kept unchanged as in Devlin et al. (2019). In our experiments, the masked sentences are generated dynamically on-the-fly.

To evaluate the models, we compute precision, recall and F1 scores on a per token basis. We report average performance of five runs.

Our implementation is based on Tensorflow library (Abadi et al., 2016)⁴⁴4Source code is available at https://github.com/trangvu/mlm4uda. We use BERT-Base model architecture which consists of 12 Transformer layers with 12 attention heads and hidden size 768 (Devlin et al., 2019) in all our experiments. We use the cased wordpiece vocabulary provided in the pretrained English model. We set learning rate to 5e-5 for both further pretraining and task tuning. Puzzle generator is a two layer feed-forward network with hidden size 256 and dropout rate 0.1.

Under the same computation budget to update the MLM, we evaluate the effect of masking strategy in the domain tuning step under various size of additional target-domain data: none, 500K and 1M. We continue pretraining BERT on a combination of unlabeled source (CoNLL2003), unlabeled target task training data and additional unlabeled target domain data (if any). If target task data is smaller, we oversample it to have equal size to the source data. The model is trained with batch size 32 and max sequence length 128 for 50K steps in 1M target-domain data and 25K steps in other cases. It equals to 3-5 epochs over the training set. After domain tuning, we finetune the adapted MLM on the source task labeled training data (CoNLL2003) for three epochs with batch size 32. Finally, we evaluate the resulting model on target task. On the largest dataset, random and POS strategy took around 4 hours on one NVIDIA V100 GPU while entropy and adversarial approach took 5 and 7 hours respectively. The task tuning took about 30 minutes.

Results are shown in Table 2. Overall, strategically masking consistently outperforms random masking in most of the adaptation scenarios and target tasks. As expected, expanding training data with additional target domain data further improves performance of all models. Comparing to random masking, prioritising content tokens over content-free ones can improve up to 0.7 F1 score in average. By taking the current MLM into account, uncertainty-based selection and adversarial learned strategy boost the score up to 1.64. Our proposed adversarial approach yields highest score in 11 out of 18 cases, and results in the largest improvement over random masking across all tasks in both UDA with and without additional target domain data.

We further explore the mix of random masking and other masking strategies. We hypothesise that the combination strategies can balance the learning of challenging tokens and effortless tokens when forming the common semantic space, hence improve the task performance. In a minibatch, 50% of sentences are masked according to the corresponding strategy while the rest are masked randomly. Results are shown in Table 3. We observe an additional performance to the corresponding single-strategy model across all tasks.