跳到正文
OpenReview · State space models· Hanyang Li; Shao Tang; Daniel Thomas Braithwaite; Gregory Dexter; Leonardo Neves; Aman Gupta; Hiroto Udagawa; Abhishek Shivanna; Daniel Silva; Rohan Ramanath·· 9 天前精选AI 评分79

在预条件器空间中舍入:重新设计 AdamW 的 4-bit 优化器状态量化

Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization

AI 导读

基于摘要分析。本文研究 AdamW 优化器状态量化误差如何经动量递推影响后续自适应更新,核心贡献是重新选择量化舍入所处的坐标空间,并提出 ZIP-SR 与 ZE-EDEN 两种 4-bit 状态量化方案。主要研究对象是 GPT 与 Llama 风格语言模型的预训练及全参数监督微调,而非 EEG 或 BCI。 机制上,作者对邻近零值的二阶矩量化单元进行局部分析,指出状态的平均误差较小,并不保证下一步预条件器的平均误差较小;一维二次函数构造进一步展示了状态空间舍入与预条件器空间舍入的不同优化动态。Zero-Inclusive Preconditioner-space Stochastic Rounding(ZIP-SR)在二阶矩码本中保留零值,但在预条件器空间计算随机舍入概率。Zero-Excluding EDEN calibration(ZE-EDEN)则使用排除零值的二阶矩码本,通过重新缩放量化后的二阶矩块,缓解正量化下限带来的预条件器失真。两种方案均对一阶矩采用 4-bit NormalFloat(NF4),并在训练最后 10% 对 LM-head 的一阶矩实施定向随机舍入。 作者报告,在参数规模为 130M 至 2.7B 的 GPT 与 Llama 风格预训练实验中,两种方法在每个已评估规模上都缩小了 TorchAO 4-bit AdamW 相对 32-bit AdamW 的验证损失差距,最大报告差距缩减为 70%;这不是验证损失本身下降 70%。在全参数监督微调中,作者报告两种方案的验证损失低于 TorchAO,且下游任务表现接近 32-bit AdamW。具体数据集、划分、任务指标、消融、统计信息以及实际显存与吞吐收益尚未验证。 平台推测(待验证):其可迁移机制是降低优化器状态存储,同时减少量化对自适应更新的扰动;假设 EEG 基础模型训练受到优化器状态显存约束,这类量化器与二阶矩块校准可能具有复用价值,但 LM-head 专用策略需按 EEG 输出头重新评估。低信噪比、跨被试梯度差异及小数据条件可能改变近零二阶矩的分布,使语言模型上的收益不再成立。可先在固定数据划分、训练预算和随机种子的 EEG 训练实验中,对比 32-bit AdamW、TorchAO 4-bit AdamW 与两种方案,核对任务指标、训练稳定性和峰值显存。已有 EEG 相关工作尚未检索确认。

阅读价值

值得阅读其将二阶矩量化的舍入坐标从状态空间转向预条件器空间的机制分析,以及相对 TorchAO 4-bit AdamW 的验证损失差距改善;能否用于 EEG 模型训练尚未验证。

来源:OpenReview · State space models · openreview.net