PaperlayerView exact revision on arXiv

Source-linked AI summary

FAdam: Adam is a natural gradient optimizer using diagonal empirical Fisher information

Dongseong Hwang

arXiv:2405.12807v11cs.LGcs.AIcs.ITmath.IT

TL;DR

Natural gradient descent accounts for loss-landscape curvature through the Fisher information matrix, whose computation can be prohibitively expensive for large models, while Adam's theoretical connection to natural gradient descent has lacked a comprehensive explanation. This paper analyzes Adam's diagonal empirical FIM and proposes Fisher Adam (FAdam), which performs across LLM, ASR, and VQ-VAE experiments and achieves state-of-the-art ASR results.

  • Problem

    Natural gradient descent's FIM can be prohibitively expensive for large models, and a comprehensive theoretical account of Adam's square-root FIM connection has remained elusive.

    Source: §1 · p. 1

  • Method

    The paper explains Adam through Riemannian and information geometry, analyzes its diagonal empirical FIM and approximations, and proposes FAdam with optimizer and weight-decay corrections.

    Source: §1 · pp. 1–2

  • Results

    FAdam performs across LLM, ASR, and VQ-VAE experiments and achieves state-of-the-art ASR results, including state-of-the-art WER on LibriSpeech for 600M-parameter models.

    Source: §1 · p. 2, §4.2 · p. 10

  • Takeaways & Limitations

    The analysis clarifies constraints on loss selection and Adam's approximations, providing groundwork for future optimizers that mitigate those approximations.

    Source: §5 · p. 10

    Adam's EMA averages FIM components across parameter locations with potentially changing bases, an operation the paper says is mathematically invalid outside Euclidean space.

    Source: §3.4.1 · p. 7

Abstract

from arXiv · show

This paper establishes a mathematical foundation for the Adam optimizer, elucidating its connection to natural gradient descent through Riemannian and information geometry. We provide an accessible and detailed analysis of the diagonal empirical Fisher information matrix (FIM) in Adam, clarifying all detailed approximations and advocating for the use of log probability functions as loss, which should be based on discrete distributions, due to the limitations of empirical FIM. Our analysis uncovers flaws in the original Adam algorithm, leading to proposed corrections such as enhanced momentum calculations, adjusted bias corrections, adaptive epsilon, and gradient clipping. We refine the weight decay term based on our theoretical framework. Our modified algorithm, Fisher Adam (FAdam), demonstrates superior performance across diverse domains including LLM, ASR, and VQ-VAE, achieving state-of-the-art results in ASR.

1 Introduction

Natural gradient descent uses the Fisher information matrix (FIM) to account for loss-landscape curvature, but computing it can be prohibitively expensive for large models. This paper explains Adam as an approximation to natural gradient descent and uses that framework to motivate FAdam.

1 Introduction

  • Natural gradient descent uses the FIM as a Riemannian metric to account for loss-landscape curvature, but its computation can be prohibitively expensive for large models.Source: §1 · p. 1
  • The paper provides an accessible mathematical account of Adam through Riemannian and information geometry, showing it as an approximation to natural gradient descent.Source: §1 · p. 1
  • Adam's diagonal empirical FIM approximations motivate using log-probability losses based on discrete distributions, including categorical cross-entropy rather than L2 loss in the image domain.Source: §1 · p. 1
  • FAdam combines proposed Adam corrections with geometrically refined weight decay and performs across LLM, ASR, and VQ-VAE experiments, achieving state-of-the-art ASR results.The proposed corrections include natural-gradient momentum, adjusted bias corrections, adaptive epsilon, and gradient clipping.Source: §1 · p. 2

2 Background

The background introduces the geometric notation used in the paper and relates gradients and Fisher information to the statistical manifold. It presents the FIM as a metric underlying natural gradient descent.

2 Background

  • The paper adopts differential-geometry notation as introduced in Appendix A.1.Source: §2 · p. 2

2.1 Gradient on Riemannian manifold

  • A scalar field's differential is a covector; the gradient maps it into the parameter vector space, with its definition tied to the inner product and directional derivative.The paper follows machine-learning notation for gradient and natural gradient.Source: §2.1 · p. 2

