Natural GaLore: Accelerating GaLore for memory-efficient LLM Training and Fine-tuning

## Abstract

Abstract Training LLMs presents significant memory challenges due to growing size of data, weights, and optimizer states. Techniques such as data and model parallelism, gradient checkpointing, and offloading strategies address this issue but are often infeasible due to hardware constraints. To mitigate memory usage, alternative methods like Parameter-Efficient-Fine-Tuning (PEFT) and GaLore approximate weights or optimizer states. PEFT methods, such as LoRA, have gained popularity for fine-tuning LLMs, though they require a full-rank warm start. In contrast, GaLore allows full-parameter learning while being more memory-efficient. This work introduces Natural GaLore , a simple drop in replacement for AdamW, which efficiently applies the inverse Empirical Fisher Information Matrix to low-rank gradients using Woodbury’s Identity. We demonstrate that incorporating second-order information speeds up optimization significantly, especially when the iteration budget is limited. Empirical pretraining on 60M, 130M, 350M, and 1.1B parameter Llama models on C4 data demonstrate significantly lower perplexity over GaLore without additional memory overhead. By fine-tuning RoBERTa on the GLUE benchmark using Natural GaLore , we demonstrate significant reduction in gap 86.05% vs 86.28% for full-finetuning. Furthermore, fine-tuning the TinyLlama 1.1B model for function calling using the TinyAgent framework shows that Natural GaLore achieving 83.09% accuracy on the TinyAgent dataset, significantly outperforms 16-bit LoRA at 80.06% and even surpasses GPT4-Turbo by 4%, all while using 30% less memory. 1 1 1 All code to reproduce the results are available at: https://github.com/selfsupervised-ai/Natural-GaLore.git

## 1 Introduction

Large Language Models (LLMs) have achieved remarkable performance across various disciplines, including conversational AI and language translation. However, training and fine-tuning these models demand enormous computational resources and are highly memory-intensive. This substantial memory requirement arises from storing billions of trainable parameters along with associated gradients and optimizer states.

To quantify this, consider a model with $\Psi$ parameters which is being trained using the Adam optimizer. In this case, storing parameters and their gradients in 16-bit precision formats like FP16 or BF16 requires $2\Psi$ bytes each. The associated optimizer states are typically stored in 32-bit precision (FP32) for numerical stability, necessitating an additional $4\Psi$ bytes for each parameter, gradient momentum, and variance, amounting to $12\Psi$ bytes. Therefore, the total memory requirement sums up to $16\Psi$ bytes. When accounting for model-dependent memory, such as activations during forward and backward passes, and residual memory, like temporary buffers and memory fragmentation, the overall memory footprint can easily exceed $18\Psi$ bytes (Raffel et al., ; Touvron et al., ; Chowdhery et al., ).

This enormous memory demand poses significant challenges, especially when training LLMs on hardware with limited memory capacity. As models continue to scale, efficient memory utilization becomes critical for making training feasible and accessible. In this work, we develop an efficient adaptation to the GaLore algorithm (Zhao et al., ), which significantly reduces the memory footprint during training and fine-tuning of LLMs by approximating the optimizer state. Our approach, Natural GaLore, leverages the low-rank structure of gradients and incorporates second-order information to achieve faster convergence and higher performance without additional memory overhead and can be used as a drop in replacement to standard optimization algorithms like Adam and AdamW.

#### Parallel and Distributed Training Techniques

Researchers have developed various distributed computing techniques that leverage system-level optimizations and hardware resources to mitigate the substantial memory requirements in training LLMs.

One prominent framework is Distributed Data-Parallel (DDP) that combines data parallelism where the training dataset is partitioned across multiple devices or nodes, with efficient gradient synchronization mechanisms, minimizing communication overhead. While data parallelism efficiently utilizes multiple GPUs, it can still face memory bottlenecks when model sizes exceed the memory capacity of a single device.

Model parallelism addresses this limitation by partitioning the model across multiple devices, allowing for the training of models that are too large to fit into the memory of a single GPU. Techniques like pipeline parallelism (Huang et al., ) and tensor parallelism (Shoeybi et al., ) enables the distribution of different layers or partitions of layers across devices. However, model parallelism introduces communication overhead and can be complex to implement effectively.

Another effective technique is gradient checkpointing (Chen et al., ), which reduces memory usage by selectively storing only a subset of activations during the forward pass and recomputing them during the backward pass as needed. This approach trades increased computational overhead for reduced memory consumption, enabling the training of deeper models without exceeding memory constraints.

Memory offloading strategies, such as those implemented in ZeRO-Offload (Rajbhandari et al., ), move optimizer states and gradients to CPU memory when not actively in use, freeing up GPU memory for other operations. ZERO can also partition optimizer states and gradients across DDP processes, eliminating redundancy and significantly reducing memory footprint. Fully Sharded Data Parallel (Zhao et al., ) extends this concept by sharding model parameters in addition to optimizer states and gradients.

These system-level optimizations have been instrumental in training state-of-the-art LLMs such as LLaMA3 (Touvron et al., ), GPT-3 (Brown et al., ), Mistral (Jiang et al., ), and Gopher (Rae et al., ) on multi-node, multi-GPU clusters.

While these distributed computing solutions enable the training of large models by leveraging extensive hardware resources, they come with increased system complexity and operational costs. Therefore, there is a pressing need for alternative approaches that reduce memory consumption without relying solely on distributed computing resources. Optimization techniques that approximate parameters or optimizer states offer a promising direction for making LLM training more accessible and efficient.

#### Parameter-Efficient Fine-Tuning

PEFT techniques efficiently adapt pre-trained language models to various downstream applications without fine-tuning all the model’s parameters (Ding et al., ), significantly reducing the computational and memory overhead.

Among these techniques, the popular LoRA (Hu et al., ) parametrizes a weight matrix $W\in\mathbb{R}^{n\times m}$ as:

$$ $W=W_{0}+BA,$ (1) $$

where $W_{0}$ is a frozen full-rank pre-trained weight matrix, and $B\in\mathbb{R}^{n\times r}$ and $A\in\mathbb{R}^{r\times m}$ are trainable low-rank adapters to be learned during fine-tuning. Since the rank $r\ll\min(m,n)$, the adapters $B$ and $A$ contain significantly fewer trainable parameters, reducing memory requirements for both parameter and optimizer states.

