This repository contains an implementation of the MiniLMv2 distillation technique. Please refer to the original paper cited below for more details about the method.

This code builds on the Hugging Face Transformers library.

• Quick Start
• Hyperparameters
• Citation
• Results
• License
• Code of Conduct
• Security Vulnerability Reporting

1. Collect the data which the teacher was trained on (e.g., for bert_base: Wikipedia and Bookcorpus).
2. Update minilmv2/train_config.json to point to your data. The path can be a local path or a publicly accessible HTTP URL.
3. In run.sh, update the input_model_dir to point to your teacher model (loadable with AutoModel.from_pretrained) and make any necessary changes to launch a distributed job. It is recommended to run this on multiple GPUs.
4. Make any necessary changes to save the student model after training (e.g., saving student state dict).
5. Start the distillation by running run.sh.

For distillation, the run.sh script contains the hyperparameters used to distill our models and can be modified for other settings as required. The run_eval.sh script contains the evaluation hyperparameters and we hypertune our models over different learning rates (1e-5 - 4e-5) and num epochs (3,5,10).

The MiniLMV2 technique was originally proposed by:

  @inproceedings{wang-etal-2021-minilmv2,
    title = "{M}ini{LM}v2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers",
    author = "Wang, Wenhui  and
      Bao, Hangbo  and
      Huang, Shaohan  and
      Dong, Li  and
      Wei, Furu",
    booktitle = "Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021",
    month = aug,
    year = "2021",
    address = "Online",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2021.findings-acl.188",
    doi = "10.18653/v1/2021.findings-acl.188",
    pages = "2140--2151",
}

The table below summarizes our results for two models (6x384 and 6x768) distilled from bert-base. GLUE results were obtained by running the hf_run_evaluation.py, which is a near replica of the run_glue script from Hugging Face.

Copyright 2022 Bloomberg Finance L.P. Licensed under the Apache License, Version 2.0.

This project has adopted a Code of Conduct. If you have any concerns about the Code, or behavior which you have experienced in the project, please contact us at [email protected].

If you believe you have identified a security vulnerability in this project, please send an email to the project team at [email protected] that details the suspected issue and any methods you've found to reproduce it.

Please do NOT open an issue in the GitHub repository, as we'd prefer to keep vulnerability reports private until we've had an opportunity to review and address them.