Gefen: Optimized Stochastic Optimizer

2026-06-11

AI Summary

This work addresses the high memory overhead of the AdamW optimizer in modern deep learning, which stems from storing momentum buffers twice the size of the model parameters. The authors propose Gefen, a novel optimizer that integrates theoretically motivated Hessian-aligned parameter grouping with hyperparameter-free adaptive quantization. Gefen shares second-moment estimates across parameter blocks and automatically infers block structure from initial squared gradients. It employs dynamic programming to construct a histogram-based codebook for precise quantization of first-moment estimates, while reusing the inferred block structure for efficient scaling. Across diverse training settings, Gefen matches AdamW’s performance while reducing optimizer memory by 6.5 GiB per billion parameters—achieving approximately 8× memory compression—and substantially increasing micro-batch sizes and training throughput.

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

Problem

• memory-efficient optimization
• AdamW
• optimizer memory footprint
• large-scale training
• deep learning

Innovation

• memory-efficient optimizer
• moment sharing
• gradient quantization
• Hessian-aligned blocks
• AdamW alternative