This paper, just accepted at ICLR's GRaM workshop, asks a simple question:
Does gradient descent systematically take the wrong step in activation space?
It is shown:
Parameters take the step of steepest descent; activations do not
The paper mathematically demonstrates this for simple affine layers, convolution, and attention.
The work then explores solutions to address this.
The solutions may consequently provide an alternative mechanistic explanation for why normalisation helps at all, as two structurally distinct fixes arise: existing (L2/RMS) normalisers and a new form of fully connected layer (MLP).
Derived is:
- A new form of affine-like layer (a.k.a. new form for fully connected/linear layer). featuring inbuilt normalisation whilst preserving DOF (unlike typical normalisers). Hence, a new alternative layer architecture for MLPs.
- A new family of normalisers: "PatchNorm" for convolution, opening new directions for empirical search.
Empirical results include:
- This affine-like solution is not scale-invariant and is not a normaliser, yet it consistently matches or exceeds BatchNorm/LayerNorm in controlled MLP ablation experiments—suggesting that scale invariance is not the primary mechanism at work—but maybe this it is the misalignment.
- The framework makes a clean, falsifiable prediction: increasing batch size should hurt performance for divergence-correcting layers. This counterintuitive effect is observed empirically and does not hold for BatchNorm or standard affine layers. Corroborating the theory.
Hope this is interesting and worth a read.
- I've added some (hopefully) interesting intuitions scattered throughout, e.g. the consequences of reweighting LayerNorm's mean & why RMSNorm may need the sqrt-n factor & unifying normalisers and activation functions. Hopefully, all surprising fresh insights - please let me know what you think.
Happy to answer any questions :-)
[link] [comments]
