纠正预训练优化器的随机更新偏差，提升语言模型训练效果。
Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers
- 用独立小批量交叉估计梯度与预条件矩阵，消除耦合偏差。
- 通过微批量方差修正逆运算偏差，使优化更精准。
- 适用于AdamW、Sophia等主流优化器，实测性能提升显著。
预条件优化器在语言模型训练中至关重要，但其随机更新规则常被视为总体预条件下降的直接近似。我们指出这种观点忽略了两个有限样本偏差：一是梯度与预条件矩阵通常来自同一小批量，引入梯度-预条件耦合偏差；二是即使预条件估计无偏，其逆或逆根因非线性运算仍存在偏差。为此提出单批次偏差校正框架：交叉拟合预条件从独立微批量组中估计分子与预条件矩阵；方差修正反演利用微批量变异性减去主导的δ方法偏差项。该框架适用于对角动量、对角曲率及矩阵预条件方法，分别在AdamW、Sophia和Shampoo中实现。在Qwen2.5-0.5B上，偏差校正使保留损失分别降低0.15、0.07和0.11纳特；在混合质量预训练与下游指令微调中效果均呈中性至正向。结果表明，偏差校正是降低有限样本更新偏差、提升预条件优化器性能的实用机制。
原文摘要 · Abstract (English)
Preconditioned optimizers are central to language model training, but their stochastic update rules are usually treated as direct approximations to population preconditioned descent. We show that this view misses two finite-sample biases. First, the gradient and preconditioner are typically estimated from the same minibatch, introducing gradient--preconditioner coupling bias. Second, even when the preconditioner estimate is unbiased, its inverse or inverse-root is generally biased because inversion is nonlinear. We propose a single-batch bias-correction framework that addresses both effects: cross-fitted preconditioning estimates the numerator and preconditioner from independent microbatch groups, while variance-corrected inversion uses microbatch variability to subtract the leading delta-method bias term. The framework applies to diagonal moment, diagonal curvature, and matrix preconditioning methods, instantiated in AdamW, Sophia, and Shampoo. Bias correction reduces held-out pretraining loss on Qwen2.5-0.5B by $0.15$, $0.07$, and $0.11$ nats, respectively; the effects on mixed-quality pretraining and downstream instruction tuning are consistently neutral-to-positive. Together, these results establish bias correction as a practical mechanism for reducing finite-sample update bias and improving the performance of preconditioned optimizers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。