M+Adam: Low-Precision Training via Additive–Multiplicative Optimization# Optimization

Summary (Overview)

• Core Contribution: M+Adam is a novel optimizer that combines additive (Adam-style) and multiplicative (Madam-style) update mechanisms to enable stable training with low-precision master weights (BF16, FP8, NVFP4) without requiring FP32 master weights or stochastic rounding.
• Key InsightPodcast: The method addresses complementary failure modes: additive updates get rounded to zero at large weight magnitudes, while multiplicative updates cannot change signs or escape zero.
• Main Results: M+Adam consistently outperforms AdamW (and AdamW+SR, AdamW+Kahan) across 60M–1B LLaMA-style models, 1–8× Chinchilla budgets, and multiple precision regimes, with the largest gains in the most aggressive low-precision settings.
• Theoretical Guarantee: The paper proves a monotone descent result under a local smoothness assumption in the combined additive–exponent coordinate system.
• Practical Efficiency: The additional optimizer state can be compressed via low-rank approximations (Apollo-E), keeping memory footprint close to standard AdamW.

---

Introduction and Theoretical Foundation

Background: The Low-Precision Training Problem

Standard mixed-precision training uses FP32 master weights for stable accumulation while performing forward/backward passes in reduced precision (BF16/FP8). However, storing master weights directly in low precision causes significant degradation because floating-point number systems have nonuniform resolution: the spacing between representable values grows with magnitude# magnitude.

Failure Modes of Additive vs. Multiplicative Updates

Additive updates (AdamW) suffer from a critical failure at large weight magnitudes:

• When $|w_t|$ is large, the local quantization bin size grows, so an additive update $u^{(a)}$ can be smaller than the bin width and round back to zero
• This stalls training progress entirely

Multiplicative updates (Madam) have complementary failures:

• They are sign-preserving—cannot flip the sign of a weight
• They are zero-absorbing—cannot escape from $w = 0$

The key insight is that these failure modes are complementary, motivating the combined approach.

Mathematical Setup

The paper models weights in base-2 floating-point decomposition:

$$w_t = m_t \cdot 2^{e_t}$$

where $m_t$ is the mantissa and $e_t$ is the exponent. The gradient with respect to the exponent simplifies elegantly:

$$g_t^{(e)} = \frac{\partial \mathcal{L}}{\partial e} = (\ln 2) \cdot w_t \cdot g_t$$

This allows computing exponent-space gradients without explicit exponent decomposition.

---

Methodology

The M+Adam Update Rule

The core update combines two branches:

Additive branch (Adam-style):

$$u_t^{(a)} = -\eta_a \frac{\hat{u}}{\sqrt{\hat{v}} + \epsilon}$$

Multiplicative branch (Madam-style):

$$\tilde{u}_t^{(m)} = -\eta_m \frac{g_t^{(e)}}{\sqrt{\hat{v}_t^{(e)}} + \epsilon}$$

with normalization:

$$\rho_t = \max(|w_t|, \tau_t), \quad u_t^{(m)} = \text{clip}(\tilde{u}_t^{(m)} / \rho_t)$$

Combined update:

$$w_{t+1} = w_t + w_t u_t^{(m)} + u_t^{(a)}$$

Algorithm 1: M+Adam

Input: learning rates η_a, η_m; moments β₁, β₂; stabilizer ε; threshold τ
State: additive moments (u, v) ← 0, multiplicative moment v^(e) ← 0

repeat
  g ← ∇_w L(w_t)
  
  # Additive branch (Adam-style)
  u ← β₁u + (1-β₁)g
  v ← β₂v + (1-β₂)g²
  û ← u/(1-β₁ᵗ), v̂ ← v/(1-β₂ᵗ)
  u^(a) ← -η_a · û/(√v̂ + ε)
  
  # Multiplicative branch
  g^(e) ← (ln 2) · w_t · g
  v^(e) ← β₂v^(e) + (1-β₂)(g^(e))²
  v̂^(e) ← v^(e)/(1-β₂ᵗ)
  ũ^(m) ← -η_m · g^(e)/(√v̂^(e) + ε)
  ρ ← max(|w_t|, τ)
  u^(m) ← clip(ũ^(m)/ρ)
  
  # Combine
  w_{t+1} ← w_t + w_t·u^(m) + u^(a)
