PaperlayerView exact revision on arXiv

Source-linked AI summary

Correcting the Dropout-LayerNorm Expectation Gap Improves Protein Structure Models

Isaac Ellmen, David Errington, Matthew I. J. Raybould, Charlotte M. Deane

arXiv:2609.32062v1q-bio.BMcs.LG

TL;DR

Dropout’s unbiased expectation does not guarantee an unbiased output when followed by LayerNorm, a pattern used in protein structure predictors. The paper derives a closed-form correction for this gap and reports accuracy improvements across ten tested models, with larger gains in some antibody-specific models.

  • Problem

    Dropout followed by LayerNorm can create an evaluation-time expectation bias because LayerNorm uses shared channel statistics, which can propagate into translation predictions in AlphaFold2-style models.

    Source: §1 · p. 1, §2.1 · p. 2

  • Method

    The paper derives a closed-form correction for dropout followed by LayerNorm and applies it at post-LayerNorm sites directly preceded by dropout.

    Source: §2.3 · p. 3, §3 · p. 4

  • Results

    The correction reliably improves performance across all 10 tested models, with full-Fv RMSD improvements of 10.9% for FlashABB and 13.3% for ABB3.

    Source: §4.3 · p. 5, §5 · p. 9

  • Takeaways & Limitations

    The computationally negligible correction can improve pretrained-model accuracy, with effects ranging from small but consistent in ESMFold to large enough to expand the antibody-prediction Pareto front in ABB3.

    Source: §5 · p. 9

    For Genie3, corrected and uncorrected scores undershoot true noise, raising the possibility of overly conservative structures; design impacts remain for future investigation.

    Source: §4.6 · p. 8

Abstract

from arXiv · show

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.

1 INTRODUCTION

Standard dropout scaling makes dropout itself unbiased, but that correction assumes the following layer is variance-invariant. In AlphaFold2-style blocks, LayerNorm can instead pass an expectation bias into translation updates.

1 INTRODUCTION

  • Dropout’s usual 1/(1 − p) training-time scaling restores E[dropout(x)] = x, but assumes the next layer is invariant to variance changes.Dropout followed by residual+LayerNorm or LayerNorm alone does not satisfy that assumption.Source: §1 · p. 1
  • AlphaFold2-style blocks place Dropout→LayerNorm→Linear→Translation at the end of each StructureModule block, propagating LayerNorm bias into predicted translation.The paper reports an approximately 5−6% translation overshoot and presents a closed-form correction that improves accuracy without retraining.Source: §1 · p. 1
  • The paper identifies a systematic bias when a dropout layer directly precedes LayerNorm.Source: §1 · p. 2

2. We provide a closed form fix which can be patched directly into LayerNorm

The paper tests its fix on ten biomolecular structure models in a therapeutically relevant immunoglobulin context and reports improved accuracy in almost all cases.

2. We provide a closed form fix which can be patched directly into LayerNorm

  • The fix improves accuracy in almost all ten biomolecular structure models tested in the therapeutically relevant immunoglobulin context.Source: §2 · p. 2

2 THEORY

Dropout is unbiased per channel, but LayerNorm’s shared channel statistics make its expected output after dropout differ from LayerNorm applied directly. The paper derives corrections for zero- and non-zero-mean inputs and discusses extensions to residual-stream normalization.

2 THEORY

  • For the toy input, dropout’s expectation recovers x, but LayerNorm after dropout shrinks its expected output to approximately √q · x.LayerNorm alone is a no-op on this particular zero-mean, unit-standard-deviation input.Source: Figure 1 · p. 2

2.1 SETUP

  • LayerNorm shares mean and standard deviation across channels, making that shared normalization the source of evaluation-time bias after dropout.In most cases, the mean is approximately zero, so changes in standard deviation directly rescale the normalized output.Source: §2.1 · p. 2

2.2 SIMPLIFIED DERIVATION (ZERO-MEAN CASE)

  • For approximately zero-mean inputs, the correction reduces to a global factor c = √q, independent of µ and σ.The correction preserves β at full strength while shrinking only the γ · norm(x) component.Source: §2.2 · pp. 2–3

2.3 GENERAL (µ̸ = 0) CORRECTION

  • For non-zero µ, the correction factor is approximated by q/(1 + p(µ/σ)2), approaching √q as µ approaches zero.The paper also derives a residual-stream correction, but reports that its practical effect was small in tested models; RMSNorm and BatchNorm corrections are not empirically tested.Source: §2.3 · p. 3

3 METHODS

