Rethinking Activation Functions: Reducing Bias Can Improve Learning

TL;DR

Peng, Zheng & Liu (Tsinghua / MetaCircle) argue that the nonzero post-activation mean built into ReLU/GELU/Swish/Sigmoid — what they call activation bias — is a hidden stability knob that trades early learning speed for worse late-stage loss. In a controlled residual-MLP teacher–student setup, subtracting the activation’s expected post-activation output in every layer slows early learning but overtakes the unshifted baseline, produces smoother hidden-state trajectories, drives Hessian eigenvalues ~10× smaller, and is associated with distinct parameter-update directions. An adaptive-shifting variant (per-step, per-layer mean subtraction) scales the intervention to ViT-B / DiT-B / GPT-2 Small and yields small but consistent late-stage gains across ImageNet classification, ImageNet-256 generation, and OpenWebText language modeling.

Key claims

The paper names activation bias — the expected value $\mathbb{E}[\sigma(x)]$ of commonly used activations — as a first-class property worth controlling, and shows that reducing it slows early loss reduction but yields a lower final loss than the matched unshifted baseline in a controlled teacher–student residual-MLP [§2.2, Fig. 2].
Reduced activation bias changes where in parameter space the student moves: on the plane spanned by the two students’ initial gradients, the trajectories are nearly orthogonal; the shifted student’s update develops primarily outside this plane, and its on-plane loss varies little [§2.3, Fig. 3].
Reduced activation bias is associated with a markedly flatter local loss landscape: the top Hessian eigenvalues stay substantially smaller throughout training, and in a simplified down-projection-only case $\mathbb{E}[z] = 0$ provably minimizes the largest Hessian eigenvalue [§2.4, Fig. 4, Eq. 7].
In a VGG-style autoencoder, decoding the activation-bias component $(\mathbb{E}[z])$ of the hidden state alone from the unshifted model recovers recognizable image structure early in training and then diverges from the full reconstruction — a “transient use” signature that is largely absent in the shifted model [§3.2, Fig. 5].
Adaptive shifting (subtract per-step per-layer post-activation mean per neuron) scales the intervention to ViT-B on ImageNet-1K, DiT-B/2 on ImageNet-256, and GPT-2 Small on OpenWebText: lower final test loss and higher top-1 on ViT-B despite higher training loss; comparable-or-lower training loss and consistently lower FID on DiT-B/2; ~0.01 absolute loss and ~1% relative perplexity reduction on GPT-2 Small [§4, Fig. 7, Eq. 9, Appendix E].
The hidden-state overshoot signature (a non-monotonic residual-stream magnitude trajectory during the window of fastest loss reduction) that the controlled model displays also appears in ViT-B pretraining and is substantially suppressed under adaptive shifting [§4.1, Fig. 7(a)].

Method

The intervention is a per-layer subtraction applied to the activation function. In the controlled setup, it is a fixed scalar: $\sigma_s(x) = \sigma(x) - s$ , with $s$ chosen as $\mathbb{E}[\sigma(x)]$ at initialization (e.g. $s=0.2$ for ReLU, $0.115$ for GELU, $0.077$ for Swish) so the expected post-activation stays near zero over a long training window. The architecture is a pre-normalized MLP residual stack — attention and output head stripped to isolate the activation’s role — trained in a teacher–student setup where both students share initialization, data, and optimizer and differ only in whether the shift is applied. On the initial-gradient plane spanned by $\{\hat{g}_u, (g_s - (g_s \cdot \hat{g}_u)\hat{g}_u)/\|\cdot\|\}$ , each student’s checkpoint trajectory is overlaid on a heatmap of the projected loss (parameters reconstructed by $\theta(a,b)=\theta_0 + a e_1 + b e_2$ ). To visualize the shifted student’s out-of-plane movement, the authors Gram–Schmidt the component of its parameter displacement orthogonal to the initial-gradient plane and re-visualize.

For scale-up (ViT-B, DiT-B/2, GPT-2 Small), fixed shifts are insufficient because the per-layer post-activation mean evolves rapidly during training. The authors replace them with adaptive shifting: at every forward pass, subtract the current across-neuron post-activation mean from each layer’s output, so the shifted layer has exactly zero mean across neurons by construction [Eq. 9]. All other training knobs are held fixed; only the shift is toggled.

Results

Controlled model (residual MLP, teacher–student). Unshifted student learns faster for roughly the first $\sim 10^3$ steps, then the shifted student overtakes and converges to lower loss [Fig. 2 top]. Hidden-state magnitude of the unshifted student overshoots around step $10^3$ — the same window where it leads on loss — while the shifted student’s hidden state evolves monotonically [Fig. 2 bottom]. Top Hessian eigenvalues of the shifted student are substantially smaller throughout training [Fig. 4].
ViT-B / ImageNet-1K. Adaptive shifting trails the unshifted baseline by a substantial margin early on but ends with both lower test loss and higher top-1 accuracy, despite ending with higher training loss [Fig. 6]. The residual-stream overshoot signature is markedly smaller under adaptive shifting [Fig. 7(a)].
DiT-B/2 / ImageNet-256. Adaptive shifting gives comparable or lower training loss throughout, and lower FID at every evaluated checkpoint [Fig. 7(b), Appendix E].
GPT-2 Small / OpenWebText. Adaptive shifting lowers loss by ≈0.01 in absolute terms and perplexity by ≈1% relative to the unshifted baseline, consistently across optimization steps [Fig. 7(c)]. Notably, the unshifted model does not show the early speed advantage in this setting — the authors attribute this to GPT-2’s small init scale.

Why it’s interesting

Sits in exactly the residual-stream-surgery family the wiki has been tracking under Training stability at scale — A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training‘s outlier-driven implicit rescaling, mHC: Manifold-Constrained Hyper-Connections‘s residual-mixing manifold constraint, Attention Residuals (AttnRes / Block AttnRes) — Technical Report‘s AttnRes — but relocates the stability lever one step upstream: in the activation function itself, with zero added parameters and a one-line code change. The mechanism it names (bias → sharper Hessian → overshoot → worse late-stage loss) is a natural companion to How to Set the Learning Rate for Large-Scale Pre-training?‘s “QK-Norm did the job µP was designed for” story — both argue a seemingly cosmetic architectural property is doing the load-bearing stability work the field had attributed elsewhere. The LeCun quote-reply (reproducing a NeurIPS 1990 paper) reinforces that this is a classical pathology being rediscovered at modern scale rather than a new one, but the scale-up evidence on ViT-B / DiT-B / GPT-2 is new.