2.2 Fisher Information Matrix (FIM) and Natural Gradient

  • The FIM is the statistical manifold's metric and defines the natural gradient; locally, it also gives the distance metric associated with KL divergence.Fisher information is described as the expected Hessian and as curvature of the log-likelihood.Source: §2.2 · p. 3

3 Toward FAdam

The paper develops FAdam from a geometric account of Adam’s diagonal empirical Fisher, addressing how the Fisher is estimated and how its geometry should shape updates. It recommends discrete log-likelihood losses and proposes changes to momentum and weight decay; the authors state that FAdam retains convergence guarantees under the cited assumptions.

3.1 Loss and FIM; 3.2 Diagonal Fisher Information

  • Natural-gradient optimization requires a log-likelihood loss, while Adam’s diagonal FIM reduces the quadratic scaling of full-FIM methods to linear memory complexity.The diagonal approximation retains variances but discards covariance information.Source: §3.1 · p. 3, §3.2 · p. 4

3.3 Empirical Fisher Information

  • Because sampling from the model distribution is difficult, the paper estimates FIM from the empirical training distribution, using minibatches and EMA to manage computational cost and uncertainty.Adam further reuses the minibatch gradient to estimate FIM.Source: §3.3 · pp. 4–5

3.3.1 Discrete vs Continuous probability distributions

  • Discrete distributions support more accurate empirical Fisher approximations than continuous ones, motivating categorical cross-entropy instead of L2 loss for images.The proposed image setup uses one-hot encoding and predicts 256 values per RGB channel.Source: §3.3.1 · p. 6

3.4.1 Reciprocal vs Reciprocal Square-root for FIM

  • To address EMA’s averaging across changing tangent-space bases, the paper constructs an equal-length Euclidean gradient for averaging and identifies the square-root FIM as optimal in its ablation.The paper calls this constructed vector the invariant natural gradient.Source: §3.4.1 · p. 7

3.4.2 Momentum

  • FAdam averages the invariant natural gradient rather than raw gradients and omits zero bias correction for momentum, which the authors describe as an informed warmup.The paper distinguishes Adam’s second accumulator as FIM rather than momentum.Source: §3.4.2 · p. 8

3.4.3 Weight decay on manifold

  • Weight decay is kept outside the log-likelihood loss to avoid affecting FIM estimation, then transformed into a natural-gradient update using the existing FIM.The paper frames this as a rule for auxiliary losses: include them only when related to log-likelihood.Source: §3.4.3 · p. 8

3.4.4 FAdam; 3.4.5 Convergence analysis

  • FAdam makes minimal changes to Adam and retains a convergence guarantee: as iterations grow, the expected gradient becomes sufficiently small under the stated assumptions.The authors refer to their convergence analysis of Adam as the basis for this guarantee.Source: §3.4 · p. 6, §3.4.5 · p. 9, §3.4.4 · p. 9

4 Experiment

Experiments compare FAdam variants with Adam-based optimizers across language, speech, and image tasks, with the paper reporting gains and a new LibriSpeech WER state of the art. An ablation study of FAdam hyperparameters is reported in the appendix.

4 Experiment

  • FAdam outperforms the tested Adam variants across LLM, ASR, and VQ-VAE experiments, with new state-of-the-art WER on LibriSpeech.The experiments replaced existing optimizers while keeping other hyperparameters unchanged; Figure 1 summarizes the domain comparisons.Source: Figure 1 · p. 9, §4 · p. 9, §4.2 · p. 10

4.1 LLM (text domain)

  • FAdafactor outperforms Adafactor on evaluation loss for the 1B-parameter PaLM model pretrained on C4.Both used β1=0.9, β2=0.99, and λ=0.001; epsilon was 1e-8 for FAdafactor and 1e-30 for Adafactor.Source: Figure 1 · p. 9, §4.1 · p. 10

