Gradient-Aware Learnable Temporal Truncation for Back-Propagation Through Time in Spiking Neural Networks

Spiking Neural Networks (SNNs) have demonstrated promising data-processing capabilities with low energy consumption. Back-Propagation Through Time (BPTT) enables training based on Spatio-Temporal Back-Propagation (STBP) by propagating errors through temporal dynamics. Conventional BPTT aggregates gradients calculated at each time step and updates network parameters using the aggregated gradients, even though the status of SNNs varies across time steps. This shortcoming causes two problems: temporal mismatch in gradient updates and high memory usage during the backward pass. Conventional studies have suggested Truncated BPTT (T-BPTT) to overcome these problems. However, they employ a predetermined truncation strategy that fails to account for temporal variations in gradient behavior, resulting in poor learning performance. To address these problems, we investigate how temporal truncation affects the gradients used to train SNNs and find that effective truncation configurations consistently exhibit smaller gradient dispersion. Guided by this finding, we propose Learnable Truncated BPTT (LT-BPTT), which jointly adjusts the number and lengths of truncated temporal segments. LT-BPTT generates candidate temporal partitions, evaluates their gradient dispersion using a probe mini-batch, and selects the partition with the lowest gradient dispersion. Extensive experiments across diverse SNN architectures and datasets demonstrate that LT-BPTT achieves the highest classification accuracy in most evaluated settings while reducing memory consumption compared with BPTT and benchmarks to improve learning performance of SNNs.