STQuant: Spatio-Temporal Adaptive Framework for Optimizer Quantization in Large Multimodal Model Training

📄 arXiv: 2604.06836v1 📥 PDF

作者: Minglu Liu, Cunchen Hu, Liangliang Xu, Fengming Tang, Ruijia Wang, Fu Yu

分类: cs.LG

发布日期: 2026-04-08


💡 一句话要点

提出STQuant框架,通过时空自适应优化器量化降低大模型训练内存占用。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 优化器量化 大模型训练 多模态模型 动态精度 分布式训练

📋 核心要点

  1. 现有优化器量化方法采用固定精度策略,忽略了优化器状态分布在不同层和训练步骤中的差异性,导致模型精度下降。
  2. STQuant框架通过动态调整优化器状态的量化精度,在不同层、状态变量和训练步骤上进行自适应分配,从而降低内存占用。
  3. 实验表明,STQuant能显著降低优化器状态内存占用,在GPT-2和ViT上实现了优于现有方法的压缩效果,且计算开销和额外空间需求很小。

📝 摘要(中文)

本文提出STQuant,一个分布式训练框架,旨在通过跨层、状态变量和训练步骤的动态精度分配,减少优化器状态的内存占用,同时保持模型质量。现有方法采用固定精度策略,忽略了优化器状态分布在不同层和训练步骤中显著变化的事实,导致精度下降。STQuant通过两种关键技术解决动态量化带来的挑战:1) 一种可证明的近优因子选择策略,准确识别用于精度调整的最具影响力的因子;2) 一种动态转换决策算法,将搜索成本从指数级降低到线性复杂度。在GPT-2和ViT上的实验表明,与现有解决方案相比,STQuant将优化器状态内存减少了84.4%,平均位宽低至5.1位,且计算开销仅为O(N/K),额外空间需求为O(1)。

🔬 方法详解

问题定义:论文旨在解决大规模多模态模型训练中,优化器状态占用大量内存的问题。现有方法采用固定精度的量化策略,无法适应不同层、不同状态变量以及不同训练阶段优化器状态分布的差异,导致量化噪声引入,影响模型精度。

核心思路:STQuant的核心思路是根据优化器状态在不同层、不同状态变量和不同训练步骤中的重要性,动态地调整量化精度。对影响较大的因子分配更高的精度,对影响较小的因子分配较低的精度,从而在保证模型精度的前提下,最大限度地降低内存占用。

技术框架:STQuant是一个分布式训练框架,主要包含两个关键模块:1) 近优因子选择策略:用于准确识别对模型精度影响最大的因子,以便进行精度调整。2) 动态转换决策算法:用于降低动态量化策略搜索的复杂度,将指数级搜索空间降低到线性复杂度。整体流程包括:首先,使用近优因子选择策略确定需要调整精度的因子;然后,使用动态转换决策算法确定每个因子的最佳量化精度;最后,根据确定的精度进行量化,并更新优化器状态。

关键创新:STQuant的关键创新在于提出了可证明的近优因子选择策略和动态转换决策算法。近优因子选择策略能够准确地识别对模型精度影响最大的因子,避免了盲目调整精度带来的负面影响。动态转换决策算法能够有效地降低动态量化策略搜索的复杂度,使得在实际训练中应用动态量化成为可能。

关键设计:STQuant的关键设计包括:1) 近优因子选择策略的具体实现,例如基于梯度信息的因子重要性评估方法。2) 动态转换决策算法的具体实现,例如基于动态规划的搜索算法。3) 量化和反量化的具体实现,例如采用均匀量化或非均匀量化方法。论文中可能还涉及一些超参数的设置,例如精度调整的频率、量化位宽的范围等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

STQuant在GPT-2和ViT模型上进行了实验,结果表明,与现有方法相比,STQuant能够将优化器状态内存减少84.4%,平均位宽降低至5.1位。同时,STQuant仅引入O(N/K)的计算开销和O(1)的额外空间需求,表明其具有较高的效率和实用性。

🎯 应用场景

STQuant可应用于各种大规模多模态模型的训练,尤其是在资源受限的环境下,如移动设备或边缘计算平台。通过降低优化器状态的内存占用,STQuant能够支持更大规模的模型训练,并加速模型迭代过程。该研究对于推动大模型在实际场景中的应用具有重要意义。

📄 摘要(原文)

Quantization is an effective way to reduce the memory cost of large-scale model training. However, most existing methods adopt fixed-precision policies, which ignore the fact that optimizer-state distributions vary significantly across layers and training steps. Such uniform designs often introduce noticeable accuracy degradation. To move beyond fixed quantization, we propose STQuant, a distributed training framework that reduces the memory footprint of optimizer states via dynamic precision allocation across layers, state variables, and training steps, while maintaining model quality. Naively applying dynamic quantization during training is challenging for two reasons. First, optimizer states are numerically sensitive, and quantization noise can destabilize quality. Second, jointly considering multiple states and layers induces a large combinatorial search space. STQuant addresses these challenges with two key techniques: 1) a provably near-optimal factor selection strategy that accurately identifies the most influential factors for precision adaptation. 2) a dynamic transition decision algorithm that reduces the search cost from exponential to linear complexity. Experiments on GPT-2 and ViT show that STQuant reduces optimizer-state memory by 84.4%, achieving an average bit-width of as low as 5.1 bits, compared with existing solutions. Moreover, STQuant incurs only O(N/K) computational overhead and requires O(1) extra space.