4.2 ASR (speech domain)

  • FAdam outperforms Adam and establishes a new state-of-the-art WER on LibriSpeech for 600M-parameter Conformer models.The comparison used identical hyperparameters for both optimizers; the models were fine-tuned semi-supervisedly on LibriLight data.Source: Figure 1 · p. 9, §4.2 · p. 10, Table 1 · p. 10

4.3 VQ-VAE (image domain)

  • FAdam outperforms AdamW on ImageNet generation, while categorical cross-entropy outperforms L2 loss and removes the need for logit-laplace transformation.The experiment trained a 100M-parameter ViT VQ-GAN model.Source: Figure 1 · p. 9, §4.3 · p. 10

4.4 Ablation study

  • The paper reports hyperparameter ablation studies for FAdam, with results placed in Appendix C.1.The main paper omits the results due to page limitations.Source: §4.4 · p. 10

5 Conclusion

The conclusion presents a mathematical foundation for Adam, clarifying loss-function constraints and derivation approximations, and reports FAdam's effectiveness across diverse domains. The appendix introduces geometric notation for the analysis.

5 Conclusion

  • The paper develops a mathematical foundation for Adam that clarifies loss-function constraints and the approximations in its derivation.It presents FAdam based on this framework and reports effectiveness across diverse domains.Source: §5 · p. 10

A.1 Tensor Calculus Notation

  • The appendix defines tangent-space vectors through basis vectors and components, and introduces covariant differentiation for vector and tensor fields.The covariant derivative extends directional differentiation from scalar fields.Source: §A.1 · p. 15
  • The appendix places model parameters in a vector space and loss differentials in a covector space, motivating their transformation for optimization operations.The directional derivative represents the loss derivative in a vector's direction.Source: §A.1 · p. 15

A.2 The proof of Fisher is negative Hessian

This appendix presents a proof relating Fisher information to the expectation of the negative Hessian.

A.2 The proof of Fisher is negative Hessian

  • The appendix sets out to prove the relationship between Fisher information and the negative Hessian's expectation.The following derivation is identified as the proof of that relationship.Source: §A.2 · p. 15
  • The appendix frames its proof using covariant differentiation, whose scalar-field case is the directional derivative.The relationship being established is between Fisher information and the expected negative Hessian.Source: §A.2 · p. 15
  • For scalar fields, the covariant derivative is called the directional derivative.The passage distinguishes this scalar-field use from covariant derivatives of vector and tensor fields.Source: §A.2 · p. 15

A.3 The proof of KL approximation

The appendix proves the KL approximation by expanding the divergence around an infinitesimal parameter displacement. The Fisher information emerges in the second-order term.

A.3 The proof of KL approximation

  • A second-order Taylor expansion is used to approximate KL divergence for an infinitesimal displacement.The proof follows the approach presented in Kristiadi.Source: §A.3 · p. 16
  • At zero displacement, the KL divergence is zero, and the first-order approximation term also vanishes.The proof shows the first-order term vanishes through the derivative calculation.Source: §A.3 · p. 16
  • The Fisher information appears in the second-order approximation term, completing the proof of the stated KL approximation.The derivation uses the Hessian property in Eq. (8).Source: §A.3 · p. 16

B.1 Compared to Newton’s second-order method

The appendix contrasts natural-gradient optimization with Newton’s method and details the approximations and algorithmic choices behind Adam and its Fisher variants. It also describes clipping, adaptive epsilon, and implementation differences in FAdam and FAdafactor.

B.1 Compared to Newton’s second-order method

  • Natural gradient uses an inverse-Hessian-related update but requires log-likelihood and is described as convergent through the positive semi-definiteness of FIM, unlike Newton’s method.The appendix also notes that Newton’s method has no loss-function restriction and often uses trust-region techniques to mitigate divergence.Source: §B.1 · pp. 16–17

B.2 Conditional probability distribution for empirical Fisher Information

  • Conditional empirical FIM approximates expectations with training data and then over model labels, while reusing the gradient requires a non-principled move of the square outside the expectation.The authors say eliminating these non-trivial approximations is a direction for future research.Source: §B.2 · pp. 17–18

