[Paper Review] Principled Weight Initialization for Hypernetworks

Oscar Chang, Lampros Flokas+1 more|arXiv (Cornell University)|Dec 13, 2023

Adversarial Robustness in Machine Learning34 references35 citations

TL;DR

The paper introduces hyperfan initialization for hypernetworks, providing principled variance-based schemes (hyperfan-in and hyperfan-out) to preserve mainnet activation/gradient scales, improving stability and convergence over ad-hoc initializations.

ABSTRACT

Hypernetworks are meta neural networks that generate weights for a main neural network in an end-to-end differentiable manner. Despite extensive applications ranging from multi-task learning to Bayesian deep learning, the problem of optimizing hypernetworks has not been studied to date. We observe that classical weight initialization methods like Glorot & Bengio (2010) and He et al. (2015), when applied directly on a hypernet, fail to produce weights for the mainnet in the correct scale. We develop principled techniques for weight initialization in hypernets, and show that they lead to more stable mainnet weights, lower training loss, and faster convergence.

Motivation & Objective

• Identify why classical weight initializations fail for hypernetworks generating mainnet weights.
• Develop principled variance-based initialization rules tailored to hypernetworks.
• Theoretically and empirically show that hyperfan initializations stabilize activations/gradients and improve convergence.

Proposed method

• Use variance analysis to derive hyperfan-in and hyperfan-out initializations for hypernetworks.
• Model the hypernet as generator of mainnet weights and analyze forward/backward variance propagation.
• Propose specific variance formulas for H, h, G, g layers to preserve mainnet activation/gradient variance.
• Differentiate cases where the hypernet outputs only weights or both weights and biases.
• Provide initialization schemes that align hypernet weight variance with classical fan-in/fan-out init semantics.

Figure 1: Mainnet Activations before the Start of Training on MNIST.

Experimental results

Research questions

• RQ1 Why do classical initialization schemes fail when applied to hypernetworks generating mainnet weights?
• RQ2 How can we initialize hypernetworks so that mainnet activations and gradients retain stable variance across depth?
• RQ3 Do hyperfan-in and hyperfan-out initializations enable stable training across different mainnet architectures and tasks?

Key findings

• Classical initializations on hypernets lead to exploding activations in the mainnet.
• Hyperfan-in and hyperfan-out initializations preserve mainnet variance and enable stable training.
• Hyperfan methods produce lower initial losses and faster convergence in MNIST feedforward experiments.
• Hyperfan initializations enable training when Xavier/Kaiming on hypernet fails (e.g., CIFAR-10, ImageNet Bayesian setup).
• Both hyperfan-in and hyperfan-out work with SGD across tasks, with no strong preference in practice.

Figure 2: Evolution of Hypernet Output Layer Activations during Training on MNIST. Xavier init results in unstable mainnet weights throughout training, while hyperfan-in and hyperfan-out init result in mainnet weights that stabilize quickly.