Package: autocontrastive-gen
Auto-Contrastive GenerationCode to experiment withmulti-exittext generation, and to reproduce theAuto-Contrastive Decodingexperiments fromGera et al. (2023).Using this library you can:Run inference on multi-exit generative language models, either using aspecific model exit layer, or contrasting between model layers withAuto-Contrastive Decoding;Train new multi-exit models, either for language modeling or for a specific task;Try out new methods and algorithms for combining and contrasting the outputs of different model layers.Table of contentsQuick startSetting the Multi-Exit ConfigurationRunning language modeling benchmarksPre-trained model checkpointsReferenceLicenseQuick startInstall the library withpip install autocontrastive-genChoose your desiredmulti-exit parameters- which model exit layer(s) do you want to use, and how?Load a pre-trained multi-exit model and use it however you wish, within your own workflowFor instance, the following code will initialize the multi-exitGPT2 model we releaseto use Auto-Contrastive Decoding:fromautocontrastive_gen.modeling.configurationimportMultiExitConfigurationfromautocontrastive_gen.modeling.auto_modelimportAutoMultiExitModel# initialize a pre-trained multi-exit model to use auto-contrast between layer 24 and layer 12multi_exit_config=MultiExitConfiguration(use_original_head=False,contrast_layer_indices=(24,12))model=AutoMultiExitModel.from_pretrained("IBM/gpt2-medium-multiexit",multi_exit_config=multi_exit_config)Then, the initialized model can be used as usual through thetransformerslibrary. For example, this code will run text generation using the initialized model:fromtransformersimportAutoTokenizertokenizer=AutoTokenizer.from_pretrained("IBM/gpt2-medium-multiexit")prompt=tokenizer("humpty dumpty sat on",return_tensors='pt')generated_ids=model.generate(**prompt,max_new_tokens=15)print(tokenizer.batch_decode(generated_ids))Similarly, you cantrain the modelwith thetransformerslibrary just as you would do with any other model.Note that the model behavior in inference and/or training depends on yourchoice of parameterswhen initializing theMultiExitConfiguration.Setting the Multi-Exit ConfigurationModel behavior is determined by theMultiExitConfigurationused to initialize it. Most of the the configuration parameters are related to inference-time text generation behavior, but some are relevant for model training as well. The following parameters can be set:lm_head_layer_indices: Tuple[int, ...]: the indices of model layers which are connected to language modeling exit heads. As this is a basic characteristic of the model, this parameter only needs to be setoncefor initial pre-training of these exit heads.Otherwise (e.g., when using one of thereleased models), there is no need to specify this parameter as it is read directly from the model's config file.freeze_parameters: bool: whether to freeze the language model parameters when training the model. You may wish to set this to True if training new exit layers for an existing (single-exit) pre-trained model checkpoint, but the default (False) is applicable for most use cases.use_original_head: bool: whether to use the original language modeling head of the pre-trained checkpoint for text generation. Setting this parameter toTrueturns off all special model behavior and ignores the additional exit heads, and thus also renders the parameters below irrelevant.output_layer_index: int: choose asinglespecific model exit for generation instead of the original (top layer) exit head.contrast_layer_indices: Tuple[Union[int, str], int]: use anauto-contrastive generationsetting, contrasting between the two specified exit layers. To perform contrast with the original LM head, pass the string 'original'. For example:(24, 18)will perform contrast between the exits at layers 24 and 18, and('original', 12)will perform contrast between the original LM head and the head at exit 12.contrast_function: Callable(Advanced): enables setting a custom contrast function, that gets logits of next-token predictions from two exit heads (i.e., those specified incontrast_layer_indices) and returns a modified set of predictions. By default, thecalculate_contrasted_logitsfunction is used, which performs the contrast calculation described inGera et al. (2023).Running language modeling benchmarksFor GPT-family models, it is possible to run benchmarks for a given multi-exit model with theLanguage Model Evaluation Harnesslibrary, for different generation settings.This is done viapython -m autocontrastive_gen.evaluation.lm_eval_harness.run_lm_eval, by specifying--model multi_exit_gptand adding the desiredmulti-exit configuration settingsto the--model_argsruntime argument.For example:python-mautocontrastive_gen.evaluation.lm_eval_harness.run_lm_eval\--taskslambada_openai\--modelmulti_exit_gpt\--model_argspretrained=IBM/gpt2-medium-multiexit,use_original_head=False,contrast_layer_indices='original;12'\--output_pathmy_output_pathFor details on the benchmark tasks available and on additional runtime arguments of the evaluation script, refer tohttps://github.com/EleutherAI/lm-evaluation-harness#basic-usage.Pre-trained model checkpointsWe release the following model checkpoints:ibm/gpt2-medium-multiexit- identical to theGPT-2 Mediumpre-trained model checkpoint, but with 12 newly-trained linear exit heads.ibm/gpt-neo-125m-multiexit- identical to theGPT Neo 125Mpre-trained model checkpoint, but with 6 newly-trained linear exit heads.The new heads were trained on the English portion of theCC-100dataset. For further details, refer to Appendix A of the paper.ReferenceAriel Gera, Roni Friedman, Ofir Arviv, Chulaka Gunasekara, Benjamin Sznajder, Noam Slonim and Eyal Shnarch (2023).The Benefits of Bad Advice: Autocontrastive Decoding across Model Layers. ACL 2023.Please cite:@inproceedings{gera2023autocontrastive,
  title={The Benefits of Bad Advice: Autocontrastive Decoding across Model Layers},
  author={Gera, Ariel and Friedman, Roni and Arviv, Ofir and Gunasekara, Chulaka and Sznajder, Benjamin and Slonim, Noam and Shnarch, Eyal},
  booktitle={Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  month={july},
  address={Toronto, Canada},
  year={2023}
}LicenseThis work is released under the Apache 2.0 license. The full text of the license can be found inLICENSE.