Gefen: Optimized Stochastic Optimizer

Abstract: AdamW is a default optimizer for modern deep learning, but its first and second moment states add roughly two parameter-sized buffers to training memory. We propose Gefen, a memory-efficient optimizer that automatically shares second-moment estimates across parameter blocks and quantizes the first moment using a learned codebook, thereby reducing AdamW's memory footprint by ~8x while maintaining the same performance, corresponding to a reduction of 6.5 GiB per billion parameters. The method is motivated by a theoretical result showing that large mixed Hessian entries constrain the ratio of squared gradients toward one, suggesting that Hessian-aligned parameters are natural candidates for sharing second-moment statistics. Since computing Hessians is impractical at scale, Gefen infers block structure from the initial squared gradients, requiring no architecture-specific metadata or hyperparameters beyond AdamW defaults. Gefen learns an exact histogram-based dynamic-programming quantization codebook and reuses the same blocks for first-moment scaling. Across diverse experiments, Gefen achieves the lowest peak optimizer memory among the compared AdamW-like methods while maintaining AdamW-level performance. In FSDP and DDP training, the reduced memory footprint enables larger microbatches and improves throughput significantly over AdamW, providing a practical drop-in replacement with lower memory usage that can increase throughput and enable training larger models or using larger batch sizes. We provide the complete Python implementation, including fused CUDA kernels at https://github.com/ndvbd/Gefen

• The paper introduces a memory-optimized stochastic optimizer that reduces state memory by approximately 8× compared to AdamW using Hessian-based parameter grouping and histogram quantization.
• It employs automatic block partitioning from initial squared gradients and fused CUDA kernels to enhance training throughput and scalability.
• Empirical results on NLP and vision tasks confirm that Gefen maintains competitive convergence and performance while enabling larger models and batch sizes.

Gefen: Optimized Stochastic Optimizer — An Authoritative Analysis

Motivation and Context

Adam-type optimizers, especially AdamW, are foundational in deep learning, providing robust convergence across architectures and datasets. However, their pervasive usage imposes substantial memory overhead due to the storage of both first and second moment states, typically requiring two parameter-sized memory buffers for optimizer state. This constraint directly impacts the scalability of training large models and batch sizes, affecting both throughput and final model quality.

Gefen addresses these limitations by introducing a memory-efficient stochastic optimizer that leverages theoretical Hessian insights and efficient quantization, maintaining AdamW-level performance while reducing optimizer memory footprint by approximately $8\times$. The optimizer achieves this through block-wise shared second-moment statistics and histogram-based quantization of first moments via learned codebooks, automatically inferred without manual configuration or additional hyperparameters.

Theoretical Foundations

A central theoretical contribution is the demonstration that parameters with strong Hessian-based coupling exhibit close squared gradients. Precisely, the paper establishes that for pairs of parameters with large-magnitude mixed Hessian entries, the squared-gradient ratio contracts toward $1$, justifying the block-wise sharing of second-moment estimates. This insight provides a formal grounding for the block grouping strategy and underpins Gefen’s empirical effectiveness.

The practical constraint that direct Hessian computation is infeasible for large models is circumvented by analyzing the structure of initial squared gradients. The optimizer infers block structure solely using the first iteration gradient, which correlates with Hessian coupling and enables efficient parameter grouping without manual architecture metadata.

Figure 1: Squared-gradient ratio on log scale versus Hessian for GPT-toy layers; as the Hessian magnitude increases, the squared-gradient ratio converges toward 1, empirically validating the theoretical motif.

Algorithmic Design

Gefen’s algorithm encompasses automatic block partitioning, quantization, and block-wise sharing mechanisms:

• Automatic Block Partitioning: The optimizer determines block size by analyzing squared gradients from the initial training step. For each parameter tensor, candidate block sizes (proper divisors) are evaluated based on within-block heterogeneity of squared gradients, and the period with a pronounced improvement is selected and fixed throughout training.
• Quantization of First Moments: Unlike prior quantized optimizers (e.g., Adam8bit, Adam4bit) that utilize hand-designed codebooks and introduce block size hyperparameters, Gefen introduces an exact histogram-based dynamic programming codebook learning algorithm. This method ensures optimal quantization with forced endpoints, guaranteeing stability and preventing contraction effects seen in Lloyd-Max quantization.

Figure 2: Comparison of quantization methods; the exact dynamic programming approach yields stable and optimal codebooks, outperforming Lloyd-Max and reducing quantization error.

• Memory-Efficient State Storage: Gefen stores a single second-moment scalar per block and quantized first moments using the learned codebook, minimizing auxiliary optimizer state.
• CUDA-Accelerated Implementation: The optimizer leverages fused CUDA kernels for efficient computation, further enhancing throughput and reducing temporary memory allocations.

Empirical Results

Gefen exhibits superior performance in both NLP and vision tasks, matching AdamW in all evaluated settings—Llama 3-1.5B (C4), GPT-2-125M (OpenWebText), and ResNet architectures (ImageNet). The validation loss and accuracy trajectories demonstrate that Gefen does not compromise convergence speed or final quality.

Figure 3: Training curves for GPT-2-125M, Llama 3-1.5B, ResNet18, and ResNet101; Gefen maintains AdamW-level performance across all benchmarks.

Optimizer memory reduction is substantial: for a 1.5B-parameter LLM, Gefen reduces optimizer state from 11.2 GiB (AdamW) to 1.5 GiB, enabling larger models and batch sizes on commodity GPUs. This memory efficiency directly translates to increased training throughput in distributed settings (FSDP, DDP), with Gefen supporting larger microbatches and achieving a $56\%$ improvement over AdamW in FSDP, and enabling training where AdamW is infeasible.

Figure 4: Qualitative desiderata comparison, optimizer memory versus model size, and training throughput; Gefen achieves both lowest peak memory and highest throughput among memory-efficient AdamW-like optimizers.

Stability and Ablation Analyses

The quantization codebooks learned at the start of training remain stable throughout, as gradient distribution histograms and corresponding codebooks demonstrate minimal change over epochs. Periodic relearning does not yield performance gains, confirming the sufficiency of a single codebook-learning step.

Figure 5: Gradients and codebooks at different training stages show remarkable stability, validating the single-step codebook learning approach.

Empirical ablations using four-bit quantization show that further memory savings are possible, but performance degradation may occur on some tasks, supporting the default use of eight-bit quantization for maximum generality.

Practical and Theoretical Implications

The practical implication of Gefen is the removal of optimizer memory as a limiting factor in large-scale model training, enabling not just higher throughput but also the exploration of larger architectures and datasets. Gefen’s block-wise second-moment sharing, grounded in Hessian theory, offers a principled pathway for further research in memory-efficient adaptive optimization, potentially informing extensions toward distributed state management and more granular adaptive statistics.

On the theoretical side, the Hessian-gradient contraction theorem motivates broader study of block-wise parameter affinities beyond neural networks, opening avenues for exploiting higher-order geometry in optimizer state compression. Additionally, the use of histogram-based DP quantization sets a new standard for memory-efficient state representation, potentially influencing future optimizer design.

Conclusion

Gefen introduces a theoretically motivated, empirically validated optimizer that achieves AdamW-level performance with an $8\times$ reduction in optimizer-state memory. It obviates manual configuration, maintains generality across architectures, and provides significant throughput improvements in distributed training. The approach highlights the value of Hessian-aligned parameter grouping and exact quantization for scalable deep learning optimization, establishing new benchmarks for practical and robust memory-efficient optimizers.