LoRA has been extensively used to reduce memory usage during fine-tuning, effectively enabling large models to be adapted to new tasks with minimal additional memory overhead. There are a few variants of LoRA proposed to enhance its performance (Renduchintala et al., ; Sheng et al., ; Zhang et al., ; Xia et al., ), supporting multi-task learning (Wang et al., ), and further reducing the memory footprint (Dettmers et al., ). Its variant, ReLoRA (Lialin & Schatz, ), extends LoRA’s approach to pre-training by periodically updating the frozen weight matrix $W_{0}$ using the previously learned low-rank adapters. This incremental updating allows for continual learning without storing entire optimizer states for all parameters, leading to faster training times and lower computational costs. Furthermore, this allows for rapid adaptation of large models to multiple downstream tasks without storing separate copies of the entire model for each task.

Despite their benefits, recent works have highlighted several limitations of low-rank reparameterization approaches. LoRA does not consistently achieve performance comparable to full-rank fine-tuning, particularly in complex tasks (Xia et al., ). In pre-training from scratch, methods like ReLoRA require an initial phase of full-rank model training as a warmup before optimizing in the low-rank subspace (Lialin & Schatz, ). The shortcomings of low-rank parameter reparameterization suggest that alternative strategies are needed to achieve both memory efficiency and high performance.

#### Gradient Low-Rank Projection (GaLore)

An alternative to parameter approximation is the approximation of the optimizer states. By reducing the memory footprint associated with optimizer states, it is possible to maintain full-parameter learning—thus preserving model capacity and performance—while achieving significant memory savings.

The core idea behind GaLore (Zhao et al., ) is to exploit the slowly changing low-rank structure of the gradient matrix $g\in\mathbb{R}^{n\times m}$, rather than approximating the weights. During neural network training, gradients naturally exhibit low-rank properties, a phenomenon studied extensively in both theoretical and practical settings (Zhao et al., ; Cosson et al., ; Yang et al., ). This intrinsic low-rank structure of gradients has been applied to reduce communication costs (Wang et al., ; Vogels et al., ) and to decrease memory footprints during training (Gooneratne et al., ; Zhao et al., ).

Specifically, consider the compact SVD decomposition of the gradient matrix $\mathbf{g}=\mathbf{P}\Sigma\mathbf{Q}^{T}$, where $\mathbf{P}\in\mathbb{R}^{n\times r}$ and $\mathbf{Q}\in\mathbb{R}^{m\times r}$ are the associated semi-orthognal matrices. Then, GaLore projects the gradient matrix $\mathbf{g}$ into a low-rank form:

$$ $\mathbf{g}_{\text{low-rank}}=\mathbf{P}^{T}\mathbf{g}.$ (2) $$

Here, $r\ll\min(n,m)$ is the target rank, $n$ is the parameter count, $m$ is the batch size and $\mathbf{g}_{\text{low-rank}}$ serves as an efficient approximation of the original gradient. The projection matrix $\mathbf{P}$ is updated periodically (e.g., every 200 iterations), which incurs minimal amortized computational cost.

By operating on low-rank approximations of the gradients, GaLore significantly reduces the memory footprint, leading to up to 30% memory reduction compared LoRA (Zhao et al., ). Moreover, GaLore maintains full-parameter learning, allowing updates to all model parameters, leading to better generalization and performance than low-rank adaptation methods. Further, GaLore is agnostic to the choice of optimizer and can be easily integrated into existing optimization algorithms with minimal code modifications.

While GaLore offers significant memory savings and enables full-parameter learning, its performance has yet to match that of optimizers in full optimizer state space. Reliance on low-rank gradient approximations may not fully capture the rich optimization dynamics. These limitations suggest that while GaLore is a valuable step toward memory-efficient training, further enhancements are necessary to bridge the performance gap with standard optimizers.

#### Our Approach

In this work, we propose to bridge the gap by incorporating a second-order regularizer into the low-rank gradient estimate, which adjusts parameter updates more effectively, leading to faster convergence. We show that applying the inverse of the empirical Fisher Information Matrix (FIM) to the low-rank gradients leads to variance reduction of the gradient estimate, incorporates information about the curvature of the loss landscape, and reduces dependence on the starting point. All of these lead to significantly faster convergence, especially in a limited iteration regime.

We introduce the Natural GaLore algorithm, a matrix-free algorithm for efficiently applying the inverse FIM to the low-rank gradients, using Woodbury Identity, Cholesky Decomposition, and Matrix-Vector Products, all of which can be efficiently implemented on the GPU. Further, our approach does not require any explicit layer-wise information or significant computational overhead, as is seen in existing approaches like K-Fac (Martens & Grosse, ).

We validate the effectiveness of Natural GaLore through extensive empirical evaluations. Pre-training experiments on LLaMA models with 60M, 300M, and 1.1B parameters using the C4 dataset demonstrate that Natural GaLore achieves significantly lower perplexity than GaLore without additional memory overhead, indicating faster convergence within the same computational budget.

Furthermore, we showcase the practical benefits of Natural GaLore in fine-tuning tasks. We fine-tune the TinyLlama 1.1B model for function calling using the TinyAgent framework. Our results show that Natural GaLore significantly outperforms LoRA in this setting, achieving an accuracy of 83.09% on the TinyAgent dataset. This performance significantly surpasses 16-bit LoRA and exceeds that of GPT-4-turbo by 4%, all while using 30% less memory.

## 2 Accelerating GaLore with Natural Gradients

### 2.1 Next Token Prediction

Generative LLMs are trained to predict the next token in a sequence based solely on the previously observed tokens. This ”causal” approach respects the temporal order of language, ensuring that the model’s predictions at any point depend only on past and not future inputs.

Given a sequence of tokens $x=(x_{1},x_{2},\dots,x_{T})$, the objective is to maximize the likelihood of a sequence by decomposing it into a product of conditional probabilities:

$$ $\displaystyle\text{Prob}_{\mathbf{\theta}}(x)=\prod_{t=1}^{T}\text{Prob}_{ \mathbf{\theta}}(x_{t}\mid x_{<t})$ (3) $$

