Transformative or Conservative? Conservation laws for ResNets and Transformers

TL;DR

This paper derives and analyzes conservation laws for modern neural network architectures like ResNets and Transformers, revealing their properties and how they persist under practical training dynamics such as SGD.

Contribution

It extends the understanding of conservation laws from shallow networks to complex architectures like ResNets and Transformers, including their behavior under discrete optimization.

Findings

Conservation laws for basic building blocks are explicitly characterized.

Residual connections do not alter the conservation laws of the underlying blocks.

Conservation principles persist under stochastic gradient descent in practice.

Abstract

While conservation laws in gradient flow training dynamics are well understood for (mostly shallow) ReLU and linear networks, their study remains largely unexplored for more practical architectures. This paper bridges this gap by deriving and analyzing conservation laws for modern architectures, with a focus on convolutional ResNets and Transformer networks. For this, we first show that basic building blocks such as ReLU (or linear) shallow networks, with or without convolution, have easily expressed conservation laws, and no more than the known ones. In the case of a single attention layer, we also completely describe all conservation laws, and we show that residual blocks have the same conservation laws as the same block without a skip connection. We then introduce the notion of conservation laws that depend only on a subset of parameters (corresponding e.g. to a pair of consecutive layers, to a residual block, or to an attention layer). We demonstrate that the characterization of such laws can be reduced to the analysis of the corresponding building block in isolation. Finally, we examine how these newly discovered conservation principles, initially established in the continuous gradient flow regime, persist under discrete optimization dynamics, particularly in the context of Stochastic Gradient Descent (SGD).