The experiments apply the correction at dropout-to-LayerNorm sites in protein structure predictors. They focus on antibodies and nanobodies and include several model variants with different dropout or residual designs.

3 METHODS

  • The illustrated structure-module block sends the dropout→LayerNorm output through a Linear layer to rotation and translation updates, with bias affecting predicted translation length.Source: Figure 2 · p. 3
  • The tested predictor set includes general protein and antibody-specific models, nanobody/VHH prediction, and protein structure generation models.The cited examples include OpenFold, ESMFold, ABB3, FlashABB, Ibex, NbForge, and Genie1–3.Source: §3 · pp. 3–4
  • Some model variants avoid or mitigate the pattern through zero dropout during finetuning or residual connections in updates.ImmuneBuilder uses 0% dropout during finetuning, FrameFlow uses 0% throughout the StructureModule, and AlphaFold3-style updates contain residual connections.Source: §3 · p. 4
  • Experiments are restricted to antibodies and nanobodies, comparing full-structure and CDRH3-specific RMSDs, with correction hooks at directly preceded post-LayerNorm sites.Most models have two such sites per block; Ibex and NbForge have additional model-specific hooks.Source: §3 · p. 4

4 EXPERIMENTS

Across antibody, nanobody, generative-backbone, and docking experiments, the correction generally improves structural accuracy and brings deterministic predictions closer to dropout-on Monte Carlo means. Benefits vary by model and evaluation setting, and the full QuickBind benchmark shows no statistically significant improvement.

4.1 TRANSLATION LENGTH IS SYSTEMATICALLY WRONG, AND THE CORRECTION FIXES IT

  • The uncorrected ABB3 first-block translation overshot the Monte Carlo mean by about 5.6%, while correction nearly matched it and reduced path travel by about 18%.The corrected path-directness mean was 1.965 versus 2.313 uncorrected across eight blocks.Source: §4.1 · p. 4, Figure 3 · p. 4

4.2 THEORY MATCHES THE TRUE MONTE CARLO EXPECTATION

  • ABB3’s corrected predictions approached the Monte Carlo mean: coordinate averaging improved RMSD monotonically with sample count, while correction provided a cheap approximation without averaging structures.The corrected and Monte Carlo results had similar RMSD; the figure reports convergence to the theory-corrected line by about 30 samples.Source: §4.2 · p. 5, Figure 4 · p. 5

4.3 EXPANDING THE PARETO FRONT OF ANTIBODY STRUCTURE PREDICTION

  • 10.9% for FlashABB and 13.3% for ABB3: the correction improved whole-Fv RMSD and brought FlashABB within ∼0.02Å of Ibex while running ∼100x faster.These three models formed the speed–accuracy Pareto front; the correction substantially closed their accuracy gap.Source: §4.3 · p. 5, Figure 5 · p. 6

4.4 IMPROVED NANOBODY STRUCTURE PREDICTION

  • Whole-domain nanobody prediction improved in every tested model; CDRH3 improved in all except NbForge, whose apparent hit varied across random subsets.The 100-structure SAbDab2 test set overlaps substantially with some models’ training data, so absolute performance warrants caution.Source: §4.4 · p. 6

4.5 MORE RELIABLE CONFIDENCE PREDICTION

  • Correction improved pLDDT accuracy substantially for ABB3 and FlashABB and whole-domain Spearman correlation for NbForge, but produced little to no improvement in Ibex models.The authors note Ibex may have been finetuned with zero dropout or its pLDDT head trained using frozen evaluation weights.Source: §4.5 · pp. 6–7

4.6 THE GENIE FAMILY

  • The correction improved all three Genie models, especially Genie1 and Genie2, and closed most of Genie2’s performance gap with Genie3 on nanobodies.For Genie3, corrected score vectors matched the dropout-on Monte Carlo mean, but all models undershot true noise, potentially yielding overly conservative structures.Source: §4.6 · p. 8

4.7 QUICKBIND: PROTEIN-LIGAND DOCKING

  • QuickBind showed no statistically significant correction effect across the full benchmark, but among predictions already within 3 Å, improvement was significant (t = −3.08, p ≈0.002) with a 63/37 skew toward improvement.The authors caution that the < 2 Å filter may exclude predictions that could benefit from correction.Source: §4.7 · p. 8, Table 4 · p. 9

5 DISCUSSION

DLC is a negligible-cost correction that improves all ten tested models, with effects ranging from small but consistent to large. The discussion links model differences to architecture and considers how correcting the bias may inform future model designs.

