Deep Network Trainability via Persistent Subspace Orthogonality
Alex Massucco, Davide Murari, Carola-Bibiane Schönlieb
TL;DR
This work tackles vanishing/exploding gradients in deep networks by studying the network Jacobian and establishing both exact Jacobian orthogonality and a relaxed persistent subspace orthogonality (PSO) framework. The authors provide a constructive characterization: exact a.e. Jacobian orthogonality can be achieved with specific piecewise-affine layers (e.g., Modified ResNet and modified feedforward forms) and then generalize to PSO, which requires isometry on a nontrivial subspace shared across depth. They propose two practical initialization strategies to enforce PSO and validate them with extensive experiments on MNIST, Fashion-MNIST, and CIFAR-10, showing substantially improved trainability for very deep feedforward and residual networks, with robustness to subspace dimension and activation choices. The findings offer a principled approach to stabilize gradient propagation in deep architectures and suggest extensions to Transformers and graph neural networks as future directions.
Abstract
Training neural networks via backpropagation is often hindered by vanishing or exploding gradients. In this work, we design architectures that mitigate these issues by analyzing and controlling the network Jacobian. We first provide a unified characterization for a class of networks with orthogonal Jacobian including known architectures and yielding new trainable designs. We then introduce the relaxed notion of persistent subspace orthogonality. This applies to a broader class of networks whose Jacobians are isometries only on a non-trivial subspace. We propose practical mechanisms to enforce this condition and empirically show that it is necessary to sufficiently preserve the gradient norms during backpropagation, enabling the training of very deep networks. We support our theory with extensive experiments.