PaperlayerView exact revision on arXiv

Source-linked AI summary

GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong, W. Ronny Huang, Tom Goldstein

arXiv:2102.08098v3cs.LGcs.CLcs.CV

TL;DR

Complex and novel architectures can be unstable under existing initialization schemes that do not transfer readily across architectures. GradInit learns parameter-block scales to tailor initialization to a stochastic optimizer’s first step, and the paper reports improved vision results and Transformer training without warmup.

  • Problem

    Initialization for novel or complex architectures can be difficult to choose and may leave training unstable, while architecture-specific schemes are not always portable.

    Source: §1 · p. 1

  • Method

    GradInit learns scales for randomly initialized parameter blocks so a prescribed optimizer’s first step minimizes loss while limiting gradient norm.

    Source: §3.1 · p. 3, §3 · p. 3

  • Results

    GradInit accelerates convergence and improves test performance across vision architectures, and enables training the original Transformer without learning-rate warmup using Adam or SGD.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Visualizing learned weight norms and gradient variances can help identify potential initialization-instability causes and suggest useful scale patterns.

    Source: §1 · p. 2

Abstract

from arXiv · show

Innovations in neural architectures have fostered significant breakthroughs in language modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-parameter choices and training instability if the network parameters are not properly initialized. A number of architecture-specific initialization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on a simple heuristic; the norm of each network layer is adjusted so that a single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing a scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many convolutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer architecture for machine translation, enabling training it without learning rate warmup using either Adam or SGD under a wide range of learning rates and momentum coefficients. Code is available at https://github.com/zhuchen03/gradinit.

1 Introduction

GradInit learns architecture-agnostic initialization scales to address training instability that can persist in complex or novel architectures. The paper reports benefits across vision and language tasks, including training the original Transformer without learning-rate warmup.

1 Introduction

  • Analytically derived initialization rules are difficult to transfer to architecture search, branched or heterogeneous networks, and entirely new architectures.Even carefully designed rules may not prevent instability in deep or branched models.Source: §1 · p. 1
  • Instead of deriving distribution parameters analytically, GradInit learns a scalar for each random weight tensor to make one prescribed optimizer step reduce training loss while controlling initial gradient norms.The authors report that the method can make initialization smoother, reduce inter-sample gradient variance, and accelerate training.Source: §1 · pp. 1–2
  • GradInit accelerates convergence and improves test performance across vision architectures, enables the original Transformer to train without warmup using Adam or SGD, and raises ResNet-1202 test accuracy above ResNet-110.The paper also reports that ResNet-1202 outperforms other initialization methods that had failed to achieve that result.Source: §1 · p. 2
  • GradInit visualizations use initial weight norms and gradient variances to identify potential initialization-instability causes and reveal scale patterns that may inform initialization rules.The authors cite normalization layers as one possible source of instability.Source: §1 · p. 2

2 Related Work

Prior initialization analyses and methods address different architectural settings, while normalization, skip connections, and warmup provide other stabilization strategies. GradInit is positioned as a more efficient, stochasticity-aware approach that accounts for optimizer dynamics.

2 Related Work

  • Feed-forward initialization analyses based on fan-in [2] and activation variance [15] have architectural limits, while moment-preserving schemes [17] are also restricted to feed-forward networks.The activation-rescaling procedure fails to prevent activation blowup with depth in ResNets.Source: §2 · p. 2
  • Normalization layers [19] [20], skip connections [21], and Transformer learning-rate warmup [3] stabilize some complex models, but do not eliminate high initial gradient variance in deep networks or poorly normalized architectures [9] [10] [11] [4].The cited passage associates this residual variance with training instability.Source: §2 · p. 2
  • Skip-connection initialization methods scale layers or add learnable factors to stabilize activation variance or updates [9] [10] [11] [12] [6] [22], but are often architecture-specific.Related analyses also extend to Transformer self-attention.Source: §2 · p. 2
  • Compared with MetaInit, GradInit uses training data, avoids repeated third-order derivatives, accounts for minibatch stochasticity, and incorporates optimizer learning rate and update direction.For CIFAR-10 at batch size 64, MetaInit required three RTX 2080Ti GPUs while GradInit required one.Source: §2 · pp. 2–3

3 Method

GradInit learns scale factors for randomly initialized parameter blocks so a prescribed stochastic optimizer's first step yields low loss under a gradient-norm constraint. Its solver alternates between enforcing that constraint and optimizing loss, using two minibatches to account for stochasticity.

