PaperlayerView exact revision on arXiv

Source-linked AI summary

Robust Training of Neural Networks Using Scale Invariant Architectures

Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank J. Reddi, Sanjiv Kumar

arXiv:2202.00980v2cs.LGstat.ML

TL;DR

Training deep networks can be unstable, and adaptive methods used to address this carry memory costs, leaving open whether SGD can train robustly. The paper combines scale-invariant architecture with SGD and weight decay, optionally using relative global clipping; SIBERT trained with SGD achieves performance comparable to BERT trained with ADAM.

  • Problem

    Deep-network training can be unstable and often relies on adaptive, carefully tuned learning rates, which raises whether SGD can achieve robust training without adaptivity’s memory cost.

    Source: §1 · p. 1

  • Method

    The paper uses scale-invariant architecture with SGD and weight decay, optionally adding Relative Global Clipping designed to be invariant to specified loss, learning-rate, weight-decay, and initialization rescalings.

    Source: §1 · pp. 2–3, §3.3 · p. 7

  • Results

    SIBERT trained with simple SGD achieves performance comparable to standard BERT trained with ADAM, and the paper reports convergence dependence only logarithmic in initialization and loss scale.

    Source: §1 · p. 3, §6 · p. 13

  • Takeaways & Limitations

    The paper concludes that jointly designing neural architecture and optimization can support robust transformer training with a substantially lower memory footprint than adaptive methods.

    Source: §6 · p. 13

    With fixed ηλ, longer training does not guarantee further convergence, because weight decay can drive the norm toward zero and eventually make the dynamics unstable.

    Source: §4.1 · p. 8

Abstract

from arXiv · show

In contrast to SGD, adaptive gradient methods like Adam allow robust training of modern deep networks, especially large language models. However, the use of adaptivity not only comes at the cost of extra memory but also raises the fundamental question: can non-adaptive methods like SGD enjoy similar benefits? In this paper, we provide an affirmative answer to this question by proposing to achieve both robust and memory-efficient training via the following general recipe: (1) modify the architecture and make it scale invariant, i.e. the scale of parameter doesn't affect the output of the network, (2) train with SGD and weight decay, and optionally (3) clip the global gradient norm proportional to weight norm multiplied by $\sqrt{\tfrac{2\lambda}{\eta}}$, where $\eta$ is learning rate and $\lambda$ is weight decay. We show that this general approach is robust to rescaling of parameter and loss by proving that its convergence only depends logarithmically on the scale of initialization and loss, whereas the standard SGD might not even converge for many initializations. Following our recipe, we design a scale invariant version of BERT, called SIBERT, which when trained simply by vanilla SGD achieves performance comparable to BERT trained by adaptive methods like Adam on downstream tasks.

1 Introduction

The introduction links training instability in deep networks to homogeneous architectures and asks whether SGD can train robustly without adaptive methods. It proposes scale-invariant architecture with SGD and weight decay, optionally adding relative global clipping, and reports comparable SIBERT and BERT performance.

1 Introduction

  • In a logistic-regression example, a reparameterized 2k-homogeneous loss can make gradient descent initialization-sensitive and cause parameter norms to explode for some learning rates.The corresponding loss in the product variable is convex with bounded smoothness and admits initialization-independent step sizes.Source: §1 · p. 2
  • Adaptive methods such as ADAM are used to handle deep networks’ k-homogeneous structure and support robust training, but require more memory.The introduction asks whether non-adaptive SGD can achieve fast, robust convergence without instability.Source: §1 · p. 2
  • The proposed recipe combines scale-invariant architecture, SGD with weight decay, and optional Relative Global Clipping.The authors apply the recipe to create Scale Invariant BERT (SIBERT).Source: §1 · p. 3
  • The paper proves convergence guarantees for scale-invariant losses and reports that SGD with weight decay is robust to initialization and loss scale.The stated rates match standard rates without requiring knowledge of loss smoothness.Source: §1 · p. 3
  • Under assumptions limiting clipping bias in expected gradients, SGD with weight decay and Relative Global Clipping has improved parameter-norm convergence and similar convergence results to SGD with weight decay.The clipping analysis is presented as a novel analysis.Source: §1 · p. 3
  • SIBERT trained with simple SGD achieves performance comparable to standard BERT trained with ADAM.The authors also report verifying their theoretical claims and describe this as the first effective training of a BERT-like model using vanilla SGD.Source: §1 · p. 3

2 Related Work & Background

The related-work discussion situates the paper among adaptive optimizers, clipping methods, and scale-invariant networks. The preliminary section defines homogeneity and summarizes derivative identities and scaling equivalences used in the paper.

