针对稀疏循环网络设计高效梯度检查点,大幅延长训练序列长度
Optimal Gradient Checkpointing for Sparse and Recurrent Architectures using Off-Chip Memory
- 提出双检查点策略,优化局部内存使用
- 可在仅轻微延迟下训练超10倍长序列或4倍大网络
- 适合在异构硬件上训练稀疏神经网络的研究者
循环神经网络(RNN)因其计算高效和低内存需求,在长序列任务中备受青睐,但训练时需高内存-处理器带宽。检查点技术通过仅存储部分中间状态降低内存占用,但因重计算阶段的计算开销而很少被采用。本文针对稀疏RNN与脉冲神经网络(SNN)提出高效的梯度检查点策略。利用智能处理单元(IPU)的分布式本地内存特性,适配稀疏不规则计算负载,实现长序列下SNN的可扩展训练。研究发现,双检查点方法最优:在最小化重计算开销的同时,最大化本地内存利用率。该方法减少对慢速全局内存的依赖,使训练序列长度突破此前10倍以上,或支持4倍规模的网络,且时间开销几乎可忽略。这些技术显著提升了稀疏与循环网络在多样化硬件平台上的可扩展性与效率,凸显稀疏激活对大规模递归网络训练的关键优势。
原文摘要 · Abstract (English)
Recurrent neural networks (RNNs) are valued for their computational efficiency and reduced memory requirements on tasks involving long sequence lengths but require high memory-processor bandwidth to train. Checkpointing techniques can reduce the memory requirements by only storing a subset of intermediate states, the checkpoints, but are still rarely used due to the computational overhead of the additional recomputation phase. This work addresses these challenges by introducing memory-efficient gradient checkpointing strategies tailored for the general class of sparse RNNs and Spiking Neural Networks (SNNs). SNNs are energy efficient alternatives to RNNs thanks to their local, event-driven operation and potential neuromorphic implementation. We use the Intelligence Processing Unit (IPU) as an exemplary platform for architectures with distributed local memory. We exploit its suitability for sparse and irregular workloads to scale SNN training on long sequence lengths. We find that Double Checkpointing emerges as the most effective method, optimizing the use of local memory resources while minimizing recomputation overhead. This approach reduces dependency on slower large-scale memory access, enabling training on sequences over 10 times longer or 4 times larger networks than previously feasible, with only marginal time overhead. The presented techniques demonstrate significant potential to enhance scalability and efficiency in training sparse and recurrent networks across diverse hardware platforms, and highlights the benefits of sparse activations for scalable recurrent neural network training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。