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

## Abstract

Abstract 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 approximations to bridge the gap between the natural gradient and Adam, 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 (NGD) is a powerful optimization method that has shown promise in training machine learning models. Introduced by Amari [], and revisited recently [], NGD offers an alternative to traditional gradient descent by taking into account the curvature of the loss landscape. This is achieved through the use of the Fisher information matrix (FIM), which provides a Riemannian metric for the statistical manifold. Despite its potential benefits, NGD faces a significant computational challenge: the calculation of FIM, which can be prohibitively expensive for large models like deep neural networks.

Adam [] is the de facto standard optimizer, favored for its fast convergence and practicality. While the original paper mentions the second momentum term can be interpreted as the diagonal FIM, a comprehensive theoretical understanding of the square root FIM, akin to the approach first seen in AdaGrad [], has remained elusive []. This is in contrast to NGD, which utilizes the inverse FIM.

Our main contribution is to provide an accessible mathematical explanation of the Adam optimizer, drawing upon fundamental concepts from Riemannian geometry and information geometry, thereby demonstrating that Adam is an approximation of NGD.
This framework enables us to illuminate the origin of the square root term in Adam and to advocate for the use of a log probability function as the loss function when employing Adam.

We demonstrate that Adam utilizes the diagonal empirical Fisher information, providing a detailed explanation of both the diagonal approximation and the empirical approximation. Notably, the empirical approximation suggests that the loss function should be based on discrete distributions, which implies using categorical cross-entropy loss rather than L2 loss in the image domain. Furthermore, our analysis uncovers flaws in the original Adam, leading us to propose corrections such as averaging natural gradient by momentum, adjusting bias corrections, and introducing adaptive epsilon and gradient clipping.
We also enhance the weight decay using principles of Riemannian geometry. Our modified algorithm, named Fisher Adam (FAdam), demonstrates strong performance across various domains, as evidenced by our experiments with Large Language Models (LLMs) for text, Automatic Speech Recognition (ASR) for speech, and Vector Quantized Variational Autoencoder (VQ-VAE) for image. In particular, our ASR experiments achieve new state-of-the-art results.

## 2 Background

We use the notation from differential geometry, as introduced in Section A.1.

### 2.1 Gradient on Riemannian manifold

The inner product in the tangent space of a Riemannian manifold is defined by the Riemannian metric tensor ${\textnormal{g}}=g_{ij}d\theta^{i}\otimes d\theta^{j}$. The Riemannian metric tensor components, $g_{ij}$, are symmetric and positive definite. Therefore, $g_{ij}$ is always invertible, and its inverse is denoted as $g^{ij}$.

$$ $\displaystyle\vec{v}\cdot\vec{u}:={\textnormal{g}}(\vec{v},\vec{u})=g_{ij}d \theta^{i}\otimes d\theta^{j}(v^{k}\frac{\partial{}}{\partial{\theta^{k}}},u^{ r}\frac{\partial{}}{\partial{\theta^{r}}})=g_{ij}v^{k}u^{r}d\theta^{i}(\frac{ \partial{}}{\partial{\theta^{k}}})\otimes d\theta^{j}(\frac{\partial{}}{ \partial{\theta^{r}}})=g_{ij}v^{i}u^{j}\\$ (1) $$

The differential form of a scalar field $\phi$ (e.g., loss ${\mathcal{L}}({\bm{\theta}})$) is defined in the covector space, as shown in Eq. 33. However, since the model parameter ${\bm{\theta}}$ resides in the vector space, we need to transform the covector $d\phi$ into the vector space to perform operations involving both $d\phi$ and ${\bm{\theta}}$ (e.g., optimizing ${\bm{\theta}}$ as shown in Eq. 13).

The object transformed into the vector space is defined as the gradient $\phi$, denoted as $\nabla{\phi}$, and its definition is based on the inner product being equivalent to the directional derivative (Eq. 35), as shown in Eq. 2. By rearranging the terms with respect to $\nabla{\phi}$, we obtain the following equation for the gradient $\nabla{\phi}$ in the tangent space ${{\mathcal{T}}_{{\bm{\theta}}}{\mathcal{M}}}$, as shown in Eq. 3.(^1^11$g^{ij}$ is the inverse tensor component of $g_{ij}$, such that $g_{ij}g^{ij}=1$.)

$$ $\displaystyle\displaystyle(\nabla{\phi})^{i}v^{j}g_{ij}=\nabla{\phi}\cdot\vec{v}$ $\displaystyle:=d\phi(\vec{v})=v^{j}\frac{\partial{\phi}}{\partial{\theta^{j}}}$ (2) $\displaystyle\displaystyle(\nabla{\phi})^{i}=g^{ij}\frac{\partial{\phi}}{ \partial{\theta^{j}}}$ $\displaystyle\rightarrow\nabla{\phi}=g^{ij}\frac{\partial{\phi}}{\partial{ \theta^{j}}}\frac{\partial{}}{\partial{\theta^{i}}}$ (3) $$

In Euclidean space, the Riemannian metric $g_{ij}$ is Kronecker delta, so the gradient is usually expressed as follows:

$$ $\displaystyle\nabla{\phi}=\frac{\partial{\phi}}{\partial{\theta^{i}}}\frac{ \partial{}}{\partial{\theta^{i}}}$ (4) $$

In ML community, Eq. 4 is often denoted as the gradient $\nabla{\phi}$, while Eq. 3 is denoted as the natural gradient $\tilde{\nabla}{\phi}$. This notation conflicts with conventions in differential geometry. Since our paper is focused on ML, we will adhere to ML conventions from this point forward.

Typically, the component part of tensor operations are represented using matrix multiplication. The metric tensor component, $g_{ij}$, is represented by a matrix ${\bm{G}}$, and $\nabla{\phi}$ is a column vector.

$$ $\displaystyle\tilde{\nabla}{\phi}={\bm{G}}^{-1}\nabla{\phi}$ (5) $$

Equation 1 can also be expressed using matrix multiplication.

$$ $\displaystyle\vec{v}\cdot\vec{u}={\bm{v}}^{\top}{\bm{G}}{\bm{u}}$ (6) $$

### 2.2 Fisher Information Matrix (FIM) and Natural Gradient

The Fisher information quantifies the amount of information an observable random variable x (representing data) conveys about an unknown parameter ${\bm{\theta}}$ (representing a model parameter) influencing its probability. Given the probability mass function $P({\textnormal{x}}|{\bm{\theta}})$ (or probability density function $p({\textnormal{x}}|{\bm{\theta}})$) for the random variable x, the Fisher information is defined as the variance of the score function (i.e., the gradient of the log-likelihood), which is symmetric and positive semi-definite by definition.

$$ $\displaystyle{\bm{F}}({\bm{\theta}}):=\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{ \theta}}}[\nabla_{{\bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})\>\nabla _{{\bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})^{\top}].$ (7) $$

Fisher information is the expected value of the Hessian matrix. (Proof provided in Section A.2.) It represents the curvature of the log-likelihood on the statistical manifold where the model parameters ${\bm{\theta}}$ reside.

$$ $\displaystyle{\bm{F}}({\bm{\theta}})=-\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{ \theta}}}[\nabla_{{\bm{\theta}}}^{2}\log P({\textnormal{x}}|{\bm{\theta}})]=- \mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}[{\bm{H}}_{{\bm{\theta}}}( \log P({\textnormal{x}}|{\bm{\theta}}))].$ (8) $$

In the realm of statistical manifolds, the distance between ${\bm{\theta}}$ and ${\bm{\theta}}+{\bm{d}}$ is quantified by the Kullback-Leibler divergence $D_{\mathrm{KL}}(P({\textnormal{x}}|{\bm{\theta}})\|P({\textnormal{x}}|{\bm{
\theta}}+{\bm{d}}))$. For an infinitesimal displacement ${\bm{d}}$, a second-order Taylor series approximation reveals the Fisher information as the underlying distance metric [, , ]. A detailed proof can be found in Appendix A.3.

$$ $\displaystyle D_{\mathrm{KL}}(P({\textnormal{x}}|{\bm{\theta}})\|P({ \textnormal{x}}|{\bm{\theta}}+{\bm{d}}))\approx\frac{1}{2}{\bm{d}}^{\top}{\bm{ F}}({\bm{\theta}}){\bm{d}}.$ (9) $$

As mentioned in Eq. 6, the magnitude of a vector ${\bm{v}}$ in the tangent space ${{\mathcal{T}}_{{\bm{\theta}}}{\mathcal{M}}}$ can be expressed using the inner product, and the Fisher Information Matrix (FIM) serves as the Riemannian metric [].(^2^22Since FIM is symmetric and positive semi-definite, it produces a Pseudo-Riemannian manifold [].) Prior works [, , ] have demonstrated the importance of FIM in providing an inherent scale for learning.

$$ $\displaystyle\vec{v}\cdot\vec{v}=\lVert{\bm{v}}\rVert^{2}={\bm{v}}^{\top}{\bm{ G}}{\bm{v}}={\bm{v}}^{\top}{\bm{F}}({\bm{\theta}}){\bm{v}}$ (10) $$

In statistical manifolds, the gradient is described using the Fisher information ${\bm{F}}$ in Eq. 5, and this is called the natural gradient.

$$ $\displaystyle\tilde{\nabla}{\phi}={\bm{F}}^{-1}\nabla{\phi}$ (11) $$

The intuitive interpretation of Eq. 11 is that components with higher information undergo conservative movement, while components with lower information exhibit wider movement. (^3^33The inverse FIM exhibits a close relationship with the covariance matrix of the log likelihood, denoted as $\Sigma^{-1}_{\bm{\theta}}$. The Mahalanobis distance, $d_{M}({\bm{x}},{\mathcal{P}})^{2}$, which measures the distance between a data point ${\bm{x}}$ and a distribution ${\mathcal{P}}$, is defined as $d_{M}({\bm{x}},{\mathcal{P}})^{2}:=({\bm{x}}-{\bm{\mu}})^{\top}\Sigma^{-1}_{ \bm{x}}({\bm{x}}-{\bm{\mu}})$ []. Comparing this to Eq. 10, we discern a connection between the inverse Fisher information and the covariance matrix: higher covariance indicates lower information. This parallel echoes the principle in information theory where higher entropy corresponds to lower information.)

## 3 Toward FAdam

### 3.1 Loss and FIM