where $x_{<t}=(x_{1},x_{2},\dots,x_{t-1})$ represents all tokens before position $t$ and $\text{Prob}_{\mathbf{\theta}}(x_{t}\mid x_{<t})$ is the probability of the next token given all previous tokens and the parameter $\mathbf{\theta}\in\mathbb{R}^{n\times m}$.

This is equivalent to minimizing the Negative Log-Likelihood (NLL) of the observed sequences, which is the cross-entropy loss between the predicted probability distribution and the actual next token:

$$ $\displaystyle\Phi(\mathbf{\theta})=-\sum_{t=1}^{T}\log\text{Prob}_{\mathbf{ \theta}}(x_{t}\mid x_{<t})$ (4) $$

This loss penalizes the model more when it assigns lower probabilities to the correct next token. By minimizing this loss, the model learns to assign higher probabilities to appropriate continuations of text. However, the loss is non-convex and high-dimensional, for LLMs the dataset is also massive, making the optimization problem very challenging.

### 2.2 Low-Rank Gradient Descent

Stochastic gradient descent algorithms are iterative, where each step aims to find the optimal update direction that minimizes the loss function locally. Now in the case of GaLore, the update direction is restricted to the affine subspace $\mathbf{u}_{k}\in{\mathbf{\theta}_{k}}+\text{Range}\left(\mathbf{P}_{k}\right)$. Here $\mathbf{P}_{k}\in\mathbb{R}^{n\times r}$ is the left projection matrix, calculated using the compact SVD decomposition of the gradient matrix $\nabla_{\mathbf{\theta}}\Phi(\mathbf{\theta}_{k})=\mathbf{P}_{k}\Sigma\mathbf{
Q}_{k}^{T}$.

Then, the local neighborhood around this update can be defined using the Taylor series expansion (Lin et al., ):

$$ $\displaystyle\Phi(\mathbf{\theta}_{k}+\mathbf{P}_{k}\mathbf{u}_{k})\approx\Phi (\mathbf{\theta}_{k})+\mathbf{g}_{k}^{T}\mathbf{u}_{k}+\frac{1}{2}\mathbf{u}_{ k}^{T}\mathbf{H}_{k}\mathbf{u}_{k}$ (5) $$

where $\mathbf{g}_{k}=\mathbf{P}_{k}^{T}\nabla_{\mathbf{\theta}}\Phi(\mathbf{\theta}_
{k})$ is the low rank projected gradient and $\mathbf{H_{k}}=\mathbf{P}_{k}^{T}\nabla^{2}_{\mathbf{\theta}}\Phi(\mathbf{
\theta})\mathbf{P}_{k}$ is the Hessian matrix.

However, the Hessian matrix $\mathbf{H}_{k}$ is often computationally expensive to compute and store, especially for large-scale language models (LLMs) with billions of parameters. Fortunately, precisely under the condition that the loss function can be represented in terms of KL divergence between the actual and approximated distributions [4], then $\mathbf{H_{k}}$ can be approximated by the FIM. The FIM is defined as the expectation of the Hessian of the negative log-likelihood w.r.t. the data distribution:

$$ $\displaystyle\mathbf{F}_{k}=\mathbb{E}_{x\sim p_{\text{data}}}\left[\mathbf{H} _{k}\right]$ (6) $$

The FIM captures the curvature of the loss landscape and provides a natural metric for the optimization process. Hence, it can better adjust parameter updates according to the geometry of the parameter space. However, as the theoretical data distribution is unknown, in practice, we need to estimate it using the empirical FIM (Martens, ) defined by:

$$ $\displaystyle\mathbf{\hat{F}}_{k}=\frac{1}{h}\sum_{k=1}^{h}\mathbf{g_{k}} \mathbf{g_{k}}^{T}$ (7) $$

where $h$ is the history of gradients from past batches we would like to consider. Then, the optimal direction $\mathbf{u}_{k}^{*}$, which minimizes the loss in this local neighborhood, is given by (cite Fuji et al. paper):

$$ $\displaystyle\mathbf{u}_{k}^{*}$ $\displaystyle=$ $\displaystyle\mathbf{\hat{F}}_{k}^{-1}\mathbf{g}_{k}$ (8) $$

This leads to the optimal gradient descent update step:

$$ $\displaystyle\mathbf{\theta}_{k+1}=\mathbf{\theta}_{k}-\eta\mathbf{P}_{k} \mathbf{u}_{k}^{*}$ (9) $$

for some learning rate $\eta$.

Many popular stochastic optimization algorithms approximate the diagonal of the empirical FIM using second-moment estimates of the gradient $\mathbf{g}_{k}$, which when added with Polyak style parameter averaging (i.e., momentum), asymptotically achieve the optimal Fisher efficient convergence rate (Martens, ).

For instance, in the case of Adam (Kingma & Ba, ), the optimal update step is approximated by including the momentum term $\mathbf{m}_{k}\in\mathbb{R}^{r\times m}$ and the learning rate $\eta$ is scaled by the square root of the second moment estimate $\mathbf{v}_{k}\in\mathbb{R}^{r\times m}$. With all operations being elementwise, the update direction becomes:

$$ $\displaystyle\mathbf{m}_{k}$ $\displaystyle=$ $\displaystyle\beta_{1}\mathbf{m}_{k-1}+(1-\beta_{1})\mathbf{g}_{k}$ (10) $\displaystyle\mathbf{v}_{k}$ $\displaystyle=$ $\displaystyle\beta_{2}\mathbf{v}_{k-1}+(1-\beta_{2})\mathbf{g}^{2}_{k}$ (11) $\displaystyle\mathbf{u_{k}^{*}}$ $\displaystyle=$ $\displaystyle\mathbf{m}_{k}/\sqrt{\mathbf{v}_{k}+\epsilon}$ (12) $$

This update, when applied to [9], gives the GaLore optimization algorithm, which is memory efficient as it only requires storing the projection matrix and the costly optimizer states $\left(g_{k},m_{k},v_{k}\right)$ are now significantly reduced by a factor of $\frac{n}{r}$, where the rank $r$, can be chosen based on the tradeoff between memory limitations and performance requirements.