5 DISCUSSION

  • DLC reliably improves performance across all 10 tested models, with a small but consistent effect in ESMFold and a larger effect in ABB3.In ABB3, the effect meaningfully expands the Pareto front for antibody structure prediction.Source: §5 · p. 9
  • Shared structure-module weights in ESMFold and AlphaFold2/OpenFold may let later layers compensate for early overshoot, unlike ABB3’s separate per-layer weights.The authors hypothesize that bias in ABB3’s large early jumps may be unrecoverable by later layers.Source: §5 · p. 9
  • Higher dropout rates can degrade generalization with dropout followed by LayerNorm, while small-model experiments suggest residual connections solve the gap at 10% dropout.The authors caution that larger models or dropout rates above 10% may remain sensitive to accumulated error.Source: §5 · p. 9
  • DLC may make dropout-plus-LayerNorm models easier to explore, including in vision, language, and protein modeling, though those future benefits are not established here.The discussion notes that CogView found improved numerical stability when normalization was separated from the residual stream.Source: §5 · p. 9

A FULL (µ̸ = 0) DERIVATION

The appendix extends the LayerNorm correction beyond the zero-mean case and treats sites with residual bypasses. It also derives theoretical corrections for RMSNorm and BatchNorm, which were not tested empirically in this work.

A FULL (µ̸ = 0) DERIVATION

  • The general derivation gives a correction involving q, p, and the ratio µ/σ, with approximations dropping terms of order O(1/d).The omitted terms vanish as channel dimension d grows and are described as negligible at the widths used in the paper.Source: A FULL (µ̸ = 0) DERIVATION · p. 14

A.1 EXACT PER-CHANNEL MOMENTS OF z

  • Independent per-channel dropout masks make z_k and z_j independent for k ≠ j.Source: §A.1 · p. 14

A.2 MEAN AND VARIANCE OF THE LAYERNORM STATISTICS

  • The channelwise mean of z equals µ exactly, while the derivation approximates the LayerNorm variance statistics by dropping O(1/d) terms.The mean identity does not require the zero-mean assumption.Source: §A.2 · p. 14, §A.3 · p. 14

A.3 THE GENERAL CORRECTION

  • For general µ, the β-preserving LayerNorm correction scales the normalized output by c(µ, σ) and tends to √q when µ = 0.The derivation treats the shared variance as concentrated near its mean.Source: §A.3 · p. 14

A.4 DAMPED VARIANT FOR RESIDUAL-BYPASS SITES

  • At residual-bypass sites, the correction is damped according to the dropped-out contribution relative to the combined signal; removing the bypass recovers the general correction.When the bypass dominates, the damping factor approaches 1 and no correction is needed.Source: §A.4 · pp. 14–15

B EXTENSION TO OTHER NORMALIZATION SCHEMES

  • The appendix derives analogous corrections for RMSNorm and BatchNorm, but neither normalization scheme is used by the tested models.These extensions are theoretical and not empirically validated in this paper.Source: B EXTENSION TO OTHER NORMALIZATION… · p. 15

B.1 RMSNORM

The appendix derives a constant RMSNorm correction and a BatchNorm correction based on population statistics. It also analyzes a second LayerNorm in antibody-model pLDDT heads and validates a further correction against Monte Carlo estimates.

B.1 RMSNORM

  • For RMSNorm, the correction is √q for every µ, unlike LayerNorm’s µ-dependent general correction.The appendix attributes the difference to RMSNorm’s use of the raw second moment without mean subtraction.Source: §B.1 · p. 15

B.2 BATCHNORM

  • BatchNorm’s evaluation map uses fixed running statistics, so dropout-corrupted training statistics may mismatch dropout-free test inputs; its derived correction has LayerNorm’s form using population channel statistics.The passage states that the BatchNorm extension is not empirically validated in the tested models.Source: §B.2 · p. 16

C FURTHER PLDDT CALIBRATION

  • In ABB3-style models, a second LayerNorm removes the first correction’s rescaling, leaving the pLDDT head with a biased representation.Source: C FURTHER PLDDT CALIBRATION · p. 16

C.1 AN EFFECTIVE CORRECTION CAN BE DERIVED AND VALIDATED

  • The derived second-site correction reduced the mean gap to a 100-sample Monte Carlo estimate from 0.507 to 0.255 across 2,273 residues.A multiplier sweep found that the as-derived formula minimized the gap.Source: §C.1 · p. 17

C.2 THE CORRECT ESTIMATE IS POORLY CALIBRATED

The theoretically derived correction worsens pLDDT calibration, while directly tuning a scalar against calibration yields an approximately 7% MAE improvement. A single scalar cannot simultaneously repair absolute calibration and residue discrimination.