To utilize the natural gradient Eq. 11, we need to know both the gradient term and FIM. The gradient term requires the calculation of the loss ${\mathcal{L}}({\bm{\theta}})$, while the FIM requires the score function $\nabla_{{\bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})$. If the loss is expressed in the form of $-\log P({\textnormal{x}}|{\bm{\theta}})$, we eliminate the need to calculate the score function separately. Therefore, for using natural gradient optimizers like Adam(^4^44 We will discuss why Adam is considered a natural gradient optimizer in Section 3.4.) [], the loss function must be in the form of the log-likelihood.
We will delve deeper into the choice of loss function in Section 3.3.1.

$$ $\displaystyle\tilde{\nabla}{{\mathcal{L}}({\bm{\theta}})}={\bm{F}}^{-1}\nabla{ {\mathcal{L}}({\bm{\theta}})}=\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}} }[\nabla_{{\bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})\>\nabla_{{\bm{ \theta}}}\log P({\textnormal{x}}|{\bm{\theta}})^{\top}]^{-1}\nabla_{{\bm{ \theta}}}{-\log P({\textnormal{x}}|{\bm{\theta}})}$ (12) $$

The model parameter ${\bm{\theta}}$ is updated using the given natural gradient $\tilde{\nabla}{{\mathcal{L}}({\bm{\theta}})}$, where $\eta$ is the learning rate.

$$ $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta{\bm{F}}^{-1}\nabla{{ \mathcal{L}}({\bm{\theta}})}$ (13) $$

Natural gradient is considered a second-order method because FIM is the expected value of the Hessian, as shown in Eq. 8. A comparison with Newton’s method is provided in Section B.1.

### 3.2 Diagonal Fisher Information

One key reason for Adam []’s significant success compared to other second-order methods is its memory complexity. Adam scales linearly with the number of parameters, $O(N)$, while methods using FIM typically scale quadratically, $O(N^{2})$. For models with billions of parameters, $O(N^{2})$ is impractical.

Adam utilizes the diagonal Fisher information matrix. As discussed in 2.2, the inverse FIM approximates the covariance of the log likelihood. The diagonal FIM results in the loss of all covariance information except for the variances. Interestingly, Adam’s success story suggests that the loss of covariance information might not be detrimental in practice.

Let $\hat{{\bm{f}}}({\bm{\theta}})$ denote the diagonal FIM (^5^55Actually, from this point forward, the diagonal FIM is a vector.) obtained by diagonalizing FIM in Eq. 7. The gradient of loss function (12) is greatly simplified.

$$ $\displaystyle\displaystyle\hat{{\bm{f}}}({\bm{\theta}})$ $\displaystyle:=\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}[\nabla_{{\bm{ \theta}}}\log P({\textnormal{x}}|{\bm{\theta}})^{2}]$ (14) $\displaystyle\displaystyle\tilde{\nabla}{{\mathcal{L}}({\bm{\theta}})}$ $\displaystyle=\hat{{\bm{f}}}^{-1}\nabla{{\mathcal{L}}({\bm{\theta}})}=-\frac{ \nabla_{{\bm{\theta}}}{\log P({\textnormal{x}}|{\bm{\theta}})}}{\mathbb{E}_{{ \textnormal{x}}\sim P_{\bm{\theta}}}[\nabla_{{\bm{\theta}}}\log P({\textnormal {x}}|{\bm{\theta}})^{2}]}$ (15) $$

Amari u.a. [] proves that the off-diagonal elements of FIM are smaller than the diagonal elements by an order of $1/\sqrt{n}$, where $n$ represents the number of elements in the matrix. This finding justifies the use of the quasi-diagonal natural gradient method when the weight matrices of each layer are sufficiently large like LLM (Large Language Models).

Meanwhile, there have been efforts to capture important off-diagonal elements. For example, approximation methods utilizing low-rank approximations [, ] and Kronecker-factored approximations [, , , ] have been explored. The application of off-diagonal FIM to Adam is left for future study.

### 3.3 Empirical Fisher Information

To compute the diagonal FIM in Eq. 14, the expected value needs to be calculated with respect to the parametric probabilistic model $P({\textnormal{x}}|{\bm{\theta}})$. While data can be sampled from the parametric model, it is not always a straightforward process. Various sampling methods exist, such as Gibbs sampling, Langevin Markov chain Monte Carlo (MCMC) sampling [], Metropolis-Hastings MC sampling [] and the recent GFlowNet []. However, none of these methods are universally efficient in generating sufficient data with high fidelity and diversity.

Instead of $P({\textnormal{x}}|{\bm{\theta}})$, we utilize the true data-generating distribution $p_{data}({\textnormal{x}})$ to compute FIM. Although the exact form of $p_{data}({\textnormal{x}})$ is unknown, we have access to a training set of samples.
We compute FIM by utilizing the training set ${\mathcal{D}}$ by substituting the true distribution $p_{data}({\textnormal{x}})$ with the empirical distribution $\hat{p}_{data}({\textnormal{x}})$. This approximated FIM is referred to as the empirical FIM in the statistics community []. The training set might not contain enough samples of low-probability x, which could lead to issues with the empirical FIM.

$$ $\displaystyle\displaystyle\hat{{\bm{f}}}({\bm{\theta}})$ $\displaystyle=\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}[\nabla_{{\bm{ \theta}}}\log P({\textnormal{x}}|{\bm{\theta}})^{2}]\approx\mathbb{E}_{{ \textnormal{x}}\sim p_{data}}[\nabla_{{\bm{\theta}}}\log P({\textnormal{x}}|{ \bm{\theta}})^{2}],$ (16) $\displaystyle\approx\mathbb{E}_{{\textnormal{x}}\sim\hat{p}_{data}}[\nabla_{{ \bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})^{2}]=\frac{1}{|{\mathcal{D }}|}\sum_{{\textnormal{x}}\in{\mathcal{D}}}\nabla_{{\bm{\theta}}}\log P({ \textnormal{x}}|{\bm{\theta}})^{2}$ (17) $$

The expected value of the loss function is also obtained from the empirical distribution rather than the true distribution. Optimizing this cost function $J({\bm{\theta}})$ is referred to as empirical risk minimization, for similar reasons.

$$ $\displaystyle\displaystyle J({\bm{\theta}})$ $\displaystyle=-\mathbb{E}_{{\textnormal{x}}\sim p_{data}}[\log P({\textnormal{ x}}|{\bm{\theta}})]\approx-\mathbb{E}_{{\textnormal{x}}\sim\hat{p}_{data}}[ \log P({\textnormal{x}}|{\bm{\theta}})]=-\frac{1}{|{\mathcal{D}}|}\sum_{{ \textnormal{x}}\in{\mathcal{D}}}\log P({\textnormal{x}}|{\bm{\theta}})$ (18) $\displaystyle\displaystyle\nabla_{\bm{\theta}}J({\bm{\theta}})$ $\displaystyle=\nabla_{\bm{\theta}}\mathbb{E}_{{\textnormal{x}}\in{\mathcal{D}} }[{\mathcal{L}}({\bm{\theta}})]=\frac{1}{|{\mathcal{D}}|}\sum_{{\textnormal{x} }\in{\mathcal{D}}}\nabla_{\bm{\theta}}{\mathcal{L}}({\bm{\theta}})=-\frac{1}{| {\mathcal{D}}|}\sum_{{\textnormal{x}}\in{\mathcal{D}}}\nabla_{\bm{\theta}}\log P ({\textnormal{x}}|{\bm{\theta}})$ (19) $$

Calculating the exact cost function and FIM over the entire training set is computationally expensive. Therefore, in practice, the expected value is approximated using a minibatch ${\mathcal{B}}$. As this approximation further increases the uncertainty in the empirical FIM, FIM is typically estimated using an exponential moving average (EMA). Therefore, during training, natural gradient is computed as follows:

$$ $\displaystyle\displaystyle\tilde{\nabla}{J({\bm{\theta}})}=\hat{{\bm{f}}}^{-1} \nabla{J({\bm{\theta}})}$ $\displaystyle\approx-\mathbb{E}_{{\textnormal{x}}\in{\mathcal{B}}}[\nabla_{\bm {\theta}}\log P({\textnormal{x}}|{\bm{\theta}})]/\mathbb{E}_{EMA}[\mathbb{E}_{ {\textnormal{x}}\in{\mathcal{B}}}[\nabla_{{\bm{\theta}}}\log P({\textnormal{x} }|{\bm{\theta}})^{2}]]$ (20) $\displaystyle\approx-\mathbb{E}_{{\textnormal{x}}\in{\mathcal{B}}}[\nabla_{\bm {\theta}}\log P({\textnormal{x}}|{\bm{\theta}})]/\mathbb{E}_{EMA}[\mathbb{E}_{ {\textnormal{x}}\in{\mathcal{B}}}[\nabla_{{\bm{\theta}}}\log P({\textnormal{x} }|{\bm{\theta}})]^{2}]$ (21) $\displaystyle=-{\bm{g}}/\mathbb{E}_{EMA}[{\bm{g}}^{2}]$ (22) $$

Let the gradient of a minibatch be denoted as ${\bm{g}}$. To reuse ${\bm{g}}$ for calculating FIM, Adam makes another approximation from Eq. 20 to Eq. 21. Wang u.a. [] showed that using Eq. 20 makes Adam less sensitive to batch size. Adam variants are recommended to use large batch sizes [, ] to accurately estimate not only the gradient but also FIM.

In supervised learning, the loss function becomes conditional log-likelihood, necessitating specific considerations detailed in Section B.2.

#### 3.3.1 Discrete vs Continuous probability distributions

It’s worth noting that while Adam often outperforms SGD in the text domain, several studies have reported that its convergence point in the image domain, particularly for generative tasks or when using CNNs, can be inferior to that of SGD [, , , , ].
Empirical evidence indicates that Adam excels when dealing with discrete distributions, such as text inputs with categorical distributions. However, it may encounter difficulties when handling continuous distributions, such as image inputs with Gaussian distributions.

In the image domain, using the L2 loss is common practice due to its equivalence to the negative log-likelihood of a Gaussian distribution under the gradient. $k$ represents the partition function of the Gaussian distribution, which is equal to $1/\sqrt{2\pi\sigma^{2}}$. Since it is independent of ${\bm{\theta}}$, it is eliminated when the gradient is taken.

