跳到正文
OpenReview · State space models· Yiming Dong; Zhouchen Lin·· 2024-01-01精选AI 评分72

通过梯度空间复用降低深度网络训练的内存占用

Reducing Memory Footprint in Deep Network Training by Gradient Space Reutilization

AI 导读

基于摘要分析。本文研究深度网络训练中的内存占用问题,提出 gradient space reutilization(梯度空间复用):在梯度不再被后续计算需要时,重新利用其已分配的内存空间,以减少训练内存需求。主要贡献是优化训练过程的内存管理,而非提出新的模型结构或学习目标。 作者将该机制应用于 AdamW、Adan 和 Lion,形成 AdamW-R、Adan-R 和 Lion-R。作者称这些改动与原优化算法等价,因此能够节省内存而不影响性能;具体复用对象、梯度生命周期判断、更新顺序及等价性论证尚需正文核对,不能仅凭摘要认定在所有实现和训练配置下均成立。 作者报告,在其训练实验中,最佳情况下内存节省达到 25.60%。摘要未提供对应模型、数据集、训练配置、对照实现及内存统计口径,因此该数值只能视为所述实验中的最佳结果,不能推断为各优化器或各任务的普遍收益。实验协议、消融及统计信息尚未验证,代码与复现条件也尚未验证。 平台推测(待验证):若 EEG 模型训练受优化阶段内存占用限制,该机制可能作为优化器实现层面的资源优化使用,而非直接改善 EEG 表征或解码性能。其依赖的是训练中的内存生命周期,不依赖特定 EEG 监督标签;可尝试复用内存管理思路,但需适配具体框架和训练配置。对于小模型、小数据或内存主要由激活占用的任务,收益可能有限,且不能解决低信噪比、跨被试差异或通道不一致。一个可检验的小实验是在固定 EEG 数据划分、模型、随机种子及数值精度下,对比 AdamW 与 AdamW-R 的峰值内存、单步耗时和参数更新一致性,并核对解码指标是否保持一致。已有 EEG 相关应用尚未检索确认。

阅读价值

值得核对其梯度内存复用的实现及算法等价性依据:作者报告该工程优化可降低训练内存占用,但具体测量协议与适用条件尚未验证。

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