3 Method

  • GradInit learns parameter-block scales to minimize loss after one stochastic optimizer step while constraining the gradient norm, and can reduce gradient variance and place initialization in a smoother loss region.Rescaling scale-invariant weights, such as convolution kernels before batch normalization, can still change their effective learning dynamics.Source: §3 · p. 3

3.1 Efficient Learning-based Initialization via Constrained Optimization

  • The method multiplies fixed, randomly initialized weight matrices by learned nonnegative scales, then folds those scales into the weights before ordinary training.It minimizes post-step loss on a second minibatch using a gradient computed on the first; the prescribed optimizer and learning rate determine the update.Source: §3.1 · pp. 3–4

3.2 Solving the Constrained Problem

  • The solver updates scales to reduce gradient norm when the constraint is violated and otherwise updates them to reduce loss, with a 0.01 lower bound on each scale.The authors report convergence in fewer than 400 CIFAR-10 iterations or within 2,000 ImageNet iterations, taking less than 1% of training time; they default to a second minibatch mixing 50% samples from the first minibatch with 50% resampled training data.Source: §3.2 · p. 4

3.3 Setting and Enforcing the Constraint

  • A gradient-norm constraint prevents trivial loss reduction through an excessively large step; first-order bounds motivate ℓ2 for SGD and ℓ1 for Adam, with γ chosen so ηγ² or ηγ equals 0.1, respectively.The authors argue the constrained form avoids second-order gradients when the constraint is satisfied and avoids choosing a penalty strength λ across architectures.Source: §3.3 · p. 5

4 Experiments

The experiments evaluate GradInit across image-classification architectures and Post-LN Transformers, testing whether optimizer-aware initialization improves training stability and performance. Results include faster ImageNet convergence, gains on deep ResNets, and warmup-free Transformer training with SGD or Adam.

4.1.1 Settings

  • The evaluation spans five CIFAR-10 architectures, ImageNet ResNet-50, and IWSLT-14 DE-EN; GradInit or MetaInit takes one epoch, under 1% of total training time in the image experiments.The image experiments compare against Kaiming initialization, one-epoch constant-rate or warmup training, and MetaInit; the Transformer study also compares GradInit with Admin.Source: §4 · p. 5, §4.1.1 · p. 6

4.1.2 Results and Analysis

  • On ResNet-1202, GradInit exceeds warmup’s final test accuracy by 0.8% and achieves higher accuracy than ResNet-110.Warmup accelerates convergence and improves final accuracy relative to other initializations, but GradInit remains ahead in final accuracy.Source: §4.1.2 · p. 7
  • For ResNets with Batch Normalization, GradInit increases the norms of linear layers other than the final FC layer; within blocks, it scales the first BN up and the second down.The paired BN scaling increasingly downscales deeper residual branches, reducing activation variance; increasing the first BN scale also helps limit the second BN’s gradient-magnifying effect.Source: §4.1.2 · pp. 7–8, Figure 1 · p. 7
  • For ResNet-110 on CIFAR-10 with AdamW, GradInit raises Acc1/Accbest from 36.6 ± 4.7/94.9 ± 0.1 to 40.2 ± 0.2/95.3 ± 0.1.The experiments use a cosine learning-rate schedule with initial learning rate 3 × 10^-3 and weight decay 0.2.Source: §4.1.2 · p. 8
  • GradInit rescales both BN and other linear layers, while simply learning layer scalars during training can cause instability and worse results.The authors note that the scalars optimized by GradInit are not learned during subsequent network training.Source: §4.1.2 · p. 8, Table 5 · p. 8
  • GradInit accelerates ResNet-50 convergence and improves test accuracy on ImageNet with or without Batch Normalization, with greater acceleration than FixUp even on the FixUp-designed architecture.The comparison is reported in Table 6; the experiments used a smaller batch for GradInit because of GPU memory limits.Source: Table 6 · p. 8, §4.1.2 · pp. 8–9