$$ $\displaystyle\displaystyle\nabla_{\bm{\theta}}J({\bm{\theta}})$ $\displaystyle\approx-\mathbb{E}_{{\textnormal{x}}\sim\hat{p}_{data}}[\nabla_{ \bm{\theta}}\log p({\textnormal{x}}|{\bm{\theta}})]=-\mathbb{E}_{{\textnormal{ x}}\sim\hat{p}_{data}}\left[\nabla_{\bm{\theta}}\log k\exp-\left(\frac{{ \textnormal{x}}-{\bm{\mu}}({\bm{\theta}})}{\sqrt{2}\sigma}\right)^{2}\right]$ (23) $\displaystyle=\frac{1}{2\sigma^{2}}\mathbb{E}_{{\textnormal{x}}\sim\hat{p}_{ data}}\left[\nabla_{\bm{\theta}}({\textnormal{x}}-{\bm{\mu}}({\bm{\theta}}))^{ 2}\right]=\frac{1}{2\sigma^{2}|{\mathcal{D}}|}\sum_{{\textnormal{x}}\in{ \mathcal{D}}}\nabla_{\bm{\theta}}({\textnormal{x}}-{\bm{\mu}}({\bm{\theta}}))^ {2}$ (24) $$

We performed empirical approximations to calculate the expected value of FIM. Eq. 16 represents the empirical approximation of FIM for generative models, while Eq. 52 and Eq. 58 represent the empirical approximations for discriminative models.
We hypothesize that these empirical approximations cause significantly more problems in continuous distributions than in discrete distributions. This disparity arises from the fundamental difference in how expected values are calculated for discrete and continuous distributions. Discrete distributions rely on probability mass functions, where the softmax function often concentrates the majority of probability mass on a few top logits. This concentration allows for relatively accurate empirical approximations. (^6^66Label smoothing in classification tasks [] has been shown to enhance performance. A theoretical explanation is that it results in a more accurate estimation of the empirical FIM. This is because label smoothing leads to a better approximation of the empirical FIM by computing all possible labels (y), rather than relying on a single data sample, as shown in Eq. 57.)
(^7^77Knowledge distillation [] is employed to enhance the performance of Large Language Models [, , ]. This technique leverages the teacher model’s probability distribution to calculate the expected value of the student model’s log probability, enabling Adam to estimate FIM more accurately. While knowledge distillation is recognized for its ability to transfer dark knowledge from teacher to student, its effectiveness may also be attributed to the improved optimization facilitated by the precise FIM estimation within Adam.)
In contrast, continuous distributions require integration over their probability density functions, making the estimation of their expected values with a single sampled value an overly simplified and potentially inaccurate approximation.

Kunstner u.a. [] argue against the use of the empirical Fisher in natural gradient descent, providing examples solely with continuous inputs and distributions.(^8^88Regrettably, relying on image classification (i.e., continuous inputs) or, even more restrictively, simple regression like curve fitting (i.e., continuous distribution) to analyze Adam or FIM may limit the generalizability of findings [, , ].) Their findings support our assertion that Adam may not perform well with continuous distributions, while the empirical success of Adam suggests that the empirical Fisher is adequate for discrete distributions.

We propose utilizing log-likelihood loss based on discrete probability distributions. In the image domain, this translates to using cross-entropy (CE) loss on a categorical distribution instead of the L2 loss.(^9^99Diffusion models [] seem to perform well with L2 loss. It is likely because the expectation of the loss is taken over not only image samples but also the noisy diffusion steps, which cover the model distribution sufficiently, allowing for a more accurate empirical estimate of the FIM. This could be a significant factor contributing to the success of diffusion models.) This can be implemented by modifying predictions to utilize a one-hot encoding, predicting 256 values per RGB channel. As demonstrated in Section 4.3, this modification significantly enhances the FID (Fréchet Inception Distance) metric for VQ-VAE [].

It has also been reported that categorical CE loss improves accuracy in floating-point number regression tasks [, , , , ]. Farebrother u.a. [] reported that CE loss with a histogram-based discretization (HL-Gauss []) significantly improved the prediction of rewards and values in reinforcement learning.

The exceptional scalability of Large Language Models (LLMs) [, ] with model size can likely be attributed to two factors: the inherently discrete nature of text data and the widespread use of the Adam optimizer. As LLMs continue to evolve into foundation models [] for various modalities, image, speech, and video domains are also adopting discrete token representations [, , ]. This provides further evidence that empirical FIM estimation necessitates the use of discrete distributions.

### 3.4 Fisher Adam

Based on the discussion in this section, we need to make minimal modifications to the Adam algorithm. We call this FAdam (Fisher Adam), and it is presented in Algorithm 1.(^10^1010The parts that differ from Adam are highlighted in blue.)

Figure: Algorithm 1 Fisher Adam (FAdam)

#### 3.4.1 Reciprocal vs Reciprocal Square-root for FIM

We have discovered that the second momentum term in Adam [] closely resembles a natural gradient with a diagonal empirical FIM, as demonstrated in Eq. 22 and Eq. 60. However, there is a key distinction: while natural gradients are divided by the FIM, Adam divides gradients by the square root of the FIM, as shown in Algorithm 3.

To precisely update ${\bm{\theta}}_{t}$ at each training step as per Eq. 13, we ideally need an accurate estimation of FIM. However, the empirical FIM computed by the minibatch data ${\mathcal{B}}$ is noisy.
Relying on a diagonal empirical FIM can result in zero components, which causes the natural gradient to diverge due to division by zero. To address this and obtain a more stable diagonal empirical FIM, an exponential moving average (EMA) is employed.

The gradient and FIM in Eq. 22 vary at each point ${\bm{\theta}}$. As shown in Section 2.1, not only their components change but also their underlying basis.
However, EMA averages the components of the Fisher information tensor obtained across different ${\bm{\theta}}$ points, despite their potentially differing basis. This is a mathematically invalid operation, except for in Euclidean space.(^11^1111Properly adding tensors from different tangent spaces requires parallel transport, which relies on understanding how the manifold’s basis changes. This is determined by the Riemannian metric tensor field, but accurately knowing even the FIM at a single point is computationally expensive, making knowledge of the entire FIM field intractable.)
As $\beta_{2}$=0.999 is standard value in Adam, EMA has a half-life of 700 steps (^12^1212$\sim\log(0.5)/\log(0.999)$) (^13^131313.5 steps with $\beta_{2}$=0.95), meaning it averages FIM over roughly 1000 steps. During this time, the basis can shift significantly as the model parameter moves to a new ${\bm{\theta}}$ location on the manifold.

In differential geometry, scalars, vectors and tensors are invariant under a coordinate change. The length of a vector, derived from the inner product (Eq. 1), is a scalar.
For a given gradient, we can express its length as shown in Eq. 10, and simplify it using the diagonal FIM in Eq. 26.

We only have access to the components of the gradient vector, and the invariant quantity is the vector length. However, EMA attempts to average the vector components as if they were in Euclidean space, ignoring the changing basis. To address this, we propose constructing a gradient vector (Eq. 27) in Euclidean space that has the same length as the known vector length (Eq. 26), and then apply EMA to this artificial Euclidean vector.

We refer to it as the invariant natural gradient, denoted as $\bar{\nabla}{J({\bm{\theta}})}$. This is the reason why Adam uses the square root, as shown in Algorithm 3. In Section C.1.1, we conduct an ablation study on the exponent of the FIM term, and find that the square root is the optimal choice.

$$ $\displaystyle\displaystyle\lVert\tilde{\nabla}{J({\bm{\theta}})}\rVert$ $\displaystyle=\tilde{\nabla}{J({\bm{\theta}})}\cdot\tilde{\nabla}{J({\bm{ \theta}})}=({\bm{F}}^{-1}\nabla{J({\bm{\theta}})})^{\top}{\bm{F}}({\bm{F}}^{-1 }\nabla{J({\bm{\theta}})})=\nabla{J({\bm{\theta}})}^{\top}{\bm{F}}^{-1}\nabla{ J({\bm{\theta}})}$ (25) $\displaystyle\approx\nabla{J({\bm{\theta}})}^{\top}\hat{{\bm{f}}}^{-1}\nabla{J ({\bm{\theta}})}=\lVert\nabla{J({\bm{\theta}})/\sqrt{\hat{{\bm{f}}}}}\rVert^{2}$ (26) $\displaystyle\displaystyle\bar{\nabla}{J({\bm{\theta}})}$ $\displaystyle:=\nabla J({\bm{\theta}})/\sqrt{\hat{{\bm{f}}}}$ (27) $$

Adam variants like AdaDelta [], AdaMax [] and Yogi [] modify the second momentum to further deviate from the natural gradient. A more principled approach would focus on improving the estimation of the empirical FIM, incorporating off-diagonal elements, and investigating how to make using natural gradient, instead of invariant natural gradient. If we were able to achieve a more accurate FIM estimation with smaller $\beta_{2}$, it might be possible to eliminate the square root operation.

#### 3.4.2 Momentum

Momentum methods [] employ exponential moving average (EMA) of gradients along the optimization trajectory to mitigate the noise inherent in stochastic gradients estimated from minibatches ${\mathcal{B}}$.

At a point ${\bm{\theta}}$ on the statistical manifold, the ideal gradient is the invariant natural gradient, as shown in Eq. 27. Therefore, momentum should average the invariant natural gradient $\bar{\nabla}{J({\bm{\theta}})}$ rather than the raw gradient, as shown in Algorithm 1. In contrast, as shown in Algorithm 3, Adam [] calculates the first and second momentums separately and then combines them, lacking a clear theoretical justification. In Algorithm 4, Adafactor [], on the other hand, already aligns with our proposed approach. LaProp [] has empirically demonstrated the benefits of this modification. Furthermore, the second momentum is not momentum. We refer to it as FIM.

FAdam omits zero bias correction in EMA of momentum (^14^1414${\bm{m}}_{t}\leftarrow{\bm{m}}_{t}/(1-\beta_{1}^{t})$), because excessive gradient descent is unnecessary before the momentum is sufficiently accumulated. This approach can be viewed as built-in informed warmup schedule. Adafactor implementation [] already omits zero bias correction for momentum, although this is not explicitly mentioned in the paper [].

#### 3.4.3 Weight decay on manifold

AdamW [], which decouples the weight decay term from the loss and applies it directly to the Adam optimizer, has demonstrated general performance improvements and is now widely used as a standard practice. Our mathematical framework provides a sound theoretical explanation for this observed phenomenon. The loss function should be log-likelihood, but the weight decay term has nothing to do with the probability distribution. If the weight decay term is included in the loss, it causes problems in estimating the FIM. (^15^1515While the weight decay term can be seen as a form of prior in a Bayesian framework, it is often directly added to the loss function (equivalent to the log posterior). The AdamW [] highlights that this practice can negatively impact performance, suggesting that the commonly used L2 weight decay might not accurately represent the true prior of the model.)