C.2 THE CORRECT ESTIMATE IS POORLY CALIBRATED

  • The theoretically motivated corrections degrade pLDDT calibration on both the 34-structure test subset and an independent 100-structure training sample.The independent sample rules out an out-of-distribution explanation, according to the passage.Source: §C.2 · p. 17
  • Derived ρ2 correction gives pLDDT MAE of 3.396 on test and 3.751 on train, versus uncorrected values of 3.315 and 3.679.The passage reports test and train MAE changes of −2.4% and −2.0%, respectively.Source: §C.2 · p. 17
  • A calibration-tuned c2 ≈1.75 improves MAE by ∼7% on both test (3.315 →3.079) and train (3.679 →3.422).This optimum amplifies away from bias2, opposite to the theoretically motivated corrections.Source: §C.2 · p. 17
  • The derived correction strengthens residue-accuracy ranking, but scalar rescaling trades discrimination against absolute calibration and cannot fix both simultaneously.Spearman correlations are −0.431 for derived ρ2, −0.425 at c2=1, and −0.265 uncorrected; at c2=1.75, correlation drops to −0.384.Source: §C.2 · p. 17

D.1 FULL ANTIBODY RESULTS

Table 5 compares paired antibody models on a 34-structure test subset, reporting framework-aligned CDR-H3 RMSD and whole-structure RMSD without MD refinement.

D.1 FULL ANTIBODY RESULTS

  • Table 5 covers paired antibody models on the 34-structure test subset, with predictions evaluated without MD refinement.The caption specifies both CDR-H3 and whole-structure RMSDs.Source: Table 5 · p. 18
  • CDR-H3 RMSD is computed after alignment to heavy framework residues.Source: Table 5 · p. 18
  • Ibex results include both member 0 alone and the full eight-member deep ensemble; ESMFold and OpenFold use the AlphaFold gap trick without recycling.The caption notes uncertainty about ESMFold and Ibex training details and that ABB2 cannot be corrected because it finetunes without dropout.Source: Table 5 · p. 18

D.2 IBEX ON A RANDOM 100-STRUCTURE TRAINING-SET SUBSAMPLE

The appendix reports supplementary Ibex and recycling comparisons, tests significance of ESMFold and OpenFold shifts, sweeps the assumed keep probability, and decomposes Genie3 noise predictions.

D.2 IBEX ON A RANDOM 100-STRUCTURE TRAINING-SET SUBSAMPLE

  • On an independent random sample of 100 Ibex training-set structures, the comparison uses framework-aligned CDR-H3 and whole-Fv RMSD for member 0 and the eight-member medoid ensemble.The sample was restricted to structures overlapping the authors’ training split; Ibex finetuning details are unavailable, so their possible effect on correction is not ruled out.Source: §D.2 · p. 18, Table 6 · p. 18

D.3 ESMFOLD AND OPENFOLD WITH THEIR OWN DEFAULT RECYCLING

  • ESMFold and OpenFold are also compared on nanobody-100 with their default three recycling iterations, rather than the no-recycling convention used in Table 1.The table caption notes that ESMFold may have been finetuned without dropout.Source: Table 7 · p. 19, §D.3 · p. 19

D.4 STATISTICAL SIGNIFICANCE OF THE ESMFOLD/OPENFOLD IMPROVEMENT

  • Correction shifts are significant for ESMFold (t = −8.64, p < 0.001) and OpenFold (t = −7.48, p < 0.001), with 79/100 and 85/100 structures improved, respectively.The mean shifts are small, but most structures move in the same direction.Source: Figure 6 · p. 19, §D.4 · p. 19

D.5 CALIBRATION: SWEEPING THE ASSUMED q

  • For ABB3 and FlashABB, q = 0.9 improves RMSD over the uncorrected q = 1, although the optimal assumed q does not always equal 0.9.The authors note possible outlier noise and downstream sensitivity to variance changes; Figure 7 plots RMSD across q for CDR-H3 and whole structure on train, validation, and test splits.Source: §D.5 · p. 20, Figure 7 · p. 20

D.6 GENIE3 NOISE-PREDICTION DECOMPOSITION

  • In Genie3, uncorrected predicted noise exceeds the N=100 Monte Carlo ceiling at 97/100 timesteps, while all correction modes undershoot the true-noise reference.Predicted noise direction is nearly unchanged across modes, indicating a magnitude-only bias; local Monte Carlo averages converge toward the theory-corrected value as draws increase.Source: Figure 8 · p. 21

Related papers

Loading 2609.32062v1…