Correcting the Dropout-LayerNorm Expectation Gap Improves Protein Structure Models

2026-09-25

AI Summary

This study addresses the systematic inference bias introduced by cascading Dropout and LayerNorm, which constrains protein structure prediction accuracy. By revealing the mathematical deficiency of this composite layer, the authors derive a closed-form first-order correction formula for the expectation gap within the AlphaFold2 architecture, proposing a theoretically grounded calibration scheme termed DLC that incurs zero computational overhead. This method eliminates inference bias without additional computation while achieving performance comparable to large-scale Monte Carlo ensembling. Evaluated across ten models, the proposed approach improves prediction accuracy by 0.3%–13%, significantly enhancing antibody structure prediction performance.

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.

Problem

• Dropout
• LayerNorm
• protein structure prediction
• systematic bias
• expectation gap

Innovation

• Dropout-LayerNorm Correction
• Protein Structure Prediction
• Systematic Bias
• Closed-form Correction
• Monte Carlo Dropout