Multi-bit spiking neural networks (SNNs) have recently become a heated research spot, pursuing energy-efficient and high-accurate AI. However, with more bits involved, the associated memory and computation demands escalate to the point where the performance improvements become disproportionate. Based on the insight that different layers demonstrate different importance and extra bits could be wasted and interfering, this paper presents an adaptive bit allocation strategy for direct-trained SNNs, achieving fine-grained layer-wise allocation of memory and computation resources. Thus, SNN's efficiency and accuracy can be improved. Specifically, we parametrize the temporal lengths and the bit widths of weights and spikes, and make them learnable and controllable through gradients. To address the challenges caused by changeable bit widths and temporal lengths, we propose the refined spiking neuron, which can handle different temporal lengths, enable the derivation of gradients for temporal lengths, and suit spike quantization better. In addition, we theoretically formulate the step-size mismatch problem of learnable bit widths, which may incur severe quantization errors to SNN, and accordingly propose the step-size renewal mechanism to alleviate this issue. Experiments on various datasets, including the static CIFAR and ImageNet datasets and the dynamic CIFAR-DVS, DVS-GESTURE, and SHD datasets, demonstrate that our methods can reduce the overall memory and computation cost while achieving higher accuracy. Particularly, our SEWResNet-34 can achieve a 2.69% accuracy gain and 4.16x lower bit budgets over the advanced baseline work on ImageNet. This work will be open-sourced.

翻译：多比特脉冲神经网络（SNNs）近年来成为研究热点，旨在实现高能效与高精度的人工智能。然而，随着比特数的增加，相关的内存与计算需求急剧上升，导致性能提升与资源消耗不成比例。基于不同层具有不同重要性、额外比特可能被浪费或产生干扰的洞察，本文提出一种用于直接训练SNNs的自适应比特分配策略，实现对内存和计算资源的细粒度逐层分配，从而提升SNN的效率与精度。具体而言，我们将权重和脉冲的时间长度与比特宽度参数化，使其可通过梯度进行学习和控制。为应对可变比特宽度与时间长度带来的挑战，我们提出了改进的脉冲神经元，该神经元能够处理不同的时间长度、支持时间长度梯度的推导，并更好地适应脉冲量化。此外，我们从理论上分析了可学习比特宽度导致的步长失配问题，该问题可能给SNN带来严重的量化误差，并据此提出了步长更新机制以缓解此问题。在多个数据集上的实验，包括静态的CIFAR和ImageNet数据集，以及动态的CIFAR-DVS、DVS-GESTURE和SHD数据集，均表明我们的方法能够在降低整体内存与计算成本的同时实现更高的精度。特别地，我们的SEWResNet-34在ImageNet上相比先进基线工作，实现了2.69%的精度提升，并降低了4.16倍的比特预算。本工作将开源。