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

Minglu Liu; Cunchen Hu; Liangliang Xu; Fengming Tang; Ruijia Wang; Fu Yu

cs.LGApr 8, 2026

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

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

Organizations: Xidian University Xian, Shaanxi, China · China Telecom Cloud Computing Research Institute Beijing, China

Abstract

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.

Figures & tables

Figure 1. Spatiotemporal evolution and correlation analysis of gradients and Adam optimizer states. (a) Sharpness of the gradient distribution (CV) across 48 core weight layers in 12 Transformer blocks, including QKV projections, attention projections, MLP expansion layers, and MLP compression layers, sampled every 50 training steps over 5000 total steps; (b) Correlation between gradients and the coefficient of variation (CV) of the first-order moment $m$ , showing a very strong linear relationship (Pearson $R=0.9777$ ); (c) Correlation between gradients and the second-order moment $v$ , where, despite the squaring operation introducing nonlinear mapping and noise amplification, the correlation remains strong ( $R=0.7169$ ).

Figure 2. Overview of the STQuant framework: The system consists of three collaborative engines: (1) the Score Engine extracts spatio-temporal gradient features across GPUs; (2) the Distributed Engine synchronizes global statistics to determine optimal bit-widths; and (3) the Quantization Engine executes dual-mode block-wise compression (for $m$ and $v$ ).

Figure 3. The $n$ - $r$ decision quadrants for bit-width allocation. The decision space is partitioned into four zones based on global EMA statistics ( $N_{\text{ema}}$ and $R_{\text{ema}}$ ): (1) Critical Zone (top-right), (2) Magnitude-Dominant Zone (top-left), (3) Structural Complexity Zone (bottom-right), and (4) Redundant Zone (bottom-left).

Stage	Task	Data	Model	Metric
Pre-train	LM	OpenWebText	GPT2-XL	Loss
Pre-train	VRecon	COCO 2017	ViT-Base	Loss / Top-1
Fine-tune	NLU	MNLI	RoBERTa-Large	Accuracy
Fine-tune	Vision	COCO 2017	ViT-Base	mAP
Fine-tune	LM	Wikitext-103	GPT2-Medium	PPL
Fine-tune	ITR	COCO 2017	ViT-B/32	Recall@1

Table 1. Overview of experimental settings.

Figure 4. Pre-training on GPT2-1.5B (XL) and ViT-Base.

Figure 5. Comparison of optimizer-state memory on GPT2-XL and ViT-Base. STQuant achieves the lowest memory overhead among all compared optimizers.

Panel of Figure 6 — Figure 6. Bit-width evolution of STQuant during pre-training. (a) shows the dynamic bit-width allocation across layers over training epochs, and (b) presents the bit-width evolution of different parameter groups within a Transformer block.

Optimizer	State	Avg.	Top-1 (%)
32-bit AdamW	32-bit	32.0	31.44
8-bit AdamW (bnb)	8-bit	8.0	31.68
STQuant	Dynamic	5.1	31.64

Table 2. Linear probing of ViT-Base with different optimizers.

	Task Setting
Optimizer	Task	Data	Model	Metric	Value	Opt.State Mem Saved
32-bit AdamW	GLUE	MNLI	RoBERTa-Large	Accuracy( $\uparrow$ )	0.9060	reference
8-bit AdamW(bnb)	GLUE	MNLI	RoBERTa-Large	Accuracy( $\uparrow$ )	0.9002	75.0%
STQuant	GLUE	MNLI	RoBERTa-Large	Accuracy( $\uparrow$ )	0.9032	84.4%
32-bit AdamW	Classification	COCO 2017	ViT-Base	mAP( $\uparrow$ )	0.72669	reference
8-bit AdamW(bnb)	Classification	COCO 2017	ViT-Base	mAP( $\uparrow$ )	0.73082	75.0%

Table 3. Performance and optimizer-state memory savings of STQuant AdamW compared with full-precision and 8-bit AdamW(bnb) baselines across core fine-tuning tasks.

Method	AvgBit $\downarrow$	$\Delta$ AvgBit	PPL $\downarrow$	$\Delta$ PPL
STQuant (Full)	6.3	0.0	127.3	0.0
w/o Dual Factor	6.4	+0.1	130.2	+2.9
w/o Temporal Factor	5.8	-0.5	128.7	+1.4
w/o Spatial Factor	8.0	+1.7	125.6	-1.7

Table 4. Ablation study of STQuant on WikiText-103 with GPT-2. Lower AvgBit and PPL indicate better compression and language modeling performance, respectively.

Neural Network Quantization

Neural network quantization represents weights, activations, and states with fewer bits to reduce memory use, computation, and deployment costs.

See all 394 papers

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers
Aug 23, 2026Yan WangDeep Learning Optimization Memory-Efficient Optimization
Independent Researcher
Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets
Mar 17, 2026Kristi Topollai, Anna ChoromanskaLow-Precision Training Low-Bit Quantization
New York University
STaR-Quant: State-Time Consistent Post-Training Quantization for Diffusion Large Language Models
Jun 3, 2026Xin Yan, Aqiang Wang, Zhenglin Wan +2LLM Quantization Diffusion Model Quantization
School of Artificial Intelligence, Beijing Normal University, Beijing, China · Department of Computer Science, National University of Singapore, Singapore · Centre for Frontier AI Research, Agency for Science, Technology and Research (A*STAR), Singapore

Aug 23, 2026J/K move · Enter open · S save

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

Yan Wang

Independent Researcher

Deep Learning Optimization Memory-Efficient Optimization Low-Bit Quantization Neural Network Quantization

Open paper AskPDF arXiv

Optimizer-state quantization is commonly designed for Adam's dense, parameter-aligned first- and second-moment arrays. This abstraction breaks for memory-efficient optimizers, whose states may be factored, confidence-modulated, or maintained in a projected space, so similar reconstruction error can produce different update error. We formulate optimizer-state quantization as a joint problem over representation, topology, and update semantics. We then introduce Adaptive Log-Space (AL) quantization for non-negative states. AL fits each block's observed nonzero logarithmic interval and reserves a separate code for exact zero, enforcing

q = 0 \Leftrightarrow x = 0

; signed momentum and state precision remain independently selectable. Controlled probes show that adaptive ranges reduce update error and temporal drift, exact-zero reservation preserves dormant states, and state topology constrains useful block granularity. End-to-end language-model training evaluates the resulting policy across dense, factored, confidence, and projected optimizer states. On TinyLlama-1.1B, AL8 with uniform 8-bit momentum reaches 72.90 perplexity versus 73.54 for bitsandbytes 8-bit AdamW, with comparable optimizer-state storage and higher throughput. CAME matches reference-level final perplexity across three seeds when its non-negative states use AL16, while a semantic grouping-and-protection policy closes most of quantized Adafactor's 100K-step late-loss gap. These results make state topology and update semantics first-class design constraints for optimizer quantization.

Abstract

Neural Network Quantization

Explore similar work

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets

STaR-Quant: State-Time Consistent Post-Training Quantization for Diffusion Large Language Models