### 2.3 Natural GaLore and Fisher Efficiency

Despite clear advantages, the performance of GaLore is not on par with AdamW (Loshchilov & Hutter, ) optimization on the original space. To bridge this gap, we propose Natural GaLore, which uses the full empirical FIM, thereby incorporating the missing second-order interaction information in the optimization process.

As we now argue, this leads to a much more favorable dependence on the starting point, which means that the optimizer can make much more progress given a limited iteration budget. Further, when using a decaying learning rate schedule like with AdamW (Loshchilov & Hutter, ), the asymptotic convergence rate can be faster (Martens, ) by a significantly large constant factor.

Natural gradient descent is known (Martens, ) to be Fisher efficient, precisely for our loss function [4]. Fisher efficiency means that the natural gradient estimator asymptotically achieves the lowest possible variance among all unbiased gradient estimators.

For Natural GaLore, the gradient descent update [9] leads to a sequence of estimates $\mathbf{\theta}_{k}$ whose variance satisfies (Amari, ):

$$ $\displaystyle\text{Var}[\mathbf{\theta}_{k}]=\frac{1}{mk}\mathbf{F}_{k}^{-1}( \mathbf{\theta}_{k}^{*})+\mathcal{O}\left(\frac{1}{k^{2}}\right)$ (13) $$

which is asymptotically the smallest possible variance matrix satisfying the Cramér-Rao lower bound, that any unbiased estimator computed from $mk$ training samples can have, with $m$ being the batch size.

Here, $\mathbf{\theta}_{k}^{*}$ is the local optimum in the neighborhood defined by the Taylor series expansion [5] around the update direction. This is an important caveat, as the guarantee is only for local convergence in a convex neighborhood. The loss function is non-convex, so the property can not be stated to hold for the global optimum.

The result also relies on the computation of the exact FIM $\mathbf{F}_{k}(\mathbf{\theta}_{k})$ using the entire data distribution, which is not practical. The Fisher efficiency guarantee is, however, only approximately satisfied when using the empirical FIM $\mathbf{\hat{F}_{k}}$ instead. Nevertheless, we still get a variance reduction in the gradient estimates, leading to faster convergence and better optimization performance in the early stages of training large-scale models, making it especially valuable for training with a limited iteration budget.

Further, incorporating second-order information through the empirical FIM allows the optimizer to account for the curvature of the loss landscape, enabling natural gradient descent to take more informed steps than standard gradient descent, potentially escaping flat regions or navigating steep ravines more effectively.

In (Martens, ), it was shown that the expected update direction can be expressed as a sum of two terms, one that scales as $\mathcal{O}(1/k)$, which is independent of the starting point and another that scales as $\mathcal{O}(1/k^{2})$, which is dependent on the starting point. If momentum is applied to the gradient estimator, the first term becomes independent of the choice of FIM estimator, thereby not leading to any asymptotic improvements. However, regularizing with the empirical FIM estimate can significantly reduce the constant factor associated with the starting-point-dependent second term. This leads to practical performance gains in finite iteration regimes (although negligible for large $k$).

Finally, the Fisher efficiency result also assumes that the model can perfectly capture the data distribution, a condition known as *realizability*. However, with the growing size of LLMs, this assumption is likely to hold, thereby satisfying the conditions for the guarantee. Therefore, especially in low-resource settings, Natural GaLore can be a promising approach for training LLMs under memory constraints.

### 2.4 Natural Gradient Transform

Our Natural GaLore algorithm is designed to efficiently apply the inverse empirical FIM to low-rank gradients using Woodbury’s Identity. Most of the steps in the algorithm are similar to GaLore (Zhao et al., ), with the critical difference being the incorporation of the natural gradient transform.

In order to implement the natural gradient transform, we compute the inverse of the empirical FIM and apply it to the gradient $\mathbf{g_{k}}$ using Woodbury’s Identity, which allows us to efficiently compute the inverse of a matrix of the form $A+UBU^{T}$. Woodbury’s Identity states that:

$$ $\displaystyle(A+UBU^{T})^{-1}=A^{-1}-A^{-1}U(B^{-1}+U^{T}A^{-1}U)^{-1}U^{T}A^{ -1}$ (14) $$

Now, if we choose $\mathbf{\hat{F}}_{k}=\lambda I+GG^{T}$, $A=\lambda I$, $U=G$, and $B=I$, where $G=[\operatorname{vec}(\mathbf{g}_{k}),\operatorname{vec}(\mathbf{g}_{k-1}),
\ldots,\operatorname{vec}(\mathbf{g}_{k-s})]$ is the stacked gradient matrix over the past $s$ gradients and $\lambda$ is a small constant for Tikhonov regularization, then, the inverse of the empirical FIM applied to the gradient $\mathbf{g_{k}}$ i.e. the natural gradient $\mathbf{\tilde{g}}_{k}=\mathbf{\hat{F}}_{k}^{-1}\mathbf{g_{k}}$ can be calculated as:

$$ $\displaystyle\mathbf{\tilde{g}}_{k}=\frac{1}{\lambda}\mathbf{g_{k}}-\frac{1}{ \lambda}G\left(\lambda I+G^{T}G\right)^{-1}G^{T}\mathbf{g_{k}}$ (15) $$

To compute the above formula efficiently, let $S=I+\frac{1}{\lambda}G^{T}G\in\mathbb{R}^{s\times s}$ and $y=G^{T}\mathbf{g_{k}}$. Cholesky decomposition is used to solve for $z$ in

$$ $\displaystyle Sz=y$ (16) $$

which requires only $\mathcal{O}(s^{2})$ time. Then, the final natural gradient estimate can be computed using only matrix-vector products, which is very memory efficient:

$$ $\displaystyle\mathbf{\tilde{g}}_{k}=\frac{1}{\lambda}\mathbf{g_{k}}-\frac{1}{ \lambda^{2}}Gz$ (17) $$

This natural gradient estimate $\mathbf{\tilde{g}}_{k}$ can then be sent to the Adam optimizer [12], and the model parameters the same way as in GaLore.

## 3 Experiments

