Package: nlp-uncertainty-zoo
:robot::speech_balloon::question: nlp-uncertainty-zooThis repository contains implementations of several models used for uncertainty estimation in Natural Language processing,
implemented in PyTorch. You can install the repository using pip:pip3 install nlp-uncertainty-zooIf you are using the repository in your academic research, please cite the paper below:@inproceedings{ulmer-etal-2022-exploring,
  title = "Exploring Predictive Uncertainty and Calibration in {NLP}: A Study on the Impact of Method {\&} Data Scarcity",
  author = "Ulmer, Dennis  and
    Frellsen, Jes  and
    Hardmeier, Christian",
  booktitle = "Findings of the Association for Computational Linguistics: EMNLP 2022",
  month = dec,
  year = "2022",
  address = "Abu Dhabi, United Arab Emirates",
  publisher = "Association for Computational Linguistics",
  url = "https://aclanthology.org/2022.findings-emnlp.198",
  pages = "2707--2735",}To learn more about the package, consult the documentationhere,
check a Jupyter notebook demohereor a Google
collabhere.Included modelsThe following models are implemented in the repository. They can all be imported by usingfrom nlp-uncertainty-zoo import <MODEL>.
For transformer-based model, furthermore a version of a model is available that uses a pre-trained BERT from the HuggingFacetransformers.NameDescriptionImplementationPaperLSTMVanilla LSTMLSTMHochreiter & Schmidhuber, 1997LSTM EnsembleEnsemble of LSTMsLSTMEnsembleLakshminarayanan et al., 2017Bayesian LSTMLSTM implementing Bayes-by-backpropBlundell et al, 2015BayesianLSTMFortunato et al, 2017ST-tau LSTMLSTM modelling transitions of a finite-state-automatonSTTauLSTMWang et al., 2021Variational LSTMLSTM with MC Dropout(Gal & Ghahramani, 2016a)VariationalLSTMGal & Ghahramani, 2016bDDU Transformer, DDU BERTTransformer / BERT with Gaussian Mixture Model fit to hidden featuresDDUTransformer,DDUBertMukhoti et al, 2021Variational Transformer, Variational BERTTransformer / BERT with MC Dropout(Gal & Ghahramani, 2016a)VariationalTransformer,VariationalBertXiao et al., 2021DPP Transformer, DPP BertTransformer / BERT using determinantal point process dropoutDPPTransformer,DPPBertShelmanov et al., 2021SNGP Transformer, SNGP BERTSpectrally-normalized transformer / BERT using a Gaussian Process output layerSNGPTransformer,SNGPBertLiu et al., 2022Contributions to include even more approaches are much appreciated!UsageEach model comes in two versions, for instanceLSTMEnsembleandLSTMEnsembleModule. The first one is supposed to be
used as an out-of-the-box solution, encapsulating all training logic and convenience functions. These include fitting
the model, prediction, getting the uncertainty for an input batch using a specific metric.model=LSTMEnsemble(**network_params,ensemble_size=10,is_sequence_classifer=False)model.fit(train_split=train_dataloader)model.get_logits(X)model.get_predictions(X)model.get_sequence_representation(X)model.available_uncertainty_metricsmodel.get_uncertainty(X)model.get_uncertainty(X,metric_name="mutual_information")In comparison, the-Moduleclass is supposed to me more simple and bare-bones, only containing the core model logic.
It is intended for research purposes, and for others who would like to embed the model into their own code base. While
the model class (e.g.LSTMEnsemble) inherits fromModeland would require to implement certain methods, anyModuleclass
sticks closely totorch.nn.Module.To check what arguments are required to initialize and use different models, checkthe documentation here.Also, check out the demo provided as a Jupyter notebookhereor a Google
collabhere.Repository structureThe repository has the following structure:models: All model implementations.tests: Unit tests. So far, only contains rudimentary tests to check that all output shapes are consistent between models and functions.utils: Utility code (see below)utils/custom_types.py: Custom types used in the repository for type annotations.utils/data.py: Module containing data collators, and data builders - which build the dataloaders for a type of task and a specific dataset. Currently, language modelling, sequence labeling and sequence classification are supported.utils/metrics.py: Implementations of uncertainty metrics.utils/samplers.py: Dataset subsamplers for language modelling, sequence labelling and sequence classification.utils/task_eval.py: Functions used to evaluate task performance.utils/uncertainty_eval.py: Function used to evaluate uncertainty quality.utils/calibration_eval.py: Function used to evaluate calibration quality.config.py: Define available datasets, model and tasks.defaults.py: Define default config parameters for sequence classification and language modelling (Note: These might not be very good parameters).Other featuresWeights & Biases integration: You can track your experiments easily with weights & biases by passing awandb_runargument tomodel.fit()!Easy fine-tuning via HuggingFace: You can fine-tune arbitrary BERT models using their name from HuggingFace'stransformers.ContributingThis repository is by no means perfect nor complete. If you find any bugs, please report them using the issue template,
and, if you also happen to provide a fix, create a pull request! A GitHub template is provided for that as well.You would like to make a new addition to the repository? Follow the steps below:Adding a new model: To add a new model, add a new module in themodelsdirectory. You will also need to implement
a correspondingModelandModuleclass, inheriting from the classes of the same name inmodels/model.pyand implementing all
required functions.Modelis supposed to be an out-of-the-box solution that you can start experimenting right away, whilModuleshould only include the most basic model logic in order to be easy to integrate into other codebases and allow tinkering.Adding a new uncertainty metric: To add a new uncertainty metric, add the function toutils/metrics.py. The function should take
the logits of a model and output an uncertainty score (the higher the score, the more uncertain the model). The function should output
a batch_size x sequence_length matrix, with batch_size x 1 for sequence classification tasks. After finishing the implementation, you can
add the metric to thesingle_prediction_uncertainty_metricsof themodels.model.Modelclass andmulti_prediction_uncertainty_metricsofmodels.model.MultiPredictionMixin(if applicable).You would like to add something else? Create an issue or contact me at dennis {dot} ulmer {at} mailbox {dot} org!