until converged

Theoretical Analysis

The paper proves a monotone descent theorem under a local smoothness assumption in the combined coordinate system:

Assumption (Eq. 5): For all $(d, z)$ in a neighborhood:

$$\mathcal{L}(w_t 2^z + d) \leq \mathcal{L}(w_t) + \langle g, d \rangle + \langle g^{(e)}, z \rangle + \frac{L_a}{2}\|d\|^2 + \frac{L_e}{2}\|z\|^2 + L_{ae}\|d\|\|z\|$$

Theorem 4.1: If the branch outputs satisfy alignment conditions and the curvature parameters satisfy:

$$L_a \eta_a + L_{ae} \eta_m < 2, \quad L_e \eta_m + L_{ae} \eta_a < 2$$

then the update guarantees descent:

$$\mathcal{L}(w_{t+1}) \leq \mathcal{L}(w_t) - \frac{\eta_a}{2}(2 - L_a\eta_a - L_{ae}\eta_m)\|g\|^2 - \frac{\eta_m}{2}(2 - L_e\eta_m - L_{ae}\eta_a)\|g^{(e)}\|^2$$

---

Empirical Validation / Results

Toy Diagnostic: Matrix Fitting

A controlled experiment fitting $\mathcal{L}(W) = c\|W - W^*\|_F^4$ with BF16-stored weights isolates four regimes:

Main Results (Table 2: 1× Chinchilla, validation perplexity)

Key observations:

• M+Adam improves over AdamW in every regime and model size
• Largest relative gains occur in the most aggressive precision regime (NVFP4/FP8): 8.7–16.6% perplexity reduction
• M+Adam beats both AdamW+SR and AdamW+Kahan in BF16/BF16, showing benefits beyond additive update improvements alone
• At 350M with FP8/FP8, M+Adam (19.395) approaches BF16-weight AdamW (19.156) despite coarser weights

Scaling Results

• Improvements persist across 1–8× Chinchilla budgets (not just early-training effects)
• FP8 compute degrades absolute perplexity for both optimizers, but M+Adam maintains its advantage
• 1B results confirm gains extend beyond the main scaling sweep

Optimizer-State Compression (Table 4)

Apollo-E compression keeps memory close to BF16-state AdamW while retaining most of M+Adam's perplexity gains.

---

Theoretical and Practical Implications

Implications for Optimizer Design

The results suggest that optimizer geometry should match the numerical structure of the storage format. When weights live on a nonuniform grid, purely additive updates are insufficient—the optimizer needs scale-aware relative updates to make progress at large magnitudes.

Practical Benefits

1. Eliminates FP32 master weights: Enables true end-to-end low-precision training
2. No stochastic rounding needed: M+Adam trains stably without SR, simplifying implementation
3. Compatible with existing hardware: Operates within standard floating-point arithmetic (unlike log-based approaches)
4. Compressible state: Additional optimizer state can be low-rank approximated

Limitations

• Runtime overhead: ~6–9% at 1B scale from the extra multiplicative branch
• Simulation cost: FP8/FP4 storage currently simulated via explicit quantize/dequantize
• Theoretical scope: Analysis covers idealized setting without stochastic gradients or clipping
• Scale: Evaluated up to 1B parameters; larger-scale validation remains future work

---

Conclusion

M+Adam demonstrates that additive–multiplicative update geometry is particularly effective for low-precision training. By combining Adam-style local corrections with Madam-style relative scale changes, it addresses the fundamental mismatch between additive updates and nonuniform floating-point grids. The approach provides:

1. A principled theoretical framework (monotone descent guarantee)
2. Consistent empirical improvements across scales, budgets, and precisions
3. A practical path toward eliminating FP32 master weights entirely

Future directions include combining M+Adam with stochastic rounding or Kahan compensation, larger-scale validation, native low-precision kernel support, and extending the theoretical analysis to stochastic settings with adaptive moments.