AdaFRUGAL: Adaptive Memory-Efficient Training with Dynamic Control

1. Problem Statement

Training Large Language Models (LLMs) is hindered by exorbitant memory requirements, particularly due to the storage of optimizer states (e.g., momentum and variance vectors for Adam/AdamW). For large models, these states can consume over 64 GB of GPU memory, forming a barrier to entry.

While the FRUGAL framework addressed this by splitting parameters into a small "state-full" subspace (updated with AdamW) and a large "state-free" subspace (updated with memoryless SignSGD), it relies on static hyperparameters:

• $\rho$ (State-full ratio): The proportion of parameters kept in the state-full subspace.
• $T$ (Update frequency): How often the subspace is redefined.

The Limitation: Static values for $\rho$ and $T$ are suboptimal because training dynamics change over time. Early training requires robust, frequent updates (high $\rho$, low $T$), while later phases (convergence) require less memory and fewer updates. A fixed configuration leads either to wasted memory or suboptimal convergence.

2. Methodology: The AdaFRUGAL Framework

The authors propose AdaFRUGAL, an adaptive extension that replaces static hyperparameters with two dynamic control mechanisms to automate resource management.

A. Dynamic Control of the State-full Ratio (Dynamic $\rho$)

• Motivation: Early training requires a large state-full subspace for stability and acceleration. Later training can rely on the simpler SignSGD optimizer, allowing the state-full subspace to shrink.
• Mechanism: A linear decay schedule adjusts $\rho$ over training steps ($k$):
  $$\rho(k) = \max\left(\rho_{end}, \rho_{start} - (\rho_{start} - \rho_{end}) \cdot \frac{k}{K_{total}}\right)$$
  This allows the model to start with a robust subspace and gradually reduce it over training to save memory.

B. Adaptive Control of Update Frequency (Dynamic $T$)

• Motivation: Frequent subspace updates ($T$) are computationally intensive but necessary when gradients change rapidly. Once the model converges, frequent updates become redundant.
• Mechanism: A loss-aware adaptive schedule monitors the validation loss ($L_{val}$).
  1. Calculate the relative change in validation loss ($\Delta L_{rel}$) over a window.
  2. If $\Delta L_{rel}$ falls below a stability threshold ($\tau_{low}$), the system assumes the model has stabilized.
  3. The update interval $T$ is multiplicatively increased ($T_{new} = \min(T_{max}, T_{current} \cdot \gamma_{increase})$).
     This automatically reduces computational overhead in later phases without manual intervention.

C. Integrated Algorithm

The algorithm dynamically calculates $\rho_k$ and $T_k$ at each step. It redefines the projection matrix $P_k$ only when $k \mod T_k = 0$, ensuring the state-full subspace is updated only when necessary, while the proportion of parameters within this subspace shrinks over time.

3. Main Contributions

1. Dynamic $\rho$ Policy: Introduces a mechanism to dynamically trade performance for memory savings, starting with high memory usage for stability and ending with low usage for efficiency.
2. Loss-aware Dynamic $T$: Automates the trade-off between adaptability and computation by increasing update intervals only when model convergence stagnates.
3. Comprehensive Evaluation: Validated across three distinct scenarios:
   • Large-scale English Pre-training (C4).
   • Large-scale Vietnamese Pre-training (VietVault).
   • Downstream Fine-tuning (GLUE Benchmark).

4. Experimental Results

Experiments were conducted on LLaMA-130M (Pre-training) and RoBERTa-Base (Fine-tuning) using NVIDIA RTX A6000 GPUs.

Performance vs. Memory Trade-offs

• AdaFRUGAL-Dynamic-$T$: Achieved the best final perplexity (18.58 on C4) among memory-efficient methods, outperforming static FRUGAL (18.60) and GaLore (20.01). It also reduced training time by ~15% compared to the static baseline ($T=200$).
• AdaFRUGAL-Dynamic-$\rho$: Demonstrated significant memory savings, reducing peak GPU memory over time from 0.52 GB to 0.37 GB. Although perplexity was slightly higher (18.85), it offers a critical advantage for hardware-constrained environments.
• AdaFRUGAL-Combined: Balanced both metrics, achieving competitive perplexity (18.83) with dynamic memory reduction.

Robustness and Generalization

• Cross-lingual: The same hyperparameters worked effectively on both English (C4) and Vietnamese (VietVault) datasets without domain-specific tuning.
• Fine-tuning: On the GLUE Benchmark, the Dynamic-$T$ variant achieved the highest average score (86.8), outperforming LoRA (86.1) and static FRUGAL (86.7), proving that adaptive stability supports convergence in sensitive fine-tuning tasks.

Scalability Analysis

The authors extrapolated that memory savings via dynamic $\rho$ scale quadratically with hidden size ($h^2$). For a 7-billion-parameter model, memory savings could reach ~5.7 GB, potentially enabling training on hardware that would otherwise lack sufficient memory.

5. Significance

• Automation: Eliminates the need for manual, task-specific tuning of $\rho$ and $T$, making memory-efficient training more accessible.
• Efficiency: Optimizes both memory footprint (via Dynamic $\rho$) and training time (via Dynamic $T$), addressing the two main bottlenecks in LLM training.
• Practicality: Offers a "best-of-both-worlds" solution where researchers can achieve near-full-rank performance (AdamW) with significantly reduced resource requirements, enabling large-scale training on consumer-grade or restricted enterprise hardware.

In summary, AdaFRUGAL represents a strategic shift from static optimization to adaptive, phase-aware resource management, providing a robust framework for the next generation of efficient LLM training.