Weight decay is a great example of how auxiliary losses should be handled. If the auxiliary loss is related to the log-likelihood, it can be included in the loss. Otherwise, it should be bypassed by the Adam optimizer, as in the case of weight decay.

Furthermore, following Eq. 11, weight decay should also be applied as a natural gradient. Since we already have the FIM, we can express the weight decay gradient in a similar way to Eq. 27.

$$ $\displaystyle\displaystyle{\bm{g}}_{w}={\bm{\theta}}/\sqrt{\hat{{\bm{f}}}}$ (28) $$

Intuitively, components with low Fisher information can be pushed closer to zero without significantly impacting the model performance. Elastic weight consolidation [] leverages a similar concept, utilizing Fisher information to regularize the change of ${\bm{\theta}}$.(^16^1616EWC can directly utilize FIM from FAdam, eliminating the need to compute it separately.)

It has been reported that for training large-scale models like LLMs, decoupling weight decay from the learning rate helps stabilize training []. However, since our modifications make the weight decay adaptable to the loss surface, such workarounds might not be necessary. Further research is needed to confirm this hypothesis.

#### 3.4.4 FAdam

Fisher Adam (FAdam), incorporating all the discussed modifications, is presented in Algorithm 1. Fisher Adafactor (FAdafactor) modification of Adafactor is presented in Algorithm 2 in Section B.4.

Additionally, we have specified that FIM (${\bm{f}}_{0}$) is initialized to 1, not 0, in Algorithm 1. This is because FIM represents a Riemannian metric, which defaults to the identity matrix in flat space. However, this change does not affect the logic of the algorithm because bias correction ignores the initial value. The bias correction for FIM is adopted from Adafactor [], which is agnostic to the initial value.

#### 3.4.5 Convergence analysis

Algorithm 1 does not deviate from the assumptions of the convergence analysis presented in Eq. (15) of the latest Adam convergence proof by  Défossez u.a. []. Therefore, FAdam achieves the following convergence guarantees. Specifically, as the number of iterations N becomes sufficiently large, the expected value of the gradient becomes sufficiently small. The detailed explanation is provided in Section B.6.

$$ $\displaystyle\displaystyle\mathbb{E}\left[\lVert\nabla F(x_{\tau})\rVert^{2} \right]\leq O(d\ln(N)/\sqrt{N})$ (29) $$

## 4 Experiment

Figure: (a) LLM 1B (text domain)
Refer to caption: x1.png

In our experiments, we took existing state-of-the-art models from various domains and simply replaced their optimizers with FAdam, keeping all other hyperparameters unchanged. Despite this, FAdam consistently outperformed the original Adam variants. Our results demonstrate that FAdam can be effectively used as a drop-in replacement for Adam in real-world applications.

### 4.1 LLM (text domain)

We pretrained the 1B parameter LLM model from PaLM [] on C4 dataset [].
FAdafactor and Adafactor shared the same hyperparameters, $\beta_{1}$=0.9, $\beta_{2}$=0.99, and $\lambda$=0.001. However, the epsilon value used was $\epsilon$=1e-8 for FAdafactor and $\epsilon$=1e-30 for Adafactor.
As demonstrated in Fig. 1(a), FAdafactor outperforms Adafactor.

### 4.2 ASR (speech domain)

The 600M parameter Conformer [] model from the w2v-BERT [] is one of the lowest WER (word error rate) achieving models on the LibriSpeech dataset []. This model was pretrained using w2v-BERT and then finetuned on LibriSpeech data using the RNNT loss []. The WER was further improved to SoTA levels using noisy student semi-supervised finetuning [] on LibriLight [].
We compared Adam and FAdam using models that were fine-tuned in a semi-supervised fashion on LibriLight data, paired with semi-supervised pseudo labels.

FAdam and Adam have identical hyperparameters as follows: $\beta_{1}$=0.9, $\beta_{2}$=0.98, $\epsilon$=1e-8, and $\lambda$=0.001.
As demonstrated in Fig. 1(b) and Table 1, FAdam not only outperforms Adam but also establishes a new state-of-the-art (SoTA) Word Error Rate (WER) on LibriSpeech for 600M parameter models.

**Table 1: LibriSpeech WERs**
| LibriSpeech WERs | dev | dev-other | test | test-other | avg |
| --- | --- | --- | --- | --- | --- |
| Adam (w2v-BERT paper []) | 1.30 | 2.60 | 1.40 | 2.70 | 2.00 |
| Adam | 1.30 | 2.54 | 1.33 | 2.59 | 1.93 |
| FAdam | 1.29 | 2.49 | 1.34 | 2.49 | 1.89 |

### 4.3 VQ-VAE (image domain)

We trained a 100M parameter ViT VQ-GAN model [] on the ImageNet dataset []. To verify the hypothesis from Section 3.3.1 that categorical cross-entropy (CE) loss is superior to L2 loss on Adam, we exclusively used either CE loss or L2 loss + logit-laplace during training.(^17^1717Since our primary focus is on replacing the L2 loss, we omitted GAN and perceptual losses, making the model effectively a VQ-VAE [] rather than a VQ-GAN.) The VQ-GAN paper [] introduced logit-laplace to adjust the output scale when using L2 loss.

FAdam and AdamW [] have identical hyperparameters as follows: $\beta_{1}$=0.9, $\beta_{2}$=0.99, $\epsilon$=1e-8, and $\lambda$=1e-4.
As demonstrated in Fig. 1(c), FAdam outperforms AdamW, and categorical CE loss not only outperforms L2 loss but also eliminates the need for the complex logit-laplace transformation.

### 4.4 Ablation study

We conducted ablation studies on each hyperparameter of FAdam. Due to page limitations in the main paper, the results are presented in Section C.1.

## 5 Conclusion

In this work, we have revealed the mathematical foundation of the Adam optimizer, clarifying the constraints on loss function selection and the approximations involved in its derivation. This groundwork opens avenues for future research to develop improved optimizers by mitigating these approximations. Building upon our theoretical foundation, we have proposed an enhanced algorithm, FAdam, and demonstrated its effectiveness across diverse domains.

## Acknowledgement

The Riemannian geometry used in this paper was learned from the tensor calculus videos on the YouTube channel eigenchris. The intuition that Adam is related to information geometry was sparked by lectures on the YouTube channel Enjoying Math. Less P Wright Jr. at Meta AI created the PyTorch implementation(^18^1818[https://github.com/lessw2020/FAdam_PyTorch](https://github.com/lessw2020/FAdam_PyTorch)) and conducting experiments demonstrating that it outperforms AdamW on LLMs. Anirbit Mukherjee at the University of Manchester provided valuable feedback on the convergence proof. Zhouyuan Huo at Google provided insightful review of the code and paper.

## Appendix A Background

### A.1 Tensor Calculus Notation

In this paper, we will derive the Adam algorithm from the perspective of information geometry. Since statistical formulas are frequently used in machine learning papers, but Riemannian geometry formulas are not, we will first clarify the notation before proceeding. We follow the notation of  Petersen []’s Riemannian geometry textbook.

We define the tangent space at a point ${\bm{\theta}}$ on a manifold as ${{\mathcal{T}}_{{\bm{\theta}}}{\mathcal{M}}}$. Given a vector $\vec{v}=v^{i}{\bm{e}}_{i}$ on ${{\mathcal{T}}_{{\bm{\theta}}}{\mathcal{M}}}$, the basis vectors are defined as partial derivatives as follows:

$$ $\displaystyle{\bm{e}}_{i}:=\frac{\partial}{\partial\theta^{i}}$ (30) $$

Therefore, a vector $\vec{v}$ can be expressed using Einstein notation as follows, with $v^{i}$ representing the vector components and $\frac{\partial{}}{\partial{\theta^{i}}}$ representing the basis vectors:

$$ $\displaystyle\vec{v}={\bm{v}}=v^{i}\frac{\partial{}}{\partial{\theta^{i}}}$ (31) $$

From an ML practitioner’s perspective, the reason for understanding vector is that the parameters ${\bm{\theta}}$ of probability density (or mass) functions $p({\bm{\theta}})$ (i.e., scalar fields) reside within a vector space.

Given a covector ${\bm{w}}=w_{i}{\bm{e}}^{i}$ in the dual space of a vector space, the covector basis is defined using a differential one-form.

$$ $\displaystyle{\bm{e}}^{i}:=d\theta^{i}$ (32) $$

From an ML practitioner’s perspective, the reason for understanding covector is that the differential form of the loss, $d{\mathcal{L}}({\bm{\theta}})$, reside within a covector space.

$$ $\displaystyle d{\mathcal{L}}({\bm{\theta}})=\frac{\partial{\mathcal{L}}({\bm{ \theta}})}{\partial\theta^{i}}d\theta^{i}$ (33) $$

The operation between the vector basis and the covector basis results in the Kronecker delta. The operation between vector $\vec{v}$ and covector $d{\mathcal{L}}$ is as shown in Eq. 35. The meaning of Eq. 35 is the directional derivative, representing the derivative of the loss ${\mathcal{L}}$ in the direction of $\vec{v}$.

$$ $\displaystyle\displaystyle d\theta^{i}(\frac{\partial{}}{\partial{\theta^{j}}})$ $\displaystyle=\delta^{i}_{j}$ (34) $\displaystyle\displaystyle d{\mathcal{L}}(\vec{v})=\frac{\partial{\mathcal{L}} ({\bm{\theta}})}{\partial\theta^{i}}v^{j}$ $\displaystyle d\theta^{i}(\frac{\partial{}}{\partial{\theta^{j}}})=\frac{ \partial{\mathcal{L}}({\bm{\theta}})}{\partial\theta^{i}}v^{i}$ (35) $$

The covariant derivative is an extension of the directional derivative operation from a scalar filed to vector and tensor fields. (^19^1919When applied to a scalar field (i.e., a scalar function), the covariant derivative is called the directional derivative.) The covariant derivative of a vector field $\vec{w}$ along a vector $\vec{v}$ is denoted as follows, where $\Gamma^{k}_{ij}$ represents the Christoffel symbols:

$$ $\displaystyle D_{\vec{v}}\vec{w}=v^{i}\frac{\partial{}}{\partial{\theta^{i}}}( w^{j}\frac{\partial{}}{\partial{\theta^{j}}})=v^{i}\frac{\partial{}}{\partial{ \theta^{i}}}{w^{j}}\frac{\partial{}}{\partial{\theta^{j}}}+v^{i}w^{j}\frac{ \partial{}}{\partial{\theta^{i}}}\frac{\partial{}}{\partial{\theta^{j}}}=(v^{i }\frac{\partial{}}{\partial{\theta^{i}}}{w^{k}}+v^{i}w^{j}\Gamma^{k}_{ij}) \frac{\partial{}}{\partial{\theta^{k}}}$ (36) $$

### A.2 The proof of Fisher is negative Hessian

The relationship between Fisher information and the negative Hessian’s expectation can be proven as follows:

$$ $\displaystyle\displaystyle\nabla_{{\bm{\theta}}}^{2}\log P({\textnormal{x}}|\theta)$ $\displaystyle=\frac{\nabla_{{\bm{\theta}}}^{2}P({\textnormal{x}}|\theta)}{P({ \textnormal{x}}|\theta)}-\nabla_{{\bm{\theta}}}\log P({\textnormal{x}}|{\bm{ \theta}})\>\nabla_{{\bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})^{\top}$ (37) $\displaystyle\displaystyle\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}} \left[\frac{\nabla_{{\bm{\theta}}}^{2}P({\textnormal{x}}|\theta)}{P({ \textnormal{x}}|\theta)}\right]$ $\displaystyle=\int{P({\textnormal{x}}|\theta)\frac{\nabla_{{\bm{\theta}}}^{2}P ({\textnormal{x}}|\theta)}{P({\textnormal{x}}|\theta)}}d{\textnormal{x}}= \nabla_{{\bm{\theta}}}^{2}\int{P({\textnormal{x}}|\theta)}d{\textnormal{x}}= \nabla_{{\bm{\theta}}}^{2}1=0$ (38) $\displaystyle\displaystyle\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}} \left[\nabla_{{\bm{\theta}}}^{2}\log P({\textnormal{x}}|\theta)\right]$ $\displaystyle=-\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}\left[\nabla_{ {\bm{\theta}}}\log P({\textnormal{x}}|{\bm{\theta}})\>\nabla_{{\bm{\theta}}} \log P({\textnormal{x}}|{\bm{\theta}})^{\top}\right]$ (39) $$

