新优化器AK-Momentum：用激活值控制动量衰减，67M模型预训练最多省46%步数
arXiv论文提出AK-Momentum，将线性层输入激活作为“键”更新动量缓冲，替代EMA方案；AK-AdamW在FineWeb-Edu上以更少步数达到AdamW验证损失。
AI解读：传统优化器（如AdamW）的动量是过去梯度的指数移动平均，所有方向按同一速率遗忘。但深度网络训练时输入往往高度各向异性——少数方向频繁出现，多数很少见。AK-Momentum的出发点是：梯度可拆成输入激活（键）和输出侧误差（值），按激活出现频率决定每个方向的遗忘速率，不改变动量更新本身以外的机制，可直接替换任何优化器的动量缓冲，无需额外矩阵求逆或持久内存。实测中，AK-AdamW在67M和370M规模的FineWeb-Edu预训练上，分别以最多约46%和22%更少的步数达到AdamW的验证损失，1B参数且在Chinchilla最优预算下增益仍存在；额外计算量仅占gated-MLP块线性成本的22.2%至25.0%。对研究者直接的意义是：在不改模型结构、不换优化器大框架的前提下，可能以更少算力达到同等损失，尤其适合预训练成本敏感的团队。需要自行复现验证，因为论文结果基于特定数据集和规模，尚未见独立第三方复现。
卡内基梅隆大学研究者Euijin Hong和Guannan Qu在arXiv发表论文，提出Activation-Keyed Momentum（AK-Momentum）优化器，它按输入激活的出现频率调整每个方向的动量遗忘速率，替代传统指数移动平均（EMA）机制。作者称AK-Momentum是“任何优化器的即插即用替换件”，在FineWeb-Edu预训练中，AK-AdamW达到AdamW验证损失所需的步数在67M规模最多减少46.39%±4.32%，370M规模减少22.12%±0.80%（三颗种子）。
论文v2版本将算法从DeltaMomentum更名为AK-Momentum，以避免与现有算法混淆。
方法：用激活值作为键的动量更新
大多数现代优化器将动量形成过往梯度的指数移动平均（EMA），以固定速率遗忘所有方向。论文指出，深度网络训练输入可能高度各向异性：少数方向被频繁查询，大多数很少出现。预处理方法（如自适应学习率）只在该缓冲外增加额外处理，不改变动量更新本身。
AK-Momentum将线性层梯度拆分为输入激活（作为键）和输出侧误差（作为值），基于激活来更新动量缓冲，采用delta规则：每个方向的遗忘速率由其出现频率决定。作者证明该更新是有效的动量，无需矩阵求逆即可实现输入侧曲率校正，并在固定或漂移的最优点下均比EMA更快清除过时方向。
AK-Momentum可替换任何优化器的动量缓冲，其系数在μP下可跨宽度迁移，额外计算量保持在gated-MLP块线性成本的22.2%至25.0%之间，且无持久内存。
实验：预训练与图像任务上的结果
在FineWeb-Edu预训练中，AK-AdamW（AdamW+AK-Momentum）达到AdamW验证损失的步数减少比例：67M规模最多46.39%±4.32%，370M规模22.12%±0.80%（基于三颗种子），1B规模在Chinchilla最优预算下增益持续。
同一协议下调优的Muon基线在两种语言模型规模上均不如AK-AdamW。在CIFAR-10上，增益在SGD、ResNet-18和ViT-Tiny上均保持。训练时诊断确认了预测机制、更好的梯度追踪和更健康的输入方向。
论文为arXiv预印本v2（2026年9月3日更新），作者未提供代码或完整实验细节的公开链接。