ScaleGrad

Source code for Straight to the Gradient: Learning to Use Novel Tokens for Neural Text Generation

Xiang Lin, Simeng Han and Shafiq Joty

Accepted at 38th International Conference on Machine Learning (ICML'21).

The repo is adapted from Neural Unlikelihood Training. You could either follow the original repo or the instruction below to finish installing the dependencies.

Setup

Dependencies

The implementation is a custom fairseq module. Download and install fairseq:

git clone https://github.com/pytorch/fairseq.git
cd fairseq
pip install --editable .

Install other dependencies:

pip install nltk
pip install pytorch-transformers  

'Installing' the a customised module for ScaeGrad

Copy the custom directory in this repo into the fairseq repo that you downloaded above:

export FAIRSEQ_DIR=/path/to/fairseq
export SCALEGRAD_DIR=/path/to/ScaleGrad

cp -r $SCALEGRAD_DIR/custom $FAIRSEQ_DIR/fairseq

Finetuning GPT-2 for Language Modeling and Auto Completion

We assume that you are in the fairseq base directory.

Dataset

Download and unpack the BPE-tokenized WikiText:

wget https://dl.fbaipublicfiles.com/unlikelihood/wikitext-103-bpe_v0.tar.gz
tar -xzvf wikitext-103-bpe_v0.tar.gz
mv wikitext-103-bpe_v0 data-bin/

Finetuning with ScaleGrad/MLE:

gamma is the only hyper-parameter needed to be pre-determined. gamma=1 is equivalent to performing MLE finetuning. In our paper, we experimented with three different choices from {0.2,0.5,0.8}.

python fairseq/custom/gpt2/run_gpt2.py  \
--data-base ./data-bin/wikitext-103-bpe_v0    \
--output-dir ./checkpoint/gpt-2   \
--eval-split valid    \
--train-n-steps 35000   \
--validate-every 1000    \
--mode train \
--train-batch-size 300    \
--gamma 0.2 \
--learning-rate 2e-5 \
--seed 60 

Evaluation

python fairseq/custom/gpt2/run_gpt2.py  \
    --data-base ./data-bin/wikitext-103-bpe_v0 \
    --output-dir ./checkpoint/gpt2/sg_output \
    --eval-split test \
    --model-load-dir ./checkpoint/gpt2/sg/best \
    --mode eval-both

We used a single GTX 2080Ti gpu.

Citation

Please cite our work if you found the resources in this repository useful:

@InProceedings{lin-2021-straight,
  title = 	 {Straight to the Gradient: Learning to Use Novel Tokens for Neural Text Generation},
  author =       {Lin, Xiang and Han, Simeng and Joty, Shafiq},
  booktitle = 	 {Proceedings of the 38th International Conference on Machine Learning},
  pages = 	 {6642--6653},
  year = 	 {2021},
  editor = 	 {Meila, Marina and Zhang, Tong},
  volume = 	 {139},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {18--24 Jul},
  publisher =    {PMLR},
  pdf = 	 {http://proceedings.mlr.press/v139/lin21b/lin21b.pdf},
  url = 	 {http://proceedings.mlr.press/v139/lin21b.html},
  abstract = 	 {}
}