4.2 Training the Original Transformer Model without Warmup

  • On IWSLT-14 DE-EN, GradInit trains the original Post-LN Transformer without warmup using SGD or Adam; with 780 initialization steps, it scores 36.0 BLEU versus 35.4 for 780-step warmup at ηmax = 2 × 10^-4.The warmup comparison is reported for the specified Adam learning-rate setting; the paper also reports warmup-free SGD performance close to Adam’s under the same schedule.Source: §4.2 · pp. 9–10, Table 7 · p. 9
  • GradInit generally produces more stable warmup-free Transformer training than Admin, while their best performance numbers are almost the same; at ηmax = 5 × 10^-4 and β2 = 0.995, wskip raises average BLEU to 36.0 from 8.9 without it.Without wskip, only one of four experiments exceeded BLEU 35 in this setting.Source: §4.2 · pp. 9–10

5 Conclusion

GradInit learns parameter-block scales by minimizing loss after one optimizer step, accounting for optimizer stochasticity, learning rate, and direction. The paper reports improved vision performance and warmup-free Post-LN Transformer training, and presents the learned scales as a guide for initialization design.

5 Conclusion

  • GradInit accelerates convergence and improves test performance across image-classification architectures, and enables Post-LN Transformer training without learning-rate warmup, even with SGD.The authors also note that learned scaling coefficients can guide initialization design for complex architectures and shorten training schedules.Source: §5 · p. 10
  • GradInit learns a scale factor for each randomly initialized parameter block by minimizing training loss on a different minibatch after one optimizer step.The design takes the optimizer's stochasticity, learning rate, and direction into account.Source: §5 · p. 10
  • GradInit initialization often decreases gradient variance for most parameter blocks.Source: §5 · p. 10
  • Analyzing learned scales and their effects on gradient variance can guide initialization design for complex architectures and help shorten training schedules and save energy.Source: §5 · p. 10

A.1 On CIFAR-10

The CIFAR-10 experiments cover several convolutional architectures, with and without Batch Normalization, and include a comparison to FixUpResNet. The appendix specifies the training and GradInit settings used.

A.1 On CIFAR-10

  • The CIFAR-10 experiments include VGG-19, ResNet-110 and its 1202-layer version, WRN-28-10, DenseNet-100, and FixUpResNet.For three networks, the experiments also remove Batch Normalization and add learnable biases; FixUpResNet replaces Batch Normalization with biases and a learnable block scale.Source: §A.1 · p. 13
  • Models use batch size 128, except DenseNet-100 uses recommended batch size 64; training uses random cropping, random flipping, and cutout, with no dropout and weight decay 10^-4 in all cases.Source: §A.1 · p. 13
  • GradInit scale factors start at 1 with lower bound α = 0.01, and the reported results use four runs with the best τ selected from a tested grid.The tested τ values range from 10^-3 to 10^-1, with the selected value varying by architecture and Batch Normalization setting.Source: §A.1 · p. 13

A.2 On ImageNet

The ImageNet appendix describes data augmentation, training duration, and GradInit's short initialization run and hyperparameter selection.

A.2 On ImageNet

  • ImageNet models use random cropping and flipping, train for 90 epochs, and decay the learning rate tenfold every 30 epochs.Source: §A.2 · p. 13
  • GradInit runs for 2,000 iterations, less than half an epoch, and uses γ = 3 for ImageNet.The authors motivate γ by noting that random-guess cross-entropy on ImageNet is three times that on CIFAR-10.Source: §A.2 · p. 13
  • The ImageNet τ is selected from four tested values, with τ = 3 × 10^-3 chosen because it maximizes first-epoch test accuracy.Source: §A.2 · p. 13

A.3 On Machine Translation

The machine-translation appendix reports separate settings for SGD and Adam, including learning-rate searches and the GradInit iteration budget.

A.3 On Machine Translation

  • For SGD, GradInit uses η = 0.15 and γ = 1, while Admin's learning rate is searched over a grid and its best result occurs at 0.06.Admin diverges when ηmax > 0.06; the prescribed SGD learning rate is searched from 0.05 to 0.2.Source: §A.3 · pp. 13–14
  • For Adam, GradInit uses η = 5 × 10^-4 and γ = 10^3, with ηmax and β2 settings listed in Figure 3.Source: §A.3 · p. 14
  • GradInit's maximum run is 780 iterations, compared with a usual 4,000-iteration warmup stage; BLEU is evaluated each epoch and the best score during training is reported.Source: §A.3 · p. 14

B Mini-batching, continued: Choice of ˜S

The minibatch used to evaluate the post-update objective affects GradInit: 50% overlap balances objective variance and optimizer stochasticity. The BN analysis explains how scaling pre-BN weights can reduce backward gradient variance.