We evaluate Natural GaLore on pre-training and fine-tuning tasks for LLMs. All experiments are conducted on a single node with 8 NVIDIA A100 GPUs to leverage high-performance computing capabilities, yet stay within reasonable limits.

### 3.1 Pre-training on the C4 Dataset

To assess the effectiveness of Natural GaLore, we apply it to pre-train LLaMA-based language models of sizes ranging from 60 million to 1.1 billion parameters, on the C4 dataset. The C4 dataset is a colossal, cleaned version of the Common Crawl Corpus, primarily intended for pre-training language models and word representations (Raffel et al., ). It provides a diverse and extensive corpus, making it suitable for evaluating pre-training methods in realistic scenarios.

We adopt the experimental setup from Lialin & Schatz (), utilizing a LLaMA-based(^2^22LLaMA materials in our paper are subject to the LLaMA community license.) architecture with RMSNorm and SwiGLU activations (Shazeer, ; Touvron et al., ). We maintain the same set of hyperparameters for each model size across all methods, except for the learning rate, which is tuned individually to ensure optimal performance. All experiments use the BF16 format to reduce memory usage without compromising computational efficiency, the same computational budget and the best validation perplexity is reported.

Figure: Figure 1: Training and Validation log Perplexity for Llama 1.1B
Refer to caption: x1.png

**Table 1: Comparison of Natural GaLore with other low-rank algorithms on pre-training various sizes of LLaMA models on the C4 dataset. Validation log perplexity is reported (averaged over 5 runs), along with a memory estimate (in gigabytes) of the total parameters and optimizer states based on BF16 format.**
|  | 60M | 130M | 350M | 1.1B |
| --- | --- | --- | --- | --- |
| Full-Rank | 3.52 (0.36G) | 3.22 (0.76G) | 2.93 (2.06G) | 2.72 (7.80G) |
| Natural GaLore | 3.53 (0.24G) | 3.22 (0.52G) | 2.93 (1.22G) | 2.80 (4.38G) |
| GaLore | 3.56 (0.24G) | 3.24 (0.52G) | 2.95 (1.22G) | 2.90 (4.38G) |
| Low-Rank | 4.35 (0.26G) | 3.82 (0.54G) | 3.62 (1.08G) | 4.96 (3.57G) |
| LoRA | 3.55 (0.36G) | 3.52 (0.80G) | 3.24 (1.76G) | 2.96 (6.17G) |
| ReLoRA | 3.61 (0.36G) | 3.38 (0.80G) | 3.37 (1.76G) | 2.91 (6.17G) |
| Rank $r/d_{\text{model}}$ | 128 / 256 | 256 / 768 | 256 / 1024 | 512 / 2048 |
| Training Tokens | 1.1B | 2.2B | 6.4B | 13.1B |

Table 1 presents the validation perplexity and memory consumption for models trained with different methods and Figure 1 shows the training run for the Llama 1.1B model. Our proposed Natural GaLore consistently outperforms GaLore (Zhao et al., ) across all model sizes, achieving validation perplexities closer to the full-rank baseline while maintaining significant memory savings. Furthermore, Natural GaLore exhibits lower perplexities and greater memory consumption compared to other low-rank adaptation methods like LoRA and ReLoRA, due to their less efficient use of low-rank structures and the need for additional optimizer states.

### 3.2 Fine-Tuning RoBERTa-Base on the GLUE Benchmark

To further evaluate the effectiveness of Natural GaLore, we conduct experiments on the General Language Understanding Evaluation (GLUE) benchmark using the pre-trained RoBERTa-Base model. The GLUE benchmark is a collection of nine natural language understanding tasks, including single-sentence tasks like CoLA (Warstadt et al., ), similarity and paraphrase tasks like MRPC (Dolan & Brockett, ) and STS-B (Cer et al., ), and inference tasks like RTE (Dagan et al., ), MNLI (Williams et al., ), and QNLI (Rajpurkar et al., ). This benchmark is widely used to assess the performance of language models on diverse linguistic phenomena.

In our experiments, we fine-tune the RoBERTa-Base model using Natural GaLore and compare its performance with full fine-tuning and LoRA (Hu et al., ). We focus on memory-efficient fine-tuning methods to reduce the computational footprint while maintaining high performance. For each method, we report the average score across all GLUE tasks and individual task scores.

We use the same training hyperparameters across all methods for a fair comparison. The batch size is 32, and we fine-tuned each model for three epochs. The learning rate is selected from {1e-5, 2e-5, 3e-5} based on the best validation performance for each task. For Natural GaLore and LoRA, we experiment with rank values of 4 and 8 to study the trade-off between performance and memory efficiency.

Table 2 presents the results of our experiments. Natural GaLore consistently achieves comparable or better performance than LoRA across most tasks while using less memory. Precisely, with a rank of 4, Natural GaLore attains an average score of 86.05, closely matching the complete fine-tuning baseline of 86.28 and outperforming LoRA’s average score of 85.61. This demonstrates that Natural GaLore can effectively fine-tune large models with reduced memory consumption without sacrificing performance.

**Table 2: Evaluating Natural GaLore for memory-efficient fine-tuning on the GLUE benchmark using pre-trained RoBERTa-Base. We report the average score of all tasks. Memory consumption is reported in millions of parameters (M).**
|  | Memory | CoLA | STS-B | MRPC | RTE | SST-2 | MNLI | QNLI | QQP | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Full Fine-Tuning | 747M | 62.24 | 90.92 | 91.30 | 79.42 | 94.57 | 87.18 | 92.33 | 92.28 | 86.28 |
| Natural GaLore (rank=4) | 253M | 61.50 | 90.80 | 92.10 | 79.50 | 94.20 | 87.05 | 92.30 | 91.15 | 86.05 |
| GaLore (rank=4) | 253M | 60.35 | 90.73 | 92.25 | 79.42 | 94.04 | 87.00 | 92.24 | 91.06 | 85.89 |
| LoRA (rank=4) | 257M | 61.38 | 90.57 | 91.07 | 78.70 | 92.89 | 86.82 | 92.18 | 91.29 | 85.61 |
| Natural GaLore (rank=8) | 257M | 61.70 | 90.90 | 92.25 | 79.80 | 94.40 | 87.20 | 92.35 | 91.25 | 86.23 |
| GaLore (rank=8) | 257M | 60.06 | 90.82 | 92.01 | 79.78 | 94.38 | 87.17 | 92.20 | 91.11 | 85.94 |
| LoRA (rank=8) | 264M | 61.83 | 90.80 | 91.90 | 79.06 | 93.46 | 86.94 | 92.25 | 91.22 | 85.93 |

