Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers
纠正预条件语言模型优化器中的随机更新偏差
🧭 预条件优化器偏差校正 ★ 深度解读
Fastino Labs
通过微批次交叉拟合与Delta方法逆修正，消除梯度-预条件子耦合和逆预条件有限样本偏差，提升AdamW、Sophia、Shampoo预训练效果
📌 工作参考价值（5 点）
- 可用于生成式推荐/大模型训练中的优化器改进：当使用AdamW、Sophia、Shampoo等自适应优化器时，将批次拆分为独立子批，分别估计梯度和预条件子，可移除耦合偏差；同时利用子批方差对逆预条件子做泰勒/Jackknife校正，降低逆运算的非线性偏差。代码级实现只需对现有优化器做包装（Algorithm 1），无需改动模型。
- 在噪声较大、预条件子估计不稳定的场景（如小批量、低EMA系数）效果更明显：电商推荐模型中若使用对角二阶矩（AdamW）或曲率估计（Sophia），可借鉴本文的Delta方法逆校正，公式为 \(\tilde{T}(\bar{p}_k) = 1/(\bar{p}_k + \lambda) - \widehat{\text{Var}}(\bar{p}_k)/(\bar{p}_k + \lambda)^3\)，仅需维护子批方差，计算开销极小。
- 矩阵预条件子（如Shampoo）的校正可在特征基下进行：对于推荐系统中可能的矩阵预条件（如K-FAC），可利用平均预条件子的特征分解，将每个子批投影到该基下，对特征值逐元素校正。避免反复特征分解，工程上可行。
- 在Agent多智体训练中，若每个Agent的本地更新可视为微批次，可利用方差校正来提升全局聚合的稳定性：将各Agent本地计算的预条件子统计量视为独立子批，使用Jackknife或方差校正再聚合，降低系统估计偏差，可能提升收敛速度和最终策略质量。
- 注意拆分方式要保留原优化器的分母规模：实验显示两折拆分（各用一半数据）会因分母样本减半而增大方差，推荐使用留一法（LOO）构造分母，保持分母接近全批量规模，从而在移除耦合的同时不显著增加噪声。