The Affine Divergence: Aligning Activation Updates Beyond Normalisation

TL;DR

The paper identifies an affine divergence between the ideal activation gradient and the actual gradient propagated through parameters during gradient descent, showing that the effective activation update scales as $||x||^2+1$. It derives two correction families—norm-like normalisers and a novel affine-like correction—and introduces PatchNorm for convolution to restore alignment to first order, with additional gradient-only learning-rate adjustments. Empirical results on CIFAR-10 demonstrate that the affine-like correction often outperforms traditional normalisers across several architectures, while batch-size analyses reveal a counterintuitive negative correlation for the structural corrections, supporting the misalignment mechanism. Overall, the work reframes normalization as a principled outcome of aligning ideal and effective updates, and it points to broader implications for activation-function design and architecture choices, including potential extensions to residuals and attention.

Abstract

A systematic mismatch exists between mathematically ideal and effective activation updates during gradient descent. As intended, parameters update in their direction of steepest descent. However, activations are argued to constitute a more directly impactful quantity to prioritise in optimisation, as they are closer to the loss in the computational graph and carry sample-dependent information through the network. Yet their propagated updates do not take the optimal steepest-descent step. These quantities exhibit non-ideal sample-wise scaling across affine, convolutional, and attention layers. Solutions to correct for this are trivial and, entirely incidentally, derive normalisation from first principles despite motivational independence. Consequently, such considerations offer a fresh and conceptual reframe of normalisation's action, with auxiliary experiments bolstering this mechanistically. Moreover, this analysis makes clear a second possibility: a solution that is functionally distinct from modern normalisations, without scale-invariance, yet remains empirically successful, outperforming conventional normalisers across several tests. This is presented as an alternative to the affine map. This generalises to convolution via a new functional form, "PatchNorm", a compositionally inseparable normaliser. Together, these provide an alternative mechanistic framework that adds to, and counters some of, the discussion of normalisation. Further, it is argued that normalisers are better decomposed into activation-function-like maps with parameterised scaling, thereby aiding the prioritisation of representations during optimisation. Overall, this constitutes a theoretical-principled approach that yields several new functions that are empirically validated and raises questions about the affine + nonlinear approach to model creation.

Figures (9)

• Figure 1: This plot displays the effect of various normalisers. This is depicted by drawing a thousand (green) input points drawn from a standard multivariate normal distribution in $\mathbb{R}^2$. Then, each plot depicts these initial input points in green and the resultant output distribution for each normaliser in a contrasting colour. The left demonstrates BatchNorm's output in red, where statistics are computed over all $1000$ input points and therefore produce little change in the output distribution, as the distribution is already standard-normal. Next is layernorm, which can be seen to collapse the entire distribution into two single blue clusters. This is because the LayerNorm map follows $\mathbb{R}^n\rightarrow S^{n-2}\hookrightarrow \mathbb{R}^2$. For $n=2$, this results in $S^0$, which preserves only sign information. Typically, this would be a hypersphere orthogonal to $\vec{1}$ in higher dimensions. Interestingly, the standard mean can be reweighted to any weighted mean, including one-hot, represented by $\vec{x}-(\vec{x}\cdot\hat{n})\hat{n}$, with geometrically minimal effect except for reorientating the hypersphere to be orthogonal to $\hat{n}$ (unless algebraic anisotropic Bird2025a function's direction distinctions interact). This geometrically undermines the layer mean-statistic argument, which sometimes supports LayerNorm, since the one-hot weighting statistic is unaffected by layer width. This supports the transition from statistical interpretation to a geometrical one, which better aligns with the activation function decomposition suggested in App.\ref{['App:ParameterisedNormalisation']}. Centre is $L_2$-norm, which projects the distribution into an $S^{n-1}$ purple hypersphere, similarly to the next right, which is RMSNorm, which projects to a $\sqrt{n}$ scaled hypersphere. Finally, a depiction of Affine-Like, for $\mathbf{W}=\mathrm{I}_{n\times n}$ and $\vec{b}=0$, shows how the distribution fills the volume of the hypersphere without projecting out the radial representation degree of freedom.
• Figure 2: Displays the performance of Standard Tanh fully-connected sequential networks of various architectures. Immediately, one can observe that the affine-like correction (red) outperforms all other normalisers, except for $3$-layer $16$-width, where norm-like marginally outperforms it. Moreover, the affine-like and norm-like (half and full $\eta$) performances are consistently substantially more performative than all other normalisers, having a wide margin between these normalisers and all others, especially for deeper networks. LayerNorm typically follows this, then BatchNorm, then no normaliser and then RMSNorm (showing the greatest and significant variability). It is notable that LayerNorm learns and stagnates the quickest in performance, whereas other normalisers take more time to build accuracy, with affine-like overtaking LayerNorm consistently around epoch $20$. The surprising underperformance of RMSNorm may be due to both the architecture and the parameterless formulation.
• Figure 3: Displays the performance of Leaky-ReLU fully-connected sequential networks of various architectures. This plot is more nuanced: early learning is considerably faster across all normalisers, with a subsequent clear drop in performance for BatchNorm and LayerNorm only, before stagnation. This is sometimes observed in other normalisers as well, on longer timescales. As before, the affine-like correction typically performs best, particularly at larger widths and depths, where it exhibits a more distinct separation from all other normalisers. At very narrow networks, $n=16$, no normaliser performs particularly well, followed by the affine-like correction. This is perhaps because the other normalisers remove representational degrees of freedom, which constitute a larger percentage of the representation volume in narrower-width networks. The clearest separation in performance is observed in the single hidden-layer $n=128$ plot, whereas the other plots are more overlapping. However, the standard error indicates that both affine-like and norm-like corrections consistently outperform alternatives.
• Figure 4: Plotted is performance against logarithmic layer-width for standard-tanh (left) and Leaky-ReLU (right), for layer widths ranging from $1$ through $256$. Training is otherwise identical to that in the other plots. Standard-tanh results show better performance separation, with the affine-like and norm-like structural solutions demonstrating significantly higher performance that grows with width, then begins to stagnate around a width of $128$. The affine-like solution also performs notably better than alternatives even at a width of $1$ neuron. Both no normaliser and RMSNorm reach a peak at $n=52$ and then decline again in performance. The ordering of normalisers follows from before: affine-like, then the two norm-like, then LayerNorm, BatchNorm, no normaliser and then RMSNorm. For Leaky-ReLU, the differences are slight, but Affine-like shows better performance at larger widths, while no normaliser outperforms at shallower widths.
• Figure 5: One can see across the results that both BatchNorm and no normaliser perform poorly, and similarly, this is because they are manifestly linear, so can achieve at most a linear fit. This is not the case for all other normalisers, which are nonlinear activation functions mathematically. There is a tendency for LayerNorm to be in the middle in performance, likely due to its additional losses in representational degrees of freedom. Affine-like follows this with some variability, then Norm-like, and RMS-like, which often perform similarly well. This may be because the latter two are more classical activation-function-like operations, which are composable with the affine layer, rather than being inseparable, as with the affine-like. Affine-like does enable folding of the activation distribution as can be seen from Eqn.\ref{['Eqn:AffineCorrectionSoftBound']}.