The Affine Divergence: Aligning Activation Updates Beyond Normalisation

TL;DR

This paper introduces the Affine Divergence, a new perspective on activation updates and normalization, proposing alternative normalization methods that outperform traditional ones and offer a fresh mechanistic understanding.

Contribution

It presents a novel theoretical framework for activation updates and normalization, including the PatchNorm method, challenging conventional affine normalization approaches.

Findings

PatchNorm outperforms traditional normalizers in tests.

Theoretical analysis links normalization to activation function maps.

Alternative normalization methods are empirically validated.

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, incidentally, derive normalisation from first principles despite motivational independence. Consequently, such considerations offer a fresh, conceptual reframe of normalisation's action, with auxiliary experiments bolstering this mechanistic interpretation. Moreover, this analysis makes clear a second possibility: a solution that is functionally distinct from modern normalisations, without scale invariance, yet remains empirically successful -- an alternative to the affine map. This outperforms conventional normalisers across several tests. This generalises to convolution via a new functional form, ``PatchNorm'', a compositionally inseparable normaliser. Together, these provide an alternative mechanistic framework that both 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. Overall, this constitutes a theoretically principled approach that yields new functions with empirical validation and raises questions about the affine + nonlinear approach.