### 3.3 Fine-Tuning TinyLlama 1.1B for Function Calling in Advanced Agentic Systems

Advanced Agentic Systems (AAS) require language models that can understand and generate code snippets to integrate various tools and APIs, fulfilling user queries through function-calling. We utilize the TinyAgent framework, which provides an end-to-end pipeline for training and deploying task-specific LLM agents capable of efficient and accurate function-calling (Erdogan et al., ) to drive agentic systems at the edge.

Given a natural language query, the LLM agent must generate a sequence of pre-defined function-calls that accomplish the desired tasks. The challenge lies in determining the appropriate arguments, to call the correct functions, in the right order while respecting interdependencies among the functions.

LLMCompiler Kim et al. (), is a framework that enables language models to perform function-calling by first generating a function-calling plan, which includes the required functions and arguments. The LLMCompiler then compiles this plan into an executable sequence of function-calls. The critical aspect is training the model to produce a function-calling plan with the correct syntax and dependencies.

The off-the-shelf pre-trained TinyLlama 1.1B (instruct-32k) model performs poorly on this task. The model generates incorrect sets of functions, hallucinated function names, fails to respect dependencies, and passes arguments incorrectly. This underperformance is expected, as the model was initially trained on datasets like SlimPajama and StarCoder, which are not specific to function-calling tasks. To address this, we follow the TinyAgent framework (Erdogan et al., ) and fine-tune the TinyLlama 1.1B model on a high-quality, curated dataset designed for function-calling.

#### TinyAgent Dataset

The TinyAgent dataset (Erdogan et al., ) is a meticulously curated collection aimed at building a local agentic system for function-calling on Apple MacBooks for day-to-day tasks. It contains 40K examples of natural language queries and corresponding function-calling plans. The dataset is divided into 38K training examples, 1K validation examples, and 1K test examples. It encompasses 16 tasks, including Email, Contacts, SMS, Calendar, Notes, Reminders, File Management and Zoom Meetings. Each task has predefined scripts that the model needs to generate. The dataset is intentionally challenging, requiring the model to understand dependencies between function-calls and the arguments to be passed.

#### Fine-Tuning Procedure

We fine-tune the TinyLlama 1.1B model on the TinyAgent dataset for three epochs using a batch size of 32. The learning rate is set to $7\times 10^{-5}$. After each epoch, the model is evaluated on the validation set, and the best-performing model is selected based on validation performance to be evaluated on the test set.

During fine-tuning, the prompt includes descriptions of the ground truth functions and irrelevant functions serving as negative samples. This strategy encourages the model to learn to select the correct functions rather than merely memorizing the ground truth. Additionally, several in-context examples demonstrate how queries are translated into function-calling plans. These examples are selected using a Retrieval-Augmented Generation (RAG) process based on the user’s query from the training data and a DeBERTa-v3-small model (He et al., ) fine-tuned for multi-label classification for retrieval among the 16 tools.

The training objective is then to maximize the accuracy of the generated function-calling plans. Success is defined by the model generating the correct plan with the proper set of function-calls, correct arguments, and the appropriate order of function-calls. Verifying the selection of the correct set of functions involves straightforward set comparison. However, ensuring the correctness of arguments and the order of function-calls is more complex and requires constructing the associated Directed Acyclic Graph to check for equality.

#### Results and Discussion

**Table 3: Latency, size, and success rate of TinyAgent models before and after quantization. Latency is the end-to-end latency of the function calling planner, including the prompt processing time and generation.**
| Model | Weight Precision | Latency (seconds) | Model Size (GB) | Success Rate (%) |
| --- | --- | --- | --- | --- |
| GPT-3.5 | Unknown | 3.2 | Unknown | 65.04 |
| GPT-4-Turbo | Unknown | 3.9 | Unknown | 79.08 |
| TinyAgent-1.1B | 16-bit (Natural GaLore) | 3.9 | 2.2 | 83.09 |
| 16-bit (LoRA) | 3.9 | 2.2 | 80.06 |  |
| TinyAgent-7B | 16-bit (Erdogan et al., ) | 19.5 | 14.5 | 84.95 |

After fine-tuning, the TinyLlama 1.1B model’s success rate on the test set improved significantly. Table 3 presents the latency, model size, and success rate of various models on the TinyAgent dataset. As shown, Natural GaLore improves the success rate of the 1.1B model from 80.06% (16-bit LoRA) to 83.09%, also surpassing GPT-4-Turbo by 4% and approaching the performance of the larger TinyAgent-7B model, which achieves 84.95%.

These results demonstrate that Natural GaLore not only enhances the performance of smaller models like the 1.1B parameter TinyLlama but also makes them competitive with significantly larger models. By efficiently incorporating second-order information through low-rank natural gradient updates, Natural GaLore enables smaller models to achieve higher accuracy without additional memory overhead.

## 4 Conclusion

We have introduced Natural GaLore, a memory-efficient pre-training and fine-tuning strategy for large language models. Natural GaLore significantly reduces memory usage—by up to 65.5% in optimizer states—while maintaining or even improving performance in large-scale LLM pre-training and fine-tuning tasks. By incorporating second-order information through an efficient approximation of the inverse Empirical Fisher Information Matrix, Natural GaLore enhances convergence rates, especially in regimes with a limited iteration budget.

Importantly, Natural GaLore can serve as a *drop-in replacement* for standard optimizers like AdamW and integrates seamlessly into existing training pipelines. Our experimental results highlight the reproducibility and effectiveness of Natural GaLore across various tasks, including pre-training LLaMA models and fine-tuning on the GLUE benchmark, as well as the TinyAgent function calling tasks. This makes it a compelling choice for large-scale pre-training scenarios where both memory efficiency and model performance are critical.