### A.3 The proof of KL approximation

As seen in Eq. 9, the Kullback-Leibler (KL) divergence with an infinitesimal displacement ${\bm{d}}$ can be approximated by the Fisher information. We provide a proof of this relationship below, following the approach presented in Kristiadi [].

$$ $\displaystyle D_{\mathrm{KL}}(P({\textnormal{x}}|{\bm{\theta}})\|P({ \textnormal{x}}|{\bm{\theta}}+{\bm{d}}))\approx\frac{1}{2}{\bm{d}}^{\top}{\bm{ F}}({\bm{\theta}}){\bm{d}}.$ (40) $$

By using a second-order Taylor series approximation, the KL divergence can be approximated as follows:

$$ $\displaystyle D_{\mathrm{KL}}(P_{\bm{\theta}}\|P_{{\bm{\theta}}+{\bm{d}}}) \approx D_{\mathrm{KL}}(P_{\bm{\theta}}\|P_{{\bm{\theta}}})+(\nabla_{{\bm{ \theta}}^{\prime}}D_{\mathrm{KL}}(P_{\bm{\theta}}\|P_{{\bm{\theta}}^{\prime}}) |_{{\bm{\theta}}={\bm{\theta}}^{\prime}})^{\top}{\bm{d}}+\frac{1}{2}{\bm{d}}^{ \top}\nabla^{2}_{{\bm{\theta}}^{\prime}}D_{\mathrm{KL}}(P_{\bm{\theta}}\|P_{{ \bm{\theta}}^{\prime}}){\bm{d}}$ (41) $$

$D_{\mathrm{KL}}(P_{\bm{\theta}}\|P_{{\bm{\theta}}})$ is 0 by the definition of KL divergence. The first-order approximation term also becomes 0 through the following process.

$$ $\displaystyle\displaystyle\nabla_{{\bm{\theta}}^{\prime}}D_{\mathrm{KL}}(P_{ \bm{\theta}}\|P_{{\bm{\theta}}^{\prime}})$ $\displaystyle=\cancel{\nabla_{{\bm{\theta}}^{\prime}}\mathbb{E}_{{\textnormal{ x}}\sim P_{\bm{\theta}}}\left[\log P_{\bm{\theta}}\right]}-\nabla_{{\bm{\theta }}^{\prime}}\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}\left[\log P_{{ \bm{\theta}}^{\prime}}\right]=-\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta} }}\left[\nabla_{{\bm{\theta}}^{\prime}}\log P_{{\bm{\theta}}^{\prime}}\right]$ (42) $\displaystyle=-\int P_{\bm{\theta}}\nabla_{{\bm{\theta}}^{\prime}}\log P_{{\bm {\theta}}^{\prime}}|_{{\bm{\theta}}={\bm{\theta}}^{\prime}}d{\textnormal{x}}=- \int P_{\bm{\theta}}\frac{\nabla_{{\bm{\theta}}^{\prime}}P_{{\bm{\theta}}^{ \prime}}}{P_{{\bm{\theta}}^{\prime}}}|_{{\bm{\theta}}={\bm{\theta}}^{\prime}}d {\textnormal{x}}$ (43) $\displaystyle=-\int\nabla_{{\bm{\theta}}}P_{{\bm{\theta}}}d{\textnormal{x}}=- \nabla_{{\bm{\theta}}}\int P_{{\bm{\theta}}}d{\textnormal{x}}=-\nabla_{{\bm{ \theta}}}1=0$ (44) $$

The Fisher information emerges from the second-order approximation term, as shown by utilizing the Hessian property in Eq. 8.

$$ $\displaystyle\displaystyle\nabla^{2}_{{\bm{\theta}}^{\prime}}D_{\mathrm{KL}}(P _{\bm{\theta}}\|P_{{\bm{\theta}}^{\prime}})$ $\displaystyle=\cancel{\nabla_{{\bm{\theta}}^{\prime}}^{2}\mathbb{E}_{{ \textnormal{x}}\sim P_{\bm{\theta}}}\left[\log P_{\bm{\theta}}\right]}-\nabla_ {{\bm{\theta}}^{\prime}}^{2}\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}} \left[\log P_{{\bm{\theta}}^{\prime}}\right]$ (45) $\displaystyle=-\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}\left[\nabla_{ {\bm{\theta}}^{\prime}}^{2}\log P_{{\bm{\theta}}^{\prime}}|_{{\bm{\theta}}={ \bm{\theta}}^{\prime}}\right]=-\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta} }}\left[\nabla_{{\bm{\theta}}}^{2}\log P_{{\bm{\theta}}}\right]$ (46) $\displaystyle={\bm{F}}({\bm{\theta}})$ (47) $$

Therefore, equation Eq. 9 is proven.

## Appendix B Toward FAdam

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

As shown in Eq. 8, FIM is equivalent to the Hessian of the loss.

$$ $\displaystyle{\bm{F}}({\bm{\theta}})=-\mathbb{E}_{{\textnormal{x}}\sim P_{\bm{ \theta}}}[{\bm{H}}_{{\bm{\theta}}}(\log P({\textnormal{x}}|{\bm{\theta}}))]= \mathbb{E}_{{\textnormal{x}}\sim P_{\bm{\theta}}}[{\bm{H}}_{{\bm{\theta}}}({ \mathcal{L}}({\bm{\theta}}))].$ (48) $$

Therefore, ${\bm{\theta}}$ update Eq. 13 incorporates the Hessian term.

$$ $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta\mathbb{E}_{{ \textnormal{x}}\sim P_{\bm{\theta}}}[{\bm{H}}_{{\bm{\theta}}}({\mathcal{L}}({ \bm{\theta}}))]^{-1}\nabla{{\mathcal{L}}({\bm{\theta}})}$ (49) $$

As FIM is the expected value of the Hessian, natural gradient optimization is considered a second order method. Since the Fisher Information Matrix (FIM) is positive semi-definite, the Hessian term is also positive semi-definite, ensuring a convex optimization.

Meanwhile, Newton’s method is a second-order optimization method that is effective only when the loss function is strongly convex with a Lipschitz continuous Hessian. The update equation for Newton’s method [] is as follows.

$$ $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta{\bm{H}}_{{\bm{\theta}} }({\mathcal{L}}({\bm{\theta}}))^{-1}\nabla{{\mathcal{L}}({\bm{\theta}})}$ (50) $$

Although the two methods were derived through vastly different processes, they both involve the inverse of the Hessian matrix, as shown in Eq. 49 and Eq. 50. However, there is a significant difference between the two methods. Natural gradient descent requires the loss function to be the log-likelihood, while Newton’s method has no such restriction on the choice of loss function. Natural gradient descent guarantees convergence due to the positive semi-definiteness of FIM (the expected value of the Hessian), whereas Newton’s method does not.

Consequently, various complex techniques are employed when using Newton’s method, such as the Gauss-Newton algorithm, conjugate gradient method, and trust region method []. These techniques ensure that each step update is confined within a trust region, mitigating the risk of divergence.

This is why second-order optimization methods based on FIM have demonstrated faster convergence compared to Newton’s method in practice [, , ]. Recent second-order optimization methods often utilize FIM instead of the Hessian of the loss function [, ]. However, these methods require the loss function to be the log-likelihood and are subject to certain constraints discussed in Section 3.2 and Section 3.3.

### B.2 Conditional probability distribution for empirical Fisher Information

In supervised learning, the loss function typically becomes a joint distribution $-\log P({\textnormal{x}},{\textnormal{y}}|{\bm{\theta}})$.
Since the expected value over $P({\textnormal{x}},{\bm{\theta}})$ is generally intractable, it is replaced with the empirical distribution $\hat{p}_{data}({\textnormal{x}})$ in Eq. 52. As noted in Eq. 16, this is referred to as the empirical Fisher.
The training set also provides the ground truth label y. Therefore, the loss function becomes the cross-entropy calculated using the conditional log-likelihood, as shown in Eq. 53.

