UDA(Unsupervised Data Augmentation) with BERT

This is re-implementation of Google's UDA [paper][tensorflow] in pytorch with Kakao Brain's Pytorchic BERT[pytorch].

Model  UDA official This repository
UDA (X) 68%  
UDA (O) 90% 88.45%

(Max sequence length = 128, Train batch size = 8)


UDA(Unsupervised Data Augmentation) is a semi-supervised learning method which achieves SOTA results on a wide variety of language and vision tasks. With only 20 labeled examples, UDA outperforms the previous SOTA on IMDb trained on 25,000 labeled examples. (BERT=4.51, UDA=4.20, error rate)

  • Unsupervised Data Augmentation for Consistency Training (2019 Google Brain, Q Xie et al.)

- UDA with BERT

UDA works as part of BERT. It means that UDA act as an assistant of BERT. So, in the picture above model M is BERT.

- Loss

UDA consist of supervised loss and unsupervised loss. Supervised loss is traditional Cross-entropy loss and Unsupervised loss is KL-divergence loss of original example and augmented example outputs. In this project, I used Back translation technique for augmentation.
The supervised loss and unsupervised loss are added to form a total loss and then total loss is descent. To be careful is loss doesn't descent trough original example route. Only by labeled data and augmented unlabeled data Model's weights are updated.

- TSA(Training Signal Annealing)

There is a large gap between the amount of unlabeled data and that of labeled data. So, it is easy to overfit to labeled data. Therefore, TSA technique mask out the examples that predicted probability is bigger than threshold. The threshold is scheduled by log, linear or exponential function.

- Sharpening Predictions

The KL-divergence loss(ori, aug) is too small to just use. It can cause that the total loss is dominated by supervised loss. Therefore, Sharpening Prediction techniques is needed.


UDA : python > 3.6, fire, tqdm, tensorboardX, tensorflow, pytorch, pandas, numpy



- Download pre-trained BERT model and unzip IMDb data

First, you have to download pre-trained BERT_base from Google's BERT repository. And unzip IMDb data

bash download.sh

After running, you can get the pre-trained BERT_base_Uncased model at /BERT_Base_Uncased director and /data

I use already pre-processed and augmented IMDb data extracted from official UDA. If you want to use your raw data, change need_prepro = True.

Example usage

This project are broadly divided into two parts(Fine-tuning, Evaluation).
Caution : Before runing code, you have to check and edit config file

  1. Fine-tuning
    You can choose train mode(train, train_eval) on non-uda.json or uda.json (default : train_eval).

    • Non UDA fine-tuning

      python main.py \
          --cfg='config/non-uda.json' \
    • UDA fine-tuning

      python main.py \
          --cfg='config/uda.json' \
  2. Evaluation

    • Basically evaluation code, dump out results file. So, you can change dump option in main.py There is two mode (real_time print, make tsv file)

      python main.py \ --cfg='config/eval.json' \ --model_cfg='config/bert_base.json'


Thanks to references of UDA and Pytorchic BERT, I can implement this code.


  1. It is known that further training(more pre-training by the specific corpus on already pre-trained BERT) can improve performance. But, this repository does not have pretrain code. So, pretrain code will be added. If you want to further training you can use Pytorchic BERT 's pretrain.py or any BERT project.

  2. Korean dataset version