论文
随笔
Tags
Slides
⤴
Search
## Pushing the Limits of Low-Bit Optimizers with a Focus on EMA Dynamics
<h2 id="background">Background</h2> <div class="slide-highlight"> 模型飞速膨胀 vs. 吃紧的硬件设备 </div> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312203012.png" alt="Image" style="max-width: 65%; height: auto; margin: 0 auto;"> </div> <ul> <li>可能的一些解决方案: <ul> <li>MoE, LoRA; ZeRO, FSDP;</li> <li>Network Quantization; <span style="color: red;">Lightweight Optimizers</span></li> </ul> </li> </ul>
<h2 id="background-1">Background</h2> <p>⚙️ Optimizer States (2x model size):</p> $$ m_{t+1} \leftarrow \beta_1 \cdot m_t + (1 - \beta_1) \cdot g, \\ v_{t+1} \leftarrow \beta_2 \cdot v_t + (1 - \beta_2) \cdot g^2 $$<ul> <li>Lightweight Optimizers: <ul> <li><strong>重新设计:</strong> Lion, Muon …</li> <li><strong>状态共享:</strong> Adafactor, SM3, Adam-Mini …</li> <li><strong>降维/稀疏化:</strong> GaLore, MicroAdam</li> <li><strong>低精度:</strong> 1-bit SGD/Adam, <u>16/8/4-bit Optimizers</u>, Q-GaLore, 8-bit Muon</li> </ul> </li> </ul> <div class="slide-highlight"> Why Low-Bit Optimizers? </div>
<h2 id="why-low-bit-optimizers">Why Low-Bit Optimizers?</h2> <ul> <li> <p><strong>泛化性:</strong> ✅无需额外调参 ✅适用任意场景</p> </li> <li> <p><strong>灵活性:</strong> ✅非环境依赖</p> </li> <li> <p><strong>成功的工程实践:</strong> DeepSeek-v3 训练框架 ($g \overset{\text{BF16}}{\rightarrow} m,v \overset{\text{FP32}}{\rightarrow} \theta$)</p> </li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312204230.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">DeepSeek-AI. DeepSeek-V3 Technical Report, 2024.</p> </div>
<h2 id="quantization-and-dequantization">Quantization and Dequantization</h2> <ul> <li> <p>Quantization:</p> $$ q = Q(x) := \mathop{\text{argmin}} \limits_{k=0}^{2^b - 1} \big|\frac{x}{\textcolor{red}{\Delta}} - \textcolor{red}{y_k} \big|. $$</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312205652.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <ul> <li> <p>Dequantization:</p> $$ \tilde{x} = Q^{\dagger}(q) := y_{q} \cdot \Delta. $$<p>
<h2 id="stateful-optimizers-in-ultra-low-bits"><u>S</u>tateful <u>O</u>ptimizers in Ultra-<u>LO</u>w Bits</h2> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617110019.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h2 id="challenges-in-ultra-low-bit-cases">Challenges in Ultra-Low-Bit Cases</h2> <ul> <li> <p><strong>表示精度:</strong> 42 亿 (32-bit) vs. 8 (3-bit) vs. 4 (2-bit)</p> </li> <li> <p><strong>量化范围:</strong> 如何将尽可能多的元素一起量化?</p> </li> <li> <p><strong>一阶/二阶动量:</strong></p> <ul> <li>(Signed) 一阶动量 ($m$): 决定参数更新方向</li> <li>(Unsigned) 一阶动量 ($m$): 决定参数更新步长</li> </ul> </li> </ul> <div class='slide-highlight'> 关键: EMA Dynamics </div>
<h2 id="quantization-for-unsigned-ema-update">Quantization for Unsigned EMA Update</h2> <ul> <li><em>Signal Swamping</em> (<u>large-to-small number addition</u>)</li> </ul> $$ \text{EMA update: } \hat{x}_{t+1} \leftarrow \beta \cdot \tilde{x}_t + \underbrace{\textcolor{red}{(1 - \beta) \cdot z_{t + 1}}}_{\text{very small as } \beta \rightarrow 1}. $$<div style="text-align: center; margin-top: 50px; margin-bottom: -80px; padding: 0"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250722100933.png" alt="Image" style="max-width: 70%; height: auto;margin: 0 auto;"> </div> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Higham N. J. The Accuracy of Floating Point Summation. SIAM Journal on Scientific Computing. 1993.</p> </div>
<h2 id="signal-swamping">Signal Swamping</h2> <p>💡 总结</p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250312212039.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <div class='slide-highlight'> <p>❎Unsigned ❎ $\beta \uparrow$ ❎ $b \downarrow$</p> </div>
<h2 id="case-study">Case Study</h2> <div class="slide-cols"> <div class="slide-col-6"> <p> </p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250722101123.png" alt="Image" style="max-width: 90%; height: auto;margin: 0 auto;"> </div> </div> <div class="slide-col-4"> <p> </p> <ul> <li> <p>一定<span style="color: red">条件</span>下:</p> <ul> <li>Linear 下全部不更新</li> <li>DE 下部分更新</li> </ul> </li> <li> <p>实际上 $\beta \ge 0.9$ 为<span style="color: red">相当常见的 setting</span></p> </li> </ul> </div> </div>
<h2 id="case-study-1">Case Study</h2> <div class="slide-cols"> <div class="slide-col-4"> <p> </p> <ul> <li> <p>随机信号:</p> <ul> <li>$X \in \mathbb{R}^{1000}$</li> <li>$Z \sim \mathcal{U}[0, 1]$</li> </ul> </li> <li> <p>Relaxed 条件:</p> </li> </ul> <p> <span style="color: red;"><em>X</em></span> Fixed $\Delta$</p> <p> <span style="color: red;"><em>X</em></span> $z \le \Delta$</p> <ul> <li>理论收敛至: $0.5$</li> </ul> </div> <div class="slide-col-6"> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250908140511.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> </div> </div>
<h2 id="solution-12-stochastic-rounding">Solution (1/2): Stochastic Rounding</h2> <ul> <li> <p>假设 $y_{k-1} \le x / \Delta \le y_k$:</p> $$ Q_{sr}(x) := \left \{ \begin{array}{ll} k-1 & w.p. \quad \frac{y_k - x / \Delta}{ y_k - y_{k-1}}, \\ k & w.p. \quad \frac{x / \Delta - y_{k-1}}{ y_k - y_{k-1}}. \end{array} \right . $$</li> <li> <p>High variance:</p> </li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250722101414.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h2 id="solution-22-logarithmic-quantization">(Solution 2/2) Logarithmic Quantization</h2> <div class='slide-highlight'> $1 \overset{\text{more levels}}{\Longrightarrow} 0$ </div> <ul> <li>3-bit quantization levels (Linear vs. Dynamic Exponent vs. Ours):</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20251210214923.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div>
<h2 id="logarithmic-quantization">Logarithmic Quantization</h2> <ul> <li>2-bit quantization illustration</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250313113535.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h2 id="logarithmic-quantization-1">Logarithmic Quantization</h2> <p>✅ Easy to implement</p> <p>✅ State decay alignment</p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250313115306.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h2 id="quantization-for-signed-ema-update">Quantization for Signed EMA Update</h2> <p>😄 <span style="color: gray">No Signal Swamping</span></p> <p>😞 <strong>额外的符号表示 (1 bit)</strong></p> <p>😞 <strong>直接决定更新方向 (误差敏感)</strong></p> <p>💡 总结:</p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250314115701.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h2 id="quantization-errors--gradient-variance">Quantization Errors $\Rightarrow$ Gradient Variance</h2> <div class="slide-cols"> <div class="slide-col-6"> <p> </p> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250722101739.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div> </div> <div class="slide-col-4"> <p> </p> <p>$\rightarrow$ <span style="color: red">Bits $\downarrow$</span> or <span style="color: red">$\beta \uparrow$ </span></p> <p>$\rightarrow$ Quantization errors <span style="color: red">$\uparrow$</span></p> <p>$\rightarrow$ gradient variance <span style="color: red"> $\uparrow$ </span></p> <p>$\rightarrow$ <span style="color: red"> worse </span> convergence</p> </div> </div> <div class='slide-highlight'> 不稳定性难以在量化算法层面避免! </div> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Li H., et al. Convergence of Adam under Relaxed Assumptions. NeurIPS, 2023.</p> </div>
<h2 id="momentum-adjustment">Momentum Adjustment</h2> <ul> <li><strong>方差控制:</strong> 选择 $\beta'$ 满足:</li> </ul> $$ \underbrace{\frac{\textcolor{gray}{\beta'}}{1 - \textcolor{gray}{\beta'}} r_{\text{median}}(b')}_{\textcolor{gray}{\text{undetermined}}} \le \underbrace{\frac{\beta}{1 - \beta} r_{\text{median}}(b)}_{\textcolor{green}{\text{valid setup}}}. $$<ul> <li><strong>查表:</strong> (<u>灰色区域代表了经验可行的参数推荐</u>)</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250314121510.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h2 id="experiments">Experiments</h2> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617111401.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div> <p>😒 <strong>传统方法:</strong> $\underset{\text{Training from scratch}}{\xrightarrow{\text{Ultra-Low-Bit}}}$ degeneration/collapse</p> <p>😊 <strong>SOLO:</strong> Robust to bits/tasks/models</p>
<h2 id="experiments-giant-models">Experiments (Giant Models)</h2> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617112643.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h2 id="loss">Loss</h2> <ul> <li>损失正常收敛</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250319170139.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h2 id="quantile">Quantile $x_p$</h2> <ul> <li>基本上 $p \in [0.05, 0.3]$ 都有不错的性能</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250319170604.png" alt="Image" style="max-width: 55%; height: auto;margin: 0 auto;"> </div>
<h2 id="beta-block-size">Beta, Block size</h2> <ul> <li>Lower-bit SOLO needs a smaller $\beta$</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250407200935.png" alt="Image" style="max-width: 95%; height: auto;margin: 0 auto;"> </div>
<h2 id="state-changes">State Changes</h2> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617112843.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h2 id="generalizability-of-solo">Generalizability of SOLO</h2> <ul> <li>AdaBelief</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617113113.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <ul> <li>Larger-scale models:</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250617113029.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div>
Thanks!