$$ $\displaystyle\displaystyle\nabla_{\bm{\theta}}J({\bm{\theta}})$ $\displaystyle=-\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim p_{data}}[ \nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{\theta}})+ \nabla_{\bm{\theta}}\log P({\textnormal{x}}|{\bm{\theta}})]$ (51) $\displaystyle\approx-\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim\hat{p}_ {data}}[\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{ \theta}})+\cancel{\nabla_{\bm{\theta}}\log\hat{p}_{data}({\textnormal{x}})}]$ (52) $\displaystyle=-\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim\hat{p}_{data} }[\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{\theta}})] :={\bm{g}}$ (53) $$

The FIM, similar to the cost function, also requires the calculation of an expected value over the joint distribution. By approximating the expectation with the empirical distribution $\hat{p}_{data}({\textnormal{x}},{\textnormal{y}})$ and moving out the square as shown in Eq. 21, we obtain the following equation:

$$ $\displaystyle\displaystyle\hat{{\bm{f}}}({\bm{\theta}})$ $\displaystyle=\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim P({\textnormal {y}}|{\textnormal{x}},{\bm{\theta}})P({\textnormal{x}}|{\bm{\theta}})}[(\nabla _{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{\theta}})+\nabla_{ \bm{\theta}}\log P({\textnormal{x}}|{\bm{\theta}}))^{2}]$ (54) $\displaystyle\approx\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim P({ \textnormal{y}}|{\textnormal{x}},{\bm{\theta}})\hat{p}_{data}({\textnormal{x}} )}[(\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{\theta}} )+\cancel{\nabla_{\bm{\theta}}\log\hat{p}_{data}({\textnormal{x}})})^{2}]$ (55) $\displaystyle\approx\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim P({ \textnormal{y}}|{\textnormal{x}},{\bm{\theta}})\hat{p}_{data}({\textnormal{x}} )}[\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{\theta}}) ^{2}]$ (56) $$

As shown in equation Eq. 22, generative models are able to reuse the gradient ${\bm{g}}$ for calculating FIM. However, this reuse poses a challenge when dealing with conditional distributions.
This is because Equation Eq. 53 calculates the expected value over the label y, while Equation Eq. 56 calculates the expected value over the conditional distribution of the model.
To address this, we introduce an additional approximation by calculating the expected value over the label y in Eq. 57. This is commonly referred to as the empirical FIM in the ML community [], while the statistics community refers to the approximation in Eq. 55 as the empirical FIM, as explained in Eq. 16.
To avoid confusion, we will refer to this as the conditional empirical FIM.

$$ $\displaystyle\displaystyle\hat{{\bm{f}}}({\bm{\theta}})$ $\displaystyle\approx\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim\hat{p}_{ data}}[\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{ \theta}})^{2}]$ (57) $\displaystyle\approx\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\sim\hat{p}_{ data}}[\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{ \theta}})]^{2}={\bm{g}}^{2}$ (58) $$

As seen in Eq. 21, to reuse ${\bm{g}}$ in Eq. 53, a non-principled approximation is made by taking the square outside the expectation in Eq. 58.
Similar to the generative model in Eq. 22, we can obtain the natural gradient with respect to a minibatch ${\mathcal{B}}$ as follows:

$$ $\displaystyle\displaystyle\tilde{\nabla}{J({\bm{\theta}})}=\hat{{\bm{f}}}^{-1} \nabla{J({\bm{\theta}})}$ $\displaystyle\approx-\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\in{\mathcal {B}}}[\nabla_{\bm{\theta}}\log P({\textnormal{y}}|{\textnormal{x}},{\bm{\theta }})]/\mathbb{E}_{EMA}[\mathbb{E}_{{\textnormal{x}},{\textnormal{y}}\in{ \mathcal{B}}}[\nabla_{{\bm{\theta}}}\log P({\textnormal{y}}|{\textnormal{x}},{ \bm{\theta}})]^{2}]$ (59) $\displaystyle=-{\bm{g}}/\mathbb{E}_{EMA}[{\bm{g}}^{2}]$ (60) $$

Despite the numerous non-trivial approximations made to FIM, it is remarkable that Adam has achieved such great success. Eliminating these non-trivial approximations could be a promising direction for future research.

### B.3 Clipping and Epsilon

Although not mentioned in the original Adam paper, it was later discovered that gradient clipping is essential for Adam [, ] and is now used as a standard practice. Adafactor [] incorporates clipping [] in the algorithm, because the diagonal empirical FIM can have zero components.
Therefore, clipping is applied to the invariant natural gradient, and then the clipped gradient is used to calculate the momentum, as demonstrated in Algorithm 1. The Adafactor implementation [] already incorporates this approach, although it is not explicitly mentioned in the original paper [], as shown in Algorithm 4.
Since clipping is already incorporated into the optimizer, global clipping [] is not necessary.(^20^2020This feature is often included in ML frameworks as a legacy practice without a clear justification.) In our experiments, we did not observe any benefit from applying global clipping.

To prevent division by zero before clipping, $\epsilon$ is added.
The standard value for $\epsilon$ in Adam is 1e-8. Adafactor [] uses 1e-30, but since it is added inside the square root, this translates to 1e-15 on the Adam scale. Wortsman u.a. [] recommends using a smaller epsilon value of 1e-15 for large-scale models like LLMs due to the empirical observation that the root mean square (RMS) value of gradients tends to decrease as models get larger and training progresses.

To mitigate the need for manual tuning of $\epsilon$ across different models, we propose an adaptive mechanism that adjusts $\epsilon$ based on the gradient RMS values: $\hat{\epsilon}\leftarrow\min(10^{-8},0.01\text{RMS}({\bm{g}}_{t}))$, as shown in Algorithm 1.
Our experiments in Section C.1.2 demonstrate that this adaptive approach performs comparably to manual tuning, where optimal epsilon values vary across model sizes and domains.(^21^2121We attempted to adjust $\epsilon$ based on the weight size using the formula $\hat{\epsilon}=\epsilon/\text{size(}{\bm{\theta}})$. However, this approach did not improve performance, likely due to the resulting epsilon values for bias and layer normalization weights becoming unsuitable.)

Moreover, in Algorithm 1, epsilon is raised to the power of $2\rho$ to ensure invariance with respect to the FIM exponent (in Section 3.4.1). To keep the epsilon values in FAdam comparable to those used in Adam, the exponent is multiplied by 2. (^22^2222Since the default $\rho$ is 0.5, $\epsilon^{2\rho}$ is typically equal to $\epsilon$.)
EAdam [] and Adafactor [] modified the algorithm to accumulate $\epsilon$ within EMA, but we did not observe any benefits from this approach in our experiments.

### B.4 FAdafactor

Adafactor [] is modified as shown in Algorithm 2, and we refer to this variant as Fisher Adafactor (FAdafactor).

Figure: Algorithm 2 Fisher Adafactor (FAdafactor)

### B.5 Adam and Adafactor

To facilitate comparison between the original algorithms, we present Adam (Algorithm 3) and Adafactor (Algorithm 4), alongside FAdam (Algorithm 1) and FAdafactor (Algorithm 2). Since the full Adafactor algorithm is not explicitly detailed in the original paper [], we refer the original implementation [].

Remarkably, Adafactor [] had already empirically discovered and implemented several of the results we derived from our mathematical framework: applying natural gradient to momentum (Section 3.4.2), omitting bias correction for momentum (Section 3.4.2), and clipping gradients (Section B.3). This is akin to the invention of the steam engine before the establishment of thermodynamics, and these empirical findings significantly bolstered our confidence in developing our theory. However, it is peculiar that the momentum calculation incorporates the learning rate, while weight decay remains independent of the learning rate. Accumulating $\epsilon$ through EMA is also an unconventional approach.

Figure: Algorithm 3 Adam []

Figure: Algorithm 4 Adafactor [, ]

### B.6 Convergence proof

While the original Adam paper [] presented a convergence proof, Reddi u.a. [] later demonstrated that this proof was flawed. Subsequently, Défossez u.a. [] provided a corrected convergence proof for Adam. Our analysis follows the approach of Défossez u.a. [].

To the best of our knowledge, there is no existing convergence proof for Adam with clipping, despite its widespread use in practice. This is likely due to the non-trivial nature of incorporating clipping into the analysis. Existing convergence proofs only cover SGD with clipping []. Therefore, we also present a convergence proof that excludes the effects of clipping.

The convergence proof for FAdam is simpler than that of Adam. Adam’s convergence proof needs to consider both the 1st and 2nd momentums simultaneously, as they both influence the gradient update. In contrast, FAdam first computes the natural gradient using FIM and then applies momentum to the natural gradient. This process of applying momentum to the natural gradient is analogous to Polyak momentum applied to SGD. It is well-known that Polyak momentum tightens the convergence bound []. Since FAdam’s momentum is analogous to Polyak momentum, FAdam’s momentum also tightens the convergence bound. Therefore, the convergence bound for the natural gradient without momentum is looser than the convergence bound for FAdam.
For the sake of simplicity, we focus on deriving the convergence bound for the natural gradient without momentum. A tighter theoretical convergence bound, incorporating the effects of momentum, is left for future work.

FAdam without the effects of momentum is equivalent to the Adam algorithm with $\beta_{1}=0$. The convergence of Adam under this condition is shown in Eq. (15) of Défossez u.a. [] as follows:

$$ $\displaystyle\displaystyle\displaystyle\mathbb{E}\left[\lVert\nabla F(x_{\tau} )\rVert^{2}\right]$ $\displaystyle\leq\frac{F(x_{0})-F_{*}}{\alpha_{1}\sqrt{N}}+\frac{1}{\sqrt{N}}( 4dR^{2}+\alpha_{1}dRL)\left(\ln\left(1+\frac{RN}{\epsilon}\right)+\frac{N}{N-1 }\right)$ (61) $\displaystyle\sim O(d\ln(N)/\sqrt{N})$ (62) $$

In Eq. 61, $x\in\mathbb{R}^{d}$ represents the model parameters ${\bm{\theta}}$, $\alpha_{1}$ is the learning rate, $N$ represents the number of iterations, and $\tau_{N}$ denotes a random index within the set $\{0,\dots,N-1\}$. To derive the equation, the following assumptions are made:

The function $F$ is bounded below by $F_{*}$, that is:

$$ $F(x)\geq F_{*}\text{ for all }x\in\mathbb{R}^{d}.$ $$

The $l_{\infty}$ norm of the stochastic gradients is uniformly almost surely bounded, meaning there exists $R\geq\sqrt{\epsilon}$ such that:

