Adam优化器状态压缩新法：共享梯度界守护更新精度

训练大模型时，Adam优化器要为每个参数维护两个FP32矩（一阶和二阶动量），其存储开销是FP32权重本身的两倍。压缩这些矩看似能节省大量显存和存储，但风险很高：Adam的更新规则中要用二阶矩的平方根做除法，微小的重构误差可能被放大，显著改变恢复后的计算行为。近期一项研究提出“守护式矩压缩”（Guarded Moment Compression），通过共享梯度界来保证压缩后解码的矩能够复现原生FP32的更新，为优化器状态压缩提供了新的思路。

为什么压缩Adam矩很危险？

Adam优化器的每个参数对应两个FP32矩：一阶矩估计梯度的均值，二阶矩估计梯度的平方均值。更新时，参数更新量等于学习率乘以一阶矩除以二阶矩的平方根（加上小常数）。这种除法结构意味着二阶矩的微小误差会直接导致更新量产生较大偏差。更具体地说，如果二阶矩被低估，更新量会被放大；如果被高估，更新量会被缩小。在极端情况下，压缩误差可能使训练发散。

现有压缩方法（如块量化、BF16存储）通常只关注矩的存储精度，而忽略了压缩后更新是否仍与原始更新一致。这项研究则提出了一个更严格的目标：对于每个坐标，只有当解码后的矩能够在给定的梯度域内复现原生FP32 Adam更新（误差在容差τ内，单位为学习率）时，才接受该压缩；否则，该坐标精确存储。

Adam优化器中的一阶矩和二阶矩本质上是梯度及其平方的指数移动平均（EMA），其更新公式为：m_t = β₁·m_{t-1} + (1-β₁)·g_t，v_t = β₂·v_{t-1} + (1-β₂)·g_t²，其中g_t为当前梯度，β₁和β₂为衰减率（通常为0.9和0.999）。由于训练初期矩估计偏向零，Adam还引入了偏差校正：m̂_t = m_t/(1-β₁^t)，v̂_t = v_t/(1-β₂^t)，最终参数更新为θ_t = θ_{t-1} - η·m̂_t/(√v̂_t + ε)。这种结构使得二阶矩v_t的平方根直接出现在分母中，任何压缩导致的v_t误差都会被非线性放大——例如，若v_t被低估10%，则√v_t约被低估5%，更新量相应放大约5.3%。此外，矩的存储通常与模型参数同精度（FP32），在数十亿参数模型中，矩的显存占用可达权重的两倍，成为训练时的关键瓶颈。因此，压缩矩不仅关乎存储效率，更需确保更新行为的数值一致性。

自然证书为何失效？

直观上，可以为每个更新单独计算一个证书：让梯度在区间内变化，检查更新量的范围是否在容差内。但研究者发现，这种“自然证书”在100万个训练过的OPT-125M坐标上严重失效：98.9%的坐标无法通过验证，导致负载增长到原始FP32的126.6%。原因是每个更新单独考虑时，梯度区间被过度放大，因为两个更新其实共享同一个梯度。

共享梯度界：关键观察

研究者的核心观察是：一阶矩和二阶矩的更新消耗的是同一个梯度和同一个参数。与其分别界定两个更新，不如直接界定它们之间的差异。具体来说，他们使用FP32舍入误差允许量和一个固定的混合线性-对数梯度划分，直接计算两个更新之差的边界。这样做的效果非常显著：在τ=0.01时，逃逸坐标比例从98.9%降至0.53%，负载从126.6%降至28.3%。

编码与恢复流程

守护式压缩的流程包括两个阶段：编码时，对每个矩使用8位对数编码，并检查解码后的更新是否在容差内；如果不在，则存储原始FP32值。恢复时，解码矩并计算更新。整个过程保证了恢复后的第一次更新与原始更新在容差内一致。需要注意的是，该保证覆盖的是在固定的Torch操作模型下的第一次恢复更新；后续训练行为则通过实验测量。

“自然证书”失效的根本原因在于它独立处理一阶矩和二阶矩的更新，而实际上这两个更新共享同一个梯度样本。在Adam中，一阶矩和二阶矩的更新依赖于同一梯度g_t，但自然证书将g_t视为区间内的任意值，并分别计算两个更新的可能范围，导致区间被过度保守地放大。共享梯度界则利用这一关联：它直接考虑两个更新之差（即一阶矩更新与二阶矩更新之间的差异）作为整体进行界定，并引入FP32舍入误差允许量以及混合线性-对数梯度划分来收紧边界。混合线性-对数划分是指将梯度域划分为线性区间（用于小梯度）和对数区间（用于大梯度），以更精确地覆盖实际梯度分布。通过这种联合界定，证书能够更准确地识别哪些坐标的压缩是安全的，从而大幅降低逃逸率。

8位对数编码是一种将浮点数转换为8位对数值的压缩格式，通常包含符号位、指数位和尾数位。例如，可采用1位符号、4位指数和3位尾数的配置，以较少的比特数表示较宽的动态范围，但精度相对较低。在守护式压缩中，对每个矩独立进行8位对数编码后，立即解码并计算更新量，并与原始FP32更新比较。若误差超过容差τ（以学习率为单位），则将该坐标的矩回退为精确FP32存储。这种“编码-验证-回退”的机制确保了每个坐标的压缩安全性。恢复时，只需解码存储的矩（8位或FP32）并执行一次Adam更新步骤，即可保证与原始更新的偏差在τ内。该过程不依赖于未来梯度，因此证书是静态可计算的。

实验验证：压缩率与精度

研究者在OPT-125M和ResNet-18上进行了广泛实验，包括120次完整检查点恢复，使用四个训练种子和每个模型三个训练比例。结果显示，守护式8位对数编码相对于FP32将矩负载分别减少69.5%和71.9%，相对于同等守护的BF16减少40.8%和45.4%，相对于同等守护的块状8位减少29.3%和37.6%。所有恢复的第一次更新都保持在容差内。

编码和恢复的总时间比守护式BF16长2.3-2.6倍，但比守护式块状8位更短。此外，联合搜索两个矩的编码在码率上仅多节省不到0.001%，但编解码时间增加6%，因此逐矩舍入已经足够。

对AI训练的启示

这项研究为优化器状态压缩提供了一个有精度保证的框架。随着模型规模增长，优化器状态成为显存瓶颈之一，而守护式压缩能够在保持训练稳定性的前提下显著降低矩的存储开销。不过，当前保证仅覆盖第一次恢复更新，后续训练（如多步更新）的稳定性仍依赖经验验证。未来工作可能需要扩展守护范围到多步更新，或探索更高效的编码方案。

对于实践者而言，该方法的核心价值在于：压缩不再是盲目追求比特数，而是以更新容差为约束，确保压缩后的训练与原生训练在第一步就对齐。这为安全地压缩Adam状态提供了可量化的依据。