Flambé: A Customizable Framework for Machine Learning Experiments

Abstract

Flambé is a machine learning experimentation framework built to accelerate the entire research life cycle. Flambé's main objective is to provide a unified interface for prototyping models, running experiments containing complex pipelines, monitoring those experiments in real-time, reporting results, and deploying a final model for inference. Flambé achieves both flexibility and simplicity by allowing users to write custom code but instantly include that code as a component in a larger system which is represented by a concise configuration file format. We demonstrate the application of the framework through a cuttingedge multistage use case: fine-tuning and distillation of a state of the art pretrained language model used for text classification. 1

Autoregressive Knowledge Distillation through Imitation Learning

“…All models are trained using the Adam optimizer (Kingma and Ba, 2014) with an inversesquare-root learning rate scheduler and learning rate warmup (Vaswani et al, 2017). Our experiments were conducted using Flambé, a PyTorch-based model training and evaluation library (Wohlwend et al, 2019). More implementation details such as hyperparameter settings are provided in Appendix A.…”