$$ $\lVert\nabla F(x)\rVert_{\infty}\leq R-\sqrt{\epsilon}\text{ for all }x\in \mathbb{R}^{d}.$ $$

The objective function is smooth, specifically, its gradient is L-Lipschitz continuous with respect to the $l_{2}$-norm:

$$ $\lVert\nabla F(x)-\nabla F(y)\rVert_{2}\leq L\lVert x-y\rVert_{2}\text{ for all }x,y\in\mathbb{R}^{d}.$ $$

Therefore, from Eq. 62, as the number of iterations $N$ increases, the expected value of the squared norm of the gradient diminishes, indicating that $F$ converges to the optimal value $F_{*}$.

## Appendix C Experiment

### C.1 Ablation study

#### C.1.1 Exponent of FIM

In Section 3.4.1, we provided a theoretical explanation for the use of the square root in FIM, as shown in Eq. 27. To explore alternative exponent values, we conducted experiments, the results of which are presented in Fig. 1(d) and Table 2.
Exponents above 0.5 exhibit relative stability, while values below 0.3 demonstrate a sharp decline in performance. The standard value of 0.5 for the FIM exponent in Adam [] appears to be a well-justified choice.

**Table 2: Effects of Varying FIM Exponent on LLM and ASR Model Performance**
| Model | Metric | Steps | $\rho$=0.3 | $\rho$=0.4 | $\rho$=0.5 | $\rho$=0.6 | $\rho$=0.7 | $\rho$=0.8 | $\rho$=0.9 | $\rho$=1.0 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLM 1B | loss | 1.5k | 5.4 | 3.9 | 3.8 | 3.9 | 4.2 | 4.3 | 4.3 | 4.3 |
| ASR 100M | WER | 8k | 43.9 | 11.2 | 6.04 | 6.08 | 6.04 | 6.08 | 6.23 | 6.33 |

#### C.1.2 Epsilon

We conducted experiments to evaluate the effectiveness of an adaptive epsilon strategy (as mentioned in Section B.3) and to investigate the impact of different (non-adaptive) epsilon values on the performance of LLM (1B), ASR (600M), and VQ-VAE (100M) models. The results are presented in Table 3.

While the optimal epsilon value varies across different size and domain models, the adaptive epsilon approach demonstrated competitive performance without requiring manual tuning, highlighting its potential for broader applicability. Notably, manual tuning sometimes led to divergence issues with smaller epsilon values like 1e-20 or 1e-30.

**Table 3: Impact of Epsilon on Model Performance (Lower is Better)**
| Model | Metric | Steps | Adaptive | 1e-8 | 1e-12 | 1e-15 | 1e-20 | 1e-30 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLM 1B | loss | 80k | 2.39 | 2.47 | 2.42 | 2.46 | NaN | NaN |
| ASR 600M | WER | 8k | 2.04 | 2.08 | 2.07 | 2.08 | 2.05 | 2.13 |
| VQ-VAE 100M | FID | 50k | 19.5 | 20.9 | 19.9 | 19.6 | 19.7 | NaN |

#### C.1.3 Epsilon2

In Algorithm 1, when the root mean square (RMS) of the gradient is very small, $\hat{\epsilon}$ is calculated by multiplying the RMS(grad) by $\epsilon_{2}$. As shown in Table 3, a $\epsilon_{2}$ value of 0.01 consistently yielded good performance across all domains, with a $\epsilon$ value of 1e-8.

**Table 4: Impact of Epsilon2 on Model Performance (Lower is Better)**
| Model | Metric | Steps | 1e-1 | 1e-2 | 1e-3 | 1e-4 | 1e-6 | 1e-8 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLM 1B | loss | 60k | 2.63 | 2.41 | 2.48 | 2.45 | 2.60 | 2.41 |
| ASR 600M | WER | 40k | 1.98 | 1.89 | 1.92 | 1.93 | 1.90 | 1.91 |
| VQ-VAE 100M | FID | 50k | 19.9 | 19.5 | 19.6 | 19.5 | 19.7 | 19.8 |

#### C.1.4 Weight decay

As shown in Eq. 28, the most significant algorithmic change in FAdam compared to Adam is the weight decay mechanism. Therefore, we conducted experiments for weight decay.

In the ASR experiments described in Section 4.2, the baseline model utilized Adam, and we found that the weight decay parameter ($\lambda$) used in Adam could generally be reused for FAdam. Our enhanced weight decay mechanism played a significant role in improving FAdam’s performance, as shown in Table 5.

**Table 5: LibriSpeech WERs of 600M ASRs with Varying Weight Decay Parameter ($\lambda$)**
| LibriSpeech WERs | dev | dev-other | test | test-other | avg |
| --- | --- | --- | --- | --- | --- |
| Adam $\lambda$=0 | 1.31 | 2.55 | 1.34 | 2.57 | 1.94 |
| Adam $\lambda$=0.001 | 1.30 | 2.54 | 1.33 | 2.59 | 1.93 |
| FAdam $\lambda$=0 | 1.30 | 2.43 | 1.34 | 2.63 | 1.92 |
| FAdam $\lambda$=0.001 | 1.29 | 2.49 | 1.34 | 2.49 | 1.89 |

In the LLM experiments described in Section 4.1, the baseline model utilized Adafactor, which decouples weight decay from the learning rate, as shown in Algorithm 4. Interestingly, despite the fact that Adafactor and FAdafactor have quite different weight decay mechanisms, the optimal weight decay value of 1e-3 for Adafactor also is optimal for FAdafactor in Table 6.

**Table 6: Eval loss of 1B LLMs with Varying Weight Decay Parameter ($\lambda$)**
| $\lambda$ | 1e-2 | 1e-3 | 1e-4 | 1e-5 |
| --- | --- | --- | --- | --- |
| LLM 1B | 3.007 | 2.875 | 2.998 | 3.004 |

### C.2 Experiments Compute Resources

The LLM 1B model was trained on 16 TPUv5 [] devices (80GB HBM) for one day. Each training example consisted of 2k tokens with a global batch size of 256.
The ASR 600M model was fine-tuned on 32 TPUv3 devices (8GB HBM) for half a day. Each utterance had an average duration of 12 seconds with a global batch size of 512.
The VQ-VAE 100M model was trained on 16 TPUv4 devices (16GB HBM) for one day. Training was performed using images with a resolution of 256x256 and a global batch size of 256.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: See weaknesses section.
soundness: 3
strengths: 1. Understanding Adam is an important problem. Using fisher information and natural gradient descent to understand Adam is novel.

2. The presentation of this paper is good in general.
summary: This paper studies the connection between Adam optimizer and natural gradient descent by leveraging techniques from Riemannian geometry. Based on this, the authors propose a modified algorithm named Fisher Adam (FAdam). The convergence analysis of FAdam is provided and the algorithm is tested by large language model (LLM) experiments.
weaknesses: 1. The theoretical analysis of FAdam is weak. It directly follows the paper (Defossez et al. 2020) and requires strong assumptions (e.g., $\beta_1$=0, bounded gradient). So it does not analyze the algorithm's momentum and is worse than the state-of-the-art analysis of Adam in the literature. 

2. I do not find any rigorous presentation of Adam's flaws in the paper as claimed in the abstract by the authors. For example, the paper does not have any clear negative results of the vanilla Adam experimentally or theoretically. Therefore, the motivation of designing a new variant of Adam such as FAdam is unclear to me.

3. The description of the experiment is unclear. Lots of details are missing, such as training/test learning curve comparison, learning rate of the optimizer, batch size, and memory costs. Also, the experiment is only run once, and the algorithm's robustness is unclear.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - What is the exact argument for the "invariant natural gradient"?
- Kunstner et al. (2019) explicitly critique the interpretation of Adam as approximate NGD. What is your response to their arguments? (E.g., degeneracy of the empirical FIM for overparametrized models, sensitivity to model misspecification, no relationship between empirical and true FIM far from an optimum.)
- How were the $\epsilon$ values in the experiment chose?
soundness: 2
strengths: - The argument relating decoupled weight decay to the information-geometric interpretation is interesting. It clarifies that the gradients used to compute $v$ (the diagonal empirical FIM) must be gradients of the log likelihood of a probabilistic model to match the definition of the FIM and therefore must not contain regularizers or auxiliary losses.
- Averaging the preconditioned gradients (versus preconditioning the averaged gradient) is an interesting variant.
summary: The paper reiterates and expands the motivation of Adam as approximate natural gradient descent. It derives multiple modifications to the Adam algorithm based on that interpretation. The resulting method (FAdam) is evaluated experimentally.
weaknesses: - The paper presents as an original finding that it "establishes a mathematical foundation for the Adam optimizer" in terms of NGD with the empirical Fisher information matrix. This is misleading. This motivation has been given in the original Adam paper and has since been discussed and critiqued in various papers, including but not limited to Kunstner et al. (2019). This should be made transparent in the discussion of related work.
- The paper states that "for using natural gradient optimizers [...] the loss function must be in the form of the log-likelihood". This is not a factual statement and should be adjusted. Preconditioning with the Fisher information matrix adapts to the geometry induced by a certain probabilistic model. The negative log likelihood under said model may be a "natural" objective function to optimize, but NGD can meaningfully be applied to any other objective. In fact, in Section 3.4.3, the authors advocate for preconditioning an additional loss term with the FIM. 
- The argument in Section 3.4.1 regarding the use of the square-root on the preconditioner is not stringent. If I am understanding correctly, the argument is that $\Vert \nabla J/\sqrt{f} \Vert^2_2 \approx \Vert F^{-1} \nabla J\Vert_F^2$, i.e., preconditioning with the square-root makes the Euclidean norm of the resulting update equal the  "Fisher norm" of the natural gradient. However, there is no discernible argument why it would be desirable to match these two quantities and, if so, why one would want to achieve this by changing the preconditioner rather than, say, scale the update with a scalar factor? (Minor: The notation should also be improved in Eq. (25) - (27), since $\Vert\cdot\Vert$ is used to refer to both the Euclidean norm and the "Fisher norm".)
- The paper briefly cites Kunstner et al. (2019), which is an explicit critique of the interpretation of Adam as NGD, but does not really engage with the arguments in that paper.
- Overall, the paper combines various components, that are somewhat independent of each other:
    a) introduce gradient clipping,
    b) apply momentum after preconditioning,
    c) apply preconditioning to the weight decay gradient.