2 Related Work & Background

  • Adaptive methods including ADAGRAD, ADAM, and RMSPROP are established in deep learning, while reduced-adaptivity alternatives can still require momentum for stability in large models.The discussion also relates the paper to normalized and projected gradient descent, which it says has seen limited success.Source: §2 · pp. 3–4
  • Gradient clipping is widely used for stability or privacy, but constant thresholds require tuning and can degrade performance; adaptive clipping adds computation and memory costs.The paper frames its clipping rule against these threshold-selection and efficiency concerns.Source: §2 · p. 4
  • Scale invariance commonly comes from normalization schemes, and prior work studies its interaction with SGD and weight decay, including norm convergence and effective learning rates.The discussion also cites convergence results for GD and SGD on normalized networks.Source: §2 · p. 4

2.1 Preliminary

  • A function is scale invariant when it is 0-homogeneous; the preliminaries also state Euler’s identity for homogeneous functions and derivative homogeneity.The scaling lemmas describe rescalings that preserve GD iterates or their directions.Source: §2.1 · pp. 4–5

3 Methods

The paper’s training recipe combines scale-invariant architecture with SGD and weight decay, optionally adding Relative Global Clipping. It applies this design to BERT as SIBERT, addressing instability and the difficulty of choosing clipping thresholds under rescaling.

3.1 Designing Scaling Invariant Architectures

  • SIBERT makes both transformer residual blocks and attention scale invariant, using a PreNorm residual variant and a novel Scale Invariant Attention design.The attention design replaces softmax with ReLU followed by row-wise normalization, making the score computation scale invariant under the stated condition on Q, K, and V.Source: §3.1 · pp. 5–6

3.2 Training Algorithm: SGD + WD

  • Scale invariance alone does not ensure robust, efficient SGD: parameter norms can grow and slow training, so the method uses weight decay to reduce the norm.With mild assumptions, prior results cited in the passage show norm stabilization in O(1/(ηλ)) steps and effective learning rate proportional to Θ(√λη).Source: §3.2 · p. 6
  • SGD+WD is insensitive to loss rescaling, initialization rescaling, and reparameterization, with convergence requiring up to |log c|/(λη) extra steps under scaling by c.The passage says this reduces hyperparameter-tuning effort and improves robustness when moving between codebases or frameworks with different default scalings or parametrizations.Source: §3.2 · p. 6

3.3 Relative Global Clipping

  • Relative Global Clipping sets a scale-consistent clipping threshold for SGD+WD, aiming to remove only extremely large gradients without triggering often.The threshold is designed to preserve behavior under the stated loss, learning-rate, weight-decay, and initialization rescalings; under mild assumptions on clipping bias, the recipe converges to approximate stationary points with logarithmic dependence on initialization and loss scale.Source: §3.3 · p. 7

4 Theoretical Analysis

The analysis studies convergence of weight-decayed optimization for scale-invariant losses, including deterministic GD, stochastic SGD, and SGD with Relative Global Clipping. It establishes robustness to initialization and loss scaling, while identifying conditions and a fixed-parameter limitation for the results.

4 Theoretical Analysis

  • After a scaling-dependent initial period, convergence to an approximate stationary point has a rate depending only on ηλ and unaffected by loss rescaling, initialization rescaling, or reparameterization.The initial period changes logarithmically with the loss and initialization scale.Source: §4 · p. 8

4.1 Convergence of GD +WD

  • GD+WD can tolerate broad choices of loss scaling, learning rate, and initial norm when ηλ is set relative to the training budget, but fixed ηλ does not guarantee further convergence with longer training.The paper attributes the limitation to weight decay dominating after directional convergence, shrinking the norm and eventually making the dynamics unstable.Source: §4.1 · p. 8

4.2 Convergence of SGD+WD

  • SGD+WD matches the standard O(T^-1/4) non-convex SGD rate and depends only logarithmically on the initialization scale under the stated assumptions.The result assumes bounded stochastic gradient norm and the conditions described for Theorem 4.5.Source: §4.2 · pp. 9–10

4.3 Convergence of SGD with Relative Global Clipping

  • Clipped SGD+WD’s norm-convergence result needs no assumption on the maximum gradient norm, unlike the unclipped result.Under an additional assumption, the analysis also yields convergence to stationary points; the authors describe that condition as reasonable given SIBERT’s 1.5% clipping frequency.Source: §4.2 · p. 9, §4.3 · p. 11

5 Experiments

The experiments report competitive SIBERT performance with SGD+WD, including downstream results close to BERT trained with ADAM. They also test training stability across initialization scales and the effects of weight decay and clipping.

Experimental Setup

  • The experiments assess competitive training performance, low memory use, robustness to initialization scale, and norm convergence with clipping.The setup uses base-size models unless otherwise specified; SIBERT's scale-invariant portion is trained with SGD+WD, while non-scale-invariant parts use LAMB.Source: §5 · p. 11

