Poly-encoders

This repository is an unofficial implementation of Poly-encoders: Transformer Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring.

How to use

1. Download and unzip the ubuntu data https://www.dropbox.com/s/2fdn26rj6h9bpvl/ubuntudata.zip?dl=0

2. Prepare a pretrained BERT (https://github.com/huggingface/transformers)

3. pip3 install -r requirements.txt

4. Train a Poly-encoder:

   python3 train.py -bert_model /your/pretrained/model/dir --output_dir /your/ckpt/dir --train_dir /your/data/dir --use_pretrain --architecture poly --poly_m 16

5. Train a Bi-encoder:

   python3 train.py -bert_model /your/pretrained/model/dir --output_dir /your/ckpt/dir --train_dir /your/data/dir --use_pretrain --architecture bi

Results

The experimental settings and results are shown as follows:

• Dataset: Ubuntu
• Device: GTX 1060 6G x1
• Pretrained model: BERT-small-uncased (https://github.com/sfzhou5678/PretrainedLittleBERTs or https://storage.googleapis.com/bert_models/2020_02_20/all_bert_models.zip)
• Batch size: 32
• max_contexts_length: 128
• max_context_cnt: 4
• max_response_length：64
• lr: 5e-5
• Epochs: 3

Model	R@1/10	Training Speed	GPU Mem Consumption
Bi-encoder	0.6714	3.15it/s	1969 Mb
Poly-encoder 16	0.6938	3.11it/s	1975 Mb
Poly-encoder 64	0.7026	3.08it/s	2005 Mb
Poly-encoder 360	0.7066	3.05it/s	2071 Mb

Different with the original paper, this experiment uses a bert-small-uncased model (from https://github.com/sfzhou5678/PretrainedLittleBERTs or https://storage.googleapis.com/bert_models/2020_02_20/all_bert_models.zip) rather than the bert-base. Besides, this experiment only uses batch_size =32, max_length = 128, and max_history=4 (which means select up to 4 context texts). All these settings lead to lower results but faster training speed. One can modify these settings for a better result.

Some Improvements

• Thanks to @chijames, this implementation is closer to the original paper and has achieved better performance.

BTW, If you have any suggestions or questions, please feel free to reach me out!