9 papers
Beyond Shadow Weights: Quantization-Aware Training as Quantized-Endpoint Descent
Sheng-An Xu, Hanyang Li, Jianhao Ma +1
Quantization-aware training (QAT) updates a full-precision shadow weight but deploys the quantized endpoint . Existing explanations for QAT largely view…
A Lower Bound for the Heavy-Ball Method on Smooth Convex Functions
Jianhao Ma, Jingzhao Zhang
Can the classical Heavy-Ball method, with arbitrary horizon-dependent parameters chosen in advance, achieve Nesterov's last-iterate rate on every smooth convex objectiv…
A lower bound for stepsize-based acceleration of gradient descent
Jianhao Ma, Yuxin Chen
Recent work has shown that, for smooth convex optimization, plain gradient descent can be accelerated from its textbook convergence rate of (where denotes the numbe…
WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training
Jianhao Ma, Yuxin Chen
Standard learning rate schedules such as cosine annealing are tied to a fixed training horizon, limiting their ability to accommodate post hoc horizon extension. Warmup-stable-deca…
Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective
Tianqi Shen, Jinji Yang, Runze Shi +3
Recently, Muon has gained substantial attention as an appealing alternative to Adam-like optimizers, with many works highlighting its advantages through spectral normalization and…
Preconditioning Benefits of Spectral Orthogonalization in Muon
Jianhao Ma, Yu Huang, Yuejie Chi +1
The Muon optimizer, a matrix-structured algorithm that leverages spectral orthogonalization of gradients, is a milestone in the pretraining of large language models. However, the u…
LIBRA: Language Model Informed Bandit Recourse Algorithm for Personalized Treatment Planning
Junyu Cao, Ruijiang Gao, Esmaeil Keyvanshokooh +1
We introduce a unified framework that seamlessly integrates algorithmic recourse, contextual bandits, and large language models (LLMs) to support sequential decision-making in high…
Quantization through Piecewise-Affine Regularization: Optimization and Statistical Guarantees
Jianhao Ma, Lin Xiao
Optimization problems over discrete or quantized variables are very challenging in general due to the combinatorial nature of their search space. Piecewise-affine regularization (P…
PARQ: Piecewise-Affine Regularized Quantization
Lisa Jin, Jianhao Ma, Zechun Liu +3
We develop a principled method for quantization-aware training (QAT) of large-scale machine learning models. Specifically, we show that convex, piecewise-affine regularization (PAR…