It would be highly desirable to perform ablation studies to understand which of these changes actually matter and how they interact.
- The quality of the empirical evaluation is a bit lacking. No error bars are given. The hyperparameter tuning protocol is somewhat unclear, e.g., FAdam uses a different epsilon value and it is not stated how this value was obtained.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - Can the authors give a citation regarding using an EMA for fisher info on line 228? I’m not aware that’s been used prior to Adam.
- How does the proposed framework justify clipping? In B.3 clipping  and epsilon is mentioned  through related work, but this step that has been added to the algorithm does not apppear to be justified by the theoretical framework.
- What norms are being used in equations (25)-(27)? I’m assuming the first norm is the one induced by the Fisher Information Matrix, but then what is the other one? Euclidian?
- Has the author tried to quantify how accurate the approximations (F)Adam is using are in a simple setting? This can help figure out if those approximations are in fact reasonable, which needs to be the case in order to claim it’s really natural gradient in disguise.
soundness: 1
strengths: The approach of  analyzing Adam from a statistical  viewpoint is interesting, and while not being new (this interpretation was mentioned in  the origional  Adam paper) it could  deserve a second  look. The authors additionally show some empiracle improvments in a few settings.
summary: This work takes a statistical viewpoint of the Adam optiumization algorithm and attempts to both explain it’s performance and add improvlemtents through the lense of the natural graident algorithm. The authors argue that Adam is effectivly preconditioning  with the diagonal of the Fisher Information Matrix,  which leads to  the  algorithm’s superior performance.
weaknesses: While section 2.1 is likely relevant for  doing a detailed  analysis of Adam in the proposed framework, as far  as I can tell that analysis does not actually take place in the paper or appendix. Given this this  section  feels quite out of place to me? I’m unsure of  it’s value for the main message of the paper. Most readers in the optimization or  statistics community are familiar with Fisher Information Matrix based methods and their connections to second order newton style algorithms,  so I’m unsure of the value of introducing them using ideas form differential geometry on manifolds.

Throughout this paper is repeatedly claimed that Adam is preconditioning with the diagonal  of the Fisher Information Matrix, the approximation used in Adam is not the same thing  in general. Adam has been connected  to second order like or natural gradient like algorithms, but it is known in general that the gradient squared is not an approximation to  the diagonal of the Fisher Information Matrix. While it is in expectation, in the finite data regime there  are clear counter examples to this, such as in [1]  which the author cites.

On line  250 the authors claim that Adam excels  in classification tasks such as next  token prediction, but this seems somewhat  contradictory to the previous line where it claimed that CNNs are often better when trained with SGD. I agree that in many vision tasks SGD matches Adam in performance, but what is left out of  the text is that that is true in most classification problems  in vision, which is again a discrete output space. I’m additionally not  sure   of the strength of the claim that Adam is less strong in the generative  setting. I’m aware of works  such as  [2] that claim the oposite (which the author cites) attempting to answer the question “What factors explain that Adam produces better quality solutions than SGDA when training GANs”, and propose modifications to SGD to help it compete with Adam.

Overall several assertions are made that Adam fails on continuous regression targets, but I feel like there  is not sufficient citation or experimentation to back that up. Adam excelling in discrete output spaces (which again is not always true, training  ResNets with SGD is still very common) is not the same thing as Adam failing on continuous tasks, and needs to be justified if it is being claimed. Examples counter to this idea exist in the literature, such as the quadratic function minimized in figure 6 of [3] where Adam handily outperforms gradient descent.

The notation of some of the equations is a but unclear, for example in equation 15 while I understand the division is being coordinate  wise, this should be explicit in the notation,  otherwise a less familiar reader may thing we’re trying to divide a vector by another vector which is ill defined. 

Cosmetically, the citation style is very   non-standard and makes reading difficult, I would suggest the authors use a  more standard method of in text citation.

Minor but Adam Algorithm written in B.5 is in fact AdamW and has clipping added which was  not included in the origional algorithm.

The central weakness of this paper in my opinion is it misunderstands how approximate the approximations in Adam are. The idea of the Adam update being connected to the diagonal of the Fisher information matrix is not new, it was mentioned in Kingma and Ba (2014). The optimization community has tried very hard to understand why Adam works (another weakness of this paper is there is no related work regarding the vast amount of research into understanding Adam) and this approach has not appeared to yield progress. The authors acknowledge the significance of these approximations in appendix B.3, but given the amount of work showing that these approximations are often very poor I don’t think the community can comfortably understand Adam as a natural gradient algorithm.


[1]
Frederik Kunstner, Lukas Balles, Philipp Hennig

Limitations of the Empirical Fisher Approximation for Natural Gradient Descent

https://arxiv.org/abs/1905.12558

[2]
Samy Jelassi, Arthur Mensch, Gauthier Gidel, Yuanzhi Li\

Adam is no better than normalized SGD: Dissecting how adaptivity improves GAN performance

https://openreview.net/pdf?id=D9SuLzhgK9

[3]
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, Alberto Bietti

Heavy-tailed class imbalance and why adam outperforms gradient descent on language models

https://arxiv.org/pdf/2402.19449

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 3
questions: Please see the weaknesses for questions and improvements.
soundness: 2
strengths: 1. The authors have identified an important and interesting connection between the success of the Adam optimization method and optimization by the natural gradient method.
2. The authors proposed an explanation for why Adam's advantages primarily emerge in problems with discrete distributions.
3. The authors established principles for using momentum, weight decay, and clipping in optimization with invariant gradients. Based on this analysis, they proposed a new method — FAdam, which demonstrates improved performance compared to traditional Adam.
summary: The authors proposed a connection between the Adam optimizer and natural gradient optimization, treating the moving average of squared gradients in Adam as an estimate of the diagonal elements of the Fisher information matrix. They hypothesised that Adam's advantage over other methods might be due to its use of natural gradients; the advantage is particularly noticeable in tasks with discrete distributions, since they allow for a tighter approximation of the Fisher matrix.

The authors also offered a justification for the necessity of normalization by the square root of squared gradients to ensure basis invariance when averaging gradients in Adam. Additionally, they analyzed how momentum, weight decay, and clipping should function in the context of natural gradients and proposed new variants of Adam and Adafactor — FAdam and FAdafactor. The proposed FAdam method demonstrates superior performance for models like LLMs and VQ-VAEs, as well as in ASR tasks.
weaknesses: 1. The authors did not provide an analysis to assess the accuracy of the approximations and simplifications used in this method.
  
a) Why is the transition from sampling from $p(x|\theta)$ to sampling from $p_{data}$ valid? In the case of an undertrained model, the distribution $p(x|\theta)$ can differ significantly from the marginalized $p_{data}$ distribution.     
  
b) How accurate is the transition from Eq. (20) to Eq. (21) and why it is not critical to the method's effectiveness?
  
c) How accurate is the FIM approximation throughout the hundreds of optimization steps for EMA? 

If the authors could provide ablation experiments comparing Adam, FAdam, a true natural gradient method, and other methods incorporating intermediate transitions on simple tasks (e.g., CIFAR-10), it would significantly increase confidence in the results.

2. In the theoretical justification for preferring discrete distributions: due to uniform sampling from $p_{data}$, discrete distributions can also provide a poor approximation of the FIM. This is because the concentration of the distribution may lie in false logits, which is common in yet-not-fully-trained networks. The score might not be large enough to yield a good approximation.

3. While Amari et al. (2019) prove that unit-wise block diagonal FIM has off-diagonal blocks smaller by $\frac{1}{\sqrt{n}}$, the authors' interpretation appears to extend beyond the original result. Their derived claim about individual diagonal weights dominating off-diagonal weights by $\frac{1}{\sqrt{n}}$ (lines 186-188) may need additional justification, as it's not directly supported by Amari's work.

4. The absence of confidence intervals in the experimental results prevents from being fully certain of FAdam superiority, due to marginal score improvements. Additionally, providing further experiments on a broader range of domains would strengthen the evidence of the proposed method's improved performance.

5. Minor typos: in Eq. (25), it should be the square of the norm; in Eq. (19), "approx." should replace the equality sign.

To summarize, I believe this paper relies too heavily on unjustified approximations and is not yet ready for the conference. However, if the authors provide additional experimental and theoretical validation, I might increase my score.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 1
questions: I suggest the authors rethink their major contribution of this paper. As a paper to propose a new optimizer, it might be better to first introduce the new algorithm and present the pros of this algorithm (they could be empirical results of theoretical guarantees). Although I understand some basic illustrations about the preliminary knowledge or motivation of some terms are necessary, at least for this paper, I believe the discussion in the current manuscript should be refined. For example,  I do not get any interesting insights from the discussion about the connection between log-likelihood of Gaussian distribution and $\ell_2$ loss, as it is a basic knowledge of statistics, and seems not deeply correlated with the Fadam.
soundness: 2
strengths: This paper provides a comprehensive discussion of the previous works, and the empirical results seem supportive.
summary: In this paper, the authors provide an explanation of the second-order moment $v_t$ of Adam from the perspective of the diagonal Fisher information, and propose a new optimizer, FAdam, by utilizing this perspective.
weaknesses: At least from my perspective, I can not grasp the main contribution of this paper. I feel it is more like a technical report or a review instead of a paper. I will list my concerns as follows:

1. The main conclusion of this paper is somewhat unclear, and the writing is difficult to follow. From my understanding, the authors aim to claim that Adam is a variant of natural gradient descent and introduce a new optimizer, Fadam, as I outlined in the summary. However, they spend over half of the paper discussing basic statistical properties and formulas related to Fisher information, along with extensive reviews of previous works, without presenting their own results or conclusions. In contrast, the descriptions of the algorithms and the theoretical convergence results are glossed over. While I acknowledge that some discussion of prior works is necessary, I believe it should be integrated with the proposed methods and conclusions of this paper. In summary, the lengthy review of existing literature and preliminary knowledge renders the current manuscript confusing and unappealing.


2. The technical contribution of this paper is relatively insufficient. As a work proposing a new optimizer, the authors fail to provide a rigorous theoretical guarantee of convergence. The current version's convergence analysis disregards the effects of momentum, and even this incomplete result is derived directly from another paper. Furthermore, the statement that _“Since FAdam’s momentum is analogous to Polyak momentum, FAdam’s momentum also tightens the convergence bound. Therefore, the convergence bound for the natural gradient without momentum is looser than the convergence bound for FAdam,”_ is presented without adequate justification and is not convincing. It is unreasonable to assert that the convergence bound of one optimizer is looser than that of another without rigorous derivation.