Memory-Efficient Exact Backpropagation
Training deep neural networks is constrained by the memory footprint of the computational graphduring automatic differentiation. This memory requirement scales linearly with model depth andsequence length, while the final projections contribute significantly to the spatial lower bound.Existing optimization techniques, such as activation checkpointing, fail to circumvent the memorybottlenecks caused by global graph construction and high-dimensional logits. We introduce abackpropagation algorithm that computes exact gradients without maintaining a global computationalgraph. Our approach limits peak memory usage to just a single model layer at a time. It achieves thisby saving only the boundaries between layers during the forward pass, calculating the final loss insmaller chunks, and recomputing the math for each layer one by one during the backward pass. Thisformulation analytically reduces overall spatial complexity from O(LBT (HT + Dinter) + BT V ) toO(LBT D + BT (HT + Dinter) + T V ). Empirical evaluation on a 60-million parameter Transformertrained on the Ultra-FineWeb dataset demonstrates a 62.9% reduction in peak memory (from 6.45 GBto 2.39 GB), identical loss trajectories and a negligible 0.9% reduction in training throughput. Thesefindings indicate that method is maximizing hardware utilization without compromising optimizationexactness.The code is available at https://github.com/VladimerKhasia/localprop