In the future we want to explore (1) further enhancing memory efficiency by employing low-memory and structured projection matrices, and (2) more extensive empirical evaluation on fine-tuning AAS on a wide variety of tasks. We also hope that our work will inspire future research on memory-efficient training methods from the perspective of optimizer state approximation. We believe that Natural GaLore will be a valuable tool for the community, enabling the training of large-scale models on consumer-grade hardware with limited resources.

## Impact Statement

This work aims to improve the memory efficiency of training LLMs, thereby reducing the environmental impact of LLM pre-training and fine-tuning. By enabling the training of larger models on hardware with lower memory requirements, our approach helps to minimize energy consumption and carbon footprint associated with training LLMs. Furthermore, by making advanced model training more accessible, we contribute to democratizing AI research and development, allowing a broader community to engage with large-scale models without the need for expensive computational resources.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: Could you add some discussion about the generality and scalability of this method?
soundness: 2
strengths: - The integration of natural gradient descent with GaLore is an interesting and promising approach. It maintains the same memory usage as GaLore while achieving better performance, demonstrating the method’s effectiveness.

- The method is applied to a TinyLlama task as a case study, showcasing its practical utility.
summary: This paper proposes Natural GaLore, which integrates natural gradient descent with GaLore for training, allowing GaLore to leverage missing second-order information in the optimization process to enhance performance. Specifically, Natural GaLore incorporates second-order information through the inverse Empirical Fisher Information Matrix (FIM) applied to low-rank gradients using Woodbury’s Identity.
weaknesses: The paper dedicates a large space to past works and background information, with only about one and a half pages (sections 2.3 and 2.4) covering the main technique. Even within these sections, it is difficult to distinguish the novel contributions from related work, as the content is intermingled.

Given that this is a training optimization method, the paper should include a convergence analysis. Additionally, a discussion on the generality and scalability of this method, such as how the performance curve behaves with larger steps, would be beneficial.

The paper should include ablation studies, particularly showing the training curve of the natural gradient descent. This would help in understanding the contributions of this paper.

For the TinyLlama results, the paper does not compare with the full-rank results for the 1.1B model, which would provide a clearer benchmark.

Besides memory savings, the authors are suggested to analyze the computational cost and discuss the ease of integrating the method into existing training frameworks.

The introduction contains extensive content on parallel and distributed training techniques, which seems unnecessary as these methods are orthogonal to the proposed method. The introduction should be focused on your research question and contributions.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Q1: How significant are performance improvements of Natural GaLore (thereby its novelty) with respect to GaLore? Table 1: 3.53 vs 3.56, 3.22 vs 3.24, etc  and Table 2 improvements seems marginal thereby highlighting superiority of the existing GaLore compared to contributions of the newer version in Natural GaLore.

Q2: What is additional cost associated with Natural GaLore compared to GaLore for pre-training and fine-tuning tasks and its tradeoff to above minor performance improvements?

Q3: What is special about task CoLa in Table 2 that favors LoRA compared to GaLore and Natural GaLore?

Q4: How are average performance values calculated in Table 2 across tasks?

Q5: What is performance of GaLore on TinyAgent and how it compares to others in Table 3 ?

Q6: It would be helpful if Section 1: Introduction could be presented with motivation, related work, (itemized) contributions and a summary table comparing LoRa vs GaLore vs Natural Galore (with say check marks and crosses indicating how each of these are distinct along supporting fine-tuning, pretraining,  low-rank approximation on which part, quantifying memory cost and computation cost, etc)

Q7: It would help understand Natural Galore better if a pseudocode of GaLore and Natural Galore could be presented, and help quantify additional cost during training while maintaining memory cost.

Q8: Under PEFT paragraph in Introduction, please define m and n. 

Q9: Line 231: (cite Fuji et al. paper) - A reference is missing.

Q10: Linespace missing between Line 365 and Line 366.
soundness: 2
strengths: 1. The paper adds to the body of work (such as GaLore) focussing on optimizer/gradient efficient techniques for pre-training and fine-tuning LLMs, orthogonal to parameter efficient "fine-tuning" methods like LoRa/Q-Lora. Hence, this work helps bring memory-efficient pre-training (and fine-tuning) to LLMs. In addition, it helps paint the various related work in parallel and distributed training techniques

2. The work seeks to cover the performance gap between low-rank gradient i.e. GaLore and Full-rank pertaining and fine-tuning of LLMs by leveraging the second order term from Taylor series expansion of the update rule. It is well understood fact that second-order Hessian term (from double derivative of loss function) helps capture the curvature of the optimization landscape. The work finds a way to not ignore this term for sake of compensating for the performance loss from optimizer approximation in GaLore.

3. Natural GaLore maintains the significant memory-savings achieved by its underlying technique GaLore compared to LoRA - acting as drop-in replacement
summary: The paper proposed Natural GaLore which is a modified version of the ICML 2024 Oral work "Gradient Low-Rank Projection" (GaLore) for improving the performance of pre-training and fine-tuning of LLMs. In particular, Natural GaLore introduces the second order regularizer information to the existing low-rank gradient estimate (of GaLore) with intention to better capture the optimization dynamics that was left wanting in GaLore due to low-rank approximation of the optimizer states. Natural Galore applies the inverse Empirical Fisher Information Matrix to low-rank gradients using Woodbury's Identity to represent the computationally expensive second order (quadratic) Hessian matrix term in the low-rank gradient descent.
weaknesses: 1. The main contribution of the work is to improve performance of GaLore while maintaining the memory-efficiency. As a result, the presentation of the work seems to be incremental in its current form.

2. The performance results comparing Natural GaLore to LoRA do not hold much relevance as the significant improvements compared to LoRA was achieved by the root work i.e.  ICML 2024 GaLore. When comparing Natural Galore to its predecessor GaLore, which should be considered as the main point of comparison, the improvements are not as significant reported in Table 1 (Pretraining LlaMa, 3.53 vs 3.56, 3.22 vs 3.24,.. etc)  and Table 2 (Finetuning RoBERTa-Base).

3. It seems some reported results in Table 2 have inconsistencies (bold-faced) for the best performing methods for some tasks, eg: 
- CoLa, Natural Galore 61.5 > 61.38 but LoRA is bold-faced, 
- MRPC GaLore is better 92.25 > 92.10 than proposed Natural Galore , but latter is bold-faced. In addition, it is unclear why Natural galore performance with improved optimization dynamics degrades here compared to GaLore?

