Correcting dropout and normalization improves protein models accuracy
Correcting the Dropout-LayerNorm Expectation Gap Improves Protein Structure Models
Machine Learning
Summary
The authors found that a common way of combining dropout and LayerNorm layers in protein structure prediction models causes a consistent error when these models are used for evaluation. This mismatch makes the models less accurate. They developed a simple mathematical fix that adjusts for this error without extra computation. This fix improves accuracy across many protein models, especially for antibody structures, making predictions more reliable.
What this means in practice
- •For protein modeling teams: Improve the accuracy of protein structure predictions by applying a correction to pretrained models using dropout and LayerNorm layers.
- •For drug discovery teams: Enhance protein-ligand docking predictions by using the dropout-LayerNorm correction to reduce bias in evaluation-time structure models.
Authors
Isaac Ellmen, David Errington, Matthew I. J. Raybould, Charlotte M. Deane
Abstract
Although $\mathbb{E}[\mathrm{dropout}(x)] = x$, here we show that $\mathbb{E}[\mathrm{LayerNorm}(\mathrm{dropout}(x))]$ is not equal to $\mathrm{LayerNorm}(x)$. Accordingly, the pattern of a Dropout layer followed by a LayerNorm, which is common to many AlphaFold2-based protein structure predictors, produces a systematic bias at evaluation time that can hamper performance. To address this, we derive a closed-form, first-order correction for this gap, which we call a Dropout-LayerNorm Correction (DLC). DLC empirically matches the performance boost of large Monte Carlo dropout ensembles. We evaluate its effect across nine protein structure models (ESMFold, OpenFold, ABB3, FlashABB, Ibex, NbForge, Genie1, Genie2, Genie3) on both paired and single-chain antibody structures as well as one protein-ligand docking model (QuickBind). The correction is computationally negligible and improves accuracy in all ten models tested ($\sim0.3\%-13\%$), with a modest but consistent improvement in ESMFold and OpenFold and a substantial improvement in antibody-specific models. This work identifies the mathematical consequence of chaining together Dropout and LayerNorm and provides a free, principled adjustment to improve the evaluation performance of many pretrained models. Code to reproduce the experiments is available at https://github.com/oxpig/DLC.