梯度检查点是一种用于深度学习的技术，旨在减少训练神经网络时的内存占用。在标准反向传播过程中，网络必须存储前向传播中计算的所有中间激活值，以便在反向传播中计算梯度。对于非常深的模型，如大型语言模型和变换器，这种存储可能超出可用硬件的内存容量。梯度检查点通过不存储所有激活值来解决此问题；它仅保留子集，并在反向传播过程中按需重新计算被丢弃的激活值。这以增加计算成本换取显著降低的内存使用，从而能够在相同硬件上训练更大的模型或使用更大的批次大小。
该技术由卡内基梅隆大学和OpenAI的研究人员于2016年在一篇题为“Training Deep Nets with Sublinear Memory Cost”的论文中提出。作者包括Tianqi Chen、Bing Xu、Chiyuan Zhang和Carlos Guestrin，他们证明通过仅在特定检查点（例如每隔几层）存储激活值并重新计算其余部分，对于一个具有n层的网络，训练的内存成本可以从O(n)降低到O(sqrt(n))，代价是大约额外进行一次前向传播。这项基础性工作自此成为机器学习社区的标准工具，尤其是在模型规模大幅增长的背景下。
标准反向传播如何使用内存
在传统的训练循环中，前向传播计算网络每一层的激活值。这些激活值被存储在内存中，因为反向传播需要通过链式法则计算梯度时用到它们。对于一个具有L层的网络，这需要存储L组激活值，每组可能很大。例如，一个具有数百层的Residual Network (ResNet)或具有数十个注意力块的transformer，对于单个训练样本可能积累数吉字节的激活数据。当使用大批次训练时，内存需求随批次大小线性增长，通常成为主要瓶颈。
检查点策略
梯度检查点将网络划分为若干段，每段边界设有一个检查点。在前向传播期间，仅将检查点边界处的激活值保存到内存中。段内的所有其他中间激活值都被丢弃。当反向传播到达某段时，它使用保存的检查点激活值重新计算该段的前向传播，重新生成计算梯度所需的中间激活值。这种重计算增加了计算开销，通常相当于每个训练步骤额外一次前向传播，但显著降低了峰值内存使用量。
检查点的选择是一个权衡。更多的检查点意味着更少的重计算但更高的内存使用；更少的检查点意味着更低的内存但更多计算。对于一个有n层的网络，最佳检查点数量约为sqrt(n)，这平衡了内存和计算。在实践中，像PyTorch和TensorFlow这样的框架允许用户指定检查点间隔或使用自动启发式方法。
变体与改进
对原始技术进行了若干改进。一种常见的变体是选择性检查点，只对某些层类型（如注意力块或卷积层）进行检查点，而其他层则正常存储。另一种方法，称为内存高效梯度检查点，使用更复杂的调度方案，在多粒度级别存储激活值，以额外的重计算为代价进一步减少内存。一些框架还实现了“卸载”，将检查点移动到CPU内存或磁盘，但这会引入数据传输开销。
在Transformer (architecture)模型的背景下，梯度检查点通常与其他节省内存的技术结合使用，如Mixed Precision Training和Gradient Clipping。例如，训练像GPT-3这样具有1750亿参数的模型，如果没有这些优化将不可能实现。该技术也用于微调大型模型，内存节省使从业者能在单个GPU而不是集群上运行。
实际实现
在现代深度学习框架中，梯度检查点通常以简单的API形式提供。例如，在PyTorch中，torch.utils.checkpoint模块提供了一个checkpoint函数，用于包装模块或一系列操作。当被包装的模块执行时，其激活值不会被保存；相反，它们会在反向传播期间被重新计算。TensorFlow通过tf.recompute_grad提供了类似的功能。这些实现自动处理簿记工作，使研究人员无需修改模型架构即可轻松采用该技术。
梯度检查点的计算开销并非微不足道。对于一个具有sqrt(n)个检查点的网络，训练期间的总前向计算比标准训练增加约30-40%。然而，这种成本通常可以接受，因为替代方案，，减小批大小或模型大小，，可能会损害收敛性或模型质量。在许多情况下，使用更大批大小带来的速度提升超过了重新计算的开销。
对大型模型训练的影响
梯度检查点已成为训练超大规模模型的基石。像OpenAI、Anthropic和Google DeepMind这样的公司依赖它来训练具有数千亿参数的模型。例如，在单节点8个GPU上训练一个700亿参数的模型，如果没有检查点，需要存储的激活值将超过这些GPU的合计内存。通过使用梯度检查点，这些组织可以在可用硬件上完成训练任务，尽管训练时间较长。
该技术对于涉及长序列的生成式AI应用（如文档摘要或代码生成）也至关重要。在这些情况下，激活内存随序列长度增长，而检查点允许在不超过内存限制的情况下处理更长上下文。这直接推动了具有更大上下文窗口（如数千甚至数百万个标记）的模型的发展。
与其他内存优化的关系
梯度检查点常与其他技术结合使用。Batch Normalization和Layer Normalization不会直接减少内存，但它们可以改善训练稳定性，从而补充检查点的效果。模型剪枝减少参数数量，但激活值仍是瓶颈，因此仍需要检查点。数据增强增加有效数据集大小，但不影响激活内存。在分布式训练中，梯度检查点可以与流水线并行结合，其中网络层分配到不同设备，以进一步减少每设备内存压力。
一个值得注意的替代方案是梯度累积，它通过多个较小批次累积梯度来模拟大批次训练。这减少了优化器状态的内存，但不减少激活内存，因此不能替代检查点。另一个相关想法是可逆层，如某些残差网络中使用的，其中激活值可以从输出重构，但这需要对架构进行修改，且不如检查点通用。
局限性与权衡
梯度检查点的主要限制是每个训练步骤增加了墙钟时间。对于已经计算密集的模型，额外的前向传播会使训练时间增加约20-30%。此外，该技术不能减少模型参数或优化器状态的内存占用，而这些在大型模型中也可能很大。因此，检查点通常与其他技术（如模型并行和分布式训练）结合使用，以全面管理内存。
另一个细微问题是重计算可能引入数值差异，尽管这些差异通常小到可以忽略。然而，在训练过程中，这些微小的差异可能会累积，导致与标准训练略有不同的优化轨迹。尽管存在这些挑战，梯度检查点仍然是管理深度学习中内存需求的一种可靠且广泛采用的方法。
未来方向
随着模型继续增长，研究人员正在探索更高效的检查点策略。一个活跃的研究方向是自动学习检查点放置，根据模型架构和硬件约束动态决定在哪里存储激活值。另一个方向是将检查点与低精度训练相结合，后者以数值精度换取更小的激活表示。此外，分层或混合方法正在开发中，其中检查点与诸如offloading或模型并行等技术结合，以扩展到越来越多的参数。
结论
梯度检查点通过允许在GPU内存有限的情况下训练更深、更大的模型，彻底改变了深度学习格局。它是一种简单而有效的技术，平衡了计算和内存使用，使研究人员和工程师能够突破硬件限制。随着模型规模的持续增长，梯度检查点将继续成为训练现代人工智能系统的基础工具。