B Mini-batching, continued: Choice of ˜S

  • With r = 0.5, GradInit gives the most favorable reported VGG-19 (w/o BN) outcome; r = 0 can yield 10% first-epoch accuracy, below baseline.At r = 1, GradInit improves first-epoch variance and final test accuracy over Kaiming, but not as much as at r = 0.5.Source: B Mini-batching, continued: Choice of ˜S · p. 14
  • With no gradient-norm constraint, r = 1 encourages greedy updates that can cut same-minibatch loss by over 50% in one iteration but may blow up initial gradient norms.For r = 0.5, GradInit improves the baseline without the constraint; at r = 1, the smallest tested τ was needed for results not significantly worse than Kaiming.Source: B Mini-batching, continued: Choice of ˜S · p. 14, Table 8 · p. 14

C Magnification Effect of BN

  • When the pre-BN weights are rescaled, the normalized output can remain unchanged while the input variance changes, affecting gradient variance during backpropagation.The scale invariance claim assumes zero bias and small ε.Source: C Magnification Effect of BN · p. 15
  • Increasing the pre-BN weight scale can reduce gradient variance after BN, explaining why GradInit enlarges VGG-19 convolutional weights with BN.The analysis also describes BN's magnification effect when input feature variance is below 1.Source: C Magnification Effect of BN · pp. 14–15

D Improved Implementation of MetaInit

For its MetaInit comparison, the paper uses training data and Adam with gradient clipping, rather than MetaInit's original random-input, signSGD setup.

D Improved Implementation of MetaInit

  • MetaInit was originally task-agnostic and initialized using random inputs, while this comparison feeds it training data to tailor it to the task.The authors expected real training data to improve MetaInit for the specific task.Source: D Improved Implementation of MetaInit · pp. 15–16
  • The comparison optimizes MetaInit's weight norms with Adam and the same gradient clipping used for GradInit, instead of MetaInit's original signSGD with momentum.The authors report that Adam with their stated hyperparameters gave better MetaInit results than the original optimizer.Source: D Improved Implementation of MetaInit · p. 16
  • Table 9 compares MetaInit versions using Acc1 and Accbest across four runs, distinguishing random-data and real-data inputs.The supplied passage identifies the metrics and variants but does not provide their values.Source: Table 9 · p. 16

E Additional Experimental Results

Additional CIFAR-10 results cover WRN-28-10, FixUpResNet, and DenseNet-100. The weight-norm visualizations show architecture- and normalization-dependent patterns in how GradInit reduces gradient variance.

E Additional Experimental Results

  • Additional CIFAR-10 results include WRN-28-10, FixUpResNet, and DenseNet-100; the table reports first-epoch and best test accuracy over four runs.For WRN, MixUp was used in training, but not for GradInit, to test transferability across training regularizations.Source: E Additional Experimental Results · p. 16, Table 10 · p. 16

F Weight norms and gradient variances

  • With BN, VGG-19 scales up convolutional and most BN layers, while scaling down the last two BN layers and final fully connected layer.The figure description attributes the scaling to reducing BN's magnification effect and forward-pass variance.Source: Figure 4 · p. 17
  • With BN, ResNet-110 reduces gradient variance across layers by scaling down deeper residual blocks' second BN layers and scaling up first BN and convolutional layers.The caption links the scaling patterns to reducing forward-pass variance growth and BN's backpropagation magnification effect.Source: Figure 5 · p. 17
  • Without BN, GradInit downscales weights throughout VGG-19 and ResNet-110; the ResNet's uninitialized gradient variance is high from cumulative skip-connection effects.VGG-19 without BN also shows increasing gradient variance at convolutional layers with different input and output dimensions.Source: Figure 6 · p. 17
  • In DenseNet-100 without BN, GradInit reduces gradient variance by reducing weights across layers; with BN, the main reduction comes from downscaling the final fully connected layer.The no-BN caption says the reduction makes training without gradient clipping possible.Source: Figure 7 · p. 18, Figure 8 · p. 18
  • For the Post-LN Transformer, Adam-based GradInit reduces variance broadly except for some decoder query- and key-projection weights; SGD mostly reduces weights in the final decoder block.With Adam, major variance reduction comes from downscaling final decoder layer-normalization weights and residual-branch linear weights.Source: Figure 9 · p. 19, Figure 10 · p. 20
Loading 2102.08098v3…