Performance

  • SIBERT's pretraining curve closely follows BERT trained by ADAM, while standard BERT with SGD either trains slowly at small learning rates or diverges.The reported comparison attributes improved training stability to the scale-invariant architecture.Source: §5 · p. 12, Figure 3 · p. 12
  • SIBERT trained with SGD+WD performs comparably to BERT trained with ADAM on downstream tasks, and the gap narrows when SIBERT gets twice the training budget.The table caption also reports that BERT with ADAM uses 3X more memory than SGD.Source: §5 · pp. 12–13, Table 1 · p. 13

Training Stability: Insensitivity to the scale of initialization

  • Across initialization scales spanning two orders of magnitude, the proposed approach converges well, whereas SGD alone struggles even with a scale-invariant architecture.The experiments also describe equivalent rescalings of loss, initialization, learning rate, and weight decay for scale-invariant models trained with SGD+WD, with or without clipping.Source: §5 · p. 13

Training Stability: Improvement in parameter norm convergence

  • SGD+WD, with or without clipping, quickly moves parameter norms toward desired ranges from different initial scales; SGD struggles when initialization and learning rate are misaligned.Figure 2 reports that relative global clipping reduces spikes in the encoder parameter-norm curve.Source: Figure 2 · p. 12, §5 · p. 13

6 Conclusion

The conclusion presents an architecture-aware procedure for robust transformer training with non-adaptive methods. It reports low memory use relative to adaptive methods and points to joint architecture-optimizer design as a future direction.

Conclusion

  • The paper presents a method for robustly training transformers with non-adaptive methods such as SGD.It combines a scale-invariant architecture with a tailored optimization procedure.Source: §6 · p. 13
  • The authors report robust neural-network training with substantially lower memory footprint than adaptive methods.Source: §6 · p. 13
  • The conclusion identifies joint design of neural architectures and optimizers as a direction for future training procedures.Source: §6 · p. 13

A Design Details of Scale Invariant BERT

SIBERT makes its encoder scale invariant while retaining a separate classification head. Its design uses homogeneous building blocks, including a PreNorm residual structure and a scale-invariant attention mechanism.

Homogeneity

Overview of SIBERT structure

  • SIBERT keeps the encoder and classification head as separate parts, making only the encoder scale invariant and training it with SGD+WD; the head uses LAMB and is used only for pretraining.Source: A Design Details of Scale Invariant BERT · p. 17

(2;2)-homogeneous encoder layer

  • The encoder uses a PreNorm residual block and replaces BERT's GeLU with ReLU, which is 1-homogeneous.The passage says the ReLU can be omitted from the design description without affecting final scale invariance.Source: A Design Details of Scale Invariant BERT · p. 17

Building blocks

  • The design's figures identify a (2;2)-homogeneous encoder layer, a (0;2)-homogeneous feedforward layer, and an attention block described as (x, x, x; x + 2)-homogeneous.Table 2 summarizes the homogeneity of SIBERT building blocks.Source: Table 2 · p. 17, Figure 5 · p. 18, Figure 6 · p. 18

Scale Invariant Attention

  • The proposed Scale Invariant Attention replaces softmax with ReLU followed by normalization; the passage says this makes the attention-score computation scale invariant when Q, K, and V are positively homogeneous.Source: A Design Details of Scale Invariant BERT · p. 18

B Introduction examples analysis

The appendix analyzes an initialization-sensitive failure mode for homogeneous models and collects scale-invariance and probability results used in the analysis.

B Introduction examples analysis

  • In the introductory example, an initialization-dependent learning-rate condition can make the iterates increase monotonically and explode when the initial value exceeds the finite optimum.Source: B Introduction examples analysis · pp. 18–19

C.1 Scale Invariance

  • The scale-invariance appendix gives smoothness results, including a multi-group version, and a gradient-norm bound.Source: §C.1 · p. 19

C.2 Probablity

  • The probability appendix states conditional Chernoff and Azuma bounds, alongside a sub-Gaussian sum result.Source: §C.2 · p. 20

D Omitted Proofs for the Convergence of GD

The omitted proof uses scale invariance to relate gradient behavior to parameter norms, then bounds the iterates to complete the convergence argument. Its intermediate lemmas control when the norm remains small and how it can grow.

E Omitted Proofs for Convergence Rate of SGD

These omitted proofs establish norm bounds used in the SGD convergence analysis under Condition 4.4. The bounds are conditioned on specified events and use additional restrictions on the learning-rate and weight-decay product.

F Omitted Proofs for Convergence of SGD with Relative Global Clipping

The omitted proofs analyze clipped SGD through properties of a function G_P,C and bounds on the parameter norm. Under Assumption 4.7, the warm-start results bound the norm, and Theorem F.9 addresses its convergence.

G Convergence of SGD for multi-group scale invariant functions

The section extends SGD-with-weight-decay analysis to multi-group scale-invariant functions, a setting that includes feedforward networks with normalization after each layer. The extension requires a corresponding lower bound on expected gradient norm because the single-group assumptions can fail.

Related papers

Loading 2202.00980v2…