Technical Summary: GEFEN: OPTIMIZED STOCHASTIC OPTIMIZER

Problem Statement

Modern deep learning relies heavily on the AdamW optimizer due to its robust performance across diverse architectures. However, AdamW maintains first and second moment moving averages for every parameter, typically in FP32. This results in an optimizer state memory footprint of approximately twice the size of the model parameters (roughly 8 bytes per parameter). For large-scale models (e.g., 13B parameters), this state alone can consume ~97 GiB of memory, severely constraining trainable model sizes, maximum batch sizes, and training throughput on memory-constrained hardware.

Existing memory-efficient alternatives face significant trade-offs:

• Quantized Optimizers (e.g., Adam8bit, Adam4bit): Reduce memory but often rely on manually crafted, suboptimal codebooks, fixed block sizes (introducing hyperparameters), and may not be universally applicable (e.g., avoiding embedding quantization).
• Structured Optimizers (e.g., Adam-mini): Share second-moment statistics based on manual rules mimicking Hessian block-diagonal structures. However, these require architecture-specific metadata (e.g., number of attention heads) and lack strong theoretical justification for why Hessian-aligned grouping is beneficial.

The goal is to design an optimizer that reduces memory footprint significantly (beyond Adam-mini) while maintaining AdamW-level performance, requiring no manual configuration, additional hyperparameters, or learning-rate retuning.

Methodology

The proposed method, Gefen, achieves memory reduction through two primary mechanisms: automatic second-moment sharing and learned first-moment quantization.

1. Theoretical Motivation

The approach is grounded in a theoretical result (Theorem 4.1) demonstrating that for parameters with large mixed Hessian entries, the ratio of their squared gradients is constrained toward one. This implies that parameters with high Hessian affinity naturally possess similar squared gradient magnitudes, making them ideal candidates for sharing second-moment estimates. While computing the Hessian directly is intractable at scale, the authors observe that the structure of squared gradients in the initial training step correlates with this Hessian structure.

2. Automatic Block Partitioning

Instead of relying on manual rules or Hessian computation, Gefen infers the optimal block structure from the squared gradients of the first optimization step ($g_0$).

• Algorithm: For a flattened parameter tensor, the algorithm tests all proper divisors $p$ of the tensor size $n$. It reshapes the squared gradients into blocks of size $p$ and calculates the within-block heterogeneity $E(p)$ (the variance of squared gradients within blocks).
• Selection Strategy: The algorithm selects the period $p$ that yields the first pronounced drop in $E(p)$ as block size increases, favoring larger blocks for memory efficiency while ensuring gradient similarity. If no meaningful drop is found or the optimal $p < 8$, it defaults to $p=1$ (no sharing).
• Usage: This single partitioning scheme is used for both sharing the second moment and quantizing the first moment.

3. Exact Dynamic Programming Quantization

To reduce the memory of the first moment (momentum), Gefen employs a learnable, exact quantization codebook rather than fixed or heuristic schemes.

• Codebook Learning: Using the initial gradients, the method constructs a histogram of normalized gradient values. It then solves the optimal scalar quantization problem using Dynamic Programming (DP) (inspired by Wang & Song, 2011) to find the exact optimal codebook for the given distribution and bit-width.
• Stability: Unlike Lloyd-Max, which can suffer from "contraction" (shrinking quantization range over time) and sensitivity to initialization, Gefen forces the codebook endpoints to be fixed at $-1$ and $1$. This ensures the quantization range is preserved throughout training.
• Efficiency: The codebook is learned once at the beginning of training (using CUDA for histogramming and CPU/Numba for DP) and reused for the entire training run.

4. Algorithm Flow

Gefen operates as a drop-in replacement for AdamW:

1. Initialization: Partition parameters based on $g_0$ and learn the quantization codebook.
2. Update Step:
   • Compute gradients.
   • Update the first moment (momentum) by dequantizing the stored codebook values, applying momentum, re-quantizing, and updating the block-wise absolute maximum.
   • Update the second moment by computing the mean of squared gradients within each block and applying exponential moving average.
   • Apply bias correction and update parameters using the shared second moment and quantized first moment.

Key Contributions

1. Memory Efficiency: Gefen reduces the optimizer memory footprint by approximately 8× compared to AdamW (reducing from ~2.0× parameter size to ~0.25×), corresponding to a reduction of ~6.5 GiB per billion parameters.
2. Automatic Adaptation: The method requires no architecture-specific metadata, manual grouping rules, or additional hyperparameters beyond standard AdamW defaults. It automatically infers block structures from data.
3. Theoretical Grounding: Provides a theoretical link between Hessian entries and squared gradient ratios, justifying the sharing of second-moment statistics for Hessian-aligned parameters.
4. Optimal Quantization: Introduces an exact, histogram-based dynamic programming approach for quantizing momentum, avoiding the suboptimality and instability of iterative methods like Lloyd-Max.
5. Performance Parity: Demonstrates that these reductions do not come at the cost of model quality, matching AdamW performance across various tasks.

Experimental Results

The authors evaluated Gefen against AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor, and SM3 across CNNs, ResNets, GPT-2, and Llama 3 models.

• Peak Memory Footprint: Gefen consistently achieved the lowest peak memory usage among compared methods. For a 1.5B parameter model, Gefen reduced the optimizer state from 11.2 GiB (AdamW) to 1.5 GiB.
• Performance (Quality): Across all settings (OpenWebText, C4, ImageNet, MNIST), Gefen achieved validation losses and accuracies comparable to AdamW. Notably, it outperformed Adam4bit and Adam-mini in several metrics while maintaining the same learning rate schedule.
• Throughput: The reduced memory footprint allowed for larger microbatches in distributed training:
  • FSDP (Llama 3-1.5B): Gefen enabled a microbatch size of 2 (vs. 1 for others), resulting in a 56% increase in throughput (tokens/sec).
  • DDP (GPT-2-1B): Gefen enabled training with a microbatch size of 2, whereas AdamW could not fit even a microbatch size of 1. This resulted in a 21% throughput improvement over Adam-mini.
• Stability: Ablation studies confirmed that the learned codebook remains stable throughout training, validating the decision to learn it only once at initialization.

Significance and Claims

The paper positions Gefen as a practical, "drop-in" replacement for AdamW that addresses the memory bottleneck in large-scale training without sacrificing performance or requiring manual tuning.

• Practical Impact: By reducing the optimizer state, Gefen enables training larger models or using larger batch sizes on fixed hardware, directly improving training throughput and sample efficiency.
• Generality: Unlike previous methods that rely on manual heuristics or specific tensor names, Gefen is architecture-agnostic and automatically adapts to the data.
• Complementarity: The authors note that Gefen's optimizer-state compression is orthogonal to gradient communication compression (e.g., 1-bit Adam) and distributed training strategies (FSDP/DDP), meaning it can be combined with these techniques for further gains.

The authors provide a complete Python implementation with fused CUDA kernels, emphasizing reproducibility and ease of integration into existing training pipelines.