Gradient checkpointing is a technique used in Deep learning to reduce the memory footprint of training neural networks. During standard backpropagation, a network must store all intermediate activations computed in the forward pass to calculate gradients in the backward pass. For very deep models, such as large language models and transformers, this storage can exceed the memory capacity of available hardware. Gradient checkpointing addresses this by not storing every activation; instead, it keeps only a subset and recomputes the discarded ones on demand during the backward pass. This trades increased computational cost for significantly lower memory usage, enabling the training of larger models or the use of larger batch sizes on the same hardware.
The technique was introduced in 2016 by researchers at Carnegie Mellon University and OpenAI in a paper titled "Training Deep Nets with Sublinear Memory Cost." The authors, including Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin, demonstrated that by storing activations only at certain checkpoints (for example, every few layers) and recomputing the rest, the memory cost of training a deep network could be reduced from O(n) to O(sqrt(n)) for a network with n layers, at the cost of roughly one extra forward pass. This foundational work has since become a standard tool in the machine learning community, particularly as model sizes have grown dramatically.
How Standard Backpropagation Uses Memory
In a conventional training loop, the forward pass computes activations for every layer of the network. These activations are stored in memory because the backward pass needs them to compute gradients via the chain rule. For a network with L layers, this requires storing L sets of activations, each of which can be large. For example, a Residual Network (ResNet) with hundreds of layers or a transformer with dozens of attention blocks can accumulate gigabytes of activation data for a single training example. When training with large batch sizes, the memory requirement scales linearly with the batch size, often making it the primary bottleneck.
The Checkpointing Strategy
Gradient checkpointing divides the network into segments, with a checkpoint at the boundary of each segment. During the forward pass, only the activations at these checkpoint boundaries are saved to memory. All other intermediate activations within a segment are discarded. When the backward pass reaches a segment, it recomputes the forward pass for that segment using the saved checkpoint activation, regenerating the intermediate activations needed for gradient calculation. This recomputation adds computational overhead, typically equivalent to one extra forward pass per training step, but it dramatically reduces peak memory usage.
The choice of checkpoint placement is a trade-off. More checkpoints mean less recomputation but higher memory usage; fewer checkpoints mean lower memory but more computation. The optimal number of checkpoints for a network with n layers is approximately sqrt(n), which balances memory and compute. In practice, frameworks like PyTorch and TensorFlow allow users to specify checkpoint intervals or use automatic heuristics.
Variants and Improvements
Several refinements to the original technique have been developed. One common variant is selective checkpointing, where only certain layer types (such as attention blocks or convolutional layers) are checkpointed, while others are stored normally. Another approach, called memory-efficient gradient checkpointing, uses a more sophisticated schedule that stores activations at multiple levels of granularity, further reducing memory at the cost of additional recomputation. Some frameworks also implement "offloading," where checkpoints are moved to CPU memory or disk, though this introduces data transfer overhead.
In the context of Transformer (architecture) models, gradient checkpointing is often combined with other memory-saving techniques like Mixed Precision Training and Gradient Clipping. For instance, training a model like GPT-3, which has 175 billion parameters, would be impossible without such optimizations. The technique is also used in fine-tuning large models, where the memory savings allow practitioners to run on a single GPU instead of a cluster.
Practical Implementation
In modern deep learning frameworks, gradient checkpointing is typically exposed as a simple API. In PyTorch, for example, the torch.utils.checkpoint module provides a checkpoint function that wraps a module or a sequence of operations. When the wrapped module is executed, its activations are not saved; instead, they are recomputed during the backward pass. TensorFlow offers a similar feature through tf.recompute_grad. These implementations handle the bookkeeping automatically, making it easy for researchers to adopt the technique without modifying their model architecture.
The computational overhead of gradient checkpointing is not negligible. For a network with sqrt(n) checkpoints, the total forward computation during training increases by roughly 30-40% compared to standard training. However, this cost is often acceptable because the alternative - reducing batch size or model size - can hurt convergence or model quality. In many cases, the speedup from using a larger batch size outweighs the recomputation overhead.
Impact on Large Model Training
Gradient checkpointing has become a cornerstone of training very large models. Companies like OpenAI, Anthropic, and Google DeepMind rely on it to train models with hundreds of billions of parameters. For example, training a 70-billion-parameter model on a single node with 8 GPUs would require storing activations that exceed the combined memory of those GPUs without checkpointing. By using gradient checkpointing, these organizations can fit the training job into available hardware, albeit with longer training times.
The technique is also essential for Generative AI applications that involve long sequences, such as document summarization or code generation. In these cases, the activation memory grows with the sequence length, and checkpointing allows for longer contexts without exceeding memory limits. This has directly enabled the development of models with context windows of 100,000 tokens or more.
Relationship to Other Memory Optimizations
Gradient checkpointing is often used alongside other techniques. Batch Normalization and Layer Normalization do not directly reduce memory, but they can improve training stability, which complements checkpointing. Model Pruning reduces the number of parameters, but activations remain a bottleneck, so checkpointing is still needed. Data Augmentation increases the effective dataset size but does not affect activation memory. In distributed training, gradient checkpointing can be combined with pipeline parallelism, where different layers are assigned to different devices, to further reduce per-device memory pressure.
One notable alternative is Gradient Accumulation, which simulates a larger batch size by accumulating gradients over multiple smaller batches. This reduces memory for optimizer states but does not reduce activation memory, so it is not a substitute for checkpointing. Another related idea is reversible layers, as used in some Residual Network (ResNet) variants, where activations can be reconstructed from the output, but this requires architectural changes and is less general than checkpointing.
Limitations and Trade-offs
The primary limitation of gradient checkpointing is the increased wall-clock time per training step. For models that are already compute-bound, the extra forward pass can slow training by 20-40%. In addition, the technique does not reduce memory for the model parameters or optimizer states, which can also be substantial for large models. For extremely large models, practitioners may need to combine checkpointing with parameter offloading or use specialized hardware like AWS Trainium or Groq that have larger on-chip memory.
Another subtle issue is that recomputation can introduce numerical differences, though these are usually negligible in practice. The technique also requires careful implementation to avoid recomputing the same activations multiple times, which would increase overhead. Despite these challenges, gradient checkpointing remains a widely used and reliable method.
Future Directions
As models continue to grow, researchers are exploring more efficient checkpointing strategies. Some recent work uses learned heuristics to decide which activations to store, based on the model's structure and the hardware's memory profile. Others are investigating checkpointing at the level of individual operations rather than layers, allowing finer-grained control. There is also interest in combining checkpointing with low-precision-training to further reduce memory, though this introduces accuracy trade-offs.
The technique is likely to remain relevant as long as memory is a constraint in deep learning. With the rise of large language models and multi-head attention architectures, the demand for memory-efficient training methods will only increase. Gradient checkpointing, along with other optimizations, will continue to be a key enabler of progress in Artificial intelligence.
Conclusion
Gradient checkpointing is a simple yet powerful idea: by selectively storing activations and recomputing them during backpropagation, it allows deep networks to be trained with far less memory. Introduced in 2016, it has become a standard practice in the field, enabling the training of models that would otherwise be impossible on available hardware. While it adds computational overhead, the trade-off is often worthwhile, especially for large-scale models. As the field moves toward even larger architectures, gradient checkpointing will remain an essential tool in the machine learning toolbox.