4. For results on TinyAgent models, GaLore performance has not been reported which could have helped our understanding of LoRA vs GaLore vs Natural Galore, and how much improvement Natural Galore adds to existing GaLore.

5. The presentation of the manuscript could be improved with explicit section/sub sections on Motivation, Related work, Contributions, Algorithm/Pseudocode of Natural Galore with "new" components on exisiting GaLore code, Table of comparison between LoRa, GaLore and Natural Galore (with check marks and crosses indicating how each of these are distinct along supporting fine-tuning, pretraining,  low-rank approximation on which part, memory cost, computation cost, etc)

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Line 230: Can you please elaborate what are past batches here?
- Lines 328-330: What is a typical value of $s$ in $O(s^2)$? 
- Lines 233 and 333: It seems like two different formulations of optimization steps. Can you please clarify step-by-step how do you compute an update in Natural Galore? For example, do you use a momentum when computing an update? Is your update just a product of the inverse of fisher matrix and gradient?
 - (related to the previous question) What are optimizer states in Natural Galore, is it just a gradient and a pseudo Fisher matrix? If you estimate memory as in the paragraph between lines 040-047, how much bytes do you need to perform your update? And how does it compare to Galore?
- Line 402: Does it mean that each optimization method has the same lr? If yes, wouldn't it be fair to compare the algorithms by choosing the learning rate separately for each method as each algorithm may exhibit different training dynamics and thus benefit from different step-sizes?
- In Table 2, can you please provide an std as well, as for some training the improvement is only marginal and thus it is not clear if the result is statistically significant enough?
- Line 459: You mention that the best-performing model is selected, can you please elaborate on that? Do you mean that you train multiple models at the same time?
- In Section 3.3, why don't you compare with Galore for this task?
soundness: 3
strengths: The authors have well motivated their algorithm both from the memory efficiency and optimization perspectives.
To perform efficiently the update based on Fisher Information Matrix, they have introduced an efficient technique to compute the inverse of it, which is discussed in Section 2.4.

To show the effectiveness of Natural Galore, they have evaluated their method on a wide set of NLP tasks. The improvement on the function calling task in advanced agentic systems looks especially interesting.

Overall, the paper is well written and easy to follow.
summary: The paper introduces a novel memory efficient optimization scheme called Natural Galore. It draws inspiration from Galore, which significantly reduces the memory load caused by gradients and optimizer states during the training, and natural gradient descent, which belongs to approximate second-order methods and thus has a faster convergence than a first order method. Natural Galore is empirically evaluated on a set of NLP tasks by training neural networks of different sizes showing the competitiveness of the method with respect to the state-of-the-art algorithms.
weaknesses: There are no theoretical guarantees provided in the paper, only an extensive discussion on why this algorithm should work. At the same time, the experiments have been only executed on the NLP tasks and thus this empirical evaluation may not suffice to conclude that the method is competitive with other optimization algorithms on other tasks, for example, in computer vision. Moreover, there are some unclear moments related to evaluation and comparison methodology (see questions below).

It can be also argued that the ideas proposed in the paper lack novelty, as both low-rank based methods and natural gradient descent methods have already existed for quite awhile. So without significant evidence of superiority of Natural Galore over other low-rank based methods, the contribution is somewhat incremental.

Minor remarks:

line 231: a missing citation

As you have done a small survey on the memory efficient training in the introduction of the paper, consider to mention activation offloading 

Rhu, M., Gimelshein, N., Clemons, J., Zulfiqar, A., and Keckler, S. W. (2016). vDNN: Virtualized Deep Neural Networks for Scalable, MemoryEfficient Neural Network Design

and combined activation checkpointing and offloading:

Beaumont, O., Eyraud-Dubois, L., and Shilova, A. (2021a). Efficient Combination of Rematerialization and Offloading for Training DNNs.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Most are mentioned in Weakness.

Another suggestion:

-Mistake of Bold fonts (STS-B, Natural Galore rank=4)
soundness: 4
strengths: -Efficient training is important for LLMs. How to accelerate model convergence with less memory is a hot topic for both academic and industrial circles.

-This paper proposes an insightful method to accelerate convergence over baseline GaLore. Extensive experiments are shown.

-The paper is easy to follow.
summary: This paper proposes a memory-efficient training method for LLMs, named Natural GaLore. Compared with previous work, GaLore, this paper applies second-order information through inverse Empirical Fisher Information Matrix to reach better model convergence. There is no memory overhead over GaLore for this method.
weaknesses: -My major concern is the experimental improvement of Natural GaLore over the baseline Galore. For example, in Table 1, GaLore is a baseline close to full-rank one, so the improvement of Natural GaLore over it seems to be limited.

-Training throughputs of different methods should be compared, as it is an important metric used to be traded-off for memory efficiency.

-Results of different runs (e.g. Table 1) should be reported to prove the method reduces dependence on the starting points, which is one of the main advantages of the proposed method. 

-Some background of domain knowledge is missing, such as the explanation of Fisher efficiency and Empirical Fisher Information Matrix, thus making part of the paper hard to follow.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 3
questions: 1. In Section 2.4, it is mentioned that G is the stacked gradient matrix over the past s gradients. Would storing past gradients increase memory consumption?  
2. How much additional time does this work consume compared to GaLore? Could GaLore train more steps or epochs within the same time, thereby achieving better performance than this work?
soundness: 2
strengths: 1. The authors provide a detailed mathematical derivation process.  
2. The authors offer a comprehensive introduction to memory efficiency improvements in LLM training.
summary: This paper builds upon GaLore by incorporating second-order terms and using the Fisher Information Matrix to estimate the computationally expensive Hessian matrix for acceleration. Experiments show that this work achieves improved performance compared to GaLore without increasing memory consumption.
weaknesses: 1. This work merely adds second-order terms to GaLore, lacking sufficient innovation.  
2. The performance improvement over GaLore is not significant, and it introduces additional time consumption. However, the authors did not specify the extent of the increase in time consumption, raising concerns about the practicality of this work.