Protein Language Model Zero-Shot Fitness Predictions are Improved by Inference-only Dropout

TL;DR

This paper demonstrates that applying dropout during inference in protein language models enhances zero-shot protein fitness prediction accuracy without retraining, using a simple averaging technique.

Contribution

Introducing inference-only dropout in PLMs improves zero-shot protein property predictions without retraining or fine-tuning.

Findings

Inference-only dropout increases prediction accuracy.

A dropout rate of 0.1 is generally effective.

Method works across multiple PLMs without retraining.

Abstract

Protein Language Models (PLMs) such as ESM2 have been shown to be capable of zero-shot prediction of critical scalar properties of proteins (fitness). In this work, we show that injecting a dropout layer at inference time between a PLM's featurizer/embedding layer and its transformer, and averaging its output akin to Monte-Carlo dropout increases zero-shot performance on a subset of the ProteinGym dataset. This is the case even when the model was not trained with dropouts to begin with, and does not require retraining or finetuning of the PLM. A dropout of 0.1 seems performant across all models.