Gefen: Optimized Stochastic Optimizer

• URL: http://arxiv.org/abs/2606.13894v1
• Date: Thu, 11 Jun 2026 20:38:09 GMT
• Title: Gefen: Optimized Stochastic Optimizer
• Abstract summary: AdamW is a default for modern deep learning, but its first and second moment states add roughly two parameter-sized buffers to training memory.<n>We propose Gefen, a memory-efficient codebook that shares second-moment estimates across parameter blocks and quantizes the first moment using a learned codebook.<n>We show that Gefen reduces AdamW's memory footprint by 8x while maintaining the same performance, corresponding to a reduction of 6.5 GiB per billion parameters.
• Score: 32.771995350910125
• License: http://creativecommons.org/licenses/by/4.0/
• 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

Related papers

• PowerStep: Memory-Efficient Adaptive Optimization via $\ell_p$-Norm Steepest Descent [29.903629736539013]
  Adaptives have become the default standard for training large-scale neural networks such as Transformers.<n>We introduce PowerStep, a memory-efficient that achieves a momentum-wise adaptivity without second-moment statistics.<n>Experiments on Transformer models demonstrate that PowerStep matches Adams convergence speed while halving memory.
  arXiv  Detail & Related papers  (2026-05-11T10:36:27Z)
• Pushing the Limits of Low-Bit Optimizers: A Focus on EMA Dynamics [64.62231094774211]
  Statefuls (e.g., Adam) maintain auxiliary information even 2x the model size in order to achieve optimal convergence.<n>SOLO enables Adam-styles to maintain quantized states with precision as low as 3 bits, or even 2 bits.<n>SOLO can thus be seamlessly applied to Adam-styles, leading to substantial memory savings with minimal accuracy loss.
  arXiv  Detail & Related papers  (2025-05-01T06:47:45Z)
• APOLLO: SGD-like Memory, AdamW-level Performance [61.53444035835778]
  Large language models (LLMs) are notoriously memory-intensive during training.<n>Various memory-efficient Scals have been proposed to reduce memory usage.<n>They face critical challenges: (i) costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial memory overhead to maintain competitive performance.
  arXiv  Detail & Related papers  (2024-12-06T18:55:34Z)
• Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss [59.835032408496545]
  We propose a tile-based strategy that partitions the contrastive loss calculation into arbitrary small blocks. We also introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems. Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed.
  arXiv  Detail & Related papers  (2024-10-22T17:59:30Z)
• AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
  We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models. AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
  arXiv  Detail & Related papers  (2023-10-16T09:04:28Z)
• CAME: Confidence-guided Adaptive Memory Efficient Optimization [20.009302737137787]
  Adaptive gradient methods have demonstrated excellent performance in the training of large language models. The need for maintaining second-moment estimates requires maintaining a high cost of extra memory overheads. Several memory-efficients have been proposed to obtain a drastic reduction in auxiliary memory usage, but with a performance penalty. We propose CAME to simultaneously achieve two goals: fast convergence as in traditional adaptive methods, and low memory usage as in memory-efficient methods.
  arXiv  Detail & Related papers  (2023-07-05T06:05:36Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)
• 8-bit Optimizers via Block-wise Quantization [57.25800395197516]
  Statefuls maintain statistics over time, e.g., the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past values. This state can be used to accelerate optimization compared to plain gradient descent but uses memory that might otherwise be allocated to model parameters. In this paper, we develop first gradients that use 8-bit statistics while maintaining the performance levels of using 32-bit gradient states.
  arXiv  Detail & Related papers  (2021-10-06T15:43:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.