Parallel Layer Normalization for Universal Approximation

TL;DR

The paper extends the universal approximation framework to neural networks that include normalization by proving that a PLN-Net, formed from parallel layer normalizations and linear layers, is a universal approximator and can approximate any $L$-Lipschitz function with width bounds that scale as $d(\lfloor L/(2\varepsilon)\rfloor+1)$. It introduces Layer Scaling as a simplification of LN and shows LS-d preserves universal approximation with the same width form, arguing centering is not essential for approximation. The authors provide theoretical comparisons against traditional activations (Sigmoid, Tanh, ReLU) and demonstrate PLN’s favorable capacity and optimization properties through shallow and deep experiments, including image classification with VGG/ResNet sans BN and Transformer translation tasks, where PLN-8 often matches or surpasses LN-based baselines. The work suggests PLN can replace or complement normalization/activation roles in modern architectures, offering improved trainability and performance, especially in regimes where normalization aids optimization. These results open avenues for designing architectures that jointly exploit normalization and nonlinear activation for universal function approximation and practical gains.

Abstract

Universal approximation theorem (UAT) is a fundamental theory for deep neural networks (DNNs), demonstrating their powerful representation capacity to represent and approximate any function. The analyses and proofs of UAT are based on traditional network with only linear and nonlinear activation functions, but omitting normalization layers, which are commonly employed to enhance the training of modern networks. This paper conducts research on UAT of DNNs with normalization layers for the first time. We theoretically prove that an infinitely wide network -- composed solely of parallel layer normalization (PLN) and linear layers -- has universal approximation capacity. Additionally, we investigate the minimum number of neurons required to approximate $L$-Lipchitz continuous functions, with a single hidden-layer network. We compare the approximation capacity of PLN with traditional activation functions in theory. Different from the traditional activation functions, we identify that PLN can act as both activation function and normalization in deep neural networks at the same time. We also find that PLN can improve the performance when replacing LN in transformer architectures, which reveals the potential of PLN used in neural architectures.

Figures (11)

• Figure 1: (a) and (b) show an intuitive difference between LN-Net and PLN-Net. (c) demonstrates that, PLN is used on hidden neurons and divides neurons of a layer into groups and performs LN in each group in parallel.
• Figure 2: (a) shows that $G({\mathbf{x}})\in\mathcal{G}(N;\varphi)$ is essentially a network. Here $n=4,d=2$ and $N=3$. Tradition activation functions act on each neuron, an example with $N=6$, $d=1$ and $\varphi=\textit{Sigmoid}$ is shown in (b). PLN requires a group of neurons to activate, an example with $N=3$, $d=2$ and $\varphi=\textit{LN}$ is shown in (c).
• Figure 3: The intuitive method to approximate $f(x)$ by $G(x)$. The function in (a) is that $G(x)=1,0{LN}(1,-1^\top x)$, demonstrating a special PLN-Net. As we can see, (b) demonstrates a classic and intuitive method to calculate the definite integral by dividing the area into small pieces. By observe the boundaries, we intuitively believe a continuous function can be approximated by the sum of multiple step functions, as shown in (c). This is the initial support for why we guess that the PLN-Net has a universal approximation property.
• Figure 4: The results of logarithmic loss of PLN-Net with different widths, norm sizes and dimensions of the inputs.
• Figure 5: The results of logarithmic loss of PLN-Net with different activations and dimensions of the inputs.

Theorems & Definitions (26)

• Definition 1: Parallel Layer Normalization
• Definition 2: Parallel $\varphi$-Net
• Definition 3: Minimum Approximation Width of $\mathcal{G}(N;\varphi)$
• Theorem 1: Minimum Approximation Width of PLN-Net
• Theorem 2: PLN for Universal Approximation
• Theorem 3: PLS for Universal Approximation
• Proposition 1: Minimum Approximation Width of $\mathcal{G}(N;\textit{Sigmoid})$
• Proposition 2: Minimum Approximation Width of $\mathcal{G}(N;\textit{ReLU})$
• Proposition 3: Minimum Approximation Width of $\mathcal{G}(N;\textit{LN-2})$
• Proposition 4: Minimum Approximation Width of $\mathcal{G}(N;\textit{LS-1})$