B.3 Clipping and Epsilon

  • FAdam clips the invariant natural gradient before momentum and adapts epsilon as min(10^-8, 0.01RMS(g_t)) to reduce manual tuning across models.The appendix reports no benefit from global clipping and says the adaptive epsilon performs comparably to manual tuning.Source: §B.3 · p. 18

B.4 FAdafactor

  • The appendix presents a Fisher Adafactor variant, FAdafactor, as a modification of Adafactor.The full algorithm is provided in Algorithm 2.Source: §B.4 · p. 18, §B.5 · p. 19

B.5 Adam and Adafactor

  • Adafactor had already implemented natural-gradient momentum, omitted momentum bias correction, and clipped gradients; FAdam’s algorithm additionally adapts epsilon and applies geometry-based weight decay.The appendix gives the FAdam update sequence and notes that the original Adafactor paper does not fully detail its algorithm.Source: §B.5 · p. 19

B.6 Convergence proof

The convergence analysis follows prior corrected Adam analysis and proves a convergence result for FAdam without clipping. It derives the bound without momentum for simplicity, while a tighter momentum-inclusive bound is left for future work.

B.6 Convergence proof

  • The convergence proof excludes clipping, which the authors identify as widely used in practice but not covered by existing Adam convergence proofs.They state that their proof excludes clipping effects.Source: §B.6 · p. 19
  • For simplicity, the analysis derives the convergence bound without momentum, leaving a tighter bound that incorporates momentum for future work.The authors relate FAdam’s natural-gradient-then-momentum procedure to Polyak momentum applied to SGD.Source: §B.6 · p. 20
  • The stated convergence result assumes bounded stochastic gradients, a smooth objective with an L-Lipschitz gradient, and an objective bounded below.The result says the expected squared gradient norm diminishes as iterations increase, indicating convergence to the optimal value.Source: §B.6 · p. 20

C.1 Ablation study

The ablations examine the FIM exponent and epsilon choices across LLM, ASR, and VQ-VAE settings. They support the standard exponent and report competitive performance from adaptive epsilon without manual tuning.

C.1.1 Exponent of FIM

  • FIM exponents above 0.5 are relatively stable, while values below 0.3 sharply reduce performance; the standard 0.5 exponent appears well justified.Table 2 reports the effects of varying the exponent on LLM and ASR performance.Source: §C.1.1 · p. 21, Table 2 · p. 21

C.1.2 Epsilon

  • Adaptive epsilon performs competitively without manual tuning, although optimal fixed epsilon varies by model size and domain and very small values can cause divergence.The tested models include LLM (1B), ASR (600M), and VQ-VAE (100M).Source: §C.1.2 · p. 21, Table 3 · p. 21

C.1.3 Epsilon2

  • An epsilon2 value of 0.01 consistently yielded good performance across domains when epsilon was 1e-8.Table 4 reports the impact of epsilon2 on model performance.Source: §C.1.3 · p. 21, Table 4 · p. 21

C.1.4 Weight decay

The weight-decay ablations assess the revised mechanism in ASR and LLM settings. The appendix also specifies the compute resources and training setups used for the three model domains.

C.1.4 Weight decay

  • The enhanced weight-decay mechanism significantly improved FAdam’s performance in the ASR experiments, where Adam’s weight-decay parameter could generally be reused.Table 5 reports LibriSpeech WERs for 600M ASR models across weight-decay values.Source: §C.1.4 · pp. 21–22, Table 5 · p. 22
  • For 1B LLMs, the optimal weight-decay value of 1e-3 for Adafactor was also optimal for FAdafactor despite their different mechanisms.Table 6 reports evaluation loss across weight-decay values.Source: §C.1.4 · p. 22, Table 6 · p. 22

C.2 Experiments Compute Resources

  • The experiments trained the 1B LLM and 100M VQ-VAE for one day, while the 600M ASR was fine-tuned for half a day.The setups used TPUv5, TPUv4, and TPUv3 devices, respectively.Source: §C.2 · p. 22

Related papers

Loading 2405.12807v11…